Qwen3.5 text generation attends over the full KV capacity every step, so decode speed drops with max_length

[mozophe]

Feature Idea

Qwen3.5 decoding gets slower as max_length grows, even when the answer itself is short, because every decode step runs attention over the whole preallocated KV cache instead of only the filled part.

Qwen3.5 uses FixedKVBias for its full-attention layers (Qwen35.init_kv_cache, and mtp_kv in Qwen35._generate_mtp). FixedKVBias.decode computes q @ key over the full capacity and masks the unused tail with a bias. In _generate_mtp the capacity is prompt + max_length + 7, so a 3,000-token answer with max_length=24000 pays for ~24,000 slots per layer per step. The plain generate path is affected too, since it also gets its cache from Qwen35.init_kv_cache.

llama.py already has a length-aware path: when init_kv_cache sets up a FixedKV cache, single-token decode calls comfy_kitchen.flash_attention_decode with the device-side seqlen. That kernel is FlashAttention 2's split-KV kernel with seqused_k = kv_lengths, so it only reads the filled part and stays CUDA-graph friendly. Qwen3.5 can't use it today because it is built for one case only (comfy_kitchen/backends/cuda/ops/flash_decode.cu):

1. head_dim 128 only (Flash_fwd_kernel_traits<128, …>, params.d = 128), while Qwen3.5's full-attention layers use 256, and
2. non-causal, one query token per sequence (the query rows carry the packed GQA groups), while MTP verification sends 1 + depth queries (3–6), each seeing the cache plus the drafts before it.

FlashAttention 2 already ships the split-KV variant this needs (flash_fwd_split_hdim256_bf16_causal_sm80.cu).

Request: extend flash_attention_decode to head_dim 256 and a small causal multi-query window (up to 6 queries, matching FixedKVBias.shared's rows = 6), and use it for Qwen3.5 in place of FixedKVBias. That would make decode cost follow the actual sequence length rather than max_length.

Measurements

Qwen-Image 2.1 prompt enhancer (Qwen3.5 9B, Comfy-Org int8 checkpoint), RTX 4090 Laptop GPU 16 GB, sampling with the official PE settings, one input image for the edit case. Answers are typically 2,000–4,000 tokens, so only the preallocated capacity changes between rows.

Mode	max_length	MTP off	MTP on
Text-to-image	16256	27 tok/s	38 tok/s
Text-to-image	8192	35 tok/s	47 tok/s
Editing	24000	21 tok/s	36 tok/s
Editing	8192	31 tok/s	52 tok/s

Lowering max_length from the PE defaults to 8192 is 24–48% faster, with MTP on or off. Users currently have to trade that speed against the risk of a truncated answer.

Environment: ComfyUI 6ceb5de6 (2026-09-22; FixedKVBias usage is unchanged on master as of 2026-09-23), comfy-kitchen 0.2.35, torch 2.12.0+cu132, Windows 11.

Existing Solutions

No existing issue or PR covers this (searched ComfyUI and comfy-kitchen).

FlashAttention's flash_attn_with_kvcache already covers this exact case (device-side cache_seqlens, several queries with a bottom-right-aligned causal mask, GQA, head dims up to 256), so it can serve as a reference for the kitchen kernel.

Other

Found while building ComfyUI-Qwen-Image-2.1-PromptEnhancer-MTP, which runs the PE through Qwen35._generate_mtp. Happy to test a branch.