[CPU][MLA] Fix prefill backend selection so MLA runs end-to-end on CPU

[maobaolong]

Purpose

Make DeepSeek-V2-Lite run end-to-end on CPU when MLA is enabled, including the
contextful prefill path needed by external KV reloads (for example LMCache).
This is the missing follow-up to #49453: that PR brought up the CPU MLA decode
path and the basic CPU MLA backend plumbing, but it still left CPU MLA unable
to run real external-hit prefill flows.

Why #49453 was insufficient

#49453 added the following pieces:

• CPUMLABackend / CPUMLAImpl under vllm/v1/attention/backends/mla/
• CPU MLA decode via torch.ops._C.mla_decode_kvcache
• CPU MLA KV-cache writes via the PyTorch fallback path
• CPU platform plumbing (block_size=16, disable chunked prefill / prefix
  caching for MLA)
• ARM correctness fixes in the CPU C++ code

That was enough to make the CPU MLA decode path exist, but not enough to make a
DeepSeek MLA model actually run through the full prefill lifecycle.

The remaining gaps were:

1. CPU MLA prefill was still mis-routed to a GPU-oriented backend instead of a
   CPU-native one.
2. The CPU MLA prefill backend only handled the "new tokens only" path. It did
   not implement:
   • return_softmax_lse=True for suffix prefill
   • run_prefill_context_chunk(...) for prefill rows that already have context
3. The shared MLA layer still assumed the custom cache gather op
   _C_cache_ops.gather_and_maybe_dequant_cache existed, but that op is not
   available on CPU.
4. After fixing the above, the context/suffix merge still fell through to the
   Triton merge kernel, which is not usable in this CPU environment.
5. The UT added in [CPU] Add MLA backend so DeepSeek-V2/V3 can run on CPU #49453 did not exercise this failure mode. It covered
   backend selection and the cold/no-context CPU MLA path, but it never created
   a prefill row with num_computed_tokens > 0, so it never executed the
   LMCache-style external-hit/contextful prefill path.

Because LMCache / KV connectors create prefill rows with
num_computed_tokens > 0, they necessarily exercise the contextful MLA prefill
path. That is why #49453 could merge while still failing for
DeepSeek-V2-Lite + CPU + LMCache.

What this PR adds

1. A CPU-native MLA prefill backend,
   CPUSDPAMLAPrefillBackend
   (vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py).
2. CPU support for contextful MLA prefill in that backend:
   • run_prefill_new_tokens(..., return_softmax_lse=True)
   • run_prefill_context_chunk(...)
3. A CPU fallback for MLA context KV gather in
   vllm/model_executor/layers/attention/mla_attention.py, avoiding the
   missing _C_cache_ops.gather_and_maybe_dequant_cache dependency.
4. A pure PyTorch CPU fallback for merge_attn_states(...) in
   vllm/v1/attention/ops/merge_attn_states.py, so prefix/suffix partial
   results can be merged without Triton.
5. Stronger CPU MLA tests, including an end-to-end external-KV-hit smoke test
   using ExampleConnector with dummy weights and token-id prompts, so it
   does not need to load real model weights and does not depend on tokenizer
   initialization.

Why this is not duplicate work

This PR is not duplicating #49453. #49453 brought up the CPU MLA backend and
made the decode path possible, but it left CPU MLA prefill incomplete and did
not cover the external-hit/contextful prefill path at all. This PR closes the
remaining CPU MLA correctness gaps so DeepSeek-V2-Lite can run through cold
prefill, external KV reload, and prefix/suffix merge on CPU.

I also checked current open PRs for overlapping work before proceeding. I did
not find an open PR that fixes this CPU MLA contextful prefill path.

Tests

Added / updated coverage in tests/v1/attention/test_cpu_mla_backend.py:

• test_kv_cache_cpu_write
  Verifies the CPU MLA KV-cache write path still stores MLA latent KV rows in
  the layout expected by the CPU decode kernel.
• test_cpu_mla_prefill_backend_selected
  Verifies CPU MLA prefill selects the CPU backend.
• test_cpu_mla_prefill_new_tokens
  Verifies the CPU SDPA MLA prefill backend matches a reference ragged causal
  attention implementation, including return_softmax_lse=True.
• test_cpu_mla_prefill_context_chunk
  Verifies the CPU SDPA MLA backend correctly handles prefill context chunks
  and returns the expected per-head LSEs.

Verification

1. Targeted pytest coverage

Ran:

/Users/mbl/projects/vllm/.venv/bin/python -m pytest \
  tests/v1/attention/test_cpu_mla_backend.py -v

Result:

5 passed in 17.85s

2. Local shell e2e: CPU TP=2 + LMCache + no prefix cache

This validation uses:

• load_format="dummy" so no real model weights are loaded
• hf_overrides to shrink DeepSeek-V2-Lite to 2 layers / 4 routed experts
• skip_tokenizer_init=True plus token-id prompts so the offline path does
  not depend on tokenizer startup
• tensor_parallel_size=2
• LMCacheMPConnector
• --no-enable-prefix-caching

On this single-NUMA local machine I also had to set
VLLM_CPU_SIM_MULTI_NUMA=1, otherwise CPU worker autobinding refused to launch
two local workers. I also had to set gpu_memory_utilization=0.25, because on
the CPU backend that flag controls the fraction of system memory reserved per
worker.

On my local machine, TCP port 6555 and LMCache's default HTTP port 8080
were already in use, so the verified run below uses 6556 and 18080.

Started LMCache:

source ~/.venv-lmcache/bin/activate
lmcache server \
  --host 127.0.0.1 \
  --port 6556 \
  --chunk-size 256 \
  --l1-size-gb 4 \
  --eviction-policy LRU \
  --max-workers 2 \
  --http-port 18080

Started vLLM:

source ~/.venv-lmcache/bin/activate
VLLM_HOST_IP=127.0.0.1 \
VLLM_LOOPBACK_IP=127.0.0.1 \
VLLM_CPU_SIM_MULTI_NUMA=1 \
vllm serve deepseek-ai/DeepSeek-V2-Lite \
  --load-format dummy \
  --trust-remote-code \
  --no-enable-prefix-caching \
  --max-model-len 1024 \
  --gpu-memory-utilization 0.25 \
  --block-size 64 \
  --tensor-parallel-size 2 \
  --port 8024 \
  --hf-overrides '{"num_hidden_layers":2,"n_routed_experts":4,"first_k_dense_replace":0,"num_experts_per_tok":2}' \
  --kv-transfer-config '{"kv_connector":"LMCacheMPConnector","kv_connector_module_path":"lmcache.integration.vllm.lmcache_mp_connector","kv_role":"kv_both","kv_connector_extra_config":{"lmcache.mp.host":"127.0.0.1","lmcache.mp.port":6556,"lmcache.mp.mp_transfer_mode":"lmcache_driven"}}'

Prepared a prompt with 300 token so LMCache has at least one full
256-token chunk to store and later reload:

Sent the first request to warm/store KV:

curl -s http://127.0.0.1:8024/v1/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "deepseek-ai/DeepSeek-V2-Lite",
"prompt": "Attention mechanisms are the foundational building block of the Transformer architecture, and understanding them requires tracing their motivation back to the limitations of recurrent and convolutional models. Before Transformers, sequence-to-sequence tasks were dominated by recurrent neural networks that processed tokens one step at a time, which made parallelization across sequence length impossible and caused information from early tokens to degrade as it was carried through many recurrent transitions. Convolutional networks offered some parallelization but still struggled to relate positions that were far apart unless the network was made very deep. The self-attention layer solves this by letting every token attend to every other token directly, so the path length between any two positions is always exactly one step regardless of how long the sequence is. Concretely, each input token is mapped to a query, a key, and a value vector through learned linear projections, and the output for a given token is a weighted sum of the value vectors of all tokens, where the weights come from the dot product between its query and the keys of others, scaled by the square root of the key dimension and passed through a softmax. Multi-head attention extends this idea by running several attention operations in parallel with different projected subspaces, allowing the model to attend to different kinds of relationships simultaneously, such as syntax at one head and coreference at another. Positional encodings are added to the token embeddings because attention itself is permutation invariant and would otherwise be blind to word order, and the original paper used fixed sinusoidal functions while later variants learned the positions directly. The encoder stack applies self-attention followed by a feed-forward network with residual connections and layer normalization, while the decoder inserts an additional masked attention sub-layer so that predictions for a position can only depend on earlier positions and never on future ones. Training such models at scale revealed that depth, data, and compute matter enormously, and that careful learning rate schedules with warmup prevent the early instability that raw Adam often shows on these architectures. Beyond machine translation where the Transformer first proved itself, attention became the substrate for large language models, vision transformers that treat image patches like tokens, and multimodal systems that align text and pixels in a shared space. Researchers have since studied the inductive biases of attention, its sample efficiency relative to convolutions, and the ways in which it can fail on rare patterns or long-range dependencies that exceed the effective context actually used during training. Efficient variants such as sparse attention, linear attention, and low-rank approximations try to reduce the quadratic cost of attending over the full sequence so that models can handle documents with tens of thousands of tokens without running out of memory. In practice, the choice of attention variant is a tradeoff between quality, speed, and the hardware available, and production systems frequently combine multiple tricks such as fused kernels,Flash style exact attention, and KV caching so that decoding new tokens reuses previously computed key and value vectors instead of recomputing them. This caching is exactly why serving engines care so deeply about memory layout and transfer, because the KV cache grows linearly with sequence length and batch size and quickly becomes the dominant memory consumer during long conversations. When you reason about deploying these models, you should think about how the prompt length, the number of concurrent requests, and the chosen block size all interact to determine how many blocks of KV cache fit in a given device, and how connectors that move that cache between processes or machines can smooth out load spikes. The deeper point is that attention is not just a mathematical operation but an architectural choice with system level consequences, and the cleanest mental model is to always ask what each token is allowed to look at, how that view is paid for in compute and memory, and where the resulting state is stored and moved as the sequence grows.",
"max_tokens": 8,
"temperature": 0
}'

Sent the second identical request to exercise the external KV hit path:

Observed results:

• both requests returned HTTP 200 with successful completion payloads
• the API server logged:

External prefix cache hit rate: 42.7%

• the LMCache server logged:

Stored 256 tokens in 0.003 seconds
Retrieved 256 tokens in 0.005 seconds
Retrieved 256 tokens in 0.005 seconds

This second identical prompt is the important case: it confirms the external KV
reload path no longer crashes on CPU, and that DeepSeek-V2-Lite + CPU MLA + TP=2 + LMCache + no prefix cache runs end-to-end.

• LMCache server also logged:

[2026-08-10 21:51:07,963] LMCache INFO: Engine KV Format: 3 NL x [NB, BS, HS] (detection.py:69:lmcache.v1.gpu_connector.kv_format.detection)
[2026-08-10 21:51:07,963] LMCache INFO: Engine KV Format: 3 NL x [NB, BS, HS] (detection.py:69:lmcache.v1.gpu_connector.kv_format.detection)
[2026-08-10 21:51:07,963] LMCache INFO: Group 0 first-layer tensor: layer_idx=0 shape=(409149, 16, 576) stride=(9216, 576, 1) is_contiguous=True dtype=torch.bfloat16 device=cpu storage_offset=0 numel=3770717184 storage_nbytes=7541434368 padding_per_block=0 (utils.py:596:lmcache.v1.gpu_connector.utils)
[2026-08-10 21:51:07,963] LMCache INFO: KV layer groups: ---
KernelGroupInfo(layers=2, indices=0-1, shape_desc=(kv=1, nl=2, nb=409149, bs=16, nh=1, hs=576, element_size=2, block_stride_elems=9216), dtype=torch.bfloat16, tokens_per_block=16, slots_per_block=16, engine_group_idx=0, sw_size_tokens=-1)
--- (kv_layer_groups.py:437:lmcache.v1.kv_layer_groups)
[2026-08-10 21:51:07,963] LMCache INFO: CPUCacheContext: 2 layers, 409149 blocks, dtype=torch.bfloat16 (shm-backed) (cache_context.py:186:lmcache.v1.platform.cpu.cache_context)
[2026-08-10 21:51:07,964] LMCache INFO: Registered KV cache for GPU ID 1336549345371966953 with 2 layers (lmcache_driven_transfer.py:963:lmcache.v1.multiprocess.modules.lmcache_driven_transfer)
[2026-08-10 21:51:07,968] LMCache INFO: Engine KV Format: 3 NL x [NB, BS, HS] (detection.py:69:lmcache.v1.gpu_connector.kv_format.detection)
[2026-08-10 21:51:07,968] LMCache INFO: Engine KV Format: 3 NL x [NB, BS, HS] (detection.py:69:lmcache.v1.gpu_connector.kv_format.detection)
[2026-08-10 21:51:07,968] LMCache INFO: Group 0 first-layer tensor: layer_idx=0 shape=(409149, 16, 576) stride=(9216, 576, 1) is_contiguous=True dtype=torch.bfloat16 device=cpu storage_offset=0 numel=3770717184 storage_nbytes=7541434368 padding_per_block=0 (utils.py:596:lmcache.v1.gpu_connector.utils)
[2026-08-10 21:51:07,969] LMCache INFO: KV layer groups: ---
KernelGroupInfo(layers=2, indices=0-1, shape_desc=(kv=1, nl=2, nb=409149, bs=16, nh=1, hs=576, element_size=2, block_stride_elems=9216), dtype=torch.bfloat16, tokens_per_block=16, slots_per_block=16, engine_group_idx=0, sw_size_tokens=-1)
--- (kv_layer_groups.py:437:lmcache.v1.kv_layer_groups)
[2026-08-10 21:51:07,969] LMCache INFO: CPUCacheContext: 2 layers, 409149 blocks, dtype=torch.bfloat16 (shm-backed) (cache_context.py:186:lmcache.v1.platform.cpu.cache_context)
[2026-08-10 21:51:07,969] LMCache INFO: Registered KV cache for GPU ID 554121685577093483 with 2 layers (lmcache_driven_transfer.py:963:lmcache.v1.multiprocess.modules.lmcache_driven_transfer)
[2026-08-10 21:51:55,123] LMCache INFO: AffinityThreadPool: affinity_key=6258934065506746361 assigned to worker slot 0 of 2 (thread affinity-pool-0-0); 1 distinct key(s) now bound (affinity_pool.py:108:lmcache.v1.multiprocess.affinity_pool)
[2026-08-10 21:51:55,125] LMCache INFO: Stored 512 tokens in 0.002 seconds (lmcache_driven_transfer.py:1217:lmcache.v1.multiprocess.modules.lmcache_driven_transfer)
[2026-08-10 21:51:57,594] LMCache INFO: Prefetch request completed (L1+L2): 2/2 retained keys (2 L1, 0 L2) in 0.5 ms (external_request_id=cmpl-8906388402e945f5-0-82c8061c, prefetch_request_id=-1) (storage_manager.py:716:lmcache.v1.distributed.storage_manager)
[2026-08-10 21:51:57,596] LMCache INFO: AffinityThreadPool: affinity_key=6362481149236448726 assigned to worker slot 1 of 2 (thread affinity-pool-0-1); 2 distinct key(s) now bound (affinity_pool.py:108:lmcache.v1.multiprocess.affinity_pool)
[2026-08-10 21:51:57,597] LMCache INFO: Retrieved 512 tokens in 0.001 seconds (lmcache_driven_transfer.py:1426:lmcache.v1.multiprocess.modules.lmcache_driven_transfer)
[2026-08-10 21:51:57,597] LMCache INFO: Retrieved 512 tokens in 0.001 seconds (lmcache_driven_transfer.py:1426:lmcache.v1.multiprocess.modules.lmcache_driven_transfer)

Model evaluation

No model eval was run.

Reason:

• this change fixes a previously crashing CPU MLA serving/runtime path
• validation here is correctness / runtime-plumbing oriented, not output-quality
  oriented
• all local end-to-end checks were run with load_format="dummy" specifically
  to avoid loading real weights during development

Notes

• Performance is intentionally not the goal here. The new CPU MLA prefill
  backend and fallbacks are correctness-first implementations.
• This PR does not touch sparse DSA / CPUSparseMLAImpl; that is separate work.
• This change was developed with AI assistance.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[maobaolong]

@bigPYJ1151 Hi Thanks for the review and help on #49453, would you like to take a look at this PR also?

After this PR, we can run vllm + deepseek-v2-lite end2end test.

[MatthewBonanni]

Seems reasonable to me, but please add a GSM8k eval result as well as vllm bench serve results in comparison to the pure decode pathway

[maobaolong]

Thanks @MatthewBonanni for your review! I now have local validation with real DeepSeek-V2-Lite weights on CPU for both paths:

• pure vLLM CPU TP=2: GSM8K 3-question smoke eval Accuracy: 0.667, Invalid responses: 0.000
• vLLM CPU TP=2 + LMCache: GSM8K 3-question smoke eval Accuracy: 0.667, Invalid responses: 0.000

Both runs used --no-enable-prefix-caching. For the pure vLLM run, the server log also confirms that it is still using the CPU MLA path (Using CPU SDPA MLA prefill backend.).

PTAL 🙏

[maobaolong]

@MatthewBonanni I'm so sorry to ping @MatthewBonanni

[maobaolong]

Hi @bigPYJ1151 , would you like to help to take another look at this PR? Thanks a lot!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @maobaolong.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[maobaolong]

Sorry to bother you, but could you please keep an eye on this PR for me? Thank you so much. @MatthewBonanni @bigPYJ1151

[Isotr0py]

Overall look reasonable, have some minor comments. PTAL!

[sunlei1992]

Apple Silicon (M3) validation — CPU SDPA MLA prefill backend works

Thanks for this fix! I ported the PR's changes (cpu_sdpa.py, registry.py CPU entry, selector.py CPU branch, _custom_ops.py gather_mla_context_cache_cpu, and the mla_attention.py CPU gather fallback) onto a v0.28.0-based tree and validated end-to-end on Apple Silicon, which is a platform the current CI/reference backend path doesn't cover.

Environment: Mac M3, macOS 15.7.4, 24 GB, CPU backend, dummy weights.

Test 1 — official smoke scenario (mirrors test_cpu_mla_backend_smoke: DeepSeek-V2-Lite + hf_overrides shrink to 2 layers / 4 experts / all-MoE):

Using CPU SDPA MLA prefill backend.      ← selector now picks the CPU backend
Using CPU Unquantized MoE backend
PASSED: both requests generated 4 tokens

Test 2 — GLM-5.2 (glm_moe_dsa) truncated to 3/78 layers:

Using CPU SDPA MLA prefill backend.
Avg latency: 2.51s per batch (batch=4, 16+16 tokens)

Both pass with only the PR's changes applied (no other patches). The log confirms the selector routes to CPU_SDPA_MLA and prefill runs through the new backend.

Notes / suggestions:

1. Out of scope, not covered by this PR: sparse-attention models on CPU still fail with NotImplementedError: Sparse Attention is not supported on CPU (platforms/cpu.py:89), because e.g. GlmMoeDsaConfig hard-codes index_topk=2048 so DeepseekV2.is_v32 (hasattr(config, "index_topk")) is always True. I had to disable that locally for the GLM-5.2 test. Worth a follow-up issue/PR (config-flag or clear error path).
2. I validated the fresh-prefill path (no num_computed_tokens > 0). The contextful/LMCache path (run_prefill_context_chunk, gather_mla_context_cache_cpu, merge) was ported but not exercised here — would be good to see a CPU CI test for that too.
3. Minor: the is_available()/supports_* classmethods on CPUSDPAMLAPrefillBackend are inherited defaults — fine since the selector special-cases CPU, but supports_mla_dimensions could be declared to fail fast on unsupported dims like the FlashAttn backend does.

Happy to run any additional scenario you'd like checked on Apple Silicon.

[maobaolong]

@sunlei1992 Thanks a lot for the Apple Silicon validation! This is very helpful.

The sparse-attention / GLM-5.2 issue looks out of scope for this PR, but it can be a followup work. The main focus here is the CPU SDPA MLA prefill path for DeepSeek-style MLA models, and I have separately validated the repeated request for LMCache path with DeepSeek-V2-Lite.

Really appreciate the detailed check!

[maobaolong]

@bigPYJ1151 Thanks for the review and suggested changes, addressed all the comments i guess, would you like to take another look? Thanks!