Size MoE activation kernel grid to the row count (+3.79% decode throughput)

[shanmugamr1992]

What changed and why

The row-parallel Triton activation kernels in
megatron/core/inference/moe/activations.py launched a hardcoded
NUM_BLOCKS = min(M, 512) CTAs. At MoE decode shapes that starves the SM array:
on a GB200 (148 SMs), 512 CTAs × 128 threads = 65,536 threads against a capacity
of ~303,000, so roughly 22% occupancy.

The consequence is that _silu_mul_bounded_kernel — a pure elementwise op with no
algorithm to be slow at — ran about 10× its memory-bandwidth floor: 15.8 µs per
layer for the 9.44 MB it actually moves, against a 1.57 µs floor at this machine's
measured 6.081 TB/s streaming ceiling. Other tuned elementwise kernels in the same
trace cost 1.26–2.60 µs, which is what identified this as launch geometry rather
than an inherent cost.

This sizes the grid to the row count instead, one CTA per row, capped at a power of
two.

Two constraints shape the fix:

• NUM_BLOCKS is a tl.constexpr, so every distinct value costs a fresh JIT
  compile. The cap therefore stays a small set of fixed powers of two rather than
  the raw row count.
• The grid is baked into the CUDA graph at capture, so it is derived from the
  capture-time buffer bound M, never from a per-step token count. Graph replay is
  unaffected.

Measured gain

+3.79% end-to-end decode throughput, 22833.56 → 23698.83 tok/s.

Metric	OFF	ON	Δ
Throughput (tok/s)	22833.56	23698.83	+865.27 (+3.79%)
Avg latency (ms)	11209.7	10789.4	−420.3
TPOT (ms/tok)	11.2116	10.8023	−0.4093 (−3.65%)

Baseline is the OFF arm of the same allocation, not a previously recorded number
(see Protocol).

Protocol

• 1 node × 4×GB200 (OCI oci-hsg), Qwen3-30B-A3B, BF16.
• EP4 / TP1 / PP1, --transformer-impl inference_optimized, NVLS dispatcher,
  full-iteration CUDA graphs.
• Batch size 256, OSL 1024, gsm8k, num_input_tokens_avg = 60.86.
• 2 warmup + 5 timed iterations per arm, so graph capture is excluded from all
  timed numbers.
• Four arms OFF / ON / OFF / ON, back to back in one Slurm allocation (job
  5931387), server restarted per arm.

Arm	Config	Throughput (tok/s)	TPOT (ms)
1	OFF	22803.99	11.2261
2	ON	23730.39	10.7879
3	OFF	22863.13	11.1971
4	ON	23667.27	10.8166

Arm separation: min(ON) = 23519.8 > max(OFF) = 22882.8 over all 20 timed
iterations — every ON iteration beats every OFF iteration, so the arms do not
overlap. Pairwise deltas are +4.06% (1→2) and +3.52% (3→4).

Arms were run in one allocation deliberately: this same code measured 0.66%
slower in a different session/node for identical binaries, which is the same order
as many candidate wins.

Where the time went

Confirmed with a second Nsight Systems capture under the change:

	Baseline	This PR	Δ
Launch geometry	grid 512 × block 128	grid 8192 × block 128	—
Launches in trace	83520	83712	+0.2% (step-count noise)
Avg per launch	17.158 µs	7.644 µs	−9.514 µs (−55.4%)

Launch count is unchanged and only gridX moved, which is the signature of a
geometry retune rather than a fusion.

Predicted-vs-measured conversion:

Step	Value	Ratio
Decision-gate prediction	613 µs/step	—
Kernel saving × 48 layers	457 µs/step	0.75 of the gate
Measured end-to-end (TPOT)	409 µs/step	0.90 of the kernel saving

The kernel-to-e2e conversion is 0.90, consistent with the kernel sitting on the
serial chain. The shortfall against the gate was the gate's optimism about
achievable kernel time — it assumed ~2× the bandwidth floor and the kernel landed
at 4.9×.

Correctness

Bit-exact. The kernel is elementwise with no cross-CTA reduction, so the
row-to-CTA mapping cannot change per-element math; bit-exactness is structural
rather than empirical.

Verified anyway:

• 30/30 shape × seed cases produced byte-identical output to the old geometry,
  zero differing elements. 5 shapes × 6 seeds, covering M = 1024/2048/4096/16384
  and live-row counts both below and above the old 512 cap.
• Coherence: four fixed temperature-0 prompts produced byte-identical output
  across all four A/B arms (md5 52f4690b327f59de2e0689c0c3b63b64), so ON is
  indistinguishable from OFF end to end.

No ulp budget is needed because there is no numerical difference to bound.

Kill switch

MCORE_MOE_ACTIVATION_GRID_CAP=512 restores the previous geometry exactly. The
variable accepts any power of two; the default is 8192.

Scope and risks

• Measured only on the bounded_silu_mul path (--swiglu MoE with
  --transformer-impl inference_optimized) at BS256 decode on GB200.
• _squared_relu_kernel shares the same min(M, 512) pattern but is not
  changed here, to keep one mechanism per PR. It is the obvious follow-up.
• Not measured on non-GB200 hardware, in training, or with EP != 4. The change is
  hardware-neutral in principle — it removes a cap, it does not tune to an SM
  count — but the 22%-occupancy figure that motivates it is GB200-specific.
• Larger grids mean more CTAs that early-exit on the n_rows bound at small batch.
  At very small batch the old geometry may be marginally better; the cap is the
  escape hatch if that ever shows up.

Artifacts

• Ledger entry: EXP-01 in skills/run-qwen-model/EXPERIMENTS.md.
• A/B: job 5931387, runs/exp01-swiglu-grid-20260806-120408/arm_{1,2,3,4}_{OFF,ON}.*.log
• Correctness: jobs 5931507, 5931659
• Profiles: baseline runs/qwen-30b-nsys-20260806-113411/mcore_profile.{nsys-rep,sqlite},
  this change runs/exp01-nsys2-124241/mcore_exp01.{nsys-rep,sqlite}
• All under /lustre/fsw/portfolios/coreai/users/shanmugamr/agents-space/

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.