SM100 sparse prefill: optional per-head-group key mask (grouped prefill) - #1
Open
JaredforReal wants to merge 1 commit into
Open
JaredforReal wants to merge 1 commit into
JaredforReal wants to merge 1 commit into
Conversation
Add `head_group_mask` to `sparse_prefill_fwd` (SM100, h_q == 128 only): a uint8 [s_q, topk / 128, 128] tensor holding, per 128-key block, eight 16-byte masks. Bit j of the mask for head group g says whether key `indices[b*128+j]` may be attended by heads [g*hg, (g+1)*hg), hg = `head_group_size` (16 by default). Masked keys behave exactly like invalid indices (-inf before the row max), so the change only ANDs 64 extra bits into the existing `is_k_valid` mask per thread; the producer warp copies the 128 mask bytes of each K block into shared memory alongside `is_k_valid`. Purpose: "grouped" sparse prefill for DSA models served with few heads per rank (e.g. GLM-5.3 at TP4 = 16 heads). G consecutive query tokens x 16 heads are passed as 128 pseudo-heads sharing the union of their top-k index lists; adjacent tokens' top-k overlap heavily (union of 8 tokens ~2.4k of 16k on real prompts), so the kernel reads each KV row once per group instead of once per token. Kernel time on real 32k top-k (21k queries): 14.25 ms with a per-token kernel -> 2.86 ms grouped. Verified against an exact per-token reference (rel err 2e-3; without the mask 0.34). The new op arguments default to None / 16, so existing callers are unaffected; the small-topk and head-64 kernels reject the feature. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
There was a problem hiding this comment.
🟡 Changes recommended
The new head_group_mask behavior is not covered by the repository’s existing sparse prefill tests, and the new docstring should clarify the h_kv == 1 constraint to avoid misleading callers.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds an optional per-head-group key mask to the SM100 h_q==128 sparse prefill path to enable “grouped” sparse prefill (multiple tokens sharing a union top-k list while selectively masking keys per head-group), preserving default behavior when the mask is not provided.
Changes:
- Extend the Python and C++ APIs for
sparse_prefill_fwdto accept optionalhead_group_maskplushead_group_size. - Add shared-memory storage and masking logic in the SM100 head128 phase1 kernel to AND head-group mask bits into the existing key-valid bitmask.
- Update feature gating to advertise and require
HEAD_GROUP_MASKonly on the SM100 head128 implementation.
File summaries
| File | Description |
|---|---|
| flash_mla/flash_mla_interface.py | Exposes head_group_mask/head_group_size in the Python sparse prefill wrapper and documents the new behavior. |
| csrc/sm100/prefill/sparse/fwd/head128/phase1.cuh | Loads per-block head-group mask bytes and applies them during the existing -inf masking stage. |
| csrc/sm100/prefill/sparse/fwd/head128/config.h | Adds hg_mask shared-memory buffer (aligned) for per-head-group masking. |
| csrc/params.h | Extends SparseAttnFwdParams with head_group_mask pointer and log2_head_group_size. |
| csrc/api/sparse_fwd.h | Validates the new inputs, plumbs params, and gates dispatch via FwdFeatures::HEAD_GROUP_MASK. |
| csrc/api/api.cpp | Extends the Torch library schema for sparse_prefill_fwd with the new optional args and defaults. |
Review details
- Files reviewed: 6/6 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+215
to
+219
| head_group_mask: optional, [s_q, topk // 128, 128], uint8 (SM100, h_q == 128 only). "Grouped" sparse prefill: the | ||
| h_q heads of one q row are h_q / head_group_size consecutive tokens x head_group_size real heads sharing one | ||
| (union) index list; for key block b, bytes [16*g, 16*g+16) hold a 128-bit mask whose bit j says whether key | ||
| indices[b*128 + j] is attendable by head group g. Masked keys behave like invalid indices. | ||
| head_group_size: heads per group for head_group_mask (power of two, h_q // head_group_size <= 8). |
Comment on lines
228
to
231
| results = flash_mla_cuda.sparse_prefill_fwd( | ||
| q, kv, indices, sm_scale, d_v, attn_sink, topk_length, out | ||
| q, kv, indices, sm_scale, d_v, attn_sink, topk_length, out, | ||
| head_group_mask, head_group_size, | ||
| ) |
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.
Summary
Optional
head_group_maskforsparse_prefill_fwdon SM100 (h_q == 128): auint8 [s_q, topk/128, 128]tensor with eight 16-byte per-head-group masks per 128-key block. Bitjof groupg's mask says whether keyindices[b*128+j]may be attended by heads[g*hg, (g+1)*hg); masked keys behave exactly like invalid indices. DefaultNone→ no behaviour change for existing callers.Kernel change is minimal: the producer warp (warp 13) copies the 128 mask bytes of each K block into shared memory next to
is_k_valid, and each scale&exp thread ANDs the 64 bits of its head group intois_k_valid_lo/hibefore the existing-infmasking.alignas(16)on the new smem array;+256 Bsmem.Why
"Grouped" sparse prefill for DSA models served with few heads per rank (GLM-5.3 at TP4 = 16 heads): G consecutive query tokens × 16 heads are passed as the kernel's 128 heads sharing the union of their top-k lists (adjacent tokens' top-k overlap heavily), so each KV row is read once per group instead of once per token. On real 32k top-k (21k queries): per-token kernel 14.25 ms → grouped 2.86 ms. The consumer side is JaredforReal/vllm PR "Grouped sparse prefill with the FlashMLA head-group-mask kernel".
Tests
bench/hgmask_kernel_test.py(in the vLLM notes dir): 64 tokens × 16 heads grouped 8×, union index lists + masks vs exact per-token attention → max abs err 1.4e-3, rel 2.2e-3; without the mask 0.34.FLASH_MLA_SRC_DIR(SM100 only) and run end to end in vLLM (GLM-5.3-NVFP4, TP4) for prefill up to 64k tokens.Notes
Based on
main; therope_dimbranch does not touch these files. AI assistance (Claude Code) was used; the submitter reviewed the change.