Repository navigation
perf(dsa): cuDNN + FlashMLA kernels for GLM-5 sparse attention (model.dsa_backend) - #3829
Conversation
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 -->
| 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) |
There was a problem hiding this comment.
can wwe do if then else instead of if return here please ?
| 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" |
There was a problem hiding this comment.
lets remove this assert please
| block_I: int = 64, | ||
| num_stages: int = 2, | ||
| threads: int = 256, | ||
| backend: str = "tilelang", |
There was a problem hiding this comment.
can we use literatal here ? also why is the backend predefined should it be rather just a non optiona config
| 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) | ||
|
|
There was a problem hiding this comment.
same as before can we do if then else, can we also remove the assert
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ 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)) |
There was a problem hiding this comment.
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)
Reviewed by Cursor Bugbot for commit cb55aa6. Configure here.
…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>
…ent top-k tie behaviour" This reverts commit 06f5c60.
…-k under dsa_backend=cudnn_flashmla" This reverts commit 5724fef.
…atch without asserts Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
cb55aa6 to
0894aa8
Compare
…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 -->
… 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>
…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>


adding cudnn backward kernel for dsa, glm5 full training
model.dsa_backendflag:tilelang(default, old path) orcudnn_flashmla(FlashMLA forward + cuDNN backward).base.tomlnow usescudnn_flashmla, and a new16k-12-node.tomloverlay runs 12 nodes at 16k (cp=1, no offloading, fullgraph, full AC).tilelang)cudnn_flashmla)flash_mla_sparse_fwddsa_bwd_sm90Full GLM-5.3 (78 layers), H200, 16k.
Fake data (random init, balanced routing)
tilelang16k-12-node.toml,tilelang16k-12-node.toml,cudnn_flashmlaReal weights + data (INTELLECT-3-SFT-10K math, 12 nodes)
dsa_backendtilelangcudnn_flashmla🤖 Generated with Claude Code
Note
Medium Risk
Changes sparse-attention forward/backward kernels and gradient numerics (LSE/index flattening);
cudnn_flashmlaneeds SM90+ and extra deps, with SM100 paths noted as untested.Overview
Adds
model.dsa_backend(tilelangdefault,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 throughmodel_configintosparse_mla, which now takes abackendargument; autograd routes backward to the matching implementation and normalizes LSE to natural log across paths (TileLang converts at the boundary).Packaging adds
nvidia-cudnn-frontendon the GPU extra, aflash-mlaoptional extra (pinned x86_64 wheel), and lockfile updates.GLM-5.3 SFT examples:
base.tomlsetsdsa_backend = "cudnn_flashmla"; new16k-12-node.tomloverlay 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.