Support NoPE shapes and unaligned top-k in TileLang sparse MLA by FunJim · Pull Request #7617 · NVIDIA/Megatron-LM

FunJim · 2026-09-23T08:18:37Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Allow fused_sparse_mla_absorbed to use TileLang for 512-channel NoPE inputs and top-k slot counts that are not multiples of 64. These inputs currently return None from the adapter, selecting the fallback path even though padding can make them compatible with the existing kernel.

The adapter zero-pads 512-channel Q/K to the kernel's 576-channel layout and pads top-k slots with the existing -1 invalid-index sentinel. It keeps the caller's softmax scale, selected keys, output shape, and autograd connection to the original inputs. Existing 576-channel inputs with aligned top-k take the same path as before.

Issue tracking

Related to #7522: KPool/NoPE can produce these shapes. This adapter fix is independently testable against main; it adds no model configuration or indexer behavior and does not depend on that PR being merged.

Validation

Tests cover Q/K widths 512 and 576, top-k counts 64/65/127, and 8/16/32 query heads. They check input preservation, sentinel/zero padding, explicit scale forwarding, and output/Q/K-gradient agreement with an independent FP32 sparse-attention reference, including an all-invalid query row.

Forward bit-exact replay tests run with stream contention and are registered in the kernel manifest. The existing DSA backward exemption remains: backward uses atomic accumulation, so its tests check numerical agreement rather than bit-exact replay.

On one H200, 276 passed, 0 failed, 0 skipped: 18 adapter tests, 18 real CUDA output/gradient tests, 4 forward replay tests, and 236 manifest tests. Replacing only the adapter with the unmodified upstream implementation gives 15 expected failures and 3 passes in the adapter tests.

Environment: current dev NGC base nvcr.io/nvidia/pytorch:26.08-py3, PyTorch 2.14.0a0+4fdf77b940.nv26.08 (CUDA 13.4), TileLang 0.1.12, TVM FFI 0.1.11, Z3 4.15.4.0, and isolated test dependencies. This was targeted validation on the NGC base; the full CI image and end-to-end training were not run. No performance benchmark is claimed.

python -m torch.distributed.run --nnodes 1 --nproc-per-node 1 \
  --master-addr 127.0.0.1 --master-port 29531 -m pytest -q \
  tests/unit_tests/transformer/experimental_attention_variant/test_dsa_tilelang_shapes.py \
  tests/unit_tests/determinism/kernels/test_dsa_tilelang_shapes.py \
  tests/unit_tests/determinism/kernels/test_manifest.py
python tools/check_kernel_determinism_coverage.py --base-ref upstream/main
BASE_REF=main CHECK_ONLY=true bash tools/autoformat.sh

Kernel coverage, Black, isort, Pylint, and Ruff checks passed. autoformat.sh completed in fix and check modes; its optional mypy command was unavailable in the repository's linting dependency group.

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests (not run; this is a focused adapter change)
I have added a bit-exact forward replay test and registered it in the kernel manifest; the existing backward exemption remains
Production function typing is preserved
The padding behavior is documented next to the adapter changes
I have run tools/autoformat.sh on my PR

Created as a draft; human author review and upstream CI are pending.

Signed-off-by: FunJim <xunjin.zheng@outlook.com>

copy-pr-bot · 2026-09-23T08:18:41Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

cspades · 2026-09-24T04:51:46Z

/ok to test 3d07347

….3 SparseMLA wrapper Replace the GLM-specific _pad_for_fused_absorbed_sparse_attention wrapper in glm5_next/dsa.py with a backport of NVIDIA/Megatron-LM#7617 (approved upstream), which teaches megatron-core's tilelang_dsa.fused_sparse_mla_absorbed the same two exact paddings: 512-wide NoPE q/k zero-padded to 576, and top-k widths padded to a multiple of 64 with -1. patch_sparse_mla_nope.py rebinds tilelang_dsa.fused_sparse_mla_absorbed (its only caller resolves it as a module global) to the upstream function from the PR head, and is a no-op that warns once the pinned megatron-core contains the fix. Applied from MegatronWorker.make_megatron_module next to the other megatron-core patches. Tests: - CPU: upstream's fake-SparseMLA adapter test (padding, -1 slots, scale, unpadding, gradients), plus the GLM k-pool width 2051, through tilelang_dsa.run_fused_absorbed_sparse_attention. 25 passed. - GPU (h100 marker): upstream's real-kernel vs FP32 reference sweep, plus the GLM-5.3-Flash-shaped case through dsa_kernels vs megatron-core's dense fallback. 20 passed on 1x B200. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…#2388) ## 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 #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](https://claude.com/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. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 236ca6e. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>  --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Eric Tang <erictang000@gmail.com>

Keep the upstream top-k padding introduced by NVIDIA#7522 and retain the NoPE Q/K zero-padding and shape regression tests from NVIDIA#7617. Validation: 302 CPU adapter/manifest checks passed using the actual production function definitions; Black, isort, Pylint, Ruff and kernel determinism coverage passed. CUDA tests were not rerun on this CPU host. Signed-off-by: FunJim <xunjin.zheng@outlook.com>

FunJim · 2026-10-23T10:05:28Z

Resolved in 7c7f759 by merging upstream main at 4422adc, preserving the existing PR history.

The conflict overlapped with the top-k padding introduced by #7522. I kept the upstream implementation and retained this PR’s 512-channel NoPE Q/K padding and shape regression tests.

Validation: 302 isolated CPU adapter/manifest checks passed, along with Black, isort, Pylint, Ruff, and kernel determinism coverage. CUDA tests were not rerun for this update.

GitHub now reports the branch as mergeable, and the signed commit and DCO checks pass.

janEbert · 2026-10-27T18:11:14Z

/ok to test 7c7f759

nemo-automation-bot · 2026-10-28T11:52:57Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40683350755

Support NoPE shapes and unaligned top-k in TileLang sparse MLA

3d07347

Signed-off-by: FunJim <xunjin.zheng@outlook.com>

nemo-automation-bot Bot added the community-request label Sep 23, 2026

FunJim marked this pull request as ready for review September 23, 2026 08:23

FunJim requested review from a team as code owners September 23, 2026 08:23

svcnvidia-nemo-ci requested a review from a team September 23, 2026 08:23

cspades approved these changes Sep 24, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to test September 24, 2026 04:52 Inactive

copy-pr-bot Bot temporarily deployed to public September 24, 2026 04:55 Inactive

svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Sep 26, 2026

erictang000 mentioned this pull request Oct 1, 2026

[megatron] GLM-5.3-Flash DSA: run the fused TileLang SparseMLA kernel NovaSky-AI/SkyRL#2388

Merged

avigyabb mentioned this pull request Oct 14, 2026

[model] GLM-5.3-Flash Tracker NovaSky-AI/SkyRL#2487

Open

Phlip79 approved these changes Oct 16, 2026

View reviewed changes

Phlip79 added the 26.12 label Oct 16, 2026

guihong-nv approved these changes Oct 16, 2026

View reviewed changes

janEbert approved these changes Oct 23, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Oct 23, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Oct 23, 2026

jaredcasper approved these changes Oct 27, 2026

View reviewed changes

janEbert enabled auto-merge October 27, 2026 18:11

copy-pr-bot Bot temporarily deployed to test October 27, 2026 18:15 Inactive

copy-pr-bot Bot temporarily deployed to public October 27, 2026 18:16 Inactive

svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Oct 27, 2026

janEbert added this pull request to the merge queue Oct 28, 2026

Merged via the queue into NVIDIA:main with commit f77f360 Oct 28, 2026
144 checks passed

guihong-nv mentioned this pull request Oct 28, 2026

Move DSA, CSA, absorbed MLA and DeepSeek-v4 attention into megatron/core/ops with deprecated aliases (3/4) #8348

Draft

4 of 7 tasks

FunJim deleted the fix/tilelang-nope-topk-padding branch October 29, 2026 09:18

Navigation Menu

Repository navigation

Support NoPE shapes and unaligned top-k in TileLang sparse MLA - #7617

Support NoPE shapes and unaligned top-k in TileLang sparse MLA#7617
janEbert merged 2 commits into
NVIDIA:mainfrom
FunJim:fix/tilelang-nope-topk-padding

FunJim commented Sep 23, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 23, 2026

Uh oh!

cspades commented Sep 24, 2026

Uh oh!

FunJim commented Oct 23, 2026

Uh oh!

janEbert commented Oct 27, 2026

Uh oh!

nemo-automation-bot Bot commented Oct 28, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

FunJim commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issue tracking

Validation

Contribution process

Pre-checks

Uh oh!

copy-pr-bot Bot commented Sep 23, 2026

Uh oh!

cspades commented Sep 24, 2026

Uh oh!

FunJim commented Oct 23, 2026

Uh oh!

janEbert commented Oct 27, 2026

Uh oh!

nemo-automation-bot Bot commented Oct 28, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

FunJim commented Sep 23, 2026 •

edited

Loading