[None][perf] Bound Triton MoE JIT specializations via token-dim bucketing and warmup

[arysef]

Description

The TRITON MoE backend (OpenAI matmul_ogs grouped GEMMs, the recommended backend
for gpt-oss MXFP4 on Hopper) JIT-compiles a new kernel specialization the first
time it sees a novel total token count m: kernel meta-parameters such as
block_m are derived from m (triton_kernels/matmul_ogs_details/opt_flags.py),
and Triton additionally specializes the m-derived integer kernel arguments on
divisibility by 16. Under serving traffic with varied prompt lengths, novel m
values keep arriving for the whole lifetime of the server, and each first-seen
specialization stalls the GPU for the duration of a JIT compile (~0.2–2 s).

Measured on gpt-oss-120b (W4A16 MXFP4) on H200 under a randomized-length serving
workload (uniform input lengths in [0.8·ISL, ISL], streaming, concurrency 4–64),
these stalls account for 12.6% of wall time at ISL/OSL 1k/1k conc 64,
22% at 8k/1k conc 64, and ~34% at TP2 8k/1k conc 16 (the per-layer
allreduce propagates either rank's stall to both ranks).

This PR adds moe_config.triton_pad_token_dim (default off): the MoE token
dimension is padded to geometric buckets (8 per power-of-two octave, ≤12.5%
padding, ~6% average) before the grouped GEMMs and sliced back before the
allreduce, which makes the set of observable m values — and therefore the JIT
specialization space — provably finite (97 buckets for max_num_tokens=20000).
Engine warmup then pre-compiles exactly those buckets before the server starts
taking traffic (_run_triton_moe_prewarm), following the same pattern as the
trtllm-gen FMHA JIT warmup (#14851).

The flag is restricted to W4A16 MXFP4, where padded rows provably cannot change
the outputs of real tokens (row-independent grouped GEMMs, BF16/FP16 activations,
no dynamic activation scales); on other quantization modes it logs a warning and
stays off (dynamic per-tensor activation scales would let padded rows shift the
second GEMM's amax).

Measured impact (InferenceX harness recipe, gpt-oss-120b, H200, 1.3.0rc14+patch,
GSM8K unchanged at 0.956 vs 0.957 baseline):

cell	baseline tok/s/GPU @ tok/s/user	with this PR	Δ
1k1k TP1 conc64	2653 @ 47.2	3048 @ 48.8	+14.9% / +3.4%
8k1k TP1 conc64	1392 @ 24.8	1781 @ 28.3	+27.9% / +14.1%
8k1k TP2EP1 conc16	480 @ 69.8	723 @ 94.4	+50.7% / +35.2%

Full 20-cell sweep: +8% to +65% throughput (avg ≈ +25% attributable to this
change alone), with simultaneous interactivity (median TPOT) improvements.
A cold→warm TRITON_CACHE_DIR A/B on the unmodified code reproduces the same
gains, confirming the mechanism; warmup-grid-only (without padding) recovers
almost nothing because Triton's %16 integer specializations are not coverable
by a finite grid of raw (unpadded) token counts.

Note on scope: TritonFusedMoE is on the legacy MoE path, but the TRITON
backend exists only there, and this is a perf fix scoped to that backend —
pattern-identical to disable_finalize_fusion (CUTLASS-specific knob consumed
inside the legacy backend file).

Test Coverage

• tests/unittest/_torch/modules/fused_moe/test_triton_moe_padding.py
  • test_bucket_round_up_properties / test_bucket_iter_is_sorted_and_compact
    (CPU): exhaustive bucket-math properties over [1, 65536] — rounds up by
    ≤12.5%, fixed-point, monotonic, enumeration matches closure.
  • test_triton_moe_padding_bit_identical (@skip_no_hopper):
    torch.equal between padded and unpadded forwards on the W4A16 MXFP4 path,
    parametrized over token counts and bias.
  • test_triton_moe_padding_requires_w4a16_mxfp4 (@skip_no_hopper):
    the flag is ignored on non-W4A16 quant modes.
• Registered in l0_h100.yml (SM90 lane).

PR Checklist

• Config surface follows CODING_GUIDELINES Pydantic rules (StrictBaseModel
  field with description; nested MoeConfig field — no api_stability
  reference changes required, matching precedent [None][chore] Cherry-pick from (#7598) Make low_precision_combine as a llm arg #7898).
• No behavior change with the flag unset (default).
• pre-commit run on all changed files.
• GPU-validated end to end (import smoke, unit tests, serving A/B) on 2xH200.