fix: zero-init fp32 UE8M0 activation scales in the masked MoE down-GEMM

[yueming-yuan]

Problem

Qwen3-30B-A3B RL with --rollout-fp8 on 4xB300 dies during decode CUDA graph capture:

deep_gemm/include/deep_gemm/impls/smxx_layout.cuh:131,
condition: (values[j] & 0x807fffffu) == 0
...
tvm.error.InternalError: CUDA driver error (jit/handle.hpp:154): 719 (CUDA_ERROR_LAUNCH_FAILED, unspecified launch failure)

The 719 is only the sticky symptom. The real error is the device-side assert above: 0x807fffff masks an fp32's sign bit and mantissa, so deep_gemm requires every scale factor to be a positive power of two (UE8M0) before it packs them into its own layout.

Cause

In _varlen_deep_gemm_silu_mul_quant, the fp32 fallback branch allocates the activation-scale buffer uninitialized:

down_input_scale = torch.empty((E, N, G), device=..., dtype=torch.float32)
silu_and_mul_masked_post_quant_fwd(..., masked_m, scale_ue8m0=...)

The Triton kernel is correct — it does round to a power of two (tl.exp2(tl.ceil(tl.log2(...)))) — but it only writes rows below masked_m. deep_gemm's scale-factor layout transform validates the entire padded tensor, including the inactive rows, so whatever was left in that memory trips the assert.

Which branch runs is decided by:

if N % 4 != 0 or G % 4 != 0 or D // 8 < E:
    use_jit_ep_activation = False

The JIT fast path allocates packed int32 scales, which never go through the fp32 validation, so it is immune. Only G % 4 != 0 reaches the exposed fp32 branch:

model	moe_intermediate	G = D / 128	branch	scale dtype	affected
Qwen3-30B-A3B	768	6	fp32 fallback	fp32	yes
DeepSeek-V3 class	2048	16	JIT fast path	packed int32	no

That is why this has gone unnoticed: it is a latent uninitialized-memory bug that only models with G % 4 != 0 can reach.

Evidence

1. Minimal reproduction — same masked grouped GEMM, only the inactive rows differ:

g, mmax, k, n = 8, 128, 768, 2048     # Qwen3-30B-A3B down-proj, G = 6
masked = 8                            # only 8 of 128 rows active
asf = torch.full((g, mmax, G), 2.0 ** -5, dtype=torch.float32)   # valid UE8M0
if mode == "garbage_padding":
    asf[:, masked:, :] = 0.3          # non-power-of-two, inactive rows only
deep_gemm.fp8_m_grouped_gemm_nt_masked((afp8, asf), (bfp8, bsf), d, masked_m, masked)

• all rows valid UE8M0 -> PASS
• only inactive rows set to a non-power-of-two -> the exact assert above

So the transform does validate padding.

2. Instrumented run — printing every scale operand of both MoE GEMMs at the failure point:

operand	dtype	shape	UE8M0 violations
gateup A (from DeepEP dispatch)	int32 packed	(64, 256, 4)	n/a
gateup B (w13_scale)	int32 packed	(64, 1536, 4)	n/a
down A (down_input_scale)	fp32	(64, 256, 6)	56532 / 98304
down B (w2_scale)	int32 packed	(64, 2048, 2)	n/a

The only fp32 scale in the whole step is down_input_scale, its trailing dim is G = 6 (confirming the fallback branch), and ~58% of its entries are not powers of two.

Fix

torch.empty -> torch.zeros. 0.0 is 0x00000000, so it satisfies the assert, and the value in inactive rows cannot affect the result because masked_m excludes them from the GEMM. Cost is one 393 KB memset per call (64x256x6x4 B) against a multi-GB GEMM; under CUDA graphs it is captured as a graph node and replays correctly.

Rejected alternatives: emitting packed int32 from this branch would require changing the Triton kernel's output format, and the fast path requires G % 4 == 0 precisely because it packs four scales per int32 — G = 6 would first have to be padded to 8. Bounding the transform by masked_m inside deep_gemm is arguably the more correct fix (validating inactive rows is pointless work) but that is upstream of this repo.

Verification

Qwen3-30B-A3B RL on 4xB300, --rollout-fp8, bf16 training, fp32 optimizer, TP4/PP1/CP1/EP4, colocated, dapo-math-17k, 8192 response length. Before the fix the engine never survives decode CUDA graph capture:

	before	after
smxx_layout.cuh:131 asserts	1125	0
719 unspecified launch failure	present	0
decode CUDA graph capture	fails here	passes
rollout	never reached	completes, response_len/mean 6312
training steps completed	0	3 / 3

Three full 3-step runs, differing only in how the paused training actor is backed up, all with zero asserts and no NaN:

actor offload	train_rollout_logprob_abs_diff (steps 0/1/2)
cpu	0.0329 / 0.0336 / 0.0350
node-local disk	0.0325 / 0.0340 / 0.0349
node-local disk + optimizer-state streaming	0.0328 / 0.0336 / 0.0351

Stable across steps and identical across the three configurations, i.e. the rollout/train logprob agreement is unaffected.

(An earlier revision of this section reported only step 0. That run was truncated by an unrelated problem — the image's torch_memory_saver predated its disk-backup backend, so disk-mode offload silently restored garbage. With the matching torch_memory_saver installed, all three configurations run to completion.)

A companion PR, #32386, targets main, where the default plain-silu path has since moved to packed int32 and is no longer affected, but the gemm1_alpha branch still carries the same pattern.