Repository navigation
[WIP] rocm cdna warp size qwen35 mrope batch - #2731
Open
daemyung-lablup wants to merge 21 commits into
Open
daemyung-lablup wants to merge 21 commits into
daemyung-lablup wants to merge 21 commits into
Conversation
The .hip sources are compiled by custom hipcc commands that do not inherit the mlx target's compile definitions, so MLX_HOST_WARP_SIZE never reached their host pass and device/config.h fell back to 32. On a CDNA-only build (gfx9xx, wave64) the device side still used 64, so every host launch that sized its block with WARP_SIZE (scan, reductions) ran 32-lane blocks of 64-lane kernels. On an MI300X (gfx942) cumsum and all() returned wrong values, the gated-delta chunked prefill scan produced NaN, and Qwen3.8-27B-4bit generated only "!" for most prompts. The warp-size selection now runs before the HIP object loop and its value is passed to hipcc as -DMLX_HOST_WARP_SIZE. RDNA builds are unchanged (host and device were both 32). A multi-arch RDNA+CDNA build still pins the host value to 32 and keeps the mismatch on CDNA devices. Verified on MI300X: cumsum and all() match the CPU device, the 25 models::gated_delta tests pass (4 failed before), and Qwen3.8-27B-4bit answers correctly at 33 to 1023 prompt tokens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…code The VLM wrapper resolves MRoPE position ids for every sequence, text-only ones included, so forward_batched_decode_with_caches always saw has_mrope = true and ran each full-attention layer once per row (slice, q/k/v/o projections, MRoPE and SDPA per sequence). On Qwen3.8-27B that is 16 layers x B separate attention forwards per decode step. When every row carries position ids they are stacked to [3, B, 1] and the layer runs through forward_split_attention with a SplitAttention wrapper that applies the same interleaved MRoPE to the whole batch. The MRoPE math is factored into Qwen3NextAttention::apply_mrope and shared with the per-row path, which is unchanged. Rows with and without positions keep the per-row path. Verified on MI300X with Qwen3.8-27B-4bit: greedy output of 8 prompts is identical alone and as one batch of 8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
forward_batched_prefill returned the logits and left every row's cache lazy. The scheduler evaluates only the logits, which do not depend on the last layers' recurrent state, so each row's gated-delta state kept its layer's whole chunked-scan graph alive until the row's first decode step. On Qwen3.8-27B-4bit that was about 2.7 GB per row at 448 prompt tokens: a burst of 64 requests grew active memory by 11 GB per 4-row prefill batch and exhausted a 192 GB MI300X before any row decoded, failing the batch with "HBM full after cache reclaim". The cache arrays are now scheduled with async_eval_all together with the logits, so the shared graph runs once and its intermediates are freed. Measured on MI300X: peak memory for a 64-request burst drops from 176 GB to 42 GB, and a 128-request burst completes with no allocation failures. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…steps Every batched decode step stacked each row's gated-delta state into one batch tensor per linear layer and copied every row back out afterwards: two launches per row per layer, about 9,000 small copies per step for 64 Qwen3.8-27B rows, which kept the step launch-bound. GatedDeltaBatchReuse keeps the stacked state from the last step. When the next step runs the same sequences in the same order and every row still holds the view it was handed (checked by array identity), the stacked state is used as is, and rows receive slice views instead of copies. Any other batch restacks as before. A view keeps the whole stacked buffer alive, so rows that leave the batch and rows about to be snapshotted for the prompt cache are materialized into their own copy first. Adds two small FFI helpers: array_id (MLX array::id) and array_share (a second handle to the same array). Measured on MI300X with Qwen3.8-27B-4bit at 64 concurrent requests: median decode step 225 ms -> 195 ms. Greedy output is byte-identical to the previous build on the B1-vs-B8 parity prompts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
fast_rope_batched rotated every row on its own and joined the rows with a growing two-way concatenate, so a batch of B rows cost B rope launches and O(B^2) copied elements per call. Batched decode calls it twice per attention layer; at 64 Qwen3.8-27B rows that was about 6,000 launches per decode step. On ROCm the batch now goes through one fast::rope call with a [B] int32 offset vector, over a contiguous copy so the batch axis has its own stride (the per-row path exists to avoid confusing the two on transposed T = 1 inputs). The ROCm vector-offset kernel computes the angle exactly as the one-row kernel does, and a decode-shaped bf16 batch of 16 rows rotates to the same bytes as the per-row path. Metal and CUDA keep the per-row path; their kernels were not checked for that equivalence. Adds the fast_rope_offsets FFI. Measured on MI300X with Qwen3.8-27B-4bit at 64 concurrent requests: median decode step 195 ms -> 171 ms. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The ROCm vector (decode) SDPA kernel refused any array mask, so a masked single-query attention fell back to the unfused graph: a QK^T matmul, a select, a full-size copy and a softmax per call. kernel_sdpav_1pass now reads an optional additive mask over [B, H, qL, kL] through its strides (the frontend broadcasts the mask and turns a boolean one into 0 / -inf before the primitive sees it). A -inf entry drops the key the same way the causal check does; every lane of a 32-lane tile reads the same key, so the branch is uniform across the tile's shuffles. Any other value is added to the score in base 2. A row masked down to its own length therefore reduces exactly the keys, in exactly the order, an unmasked call over that length would. The 93 mlxcel-core attention, SDPA and attend tests pass on MI300X. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Qwen 3.5 keeps one dense KV cache per sequence, so batched decode attended row by row: an update, a fetch and an SDPA launch per row and full-attention layer, because the rows' live lengths differ. StackedKvReuse (models::stacked_kv) keeps a [B, n_kv_heads, capacity, head_dim] buffer per attention layer for the current batch. One scatter writes every row's new K/V at its own position and one masked SDPA attends the whole batch; each row's cache adopts a view of the buffer (KVCache::adopt_dense_fp16), so it stays a valid dense cache for every other consumer. The next step reuses the buffer while every row still holds its view; otherwise it restacks. Rows that leave the batch and rows about to be snapshotted are materialized first, as for the gated-delta state. Caches that are not dense FP16 from slot 0 keep the per-row path. With the ROCm vector SDPA kernel taking the length mask, each row reduces the same keys in the same order as when attended alone; greedy B1-vs-B8 parity output is unchanged from the previous build. Also folds the projection code shared by the batched paths into Qwen3NextAttention::project_batched and drops the two batched-decode methods that no longer have callers. Measured on MI300X with Qwen3.8-27B-4bit at 64 concurrent requests: median decode step 171 ms -> 140 ms; launches per step 22,000 -> about 6,800 together with the earlier gated-delta and rope changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ROCm had no gated-delta kernel: gated_delta_kernel_available() is Metal-only and the overlay's GatedDeltaUpdate primitive always falls back. Decode ran the compiled single-step graph (several passes over the float32 state) and prefill ran the 64-row chunked scan, a 63-step Horner matmul loop per linear-attention layer that hipBLAS splits into thousands of small GEMMs. On Qwen3.8-27B that scan was most of TTFT and, under load, of TPOT through interleaved prefills. hip_gated_delta_forward ports the Metal gated_delta_step kernel through fast::hip_kernel: one 32-lane group per (sequence, value head, value row) walks the tokens in order with its slice of the state in registers, for decode (T = 1) and prefill (T > 1) alike. Arithmetic is float32 and the state keeps its input dtype. Both lane reductions fold with an explicit width of 32 and the block is (32, 4, 1), so on a 64-lane wavefront each half of the wave reduces its own row; it is therefore validated on gfx942 directly instead of going through KernelPorts, whose ROCm ports are held to 32-lane devices. gated_delta_ops_with_parity takes it for the scalar-gate, unmasked shape on ROCm; chain-parity blocks keep the graph path. MLXCEL_GDN_HIP_KERNEL=0 turns it off. Against the graph reference on MI300X (bf16 q/k/v, float32 state): bit identical to the sequential loop at 1, 5 and 31 tokens, and within 2e-5 relative error of the chunked scan at 64 and 200 tokens. Measured with Qwen3.8-27B-4bit (448-token prompt): TTFT 708 -> 321 ms and TPOT 28.3 -> 24.6 ms at B=1; at B=16 TTFT 4.3 s -> 0.95 s, TPOT 138 -> 77 ms, 63 -> 125 tok/s; at B=64 77 -> 217 tok/s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rnels A Qwen 3.5 linear-attention layer launched about 43 kernels per decode step on ROCm, eight of them matrix products; the rest were 4 us launches the host could not issue fast enough (the GPU idled 26% of a B=1 step and 42% of a B=64 step on MI300X, BF16 checkpoint). hip_gated_delta_fused_conv takes the raw depthwise-conv output and does, per token, the SiLU, the no-weight RMS normalisation and scale of q and k, beta = sigmoid(b) and g = exp(-exp(A_log) * softplus(a + dt_bias)) before the recurrence, with MLX's stable sigmoid and every rounding the graph makes. Each 32-lane group recomputes its head's prep per token, so it is used for blocks of at most 8 tokens: on a 448-token prefill the 128-fold recomputation made the layer slower than the batched graph prep (TTFT 319 -> 473 ms at B=1), and prefill keeps the graph prep with the plain recurrence kernel. hip_rms_norm_gated replaces RMSNormGated's six launches (rms_norm, two casts, float32 SiLU, multiply, cast) with one, one 32-lane group per row, for prefill and decode alike. Both match the graph bit for bit on MI300X: the fused layer core at 1, 7, 64 and 300 tokens with Qwen3.8-27B head geometry, the gated norm at 48, 3,072 and 21,504 rows. MLXCEL_GDN_HIP_KERNEL=0 turns both off with the recurrence kernel. Quick ladder (one round): decode TPOT at B=1 drops 24.6 -> 20.9 ms on the 4-bit checkpoint and 27.2 -> 23.5 ms on BF16. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The fused decode kernel from f5fe633 ran one block per four value rows and recomputed its head's SiLU, q/k normalisation, gate and beta in every block: 128 copies per head per token on Qwen3.8-27B. That cost grew with the batch and made the kernel 4x slower than the plain one at 64 rows, so the B=1 gain did not reach large batches (BF16 on MI300X: GDN 13.4 ms of a B=64 step, 26.8 ms at B=128). A block now owns one (sequence, value head): 16 rows of 32 lanes, each thread keeping Dv / 16 value rows of the state in registers. Per token one 32-lane group prepares q and k, another the SiLU of every v row and one lane the gate and beta, into LDS, once. The fused path is limited to single-token decode steps; an 8-token block measured 13x slower than the unfused path. Both the fused wrapper and gated_delta_ops_with_parity copied the incoming state before the update, and Qwen 3.5 copied it again out of the cache; every consumer only reads it, so all three now take a second handle (array_share) instead of moving 3 MiB per row per layer. Against the unfused path (graph preparation plus the plain kernel) the output y differs in at most 54 of 1.57 M bf16 elements at B=128 and the state by at most 6e-8 absolute: float rounding, not bit identity as with the previous layout. Every decode row takes the same kernel, so batched and single-row decode still agree. MI300X, Qwen/Qwen3.8-27B BF16, quick ladder: TPOT p50 at B=64 92.2 -> 73.5 ms; pure B=128 decode step 89.4 -> 70.6 ms (GDN 26.8 -> 9.2 ms). B=128 TPOT is now bound by interleaved prefills. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…generation `lookahead_params` kept every run synchronous whenever `max_kv_size` was set. The server also sets it from `-c` (each slot's share of the context), so any server started with `-c` never pipelined decode, whatever the model. Without context shift the cap never trims the cache mid-decode; it only stops generation (`ContextBound::stop_due`), which the pipelined finish step applies too. Keep the pipeline in that case and decode synchronously only when a row has no room under the cap for the two speculative appends. Context shift still forces the synchronous path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…peline Qwen 3.5 keeps its gated-delta and attention caches in the model, so the scheduler pipelines it only when the model can undo a speculative step (`supports_decode_lookahead_rewind`). Record each sequence's state before a speculative decode forward (array handles, no copies: MLX arrays are immutable) for at most the two appends a teardown unwinds, and restore it in `rewind_decode_appends`: conv and recurrent state plus offset for the linear layers, a trim for the attention layers. A synchronous step drops the history instead of recording it, since it is never unwound and holding it keeps whole stacked states alive (about 19 GB a step at B=128). Prefill, restore and release forget the history. The VL wrapper delegates both hooks. `decode_lookahead_append_active` exposes the existing speculative-append scope so a model can tell the two kinds of step apart. Greedy outputs match MLXCEL_FORCE_SYNC for single, concurrent and staggered requests (the staggered case differs between two sync runs too, from batch composition). Qwen 3.8 27B BF16 on MI300X, 447-token prompt, 256 tokens: B=1 41.2 -> 46.0 tok/s, B=16 449 -> 514. Large batches gain little; see the following ROCm cache commit for B=128. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The 2 GiB cache bound (lablup#2062) was chosen for gfx1151, whose GPU shares host memory. On a discrete GPU it makes the large buffers a decode batch reallocates every step miss the cache: Qwen 3.8 27B BF16 at B=128 on MI300X frees and reallocates about 19 GB of gated-delta state per step through hipFree/hipMalloc. With the decode lookahead holding one more step of state, the device filled (rocm-smi 206 GB used while MLX counted 122 GB active) and single hipMalloc calls blocked for up to 2.9 s: 263 tok/s pipelined against 1,129 synchronous. Without the bound the pipelined run is 1,129 tok/s. The ROCm device_info now publishes `integrated`, read through the new `gpu_is_integrated` bridge call; the bound applies to integrated GPUs and when that is unknown. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…he warm-up An idle scheduler with queued prompt-cache warm-ups (lablup#1144) started one immediately, and a warm-up forward cannot be interrupted. A client that sent its next request as soon as the last one finished waited behind it: about 50 ms of TTFT for Qwen 3.8 27B BF16 on MI300X, measured as the gap between the client sending and the scheduler admitting the request. Wait up to 100 ms for a request first; warm-ups target a conversation's next turn, which rarely follows that closely, and still run once the server is idle (7 of 8 ran with 0.5 s between requests, none while requests arrived back to back). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ete GPUs The history-boundary split (lablup#1143) cleared MLX's buffer cache between the boundary segment and the suffix, so the segment's intermediates did not stay resident through the rest of the prefill. That matters where the GPU shares host memory. On a discrete GPU (ROCm reports it, see `gpu_is_integrated`) the freed buffers are reused by the suffix and the decode, and releasing them cost a hipFree per buffer plus the hipMalloc to get them back: 13 to 18 ms per prefill for Qwen 3.8 27B on MI300X. Integrated and unknown devices (Metal, CUDA) keep the clear. With the previous commit, Qwen 3.8 27B BF16 B=1 TTFT (447-token prompt, requests back to back) went from 212 to 171 ms; with the boundary snapshot disabled it is 140 ms, the rest being the suffix forward the split needs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The gated-delta recurrence, the fused decode step and the gated RMSNorm HIP kernels (986f8c5, f5fe633) chose their port by comparing the backend kind and called their holders directly, which `make verify-kernel-port-dispatch` rejects. They now resolve through ROCm-only `KernelPorts` tables, and `gated_delta_hip_kernel_available` answers from those tables (plus the existing `MLXCEL_GDN_HIP_KERNEL=0` switch). The kernels are shuffle-based, so rule 5 (lablup#2147) would hold them to 32-lane devices, which would turn them off on MI300X, where they were written and validated. Their shuffles name a width of 32 inside 32-lane block rows. The checker gains `VALIDATED_SHUFFLE_ANY_WAVE`: shuffle-based tables that were run on a 64-lane device and are allowed at any width, pinned with their HIP source like `EXPECTED_ANY_WAVE` but exempt from the lane-operation scan. This is the per-kernel decision lablup#2147 left open and is a policy change for review. The checker's fixture test covers the new pins (mark dropped, entry swapped). They have not been run on a 32-lane (RDNA) device, where they also ran before this change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
clippy::question_mark under make verify-clippy-rocm. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… to it A port table had one ROCm entry, either held to 32-lane devices or allowed at any width. A kernel written for CDNA could only run there by being allowed everywhere, including RDNA devices it never ran on, which the previous commit did for the gated-delta kernels through a pinned checker exception. `KernelPorts` gains `rocm_cdna`, the entry for 64-lane (CDNA) devices: `port_for` picks it on a 64-lane device ahead of `rocm` and never on a 32-lane one (`rocm_port_choice`, pure and tested through the bridge like `rocm_port_allowed`). `select_kernel_port` names a CDNA-only refusal. The gated-delta recurrence, fused decode step and gated RMSNorm move to `rocm_cdna`, so RDNA keeps the graph as it did before they existed, and the checker exception (`VALIDATED_SHUFFLE_ANY_WAVE`) is removed again; the checker's docs say where a CDNA kernel goes. On MI300X the three kernels still run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…te ring Every batched decode step allocated a fresh recurrent state per linear layer (about 19 GB per step for Qwen 3.8 27B at 128 rows): MLX arrays are immutable and the gated-delta kernel wrote `state_out`. On ROCm that allocation counts against the in-flight budget, so the host waited for the GPU in the middle of building each step and the decode lookahead could not overlap the two. On CDNA the fused decode kernel now has an in-place variant (`hip_gated_delta_fused_conv_into`, the same source under `GD_STATE_DST`) that writes rows [0, B) of a caller-owned buffer and returns them ordered after the kernel through `depends` and a slice, allocating nothing. `GatedDeltaBatchReuse` keeps three such buffers per layer and hands each step the one holding neither of the last two states, so the views a lookahead teardown can rewind to (two steps) stay intact; buffers grow to the batch rounded up to a power of two and survive batch changes. Rows leaving the batch and snapshots were already copied out of the stacked state, which covers the ring the same way. Other backends, and `MLXCEL_GDN_STATE_RING=0`, allocate as before. Qwen 3.8 27B BF16 on MI300X (447-token prompt, 256 tokens, same binary): B=64 959 -> 1,084 tok/s (TPOT 55.8 -> 47.5 ms), B=128 1,089 -> 1,210 (92.5 -> 81.6 ms); B=1 and B=16 unchanged. Greedy outputs are identical with the ring on and off for single, concurrent and staggered requests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Same checkpoint, same workload on both engines (447-449 input tokens, exactly 256 output tokens per request, greedy, B = 1 to 128): the ladder at each step of this branch, the per-commit A/B measurements, the decode-step profiles, and what was tried and left out. The summary file is the short version. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
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
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.
This PR is for testing and review only. Do not merge.
It collects the MI300X (CDNA 3) investigation of Qwen3.8-27B against vLLM in one branch so the changes and measurements can be read together; the pieces that should land will go in as separate, smaller PRs. ## Summary mlxcel
main(2f5fb7a5) did not run correctly on an AMD Instinct MI300X (gfx942, wave64): Qwen3.8-27B generated only!, mlxcelmain(2f5fb7a5) did not run correctly on an AMD Instinct MI300X (gfx942, wave64): Qwen3.8-27B generated only!, and B ≥ 64 ran out of memory. This branch fixes that and then closes part of the gap to vLLM on the same checkpoint and workload.Changes
Correctness and memory
b6c45909pass the host warp size to the hipcc-compiled HIP objects; on a CDNA-only build host code assumed 32 lanes anddevice code 64, which broke
cumsum/all()and produced!-only output.a5b6a75fevaluate batched-prefill caches with the logits (B=64 peak 176 → 42 GB; B=128 no longer OOMs).6cb2bbcckeep MLX's buffer-cache default on discrete GPUs; the 2 GiB bound (perf(rocm): explain and bound allocator peak memory on UMA hosts #2062, tuned on gfx1151) now applies tointegrated GPUs only. On MI300X the bound made every step
hipFree/hipMalloc~19 GB of state and, with the decodelookahead, filled HBM (B=128: 263 tok/s).
Decode
e819f35f,ba920c59,186ad7bd,84838fcf,24c91d0bkeep Qwen 3.5's batched decode batched: MRoPE attention, reusedDecode
e819f35f,ba920c59,186ad7bd,84838fcf,24c91d0bkeep Qwen 3.5's batched decode batched: MRoPE attention, reusedstacked gated-delta state, one-kernel batch RoPE, masked vector SDPA, stacked KV buffer (22,000 → ~6,800 launches per B=64
step).
986f8c5e,f5fe633b,b7c23d45gated-delta recurrence, decode-step preparation and gated RMSNorm as HIP kernels.8d15a3fckeep the decode lookahead when the KV cap only stops generation.-csets a per-slot cap, which kept everyserver started with
-coff the lookahead pipeline, for every model.5e428e84Qwen 3.5 rewinds speculative decode appends, so it can use the lookahead (B=1 and B=16 +12–15%).c5ef46beon CDNA, batched decode writes gated-delta state into a three-buffer state ring instead of allocating it eachstep (B=64 +13%, B=128 +11%; output identical with it on and off).
TTFT
c23ce6c2wait up to 100 ms for a request before an idle prompt-cache warm-up, which otherwise delayed a back-to-backrequest by ~50 ms.629bd500keep the buffer cache at the history boundary on discrete GPUs (13–18 ms per prefill). With the previous commit, B=1 TTFT 212 → 171 ms.Structure (review requested)
68e1f051,5e8f9decthe gated-delta HIP kernels go throughKernelPortstables, and the tables gain a CDNA-only entry,.rocm_cdna, selected on 64-lane devices only (rocm_port_choice, unit-tested). The gated-delta kernels live there, so RDNA keeps the graph path as before this branch.MLXCEL_GDN_STATE_RING(kill switch, documented).Measured and left out
MLX_ROCM_MAX_INFLIGHT_MBon CDNA: helped before the state ring, slower after it.MLX_ROCM_HIPBLASLT_TUNE=1: no change.Remaining gap (rocprofv3, decode step)
Testing
cargo testfor the touched areas (qwen3_5, gated_delta, recurrent_batch, scheduler lookahead, prompt cache, runtime, kernel_port, fused-norm parity on wave64).make verify-clippy-rocm,verify-fmt,verify-kernel-port-dispatch(and its fixture test),verify-kernel-dtype-keys,verify-env-var-docs,verify-versions,verify-rocm-overlay,verify-llama-compat,verify-python-tooling.Not tested: RDNA (gfx1151), Metal, CUDA. The shared Rust changes (lookahead gate, Qwen 3.5 rewind, warm-up grace) also run there.
Known issues