Skip to content

perf(dsa): cuDNN + FlashMLA kernels for GLM-5 sparse attention (model.dsa_backend) - #3829

Merged
samsja merged 11 commits into
mainfrom
perf/sparse-mla-bwd
Oct 5, 2026
Merged

samsja merged 11 commits into
mainfrom
perf/sparse-mla-bwd

Conversation

@samsja

@samsja samsja commented Oct 4, 2026 •

Copy link
Copy Markdown
Member

adding cudnn backward kernel for dsa, glm5 full training

  • New model.dsa_backend flag: tilelang (default, old path) or cudnn_flashmla (FlashMLA forward + cuDNN backward).
  • GLM-5.3 SFT example: base.toml now uses cudnn_flashmla, and a new 16k-12-node.toml overlay runs 12 nodes at 16k (cp=1, no offloading, fullgraph, full AC).
Part Old (tilelang) New (cudnn_flashmla) Kernel speedup (H200) Status
Sparse attention forward TileLang FlashMLA flash_mla_sparse_fwd 1.9–3.6× Hopper tested; Blackwell built in the wheel, untested
Sparse attention backward TileLang cuDNN dsa_bwd_sm90 3–4.7× Hopper tested; Blackwell wired in, untested

Full GLM-5.3 (78 layers), H200, 16k.

Fake data (random init, balanced routing)

setup nodes tok/s per GPU MFU
original config (cp=2, offloading), tilelang 8 465 22.3%
16k-12-node.toml, tilelang 12 701 33.6%
16k-12-node.toml, cudnn_flashmla 12 835 40.1%

Real weights + data (INTELLECT-3-SFT-10K math, 12 nodes)

dsa_backend step tok/s MFU
tilelang 31.2s 50.4k 25.2%
cudnn_flashmla 24.2s 64.9k 32.5%

🤖 Generated with Claude Code


Note

Medium Risk
Changes sparse-attention forward/backward kernels and gradient numerics (LSE/index flattening); cudnn_flashmla needs SM90+ and extra deps, with SM100 paths noted as untested.

Overview
Adds model.dsa_backend (tilelang default, cudnn_flashmla) so GLM-5 DSA sparse MLA can use FlashMLA sparse prefill forward and cuDNN CuTe-DSL backward instead of TileLang for both directions. The trainer passes the setting through model_config into sparse_mla, which now takes a backend argument; autograd routes backward to the matching implementation and normalizes LSE to natural log across paths (TileLang converts at the boundary).

Packaging adds nvidia-cudnn-frontend on the GPU extra, a flash-mla optional extra (pinned x86_64 wheel), and lockfile updates.

GLM-5.3 SFT examples: base.toml sets dsa_backend = "cudnn_flashmla"; new 16k-12-node.toml overlay documents 12-node / 16k training without CP or offloading; README expanded for both setups.

Reviewed by Cursor Bugbot for commit 0894aa8. Bugbot is set up for automated code reviews on this repo. Configure here.

samsja pushed a commit that referenced this pull request Oct 4, 2026
Removes duplicate and unused entries from `pyproject.toml` (rev-infra
P9). The set of installed packages does not change for any extra
combination (checked below).

| Removed | Reason |
|---|---|
| second `orjson>=3.10` in `dependencies` | Duplicate of
`orjson>=3.11.0` a few lines above. |
| `pandas>=2.0` | Nothing imports it: no `import pandas` / `pd.` in
`src/`, `tools/`, `scripts/`, `tests/`, `examples/`, `skills/`,
`packages/`. It was added for the nano bridge (7089ba1). It stays
installed because `datasets` depends on it. |
| `pyarrow>=21.0.0` | No direct use anywhere in those directories. The
one parquet writer, `tools/convert_traces_to_hf_dataset.py`, calls
`datasets.Dataset.to_parquet`, and `datasets` depends on `pyarrow`
itself. It stays installed. |
| `nixl-cu13==0.10.1` in the `disagg` extra | `disagg` already lists
`nixl`, and the `[[tool.uv.dependency-metadata]]` override for `nixl`
0.10.1 requires exactly `nixl-cu13==0.10.1`. The `gpu` extra keeps its
explicit pin, because `transports/weights/nixl/agent.py` imports
`nixl_cu13` directly. |
| `mamba-ssm` dependency group | Duplicate of the `gpu` extra entry,
which was added with the native Nemotron-H model (#3426). Every user of
the group also syncs `--extra gpu` or `--all-extras`, so I dropped
`--group mamba-ssm` from `Dockerfile.cuda` (2 places),
`scripts/docker-entrypoint.sh` and the `pre_run_command` in
`configs/advanced/nemotron-3-super/swe.toml`. The `[tool.uv.sources]`
git pin and the build dependencies for mamba-ssm stay. |

Behavior change:
- None for installs: same locked packages and versions.
- `uv sync --group mamba-ssm` no longer works; use `--extra gpu`, which
every repo caller already passes.
- `pandas` and `pyarrow` lose their direct lower bounds. They are still
pulled in by `datasets`.

## Checks
- `uv lock`: resolves 564 packages, the same set as before (568
`[[package]]` entries before and after). `uv lock --check` passes. The
`uv.lock` diff only touches prime-rl's own metadata. I reverted the
reordering of `[options.exclude-newer-package]` that my uv version
(0.12.10) introduced.
- I ran `uv export --frozen --all-packages` and compared the package
lists before and after, ignoring the `# via` comment lines. They are
identical for: no extras, `--extra disagg`, `--extra gpu`, `--extra
dashboard`, `--all-extras`, and `--all-extras --all-groups`. `--extra
disagg` still yields `nixl==0.10.1` and `nixl-cu13==0.10.1`, and
`--extra gpu` still yields `mamba-ssm` at the pinned rev.
- No Python changes, so no unit tests are affected. GPU: not needed. The
Docker image build in CI covers the Dockerfile and entrypoint edits.

## Conflicts
Text conflicts only in `pyproject.toml` / `uv.lock`, as with any
dependency PR. None of the open PRs that touch these files (e.g. #3829,
#3803, #3767, #3660) edit the lines removed here. A re-lock after rebase
is enough.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Low Risk**
> Metadata and sync-flag cleanup only; lockfile exports are unchanged
and runtime installs should match prior behavior.
> 
> **Overview**
> Cleans up **duplicate and redundant dependency declarations** in
`pyproject.toml` and aligns `uv.lock`, Docker, and Slurm sync commands
with the slimmer manifest. **No intended change to installed packages**
for existing extra combinations (`gpu`, `disagg`, `--all-extras`, etc.).
> 
> **Removed from core `dependencies`:** duplicate `orjson>=3.10`, unused
direct pins on `pandas` and `pyarrow` (they still come in via
`datasets`). **Removed from `disagg` extra:** explicit
`nixl-cu13==0.10.1` because `nixl` metadata already pulls that version;
the **`gpu` extra keeps** its explicit `nixl-cu13` pin for direct
imports. **Dropped the `mamba-ssm` dependency group** as a duplicate of
the `gpu` extra entry—`Dockerfile.cuda`, `scripts/docker-entrypoint.sh`,
and the Nemotron SWE `pre_run_command` no longer pass `--group
mamba-ssm`.
> 
> **Workflow note:** `uv sync --group mamba-ssm` alone is no longer
supported; use `--extra gpu` (or `--all-extras`), which callers already
use.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
f0f2ff6. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->
@samsja samsja changed the title perf(kernels): use cuDNN CuTe-DSL DSA kernel for sparse MLA backward … perf(dsa): cuDNN + FlashMLA kernels for GLM-5 sparse attention (model.dsa_backend) Oct 5, 2026

@cursor cursor Bot left a comment •

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

Comment on lines +400 to +404
if backend == "cudnn_flashmla":
arch = cudnn_backward_arch(q, kv)
assert arch is not None, "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, arch)
return tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can wwe do if then else instead of if return here please ?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in cb55aa6.

Comment on lines +401 to +402
arch = cudnn_backward_arch(q, kv)
assert arch is not None, "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lets remove this assert please

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in cb55aa6.

block_I: int = 64,
num_stages: int = 2,
threads: int = 256,
backend: str = "tilelang",

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we use literatal here ? also why is the backend predefined should it be rather just a non optiona config

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in cb55aa6.

Comment on lines +252 to +262
if backend == "cudnn_flashmla":
assert cudnn_backward_arch(q, kv) is not None, (
"dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
)
assert importlib.util.find_spec("flash_mla") is not None, (
"dsa_backend='cudnn_flashmla' requires FlashMLA; install it with `uv sync --extra flash-mla`"
)
if sm_scale is None:
sm_scale = dim_plus_tail_dim**-0.5
return flashmla_sparse_mla_forward(q, kv, indices, sm_scale, d_v)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as before can we do if then else, can we also remove the assert

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in cb55aa6.

JohannesHa
JohannesHa previously approved these changes Oct 5, 2026

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit cb55aa6. Configure here.

assert lse.shape == (B, S, H)

if backend == "cudnn_flashmla":
dq, dkv = cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, cudnn_backward_arch(q, kv))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

None arch selects SM100 backward

Medium Severity

cudnn_flashmla backward now passes cudnn_backward_arch() straight into cudnn_sparse_mla_backward. That helper returns None for a non-single KV head or an unsupported GPU, and the callee treats every non-sm90 value as SM100, so the wrong kernel runs instead of failing closed.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit cb55aa6. Configure here.

samsja and others added 11 commits October 5, 2026 04:51
…on SM90

The vendored TileLang sparse MLA backward reaches ~93-130 TFLOP/s on H200
(no pipelining, fp32 atomics for every top-k slot including sentinels).
On Hopper, dispatch to cuDNN frontend's SM90 DSA backward instead
(warp-specialized, TMA Q/dO, cp.async KV gather, skips dKV atomics for
masked slots): 3.5-4.7x faster, ~410-440 TFLOP/s at S=8k/16k, topk=2048.
Sentinel indices are remapped to -1 and the base-2 LSE is converted to
natural log. TileLang remains the fallback on other architectures.

Bumps nvidia-cudnn-frontend to >=1.30 (1.25 is incompatible with the
pinned CUTLASS DSL 4.7.1) and adds a standalone microbenchmark with an
fp32 reference check.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rnel

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ward route

Replace model.sparse_mla_backward with model.dsa_backend, selecting the DSA
sparse MLA kernels for both passes through the single prime_rl::sparse_mla op:

- tilelang (default): TileLang forward, cuDNN backward on SM90, TileLang
  backward elsewhere.
- cudnn_flashmla: FlashMLA flash_mla_sparse_fwd forward + cuDNN backward
  (SM90; the SM100/SM103 cuDNN backward is wired up but untested).
- tilelang_legacy: TileLang forward and backward.

The op now always returns a natural-log LSE. The FlashMLA route flattens the
batch, maps sentinel indices to -1 and passes a topk_length covering every
valid index so the sentinel tail is skipped.

FlashMLA (nv_dev @ b7643bd, sm_90a + sm_100f) ships as the new flash-mla
extra (part of `all`) from a prebuilt torch 2.13 / cu130 x86_64 wheel.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… overlay

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…n_flashmla

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… dsa_backend=cudnn_flashmla

On SM90 the cudnn_flashmla route now selects DSA indices with the cuDNN
frontend SM90 FP8 indexer forward (ratio=1 token-level causal mask, packed
documents as THD segments with q_causal_offsets for the CP shard) and the
cuDNN radix top-k (per-row causal lengths, chunked to bound its scratch
buffer), behind the prime_rl::cudnn_fp8_indexer custom op (AC mandatory
save). Same UE8M0 FP8 quantization as the Triton indexer; ~2x faster per
layer on H200. SM100 and the tilelang backend keep the Triton indexer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…k tie behaviour

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-k under dsa_backend=cudnn_flashmla"

This reverts commit 5724fef.
…atch without asserts

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@samsja
samsja force-pushed the perf/sparse-mla-bwd branch from cb55aa6 to 0894aa8 Compare October 5, 2026 04:52
@samsja
samsja merged commit 2b688ed into main Oct 5, 2026
18 of 20 checks passed
@samsja
samsja deleted the perf/sparse-mla-bwd branch October 5, 2026 04:54
faresobeid added a commit that referenced this pull request Oct 10, 2026
…3968)

Sets `dsa_backend = "cudnn_flashmla"` for the trainer in the GLM-5.3 RL
example (`examples/advanced/glm-5.3/swe.toml`). The SFT example already
uses it (#3829); the RL example still ran the TileLang forward and
backward.

Measured on 1x8 H200 with the trainer settings of `swe.toml` (131k seq,
cp 4, ep 8, fp8 + `deepgemm_fp8`, full AC + offload, compile). Random
init, 8 layers (3 dense + 5 MoE), fake data, 2 micro-batches per dp
rank, balanced routing:

| `dsa_backend` | step 2 / 3 | sparse attn fwd per call | sparse attn
bwd per call |
|---|---|---|---|
| `tilelang` | 16.6 s / 15.5 s | 35 ms | 190 ms |
| `cudnn_flashmla` | 12.8 s / 12.1 s | 12 ms | 53 ms |

- Entropy and mismatch KL are identical to 4 decimals.
- With a 4-layer run (5.7 s), one MoE layer costs 1.6 s instead of 2.07
s. Extrapolated to 78 layers at this shape: ~124 s vs ~160 s per step
(-22%).
- Kernel check at the CP shape (8k local queries vs 32k keys, cp rank 0
and 3, top-k 2048, sentinel slots): out / dq / dkv relative error vs
TileLang 5e-4 to 9e-4, LSE 7e-8.
- Needs the `flash-mla` extra, which the README's `uv sync --all-extras`
installs.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Low Risk**
> Single example TOML knob with no production code changes; requires the
optional flash-mla dependency for the new backend.
> 
> **Overview**
> Aligns the GLM-5.3 RL example trainer config with the SFT example by
setting **`dsa_backend = "cudnn_flashmla"`** in
`examples/advanced/glm-5.3/swe.toml`, so sparse DSA attention uses the
cuDNN/FlashMLA path instead of the default TileLang forward/backward.
> 
> This is a **performance-only** example change for long-context RL
training (131k seq, CP/EP, fp8 MoE); behavior should match TileLang
within measured numerical tolerance when the `flash-mla` extra is
installed.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
e2135c3. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->
samsja added a commit that referenced this pull request Oct 12, 2026
… example

- examples/advanced/qwen3-235b-a22b/sft-max-perf.toml: 4 x 8 B300, seq 16k,
  EP 8, full AC, fullgraph, fp32 AdamW, no offloading, balanced routing.
  Measured 4.39 s/step, 46.6% MFU. Ported from #3765 to the current schema
  (no `model.impl`, `offload = "none"`, no cluster-specific paths).
- docs/performance-trainer.md: how MFU is computed, the full setup and
  numbers for the Qwen3-235B run, and the GLM-5.3 SFT before/after for the
  cuDNN + FlashMLA sparse attention backend (#3829).
- Linked from docs nav, docs/scaling.md#benchmarking, and examples/README.md.

Co-Authored-By: Mario Sieg <mario.sieg.64@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
samsja pushed a commit that referenced this pull request Oct 14, 2026
…forward, cuDNN backward) (#4078)

Adds a `cudnn_flashmla` option to `model.dsa_backend` for DeepSeek V4
that runs the same FlashMLA sparse prefill forward vLLM uses for
DeepSeek V4 prefill, so the trainer and inference compute attention with
the same kernel. The backward is the cuDNN frontend's DSA kernel; it
runs on SM90, SM100 and SM103, and `tilelang` stays the default.

- `cudnn_dsa_patch` stops cuDNN recompiling its backward each time
sequence length or packing changes.
- On B300 (SM103), both DSv4 test files pass and the kernels match a
float32 reference; no end-to-end run there yet.
- On H200 at the B300 changes (`6a1cae560`), both DSv4 test files pass
again, and a 3-step end-to-end run (cp 8,
ep 8) compiles 3 cuDNN kernels per rank in step 1 and matches the
earlier losses, step time and peak memory.
- Fixes GLM-5's SM100 `cudnn_flashmla` backward, which omitted cuDNN's
required workspace and always raised.
- `tilelang` stays the shared `dsa_backend` default and
`configs/advanced/deepseek-v4-flash/sft.toml` opts in to
`cudnn_flashmla`, following #3829, which did the same for the GLM-5.3
configs. FlashMLA only ships an x86_64 wheel
in an optional extra, so a `cudnn_flashmla` default would break aarch64
and installs without the extra. We can flip
  the default instead if reviewers prefer.

DeepSeek-V4-Flash first 6 layers, real weights, cp 8, ep 8, one 8xH200
node. Sparse attention GPU time per step, rank 0 profiler trace at 8k
tokens/GPU:

| kernels | tilelang | cudnn_flashmla | speedup |
|---|---|---|---|
| forward, incl. recompute (48 calls) | 56.3 ms | 26.5 ms | 2.12x |
| backward (24 calls) | 71.3 ms | 36.7 ms | 1.94x |
| all sparse attention | 136.0 ms | 76.7 ms | 1.77x |

Sparse attention is under 5% of the step here, so the end-to-end gain is
smaller. Median steady-state step time:

| tokens/GPU | tilelang | cudnn_flashmla | speedup | max loss diff vs
tilelang | max loss diff, two tilelang runs |
|---|---|---|---|---|---|
| 8k | 2.86 s | 2.80 s | 1.02x | 1.0e-3 | 1.7e-3 |
| 16k | 4.59 s | 4.36 s | 1.05x | 7.3e-4 | 1.4e-3 |

Peak memory is equal, and each rank compiles 3 cuDNN kernels, all in
step 1. Baseline
[`main`](7c9d474);
[results](https://github.com/PrimeIntellect-ai/prime-rl/blob/524ec668e5684134210cef64b3575476fdf59693/benchmarks/scripts/dsv4_cudnn_flashmla_pr/results/README.md)
and
[config](https://github.com/PrimeIntellect-ai/prime-rl/blob/524ec668e5684134210cef64b3575476fdf59693/benchmarks/scripts/dsv4_cudnn_flashmla_pr/sft.toml)
at `524ec668e`.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants