lablup / lablup/mlxcel

perf(cuda): bucket the cuDNN SDPA plan-cache key across verify rounds

Open
#1,820 0 comments 0 reactions 0 assignees View on GitHub
area:core area:inference platform:linux priority:medium status:in-progress type:performance
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem / Background

MLX keys its cuDNN SDPA execution-plan cache on the exact q/k/v/mask shapes and strides, so a speculative verify round, whose key length grows every round, never hits that cache and rebuilds plans on the host every round. #1799's attribution measured the consequence on GB10 (sm_121, MLX pin `81ba1c6a`, `models/mlx/laguna-xs-2.1-nvfp4` with `models/mlx/laguna-xs-2.1-dflash`): `ScaledDotProductAttention::eval_gpu` costs 67.1 ms of host time per verify round at block 8 and 75.8 ms at block 16, against 2.8 ms per classic decode token, and the term is flat in the row count. Three shape classes (target full layers, target sliding layers with sinks, drafter layers) rebuild a plan each round at about 22 ms apiece. The same LRU's lifetime miss counter then aborted every block-2 profile past about 170 rounds, so the miss path is not only slow but ends the process. Record: `docs/benchmark_results/laguna-dflash-verify-cost-gb10-2026-09-11.md`; attribution: `TECHNICAL_REPORTS/1799-laguna-dflash-verify-cost-cuda-20260911.en.md` (commit `c6a23baf`).

PR #1817 (open, not merged) routes small masked verify blocks around cuDNN entirely and works: Laguna's best width moves from 0.72x of classic (block 6) to 1.33x (block 4), and Qwen 3.5 cross-checks at 1.26x and 1.35x with greedy text byte-identical to classic. It is a contained fix that leaves the cache-key design untouched, and it pays for that: the record measures the ops fallback at about 10 ms more GPU time per round at block 2 (44.3 ms of kernels before, 54.7 ms after) against the 69 ms of host time it removes. Making the key reusable would let the cache hit and keep cuDNN's faster flash kernel, recovering that 10 ms per round instead of trading it away. That is the general fix, and it is why this is a separate issue rather than a follow-up commit on #1799.

## Current Behavior

Read at the pinned MLX commit `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a` (`src/lib/mlx-cpp/CMakeLists.txt:122`, the single source of truth per #1047), file `mlx/backend/cuda/scaled_dot_product_attention.cpp`:

- `SDPACacheKey` (lines 126-140) holds the device id, the cuDNN dtype, the q/k/v shapes and strides, `do_causal`, the mask shape and strides, `has_sinks` and `output_logsumexp`. `build_sdpa_cache_key` (142-176) fills it, including the mask shape and strides at 165-167.
- The one-row decode path is already canonicalized. When `decoding` is set (168-174), `k_shape[2]` and `v_shape[2]` are replaced by `T_kv`, the allocated cache extent, and the k/v strides are zeroed, so consecutive decode steps share one key. The device side that makes this correct is in `sdpa_cudnn` (371-385): it unslices k/v to the full buffer and builds `seq_len_q` / `seq_len_kv`, and `build_sdpa_graph` (240-243) turns those into `set_padding_mask(true)` plus `set_seq_len_q` / `set_seq_len_kv`, so cuDNN reads only the valid prefix of a padded tensor. That machinery exists today and is gated to `q.shape(2) == 1` with no array mask by `use_cudnn_for_decoding` (72-107, `kv_cache_step = 256` at line 85).
- The cache itself is `sdpa_cache()` (178-182), an `LRUBytesKeyCache` sized by `MLX_CUDA_SDPA_CACHE_SIZE` with default capacity 256. `mlx/backend/cuda/lru_cache.h:91-98` throws the fatal `Cache thrashing` error once the lifetime miss count passes `2 * capacity_`, armed by the env-name constructor at 35-39. This is the abort documented in `docs/upstream/mlx-cuda-graph-cache-lifetime-miss-abort.md` (#818, #821).
- Dispatch: `supports_sdpa_cudnn` (315-350) accepts Ampere and later, head_dim a multiple of 8 and at most 128, f16 or bf16; `ScaledDotProductAttention::eval_gpu` (609-655) then takes cuDNN, else `sdpa_vector`.

On the mlxcel side, a Laguna multi-row append builds its own additive f32 `0 / -inf` mask of shape `[l, l + prior]` per layer (`src/models/laguna_layers.rs:165-175`, built by `create_causal_mask` / `create_causal_mask_with_window_full`, `src/lib/mlxcel-core/src/utils.rs:120-122,140-162`), and the dense KV cache grows in 256-position steps (`src/lib/mlxcel-core/src/cache.rs:499-506`), the same step MLX's decode canonicalization already relies on. So per round the k/v shape, the k/v strides and the mask shape all move while the q shape is fixed at the block width. Confirm which key fields actually differ round over round, per shape class, before designing around them.

## Proposed Solution

Make the plan-cache key reusable across rounds whose key length differs only by a small increment, by extending MLX's own decode-path canonicalization to small multi-row calls: bucket (pad) the key length to a step, canonicalize the k/v shape and stride fields in the key the way lines 168-174 already do for decoding, pad the array mask to the same bucket, and hand cuDNN the true lengths through the padding-mask and `seq_len` tensors it already supports. Land it as an overlay in `src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp`; the CMake glob at `src/lib/mlx-cpp/CMakeLists.txt:51-64` already copies that directory over the fetched tree, so no build wiring is needed. If #1817 has merged by then, this edits the same patched file; #1816 may move that path, so check before starting.

Decisions to make explicitly in the implementation, not left to the reader:

- **Bucket step.** 256 is the natural first choice because both MLX (`kv_cache_step`, line 85) and mlxcel's KV cache (`cache.rs:499-506`) already step by it, so the backing buffer is already that size and no reallocation is implied. State the step chosen and why.
- **Mask.** The array mask is in the key (165-167) and its column count grows with `k_len`, so bucketing the k/v shape alone does not stop the miss. Pad the mask to the bucket with `-inf` (blocked) columns, which matches the existing builder's sentinel contract, or express the band natively as `set_causal_mask_bottom_right` plus the sliding bound. Say which, and why the other was rejected.
- **Sinks.** Sliding Laguna layers pass per-head sinks (`laguna_layers.rs:179-187`); `has_sinks` is in the key and `build_sdpa_graph` sets a sink token at 237-239, so the padded path must preserve sink semantics rather than silently dropping them.
- **Rejected alternative, recorded so it is not retried.** Raising `MLX_CUDA_SDPA_CACHE_SIZE` (what #1817 does, defaulting it to 2000 on CUDA builds) only postpones the abort; it never produces a hit, because every round's key is new. That is the point of the doc in `docs/upstream/`.

## Scope

**In scope:** the cuDNN SDPA overlay under `src/lib/mlx-cpp/patches/mlx/backend/cuda/`, any Rust-side mirror of its gate in `src/lib/mlxcel-core/src/layers.rs`, the CUDA default in `src/lib/mlxcel-core/src/hardware.rs`, `docs/environment-variables.md`, a benchmark record under `docs/benchmark_results/` and a technical report under `TECHNICAL_REPORTS/` (`.en.md` and `.ko.md`).

**Out of scope:** the default draft block width (#1797), the MLX graph ops and byte budgets on GB10 (#1798), the Laguna exactness probe recorded as out of scope on #1799, and the per-row expert term (`qmm_sm80`, 2.5 ms per row), which #1799 attributed to expert reads and which is untouched here.

## Implementation Notes

- **Reuse.** Do not build a second canonicalization path: extend the `decoding` branch of `build_sdpa_cache_key` and the `seq_len` plumbing in `sdpa_cudnn`. Do not add new timers: the per-round split already exists on the CLI's `DFlash:` line and in `DFlashDiagnostics` (`src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs:530-538`, logged in `src/server/batch/dflash_target.rs:706-718`). Reuse `scripts/bench_block_width.sh` (interleaved widths, rotating start) under `scripts/with_indexers_paused.sh`, and the same-binary kill-switch pattern #1795 and #1817 used.
- **Blast radius, wider than #1799's.** This touches MLX's CUDA SDPA dispatch, which every CUDA model with head_dim at most 128 on Ampere and later reaches, not only speculative decoding: prefill, the trailing short chunk of a chunked prefill, and any multi-row append all pass through it. Measure at least one non-speculative workload (classic decode plus a long-prompt prefill on a head_dim at most 128 model) on the same binary and report it, not only the Laguna and Qwen 3.5 speculative arms.
- **Padding changes what the kernel reads, so prove correctness rather than asserting it.** State explicitly how the padded positions are excluded (the `seq_len` tensors, the `-inf` mask columns) and what the padded k/v region actually contains (stale cache bytes, not zeros), then prove greedy output is byte-identical by comparing token ids (`MLXCEL_PRINT_TOKEN_IDS=1`) against the pre-change binary, on Laguna and on one head_dim at most 128 non-speculative model. #1782 is the precedent: a dtype promotion there silently changed Qwen 3.5's output at the default block width, and it was invisible until someone compared ids.
- **Bucketing trades memory and bandwidth for hits, so bound both.** State the distinct-key count per generation and the steady-state resident plan count (buckets times shape classes times layer classes), the headroom against the `2 * MLX_CUDA_SDPA_CACHE_SIZE` lifetime-miss abort, and any cuDNN workspace growth. Padding to a 256 bucket also means cuDNN may read up to 255 extra key positions per call across the 45 attention calls in a Laguna round, which directly offsets the 10 ms the fix is chasing; measure it rather than assuming it is free.
- **Settle the interaction with #1817 in the same PR.** If bucketing makes cuDNN viable for these shapes again, #1817's routing-around is redundant for them: either default `MLXCEL_SDPA_FALLBACK_MAX_QUERIES` to 0, or remove the gate and its Rust mirror (`cuda_sdpa_small_query_fallback` and `sdpa_fallback_max_queries` in `layers.rs`), and say which, with the measurement that decided it. Do not leave both mechanisms in place unexamined. If the fallback is still faster at some widths, state at which and keep the gate scoped to those.
- **Edge cases.** A partial accept grows `k_len` by a variable amount per round, so the bucketing must tolerate a non-uniform stride; a rotating sliding cache exposes a key count that changes as the window fills, with speculative slack on top (`src/models/laguna_speculative.rs:64-82`); crossing a bucket boundary mid-generation is one legitimate miss and must not be read as a regression; `force_fused=True` must keep raising where no fused kernel exists; the backward primitive (`ScaledDotProductAttentionVJP::use_fallback`, 657-664) must stay out of the new path.
- **Measurement hazards.** GB10 single-run decode is bimodal by up to 25% even at pinned clocks (#755), so n at least 3 with min and max, on an idle host with the GPU held exclusively. nsys does not see graph-captured kernels without `--cuda-graph-trace=node`. An unscoped `cargo test --lib` aborts on this host with a `cudaStreamEndCapture` C++ abort unrelated to any of this, so scope the test selector.

## Acceptance Criteria

- [ ] A committed nsys/NVTX table shows `ScaledDotProductAttention::eval_gpu` host time per verify round at widths 2, 4, 8 and 16 falling to the same order as the classic step (2.8 ms per token), with the plan-build count per round measured rather than inferred.
- [ ] The throughput A/B is measured against the #1817 state as the baseline (that gate active, since it is what will be on main), not against the pre-#1817 baseline, on Laguna DFlash at widths 2 to 16, n at least 3, min and max, same binary, with a kill switch for the new behavior and the classic arm re-measured in the same session.
- [ ] The record states whether the roughly 10 ms per round of extra GPU time the #1817 fallback costs at block 2 is recovered, with the before and after kernel totals per round.
- [ ] At least one non-speculative workload (classic decode and a long-prompt prefill on a head_dim at most 128 model) is measured on the same binary and shows no regression.
- [ ] Greedy token ids are byte-identical to the pre-change binary on Laguna and on one non-speculative model, compared as ids and not as text, and the Qwen 3.5 speculative pairing still reaches its measured 1.26x and 1.35x.
- [ ] Plan-cache growth is bounded and stated: distinct keys per generation, steady-state resident plans, and headroom against the lifetime-miss abort.
- [ ] The #1817 interaction is resolved in the same PR: its gate is narrowed, removed or kept, with the measurement that decided it, and `docs/environment-variables.md` updated to match.
- [ ] The change is on by default in the real decode path for both the CLI and the server, with a same-binary kill switch, not behind an off-by-default flag.
- [ ] A benchmark record under `docs/benchmark_results/` and a technical report under `TECHNICAL_REPORTS/` (`.en.md` and `.ko.md`), following the #1782 and #1799 pair.

## Verification

```
cargo fmt --all -- --check
cargo clippy --profile test-fast --features cuda --lib --bins --tests -- -D warnings
cargo test --profile test-fast --features cuda -p mlxcel-core -- --test-threads=1 layers::
make verify-test-cuda

# Round split and plan builds (repeat at widths 2, 4, 8, 16; kill switch on and off)
nsys profile -t cuda,nvtx --cuda-graph-trace=node -o laguna_w4 \
mlxcel generate -m models/mlx/laguna-xs-2.1-nvfp4 \
--draft-model models/mlx/laguna-xs-2.1-dflash --draft-kind dflash \
--draft-block-size 4 -n 200 --temp 0
nsys stats --report nvtx_sum laguna_w4.nsys-rep | grep ScaledDotProductAttention

# Throughput sweep on an idle host, widths interleaved
./scripts/with_indexers_paused.sh ./scripts/bench_block_width.sh 2 4 6 8 10 12 16

# Identity (ids, not text): pre-change binary against post-change binary
MLXCEL_PRINT_TOKEN_IDS=1 mlxcel generate -m -p -n 200 --temp 0
```

A pass is: the per-round host time in the SDPA primitive collapsed with the plan-build count per round at or near zero after warm-up, a width table at or above the #1817 numbers on the same host, an unchanged non-speculative arm, identical token ids, and the two documents committed.

## Technical Considerations

Provenance: #1799 (the attribution this rests on, with PR #1817 in flight for its contained fix), #1782 and PR #1795 (the Qwen 3.5 attribution and fix, and the byte-identity precedent), #1798 (the GB10 graph-budget finding recorded from the same sweep), #818 and #821 with `docs/upstream/mlx-cuda-graph-cache-lifetime-miss-abort.md` (the LRU abort, whose affected-surface table already lists this SDPA cache), and #1816 (may move the overlay path).

Ordering: resolve this before #1797 fixes a default width. #1797 measures an optimum that depends on this cost structure, and after #1817 the Laguna optimum is width 4 at 1.33x with the checkpoint default of 16 still losing at 0.83x. If bucketing lands afterwards and returns cuDNN's faster kernel to these shapes, that curve moves again and #1797's number goes stale on arrival.

Note that Qwen 3.5 (head_dim 256) never enters cuDNN, so it is a no-regression cross-check here rather than evidence either way, the same role it played in #1799.

Contributor guide

Open the contributing guide

Research direction

Read the MLX overlay at src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp, then inspect the Rust gate in src/lib/mlxcel-core/src/layers.rs and cache growth in cache.rs. First confirm which key fields change across rounds, then use scripts/bench_block_width.sh with the existing diagnostics and scoped tests; done requires measured cache reuse, bounded memory, no non-speculative regression, and byte-identical token ids.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, rust
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.