Retune MoE token-count kernel block size for decode shapes

[shanmugamr1992]

What changed and why

_count_local_tokens_kernel_persistent (inference MoE dispatch) hardcoded
BLOCK_SIZE = 1024. Triton launches it with 128 threads, so each thread owns 8
(token, topk) pairs and issues its tl.atomic_add increments serially. At
decode shapes that makes the kernel cost far more than its work: 7.95 µs per
launch to count 256 tokens over 32 local experts, i.e. 381.8 µs per decode
step across 48 layers.

The mechanism is per-thread serialisation, not launch geometry or atomic volume.
Measured on GB200, three hypotheses separated with a CUDA-graph-replay harness:

Hypothesis	Test	Result
Launch geometry (152 CTAs, 2 do work)	buffer capacity 16384 → 256, grid 128 → 2	8.22 → 8.21 µs. No effect.
Atomic volume	privatised histogram, 8× fewer atomics	2.7× worse
Per-thread serialisation	sweep BLOCK_SIZE	8.08 → 2.11 µs at 1024 → 128

The kernel is flat in grid and flat in routed pairs (12.31 / 12.30 / 12.29 µs at
2048, 8192, 32768 pairs), which is the signature of a fixed per-thread cost rather
than a throughput limit. Atomics do dominate — removing them takes 4.51 → 1.56 µs —
but the fix is to spread them across CTAs, not to issue fewer.

BLOCK_SIZE=128 is the right point on the curve, not the minimum: it is 3.83× faster
at the decode shape and within 1.3% of the best value at the largest graph bucket,
where BLOCK_SIZE=64 starts to lose badly.

BLOCK	256 tok	1024 tok	4096 tok	16384 tok
64	2.05	2.35	5.51	16.07
128	2.11	2.44	4.66	13.34
256	2.93	3.03	4.56	13.26
1024 (before)	8.08	8.07	8.23	13.57

Measured gain

+2.290% throughput, 23016.42 → 23543.46 tok/s, against the mcore EP4/TP1
baseline at af734982d. TPOT 11.1225 → 10.8735 ms (−0.2490 ms/step); average
latency 11121.5 → 10866.6 ms.

Protocol

1× OCI GB200 node (4 GPU), Qwen3-30B-A3B, EP4/TP1, --transformer-impl inference_optimized, NVLS dispatcher, vLLM grouped-GEMM backend. gsm8k, batch 256,
OSL 1024, num_input_tokens_avg = 60.86, 2 warmup + 5 timed iterations per arm.
Arms ran back to back in the same allocation as OFF/ON/OFF/ON, each with its own
server launch.

Arm	i1	i2	i3	i4	i5	mean
OFF	23027.5	22998.0	23008.3	22993.6	23024.4	23010.36
ON	23580.0	23510.6	23535.9	23548.9	23522.6	23539.59
OFF	23050.6	22997.1	23032.6	23014.0	23018.1	23022.49
ON	23547.3	23517.0	23543.0	23553.0	23576.3	23547.28

Arm separation: min(ON) = 23510.6 > max(OFF) = 23050.6 — every ON iteration
beats every OFF iteration, with a 460 tok/s margin. Pooled sd is 18.0 (OFF) and 23.2
(ON) tok/s, so the win is ~23× the run-to-run noise of this configuration.

Where the time went

Attributed to _count_local_tokens_kernel_persistent, 48 launches/step (one per
layer), in the validated steady-state decode window of the Nsight Systems trace
(125 consecutive passes, 10272.19 µs/step). Launch count is unchanged, so this is a
pure GPU-busy win with no change to scheduling idle.

Predicted 381.8 → 99.7 µs/step (−282 µs). Measured TPOT fell 249 µs/step, a 0.88
predicted-to-e2e conversion. The projection held because the microbenchmark was
validated against the trace first: it reads the shipped configuration at 8.08 µs
against the trace's 7.95 µs, a 1.6% match.

Correctness

Bit-exact. The kernel is an integer reduction via atomic_add, so block size
changes only which thread issues which increment, never the multiset of increments;
integer addition is associative and commutative.

• Exact match against BLOCK_SIZE=1024 in 2280/2280 cases: 10 seeds × 3 buffer
  capacities {16384, 4096, 512} × 7 valid-token counts {0, 1, 7, 63, 256, 1000, 4096}
  × 4 local-expert offsets {0, 32, 64, 96} × blocks {128, 256, 512}.
• Non-vacuous: decode-shape counts sum to 513 over 32/32 populated experts.
• The four temperature-0 coherence prompts returned character-identical
  completions in both arms.

Kill switch

MCORE_COUNT_TOKENS_BLOCK=1024 restores the previous behaviour.

Scope and risks

• Measured only on GB200, Qwen3-30B-A3B, EP4/TP1, batch 256, topk=8, 32 local
  experts. The block-size curve depends on threads per program and on experts per
  rank, so a very different expert count may prefer another point.
• Only the persistent variant changes. The non-persistent path still uses
  BLOCK = 1024, because its grid is ceil_div(max_pairs, BLOCK) and shrinking the
  block would change its launch geometry too — an unmeasured change.
• No effect on prefill-dominated workloads beyond the ≤1.3% shown at 16384 tokens.

Artifacts

• Ledger: skills/run-qwen-model/EXPERIMENTS.md, GATE G4 / EXP-02
• A/B job 5934676, …/agents-space/runs/exp02-count-block-20260806-145539
• Gate microbenchmarks: jobs 5934194, 5934247, 5934417, 5934617
• Trace the attribution is drawn from: …/runs/exp01-nsys2-124241/mcore_exp01.{nsys-rep,sqlite}

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.