Skip to content

[Apple Silicon] Add (env-gated) MLX paged attention decode - #28335

Closed
adityavaid wants to merge 5 commits into
sgl-project:mainfrom
adityavaid:adityavaid/mlx-paged-attention-metadata
Closed

adityavaid wants to merge 5 commits into
sgl-project:mainfrom
adityavaid:adityavaid/mlx-paged-attention-metadata

Conversation

@adityavaid

@adityavaid adityavaid commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

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 flat MlxAttentionKVPool so 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

  • Add MLXPagedAttentionMetadata for decode.
  • Build page-size=1 metadata from existing req_to_token state:
    • kv_indptr
    • kv_indices
    • qo_indptr
    • slot_mapping
  • Add AOT Metal page-size=1 paged decode kernel.
  • Add sgl_kernel.metal.paged_attention_decode(...).
  • Wire the paged decode path behind SGLANG_MLX_USE_PAGED_ATTENTION=1.
  • Add bench_one_batch --skewed-input-lens for 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)
  • Update Apple Metal docs for the experimental paged-attention flag.

How the metadata is used

The decode metadata is built from req_to_token and attached to BatchedDecodeContext.

When SGLANG_MLX_USE_PAGED_ATTENTION=1, the MLX attention wrapper passes kv_indptr and kv_indices to the AOT Metal decode kernel. The kernel reads K/V directly from the flat MlxAttentionKVPool instead 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:

  • Add split-KV / partitioned long-context decode for very long contexts.
  • Add prefill/extend paged attention with qo_indptr, prefix K/V, new-token K/V, and causal masking.
  • Evaluate moving from page-size=1 flat slots to a vLLM-style block-table layout.
  • Decouple paged attention from fused RoPE by adding a standalone MLX-visible K/V pool scatter primitive, If needed.
  • Add cache movement operations such as block copy, gather, and future compaction.
  • Add support or explicit fallback coverage for sliding window, MLA, speculative decode, and quantized KV.

Accuracy Tests

N/A

Profiling of existing stack

Baseline, MLX padded SDPA:

SGLANG_USE_MLX=1 \
SGLANG_MLX_USE_PAGED_ATTENTION=0 \
.venv-mlx-dev/bin/python -m sglang.bench_one_batch \
  --model-path Qwen/Qwen3-0.6B \
  --trust-remote-code \
  --skewed-input-lens 128 512 1024 2048 \
  --output-len 64 \
  --run-name mlx_skew_padded \
  --result-filename /tmp/mlx_skew_padded.jsonl
  
  Prefill. latency: 0.79446 s, throughput:   4672.36 token/s
Decode 0. Batch size: 4, latency: 0.03999 s, throughput:    100.03 token/s
Decode 1. Batch size: 4, latency: 0.03686 s, throughput:    108.52 token/s
Decode 2. Batch size: 4, latency: 0.03758 s, throughput:    106.43 token/s
Decode 3. Batch size: 4, latency: 0.03660 s, throughput:    109.30 token/s
Decode 4. Batch size: 4, latency: 0.03766 s, throughput:    106.22 token/s
Decode.  median latency: 0.03751 s, median throughput:    106.64 token/s
Total. latency:  1.990 s, throughput:   1929.72 token/s
Benchmark ...
Skewed input lengths: [128, 512, 1024, 2048]
Prefill. latency: 0.75475 s, throughput:   4918.21 token/s
Decode 0. Batch size: 4, latency: 0.03722 s, throughput:    107.46 token/s
Decode 1. Batch size: 4, latency: 0.03739 s, throughput:    106.99 token/s
Decode 2. Batch size: 4, latency: 0.03697 s, throughput:    108.19 token/s
Decode 3. Batch size: 4, latency: 0.03693 s, throughput:    108.30 token/s
Decode 4. Batch size: 4, latency: 0.03687 s, throughput:    108.49 token/s
Decode.  median latency: 0.03731 s, median throughput:    107.22 token/s
Total. latency:  3.176 s, throughput:   1249.27 token/s
  

Baseline, MLX padded SDPA:

SGLANG_USE_MLX=1 \
SGLANG_MLX_USE_PAGED_ATTENTION=1 \
.venv-mlx-dev/bin/python -m sglang.bench_one_batch \
  --model-path Qwen/Qwen3-0.6B \
  --trust-remote-code \
  --skewed-input-lens 128 512 1024 2048 \
  --output-len 64 \
  --run-name mlx_skew_padded \
  --result-filename /tmp/mlx_skew_padded.jsonl
  
  Prefill. latency: 0.79935 s, throughput:   4643.80 token/s
Decode 0. Batch size: 4, latency: 1.81194 s, throughput:      2.21 token/s
Decode 1. Batch size: 4, latency: 0.04332 s, throughput:     92.34 token/s
Decode 2. Batch size: 4, latency: 0.03756 s, throughput:    106.50 token/s
Decode 3. Batch size: 4, latency: 0.03594 s, throughput:    111.28 token/s
Decode 4. Batch size: 4, latency: 0.03455 s, throughput:    115.77 token/s
Decode.  median latency: 0.03416 s, median throughput:    117.09 token/s
Total. latency:  3.637 s, throughput:   1055.92 token/s
Benchmark ...
Skewed input lengths: [128, 512, 1024, 2048]
Prefill. latency: 0.75851 s, throughput:   4893.82 token/s
Decode 0. Batch size: 4, latency: 0.16601 s, throughput:     24.10 token/s
Decode 1. Batch size: 4, latency: 0.03338 s, throughput:    119.82 token/s
Decode 2. Batch size: 4, latency: 0.03299 s, throughput:    121.26 token/s
Decode 3. Batch size: 4, latency: 0.03319 s, throughput:    120.50 token/s
Decode 4. Batch size: 4, latency: 0.03327 s, throughput:    120.22 token/s
Decode.  median latency: 0.03340 s, median throughput:    119.76 token/s
Total. latency:  3.000 s, throughput:   1322.60 token/s

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_USE_MLX=1 \
SGLANG_MLX_USE_PAGED_ATTENTION=1 \
python -m sglang.launch_server \
  --model-path <MODEL_ID_OR_PATH>

SGLANG_MLX_USE_PAGED_ATTENTION=1 automatically enables the fused RoPE/KV-pool writer required by the paged decode path. Users do not need to set SGLANG_MLX_USE_CUSTOM_ROPE=1 separately.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #28810467352
Latest PR Test (Extra): ❌ Run #28810467239

@adityavaid
adityavaid requested a review from yeahdongcn as a code owner June 15, 2026 19:40
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@changminbark changminbark left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@adityavaid
adityavaid force-pushed the adityavaid/mlx-paged-attention-metadata branch from 4f8f7ef to 262b4cf Compare June 23, 2026 16:56
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Jun 23, 2026
@adityavaid adityavaid closed this Jun 23, 2026
@adityavaid
adityavaid force-pushed the adityavaid/mlx-paged-attention-metadata branch from 262b4cf to ed26a10 Compare June 23, 2026 17:00
@adityavaid adityavaid reopened this Jun 23, 2026
@adityavaid adityavaid changed the title [Apple Silicon] Add MLX paged attention metadata contract [Apple Silicon] Add (env-gated) MLX paged attention decode Jun 24, 2026
Comment thread sgl-kernel/csrc/metal/metal_common.cpp Outdated
out.set_data(allocator::malloc(out.nbytes()));

auto& d = metal::device(stream().device);
auto& enc = d.get_command_encoder(stream().index);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know if this is the metal version diff or something else. Mine needs to update to metal::get_command_encoder(stream());.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done. Thanks, this was an issue with my MLX version, updated now.

Comment thread python/sglang/srt/hardware_backend/mlx/tp_worker.py Outdated
Comment thread python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
seq_lens: list[int]
max_seq_len: int
total_kv_tokens: int
page_size: int = 1

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this 1? I remembered that vllm-metal uses 16 (if I recall correctly). Is this because of the existing flat MlxAttentionKVPool layout?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@adityavaid
adityavaid force-pushed the adityavaid/mlx-paged-attention-metadata branch from a8556b6 to bd9c586 Compare June 24, 2026 15:47

@jlee5814 jlee5814 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@adityavaid

Copy link
Copy Markdown
Contributor Author

Doesn't compile on MLX 0.31.2 . 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()).

Adding metal::get_command_encoder(stream) so both kernels use the MLX 0.31.2-compatible API from one place.

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.

Fixed.

No parity test (test paged vs padded SDPA).

Adding in next commit

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.

Need a bit more clarity on this, I can add explicit correctness check before we build metadata.

model_runner_stub.py:114 revives initialize(self, pre_model_load_memory=None), Please see #28660.
done

@LijuanTang94

Copy link
Copy Markdown
Contributor

Tested on Apple M4 Pro, macOS 26.3.1, MLX 0.31.2, clang 17, at HEAD 91284d2.

Duplicate NB_MODULE is fixed ✅, but get_command_encoder still won't build on 0.31.2 — the helper kept the device-member form:

// 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 (device.h:231) builds clean and matches the CommandEncoder& return type:

return metal::get_command_encoder(stream);

Happy to run the full build + parity test on 0.31.2 once you push.

@adityavaid

Copy link
Copy Markdown
Contributor Author

Tested on Apple M4 Pro, macOS 26.3.1, MLX 0.31.2, clang 17, at HEAD 91284d2.

Duplicate NB_MODULE is fixed ✅, but get_command_encoder still won't build on 0.31.2 — the helper kept the device-member form:

// 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 (device.h:231) builds clean and matches the CommandEncoder& return type:

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 ?

@LijuanTang94

Copy link
Copy Markdown
Contributor

Re-tested on MLX 0.31.2 (M4 Pro): the C++ type-checks clean now ✅ — metal_common.cpp, rope_pool_fused.cpp, paged_attention.cpp, and bindings.cpp all pass -fsyntax-only with no errors. get_command_encoder resolves via the free function as suggested, and there's a single NB_MODULE(_metal) now, so the duplicate-symbol issue is gone. Thanks for the quick turnaround!

Removed unused imports and profiling functions from bench_one_batch.py.

@yeahdongcn yeahdongcn left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Reference: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal_kernel_backend/cache.py#L1-L62

visible_len = before_len + 1 if include_current_token else before_len
visible_lens.append(visible_len)

for pos in range(visible_len):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Reference: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal/paged_ops.cpp#L149-L255

@@ -0,0 +1,156 @@
// SGLang Apple Silicon Metal kernel: page-size-1 decode attention over a flat KV pool.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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];

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Reference: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal/paged_ops.cpp#L190-L218

local_acc[d] = 0.0f;
}

for (int32_t i = start + int32_t(lane); i < end; i += int32_t(PAGED_ATTN_TG_SIZE)) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Reference: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal/kernels_v1/pagedattention.metal#L1379-L1384

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));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Reference: https://github.com/vllm-project/vllm-metal/blob/6c4e4d0cf3135b0853f6dcea704f46f7c2450e23/vllm_metal/metal/paged_ops.cpp#L246-L255

and ctx.paged_metadata is not None
)

if use_paged_attention:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

Labels

apple-silicon documentation Improvements or additions to documentation sgl-kernel

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants