[avigyabb]

Summary

Lets GLM-5.3-Flash's DSA layers use megatron-core's fused TileLang SparseMLA kernel instead of
the dense fallback, so DSA attention memory is linear in sequence length. Split out of #2341 so it
can be compared against the cuDNN DSA backend on its own.

Why

GLM-5.3-Flash's DSA layers are NoPE MLA (q/k width 512), and its k-pool selection appends the
query's incomplete tail pool, so the top-k width is 2048 + 3 = 2051. The fused absorbed
sparse-attention kernels are specialized for the DeepSeek-V3.2 layout (q/k width 576, top-k width
a multiple of 64). They decline, and megatron-core falls back to a dense [b, heads, sq, sk] FP32
masked softmax: ~128 GiB/GPU at 32k tokens.

Change

glm5_next/dsa.py: _pad_for_fused_absorbed_sparse_attention wraps
dsa_kernels.run_fused_absorbed_sparse_attention inside _forward_with_kpool_topk:

• q/k zero-padded 512 → 576 channels in the RoPE slot. Zero channels add nothing to q·k,
  softmax_scale is passed explicitly, and the value is still key[..., :512].
• top-k indices padded to a multiple of 64 with -1, which the kernels mask out.

Exact. It needs dsa_kernel_backend="tilelang". If the backend still declines, the dense fallback
runs on the unpadded tensors as before. The first call logs whether the fused kernel was used or
declined.

Comparing against cuDNN

• k-pool path: _forward_with_kpool_topk already forces megatron-core's full fused DSA hook
  (run_fused_dsa_attention, which does its own token-level top-k) to decline, so only the
  sparse-attention step is swapped between backends.
• The padding wrapper is backend-agnostic: it wraps whichever backend's
  run_fused_absorbed_sparse_attention is selected.
• What I measured and didn't: I only measured TileLang. An earlier look suggested the cuDNN /
  FlashMLA route needs flash_mla installed and pads TP-local heads to 64. That's untested here.
• Expected impact: in [sft] Enable GLM-5.3-Flash full fine-tuning in SFT (Megatron) #2341's GLM-5.3-Flash profiles, SparseMLA forward+backward was a few
  percent of a long-context step; the k-pool indexer and inter-node communication dominated.

Testing

gpu_ci/patches/megatron/test_glm5_next_fused_sparse_attention.py: padded fused attention vs the
dense reference, forward and backward.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes the fused DSA attention path for GLM-5.3-Flash long-context training; incorrect padding would alter attention or reintroduce dense OOM, though behavior is covered by CPU/GPU parity tests and auto-disables when upstream #7617 is present.

Overview
Adds a Megatron runtime backport (NVIDIA/Megatron-LM#7617) so GLM-5.3-Flash DSA can use the TileLang SparseMLA fused path instead of megatron-core’s dense (O(L^2)) fallback.

patch_sparse_mla_nope replaces tilelang_dsa.fused_sparse_mla_absorbed at model build time (MegatronWorker.make_megatron_module). For NoPE absorbed MLA it zero-pads q/k from 512→576 and pads top-k indices to a multiple of 64 with -1 before calling the kernel, then trims head/output padding. It no-ops when upstream already contains the fix or TileLang is unavailable.

Documentation in patches/megatron/README.md describes removal once the megatron-core pin lands. CPU tests use a fake SparseMLA to assert padding/gradients; H100 GPU tests compare the real kernel and dsa_kernels.run_fused_absorbed_sparse_attention against references (including 2051-wide k-pool shapes). H100 CI runs the GPU test file.

^{Reviewed by Cursor Bugbot for commit 236ca6e. Bugbot is set up for automated code reviews on this repo. Configure here.}

[erictang000]

relevant megatron-lm pr: NVIDIA/Megatron-LM#7617

[gemini-code-assist]

Code Review

This pull request introduces a backport patch (patch_sparse_mla_nope.py) for NVIDIA/Megatron-LM#7617 to support NoPE shapes and unaligned top-k in the TileLang SparseMLA kernel, preventing a fallback to dense softmax for GLM-5.3-Flash. It also adds corresponding CPU and GPU tests and integrates the patch application into the Megatron worker. The review feedback suggests wrapping the inspect.getsource call in a try-except block to avoid potential runtime crashes in environments where the source code is unavailable.

[gemini-code-assist]

Using inspect.getsource can raise an OSError or TypeError if the source code is not available (e.g., in compiled, zipped, or packaged environments where .py files are stripped). To prevent potential startup crashes, wrap the inspect.getsource call in a try-except block and default to applying the patch if inspection fails.

Suggested change

	if _UPSTREAM_SENTINEL in inspect.getsource(tilelang_dsa.fused_sparse_mla_absorbed):
	logger.warning(
	"megatron-core already supports NoPE SparseMLA (NVIDIA/Megatron-LM#7617); "
	"delete skyrl/backends/skyrl_train/patches/megatron/patch_sparse_mla_nope.py"
	)
	return False
	try:
	upstream_source = inspect.getsource(tilelang_dsa.fused_sparse_mla_absorbed)
	has_upstream_fix = _UPSTREAM_SENTINEL in upstream_source
	except (OSError, TypeError) as e:
	logger.warning(f"Could not inspect fused_sparse_mla_absorbed source: {e}. Proceeding with patch.")
	has_upstream_fix = False
	if has_upstream_fix:
	logger.warning(
	"megatron-core already supports NoPE SparseMLA (NVIDIA/Megatron-LM#7617); "
	"delete skyrl/backends/skyrl_train/patches/megatron/patch_sparse_mla_nope.py"
	)
	return False

[greptile-apps]

Confidence Score: 4/5

[High risk] Patches megatron-core sparse attention kernel for GLM model.

The PR should not merge until the GLM-5.3-Flash training configuration selects the patched TileLang backend.

Findings

1.  Flash launcher misses TileLang ▶
2.  Upstream fix fails fixtures ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[GLM-5.3-Flash DSA] --> B{DSA backend}
  B -->|tilelang| C[Patched SparseMLA shape adapter]
  C --> D[Fused sparse attention]
  B -->|other backend| E[Existing backend or dense fallback]

Loading

_{Reviews (1) · Last reviewed commit: "[megatron] Port NVIDIA/Megatron-LM#7617 ..."}

[greptile-apps]

Flash launcher misses TileLang

This installs the patch, but the GLM-5.3-Flash training launcher never sets dsa_kernel_backend="tilelang". The new GPU test sets it explicitly. Running the provided launcher therefore does not select the patched kernel, leaving long sequences exposed to the dense attention fallback and its O(L²) memory use. Set the backend in the Flash configuration and test that configuration.

[greptile-apps]

Upstream fix fails fixtures

patch_sparse_mla_nope() intentionally returns False when megatron-core already includes the fix. This assertion, and the same one in the CPU fixture, will fail after that dependency update even if the kernel works correctly. Allow the no-op case and check the resulting behavior instead, so a routine upgrade does not break both test suites.