[ZLkanyo009]

Summary

Decode-sized QSA index scoring (the target's MTP verify and the draft, at most 256 query rows) now runs on a Triton MFMA kernel that stacks every row's 4 index heads into the M dimension. A request's compressed keys are read once per verify instead of once per row, and the target's decode scoring drops from 1130 to 320 µs a step (-6.2% of decode kernel time on Qwen3.8-Flash-Next).

Motivation

Part of the Qwen3.8-Flash-Next (amd/Qwen3.8-Flash-Next-Quark-MXFP4) optimization on MI355X / ATOM: https://github.com/wuhuikx/LLM-Booster/issues/263

At the decode work point, _qsa_paged_mqa_logits_kernel (CUDA-core, one program per query row) was the third-largest decode kernel. With MTP 3 a request contributes 4 verify rows that share the same keys, and each of them streamed the request's whole compressed-key cache. The existing MFMA kernel is not used below 256 rows because one row per request leaves M mostly empty.

Changes

• atom/model_ops/qwen4_exp/ops/qsa.py
  • New _qsa_paged_mqa_logits_stacked_kernel: a program's BMT rows x NUM_HEADS heads form one M = BMT * heads operand, so one tl.dot per key tile scores every head of every row; the per-head ReLU is applied before the heads are summed back out of M. BMT = 4 with 4 heads is one 16-row MFMA tile and, in decode, exactly one request's verify rows. Programs stride over column tiles up to the block's widest visible count; rows of different requests in one block (the draft: one row per request) walk their requests as the prefill MFMA kernel does.
  • qsa_paged_mqa_logits routes calls with at most 256 rows to it when the index heads and head dim are powers of two (heads >= 4); everything else is unchanged. Tile (BMT, BN, warps, column programs) = (4, 32, 2, 256), chosen with an L2-flushed microbench at 36 rows over 9 requests and confirmed on a live server.
• atom/utils/envs.py, docs/environment_variables.md: ATOM_QSA_STACKED_DECODE_SCORER (default 1; 0 restores the CUDA-core kernel).
• tests/model_ops/test_qsa_stacked_scorer.py (GPU): stacked vs CUDA-core vs a torch reference, 1 and 4 rows per request, 4- and 16-row pages (block 16 / 64), padded rows, uneven contexts, max_columns.

Accuracy

Selection is unchanged: logits match the CUDA-core kernel to FP32 summation order (the unit test holds them to rtol / atol 5e-5). End to end (lm_eval gsm8k 5-shot through the chat API with fewshot_as_multiturn, temperature 0, real MTP verification, MTP 3):

	GSM8K strict / flexible	MTP accepted tokens / forward
before (main), full 1319	0.9765 / 0.9765	3.165
this PR + #2652 (BF16 path), block 64, full 1319	0.9750 / 0.9750	3.176
this PR, block 16, --limit 200 (branch code)	0.985 / 0.985	3.164

(stderr ≈ 0.004 on the full set; the first 200 questions score 0.975-0.99 across all earlier full runs.)

Performance

Kernel time per step from a torch trace, MI355X, TP1, MTP 3, BF16 KV, block 16 (rocm/atom-dev:latest 566aa3dc9d5d, ATOM 53e4fb524, AITER 948dd267f). Decode = decode[bs=9/16 tok=36 d=9 spec=3] averaged over 400+ steps; prefill = prefill[bs=1 tok=16384 ctx=98304]. Both from a controlled profile (decode: 9 concurrent requests on 64k-128k prompts; prefill: the 7th chunk of a 130k prompt). "before" = this branch with ATOM_QSA_STACKED_DECODE_SCORER=0, i.e. main's kernels (mean of two runs).

µs / step	before	this PR
decode, all kernels	16,929	15,883 (-6.2%)
decode + 3 MTP drafts	19,318	18,161 (-6.0%)
target index scoring (12 layers)	1,130	320
draft index scoring (3 steps)	147	85
prefill, all kernels (not routed here)	445,272	448,818 (run-to-run ±1%)

With --block-size 64 the target's decode scoring is 285 µs (16 rows a page, fewer page-table lookups).

AgentX C16 (--scenario agentx, 3600 s, golden AL 2.32), BF16 KV, block 16:

	tput / GPU	realized prefix hit	TTFT p90	ITL p90	errors
before (main)	45,885	62.5%	16.1 s	56.6 ms	0 / 1671
this PR	52,688	68.1%	10.8 s	45.1 ms	0 / 1971

Not a clean kernel A/B: the BF16 pool (1.61 M tokens) is too small for 16 AgentX sessions, so the realized prefix hit, which depends on eviction timing in this closed-loop load, moves by several points between runs and drives most of the difference. The two runs were also on different MI355X nodes. Read it as "no regression"; the kernel gain is the decode profile above.

How to use

On by default. ATOM_QSA_STACKED_DECODE_SCORER=0 keeps the CUDA-core scorer. It serves block 16 and block 64 alike; the recommended Qwen3.8-Flash-Next config is --block-size 64 (with --kv_cache_dtype fp8 --index-cache-dtype bf16 from #2652), where block 64 also moves BF16 index prefill scoring to FlyDSL.

Test plan

• pytest tests/model_ops/test_qsa_stacked_scorer.py tests/model_ops/test_triton_dynamic_shapes.py tests/test_envs.py on MI355X: 128 passed.
• black / ruff (pre-checks.yaml versions).
• Server from this branch (+ fix(qwen4_exp): load unscaled Quark MTP weights as BF16 #2606 for the MTP weights): GSM8K --limit 200 and MTP acceptance above; kernel profile and AgentX above.

Depends on #2606 at runtime for amd/Qwen3.8-Flash-Next-Quark-MXFP4 (MTP weights), not for the code itself. #2652 (FP8 KV / FP8 / FP4 index caches for QSA) builds on this PR.

[github-actions]

🏷️ 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 2651 --add-label <label>

[ZLkanyo009]

Merged origin/main at 5fea029 (#2606) into this branch (merge 6698b9a, no rebase/force-push).

Conflict: atom/utils/envs.py only. Kept main's new ATOM_PA_FLYDSL and this PR's ATOM_QSA_STACKED_DECODE_SCORER, in file order, no duplicates. No other files needed manual edits. The branch also merges cleanly with the current main (4f6ce2f, #2684 engram).

Verified (MI355X, rocm/atom-dev:latest + CI aiter wheel):

• black/ruff (CI versions) clean on the touched files.
• CPU unit tests (.github/scripts/run_unit_tests.sh): 3 failed / 9721 passed. These are the same 3 failures as unmodified main (test_graph_marker x2, test_sparse_indexer_side_effect_op).
• GPU: test_qsa_stacked_scorer.py, test_triton_dynamic_shapes.py, test_qwen4_exp_mtp.py and test_qwen4_exp_quantization.py: 102 passed, 6 failed. The 6 failures are test_qsa_rope_cache_matches_implicit_positions (fake metadata lacks is_prefill), and main fails the same 6.
• Qwen3.8 GSM8K, full 1319, MTP 3 with real verification: 0.9765 (main + fix(qwen4_exp): load unscaled Quark MTP weights as BF16 #2606 control: 0.9788). Acceptance rate 0.717 (main 0.723).

[ZLkanyo009]

Folded into #2652. #2652 was stacked on this PR and already contains its commits (branch lingzha/qwen38-qsa-fp8-kv-index includes lingzha/qwen38-qsa-stacked-decode-scorer at 6698b9a9), so the two can be reviewed and merged as one. The description of this PR is kept as Part 1 of #2652. No code change.