fix: zero-init fp32 UE8M0 activation scales in the gemm1_alpha masked path

[yueming-yuan]

Problem

deep_gemm requires every fp32 scale factor handed to its scale-layout transform to be a positive power of two (UE8M0). It enforces this with a device-side assert:

deep_gemm/include/deep_gemm/impls/smxx_layout.cuh:131,
condition: (values[j] & 0x807fffffu) == 0

0x807fffff masks the sign bit and the mantissa, so anything with a nonzero mantissa fails. Because it is a device assert, the CUDA context is poisoned and the user-visible error is the unhelpful 719 (CUDA_ERROR_LAUNCH_FAILED, unspecified launch failure).

The transform validates the entire padded scale tensor, including rows at or beyond masked_m. Any masked kernel that allocates its fp32 scale buffer with torch.empty and only writes the active rows therefore leaves the inactive rows holding whatever was in that memory, which is almost never a power of two.

Where this is reachable on main

In _varlen_deep_gemm_silu_mul_quant, the gemm1_alpha branch:

use_packed = (
    deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
    and num_real_tokens is not None
    and G % 4 == 0
    and D % (group_size * 4) == 0
)
down_input_scale = torch.empty(
    (E, G // 4, N) if use_packed else (E, N, G),
    device=hidden_states_device,
    dtype=torch.int32 if use_packed else torch.float32,
)

When DEEPGEMM_SCALE_UE8M0 is on but any of the other three conditions fails — most easily G % 4 != 0 — the buffer is fp32 and uninitialized, and silu_and_mul_masked_post_quant_fwd writes only rows below masked_m.

The other two branches are not affected: the swiglu_limit/swizzle branch asserts G % 4 == 0, so with UE8M0 enabled it is always packed int32; and the default plain-silu path routes through per_token_group_quant(..., column_major_scales=True), which lands on the packed-int32 allocation in create_per_token_group_quant_fp8_output_scale and so never reaches the fp32 validation.

Evidence

Minimal reproduction — the same masked grouped GEMM twice, differing only in the inactive rows:

g, mmax, k, n = 8, 128, 768, 2048     # G = k / 128 = 6, i.e. G % 4 != 0
masked = 8                            # only 8 of 128 rows active
asf = torch.full((g, mmax, k // 128), 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)

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

Confirming that inactive rows are validated, and that valid-but-unwritten padding is the whole problem. Reproduced on both B200 (sm_100) and B300 (sm_103) with deep_gemm 0.1.4; the behaviour and the DEEPGEMM_BLACKWELL / DEEPGEMM_SCALE_UE8M0 / DEEPGEMM_NEED_TMA_ALIGNED_SCALES flags are identical on both, so this is not architecture-specific.

Fix

Allocate the fp32 case with torch.zeros. 0.0 is 0x00000000, which satisfies the assert, and the value in inactive rows cannot affect the result because masked_m excludes them from the GEMM. The packed int32 case keeps torch.empty since it never reaches the fp32 validation.

Rejected alternatives: bounding the transform by masked_m inside deep_gemm is arguably more correct — validating inactive rows is pointless work — but that is a change in a different project and needs per-expert row bounds. Emitting packed int32 from this branch is not possible without padding G up to a multiple of 4, which is exactly the condition the packed path already requires.

Verification, and what is not verified

The same class of bug, at the equivalent site on the sglang-miles branch (where the default plain-silu path still used the raw fp32 torch.empty), was reproduced and fixed end to end — see the companion PR #32385. Qwen3-30B-A3B RL on 4xB300 with fp8 rollout went from 1125 asserts and a hard failure during decode CUDA graph capture to zero asserts and three complete training steps, with train_rollout_logprob_abs_diff stable at 0.033-0.035 across three separate 3-step runs (cpu offload, node-local disk offload, and disk offload plus optimizer-state streaming).

I have not executed the gemm1_alpha path on main — I do not have a model configuration that exercises oai-swiglu with G % 4 != 0 on hand. The change here is the same one-line allocation swap, justified by the minimal reproduction above and by reading the branch conditions; a reviewer who runs that path should sanity-check it.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.