[xytpai]

Summary — general gfx1250 ordinary-TP route optimization

Reduce contended global route atomics with workgroup histograms, now generalized
to capacity-based dispatch for up to 4096 routed experts.

• gfx1250 only; wave32. EP, non-gfx1250, small calls and over-budget expert sets
  retain their existing paths.
• Private TP LDS size rounded to 256 buckets; 512 minimum, 4096 maximum (16 KiB).
  Existing E<=512 storage and dispatch are preserved.
• S = ceil(E / 256): enable at num_routes >= 256*S*max(4, 2*S).
  No model-name or exact Kimi-shape branch.
• Kimi-K3 TP8 retains 896 experts / top-16, with routed K=3584 and
  intermediate=384. This now reaches LDS at 512 actual input rows or more.
• Repeated real Kimi-sized full-MoE measurements show approximately
  4.9%–8.2% lower operator latency for the tested prefill cases, with
  bitwise old/new output agreement and unchanged non-route GPU dispatch.
  No E15 TP8+DCP8 full-model or MLPerf Endpoints gain is claimed.

Latest extension commit: 3dbb4307855a99589fb4ffe464123ac72754e568.

Related: https://github.com/wuhuikx/LLM-Booster/issues/188

The generalization and its correctness/performance boundaries are detailed next.
The complete earlier MiniMax report is retained afterward as historical evidence.

Generalization for wide expert sets, including Kimi-K3 (2026-10-29)

This follow-up extends ordinary TP routing beyond the original 512-expert
limit without changing EP scratch, GPU wave size, expert selection, quantization
or GEMM code. It is not a Kimi-specific shape dispatch.

Capacity and dispatch

• Private TP counter storage is rounded to 256 buckets, with a 512-bucket minimum
  and a 4096-bucket / 16 KiB maximum. Larger expert sets keep the old router.
• E≤512 reuses the original _RouteCntStorage, and its dispatch gate is unchanged.
  The EP/LUT factories and their 512-bucket capacity remain unchanged.
• Let S = ceil(E / 256). Enable the TP LDS route when:
  routes >= 256 * S * max(4, 2*S), on gfx1250 only.
• This conservative measured gate accounts for extra histogram scans as the
  expert table gets wider. No equality test on model name, hidden dimension,
  896 experts, or a particular token count.
• For Kimi-K3 E=896, top-k=16, this is M≥512 actual route-input token rows.
  Smaller ordinary calls retain the original implementation. Offered request
  concurrency is not M: speculative verify rows, padding and prefill chunks
  must be read from the actual execution.
• Three unconditional workgroup barriers and global atomic interval reservations
  are unchanged. There is no grid-wide wait, tensor/pointer cache, or persisted
  counter state.

Kimi's two shared experts run as a separate MLP; they are not in this routed
histogram. TP8 preserves all 896 routed experts and partitions intermediate
3072→384. The routed MoE hidden dimension is 3584, distinct from model hidden
7168. The tested complete MoE shape therefore uses E896 / top16 / K3584 / I384,
SiTUv2(beta=4, linear_beta=25), BF16 output.

Actual Kimi-sized full-MoE operator A/B

Real aiter.fused_moe, including routing, quantization, both GEMMs and reduction.
Identical weights/activations/top-k per pair. Not a mocked GEMM: a private
test harness selects the old or new real route launcher before graph capture,
leaving all downstream operations unchanged. Each independent process collects
11 alternating-order samples per arm, eight full-MoE calls/graph and five graph
replays/sample. GPU event timing; warm, reused allocations. Compilation,
initialization and reference computation excluded. Some runs drift within a
process, so the table retains both independent-process medians and raw
alternating samples; the percentages are observed, not guaranteed fixed gains.

Format	Actual rows	Process 1: old→new us	Reduction	Process 2: old→new us	Reduction
A4W4	2517	245.459→230.205	6.21%	243.114→231.106	4.94%
A4W4	8192	651.002→604.393	7.16%	650.111→596.705	8.21%
A8W4	2517	271.715→251.869	7.30%	270.773→253.030	6.55%

At the final Kimi gate, M512 A4W4 was 155.859→154.159 us (1.09%) in one
process, not a strong standalone gain claim. An earlier M256 trial was
143.254→143.258 us, effectively unchanged; the final dispatch keeps M256 on
the original route path. That earlier M256 trace exercised a less conservative
prototype gate; it is not evidence that the final public dispatcher selects LDS
at M256. The final-gate M512 test is recorded separately.

• New/old full-MoE outputs are bitwise identical for tested pairs, including
  changed-input graph replay.
• Both arms pass the existing independent reference with unchanged
  logits_diff < 0.01; errors are small and nonzero, not exact-reference claims.
  Maximum among the paired tests' original/changed checks: below 7.2e-6.
• Separate three-call traces per case retain exactly three routing launches.
  All non-route kernel names/grid/block configurations and counts match.
  Example repeated A4W4 M8192 route durations: old 77.195/82.682/80.279 us;
  new 16.944/16.744/16.944 us. These short traces are attribution controls,
  not the source of the full-MoE timing table.
• A8W4 E896/K3584/I384 has no matched tuned grouped CSV row on this test stack;
  both arms use the same fallback GEMM settings. A4W4 uses the existing table.

These are TP8 per-rank shape operator measurements on one gfx1250 GPU.
No two-tray TP8+DCP8 Kimi model or MLPerf Endpoints performance result is claimed.
The original MiniMax model-trace report below remains historical and is not
relabelled as a Kimi result. There is no quality/inline-accuracy signoff for Kimi.

Generality and negative results

The scalar dispatch/size rules are tested at E1..4096; E4097 is fallback.
Real route sweeps include E129/256/512/513/768/769/896/1024/1025/1536/1537/
2048/2049/3072/4096. Uniform and skewed top-16 inputs use distinct expert IDs
within each token; all-same-expert cases are separately labelled stress tests.

Initial uniform E896 route microbenchmarks (existing counter zeroing included):

• M2517: 29.704→9.658 us, 3.08x.
• M8192: 83.820→20.655 us, 4.06x.

These are route-only results, not full-MoE/model speedups. Microbenchmarks
use preallocated buffers, nine samples with alternating order, 24 calls/graph
and five replays/sample. The microbenchmark prototype and final kernel body
agree for E896; final dispatch is narrower.

Early capacity-only gate extensions were rejected:

• E2048/M512 uniform: 7.266→8.712 us.
• E4096/M1024 uniform: 10.053→13.777 us.
• E4096/M2517 uniform: 16.899→18.906 us.
• An intermediate threshold at E4096/M4096: 24.318→24.613 us.

The final sweep-scaled gate excludes those regimes. At its tested wide-expert
boundaries, uniform routes improved, e.g.:
E896/M512 9.685→6.408 us, E2048/M2048 18.766→13.537 us,
E4096/M8192 43.235→39.210 us. E≤512 existing-path controls were broadly unchanged
(small measurement variation retained).

Tests and compatibility

• Full combined regression selection: 102 passed, 45 deselected.
  Includes 98 route tests, two real Kimi SiTUv2 public-MoE comparisons and two
  existing MiniMax decode regressions.
• Route checks cover exact counts, row-slot bijections, per-expert ownership,
  nonzero initial counters, changed/poisoned graph buffers, all-dropped/negative
  IDs, distinct top-16 inputs, and both sides of the capacity-dependent gate.
• CPU scalar checks: all 4096 legal expert counts plus invalid bounds; all
  E≤512 thresholds match the earlier PR exactly.
• Test reference allocations now explicitly use CPU: another existing MoE
  test module sets the default device to CUDA. The combined run exposed this
  pre-existing test reference assumption; no numerical tolerance was relaxed.
• A8W4 test overrides are restored through monkeypatch.
• Black, Ruff and git diff --check pass.
• A separately prepared minimal backport to the recipe's aiter 7061cbd
  retains its old direct FlyDSL launcher call style. On f01/FlyDSL0.4.0 it passes
  seven route/graph cases and both real SiTUv2 A4W4/A8W4 M2517 full-MoE
  reference/exact-old-vs-new checks. This is not an execution of the
  gfx1250-atom-k3-20261027 image or the E15 TP8 server.
• The old 7061 fallback lacks the newer negative-ID guard; backport tests keep
  that fallback inside its original supported input domain.
• No gfx942/gfx950 hardware test or EP performance claim. Their code/dispatch
  is unchanged.

Reproduce the new committed tests:

HIP_VISIBLE_DEVICES=1 python3 -m pytest -q -p no:cacheprovider \
  op_tests/test_moe_route_rows.py \
  op_tests/flydsl_tests/test_flydsl_grouped_gemm.py \
  -k 'route_capacity or wide_tp_route or route_rows_graph or kimi_tp_route_histogram_public_path or grouped_swiglu_m3_tp_decode'

Validation uses f01 gfx1250, 256 CUs, physical GPU1, pinned FlyDSL0.4.0 and
image rocm/fw-bringup-mlperf:gfx1250-atom-20261023 (digest
bad23258ca8d0852fd1547c41005ded0964c24e264cbb5ffe66f8e0db183797c).
The incremental baseline is this PR's previous head 941e32b774c7.
No clocks/power or model configuration changes were made, and no other users' jobs were stopped.
Raw samples, traces and excluded failed attempts remain in
~/bench/kimi_route_general_20261029/ on f01 and local
temp/kimi_route_general/. No benchmark reports are added to docs.

Historical initial E≤512 / MiniMax-M3 report (previous head 941e32b)

Summary

Reduce global-atomic contention when building ordinary TP MoE route rows on
gfx1250 (wave32). Add a workgroup-local histogram, then reserve one global
counter interval per nonempty expert/workgroup instead of one device atomic
per route. Local atomic ranks select unique slots within each reserved interval.

Related: https://github.com/wuhuikx/LLM-Booster/issues/188

This is separate from #6894's small-TP single-block route/psum fusion, EP/LUT
routing, and the existing grouped-MoE GEMM/quantization kernels.

Dispatch and contracts

• gfx1250 only, 0 < E <= 512.
• Enable when num_routes >= 4 * ceil(E / 256) * 256, amortizing one or two
  256-bucket sweeps. Other configurations retain the original route launcher.
• Three unconditional workgroup barriers; no cross-workgroup polling,
  cooperative-launch assumption, or grid-wide barrier.
• Existing global counter initialization is unchanged. Caller-provided counters
  retain their initial offsets; negative expert IDs produce the existing
  dropped-route sentinel and claim no slot.
• Per-expert route order remains unspecified, as with the original atomic
  kernel. Expert selections, route weights, GEMM arithmetic and quantization
  algorithms are not changed.
• EP/LUT routing, route-ksplit callsites, small fused TP dispatch and gfx942/gfx950
  keep their existing paths. No clock, power or model configuration changes.
• Valid positive expert IDs remain [0,E); the new bounded implementation does
  not access counters for larger IDs, but this is not a new supported public
  input contract. Counter/row integer capacity remains the caller's obligation.

Real MiniMax-M3 trace confirmation

MiniMax-M3-MXFP4, TP4 on gfx1250, EAGLE3 with 3 speculative tokens, FP8 KV,
ptpc_fp8 non-MoE linears, prefix caching. Same pinned ATOM with the earlier
SwiGLU FP32 fix in both arms. Real speculative acceptance; no fixed-acceptance
option or output/work elision.

Two profiled requests per arm, all four ranks. Every selected prefill annotation
has 2517 tokens, context 2517. Each rank/recording retains 57 route calls.

Batch	Base rank medians, moe_route (us)	Candidate rank medians, moe_route_lds (us)
0	32.727 / 32.408 / 33.369 / 33.289	3.685 / 2.803 / 3.484 / 3.044
1	33.889 / 33.649 / 34.210 / 32.848	2.964 / 2.883 / 3.484 / 2.963

Median of these rank-local medians: 33.329 -> 3.004 us, approximately
11.09x, or 91.0% lower route-kernel duration. Concurrent ranks and the
two recordings within a service are not independent service-start replicates.

Non-route GPU kernel names, grid/block/shared-memory configurations and counts
match exactly between arms. The expected route-kernel replacement is the only
launch-signature difference in the selected prefill windows.

End-to-end result: no established model latency/throughput gain

16 unprofiled matched requests per arm, concurrency 1, 2517 input / 16 output
tokens. Inputs, seeds and token counts match; all 16 response texts match exactly.
Warmups and profiled requests are excluded from request-latency statistics.

Metric	Base	Candidate
Median first nonempty streamed text-chunk latency	88.462 ms	88.119 ms
Median complete request latency	208.597 ms	211.354 ms
Profiled GPU prefill, batch rank-median 0 / 1	83.903 / 83.225 ms	85.017 / 84.496 ms
Profiled host MoE event median, batch 0 / 1	371.136 / 371.476 us	376.407 / 376.301 us

This PR does not claim an end-to-end speedup. The kernel improvement is
visible in the real model, but this host-bound single-request workload does not
convert it into a reliable request-latency improvement. Complete request latency
was slightly worse in this one A/B. No throughput, long-context, other
concurrency or model-quality equivalence claim is made.

Operator microbenchmarks

Physical GPU1; preallocated reused buffers; nine samples per configuration,
alternating baseline/candidate order; each HIP graph contains 32 calls and each
sample replays it five times. Includes existing counter zeroing plus actual
route execution, excludes compilation and allocation. Same input per pair.
These are warm-cache route measurements, not full MoE timings or model results.

Experts	Tokens	Top-k	Distribution	Base us	LDS us	Speedup
129	2517	5	uniform	18.330	4.374	4.19x
129	2517	5	four random + shared expert	31.039	4.675	6.64x
129	2517	5	all routes one expert	137.270	4.445	30.88x
129	2517	5	every third route dropped	15.646	4.347	3.60x
129	8192	5	uniform	54.747	6.740	8.12x
129	8192	5	shared	92.616	6.840	13.54x
512	2517	5	uniform	13.917	5.325	2.61x
512	2517	5	shared	33.717	5.216	6.46x
512	2048	1	uniform	5.104	4.325	1.18x
257	2048	1	uniform	5.279	4.149	1.27x
384	2048	1	uniform	5.273	4.100	1.29x

A second alternating sweep reproduced the target shared shape at
31.033 -> 4.618 us (6.72x).

Negative results retained and used to narrow dispatch

The initial routes >= 1024 gate regressed uniform low-route/high-expert cases:
E512/1024 routes 3.934 -> 4.283 us, E257/1028 routes
3.877 -> 4.233 us, E384/1536 routes 3.934 -> 4.165 us.
The final bucket-sweep gate retains the original launcher for these cases.
Final-gate boundary checks at 2048 routes passed for E257/384/512 and all four
distributions. This is a conservative measured dispatch, not an exact M3 shape
specialization.

Correctness and testing

• 80 GPU pytest cases passed on the final source. Counts match an integer
  torch reference exactly; per-expert row slots form the expected bijection.
  Includes E1/17/129/512/513, short/tail/large route counts, uniform/shared/hot/
  negative-ID distributions, preinitialized nonzero counters, poisoned output,
  eight graph replays with changed IDs and counters.
• Actual public grouped-MoE A/B: A8W4 small; A4W4 decode1, boundary205 and
  prefill2517, using real GEMMs and the existing independent reference and
  unchanged logits_diff < 0.01 gate.
  Original outputs, references, changed-input outputs and graph outputs are
  bitwise identical across arms for all four cases.
• Caller-preallocated output checks and five invalid output-buffer checks match.
  A small A8W4 case also passed torch.compile(backend="inductor", fullgraph=False);
  this is not fullgraph or all-shape compile coverage.
• Non-route kernel launch configurations/counts match in the operator traces.
• The standard benchmark CLI ran successfully: target shared shape 4.663 us,
  count error 0, with the usual @benchmark, candidate dict, run_perftest,
  checkAllclose, and markdown table.
• Black, Ruff and git diff --check passed.
• No gfx942/gfx950 hardware run, EP runtime regression or GSM8K evaluation is
  claimed. Unchanged architecture/EP paths were checked in source.

Commands:

HIP_VISIBLE_DEVICES=1 python3 -m pytest -q -p no:cacheprovider \
  op_tests/test_moe_route_rows.py
HIP_VISIBLE_DEVICES=1 python3 op_tests/test_moe_route_rows.py \
  -b 2517 --experts 129 --topk 5 --distribution shared
python3 -m black --check aiter/ops/flydsl/kernels/moe_route_maps.py \
  aiter/ops/flydsl/moe_kernels.py op_tests/test_moe_route_rows.py
python3 -m ruff check --no-cache aiter/ops/flydsl/kernels/moe_route_maps.py \
  aiter/ops/flydsl/moe_kernels.py op_tests/test_moe_route_rows.py

Provenance and limitations

2026-10-28 UTC; MI455/gfx1250, 256 CUs/device; PyTorch
2.11.0+rocm10.1.0a20260811-bkc.20260820; HIP 7.16.26315; FlyDSL 0.4.0;
image rocm/fw-bringup-mlperf:gfx1250-atom-20261023.
Separate source mounts and JIT caches, both PYTHONPATH and AITER_META_DIR fixed
to each arm, no cross-revision .so copying.

Measured base: 4aace7695a6da9c7f13c562d9301e9ba11f97bbe.
Both modified production files are blob-identical on main
b0043b6c778f0378e8892311e575e1ff898039a6.
The final dispatch narrowing leaves the measured M3 path and kernel body
unchanged; final source was rerun through the route pytest suite and boundary
microbenchmarks. The entire newer-main model stack was not re-benchmarked.

Excluded validation attempts are retained: an initial stdin-formatting mistake
produced empty experimental files and pytest collected zero tests; sources were
restored from the pinned baseline and saved implementation, then real GPU tests
were run. A temporary pytest parameterization of the benchmark wrapper failed
collection; final pytest/CLI entrypoints were both executed successfully.

Raw samples, source/input hashes, model manifests, traces, regression outputs,
and negative experiments remain outside the repository. No docs/benchmarks files
are added. The prior combined #7080/#7082/#7090 quality caveats are not resolved
or silently reclassified by this PR.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 7095 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[xytpai]

Extended this PR in 3dbb4307855a99589fb4ffe464123ac72754e568 with general, capacity-based ordinary-TP routing:

• Up to4096 experts, private LDS allocation; EP scratch and E<=512 behavior remain unchanged.
• Kimi-K3 E896/top16 now qualifies at512 actual token rows, not based on offered concurrency.
• Two paired full-MoE runs: A4W4 M2517 about4.94–6.21% lower, M8192 about7.16–8.21% lower; A8W4 M2517 about6.55–7.30% lower.
• Original and changed-input graph outputs are bitwise identical; nonroute kernel configurations/counts match.
• Combined regression selection102passed; a separate minimal7061cbd88 backport passes seven route cases and A4W4/A8W4 realMoE reference+exact comparison checks.
• These are operator measurements with TP8 per-rank shapes on one GPU, NOT an eight-GPU/two-tray Kimi model or MLPerf result.
• Initial regressing wide-expert thresholds and all earlier MiniMax results are retained in the description.

[yadaish]

LGTM