[BUG] TRT-LLM Gen full attn. Incorrect result for `head_dim=256`

[vadiklyutiy]

If change head_dim from 128 to 256 here

flashinfer/tests/attention/test_trtllm_gen_attention.py

Line 645 in ef687e9

head_dim = 128

and run

pytest flashinfer/tests/attention/test_trtllm_gen_attention.py::test_trtllm_batch_decode

will see 756 failed tests.

FlashInfer version

uv pip show flashinfer-python
Name: flashinfer-python
Version: 0.4.1

Context: this head_dim from Qwen3-Next model.

[vadiklyutiy]

@Tom-Zheng made further investigation and figured out that bug appear only if head_dim=256 and page_size=16

[nvpohanh]

@PerkzZheng @vadiklyutiy Is this fixed by #2137 ?

[PerkzZheng]

@PerkzZheng @vadiklyutiy Is this fixed by #2137 ?

yes, feel free to give it a try and let us know if there are other issues.

[nvpohanh]

@vadiklyutiy could you try it and close it? thanks!