feat: add HND KV layout support#449
Open
hlin99 wants to merge 6 commits into
Open
Conversation
Signed-off-by: Tony Lin <tony.lin@intel.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds support for an additional paged KV-cache layout (“HND”, implemented as a permuted view) across the XPU attention stack, enabling stride-based layout auto-detection while keeping existing behavior unchanged for the current “NHD” layout.
Changes:
- Add HND/NHD layout auto-detection and layout-aware stride/shape handling in Xe2 chunk-prefill and paged-decode paths.
- Update FlashAttention XPU entrypoints to derive heads/block_size correctly under HND and to pass layout-aware KV strides to kernels.
- Expand FlashAttention varlen tests to cover both KV layouts and construct HND caches via permutation.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/flash_attn/test_flash_attn_varlen_func.py | Adds HND/NHD parametrization and builds HND test tensors via permute; updates reference attention for both layouts. |
| csrc/xpu/attn/xe_2/paged_decode.hpp | Updates documentation/comments to reflect NHD/HND stride support and caller-provided stride fields. |
| csrc/xpu/attn/xe_2/paged_decode_xe2.cpp | Detects HND via strides and computes layout-aware K/V strides for paged decode. |
| csrc/xpu/attn/xe_2/fmha_xe2.cpp | Detects HND and sets layout-aware KV strides for chunk prefill kernels. |
| csrc/xpu/attn/xe_2/collective/chunk_prefill_mainloop.hpp | Fixes tile-index computation to use page-stride tiles when clamping to the last page. |
| csrc/xpu/attn/paged_kv_utils.h | Adds HND layout detection and makes page-stride normalization/layout handling HND-aware. |
| csrc/flash_attn/flash_api.cpp | Makes paged-decode shape extraction layout-aware (num_heads_kv, block_size) for HND. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: Tony Lin <tony.lin@intel.com>
Signed-off-by: Tony Lin <tony.lin@intel.com>
Signed-off-by: Tony Lin <tony.lin@intel.com>
Collaborator
|
hi, we already have one PR WIP for this feature #382. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation:
HND KV layout is required by PD disaggregation scenario
This PR adds HND layout and keeping this other layout unchanged.