Repository navigation
Conversation
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
3 of 32 tasks
AITER's V4 NM wrapper can now pick num_kv_splits from a latency model when told the per-token KV length to plan for (`kv_len_hint`). ATOM's decode kv_indices buffers are capacity-sized, so AITER cannot derive that length itself; the builder publishes the worst case per layer type and the decode call forwards it: - CSA: window + index_topk - HCA: window + max_model_len // 128 Both depend only on config, so the pick stays fixed across CUDA graph capture and replay. With an AITER build that predates `kv_len_hint` (detected from the wrapper signature), CSA keeps the batch split table and HCA keeps AITER's own auto pick. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Replace kv_len_hint with an explicit split plan. The V4 metadata builder calls AITER's get_mla_v4_nm_split_plan once per decode step for each sparse layer type and stores (num_kv_splits, split_indptr) on the attention metadata; the fp8 ASM decode forwards both to mla_decode_fwd_v4_nm. - Plans are sized for the padded row count and the longest kv_len of the layer type (SWA: window, CSA: window + index_topk, HCA: window + max_model_len // 128), so capture and replay agree. - split_indptr lives in persistent forward_vars buffers (per ubatch prefix) that the builder rewrites in place before each forward; graphs of different sizes share them. - MTP draft steps clear the SWA plan (it is sized for the verify rows) and fall back to AITER's own pick. - AITER builds without get_mla_v4_nm_split_plan keep the CSA batch table. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…de CSR - v4_decode_split_plan (paged_decode) is the single place that knows about aiter's get_mla_v4_nm_split_plan; it returns None on older aiter, which leaves the split pick to aiter. - Each owner of a decode CSR builds its plan next to it, in its own persistent split_indptr buffer: the V4 metadata builder (SWA/CSA/HCA for verify, SWA for the MTP draft) and DSpark's DSparkIndexBuffers (built with the CSR at stage 0, read by every stage). - sparse_attn_v4_paged_decode takes one split_plan argument; the model picks it in the same ratio branch as kv_indices / kv_indptr. - Drop the legacy CSA batch table and the query_group / kv_kind parameters it needed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
This was referenced Sep 29, 2026
Contributor
Author
|
Superseded by #2462, which now carries these commits (tuned ASM FP8 decode with caller-built split plans) together with the HCA persistent-kernel path, merged with current main and with the non-GPU unit tests fixed. Closing in favour of that PR. |
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
Adapt DeepSeek-V4 native FP8 decode to the swapped gfx950 ASM decode kernel from [ASM] [MLA v4 nm] Swap gfx950 asm decode kernel aiter#5195.
Pass AITER's KV split plan ([ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge aiter#6126) to the decode ASM kernel. Whoever builds a decode CSR builds its split plan next to it, in its own persistent
split_indptrbuffer:window, CSAwindow + index_topk, HCAwindow + max_model_len // 128), sized for the padded row count.prepare_mtp_decode): SWA plan for the draft rows.DSparkIndexBuffers): built with the draft CSR at stage 0 and read by every stage (B*Trows,W + TKV).Each plan depends only on the row count and config, so CUDA graph capture and replay agree. The owner rewrites the buffer in place before each forward.
sparse_attn_v4_paged_decodetakes a singlesplit_plan; the model selects it in the same ratio branch askv_indices/kv_indptr.v4_decode_split_planis the only place that checks for AITER's helper. With an AITER build that lacks it, plans areNoneand AITER keeps its own split pick.This PR is ASM-only. It adds or changes no FlyDSL/Triton decode implementation, route, environment variable, or kernel file.
Dependencies
get_mla_v4_nm_split_plan. Without it this PR leaves the split pick to AITER.Test plan
tests/test_v4_prefill_asm_decode_dispatch.py,tests/test_paged_attention_dispatch.py,tests/test_dspark.py: pass with and without the AITER helper--cudagraph-mode FULL, real draft acceptance, with [ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge aiter#6126: graph capture, correct outputs, DSpark acceptance unchanged vs. the stock image🤖 Generated with Claude Code