[ROCm] Fuse per-token activation quant into RMSNorm for per-channel quantized attention

[Emmanuel0612]

This PR fuses the per-token FP8 activation quantization into the preceding RMSNorm for per-channel dynamic FP8 attention projections on AMD gfx95 (validated on MI355X with Kimi-K2.7-Code-MXFP4). By emitting a pre-quantized (fp8, scale, orig_dtype) tuple straight from the norm, it removes the standalone per-token quant kernel before each attention projection, improving throughput by up to ~3.5% with no accuracy change.

Motivation

For per-channel fp8 attention projections on gfx95, the per-token activation quant runs as a standalone _per_token_group_quant_8bit launch immediately before each projection (fused_qkv_a_proj_with_mqa, q_b_proj, kv_b_proj). Profiling shows this as a separate kernel that can be folded into the preceding RMSNorm.

Change

Fold the per-token fp8 quant into the preceding RMSNorm, so the projection receives a pre-quantized (fp8, x_scale[m, 1], orig_dtype) tuple consumed directly by gemm_a8w8_bpreshuffle in apply_fp8_linear. The standalone quant kernel is eliminated. The tuple carries the original activation dtype so the GEMM output preserves it (FP16 stays FP16; not promoted to BF16).

Fusion sites, each emitting a per-token whole-row scale [m, 1]:

• Entry proj (fused_qkv_a_proj_with_mqa): fold into input_layernorm, covering the fused all-reduce + RMSNorm path (custom_fused_ar_rms_quant) and the post-all-reduce path used by large-batch prefill.
• Decode q_b_proj: fold into the fused q/kv RMSNorm (fused_qk_rmsnorm, quant_type=per_Token).
• Prefill q_b_proj / kv_b_proj: fold into the single-tensor RMSNorm (rmsnorm2d_fwd_with_dynamicquant, group_size=0).

Testing

• test/registered/unit/models/test_pertoken_fp8_eligibility.py (CPU)
  — unit tests for the fold-eligibility predicate _is_per_channel_dynamic_fp8. Asserts the layout contract is read from the layer's own state: the explicit _fp8_weight_preshuffled marker (not the global gfx95 flag), input_scale is None (dynamic), and weight_scale.numel() equal to the projection's real output-channel count (output_size_per_partition else output_size), 1-D [N] or 2-D [N,1]. Covers eligible per-channel (1-D/2-D) vs. rejected per-tensor, static-input-scale, block-scale, non-fp8, missing-marker, and a scale sized to the input dim K (must be rejected).

• test/registered/kernels/test_pertoken_fp8_fold_parity.py (MI35x)
  —GEMM-level parity: the pre-quantized (fp8, scale) tuple path vs. the standard apply_fp8_linear path produce identical output, and the carried dtype is preserved (FP16 stays FP16; a 2-tuple falls back to pre_quant_output_dtype / BF16).

• test/registered/kernels/test_pertoken_fp8_quant_method_routing.py (MI35x)
  — drives the folded tuple through the real quant method rather than apply_fp8_linear directly: test_quark_scheme_routes_tuple runs the actual Quark W8A8Fp8 apply_weights for the entry (2112), q_b (12288), and kv_b (16384) widths at M=1 (decode) and M>1 (prefill), BF16/FP16, asserting parity and dtype preservation; test_unwrapped_per_token_scale_does_not_assert guards the native Fp8LinearMethod / compressed-tensors unwrap path (apply_fp8_linear(input=qx, input_scale=scale[M,1])) that previously asserted numel == 1 and failed for M>1.

• test/registered/kernels/test_pertoken_fp8_layer_routing.py (MI35x)
  —genuine layer-path regression test: builds a real ColumnParallelLinear, quantizes and preshuffles its weight via the real Quark scheme (create_weights → load → process_weights_after_loading, which sets the marker), then calls layer.forward() — exercising the full production routing (ColumnParallelLinear.forward → quant_method.apply → scheme.apply_weights → apply_fp8_linear → gemm_a8w8_bpreshuffle). Asserts the preshuffle marker landed, the folded-tuple forward matches the plain-activation forward, and dtype is preserved (q_b/kv_b, M=1 and M>1, BF16/FP16).

• test_pertoken_fp8_ar_capture_guard.py (CPU)
  — guards the P0 in fused_allreduce_rmsnorm_quant. Under AITER's TC-piecewise replay state (_IS_CAPTURING set, stream not capturing, in a piecewise graph), custom_fused_ar_rms_quant returns non-None dummy zeros that would be consumed as real (fp8, residual, scale). Using a fake communicator with the three capture predicates patched, it asserts:
  — the wrapper returns None without calling the kernel in that capture state, and
  — it still calls the kernel in the three non-capture cases (so the fold isn't silently disabled).
  Pure control-flow logic, runs on CPU CI.

All four pass on MI355X (gfx950); the CPU eligibility test also runs on CPU CI.

Accuracy test

Accuracy (GSM8K) — neutral, within run-to-run noise:
Results (MI355X / gfx950, Kimi-K2.7-Code-MXFP4)

parallel	100	300	600	1200
original	0.951	0.953	0.942	0.950
fold_quant	0.947	0.946	0.947	0.951

Speed Tests and Profiling

Throughput / GPU (tok/s) — consistent improvement:
Results (MI355X / gfx950, Kimi-K2.7-Code-MXFP4 / 8K-1K)

Conc	4	8	16	32	64	128
original	906.82	1491.25	2128.38	2955.05	3925.4	5048.96
fold_quant	932.15	1511.13	2203.48	2992.39	3983.7	5084.88
Δ	+2.79%	+1.33%	+3.53%	+1.26%	+1.49%	+0.71%

original profiling:

fold_quant profiling:

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #34840567183
Latest PR Test (Extra): ❌ Run #34840566707
Latest PR Test (AMD ROCm 10): ❌ Run #34840567029

[kkHuang-amd]

Reviewed current head 21d2f5ab. The earlier tuple-routing, exact [M,1] scale, output-dtype, preshuffle-marker, and real ColumnParallelLinear coverage findings are resolved. Three P0 issues remain in the newly added fused collective wrapper, so I cannot approve yet. The current-head CPU and MI35X jobs were also gate-skipped; rerun them after the fixes. One P1 test-hardening item is non-blocking and may be handled in a follow-up PR.

[kkHuang-amd]

[P1] The generic default can select AITER's 1-stage kernel beyond its documented 80-token limit, but the advertised Kimi/DeepSeek path is safe. For BF16 hidden size 7168, the 128 KiB threshold permits at most nine tokens, so this is not reachable in the supported checkpoint. A future/custom hidden width such as [128,512] would exceed kMaxBlocks=80 and can silently truncate work. In a follow-up PR, add token_num <= 80 to the default decision and cover the 80/81-token boundary. This does not block approval of the current Kimi path.

[Emmanuel0612]

We can separate the issue into 3 situations and discuss.

1. the batch_size is small (<= 128KB)
   it will run "1-stage fused AR+RMSNorm+Quant", 9 tokens on Kimi, smaller than 80.
2. bigger batch_size (>128KB and AR-fusion still on)
   it will run "2-stage fused AR+RMSNorm+Quant", no token cap.
3. AR-fusion off
   it will run "plain AR + fused RMSNorm+Quant", still does not encounter 80-token risk.
   So, it will not cause any problem.

[kkHuang-amd]

[P0] The quantized wrapper omits the TC-piecewise CUDA-graph capture-state handling used by fused_allreduce_rmsnorm. When AITER's global _IS_CAPTURING state is active but the current stream is not capturing, the custom wrapper can return dummy zero outputs; this non-None tuple is then consumed as real activations/residuals/scales. Mirror the existing piecewise guard (or use the communicator's real registered path) before this call, and add capture/replay correctness coverage for the per-token format.

[Emmanuel0612]

This has been solved.

[kkHuang-amd]

[P1] The broad catch exposes a real collective-safety robustness gap, but no normal supported Kimi trigger is demonstrated. Shape/width preflight is rank-symmetric and most HIP kernel failures surface asynchronously, so the current supported path is not known to diverge here. A rank-local synchronous launch fault could still make one rank return None and enter a second all-reduce while peers proceed in the fused collective. In a follow-up PR, return None only from deterministic checks before backend entry and propagate exceptions after launch. This does not block approval once the TC-piecewise P0 is fixed.

[Emmanuel0612]

If some errors go to the exception.
Then it still can't run since the asynchronous GPU error and shape error.
That means the launching will stop without sending any incorrect result.

[kkHuang-amd]

[P1] This broad import catch can silently turn the required MI35X layer-routing regression into a green skip when AITER/Quark initialization is broken. In a follow-up PR, keep normal skips for unsupported platforms, but let unexpected import or initialization errors fail on the gfx95 runner so coverage cannot disappear silently. This item does not block approval once the P0s and current MI35X run are resolved.

[Emmanuel0612]

1. The Quark/linear imports now run unguarded when _HAS_PATH=True. For gfx95 runners, if AITER/Quark initialization is broken, the import fails (the test errors red) instead of degrading into a green skip — that's the coverage this guards.
2. The try/except wraps only the lightweight, aiter-independent hardware probe. Its sole purpose is to make sure the test skips cleanly on non-ROCm/non-gfx95 environments — i.e. "keep normal skips for unsupported platforms."

[kkHuang-amd]

Severity correction after checking the supported Kimi/DeepSeek shapes and AITER contract at head 21d2f5ab:

• P0 (blocking): missing TC-piecewise CUDA-graph handling is directly reachable under the default ROCm prefill graph path. AITER returns non-None dummy zero activations/residuals/scales when _IS_CAPTURING is set but the current stream is outside capture; this wrapper accepts them as real output. Mirror the existing unquantized guard and add capture/replay parity coverage.
• P1 (follow-up): the generic one-stage rule lacks the 80-token cap, but Kimi's BF16 hidden size 7168 limits the byte-threshold path to at most nine tokens, so this is not reachable in the supported checkpoint.
• P1 (follow-up): broad exception-to-None fallback is a collective-safety robustness gap, but no normal supported Kimi trigger for rank-local synchronous divergence is demonstrated.
• P1 (follow-up): the MI35X layer-routing test can silently skip on unexpected AITER/Quark initialization errors.

The request-changes state remains because one reachable P0 is still open. Once it is fixed and current-head CPU/MI35X CI executes successfully, the remaining P1s do not block approval.

[kkHuang-amd]

Re-reviewed head d31e005a79c702368756ff3a326968b0ba1b6d3c against latest main dad6fd0f04556a9a2c09fc08388ecee45ed5a33f.

The previous blocking findings are resolved: tuple eligibility and exact [M,1] routing are tied to the real preshuffled per-channel layout, output dtype is preserved, production ColumnParallelLinear coverage is present, and the fused per-token collective now avoids AITER dummy outputs in the TC-piecewise replay state. The MI35X import guard also fails loudly on the target platform. No current PR hunk is redundant with main.

The remaining 80-token one-stage cap and post-collective broad exception fallback are non-blocking P1 follow-ups already noted in review comments.

Approving from a code-review perspective. Before merge, integrate current main and rerun the required CPU and MI35X checks; the relevant MI35X targeted job passed, but aggregate CI is currently red.

[Emmanuel0612]

@amd-bot ci-status

[Emmanuel0612]

@amd-bot ci-status

[yichiche]

@Emmanuel0612 Please fix lint test.
Lint / lint (pull_request)
Lint / lint (pull_request)Failing after 2m

[Emmanuel0612]

@Emmanuel0612 Please fix lint test. Lint / lint (pull_request) Lint / lint (pull_request)Failing after 2m

Thank you. It is done.