Repository navigation
[ASM] [MLA v4 nm] Swap gfx950 asm decode kernel - #5195
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
Co-authored-by: Cursor <cursoragent@cursor.com>
The kernel read 16 e8m0 scale bytes out of the packed Q row, but only 14 are valid: the row is 448B fp8 NoPE, 14B scale, then padding from byte 462. Byte 462/463 is the e8m0 NaN encoding when it happens to be 0xFF, so a served model -- which reuses the Q buffer and leaves junk in the padding -- produced NaN output. gsm8k on DeepSeek-V4-Pro dropped from 0.9522 to 0.0265. Synthetic tests could not see it: the quantizer builds the packed row from zeroed storage, so the padding is always 0x00, and `abs(a-b) > tol` is False for NaN, so an all-NaN output still scored zero bad rows. The KV pool shares the same row layout but its scale read was already bounded correctly; only Q was affected. Verified: op_tests/test_mla_v4_nm.py passes; the padding-poisoned sweep (Q/K/both x 0x00/0xFF/per-byte/random, 18 cases) is clean, as is an 80-config accuracy sweep repeated 4x with the padding filled with 0xFF. Kernel timing is unchanged (24 configs x 5 interleaved rounds, median -0.15%, Wilcoxon p=0.32). Co-authored-by: Cursor <cursoragent@cursor.com>
The packed Q/KV row is 448B fp8 NoPE, 14B e8m0 scale, then 50B of padding, and `_native_to_2buff_for_asm` builds it with `torch.zeros` -- so every existing test hands the kernel padding that is always 0x00. A served model reuses these buffers and the padding carries junk, which is how a kernel that read 16 scale bytes instead of 14 shipped: byte 462/463 at 0xFF is the e8m0 NaN encoding and the row came out NaN. Poisons the padding of Q, KV and both, with 0xFF and with an arbitrary fill, across the three shipped (gqa, q_seq_logical) entry points. The reference reads only the defined bytes, so its result must not move either. Checks isfinite before the tolerance compare: `nan > tol` is False, so an all-NaN result otherwise reports zero mismatching rows and passes. Co-authored-by: Cursor <cursoragent@cursor.com>
…nge and kv-tail regression tests Co-authored-by: Cursor <cursoragent@cursor.com>
The closure read gqa and the per-gqa tensors from the enclosing loop, so it would have used the last iteration's values had it outlived one. Taking them as an explicit argument removes the late binding. Co-authored-by: Cursor <cursoragent@cursor.com>
Same results as the kernel it replaces — identical cos_diff across an 80-config gqa x batch x kv sweep, and gsm8k unchanged end to end. Decode at gqa=16 and 32 gains 10-17%; at 64 and 128 no config is measurably slower. Co-authored-by: Cursor <cursoragent@cursor.com>
liyjiang
force-pushed
the
liyjiang/mla_v4_nm_perf27
branch
from
September 22, 2026 16:49
55eda3f to
0a3ed7f
Compare
amd-ruitang3
approved these changes
Sep 24, 2026
Merged
5 tasks
3 of 4 tasks
This was referenced Sep 27, 2026
This was referenced Sep 29, 2026
yhl-amd
added a commit
that referenced
this pull request
Oct 4, 2026
…l with in-kernel split and merge (#6126) * [MLA v4 nm] Cost-model num_kv_splits auto pick and kv_len_hint The gfx950 v4 nm auto split pick (ignore_total_kv=1) maximizes CU occupancy only, so it prefers whole multiples of cu_num: for gqa=128 single-token decode it picks 13-16 splits (several workgroup waves) for batches whose KV already fits one wave, paying extra partial writes and a wider stage-2 merge for no added parallelism. - On gfx950, pick num_kv_splits by minimizing a fitted latency model: workgroup waves x per-wave KV time, plus the partial-write/stage-2 cost of splitting. Other archs keep the existing pick. - Add kv_len_hint: callers that pass capacity-sized kv_page_indices (e.g. CUDA-graph decode) cannot be planned from total_kv / num_seqs, so they pass the per-seq KV length to plan for. The pick depends only on shapes and this int, so it stays fixed across graph capture and replay. - Fix the stage-2 merge for a caller-supplied split_indptr with one split per seq on a grid > 1: it took the in-place BF16 copy (keyed on split_indptr[-1] == num_seqs) although the kernel wrote FP32 partials, returning NaN. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Expose the split plan to callers instead of kv_len_hint Planning inside mla_decode_fwd_v4_nm needs a per-seq KV length, and the only one the wrapper can see, kv_page_indices.shape[0] / num_seqs, is meaningless for CUDA-graph callers whose index buffers are capacity-sized. - Add get_mla_v4_nm_split_plan(num_seqs, num_heads, kv_len, split_indptr=): returns the split count (cost model on gfx950, occupancy pick elsewhere) and fills a caller-owned persistent split_indptr in place. Callers build it in their metadata step and pass both to the wrapper. - Drop kv_len_hint; the wrapper's own auto pick is unchanged from main. - Tests: plan fill, accuracy through the wrapper, and CUDA-graph replay with different KV lengths under a captured plan. _run_varlen_point now returns its checkAllclose error so the new tests gate on it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Return the split plan as a NamedTuple, per-arch cost table - get_mla_v4_nm_split_plan returns MlaV4NmSplitPlan(num_kv_splits, split_indptr). The field names are the wrapper's kwargs, so callers pass a plan on as mla_decode_fwd_v4_nm(..., **plan._asdict()). - Latency-model coefficients live in a per-arch table; archs without an entry keep the occupancy-only pick. Only gfx950 is fitted. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Add a persistent decode kernel with in-kernel split and merge A persistent variant of the v4 nm fp8 decode kernel (gqa=128, one query token per row). Each call is one launch of grid (2P, 1, 1): the kernel plans the KV split over P partitions from kv_indptr, runs the attention and merges the split partials itself, so callers need no host split plan and no stage-2 merge, and the launch is CUDA-graph safe with varying KV lengths. - hsa/gfx950/mla_v4: mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co, symbol _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_nm_psE, registered as the ps=1 row of mla_v4_asm.csv. - csrc: mla_decode_v4_ps_asm in module_mla_v4_asm, with its own documented kernarg struct (same 21 x 16 B layout, repurposed slots marked), shape / dtype / contiguity / workspace checks and its own kernel cache. - aiter.mla: MlaV4NmPsWorkspace, get_mla_v4_nm_ps_workspace(device, num_partitions=128) and mla_decode_fwd_v4_nm_ps(...). One workspace serves calls that are ordered on the GPU; allocate it outside capture. - Tests: accuracy and LSE against the fp8-dequant reference (uniform, ragged, K=0 rows, N up to 3000, K up to 8320), kv_indptr with offset and extra entries, workspace reuse across batch sizes, partition counts, CUDA-graph replay with changing KV lengths, agreement with mla_decode_fwd_v4_nm, argument checks. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Add the generator for the persistent decode code object mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co is now reproducible from sources in the repo: nm_ps_src/gen.py patches the disassembly of the #5195 decode kernel (asm/rt.s) into the persistent kernel, splices in the in-kernel split planner and last-arriver merge (splice/*.hip, compiled with hipcc -S), renames the kernel symbol, assembles for gfx950 and can check the sha256 against the shipped .co. README.md documents provenance, rebuild steps and the kernarg layout. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Drop the unused span buffer from the persistent decode API The persistent code object never reads kernarg slot 0xb8, yet the span workspace [P, 2] was a field of MlaV4NmPsWorkspace, a required argument of mla_decode_v4_ps_asm and size-checked in C++. Remove it from the workspace, the ctypes signature and the C++ entry point, and pass nullptr at 0xb8. Document the three regions of the cnt counter buffer next to its size. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Add a guard-page out-of-bounds test for the persistent kernel Run mla_decode_fwd_v4_nm_ps in a subprocess at large N (up to the 32768 row limit) with every input and output placed at the very end of its own 2 MiB-rounded allocation under the non-caching allocator, so an access one element past q, qrope, the KV pools, kv_indptr (exactly N+1 entries), kv_page_indices, sink, out or lse faults the GPU. Covers the lse and no-lse paths and checks the merge counters return to zero. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Rename research identifiers in the persistent kernel generator Rename the spliced planner / combine kernels and their labels (t8pre -> ps_plan, t8fin -> ps_merge, T8_* -> PS_*) and regenerate the code object. .text, .rodata (kernel descriptor) and .note (metadata) are byte-identical to the previous build; only local symbol names in .symtab change. Update the expected sha256. README fixes: the row limit is N <= 32768 (out / q byte offsets wrap above it; the cnt / arange sizes come from a 65536-row layout), cnt is int32 with its three regions spelled out, and slot 0xb8 is passed as 0. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Mark kernarg slot 0x68 of the persistent kernel as unused Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Ship the persistent decode kernel as a code object only Drop hsa/gfx950/mla_v4/nm_ps_src (assembly, patch scripts and HIP splice sources). Nothing at build or run time reads it; the kernel ships as mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co, unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Satisfy ruff in the persistent decode changes dict literals for the workspace kwargs and the test inputs; one with-statement for the capture-time rejection check. No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Expose the split count as a host-only helper get_mla_v4_nm_num_kv_splits(num_seqs, num_heads, kv_len) returns the split count the plan helper picks, with no device work. The uniform split_indptr for it is [0, s, 2s, ...], and any prefix of a longer one is valid, so a caller can keep one constant buffer per split count and slice it instead of writing a split_indptr before every call. get_mla_v4_nm_split_plan now builds on it; its behaviour is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> * [MLA v4 nm] Drop get_mla_v4_nm_split_plan; the split count is the API Nothing calls the plan helper or MlaV4NmSplitPlan any more: callers take the count from get_mla_v4_nm_num_kv_splits and pair it with the uniform split_indptr [0, s, 2s, ...]. The tests build that pair with a local helper, and the mla_decode_fwd_v4_nm docstring points at the count helper. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com> --------- Signed-off-by: yihonglie <hyi@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
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
2 participants
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.
Technical Details
New pipeline for Dsv4 mla decode FP8+BF16