fix(cuda/quantized): chunk long-prompt qmm launches to avoid gridDim>65535 and int overflow#652
Merged
Merged
Conversation
…65535 and int overflow Long-prompt prefill of 4-bit models aborted with an invalid CUDA launch (cuGraphAddKernelNode / cuLaunchKernelEx "invalid argument", SIGABRT). Two distinct causes in the quantized-matmul path (#648): - MoE GatherQMM sets qmm_sm80 grid.z = l = tokens * num_experts_per_tok, which exceeds the CUDA gridDim.z limit of 65535 once tokens*top_k >= 65536 (qwen3-30b-a3b @8192, mixtral-8x7b @32768). - The dense LM head computes l = out.size() / (m * n) with int32 m,n, which overflows when tokens*vocab >= 2^31 (about 16744 tokens for a 128256 vocab), yielding l = 0 and an invalid grid.z of 0. Fix, all in the mlxcel-owned overlay quantized.cpp: - run_batch_chunked splits GatherQMM's gathered batch into <=65535 slices. - run_row_chunked splits QuantizedMatmul's M into slices of min(65535, INT32_MAX/N) so each launch keeps grid.y <= 65535 and m*n < 2^31. - int B = out.size() / M / N (sequential 64-bit division) so the overflow no longer yields B=0 and misroutes dispatch. Both early-return to the original single call when no split is needed, so normal-size matmuls are byte-identical. Validated on GB10: previously FAIL:bench long-prompt cells now pass (llama-3.1-8b-4bit @17000/@23491/@32768, qwen3-30b-a3b-4bit @8192), with coherent generation confirming the sub-array offsets are correct. CUDA-only; the Metal backend uses out.size()/M/N and is unaffected. Refs #648
inureyes
added
area:benchmark
inureyes
added a commit
that referenced
this pull request
Jul 6, 2026
…gible models (#676) * fix(cuda/attention): bound long-prompt prefill memory to O(chunk*L) (#672) gemma-4/gemma-3 (head_dim 256) disqualify both CUDA fused SDPA kernels (cuDNN caps head_dim at 128, sdpa_vector covers 64/96/128 with q_len < 4), so a long single-pass prefill fell to MLX's materializing fallback: O(heads*L^2) attention scores (~68 GiB at 32k for the 31b), full-sequence 262k-vocab logits with a same-size final_logit_softcapping copy, and two dense [L, L] f32 masks together pushed GB10 past its budget and OOM'd (#672). Four coordinated changes, each gated so short-context behavior stays byte-identical: - layers::attention gains a query-chunk gate: when a CUDA config cannot reach a fused SDPA (head_dim > 128, softcap composite, or non-f16/bf16 dtype) and the score matrix would exceed MLXCEL_ATTENTION_CHUNK_BUDGET_MB (default 1024, 0 disables), the query axis is split into near-equal chunks; masks are row-sliced per chunk and cast to the score dtype, because MLX's fallback adds the mask without casting and an f32 mask silently promotes the masked scores and softmax to f32. The maskless causal path in causal_attention routes through the same chunking with per-chunk causal masks and K/V sliced to the causal bound. - LanguageModel::forward_last_logits(last_pos) lets the single-sequence prefill in generate_streaming/generate_with_stats request only the sampled position's logits. The default implementation is the previous forward-plus-slice; the gemma4 override projects a single hidden row through the 262k-vocab LM head for prefills longer than 4096 tokens, dropping ~17 GiB of logits plus an equal softcap copy at 32k. - gemma4 leaves the prefill mask pair None for text-only prefills longer than 4096 tokens, routing full-attention layers through the causal flag path and sliding layers through the over-window prefill path instead of retaining two dense [l, l+offset] f32 masks (~8 GiB at 32k) across the whole forward; verify, padding, per-row valid-end, and vision-overlay shapes keep the dense masks unchanged. - create_causal_mask / create_causal_mask_with_window_full build from broadcast index comparisons instead of the six-buffer ones/tril/triu/multiply/greater/where chain, cutting mask-construction transients from ~24 GiB to ~7 GiB at 32k. mlxcel-bench-decode now prints the MLX allocator peak after each run. Validated on GB10 CUDA: gemma-4-12b-it-4bit completes a 32768-token single-pass prefill (previously the OOM class) with MLX peak 70.5 -> 56.8 GB across the optimization steps; gemma-4-31b-it-4bit at 16384 drops 58.4 -> 54.3 GB; forced-chunk A/B on gemma-4-31b is token-identical to the unchunked path on a natural prompt; 13 new unit tests cover chunked-vs-unchunked equivalence (mask, softcap GQA, broadcast mask, causal offset), chunk-length math, mask row slicing, and forward_last_logits; clippy -D warnings clean on default and cuda features. * feat(generate): opt-in chunked prefill via MLXCEL_PREFILL_CHUNK (#672) The single-sequence CLI/bench prefill ran the whole prompt in one pass, so even with the #672 attention/logits/mask fixes a 32k gemma-4 prefill held every sliding-window layer's full-length KV plus chunk transients for the entire pass. The server already bounds this with prefill_chunk_size, and upstream mlx-vlm/mlx-lm chunk by default (DEFAULT_PREFILL_STEP_SIZE = 2048 with a per-model chunked_prefill_policy opt-out); the CLI path was the odd one out. chunked_prefill_last_logits feeds the prompt through forward_last_logits in MLXCEL_PREFILL_CHUNK-token pieces (default 0 = single-pass, unchanged), evaluating each piece so the lazy graph never spans the whole prompt and rotating sliding-window caches trim to their window between pieces. Intermediate pieces project one hidden row through the LM head and drop the [1, 1, vocab] result. Each piece also returns freed buffers to the OS (clear_memory_cache): every piece sees a different key length, so its transients land in differently-sized allocations, and without the release the CUDA async-malloc pool accumulated each shape's high-water mark across the prompt (measured 84 GB system peak for a 32k gemma-4-31b chunked prefill whose live set is ~30 GB). Measured on GB10 CUDA at 32768 prompt tokens: gemma-4-12b-it-4bit MLX peak 56.8 -> 16.6 GB (system 23 GB) with prefill 285 -> 706 tok/s; gemma-4-31b-it-4bit completes at MLX peak 29.9 GB (system 38 GB, previously the OOM/freeze class), prefill 271.6 tok/s. At 8192 tokens decode is unchanged (13.99 vs 14.03 tok/s over 64 tokens). Kept opt-in until multi-call prefill parity is validated across model families; a follow-up can flip the default with a per-model policy hook mirroring mlx-vlm. Includes a chunked-vs-single-pass equivalence unit test over a stateful stub (token order, exactly-once feeding, final-logits equality). * chore(benchmarks): add post-fix long-prompt sweeps for #648 and #672 cuda_gb10_longprompt_2026-07-04.csv is the post-#652 rerun of the long-prompt sweep on the clean rebuilt binary: llama-3.1-8b-4bit now completes 32768 prompt tokens (previously FAIL:bench from the qmm gridDim/int-overflow crash, #648). cuda_gb10_gemma_longprompt_672_2026-07-06.csv records the #672 gemma-4 validation with memory columns: MLX allocator peak plus sampled system peak, single-pass vs MLXCEL_PREFILL_CHUNK=2048 runs, including gemma-4-31b-it-4bit completing 32768 prompt tokens at 29.9 GB MLX peak (previously OOM at ~101 GB).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Long-prompt prefill of 4-bit models aborted on the CUDA backend with an invalid launch (
cuGraphAddKernelNode/cuLaunchKernelEx"invalid argument", SIGABRT). This surfaced asFAIL:benchcells in the long-prompt sweep (epic #623/#624) and as crashes inmlxcel generateon long prompts. Two distinct causes in the quantized-matmul path:qmm_sm80setsgrid.z = l = tokens * num_experts_per_tok, which exceeds the CUDAgridDim.zlimit of 65535 oncetokens*top_k >= 65536(qwen3-30b-a3b @8192, mixtral-8x7b @32768).l = out.size() / (m * n)is computed withint32m,n, which overflows whentokens*vocab >= 2^31(about 16744 tokens for a 128256 vocab), yieldingl = 0and an invalidgrid.zof 0.Fix
All in the mlxcel-owned overlay
src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:run_batch_chunkedsplits GatherQMM's gathered batch into<= 65535slices (sub-array views over the same buffer).run_row_chunkedsplits QuantizedMatmul'sMinto slices ofmin(65535, INT32_MAX/N)so each launch keepsgrid.y <= 65535andm*n < 2^31. Gated onsingle_batch = out.size() == size_t(M)*N(the LM-head output is 3D[1, tokens, vocab], so anndim==2gate would miss it).int B = out.size() / M / N(sequential 64-bit division) so the overflow no longer yieldsB=0and misroutes dispatch.Both helpers early-return to the original single call when no split is needed, so normal-size matmuls are byte-identical. The change is CUDA-only; the Metal backend already uses
out.size()/M/Nand has no 65535 grid cap, so it is unaffected.Validation (GB10, NVIDIA sm_121)
Previously-failing long-prompt cells now pass, exit 0 with real throughput:
scripts/bench_longprompt.shruns):llama-3.1-8b-4bit@17000 / @32768,qwen3-30b-a3b-4bit@8192.mlxcel generate:llama-3.1-8b-4bit@23491 andqwen3-30b-a3b-4bit@~11k complete with coherent output, confirming the sub-array offsets are correct.Closes #648