[avigyabb]

Summary

SKYRL_DSA_QUERY_CHUNK=<tokens> (opt-in) runs GLM-5.3-Flash's fused absorbed sparse attention in query chunks, each checkpointed. Each query attends only its own top-k keys, so chunks are independent (key grads sum).

The checkpoint is what saves memory, so chunking always implies the recompute. Without it, autograd keeps every chunk's padded query (TileLang pads 8 local heads to 16 at TP8; FlashMLA pads to 64) and kernel output, as much as one unchunked call. With it, only one chunk's are alive at a time, for one extra kernel forward in backward: ~5.6 GiB less at 512k, TP8. With gradients off (e.g. inference logprobs) it only chunks.

Wraps whichever dsa_kernel_backend is selected; on tilelang the kernel takes GLM-5.3-Flash's NoPE / 2051-wide layout through patch_sparse_mla_nope (#2388). The first call logs whether the fused kernel was used or declined.

Independent of the k-pool indexer PRs (#2413, #2422). Both add an env flag in the same lines at the top of glm5_next/dsa.py, so whichever merges second needs a one-line rebase.

Testing

• gpu_ci/patches/megatron/test_glm5_next_fused_sparse_attention.py (chunked == one-shot fused kernel, fwd + bwd, ragged last chunk): passed on this branch (B200)
• tilelang vs cudnn on GLM-5.3-Flash (tiny config, TP8, 16k): fused path used on all ranks for both backends, no dense fallback; cudnn ~24% slower and +4.4 GiB there (8 local heads padded to 64). Measured on the pre-split branch.

Split out of #2341.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes GLM-5.3-Flash training/inference attention execution and autograd behavior behind an env flag, but defaults off and is covered by a GPU numerical parity test.

Overview
Adds an opt-in memory path for GLM-5.3-Flash k-pool DSA: set SKYRL_DSA_QUERY_CHUNK=<tokens> to run fused absorbed sparse attention over the query axis in chunks, with gradient checkpointing per chunk so only one chunk’s padded queries and kernel outputs stay live during backward (trades one extra kernel forward for large-sequence memory).

The wrapper _query_chunked_fused_absorbed_sparse_attention is installed during Glm5NextDSAttention._forward_with_kpool_topk by temporarily replacing dsa_kernels.run_fused_absorbed_sparse_attention; it is backend-agnostic and logs once whether the fused path was used or declined. Chunking is skipped when topk_length is set or the sequence is shorter than the chunk size; without autograd it chunks without checkpoint.

Documentation in the Megatron patches README is updated, and a new GPU CI test asserts chunked vs one-shot TileLang fused sparse attention match in forward and backward (with relaxed tolerance on key grads due to summation order).

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

[gemini-code-assist]

Code Review

This pull request introduces query-chunked fused sparse attention for GLM-5.3-Flash, which runs the fused absorbed sparse attention in checkpointed query chunks to reduce memory usage during training. It also adds a corresponding GPU test to verify that the chunked implementation matches the unchunked version. Feedback on the implementation highlights two issues with the checkpointing logic: the closure capture of the key tensor, which can prevent autograd from correctly tracking gradients, and a potential RuntimeError if checkpoint is called when neither input tensor requires gradients.

[gemini-code-assist]

There are two issues with the current checkpointing implementation:\n1. Closure Capture of key: The key tensor is captured via closure inside run instead of being passed as an explicit argument to checkpoint. In PyTorch, capturing tensors that require gradients inside a checkpointed function can prevent autograd from correctly tracking and backpropagating gradients to them.\n2. Potential RuntimeError: If torch.is_grad_enabled() is True but neither query nor key requires gradients (e.g., during validation/evaluation passes where requires_grad=False but torch.no_grad() is not used), torch.utils.checkpoint.checkpoint will raise a RuntimeError: none of the inputs have requires_grad=True.\n\nTo fix both issues, pass key as an explicit argument to run and only invoke checkpoint if at least one of the tensors requires gradients.

    def chunked(config, query, key, topk_indices, softmax_scale, v_channels, topk_length=None):\n        def run(q, k, idx, lengths):\n            return fused_fn(config, q, k, idx, softmax_scale, v_channels, topk_length=lengths)\n\n        # Each query attends only its own top-k keys, so query chunks are independent (key grads\n        # just sum). Checkpointing each chunk keeps only one chunk's padded q / kernel output\n        # alive instead of the whole sequence's -- at TP8 the kernel pads 8 local heads to 16, so\n        # that is ~13 GiB/GPU at 512k tokens -- for one extra kernel forward in backward.\n        sq = query.size(0)\n        if _DSA_QUERY_CHUNK and sq > _DSA_QUERY_CHUNK and topk_length is None:\n            outs = []\n            for q0 in range(0, sq, _DSA_QUERY_CHUNK):\n                q1 = min(sq, q0 + _DSA_QUERY_CHUNK)\n                args = (query[q0:q1], key, topk_indices[:, q0:q1], None)\n                if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad):\n                    o = torch.utils.checkpoint.checkpoint(run, *args, use_reentrant=False)\n                else:\n                    o = run(*args)\n                if o is None:\n                    return None\n                outs.append(o)\n            out = torch.cat(outs, dim=0)\n        else:\n            out = run(query, key, topk_indices, topk_length)

[greptile-apps]

Confidence Score: 4/5

[Medium risk] Adds optional query-chunked sparse attention with checkpointing.

The PR appears safe to merge, with non-blocking observability and integration-test improvements recommended.

Findings

1.  Chunk declines go unlogged ▶
2.  Production wiring remains untested ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[GLM-5.3-Flash DSA forward] --> B[Install fused-attention wrapper]
  B --> C{Chunking enabled and eligible?}
  C -- No --> D[One fused call]
  C -- Yes --> E[Checkpoint each query chunk]
  E --> F{Chunk declined?}
  F -- Yes --> G[Return None for dense fallback]
  F -- No --> H[Concatenate chunk outputs]
  D --> I[Log outcome]
  H --> I

Loading

_{Reviews (1) · Last reviewed commit: "[megatron] GLM-5.3-Flash DSA: query-chun..."}

[greptile-apps]

Chunk declines go unlogged When a fused backend declines on a chunk, this early return skips the new “declined” log. The caller then uses the much more memory-intensive dense fallback, but an opted-in run has no diagnostic explaining why the fused path was not used.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Production wiring remains untested This test constructs and calls the wrapper directly. It would still pass if Glm5NextDSAttention.forward stopped installing or invoking the wrapper, leaving the opt-in flag ineffective in training. An assertion through the attention forward path would cover that integration.

[avigyabb]

we ended up changing this structure in a previous PR. maybe we can go towards that one

[avigyabb]

this PR may not be that useful since sparse attention's saved tensors grow linearly