Skip to content

Allow force using fused kernel in scaled_dot_product_attention #3658

Closed

Description

@DT-HJ

Summary

PR #3293 added a fused full-attention path for head_dim=192/256 (steel_attention instantiation + use_fallback routing at kL>16384). The core change was reviewed favorably and validated by third parties on M2 Ultra and M3 Ultra, but the PR was closed when the author stepped away from this line of work. The last open question from @angeloskath was:

Simply put, a real world scenario would be even better. Something like mlx_lm.generate .... results in OOM ...

We have been running head_dim=256 models in production on Apple Silicon for months, and have exactly that evidence.

Setup

  • Mac Studio (Apple Silicon, 36 GB unified memory), MLX 0.31.x
  • Runtime: oMLX (OpenAI-compatible server on MLX) with chunked prefill, prefix KV cache, and a Metal memory guard
  • Model: Gemma 4 26B-A4B MoE (3-bit), head_dim=256 — 30 layers, of which only 5 attend over the full context (25 are sliding-window 1024)

Chunked prefill does not solve it

Chunked prefill avoids the single qL x kL self-attention, as noted in the PR thread. But on the unfused path each chunk still materializes fp32 scores against the whole cache, per full-attention layer:

transient ≈ n_q_heads × chunk_len × kv_len × 4 bytes

This grows linearly with kv_len, so a chunk size that is safe at 16K is not safe at 130K. In practice on our setup:

Scenario (Gemma 4 26B, 36 GB machine) Result
Cold prefill of a ~130K-token prompt, 1024-token chunks 226 s — completes, transient-bound
Default 2048-token chunks at long context Metal OOM abort (kIOGPUCommandBufferCallbackErrorOutOfMemory) — we had to patch the chunk size down and add memory guards
Multi-turn conversation growing past ~133K, memory ceiling active guard shrinks chunks to 32–512 tokens to stay alive → a single turn's prefill takes 26+ minutes

At kv_len=130K, even a 32-token chunk allocates ~665 MB for scores alone. Decode is unaffected (sdpa_vector already covers head_dim=256) — the gap is exactly the full-attention path #3293 addressed.

Note this is with only 5/30 layers attending over the full context; dense head_dim=256 models hit the same wall proportionally harder.

Control group (same hardware, same runtime)

head_dim=128 models — Granite 4.1 30B (dense), Nemotron-3-Nano 30B-A3B (hybrid MoE), and a cohere2-MoE 30B — show flat-to-linear per-turn prefill latency up to 164K context on the same setup. The wall is specific to the unfused head_dim=256 full-attention path, not the runtime or the hardware.

Ask

Could the final form of #3293 (bd=192/256 steel_attention instantiation + kL>16384 routing — unfused stays the default where it is faster) be reconsidered, either by reviving the branch or reimplementing it? head_dim=256 is common and growing across model families (Gemma 4, Qwen3.5/3.6 dense and MoE). Happy to benchmark on our workloads or test a branch.

Activity

  1. DT-HJ commented on Jun 11, 2026

    @DT-HJ
    ContributorAuthor

    Correction (2026-06-13): The Gemma attribution in this comment is incorrect. Gemma 4 26B large-context global-attention layers use head_dim=512 and n_kv=2; the head_dim=256, n_kv=8 shape benchmarked below is the bounded sliding-window shape, so it does not exercise the large-kL route in this model. The kernel-level head-dim-256 result remains valid, but it is not evidence that #3660 changes Gemma 4 26B/31B. I also incorrectly described the fallback score tensor as fp32; for this model it is bf16 (with fp32 softmax accumulation). A stock mlx-lm A/B found identical Gemma memory and prefill throughput across the two builds. See the full correction and measurements. The original comment is preserved below for context.

    Follow-up with concrete numbers: I rebased #3293's three commits onto current main (plus two small completeness fixes), built from source, and benchmarked on M4 Max (36 GB).

    Branch: https://github.com/hojin12312/mlx/tree/sdpa-256-revival (original commits preserved with authorship; fixes: missing bd=192 steel instantiation that the routing already targeted, and excluding head_dim>=192 from the NAX dispatch branch since the NAX family has no 192/256 instantiations).

    Peak memory — the actual point of the change

    n_q=16, n_kv=8 (GQA), head_dim=256, fp16 (Gemma 4 26B attention shape), inputs included:

    qL kL inputs fused peak unfused peak max abs err
    1024 32768 264M 272M 1,824M 0.0
    2048 32768 272M 288M 2,880M 0.0
    2048 65536 528M 544M 5,696M 0.0
    2048 98304 784M 800M 8,512M 0.0
    2048 131072 1,040M 1,056M 11,328M 0.0

    Fused transient is ~16 MB regardless of kL; unfused grows linearly to +10.3 GB at 131K. Outputs agree exactly.

    Speed — honest numbers

    Same shapes, min/mean of 10 iters: above the kL>16384 routing threshold the fused kernel is 0.74–0.79x of unfused on M4 Max (e.g. qL=2048, kL=65536: 275 ms fused vs 214 ms unfused). Consistent with @angeloskath's earlier measurements — per-chunk, unfused is faster when the transient fits.

    Why the memory column wins in production

    A serving runtime holds 15–25 GB of weights+KV on a 36 GB machine. It cannot absorb a 3–10 GB fp32 transient per prefill chunk, so under memory pressure it shrinks chunks (2048 → 512 → 32 tokens) to keep the transient bounded — and GPU utilization collapses: we measure 26+ minute single-turn prefills at ~133K context (Gemma 4 26B, 5/30 full-attention layers). With the fused path the same step is a few hundred ms per chunk at full chunk size.

    The kL>16384 routing already keeps unfused wherever it is faster and safe; the threshold is tunable if you'd prefer it gated on estimated transient size (n_q*qL*kL*4) instead of kL alone.

    Happy to open this as a PR if there's interest.

  2. DT-HJ commented on Jun 11, 2026

    @DT-HJ
    ContributorAuthor

    Correction (2026-06-13): The measurements below are preserved as historical oMLX results, but the Gemma improvements were incorrectly attributed to #3660. Gemma 4 26B/31B use head_dim=512 in their large-context global layers, so the new 192/256 fused route in the PR does not activate there. The reported 156K→184K admitted-context change came from relaxing the oMLX admission estimator, not from a kernel memory reduction; that relaxation under-estimated the actual bf16 fallback and must be corrected on our side. Therefore the Gemma rows below are not evidence for this PR. See the full correction and stock mlx-lm A/B. The original comment is preserved below for context.

    Update — we put the branch into production serving and measured
    end-to-end (Mac Studio 36 GB, oMLX, chunked prefill 1024, 27 GB prefill
    safety cap, Gemma 4 26B 3-bit). The serving admission estimator was the
    remaining blocker (it assumed the unfused transient and rejected long
    prompts before the kernel ever ran); with the estimator taught the
    fused O(1) transient:

    Scenario unfused fused
    Max cold single-shot prefill ~156K (estimator-bound) 184K (629 s; 204K rejected cleanly at 182K — now genuinely memory-bound)
    Multi-turn ceiling ~133K, then 26+ min/turn crawl 167K, 99% cache hit, 9–12 s/turn, crawl gone
    25K paste on a 130K conversation impractical 83 s
    Qwen3.6 35B 145K cold 280 s 332 s (0.84x — honest per-kernel cost where memory never bound)

    PR opened: #3660

  3. zcbenz commented on Jun 13, 2026

    @zcbenz
    Member

    @hojin12312 Would it be possible to measure with the mlx-lm instead of oMLX? We are not familiar with oMLX's code so it is hard for us to check if this is something that we could fix without sacrificing the performance.

  4. Thump604 commented on Jun 13, 2026

    @Thump604

    I ran a clean official mlx-lm A/B locally against current MLX main and this PR branch.

    Scope: official mlx-lm benchmark path only. Excluded: oMLX, Runtime serving,
    Jobs, draft/MTP, custom scheduler.

    Environment:

    • MLX main: a025496c8f6da9d599e2010019ea5ee328c2d119
    • PR branch: 0aa50825acca3be22c8f02c5581d17aa442ba742
    • mlx-lm: 2c008fd0252b2c569227d12568356ab88ab0560a / 0.31.3
    • Command path: python -m mlx_lm.benchmark

    Results:

    Workload main prompt_tps PR prompt_tps main peak GB PR peak GB
    Gemma 4 26B 8bit, 32K 896.247 892.607 31.078 31.078
    Gemma 4 26B 8bit, 64K 745.338 742.810 33.902 33.902
    Qwen 3.6 35B 8bit, 32K 1005.743 871.489 41.316 39.882
    Qwen 3.6 35B 8bit, 64K 802.929 624.513 44.201 40.068

    Interpretation: official mlx-lm reproduces the memory/speed tradeoff on Qwen
    35B: lower peak memory, materially slower prompt throughput. In this benchmark
    shape, Gemma 4 26B did not show a peak-memory improvement at 32K or 64K.

    I also tried an optional Qwen 35B 131K row. Main completed with
    prompt_tps=516.049, peak_memory=49.971 GB; the PR branch aborted before
    completing the trial with:

    kIOGPUCommandBufferCallbackErrorImpactingInteractivity
    

    I would treat the 131K row as an unresolved follow-up, not a valid A/B
    comparison.

    Full local artifact:
    /opt/ai-runtime/run/mlx-sdpa-3660-validation/20260613T041902Z/summary.json

  5. DT-HJ commented on Jun 13, 2026

    @DT-HJ
    ContributorAuthor

    Thanks. I re-ran this through the stock mlx-lm benchmark path — no oMLX, no serving layer: vanilla
    mlx 0.31.2 vs the PR wheel at 0aa50825, python -m mlx_lm.benchmark, Gemma 4 26B UD-MLX 3-bit and
    4-bit, 32K→131K, prefill-step 2048. (The two arms are release builds 71 commits apart, not a
    single-commit A/B; the static dispatch analysis below is what isolates the PR's behavior.)

    No difference. Peak memory is identical between the two builds in every row (to the benchmark's
    0.001 GB precision; e.g. 4-bit 131K = 28.088 GB both), and prompt_tps matches within 0.1%. This agrees
    with @Thump604's earlier Gemma result.

    Root cause — and a correction to my own earlier micro-benchmark: Gemma 4 26B/31B run their
    full-attention (global) layers at head_dim=512 (global_head_dim)
    , not 256. Only the
    sliding-window layers are head_dim 256, and those are bounded (kL ≈ window+chunk, well under the 16384
    threshold), so they never hit a large-kL full-attention SDPA. Since 512 is not in the fused dispatch
    set, both vanilla and the PR build take the same unfused matmul→softmax→matmul fallback for those
    layers
    — so the new 192/256 kernel can't change anything for these models. My earlier table labeled
    (n_q=16, n_kv=8, head_dim=256) as "the Gemma 4 26B shape", but that's actually the sliding layer;
    the real large-kL shape is (n_q=16, n_kv=2, head_dim=512). Apologies for the mislabel.

    One more correction: the fallback here materializes the score tensor in bf16 (the model's
    activation dtype; precise=True softmax only raises the accumulator to fp32), ~8 GiB at 131K with
    qL=2048 — not an fp32 transient. The peak grows about 2.8–3.0 GB per 32K of context in my runs; a
    bf16 score plus global KV predicts ~2.82 GB (with a small residual from other buffers), whereas an
    fp32 score would predict ~4.96 GB and clearly does not match the observed slope. So this is a real,
    kL-scaling materialization that vanilla does not stream — the 156K→184K admitted-context change I
    reported in production came from relaxing our serving-side admission estimator, not the kernel, and I
    now consider that estimator change an under-estimate we need to fix on our side, not evidence about MLX.

    Net for PR #3660: its new 192/256 full-attention kernel is effectively a no-op for Gemma 4 26B/31B
    (their long-context layers are 512). It can still help models whose actual full-attention head dim is
    genuinely 192 or 256; my result only shows Gemma 4 26B/31B aren't such models, so I wouldn't gate any
    Gemma-4 decision on it. Full 16-run table below.

    Full 16-run mlx-lm A/B (peak GB / prompt_tps)

    Models: unsloth/gemma-4-26b-a4b-it-UD-MLX-3bit and unsloth/gemma-4-26b-a4b-it-UD-MLX-4bit
    (download revision not recorded). mlx-lm 0.31.3; vanilla mlx 0.31.2 vs fused wheel
    0.32.0.dev20260611+0aa50825; mlx_lm.benchmark --model <dir> -p <ctx> -g 8 --num-trials 1 --prefill-step-size 2048; Mac Studio 36 GB. peak at the benchmark's 0.001 GB precision.

    model ctx van peak fused peak van pTPS fused pTPS
    26B 4-bit 32K 19.226 19.226 985.99 986.12
    26B 4-bit 65K 22.050 22.050 810.18 810.45
    26B 4-bit 98K 25.069 25.069 683.31 683.55
    26B 4-bit 131K 28.088 28.088 589.87 590.12
    26B 3-bit 32K 16.304 16.304 983.88 983.06
    26B 3-bit 65K 19.128 19.128 809.80 809.13
    26B 3-bit 98K 22.147 22.147 682.90 682.90
    26B 3-bit 131K 25.166 25.166 589.53 589.33

    (The no-difference conclusion is the prefill peak + prompt_tps above, plus the static dispatch
    analysis; decode -g 8/1-trial is too short to compare and isn't part of the claim.)

  6. zcbenz commented on Jun 14, 2026

    @zcbenz
    Member
    Workload main prompt_tps PR prompt_tps main peak GB PR peak GB
    Qwen 3.6 35B 8bit, 32K 1005.743 871.489 41.316 39.882
    Qwen 3.6 35B 8bit, 64K 802.929 624.513 44.201 40.068

    Thanks for running the benchmark. The performance hit is not small so I think we should not blindly route to fused sdpa? It might be nice if we have heuristics to estimate the needed RAM and then check memory limitation but it is going to be quite hard to get right for all cases on framework's side.

    I wonder if it makes sense adding a force_fused option to fast sdpa op and let the inference engines or users to decide. @angeloskath

  7. DT-HJ commented on Jun 14, 2026

    @DT-HJ
    ContributorAuthor

    That sounds reasonable to me. The Qwen results show that a fixed kL threshold is too blunt, since whether the memory reduction is worth the throughput cost depends on the runtime's resident weights, KV cache, and memory budget.

    Making the 192/256 full-attention fused path opt-in via force_fused seems like a cleaner boundary: MLX can expose the supported kernel, while inference engines or users with the relevant memory context decide when to select it. I would expect the default to preserve the existing MLX dispatch, with the option only requesting the fused path when the shape and other capability checks support it.

    I'm happy to update #3660 accordingly once the intended API semantics are agreed on.

  8. ddalcu commented on Jul 11, 2026

    @ddalcu
    Contributor

    Independent confirmation from a second production runtime (mlx-serve), plus data on how far the composed path can be pushed, and +1 on the force_fused opt-in.

    We ship mlx-serve (a Zig serving runtime on mlx-c) and hit exactly this class: effectively every checkpoint we serve is head_dim=256 (Gemma-4 E4B/26B/31B, Qwen 3.5/3.6 dense and MoE), so the unfused full-attention path is our default long-context reality, not an edge case.

    1. Isolated measurement (MLX 0.32.0, M4 Max 128 GB). One mx.fast.scaled_dot_product_attention call, bf16, mask="causal", GQA shapes, peak via reset_peak_memory:

    shape (heads=16, kv_heads=8) hd=128 peak Δ hd=256 peak Δ heads×qL×kL×2 B
    qL=2048, kL=32768 0.01 GB 2.25 GB 2.15 GB
    qL=8192, kL=32768 0.03 GB 8.99 GB 8.59 GB

    On 0.32.0 with bf16 inputs the fallback's materialized scores measure ≈ heads × qL × kL × 2 bytes (+~5%).

    2. Lazy eval multiplies it. Chaining N attention layers with an mx.eval every 4 layers, ~3.2 layers' score transients coexist at peak (3.44 GB measured for a 1.07 GB/layer shape); eval-per-layer bounds it to ~1× at sub-1% wall-clock cost. So a runtime's effective transient is 1–3× one layer's scores depending on its eval cadence.

    3. Real-world A/B — Gemma-4 26B-A4B (4-bit, quantized KV), cold prefill of a 102,448-token prompt, M4 Max 128 GB, byte-identical outputs:

    prefill strategy (composed path both rows) peak MLX memory prefill tok/s
    fixed 8192-token chunks, eval every 4 layers 51.6 GB 262.7
    context-aware chunk cap (scores ≤ 4 GiB/layer → 1024-token chunks here) + eval-per-layer 27.0 GB 298.8 (+14%)

    Two takeaways that I think support the direction this thread is heading:

    • A context-aware chunk cap makes the composed path both faster and ~2× lighter than naive big-chunk composed, consistent with @zcbenz's Qwen 35B benchmark showing the fused D=256 kernel trades prefill throughput for memory. The fused kernel's remaining value is the memory floor (the last few GB of bounded transient), which matters most on small-RAM machines near their limit, exactly where a throughput hit is an acceptable price.
    • That's why a static kL threshold inside MLX can't pick correctly, and force_fused as an opt-in is the right boundary: the caller is the only party that knows resident weights, KV size, per-request cache quantization, and the machine's budget. A serving engine already computing an admission estimate can flip force_fused precisely when the floor matters and eat the throughput cost knowingly.

    Happy to benchmark a revived #3293 wheel on M4 Max across Gemma-4 (E4B / 26B-A4B) and Qwen 3.6 (27B dense / 35B MoE) at 32K–262K, with dense and quantized KV, if that helps get it over the line.

  9. Thump604 commented on Jul 11, 2026

    @Thump604

    @zcbenz @angeloskath Before rebasing #3660, could you confirm the intended API and semantics for a caller-selected fused full-attention path?

    Exact local reproduction used the stock mlx-lm benchmark path, not a serving runtime:

    python -m mlx_lm.benchmark \
      --model /Users/David/ai-models/mlx_models/Qwen3.6-35B-A3B-VLM-MTP-8bit \
      --prompt-tokens 32768 \
      --generation-tokens 1 \
      --batch-size 1 \
      --num-trials 2 \
      --prefill-step-size 2048

    Against the June comparison bases (MLX main a025496 and #3660
    0aa50825), Qwen 3.6 35B measured 41.316 -> 39.882 GB peak at 32K
    and 44.201 -> 40.068 GB at 64K, while prompt throughput changed
    1005.743 -> 871.489 and 802.929 -> 624.513 tokens/s respectively.
    The #3660 arm did not complete the optional 131K row, so I am not
    claiming a result for that context.

    The relevant upstream paths are use_fallback() and
    sdpa_full_self_attention_metal() in
    mlx/backend/metal/scaled_dot_product_attention.cpp, plus the compiled
    Steel full-attention shapes in
    mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal.

    The minimal ask is not a new default threshold. It is confirmation of the
    preferred opt-in surface: for example, a caller-selected force_fused
    mode, or another API name you prefer, which requests the compatible fused
    full-attention path while preserving current automatic dispatch by default.
    That would let an inference engine make the memory/throughput decision from
    its own resident-memory and request context, while unsupported shapes retain
    the normal path.

    I will prepare a fresh, narrow PR on current main only after that API
    direction is confirmed. This is not a request to merge the current dirty
    #3660 branch, not a claim that its historical kL > 16384 threshold is the
    right global policy, and not a Gemma 4 26B/31B long-context claim: those
    models' global attention uses head_dim=512 and does not exercise this
    192/256 full-attention route.

  10. apocryphx commented on Jul 21, 2026

    @apocryphx
    Contributor

    Another real-world datapoint for reviving #3293, from a different corner of the same model family: we run Gemma-4-31B dense (from-scratch MLX C++ engine, M4 Max 128 GB, bf16 compute, Q4 weights) — 50 sliding-window layers at d=256 (32 q / 16 kv heads) plus 10 global layers at d=512. With the current routing (sdpa_full_supported_head_dim ∈ {64, 80, 128}, re-verified on 0c537a4), no layer of this model fuses at prefill at all.

    Where the OP's evidence is memory (OOM at 100K+), ours is a throughput/scaling datapoint at moderate context, isolated to single mx.fast.scaled_dot_product_attention calls (array mask, bf16):

    L=4096, one layer time
    d=256, 32q/16kv, composite fallback (today) 46.7 ms
    d=128, 32q/16kv, fused full kernel (reference) 21.6 ms

    Per-FLOP the composite path is respectable at this L — the dominant costs are the materialized [n_kv, n_rep, qL, kL] score transients (≈1 GB/layer at 4K, growing quadratically, i.e. exactly the OP's chunk × kv_len wall at long context) and the memory-bandwidth those cost. So our data supports the same conclusion from the throughput side: the fused kernel's value at d=256 is primarily the O(L) memory behavior, and it compounds with context length.

    One nuance worth having in the thread, since both reported models are Gemma-4 variants where most d=256 layers are sliding-window (25/30 for the OP, 50/60 for us): a fused kernel with an arbitrary array mask still processes the full qL×kL area unless it skips fully-masked tiles. For the sliding layers there's an algorithmic mitigation available today that composes with the requested kernel rather than substituting for it: chunk the prefill and trim K/V to the last (window + chunk) positions before the SDPA call — every dropped key is outside every query's window, so outputs are unchanged (we gate this greedy-token-exact). That took our d=256 composite cost from 46.7 → 18.1 ms/layer at L=4096 and scales as O(L·window) instead of O(L²). With #3293 revived, the fused kernel would then run on the bounded [chunk × (window+chunk)] problem — memory and compute both bounded. It does nothing for the global (full-attention) layers, which is where the fused kernel is irreplaceable — that plus the d=512 case is filed separately as #3885.

    Happy to benchmark a revived #3293 on M4 Max (both the dense 31B shapes above and isolated kernels).

  11. beyondsimulations commented on Aug 5, 2026

    @beyondsimulations

    Another data point for this issue, from a shipping macOS app (on-device
    handwriting OCR with Qwen3-VL via mlx-swift): the unfused fallback is not just
    slow — at vision-model prefill sizes it is a hard, unrecoverable crash, and it
    is silent until then.

    Qwen3-VL-8B's vision tower has hidden_size 1152 / num_heads 16 → head_dim
    72, which misses sdpa_full_supported_head_dim (64/80/96/128 on current
    main), so every ViT prefill materializes the full score matrix. Byte-exact
    example from a real 1404×4079 px document page:

    • N = 88 × 254 = 22,352 patches (16 px patches)
    • score buffer = 16 heads × 22,352² × 2 B = 15,987,580,928 B (14.89 GiB)
    • maxBufferLength on a 16 GB M-series machine = 9,534,832,640 B (8.88 GiB)

    The result is [metal::malloc] Attempting to allocate 15987580928 bytes which is greater than the maximum allowed buffer size of 9534832640 bytes raised as
    an uncatchable fatalError, with the process left spinning afterwards — a host
    app cannot catch it, fall back, or even die cleanly. Qwen3-VL-4B (head_dim 64)
    is immune at any image size, which isolates the gate as the cause.

    The model-level fix exists (mlx-swift-lm #455 and mlx-vlm both zero-pad 72→80
    onto the fused kernel — exact, since padded dims contribute nothing to the dot
    products), but every downstream port has to rediscover and re-implement the
    same trick after debugging a crash. Two asks for mlx core, either of which
    would have prevented this:

    1. dispatch the fused kernel for head dims below the largest supported target
      by padding internally (or at minimum add 72, which two model families ship);
    2. independently of routing: surface an allocation that cannot fit
      maxBufferLength as a recoverable error instead of fatalError — callers
      that could chunk their input currently get no chance to.

    Related: the silent-fallback observability gap discussed in #3885 is exactly
    what made this expensive to debug — there is no way to know you fell off the
    fused path until the allocator aborts.

  12. zcbenz commented on Aug 11, 2026

    @zcbenz
    Member

    Let's add a force_fused=True option to mx.fast.scaled_dot_product_attention API that ignores the heuristics and uses fused path, and throws when a fused kernel is not available for the shape. I would be happy to review a PR implementing it.

  13. changed the title [-]head_dim=256 unfused SDPA makes 100K+ context prefill impractical — real-world evidence for reviving #3293[/-] [+]Allow force using fused kernel in scaled_dot_product_attention[/+] on Aug 11, 2026
  14. apocryphx commented on Aug 11, 2026

    @apocryphx
    Contributor

    @zcbenz We'll take that PR, if nobody has started — with one safety pairing from the #3885 investigation that force_fused makes load-bearing: the vector dispatches got their check_kernel_threadgroup_size in #4018 and #3099, but the full-attention dispatch (sdpa_full_self_attention_metal) still launches unchecked. Under heuristic routing it only runs at shapes whose pipelines are known to fit; force_fused makes register-limited configurations reachable — and on M1-class GPUs an oversized launch is silently dropped and returns zeros (measured; details in that thread). Wide head dims are exactly where this bites: in our head_dim=256 experiments the specialized full-kernel pipelines on M4 Max already report maxTotalThreadsPerThreadgroup = 128 — precisely the 4-warp dispatch size, zero headroom — so devices with smaller register budgets would cap below the launch. The check should extend to the full dispatch together with the option.

    Proposed semantics, objections welcome:

    • force_fused=True bypasses use_fallback entirely; vector-vs-full selection still follows query length as today.
    • Throws with the shape and the supported-dims list when the selected kernel has no instantiation for the head dims, or when the kernel's mask/layout constraints aren't met — loud, never silent.
    • Metal first; CUDA consults the same flag against its own support matrix so semantics stay uniform across backends.
    • Default false; zero behavior change for existing callers.

    Vector-kernel d=512 instantiations (with the BN=16 threadgroup shape M1-class needs — three-device sweep data in #3885) would follow as a separate small PR. Since force_fused is explicit opt-in, they would ship with no routing-heuristic changes and zero behavior change for default callers — which removes the regression risk that stalled the earlier head-dim expansion.

  15. zcbenz commented on Aug 19, 2026

    @zcbenz
    Member

    Closed by #4185.

  16. ddalcu commented on Aug 19, 2026

    @ddalcu
    Contributor
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions