[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