[gbyu-amd]

Summary

#2725 started passing the pass's real row count to aiter fused_moe (num_local_tokens=scheduled_tokens_gpu) on every Mxfp4MoEMethod.apply call. aiter's inline-sort MXFP4 kernels cannot take that argument. aiter therefore drops the whole tuned row and falls back to default heuristics:

[aiter] [fused_moe] inline-sort config is unsupported (unsupported option num_local_tokens); using default heuristics

• The check is _mxfp4_inline_sort_unsupported in aiter/fused_moe.py. It is present in both the old and the new aiter, so only the ATOM caller changed.
• Kimi-K3's tuned rows for M ≤ 64 (kimik3_a4w4_tuned_fmoe.csv, gfx950, E=896, topk=16) are all inline-sort, e.g. flydsl_mxmoe_g1_a4w4_16x256x256_f16in_nt_situv2_xcd4 + flydsl_moe2_..._sbm16.
• With the count passed, those calls fall back to flydsl_mxmoe_g1_a4w4_64x256x256_situv2 / ..._t64x256x128_reduce_sbm64, and the routed MoE gets 35–45% slower for every decode batch of 4–64 rows.
• Rows ≥ 128 use the non-inline-sort tuned rows and are unaffected.

Fix: routing already sends rows at or past the real count to expert -1 (select_experts → num_valid_rows). The count is therefore only needed for a prefilled fused shared-expert column, which routing cannot mask. This PR passes num_local_tokens only when layer.num_fused_shared_experts > 0. Layers without a fused shared column go back to the pre-#2725 call and keep their tuned rows.

Impact observed

Kimi-K3 TP8 MXFP4 on MI355X, AgentX benchmark (InferenceX), image nightly_202610261558 (ATOM e893780) vs the 10-19 image:

conc	rows / decode step	ITL p50	out tok/s
C4 (k=7)	~32	+8.5%	-0.0%
C14 (k=3)	~30–60	+13%	-4.0%
C16 (k=3)	~30–64	+13.2%	-4.1%
C48 / C56 (no spec)	~190–220	-3.2% / +0.1%	+5.4% / -0.8%

Local reproduction on one MI355X node, same recipe. The old ATOM is f08cd7f + aiter 7be01d1c7; the new is the stock nightly_202610261558:

conc	ITL p50 old → new	out tok/s old → new
C4	6.4 (CI 10-19) → 6.91 ms	163.4 → 163.1
C16	11.62 → 12.86 ms (+10.7%)	527.4 → 510.7 (-3.2%)

The new C16 run has a larger KV pool (14179 vs 13259 blocks), so the slowdown is not KV-pool driven.

The C4 server log on the new image has ~60k of the warning above. The old image has 0.

We ruled out the other candidates:

• A local HERD on/off A/B at C16 gave identical results.
• The aiter commits in 7be01d1c7..41d4d8d7e are not involved:
  • CustomAR #6819 is capture-only.
  • 8-wave a8w8 #6368 only runs at M ≥ 512.
  • The topk pad-row changes are cheaper, not dearer.

Unit-level A/B (aiter fused_moe, K3 TP8 shape)

The benchmark times one aiter fused_moe call:

• Shape: E=896, topk=16, dim=3584, inter=384 (3072/TP8), Situv2, per_1x32 fp4x2 a4w4, preshuffled weights.
• Routing: real biased_grouped_topk, with AITER_FLYDSL_USE_HERD=1 AITER_FLYDSL_HERD_MIN_M=64 as in the recipe. The ids match a torch reference.
• Timing: CUDA-graph replay, 200 iterations, idle MI355X (gfx950).
• Env: AITER_SITUV2_A4W4=1 AITER_FLYDSL_STAGE2_FP8=1.
• nlt=None is this PR's and the pre-feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV #2725 behaviour. nlt=tensor is the current main behaviour.

M	aiter 41d4d8d7e nlt=None (µs)	aiter 41d4d8d7e nlt=tensor (µs)	Δ	aiter 7be01d1c7 nlt=None (µs)	aiter 7be01d1c7 nlt=tensor (µs)	Δ
1	20.9	36.1	+73.0%	21.0	36.0	+71.5%
4	34.3	48.5	+41.6%	34.1	49.6	+45.4%
8	55.2	71.2	+28.9%	55.4	72.2	+30.2%
16	86.1	122.8	+42.7%	86.4	122.2	+41.4%
32	142.5	193.3	+35.6%	142.3	193.7	+36.1%
48	188.6	250.2	+32.7%	188.6	250.0	+32.6%
56	198.5	265.5	+33.8%	198.3	266.0	+34.2%
64	195.8	260.8	+33.2%	195.9	260.9	+33.2%
128	258.9	258.9	+0.0%	258.7	259.3	+0.2%
192	300.5	300.4	-0.0%	300.5	300.3	-0.1%
256	307.6	308.6	+0.3%	308.1	308.5	+0.1%

• The two aiter versions agree within 0.5% in every cell. The regression comes entirely from the argument ATOM passes, not from any aiter kernel change.
• Every M from 1 to 64 (the inline-sort tuned rows) pays +29–73%. Rows ≥ 128 are untouched.
• The relative L2 between the two paths is ~2.7% for M ≤ 64. The tuned kernel keeps an f16 intermediate (_f16in) and the fallback an fp4 one. This is expected; it is not a correctness bug.

benchmark script

"""Kimi-K3 TP8 decode MoE: aiter fused_moe with vs without num_local_tokens.

ATOM e893780f4 (#2725) started passing num_local_tokens on every target decode
pass; aiter then discards the tuned inline-sort rows (kimik3_a4w4_tuned_fmoe.csv,
E=896 token<=64) and falls back to default heuristics. Same script runs in the
old (atom_1025) and new (atom_1026) containers.

  AITER_SITUV2_A4W4=1 AITER_FLYDSL_STAGE2_FP8=1 python bench_k3_moe_nlt.py
"""

import os
import sys

import torch

import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe
from aiter.ops.shuffle import shuffle_weight
from aiter.utility import fp4_utils

E, TOPK, DIM, INTER = 896, 16, 3584, 384  # K3 latent MoE, inter 3072 / TP8
MS = [int(m) for m in os.environ.get("MS", "4,8,16,32,48,56,64,128,192,256").split(",")]
ITERS = int(os.environ.get("ITERS", "200"))
dev = "cuda"
torch.manual_seed(0)


def fp4_weight(n, k):
    # Random packed fp4 nibbles + e8m0 scales near 1.0 (127 = 2^0); timing only.
    w = torch.randint(0, 256, (E, n, k // 2), dtype=torch.uint8, device=dev)
    s = torch.randint(118, 124, (E * n, k // 32), dtype=torch.uint8, device=dev)
    w = shuffle_weight(w.view(dtypes.fp4x2), layout=(16, 16))
    w.is_shuffled = True
    return w, fp4_utils.e8m0_shuffle(s.view(dtypes.fp8_e8m0))


w1, w1_s = fp4_weight(2 * INTER, DIM)
w2, w2_s = fp4_weight(DIM, INTER)
# bf16 like the logits: the hip topk reads correction_bias as the logits dtype.
bias = torch.randn(E, dtype=torch.bfloat16, device=dev) * 0.01


def route(x_m):
    logits = torch.randn(x_m, E, dtype=torch.bfloat16, device=dev)
    tw = torch.empty(x_m, TOPK, dtype=torch.float32, device=dev)
    ti = torch.empty(x_m, TOPK, dtype=torch.int32, device=dev)
    aiter.biased_grouped_topk(logits, bias, tw, ti, 1, 1, True, 1.0)
    return tw, ti


def run(x, tw, ti, nlt):
    return fused_moe(
        x, w1, w2, tw, ti,
        activation=ActivationType.Situv2,
        quant_type=QuantType.per_1x32,
        w1_scale=w1_s, w2_scale=w2_s,
        num_local_tokens=nlt,
    )


def bench(m, use_nlt):
    torch.manual_seed(m)  # both arms see the same x and routing
    x = torch.randn(m, DIM, dtype=torch.bfloat16, device=dev)
    tw, ti = route(m)
    nlt = torch.tensor([m], dtype=torch.int32, device=dev) if use_nlt else None
    for _ in range(3):
        out = run(x, tw, ti, nlt)
    torch.cuda.synchronize()
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        out = run(x, tw, ti, nlt)
    for _ in range(10):
        g.replay()
    a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    a.record()
    for _ in range(ITERS):
        g.replay()
    b.record()
    torch.cuda.synchronize()
    return a.elapsed_time(b) * 1000 / ITERS, out


print(f"aiter={os.path.dirname(aiter.__file__)} E={E} topk={TOPK} dim={DIM} inter={INTER}", flush=True)
print(f"{'M':>5} {'nlt=None us':>12} {'nlt=tensor us':>14} {'delta%':>8} {'relL2':>8}")
for m in MS:
    t0, o0 = bench(m, False)
    t1, o1 = bench(m, True)
    # Same math either way: a slower path must not also change results.
    diff = ((o0.float() - o1.float()).norm() / o0.float().norm()).item()  # relative L2
    print(f"{m:>5} {t0:>12.1f} {t1:>14.1f} {100 * (t1 / t0 - 1):>+8.1f} {diff:>8.3g}", flush=True)
sys.exit(0)

Tests

• New: tests/test_mega_mxfp4_method.py::test_real_row_count_only_reaches_fused_moe_for_a_shared_column. With no fused shared column the count is not forwarded; with one it is. It fails on main and passes with this change.
• Passing: tests/test_mega_mxfp4_method.py, tests/test_pass_scheduled_tokens.py, tests/model_ops/test_dp_gather_routed_rows.py, tests/model_ops/test_shared_expert_dispatch.py (40 passed).
• black and ruff are clean.

Not done yet (draft)

• End-to-end Kimi-K3 AgentX C4 / C16 rerun with this change.
• Clean (idle-GPU) rerun of the unit-level numbers (table above).

Scope: only the Mxfp4MoEMethod FlyDSL tail. The other two num_local_tokens call sites (UnquantizedFusedMoEMethod, Iq2rMoEMethod) are untouched: they are not MXFP4 a4w4, so the inline-sort rows never apply to them.

cc @valarLip (author of #2725)

---

This PR was prepared with AI assistance (Claude Code): root-cause analysis, the fix, the test, and the benchmark. All results were run and reviewed by the author.

🤖 Generated with Claude Code

[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 2756 --add-label <label>

[valarLip]

should fixed by moe side, instead of here

[gbyu-amd]

agree. Just confirm, do you mean fix in somewhere else in ATOM or fix in AITER kernel side?

[gbyu-amd]

Fixed on the aiter side instead: ROCm/aiter#7020 — the inline-sort sort skips ids outside [0, E), the BM16 sort honours num_local_tokens on device (CUDA-graph safe), and fused_moe no longer discards the tuned inline-sort row when the count is passed. With it, K3's decode MoE with the count matches the no-count timing (M=1..64 within 1%, was +29-75%), no ATOM change needed.

Converting this PR to draft; I'll close it once aiter#7020 lands.

[gbyu-amd]

Superseded by ROCm/aiter#7020 (merged as f2d9434bbc): the MXFP4 inline-sort path now accepts num_local_tokens and skips -1 ids, so ATOM keeps passing the count and K3 keeps its tuned rows. Closing as promised.