[GH-ISSUE #274] flash_attention_decode_gqa rejects single-token decode (S == 1): output stride check fails after DLPack normalizes size-1 strides #73

New issue

Open

opened 2026-10-26 22:19:58 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-26 22:19:58 +00:00

Owner

Originally created by @ntsrigaud on GitHub (Oct 23, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/274

Environment

comfy-kitchen 0.2.39. The same check is present in 0.2.40.
ComfyUI v0.42.1.
PyTorch 2.9.0a0+50eac811a6.nv25.09 (NGC nvcr.io/nvidia/pytorch:25.09-py3), CUDA 13.0, Python 3.12.3.
GPU: NVIDIA GB10 (DGX Spark), compute capability 12.1.

Symptom

The FLUX 3 Klein prompt upsampler (qwen3.5_9b_flux3_prompt_upsampler_int8_convrot.safetensors, CLIPLoader type flux3) fails in ComfyUI's TextGenerate at the first decode step:

RuntimeError: Unsupported Flash Attention tensor strides
  comfy/text_encoders/qwen35.py:396   output = past_key_value.attend(xq, xk, xv)
  comfy/text_encoders/llama.py:80     return comfy_kitchen.flash_attention_decode_gqa(xq, self.key, self.value, self.seqlen)
  comfy_kitchen/flash_attention.py:205  _cuda_backend._C.flash_attention_decode_gqa(...)

Cause

What the wrapper allocates: flash_attention_decode_gqa allocates output = torch.empty((batch, query_length, heads * head_dim)). For one-token decode with B=1, H=16, D=256, that is shape (1, 1, 4096) with torch strides (4096, 4096, 1).
What the binding receives: after _wrap_for_dlpack, the CUDA binding sees strides (1, 1, 1), because PyTorch 2.9's DLPack export gives size-1 dimensions stride 1.
Why it fails: the binding then requires output.stride(1) == heads * kHeadDim and output.stride(0) == query_length * heads * kHeadDim (dlpack_bindings.cpp, flash_attention_decode_gqa), which a size-1 dimension no longer satisfies. q is affected the same way, (4096, 256, 256, 1) becoming (1, 256, 1, 1), but its checks still pass.
The fix: the stride checks (and any stride-based indexing) should ignore the strides of size-1 dimensions, or the kernel should read the strides it relies on only where the extent is greater than 1.

Other observations

--use-ck-attention: with this flag the cache becomes Int8FixedKVCache, and decode fails with Unsupported step merge strides, which looks like the same pattern.
w4a8 build: the w4a8 build of the same upsampler fails earlier, with w4a8_codebook_mma_linear requires contiguous tensors. It may be related; cf. #248.
Workaround: decoding over the same cache with torch.nn.functional.scaled_dot_product_attention gives correct text, and matches _decode_gqa_torch to within 4.9e-4 (one token) and 4.8e-7 (three tokens).

Steps to reproduce

Queue CLIPLoader (the upsampler above, type flux3) → TextGenerate (the official flux3_klein_int8_t2i.json system prompt; max_length 3584, sampling on, temperature 1, top_k 1, top_p 0.95, presence_penalty 1.5, seed 804; prompt "Appetizing bowl of rich Thailand tom yum soup on a plain white background") → PreviewAny. It fails at the first decode step.

Originally created by @ntsrigaud on GitHub (Oct 23, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/274 ### Environment - comfy-kitchen 0.2.39. The same check is present in 0.2.40. - ComfyUI v0.42.1. - PyTorch 2.9.0a0+50eac811a6.nv25.09 (NGC `nvcr.io/nvidia/pytorch:25.09-py3`), CUDA 13.0, Python 3.12.3. - GPU: NVIDIA GB10 (DGX Spark), compute capability 12.1. ### Symptom The FLUX 3 Klein prompt upsampler (`qwen3.5_9b_flux3_prompt_upsampler_int8_convrot.safetensors`, CLIPLoader type `flux3`) fails in ComfyUI's `TextGenerate` at the first decode step: ``` RuntimeError: Unsupported Flash Attention tensor strides comfy/text_encoders/qwen35.py:396 output = past_key_value.attend(xq, xk, xv) comfy/text_encoders/llama.py:80 return comfy_kitchen.flash_attention_decode_gqa(xq, self.key, self.value, self.seqlen) comfy_kitchen/flash_attention.py:205 _cuda_backend._C.flash_attention_decode_gqa(...) ``` ### Cause - **What the wrapper allocates:** `flash_attention_decode_gqa` allocates `output = torch.empty((batch, query_length, heads * head_dim))`. For one-token decode with B=1, H=16, D=256, that is shape `(1, 1, 4096)` with torch strides `(4096, 4096, 1)`. - **What the binding receives:** after `_wrap_for_dlpack`, the CUDA binding sees strides `(1, 1, 1)`, because PyTorch 2.9's DLPack export gives size-1 dimensions stride 1. - **Why it fails:** the binding then requires `output.stride(1) == heads * kHeadDim` and `output.stride(0) == query_length * heads * kHeadDim` (`dlpack_bindings.cpp`, `flash_attention_decode_gqa`), which a size-1 dimension no longer satisfies. q is affected the same way, `(4096, 256, 256, 1)` becoming `(1, 256, 1, 1)`, but its checks still pass. - **The fix:** the stride checks (and any stride-based indexing) should ignore the strides of size-1 dimensions, or the kernel should read the strides it relies on only where the extent is greater than 1. ### Other observations - **`--use-ck-attention`:** with this flag the cache becomes `Int8FixedKVCache`, and decode fails with `Unsupported step merge strides`, which looks like the same pattern. - **w4a8 build:** the w4a8 build of the same upsampler fails earlier, with `w4a8_codebook_mma_linear requires contiguous tensors`. It may be related; cf. #248. - **Workaround:** decoding over the same cache with `torch.nn.functional.scaled_dot_product_attention` gives correct text, and matches `_decode_gqa_torch` to within 4.9e-4 (one token) and 4.8e-7 (three tokens). ### Steps to reproduce Queue `CLIPLoader` (the upsampler above, type `flux3`) → `TextGenerate` (the official `flux3_klein_int8_t2i.json` system prompt; `max_length` 3584, sampling on, temperature 1, top_k 1, top_p 0.95, presence_penalty 1.5, seed 804; prompt "Appetizing bowl of rich Thailand tom yum soup on a plain white background") → `PreviewAny`. It fails at the first decode step.

bsakaguchi commented

2026-10-26 22:19:59 +00:00

Author

Owner

@potoior commented on GitHub (Oct 26, 2026):

I'd like to take this — planning to relax the binding's stride assertions for size-1 dimensions (whose strides are irrelevant to addressing), add a regression test driving the S == 1 decode path with synthetic tensors, and verify against the reference on real hardware.

@potoior commented on GitHub (Oct 26, 2026): I'd like to take this — planning to relax the binding's stride assertions for size-1 dimensions (whose strides are irrelevant to addressing), add a regression test driving the S == 1 decode path with synthetic tensors, and verify against the reference on real hardware.

bsakaguchi referenced this issue

2026-10-26 22:20:35 +00:00

[PR #73] [MERGED] Faster turing int4 and int8. #121