Add HND paged KV cache layout support via stride-based auto-detection

[Copilot]

XPU attention kernels only supported NHD [num_blocks, block_size, num_heads, head_size] paged KV cache layout. This adds HND [num_blocks, num_heads, block_size, head_size] support — preferred for decode-heavy workloads where all heads at a single token position are contiguous.

Design

Follows vLLM CUDA's approach: no kv_cache_layout parameter added to any C++ function signature. HND is created by permuting NHD via .permute(0, 2, 1, 3); the resulting strides encode the layout and are read at runtime:

// NHD contiguous:  stride(1) = num_heads*head_size = size(2)*stride(2)  → false
// HND (permuted):  stride(1) = head_size           ≠ size(2)*stride(2)  → true
inline bool is_paged_kv_hnd_layout(const at::Tensor& key_cache) {
  return key_cache.stride(1) != key_cache.size(2) * key_cache.stride(2);
}

This is robust for cross-layer (non-contiguous stride(0)) tensors since only the intra-block strides (1) and (2) matter for layout detection.

C++ changes

• csrc/xpu/attn/paged_kv_utils.h — new is_paged_kv_hnd_layout() helper; get_paged_kv_cache_page_stride_elements() and check_paged_kv_cache_strides() now pick stride(1) vs stride(2) as seq stride based on layout
• csrc/xpu/attn/xe_2/fmha_xe2.cpp — layout-aware block_size/num_heads_kv extraction in is_paged block; layout-aware K/V stride assignment in both varlen and non-varlen paged branches; HND detection hoisted and reused
• csrc/xpu/attn/xe_2/paged_decode_xe2.cpp — same pattern; stride variables precomputed before args struct initialization; HND detection hoisted and reused
• csrc/xpu/attn/xe_2/collective/chunk_prefill_mainloop.hpp — paged decode fallback now uses page_stride_tiles so cross-layer page strides land on the correct final tile
• csrc/flash_attn/flash_api.cpp — layout-aware num_heads_kv/block_size extraction in both mix_batch and pure decode paths
• csrc/xpu/attn/xe_2/paged_decode.hpp — updated KV strides struct comment

Test changes

• ref_paged_attn gains a kv_layout parameter; HND gathering uses .transpose(1,2).reshape(...) after block index gather
• All five paged KV tests parametrized with kv_layout=["NHD","HND"]: test_varlen_with_paged_kv, test_varlen_with_interleaved_paged_kv, test_decode_with_paged_kv, test_varlen_with_cross_layer_paged_kv, test_decode_with_cross_layer_paged_kv
• HND tensors constructed via .permute(0,2,1,3) to match the vLLM allocation pattern; cross-layer HND tests assert stride(0) is preserved after permute
• FP8+HND combinations skipped (.to(fp8_dtype) produces a new contiguous tensor, discarding permuted strides)

[hlin99]

Code Review: HND paged KV cache layout support

Context

This PR adds HND [num_blocks, num_heads, block_size, head_size] layout support to the XPU paged KV cache kernels, alongside the existing NHD [num_blocks, block_size, num_heads, head_size] layout.

Design principle: No explicit kv_cache_layout parameter is added to any C++ function signature. Instead, layout is auto-detected at runtime from tensor strides, following the same approach as vLLM's CUDA reshape_and_cache_flash_kernel:

// NHD contiguous:  stride(1) = num_heads * head_size = size(2) * stride(2)  → false (NHD)
// HND (permuted):  stride(1) = head_size             ≠ size(2) * stride(2)  → true  (HND)
inline bool is_paged_kv_hnd_layout(const at::Tensor& key_cache) {
  return key_cache.stride(1) != key_cache.size(2) * key_cache.stride(2);
}

HND tensors are constructed by permuting an NHD tensor: .permute(0, 2, 1, 3). The resulting strides encode the layout without any extra bookkeeping.

---

Files changed

File	What changed
csrc/xpu/attn/paged_kv_utils.h	New is_paged_kv_hnd_layout() helper; get_paged_kv_cache_page_stride_elements() and check_paged_kv_cache_strides() pick stride(1) vs stride(2) as seq stride based on layout
csrc/xpu/attn/xe_2/fmha_xe2.cpp	Layout-aware block_size/num_heads_kv extraction in is_paged block; layout-aware K/V stride assignment in both varlen and non-varlen paged branches
csrc/xpu/attn/xe_2/paged_decode_xe2.cpp	Same pattern; stride variables precomputed before args struct initialization
csrc/flash_attn/flash_api.cpp	Layout-aware num_heads_kv/block_size extraction in both mix_batch and pure decode paths
csrc/xpu/attn/xe_2/paged_decode.hpp	Updated KV strides struct comment to document both layouts
tests/flash_attn/test_flash_attn_varlen_func.py	ref_paged_attn gains kv_layout param; all 5 paged KV tests parametrized with kv_layout=["NHD","HND"]

---

Review focus areas

Please review with these questions in mind:

1. Correctness of layout detection

• Is stride(1) != size(2) * stride(2) a sufficient and correct condition to distinguish NHD from HND?
• Does it hold correctly for cross-layer (non-contiguous) tensors where stride(0) is larger than expected?
• Any edge cases where this detection could produce a false positive or false negative?

2. Correctness of dimension extraction

In the is_paged block:

bool is_hnd = is_paged_kv_hnd_layout(key_cache);
if (is_hnd) {
    num_heads_kv = key_cache.size(1);  // HND: dim1=num_heads
    block_size   = key_cache.size(2);  // HND: dim2=block_size
} else {
    block_size   = key_cache.size(1);  // NHD: dim1=block_size
    num_heads_kv = key_cache.size(2);  // NHD: dim2=num_heads
}

• Are size(1) and size(2) assigned to the correct variables for each layout?

3. Correctness of stride assignment

// HND
args.k_stride_seq   = key_cache.stride(2);  // seq dim = dim2
args.k_stride_heads = key_cache.stride(1);  // head dim = dim1

// NHD
args.k_stride_seq   = key_cache.stride(1);  // seq dim = dim1
args.k_stride_heads = key_cache.stride(2);  // head dim = dim2

• Are k_stride_seq and k_stride_heads correctly swapped between NHD and HND?
• Does this match what the downstream kernel (chunk_prefill_mainloop.hpp, paged_decode.hpp) expects for k_stride_seq and k_stride_heads?

4. paged_kv_utils.h utility functions

• get_paged_kv_cache_page_stride_elements: NHD uses stride(0)/stride(1), HND uses stride(0)/stride(2). Is the seq stride picked correctly for each layout?
• check_paged_kv_cache_strides: does the divisibility check k_physical_page_stride % k_stride_seq == 0 still hold for HND permuted tensors?

5. Missing fix in chunk_prefill_mainloop.hpp ⚠️

This file is not included in the PR but contains a related bug.

In DecodeFwdMainloop::operator(), line ~1004:

// current code (bug):
next_tile_idx = params.max_pages_per_seq * tiles_per_page - 1;

// correct code (as fixed in vllm-project/vllm-xpu-kernels#382):
next_tile_idx = params.max_pages_per_seq * page_stride_tiles - 1;

page_stride_tiles is already computed just above:

int page_stride_tiles = params.page_stride_elements / get<1>(TileShapeQK{});

When page_stride_elements > page_size (cross-layer layout), page_stride_tiles > tiles_per_page. Using tiles_per_page here causes the "set to last page" fallback to land at the wrong tile index, which can produce incorrect attention output for cross-layer HND KV cache.

This PR fixes page_stride_elements to be layout-aware (correct for HND), but the downstream consumer of that value in the decode mainloop still has the tiles_per_page bug. Both fixes are needed for correctness under cross-layer + HND.

Should this file be patched in the same PR?

6. is_hnd computed twice in some functions

In fmha_xe2.cpp, is_paged_kv_hnd_layout(key_cache) is called once in the is_paged block (for dimension extraction) and again later in the stride assignment block. Same issue in paged_decode_xe2.cpp.

• Is this a problem (performance, consistency)?
• Should is_hnd be hoisted to a single variable computed once and reused?

7. Test construction correctness

HND tensors are constructed as:

key_cache = torch.randn(num_blocks, block_size, num_kv_heads, head_size).permute(0, 2, 1, 3)

• Does .permute(0, 2, 1, 3) correctly produce [num_blocks, num_kv_heads, block_size, head_size]?
• Is the resulting tensor non-contiguous? Does the kernel handle non-contiguous tensors correctly?

For cross-layer HND:

key_cache = key_cache_nhd.permute(0, 2, 1, 3)  # non-contiguous, large stride(0)

• Does stride(0) remain unchanged after permuting dims 1 and 2? What is its value?
• Should there be an assertion on key_cache.stride(0) after permute to validate the test setup?

8. ref_paged_attn HND gather correctness

# HND gather path
k = key_cache[block_indices].transpose(1, 2).reshape(-1, num_kv_heads, head_size)

• Does transpose(1, 2) followed by reshape correctly flatten [num_blocks_i, block_size, num_kv_heads, head_size] → [num_blocks_i * block_size, num_kv_heads, head_size]?
• Is reshape safe here on a non-contiguous tensor (should it be contiguous().reshape(...) or is .reshape() fine)?

9. Skipped combinations

The following combinations are skipped:

if kv_layout == "HND" and fp8_dtype is not None:
    pytest.skip(...)
if kv_layout == "HND" and q_dtype is not None:
    pytest.skip(...)

• Is it intentional to skip FP8 + HND? Is there a technical reason (e.g. .to(fp8_dtype) produces a new contiguous tensor, discarding permuted strides)?
• Should there be a TODO comment explaining why and when this combination will be supported?

---

Known issues to address before merge

1. chunk_prefill_mainloop.hpp line ~1004: tiles_per_page should be page_stride_tiles — this is a correctness bug for cross-layer + HND, and should be fixed in this PR
2. is_hnd computed twice in fmha_xe2.cpp and paged_decode_xe2.cpp — hoist to a single variable
3. Missing stride(0) assertion in cross-layer HND tests after permute
4. Draft PR — needs to be converted to ready for review after issues are resolved

[hlin99]

Code Review: HND paged KV cache layout support

Context

This PR adds HND [num_blocks, num_heads, block_size, head_size] layout support to the XPU paged KV cache kernels, alongside the existing NHD [num_blocks, block_size, num_heads, head_size] layout.

Design principle: No explicit kv_cache_layout parameter is added to any C++ function signature. Instead, layout is auto-detected at runtime from tensor strides, following the same approach as vLLM's CUDA reshape_and_cache_flash_kernel:

// NHD contiguous:  stride(1) = num_heads * head_size = size(2) * stride(2)  → false (NHD)
// HND (permuted):  stride(1) = head_size             ≠ size(2) * stride(2)  → true  (HND)
inline bool is_paged_kv_hnd_layout(const at::Tensor& key_cache) {
  return key_cache.stride(1) != key_cache.size(2) * key_cache.stride(2);
}

HND tensors are constructed by permuting an NHD tensor: .permute(0, 2, 1, 3). The resulting strides encode the layout without any extra bookkeeping.

Files changed

File What changed
csrc/xpu/attn/paged_kv_utils.h New is_paged_kv_hnd_layout() helper; get_paged_kv_cache_page_stride_elements() and check_paged_kv_cache_strides() pick stride(1) vs stride(2) as seq stride based on layout
csrc/xpu/attn/xe_2/fmha_xe2.cpp Layout-aware block_size/num_heads_kv extraction in is_paged block; layout-aware K/V stride assignment in both varlen and non-varlen paged branches
csrc/xpu/attn/xe_2/paged_decode_xe2.cpp Same pattern; stride variables precomputed before args struct initialization
csrc/flash_attn/flash_api.cpp Layout-aware num_heads_kv/block_size extraction in both mix_batch and pure decode paths
csrc/xpu/attn/xe_2/paged_decode.hpp Updated KV strides struct comment to document both layouts
tests/flash_attn/test_flash_attn_varlen_func.py ref_paged_attn gains kv_layout param; all 5 paged KV tests parametrized with kv_layout=["NHD","HND"]

Review focus areas

Please review with these questions in mind:

1. Correctness of layout detection

• Is stride(1) != size(2) * stride(2) a sufficient and correct condition to distinguish NHD from HND?
• Does it hold correctly for cross-layer (non-contiguous) tensors where stride(0) is larger than expected?
• Any edge cases where this detection could produce a false positive or false negative?

2. Correctness of dimension extraction

In the is_paged block:

bool is_hnd = is_paged_kv_hnd_layout(key_cache);
if (is_hnd) {
    num_heads_kv = key_cache.size(1);  // HND: dim1=num_heads
    block_size   = key_cache.size(2);  // HND: dim2=block_size
} else {
    block_size   = key_cache.size(1);  // NHD: dim1=block_size
    num_heads_kv = key_cache.size(2);  // NHD: dim2=num_heads
}

• Are size(1) and size(2) assigned to the correct variables for each layout?

3. Correctness of stride assignment

// HND
args.k_stride_seq   = key_cache.stride(2);  // seq dim = dim2
args.k_stride_heads = key_cache.stride(1);  // head dim = dim1

// NHD
args.k_stride_seq   = key_cache.stride(1);  // seq dim = dim1
args.k_stride_heads = key_cache.stride(2);  // head dim = dim2

• Are k_stride_seq and k_stride_heads correctly swapped between NHD and HND?
• Does this match what the downstream kernel (chunk_prefill_mainloop.hpp, paged_decode.hpp) expects for k_stride_seq and k_stride_heads?

4. paged_kv_utils.h utility functions

• get_paged_kv_cache_page_stride_elements: NHD uses stride(0)/stride(1), HND uses stride(0)/stride(2). Is the seq stride picked correctly for each layout?
• check_paged_kv_cache_strides: does the divisibility check k_physical_page_stride % k_stride_seq == 0 still hold for HND permuted tensors?

5. Missing fix in chunk_prefill_mainloop.hpp ⚠️

This file is not included in the PR but contains a related bug.

In DecodeFwdMainloop::operator(), line ~1004:

// current code (bug):
next_tile_idx = params.max_pages_per_seq * tiles_per_page - 1;

// correct code (as fixed in vllm-project/vllm-xpu-kernels#382):
next_tile_idx = params.max_pages_per_seq * page_stride_tiles - 1;

page_stride_tiles is already computed just above:

int page_stride_tiles = params.page_stride_elements / get<1>(TileShapeQK{});

When page_stride_elements > page_size (cross-layer layout), page_stride_tiles > tiles_per_page. Using tiles_per_page here causes the "set to last page" fallback to land at the wrong tile index, which can produce incorrect attention output for cross-layer HND KV cache.

This PR fixes page_stride_elements to be layout-aware (correct for HND), but the downstream consumer of that value in the decode mainloop still has the tiles_per_page bug. Both fixes are needed for correctness under cross-layer + HND.

Should this file be patched in the same PR?

6. is_hnd computed twice in some functions

In fmha_xe2.cpp, is_paged_kv_hnd_layout(key_cache) is called once in the is_paged block (for dimension extraction) and again later in the stride assignment block. Same issue in paged_decode_xe2.cpp.

• Is this a problem (performance, consistency)?
• Should is_hnd be hoisted to a single variable computed once and reused?

7. Test construction correctness

HND tensors are constructed as:

key_cache = torch.randn(num_blocks, block_size, num_kv_heads, head_size).permute(0, 2, 1, 3)

• Does .permute(0, 2, 1, 3) correctly produce [num_blocks, num_kv_heads, block_size, head_size]?
• Is the resulting tensor non-contiguous? Does the kernel handle non-contiguous tensors correctly?

For cross-layer HND:

key_cache = key_cache_nhd.permute(0, 2, 1, 3)  # non-contiguous, large stride(0)

• Does stride(0) remain unchanged after permuting dims 1 and 2? What is its value?
• Should there be an assertion on key_cache.stride(0) after permute to validate the test setup?

8. ref_paged_attn HND gather correctness

# HND gather path
k = key_cache[block_indices].transpose(1, 2).reshape(-1, num_kv_heads, head_size)

• Does transpose(1, 2) followed by reshape correctly flatten [num_blocks_i, block_size, num_kv_heads, head_size] → [num_blocks_i * block_size, num_kv_heads, head_size]?
• Is reshape safe here on a non-contiguous tensor (should it be contiguous().reshape(...) or is .reshape() fine)?

9. Skipped combinations

The following combinations are skipped:

if kv_layout == "HND" and fp8_dtype is not None:
    pytest.skip(...)
if kv_layout == "HND" and q_dtype is not None:
    pytest.skip(...)

• Is it intentional to skip FP8 + HND? Is there a technical reason (e.g. .to(fp8_dtype) produces a new contiguous tensor, discarding permuted strides)?
• Should there be a TODO comment explaining why and when this combination will be supported?

Known issues to address before merge

1. chunk_prefill_mainloop.hpp line ~1004: tiles_per_page should be page_stride_tiles — this is a correctness bug for cross-layer + HND, and should be fixed in this PR
2. is_hnd computed twice in fmha_xe2.cpp and paged_decode_xe2.cpp — hoist to a single variable
3. Missing stride(0) assertion in cross-layer HND tests after permute
4. Draft PR — needs to be converted to ready for review after issues are resolved

@copilot

[Copilot]

Code Review: HND paged KV cache layout support

Context

This PR adds HND [num_blocks, num_heads, block_size, head_size] layout support to the XPU paged KV cache kernels, alongside the ex...

Addressed the actionable items in d65f590: patched chunk_prefill_mainloop.hpp to use page_stride_tiles, hoisted HND layout detection in the paged paths, and added cross-layer HND stride(0) assertions. Final code review and CodeQL validation passed.