Skip to content

perf(dsv4): adapt tuned AITER ASM FP8 decode - #2411

Closed
yhl-amd wants to merge 5 commits into
ROCm:mainfrom
yhl-amd:perf/dsv4-aiter5195-asm-decode
Closed

yhl-amd wants to merge 5 commits into
ROCm:mainfrom
yhl-amd:perf/dsv4-aiter5195-asm-decode

Conversation

@yhl-amd

@yhl-amd yhl-amd commented Sep 26, 2026 •

Copy link
Copy Markdown
Contributor

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_indptr buffer:

    • V4 metadata builder, verify/decode: one plan per layer type (SWA window, CSA window + index_topk, HCA window + max_model_len // 128), sized for the padded row count.
    • V4 metadata builder, MTP draft (prepare_mtp_decode): SWA plan for the draft rows.
    • DSpark draft (DSparkIndexBuffers): built with the draft CSR at stage 0 and read by every stage (B*T rows, W + T KV).

    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_decode takes a single split_plan; the model selects it in the same ratio branch as kv_indices / kv_indptr.

  • v4_decode_split_plan is the only place that checks for AITER's helper. With an AITER build that lacks it, plans are None and 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

Test plan

🤖 Generated with Claude Code

@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2411 --add-label <label>

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>
@yhl-amd
yhl-amd marked this pull request as ready for review September 28, 2026 04:37
yhl-amd and others added 2 commits September 28, 2026 01:05
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>
@yhl-amd

yhl-amd commented Oct 4, 2026

Copy link
Copy Markdown
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.

@yhl-amd yhl-amd closed this Oct 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants