Skip to content

feat(serving): DeepSeek-V2/V3 paged-attention shim (MLA) — fixes #191 - #195

Open
drunkcoding wants to merge 12 commits into
mainfrom
feat/deepseek-mla-paged-attention
Open

feat(serving): DeepSeek-V2/V3 paged-attention shim (MLA) — fixes #191#195
drunkcoding wants to merge 12 commits into
mainfrom
feat/deepseek-mla-paged-attention

Conversation

@drunkcoding

@drunkcoding drunkcoding commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes #191. Serving DeepSeek-V2-Lite-Chat through the OpenAI continuous-batching server produced a correct first token then collapsed into garbage (e.g. "The capital of France is" → " Paris, 100000000000"). This PR implements the real DeepSeek-V2/V3 MLA paged-attention shims and fixes several serving-path bugs uncovered while making decode coherent. After this change, greedy decode is coherent and matches a plain-HF generate() reference.

Root causes (found while driving decode to coherence)

  1. No production DeepSeek paged-attention shim. DeepseekV2/V3PagedAttention existed only as test mocks, so use_paged_context=False and the paged KV cache was never consulted during decode.
  2. Backend never persisted decode-token KV. PagedAttentionBackend.forward wrote KV only on prefill, so each decode step attended to prefill tokens plus uninitialized slots.
  3. MLA head-dim / kernel support. MLA keys are qk_head_dim=192 (nope 128 + rope 64), values v_head_dim=128. FlashInfer's generic paged wrappers JIT-stall on non-standard dims and the compiled paged_attention_v1 mis-reads them.
  4. Off-by-one decode context. The serving batch used context_length = num_computed_tokens, which already counts the just-generated token, so decode used position N+1 for a token at position N (wrong RoPE + a zero KV-cache gap). The native path (generation_loop.py) uses the correct value.
  5. Per-layer KV-cache collision (primary decode bug). A single PagedAttentionBackend is shared by every decoder layer via a ClassVar, but its k_cache/v_cache had no layer dimension — all 27 layers wrote the same slots, so decode read back only the last layer's KV → garbage. Prefill was unaffected because it attends over fresh Q/K/V and never reads the cache. (This latent bug also affected Qwen3.)

Changes

  • New moe_infinity/models/deepseek_v2_paged_attention.py: DeepseekV2PagedAttention / DeepseekV3PagedAttention mirroring Qwen3PagedAttention. Reproduces HF MLA projection + RoPE, zero-pads the symmetric cache to a kernel-supported head dim (192 → 256), slices the output back to v_head_dim, passes the MLA softmax scale, and threads layer_idx.
  • runtime/attention_backend.py: per-layer KV cache (leading layer dim on k_cache/v_cache/_fi_kv_cache, layer_idx threaded through forward/write_kv/_prefill_forward/_decode_forward); decode-KV-write (write once whenever slot_mapping present, keep the prefill guard); FlashInfer gated to {64,128} (MLA uses SDPA prefill + the precompiled kernel).
  • kernel/paged_attention_ops.py: gate the compiled paged_attention_v1 to head sizes it reads correctly; other sizes use the SDPA reconstruction fallback.
  • serving/batch.py: decode context_length = num_computed_tokens - len(input_token) (fixes the off-by-one, aligning with the native path).
  • runtime/model_offload.py: install/restore swap of native Deepseek{V2,V3}Attention → paged shim (mirrors the Qwen3 mechanism).
  • entrypoints/openai/api_server_v2.py, entrypoints/big_modeling.py: MLA KV-cache spec plumbing (head_dim=256, num_kv_heads=num_attention_heads) and pass num_layers to the backend.
  • models/__init__.py, models/qwen3_paged_attention.py: lazy exports; Qwen3 shim also passes layer_idx.
  • tests/python/integration/test_flashinfer_model_attention.py: exercise the real production shims (backend routing, padded head dim, scale, fallback, KV-cache spec).

Verification

Offline parity vs HF eager (float32 + bfloat16): prefill max diff 3.9e-7 / 3.9e-3; multi-layer interleaved prefill→decode (2 layers sharing one backend) max diff ~3e-7, allclose — confirming no cross-layer collision.

Live (api_server_v2, DeepSeek-V2-Lite-Chat, greedy temperature=0), matching plain-HF generate():

Prompt Server output HF reference
The capital of France is Paris.\n\nThe official language is French.\n\nThe currency is the Paris. The official language is French. The currency is the Euro…
2 + 2 = 4\n\n## 2. 덧셈 순서… 4\n\n## 2. 덧셈 순서…
Once upon a time , there was a little girl named Sophie. She lived in a small house with her mother, father, and younger brother , there was a little girl named Sophie. She lived in a small house…

Unit tests: 15 passed, 1 skipped (CUDA-only e2e). lsp clean on all changed files.

Fixes #196 (test/code drift between the FlashInfer head-dim gate and the mocked-module fixtures; resolved in 86e7444 + d35b5b6).

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

Labels

None yet

Projects

None yet

1 participant