[Apple Silicon] Add (env-gated) MLX paged attention decode - #28335
adityavaid wants to merge 5 commits into
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
changminbark
left a comment
There was a problem hiding this comment.
Overall LGTM. It would be nice to have a quick section describing how you would be using this paged metadata. Sorry if I missed anything since I was missing for a while.
4f8f7ef to
262b4cf
Compare
262b4cf to
ed26a10
Compare
| out.set_data(allocator::malloc(out.nbytes())); | ||
|
|
||
| auto& d = metal::device(stream().device); | ||
| auto& enc = d.get_command_encoder(stream().index); |
There was a problem hiding this comment.
I don't know if this is the metal version diff or something else. Mine needs to update to metal::get_command_encoder(stream());.
There was a problem hiding this comment.
done. Thanks, this was an issue with my MLX version, updated now.
| seq_lens: list[int] | ||
| max_seq_len: int | ||
| total_kv_tokens: int | ||
| page_size: int = 1 |
There was a problem hiding this comment.
Why is this 1? I remembered that vllm-metal uses 16 (if I recall correctly). Is this because of the existing flat MlxAttentionKVPool layout?
There was a problem hiding this comment.
Yes, current MLX KV pool is flat token slots. We can change it in the future but for now this would be out of scope.
a8556b6 to
bd9c586
Compare
jlee5814
left a comment
There was a problem hiding this comment.
Doesn't compile on MLX 0.31.2
Duplicate NB_MODULE(_metal)
(bindings.cpp:8, metal_common.cpp:8), both compiled into _metal (setup_metal.py:101-105) → ld rejects the duplicate symbol. Please drop one.
d.get_command_encoder(stream().index)
(rope_pool_fused.cpp:107, paged_attention.cpp:44): no such member on 0.31.2. So the bench can't be from this tree. Please use the free metal::get_command_encoder(stream()).
No parity test (test paged vs padded SDPA).
Gate isn't a correctness boundary
(attention_wrapper.py:224): checks kernel/metadata presence, never model class, and the kernel runs unmasked full attention. The flag on a sliding-window / MLA / quantized-KV model (future-work admits unsupported) silently returns wrong tokens, no fallback. Please return None metadata for those.
model_runner_stub.py:114 revives initialize(self, pre_model_load_memory=None), Please see #28660.
Adding
Fixed.
Adding in next commit
Need a bit more clarity on this, I can add explicit correctness check before we build metadata.
|
|
Tested on Apple M4 Pro, macOS 26.3.1, MLX 0.31.2, clang 17, at HEAD Duplicate // metal_common.cpp:52
return metal::device(stream.device).get_command_encoder(stream.index);
// error: no member named 'get_command_encoder' in 'mlx::core::metal::Device'The free function ( return metal::get_command_encoder(stream);Happy to run the full build + parity test on 0.31.2 once you push. |
Fixed, Thanks @LijuanTang94. Can you run this again to confirm everything is working fine or not ? |
|
Re-tested on MLX 0.31.2 (M4 Pro): the C++ type-checks clean now ✅ — |
Removed unused imports and profiling functions from bench_one_batch.py.
yeahdongcn
left a comment
There was a problem hiding this comment.
Can you refresh the perf section? The current command in the PR description uses --skewed-input-lens, but that flag is no longer present after the cleanup commit.
| self.weight_load_mem_usage = 0 | ||
|
|
||
| def initialize(self): | ||
| def initialize(self, pre_model_load_memory: float | None = None): |
There was a problem hiding this comment.
This still looks stale relative to the current MLX stub on main / #28660. Please drop the pre_model_load_memory argument and keep def initialize(self):.
| seq_lens: list[int] | ||
| max_seq_len: int | ||
| total_kv_tokens: int | ||
| page_size: int = 1 |
There was a problem hiding this comment.
I think we should be stricter here because this is the second MLX AOT Metal kernel and it will set the pattern for the MLX AOT path. The current implementation is page-size-1 flat-KV / token-indexed decode attention over MlxAttentionKVPool, but it is exposed as paged attention. Please fix the mismatch in this PR: either implement the block-table paged-attention shape, or rename/gate the feature so it is clearly not claiming vLLM-Metal-style paged attention.
The concrete issues to resolve in this PR are: token-level kv_indices metadata instead of block tables, one 32-lane SIMD group per (batch, q_head) instead of a 256-thread block-table kernel, two full passes over visible KV instead of an online-softmax/split-KV path, head_dim <= 128, and no quantized-KV cache handling. If any of these remain unsupported, the PR should include explicit runtime guards plus tests and current perf/profiling data proving the supported path is the intended shipped scope.
I used codex to find the differences between your implementation and vllm-metal's. Please see the following comments.
| seq_lens: list[int] | ||
| max_seq_len: int | ||
| total_kv_tokens: int | ||
| page_size: int = 1 |
There was a problem hiding this comment.
Please fix the mismatch between the name and the implementation. This page_size=1 metadata path is flat-token / token-indexed decode, not the same shape as vLLM-Metal's block-table paged cache where the cache is [num_blocks, block_size, num_kv_heads, head_dim] and the attention kernel consumes block_tables + seq_lens. Either change the metadata/kernel to block-table paged attention in this PR, or rename the feature/API so it does not advertise full paged attention.
| visible_len = before_len + 1 if include_current_token else before_len | ||
| visible_lens.append(visible_len) | ||
|
|
||
| for pos in range(visible_len): |
There was a problem hiding this comment.
This builds kv_indices by expanding every visible token slot on every decode step, so metadata size/work is O(total visible tokens). vLLM-Metal passes block tables plus sequence lengths and lets the kernel walk blocks, which keeps metadata O(num blocks).
Please fix this in the PR rather than leaving it as design debt: build block-table-style metadata (block_tables, seq_lens, block_size) or rename the path away from paged attention and add tests/perf data for the flat-token metadata cost. I do not think we should merge an AOT attention kernel where O(total visible tokens) metadata construction is hidden behind the "paged" name.
| @@ -0,0 +1,156 @@ | |||
| // SGLang Apple Silicon Metal kernel: page-size-1 decode attention over a flat KV pool. | |||
There was a problem hiding this comment.
The kernel header says "page-size-1 decode attention over a flat KV pool", but the file and public API are named paged_attention. Please resolve this in the PR: either implement block-table paged attention or rename the kernel/API/env-facing description to flat-token decode attention. The distinction matters because this implementation consumes token-level kv_indices, while vLLM-Metal consumes block_tables, seq_lens, and block_size.
vLLM-Metal host dispatch for comparison: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal/paged_ops.cpp#L149-L255
| const float row_max = simd_max_32(local_max); | ||
|
|
||
| float local_sum = 0.0f; | ||
| float local_acc[PAGED_ATTN_MAX_HEAD_DIM]; |
There was a problem hiding this comment.
local_acc[128] is per thread. For HEAD_DIM=128 and 32 lanes, this can put a lot of pressure on registers/local memory and may explain why the median decode speedup is modest in my local run. vLLM-Metal uses 256-thread groups with explicit threadgroup memory sizing/merge buffers for the paged path.
Please fix this with kernel-level evidence or code changes in this PR. Either provide Metal capture/profiling data showing this does not spill heavily for Qwen3 head_dim=128, or restructure the kernel to avoid the per-thread 128-float accumulator. For a new AOT kernel, I do not think we should merge a register-heavy design without that evidence.
| local_acc[d] = 0.0f; | ||
| } | ||
|
|
||
| for (int32_t i = start + int32_t(lane); i < end; i += int32_t(PAGED_ATTN_TG_SIZE)) { |
There was a problem hiding this comment.
This second pass recomputes the QK logits instead of retaining partial logits or using an online softmax formulation. That keeps the kernel compact, but it doubles Q/K reads and dot-product work. Please fix this in the PR with either an online-softmax/block-table implementation or profiling/perf evidence that the recomputation is acceptable for the shipped path.
vLLM-Metal's v2 dispatch is the closest comparison point: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal/paged_ops.cpp#L261-L363
| tid, tg_pos); \ | ||
| } | ||
|
|
||
| INSTANTIATE(f16, half) |
There was a problem hiding this comment.
The current instantiations only cover f16/bf16/f32 KV cache tensors. That is fine for Qwen3-0.6B, but vLLM-Metal also has uchar-cache instantiations for quantized KV. Please fix the shipped capability boundary in this PR: either add quantized-KV handling or explicitly gate/raise unsupported quantized KV with tests.
| enc.set_output_array(out, 5); | ||
|
|
||
| const uint32_t batch = static_cast<uint32_t>(q.shape(0)); | ||
| enc.dispatch_threads(MTL::Size::Make(batch, nq, 32), MTL::Size::Make(1, 1, 32)); |
There was a problem hiding this comment.
This dispatch shape is the host-side counterpart of the 32-lane design in the .metal file. Please either change it to a block-table/256-thread style dispatch comparable to vLLM-Metal, or include profiling/perf evidence in this PR that one SIMD group per (batch, q_head) is sufficient for the target workload. vLLM-Metal uses dispatch_threadgroups(num_heads, num_seqs, 1) with NUM_THREADS=256.
| and ctx.paged_metadata is not None | ||
| ) | ||
|
|
||
| if use_paged_attention: |
There was a problem hiding this comment.
In the paged path, _rope_custom_aot() already scatters K/V into the shared KV pool through new_token_slots, then this loop still writes each token into the per-request layer_caches. Please fix this in the PR: if it is duplicate work, remove it; if it is required to advance offsets / keep fallback cache state consistent, make that explicit in code and include a focused test so this does not look like accidental duplicate KV writes on the hot decode path.
Motivation
This PR adds an experimental MLX PagedAttention decode path for Apple Silicon.
The current MLX batched decode path pads each request’s K/V cache to the longest sequence in the batch and then calls dense
mx.fast.scaled_dot_product_attention. That is correct, but it does extra work for skewed batches. This PR adds a page-size=1 Metal decode kernel over the existing flatMlxAttentionKVPoolso supported decode batches can read only the visible KV slots.This builds on the MLX KV-pool / AOT Metal scaffolding from #22868 and is part of the Apple Device Support roadmap (#19137).
It does not enable a new end to end attention kernel yet.
Part of the Apple Device Support roadmap (#19137).
Modifications
MLXPagedAttentionMetadatafor decode.req_to_tokenstate:kv_indptrkv_indicesqo_indptrslot_mappingsgl_kernel.metal.paged_attention_decode(...).SGLANG_MLX_USE_PAGED_ATTENTION=1.bench_one_batch --skewed-input-lensfor skewed decode benchmarking. ( Note : This was added for local paged-attention decode benchmarking by me and Can be ignored for this PR, In case reviewers think it is not needed)How the metadata is used
The decode metadata is built from
req_to_tokenand attached toBatchedDecodeContext.When
SGLANG_MLX_USE_PAGED_ATTENTION=1, the MLX attention wrapper passeskv_indptrandkv_indicesto the AOT Metal decode kernel. The kernel reads K/V directly from the flatMlxAttentionKVPoolinstead of padding per-request K/V to the batch max length.Paged attention is env-gated and falls back to the existing padded SDPA path for unsupported cases.
Future work
This PR is intentionally limited to env-gated decode over the current flat KV-pool layout. Follow-up work:
qo_indptr, prefix K/V, new-token K/V, and causal masking.Accuracy Tests
N/A
Profiling of existing stack
Baseline, MLX padded SDPA:
Baseline, MLX padded SDPA:
On a skewed MLX bench_one_batch case [128, 512, 1024, 2048], env-gated paged decode improves median decode latency by ~12% and overall throughput by ~4% versus the existing padded SDPA path.
How to enable
SGLANG_MLX_USE_PAGED_ATTENTION=1automatically enables the fused RoPE/KV-pool writer required by the paged decode path. Users do not need to setSGLANG_MLX_USE_CUSTOM_ROPE=1separately.Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #28810467352
Latest PR Test (Extra): ❌ Run #28810467239