[ATTN] support HND format

[YizhouZ]

No description provided.

[Copilot]

Pull request overview

This PR updates the Xe2 FlashAttention paged-KV implementation to support an HND-style KV cache layout where the paged KV tensors are shaped as [num_blocks, num_heads_kv, block_size, head_size] (instead of the previous block-major layout), and adjusts kernel indexing/strides and tests accordingly.

Changes:

• Update paged KV cache shape assumptions (sizes/strides) across Xe2 prefill and decode paths.
• Fix paged indexing math to account for the new layout’s per-head physical extent in the flattened K/V “seq” dimension (including interleaved KV).
• Update FlashAttention API-side shape parsing and Python tests/reference implementation to match the new KV layout.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/flash_attn/test_flash_attn_varlen_func.py	Updates test KV allocations, interleaved KV slicing/stride assertions, and the reference paged-attn gather logic for [blocks, heads, block, dim].
csrc/xpu/attn/xe_2/paged_decode.hpp	Threads the KV-head count into decode kernel mainloop args for correct paged indexing under the new KV layout.
csrc/xpu/attn/xe_2/paged_decode_xe2.cpp	Updates paged KV shape/stride interpretation and inflates total_seqlen_kv bounds to match new paged indexing.
csrc/xpu/attn/xe_2/fmha_xe2.cpp	Updates paged KV shape/stride interpretation and passes KV-head count for paged indexing in prefill.
csrc/xpu/attn/xe_2/fmha_utils.hpp	Fixes interleaved-KV detection to use the correct stride axis for the new paged KV layout.
csrc/xpu/attn/xe_2/collective/chunk_prefill_mainloop.hpp	Extends mainloop params with num_heads_kv and updates paged tile index math for HND paged layout (prefill + decode mainloops).
csrc/xpu/attn/xe_2/chunk_prefill.hpp	Adds num_heads_kv_for_paged_idx plumbing into the prefill kernel launcher args.
csrc/flash_attn/flash_api.cpp	Updates paged-KV shape parsing (num_heads_kv, block_size) to match [blocks, heads, block, dim].

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jikunshang]

seems some MLA ut failed. cc @baodii PTAL as well.

[zhenwei-intel]

@YizhouZ @zufangzhu

I have already verified this on vLLM based on Yizhou's branch, and then submitted some changes to this branch. Could you please help review?

cc @jikunshang @rogerxfeng8

[zufangzhu]

@zhenwei-intel we don't need cache change. The latest cache kernel suppose already support HND. Please help to check. On the vllm side, simply change set_kv_cache_layout(“NHD”) to “HND” in vllm/platforms/xpu.py,

[zhenwei-intel]

@zhenwei-intel we don't need cache change. The latest cache kernel suppose already support HND. Please help to check. On the vllm side, simply change set_kv_cache_layout(“NHD”) to “HND” in vllm/platforms/xpu.py,

vllm-project/vllm#47975 updated the layout.
w/o this change, will produces following errors:

reshape_and_cache_flash(
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] File "/opt/venv/lib/python3.12/site-packages/vllm/_custom_ops.py", line 2453, in reshape_and_cache_flash
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] torch.ops._C_cache_ops.reshape_and_cache_flash(
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] File "/opt/venv/lib/python3.12/site-packages/torch/_ops.py", line 1275, in call
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] return self._op(*args, **kwargs)
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] RuntimeError: new layout requires value_cache to be adjacent to key_cache
``

[zufangzhu]

@zhenwei-intel we don't need cache change. The latest cache kernel suppose already support HND. Please help to check. On the vllm side, simply change set_kv_cache_layout(“NHD”) to “HND” in vllm/platforms/xpu.py,

vllm-project/vllm#47975 updated the layout. w/o this change, will produces following errors:

reshape_and_cache_flash(
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] File "/opt/venv/lib/python3.12/site-packages/vllm/_custom_ops.py", line 2453, in reshape_and_cache_flash
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] torch.ops._C_cache_ops.reshape_and_cache_flash(
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] File "/opt/venv/lib/python3.12/site-packages/torch/_ops.py", line 1275, in call
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] return self._op(*args, **kwargs)
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] RuntimeError: new layout requires value_cache to be adjacent to key_cache
``

I see, you need to rebase the latest vllm. In the latest vllm, cache used new layout. I have support this layout in #432.

[zhenwei-intel]

@zhenwei-intel we don't need cache change. The latest cache kernel suppose already support HND. Please help to check. On the vllm side, simply change set_kv_cache_layout(“NHD”) to “HND” in vllm/platforms/xpu.py,

vllm-project/vllm#47975 updated the layout. w/o this change, will produces following errors:

reshape_and_cache_flash(
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] File "/opt/venv/lib/python3.12/site-packages/vllm/_custom_ops.py", line 2453, in reshape_and_cache_flash
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] torch.ops._C_cache_ops.reshape_and_cache_flash(
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] File "/opt/venv/lib/python3.12/site-packages/torch/_ops.py", line 1275, in call
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] return self._op(*args, **kwargs)
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=26848) ERROR 07-16 04:30:55 [core.py:1231] RuntimeError: new layout requires value_cache to be adjacent to key_cache
``

I see, you need to rebase the latest vllm. In the latest vllm, cache used new layout. I have support this layout in #432.

great. thanks

now we can support HND layout vllm-project/vllm#47975 w/o this pr

[YizhouZ]

vllm-project/vllm#44455 makes shape fixed to NHD, only the inner stride would change. this PR is not longer needed since we already supported stride FMHA.