add asym feature for chunk prefill and kv cache

[tianyuan0211]

Purpose

Add asymmetric attention support (different QK and V/O head dimensions) to the XPU kernel stack, covering chunk prefill and KV-cache reshape/copy paths.

This is required by models such as MiMo (MiniMax-M2.5) where the QK head dimension is 192 and the V/O head dimension is 128.

Specifically:

• fmha_utils.hpp – Add two new tile policies for asymmetric QK=192 / V=128:

  • chunk_policy_head192_vo128 (non-paged / block_size=32 path)
  • chunk_policy_head192_vo128_b16 (paged block_size=16 path)

• chunk_prefill.hpp – Add v_head_size field to chunk_prefill_args_t; kernel uses it to set the VO head dimension independently from QK. Falls back to head_size for symmetric models (zero value = backward-compatible default).

• fmha_xe2.cpp – Derive v_head_size from value_cache.size(-1) and pass it through to args; dispatch to the asymmetric policy when head_size == 192 && v_head_size == 128.

• flash_api.cpp – Allocate the MHA output tensor with v.size(-1) instead of q.size(-1) so the output shape is correct for asymmetric models.

• cache.cpp – Update both reshape_and_cache_kernel and reshape_and_cache_flash_kernel to handle different K and V head sizes and strides:

  • Split the single KV copy loop into separate loops for K (n_k = num_heads * head_size) and V (n_v = num_heads * v_head_size).
  • Add independent value_block_stride / value_page_stride parameters to the flash kernel so K-cache and V-cache strides are no longer required to be equal.
  • Fix a latent work-group oversubscription bug in the flash path: the work-group was sized to the raw element count but the kernel copies via vectorize_with_alignment<VEC_SIZE>, so only num_elems / VEC_SIZE threads are ever active. The new sizing is ceil(num_elems / VEC_SIZE).

• chunk_prefill_configure.cmake – Register the two new asym policies in the build configuration.

Test Plan

Run the existing unit-test suite plus the two new targeted tests:

# KV-cache asymmetric reshape tests (both vllm-style and flash-style)
pytest tests/test_cache.py::test_asym_reshape_and_cache -v
pytest tests/test_cache.py::test_asym_reshape_and_cache_flash -v

# Asymmetric QK/V attention (paged and non-paged, block_size 16 and 32)
pytest tests/flash_attn/test_flash_attn_varlen_func.py::test_varlen_with_asym_qk_vo -v

Test Result

All new unit tests pass on Intel Arc B60 (BMG) XPU:

Test	Result
test_asym_reshape_and_cache	PASS
test_asym_reshape_and_cache_flash	PASS
test_varlen_with_asym_qk_vo	PASS

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR adds support for asymmetric attention where the V/O head dimension may differ from the QK head dimension, updating output sizing, kernel argument plumbing, and KV-cache reshape/copy logic accordingly.

Changes:

• Plumb a new v_head_size parameter through chunk prefill / XE2 FMHA launch paths.
• Allocate MHA output tensors using v.size(-1) for asymmetric attention.
• Update KV-cache reshape/copy kernels (including Flash-style paged cache) to handle differing K vs V head sizes and strides.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File	Description
csrc/xpu/attn/xe_2/fmha_xe2.cpp	Derives and passes v_head_size from value_cache into chunk prefill impl args.
csrc/xpu/attn/xe_2/chunk_prefill.hpp	Adds v_head_size to args and uses it to set VO head size in kernel shape.
csrc/flash_attn/flash_api.cpp	Allocates output with last dim following V head dim instead of Q head dim.
csrc/cache.cpp	Updates reshape-and-cache kernels to copy keys/values with different head dims and supports distinct value-cache strides in flash path.

---

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

[xinyu-intel]

please add unit test

[tianyuan0211]

please add unit test

added

[tianyuan0211]

Description added, please help review, thank you! @xinyu-intel

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

csrc/cache.cpp:1014

• In the is_strided_head path, the dispatched reshape_and_cache_flash_strided_kernel assumes the value head dimension equals head_size (it indexes v_base + head_idx * head_size_ and writes V at offset + head_size_). With asymmetric attention (v_head_size != head_size), this will read past the end of value and/or corrupt the fused KV-cache layout. Add an explicit check (or extend the strided kernel) so this layout is rejected unless v_head_size == head_size.

  bool is_strided_head = (head_stride != head_size);
  if (is_strided_head) {
    TORCH_CHECK(
        reinterpret_cast<uint8_t*>(value_cache.data_ptr()) ==
            reinterpret_cast<uint8_t*>(key_cache.data_ptr()) +

[tianyuan0211]

Pre-commit checked & merged the latest commit to solve the unit-test fail. Could test again. @xinyu-intel

[tianyuan0211]

This PR only touches attention code, which is built into a separate .so that the MoE grouped-GEMM path doesn't appear to link, so the MoE binaries look identical to main. With a CI-matched toolchain on BMG, the failing e5m2 case passed 20/20 and a 20,000-run probe stayed deterministic and well within tolerance, so I wasn't able to reproduce the failure, though I can't fully rule out other factors. Could we re-trigger the BMG job and see if it reproduces? @xinyu-intel

[tianyuan0211]

Hi Xinyu, the previous unit-test fail comes from fused-moe, and my asym attention code doesn't touch that part. Please retrigger the test, thanks! @xinyu-intel