GLM-5.2 SparseMLA TileLang backward returns NaN gradients from finite inputs

Summary

While bringing up GLM-5.2 RL training with Slime, we repeatedly hit native Megatron/Slime NaN checks in the local gradient norm before data-parallel communication. The issue was eventually isolated to the GLM SparseMLA TileLang backward kernel:

• slime_plugins/models/glm5/ops/sparse_mla.py
• slime_plugins/models/glm5/ops/tilelang_sparse_mla_bwd.py

A saved single-call replay had finite inputs to sparse_mla_bwd(...), but TileLang returned nonfinite tl_dq / tl_dkv. A small PyTorch reference backward over the same active positions returned finite dq / dkv and matched TileLang on the finite overlap.

Environment / setup

• Slime commit tested: 8f732538ac78726b0b3b6b2e4e53eb610fdf1cbc
• Model family: GLM-5.2 / glm_moe_dsa
• Training checkpoint: BF16 Megatron torch_dist
• Rollout checkpoint: FP8 HF/SGLang
• Cluster shape: 16 x H200:8
• Parallelism shape where reproduced: TP=4, PP=8, CP=4, EP=16, ETP=1
• Native NaN checks enabled (check_for_nan_in_loss_and_grad=True)
• TIS disabled during the narrower repro branch, to rule out TIS ratio instability
• Router auxiliary load balancing disabled during the narrower repro branch, to rule out router aux loss as the source

This was first observed in RL training, but the final dump/replay points to the SparseMLA backward kernel rather than reward/PPO/optimizer/distributed collectives.

Symptoms

The training step fails after actor-train microbatches with errors like:

found NaN in local grad norm for bucket #0 in backward pass before data-parallel communication collective

Earlier diagnostics showed finite policy-loss inputs:

• finite advantages
• finite old/train/rollout logprobs
• ratio=1.0
• finite pg_loss_pre_reduce
• finite local loss masks on response-carrying CP shards

Module backward hooks then narrowed the first bad area to GLM DSA/MLA attention, especially around DSAMLASelfAttention and SparseMLA backward.

Strongest repro evidence

We instrumented SparseMLA.backward and dumped one bad call immediately after sparse_mla_bwd(...) returned nonfinite outputs.

For the saved bad call:

• inputs were finite:
  • q
  • kv
  • indices
  • tl_out
  • grad_output
  • tl_lse
• TileLang sparse_mla_bwd(...) returned:
  • tl_dq: 829,440 NaNs
  • tl_dkv: 3,144,384 NaNs
• PyTorch reference backward on the same active response positions returned fully finite:
  • ref_dq
  • ref_dkv
• On finite overlap, TileLang and reference matched closely; for tl_dkv the max absolute diff was about 1.5e-8.

We also checked whether this was only padded sparse-index handling. There were many indices == -1 entries in some earlier diagnostics, but the decisive dump showed NaN query rows on real response-token spans with valid top-k indices and nonzero grad_output, not just padding-only rows.

Workarounds tested

These did not remove the nonfinite TileLang backward outputs:

1. Guarding Indices == -1 loads/stores in forward/backward.
2. Starting invalid score lanes as finite zero and masking probabilities after score GEMMs.
3. Disabling router aux/load-balancing loss.
4. Trying lower CP shapes where memory allowed; CP=2 still reproduced the local-grad NaN, CP=1 OOMed before backward.
5. Disabling TIS to remove TIS ratio instability from the critical path.

A diagnostic confinement patch that replaced nonfinite tl_dq / tl_dkv with zero allowed the train step to complete, which confirmed SparseMLA backward was the first nonfinite source. That was only a diagnostic, not a correctness fix.

A more correctness-oriented fallback then did this:

1. Run fast TileLang sparse_mla_bwd(...).
2. If tl_dq / tl_dkv are finite, keep them.
3. If they contain NaN/Inf, recompute the same backward call with a PyTorch reference implementation.
4. Replace only nonfinite TileLang entries with the reference values.

With that fallback, the run completed the first train update with native NaN checks enabled and logged finite optimizer metrics, e.g.:

train/grad_norm=10.70561957699524
train/loss=0.0074499137699604034
train/pg_loss=0.0074499137699604034
train/ppo_kl=0.0
train/global_batch_size=64

However, this fallback is too slow for production training because it recomputes the PyTorch reference on large bad events.

Suspected area

The active suspect is the TileLang SparseMLA backward implementation for this GLM DSA shape, specifically sparse_mla_bwd(...) in tilelang_sparse_mla_bwd.py.

Relevant code path:

# slime_plugins/models/glm5/ops/sparse_mla.py
tl_dq, tl_dkv = sparse_mla_bwd(q, kv, tl_out, grad_output.contiguous(), indices, tl_lse, sm_scale=scaling)

The wrapper currently has a hard-coded D = 512 in tilelang_sparse_mla_bwd.py, and the kernel uses a sparse top-k/block path with TileLang GEMMs and atomic dKV accumulation. I am not sure yet whether the root cause is the D/D_tail assumption, a particular CP/local sequence shape, top-k/block masking, Lse/probability math, atomic dKV accumulation, or another TileLang kernel corner case.

Request

Could maintainers help confirm one of the following?

1. Is this a known issue in GLM-5.2 SparseMLA TileLang backward?
2. Is there an official flag/config to disable the TileLang SparseMLA backward path or use a reference/safer backend for GLM-5.2 while debugging?
3. Are there expected constraints on TP/PP/CP/EP, top-k, D/D_tail, or sequence layout for this kernel that we may be violating?
4. Would you prefer a minimal standalone repro script around a dumped sparse_mla_bwd(...) call, or a PR adding a finite-check/reference test for this kernel?

The short-term correctness fallback works, but the real fix should probably make TileLang SparseMLA backward finite directly, or provide a supported granular fallback path.

Start with slime_plugins/models/glm5/ops/sparse_mla.py and slime_plugins/models/glm5/ops/tilelang_sparse_mla_bwd.py, focusing on sparse_mla_bwd(...), its hard-coded D=512, masking, Lse/probability math, and atomic dKV accumulation. Reproduce the saved single-call inputs and compare TileLang outputs with the PyTorch reference backward. Done means finite dq/dkv for the reported GLM-5.2 shape, or a supported safer fallback with a finite-check/reference test.