perf(dsv4): DSpark decode kernel fusions (stacked on #16)

[sunxxuns]

Stacked on #16 (feat/dspark-verify-width); the first commit here is that PR. Review from the second commit on.

DeepSeek-V4.1 DSpark decode on MI355X, TP4/EP4 on 4 GPUs. The changes remove launches and host syncs from the verify/draft step and fuse small memory-bound kernels into their producers. No before/after profile of the baseline exists, so how much each change contributes to the gain has not been measured.

What changed

Verify-width runtimes (second commit; part of the baseline below, not covered by the numbers) (dspark_worker_v2.py, dspark_config.py, environ.py, scheduler / memory pool)

• Separate width-3 and bs1 (gamma 7) runtimes with their own captured CUDA graphs, so the per-step verify width from feat(spec): pick the DSpark static verify width per step #16 switches graphs instead of padding to full width.
• Verify widths 2,3 on by default for DeepSeek-V4.1 targets, with a packaged MI355X table as fallback.
• Scheduler resolves seq_lens_cpu on the forward stream; memory pool and graph batch sizing account for the bs1 draft length.

FP4 indexer and attention (fp4_indexer*.py, fp4_rope_fake_quant.py, fused_qk_norm_rope_store.py, HIP attention backends)

• Row-block FP4 indexer quantization.
• RoPE-tail fake-quant and FP4 query packing in one launch.
• Same-length DSpark expansion kernel for the verify batch.

MoE (moe_runner/aiter.py, topk.py, fill_padded_rows.py, deepseek_v2.py)

• Masked local route-reduce as the epilogue of AITER's dense EP stage 2, fused with the shared-expert add and routed scale.
• Small-M sort path.
• Padded top-k IDs and zero weights filled in one launch.

mHC / norm / main projection (mhc_post_combine_hip.py, mxfp8_dot_scaled_splitk.py, deepseek_v4*.py, fp8_hip.py)

• mHC post-combine returns the updated residual and the combined input in one launch.
• FP8 grid quantization folded into the mHC/norm producers.
• Fixed-shape DSpark main projection as a dot_scaled split-K GEMM with a separate reduce.

Router gate (rocm_router_gate.py, topk.py)

• Padded rows are masked inside the ROCm router gate launch, so top-k post-processing skips the separate pad-mask fill.

mHC post-combine norm and activation (mhc_post_combine_hip.py, silu_and_mul_clamp_hip.py, deepseek_v2_hip_act.py)

• mHC post-combine also returns the normed input in the same launch.
• HIP SiLU-mul-clamp drops the 1024-wide cap and uses exp2 on the fp8-grid path.

Draft and acceptance (logits_processor.py, dspark_verify.py, dspark_draft.py, dspark_verify_width.py, CUDA graph runner)

• Greedy argmax over the vocab-parallel shards without gathering logits.
• Draft input copies in one kernel.
• Acceptance and output tokens finalized on device with no blocking host sync.
• RoPE applied to the draft query tail while writing the full row in one kernel; preallocated verify output buffers; candidate page-table zeroing folded into the length kernel; engram gather reads strided 2-D indices without a reshape copy.

Results

Baseline is #16 plus the verify-width runtimes (commits 1-2); candidate is the branch head. The numbers therefore cover the FP4, MoE, mHC, router and draft/acceptance changes, not the verify-width runtimes.

Design, and what it does not control:

• Each run launches the baseline server, benchmarks it, then launches the candidate server on the same GPUs. The order was always baseline first, candidate second (no ABBA).
• The 4 replicate runs were two pairs running concurrently on the two 4-GPU halves of the same node (runs 2+5 on one node, 3+4 on another).
• Run 1 is the automated-search run that selected this candidate, so it is not independent. Stats are shown with and without it.
• No A/A (same build twice) control was run.
• Numbers were measured on the build with one-time debug prints (VARCTX_SCHEDULE_OPT, MOE_SORT_SMALL) in AITER and matching log lines in SGLang. The branch head and the patch below remove them; they fire once per (rank, shape) and are not in the timed path, but the head itself was not re-measured.

Throughput gain, median of 6 timed runs per launch; 95% CI and two-sided paired t:

workload	n=4 replicates (excl. selecting run)	p	n=5 (incl. selecting run)	p	per-run (run 1 first)
real-bs1	+9.6% [-6.5%, +25.7%]	0.15	+9.5% [-1.4%, +20.4%]	0.073	+9.1, +5.2, +24.6, +2.3, +6.3
real-bs8	+6.8% [+4.1%, +9.6%]	0.004	+6.8% [+5.0%, +8.7%]	<0.001	+6.9, +9.2, +5.2, +6.9, +6.1
real-bs32	+6.1% [+5.6%, +6.5%]	<0.001	+6.4% [+5.4%, +7.5%]	<0.001	+7.9, +6.5, +5.9, +5.9, +5.9
random-bs1	+9.6% [+8.7%, +10.5%]	<0.001	+9.3% [+8.4%, +10.3%]	<0.001	+8.3, +9.2, +9.1, +9.8, +10.3

real-bs1 varies most per run because the draft acceptance length differs between server launches; normalized by acceptance length (n=5) it is +7.1% [+2.4%, +11.7%].

Correctness:

• GSM8K (5-shot, 1319 questions, greedy) is a non-inferiority gate, not an exactness check. It passed in all 5 runs; candidate minus baseline across runs is +21 correct of 6595.
• Decoding is not deterministic across server launches: two launches of the same build differ on 357-405 answers, and baseline vs candidate differ on 373-391.
• Kernel-level exactness was checked only with torch.equal on selected shapes, for example varctx schedule and small-M MoE sort at M=1/4/24/32/37/128. It was not checked end to end.
• test/registered/unit/spec/test_dspark_verify_width_policy.py: 6 passed.

Reproduction pins

• Image: ghcr.io/amdpilot-org/dsv41-pr39857@sha256:fccf7f7268e97eb137ddf2df0dc94cf1af1fd60186aaafef03deb74504ccef20
• SGLang base: e2e824dc583071ea1bf62f4cebe4dd437d70ff79
• AITER base: acf8fdf9307431ece8ee275971c41cb3d1a7020b + [Bugfix] Drain FlyDSL stage-1 LDS-DMA loads before the tile barrier ROCm/aiter#5561 @ 610a95851c + [Perf] Add gfx950 DSV4.1 Flash EP4 a8w4 FMoE tuning ROCm/aiter#5562 @ a2506c5bcb
• Model revision: dba1be0a40aa45a94ad051997016db3960a90277
• Eval contract pr39857-eval-v2 (sha256 5b2efa85c0d42f9c7dfe0dd08f618bab576b5efbcb75c07602d3f778c6bae8d6), script /opt/amdpilot/contract-pr39857/pr39857_eval.py in the image.
• Workload: 4096 input / 1024 output tokens; real prompts at bs 1/8/32 and random prompts at bs 1; FP8 KV cache; DSpark speculative decoding.

Required AITER change

The small-M MoE sort path, the fused_moe.py dispatch and the FP4 pa_mqa_logits schedule need this AITER patch (against the AITER base above); without it these fall back to the existing paths. The same change is open as amdpilot-org/aiter#88 (branch perf/dsv41-small-m-moe-sort-varctx, on acf8fdf9).

• moe_sorting_small.py + fused_moe.py dispatch: a Triton sort for M<=128, 384 experts, top-6, EP mask. Its outputs match the existing path under torch.equal across M=1-128 and random/all-local/all-remote/single-expert masks, plus full-layer parity at M=24 and M=128. Microbenchmark: 17.6 to 13.5 µs at M=24, 17.7 to 14.3 µs at M=128.
• pa_mqa_logits_fp4.py: BLOCK_S for the _varctx_cta_info_kernel grid goes from 256 to 32. 32 was the fastest of 32/64/128, and outputs are unchanged under torch.equal. CUDA-graph median: 24.9 to 17.8 µs at M=24, 53.0 to 23.2 µs at M=128.

aiter patch

diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py
index 85b02d52e..9c59bbdae 100644
--- a/aiter/fused_moe.py
+++ b/aiter/fused_moe.py
@@ -44,6 +44,7 @@ from aiter.ops.flydsl.mxfp4_kname import (
     parse_g2_kname_any,
 )
 from aiter.ops.moe_mxfp4_aux import _mxfp4_moe_sort_internal_is_supported
+from aiter.ops.moe_sorting_small import moe_sorting_small
 from aiter.ops.opus import moe_stage2_a8w4 as _opus_a8w4
 from aiter.ops.opus.moe_stage1_a8w4 import (
     opus_a8w4_stage1_wrapper as _opus_a8w4_stage1_wrapper,
@@ -563,6 +564,37 @@ def moe_sorting(
     output_aux=False,
     output=None,
 ):
+    if (
+        torch.version.hip is not None
+        and not flat
+        and not output_aux
+        and not return_local_topk_ids
+        and dispatch_policy == 0
+        and num_local_tokens is None
+        and expert_mask is not None
+        and expert_mask.numel() == 384
+        and num_experts == 384
+        and topk_ids.shape[1] == 6
+        and topk_ids.shape[0] <= 128
+        and block_size == 32
+        and topk_ids.dtype == dtypes.i32
+        and topk_ids.is_contiguous()
+        and topk_weights.dtype == dtypes.fp32
+        and topk_weights.is_contiguous()
+        and expert_mask.dtype == dtypes.i32
+        and expert_mask.is_contiguous()
+    ):
+        return moe_sorting_small(
+            topk_ids,
+            topk_weights,
+            num_experts,
+            model_dim,
+            moebuf_dtype,
+            block_size,
+            expert_mask,
+            accumulate,
+            output,
+        )
     if (
         not _USE_CK_MOE_SORTING
         and _USE_FLYDSL_MOE_SORTING
diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py
index a2972fe6d..7e0ff3ab8 100644
--- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py
+++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py
@@ -137,7 +137,7 @@ def compute_varctx_schedule(
         cta_info = cta_info_out
     safe_out = torch.empty(1, dtype=torch.int32, device=dev)
     BLOCK_B = triton.next_power_of_2(max(int(B), 1))
-    BLOCK_S = 256
+    BLOCK_S = 32
     grid = (triton.cdiv(S, BLOCK_S),)
     _varctx_cta_info_kernel[grid](
         ctx_i32,
diff --git a/aiter/ops/moe_sorting_small.py b/aiter/ops/moe_sorting_small.py
new file mode 100644
index 000000000..ae6d1a354
--- /dev/null
+++ b/aiter/ops/moe_sorting_small.py
@@ -0,0 +1,193 @@
+"""Small-M expert-parallel MoE sorting for DeepSeek-V4.1-Flash decode."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _moe_sorting_small_kernel(
+    topk_ids_ptr,
+    topk_weights_ptr,
+    sorted_ids_ptr,
+    sorted_weights_ptr,
+    sorted_expert_ids_ptr,
+    num_valid_ids_ptr,
+    moe_buf_ptr,
+    expert_mask_ptr,
+    local_ids_ptr,
+    M,
+    num_local,
+    invalid_local_id,
+    moe_buf_numel,
+    E: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+    BLOCK_NK: tl.constexpr,
+    BLOCK_PAD: tl.constexpr,
+    BLOCK_BLOCKS: tl.constexpr,
+    BLOCK_ZERO: tl.constexpr,
+    HAS_MOE_BUF: tl.constexpr,
+):
+    expert = tl.program_id(0)
+    if expert != 0:
+        local_id = tl.load(local_ids_ptr + expert)
+        if local_id == invalid_local_id:
+            return
+    item_offs = tl.arange(0, BLOCK_NK)
+    item_mask = item_offs < M * TOPK
+    expert_ids = tl.load(topk_ids_ptr + item_offs, mask=item_mask, other=0).to(
+        tl.int32
+    )
+    route_local_ids = tl.load(
+        local_ids_ptr + expert_ids, mask=item_mask, other=invalid_local_id
+    )
+    histogram = tl.histogram(route_local_ids, BLOCK_L, mask=item_mask)
+
+    local_offs = tl.arange(0, BLOCK_L)
+    local_valid = local_offs < num_local
+    blocks = (histogram + BLOCK_M - 1) // BLOCK_M
+    padded = tl.where(
+        (histogram != 0) & local_valid, blocks * BLOCK_M, 0
+    )
+    inclusive = tl.cumsum(padded, axis=0)
+    exclusive = inclusive - padded
+    total_padded = tl.sum(padded, axis=0)
+    if expert == 0:
+        tl.store(num_valid_ids_ptr, total_padded)
+        tl.store(num_valid_ids_ptr + 1, M)
+        if HAS_MOE_BUF and num_local == 0:
+            for zero_start in tl.range(0, moe_buf_numel, BLOCK_ZERO):
+                zero_offs = zero_start + tl.arange(0, BLOCK_ZERO)
+                tl.store(
+                    moe_buf_ptr + zero_offs,
+                    tl.zeros((BLOCK_ZERO,), dtype=moe_buf_ptr.dtype.element_ty),
+                mask=zero_offs < moe_buf_numel,
+                )
+
+    local_id = tl.load(local_ids_ptr + expert)
+    if local_id == invalid_local_id:
+        return
+
+    count = tl.sum(tl.where(local_offs == local_id, histogram, 0))
+    start = tl.sum(tl.where(local_offs == local_id, exclusive, 0))
+    padded_count = ((count + BLOCK_M - 1) // BLOCK_M) * BLOCK_M
+
+    if HAS_MOE_BUF:
+        for zero_start in tl.range(
+            local_id * BLOCK_ZERO, moe_buf_numel, num_local * BLOCK_ZERO
+        ):
+            zero_offs = zero_start + tl.arange(0, BLOCK_ZERO)
+            tl.store(
+                moe_buf_ptr + zero_offs,
+                tl.zeros((BLOCK_ZERO,), dtype=moe_buf_ptr.dtype.element_ty),
+                mask=zero_offs < moe_buf_numel,
+            )
+
+    is_mine = item_mask & (route_local_ids == local_id)
+    rank = tl.cumsum(is_mine.to(tl.int32), axis=0) - is_mine.to(tl.int32)
+    token = item_offs // TOPK
+    slot = item_offs % TOPK
+    packed = (slot << 24) | token
+    weights = tl.load(topk_weights_ptr + item_offs, mask=item_mask, other=0.0)
+    tl.store(sorted_ids_ptr + start + rank, packed, mask=is_mine)
+    tl.store(sorted_weights_ptr + start + rank, weights, mask=is_mine)
+
+    pad_offs = tl.arange(0, BLOCK_PAD)
+    pad_mask = (pad_offs >= count) & (pad_offs < padded_count)
+    sentinel = (TOPK << 24) | M
+    tl.store(sorted_ids_ptr + start + pad_offs, sentinel, mask=pad_mask)
+    tl.store(
+        sorted_weights_ptr + start + pad_offs,
+        tl.zeros((BLOCK_PAD,), dtype=tl.float32),
+        mask=pad_mask,
+    )
+
+    block_offs = tl.arange(0, BLOCK_BLOCKS)
+    num_blocks = (count + BLOCK_M - 1) // BLOCK_M
+    tl.store(
+        sorted_expert_ids_ptr + start // BLOCK_M + block_offs,
+        local_id,
+        mask=block_offs < num_blocks,
+    )
+
+
+
+def moe_sorting_small(
+    topk_ids,
+    topk_weights,
+    num_experts,
+    model_dim,
+    moebuf_dtype,
+    block_size,
+    expert_mask,
+    accumulate,
+    output,
+):
+    """Small-M (M<=128) replacement for the production EP sorting path."""
+    device = topk_ids.device
+    M, topk = topk_ids.shape
+    max_num_tokens_padded = int(topk_ids.numel() + num_experts * block_size - topk)
+    max_num_m_blocks = int((max_num_tokens_padded + block_size - 1) // block_size)
+    sorted_ids = torch.empty(max_num_tokens_padded, dtype=torch.int32, device=device)
+    sorted_weights = torch.empty(
+        max_num_tokens_padded, dtype=torch.float32, device=device
+    )
+    sorted_expert_ids = torch.empty(max_num_m_blocks, dtype=torch.int32, device=device)
+    num_valid_ids = torch.empty(2, dtype=torch.int32, device=device)
+
+    if expert_mask is not None or accumulate:
+        moe_buf = (
+            output
+            if output is not None
+            else torch.empty((M, model_dim), dtype=moebuf_dtype, device=device)
+        )
+    else:
+        moe_buf = torch.empty((0, 0), dtype=moebuf_dtype, device=device)
+
+    cache = getattr(expert_mask, "_moe_sorting_small_cache", None)
+    if cache is None or cache[0] != expert_mask._version:
+        num_local = int(expert_mask.sum().item())
+        local_ids = torch.cumsum(expert_mask, dim=0, dtype=torch.int32) - expert_mask
+        block_l = max(4, triton.next_power_of_2(num_local + 1))
+        invalid_local_id = block_l - 1
+        local_ids = local_ids.masked_fill(expert_mask == 0, invalid_local_id)
+        local_ids = local_ids.contiguous()
+        expert_mask._moe_sorting_small_cache = (
+            expert_mask._version,
+            local_ids,
+            num_local,
+            invalid_local_id,
+        )
+    else:
+        _, local_ids, num_local, invalid_local_id = cache
+    block_nk = triton.next_power_of_2(max(M * topk, 1))
+    max_padded = ((M * topk + block_size - 1) // block_size) * block_size
+    max_blocks = (M * topk + block_size - 1) // block_size
+    _moe_sorting_small_kernel[(num_experts,)](
+        topk_ids,
+        topk_weights,
+        sorted_ids,
+        sorted_weights,
+        sorted_expert_ids,
+        num_valid_ids,
+        moe_buf,
+        expert_mask,
+        local_ids,
+        M,
+        num_local,
+        invalid_local_id,
+        moe_buf.numel(),
+        E=num_experts,
+        TOPK=topk,
+        BLOCK_M=block_size,
+        BLOCK_L=max(4, triton.next_power_of_2(num_local + 1)),
+        BLOCK_NK=block_nk,
+        BLOCK_PAD=triton.next_power_of_2(max(max_padded, 1)),
+        BLOCK_BLOCKS=triton.next_power_of_2(max(max_blocks, 1)),
+        BLOCK_ZERO=1024,
+        HAS_MOE_BUF=moe_buf.numel() > 0,
+        num_warps=2,
+    )
+    return sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf

[jhinpan]

DeepSeek-V4.1 AMD support is on main now (sgl-project#41308), and #16 and #18 went to sgl-project#41994. We checked whether this PR's AMD-only kernel commits are worth porting to main: FP4 indexer and attention launch fusions, MoE local route-reduce epilogue and padded-row fill, fp8-grid quantization folded into the mHC/norm producers with the split-K main projection, the router-gate pad mask, and the mHC post-combine norm / SiLU-mul-clamp changes.

They don't apply to main as they are. Each of these commits conflicts with main, and main already covers part of the same ground: topk._fold_pad_into_append folds the padded fill into the aiter append+remap kernel, and sgl-project#41021 fused the mHC boundary with the all-reduce and added mHC post kernels.

The remaining budget on main is small. MI355X x4, TP4/EP4, Low-Latency cell, main a702970ecb:

• Batch 32 verify step: 22.9 ms of GPU time. The kernels these commits target take about _fill_padded_rows 0.35 ms, _fp8_grid_quant 0.40 ms, _silu_and_mul_clamp 0.18 ms and moe_reduction 0.20 ms per step (kineto, so upper bounds).
• Knockout ceiling, with simulated acceptance fixed at 3.2 in both arms and ABBA order: we skipped the padded-row fill and made the fp8-grid fake-quant return its input. That removes the work entirely while keeping the values: the fill is a no-op at these graph sizes, and the fake-quant output differs only by fp8 rounding. It speeds up verify cycles by +2.46% [+2.09, +2.92] at batch 1 and +0.80% [+0.45, +1.23] at batch 32.
• A fusion keeps part of that work, so on main these commits are worth at most about 2% at batch 1 and under 1% at batch 32.

Upstreaming would also mean a rewrite. The code uses module-level mutable epilogue state, @dataclass, two identical Triton kernels, NotImplementedError for dispatch, and kernels defined inside runner/backend files. The small-M sort path also needs amdpilot-org/aiter#88.

Our guess is that most of the gain measured here comes from the shared-path commits (verify-width runtimes, sharded argmax and non-blocking acceptance, draft RoPE and preallocated verify outputs) and from the AITER patch; we did not measure those separately. They change behavior on other platforms too, e.g. the --speculative-dspark-block-size semantics, verify widths on by default for every V4.1 target with an MI355X table, and the bs=1 verify block. So they would need their own design discussion before going to sgl-project/sglang.