feat(serving): DeepSeek-V2/V3 paged-attention shim (MLA) — fixes #191 - #195
Open
drunkcoding wants to merge 12 commits into
Open
feat(serving): DeepSeek-V2/V3 paged-attention shim (MLA) — fixes #191#195drunkcoding wants to merge 12 commits into
drunkcoding wants to merge 12 commits into
Conversation
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
…e + decode context fix (#191)
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.
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)
DeepseekV2/V3PagedAttentionexisted only as test mocks, souse_paged_context=Falseand the paged KV cache was never consulted during decode.PagedAttentionBackend.forwardwrote KV only on prefill, so each decode step attended to prefill tokens plus uninitialized slots.qk_head_dim=192(nope 128 + rope 64), valuesv_head_dim=128. FlashInfer's generic paged wrappers JIT-stall on non-standard dims and the compiledpaged_attention_v1mis-reads them.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.PagedAttentionBackendis shared by every decoder layer via a ClassVar, but itsk_cache/v_cachehad 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
moe_infinity/models/deepseek_v2_paged_attention.py:DeepseekV2PagedAttention/DeepseekV3PagedAttentionmirroringQwen3PagedAttention. Reproduces HF MLA projection + RoPE, zero-pads the symmetric cache to a kernel-supported head dim (192 → 256), slices the output back tov_head_dim, passes the MLA softmax scale, and threadslayer_idx.runtime/attention_backend.py: per-layer KV cache (leading layer dim onk_cache/v_cache/_fi_kv_cache,layer_idxthreaded throughforward/write_kv/_prefill_forward/_decode_forward); decode-KV-write (write once wheneverslot_mappingpresent, keep the prefill guard); FlashInfer gated to{64,128}(MLA uses SDPA prefill + the precompiled kernel).kernel/paged_attention_ops.py: gate the compiledpaged_attention_v1to head sizes it reads correctly; other sizes use the SDPA reconstruction fallback.serving/batch.py: decodecontext_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 nativeDeepseek{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 passnum_layersto the backend.models/__init__.py,models/qwen3_paged_attention.py: lazy exports; Qwen3 shim also passeslayer_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, greedytemperature=0), matching plain-HFgenerate():Paris.\n\nThe official language is French.\n\nThe currency is theParis. The official language is French. The currency is the Euro…4\n\n## 2. 덧셈 순서…4\n\n## 2. 덧셈 순서…, 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).lspclean on all changed files.Fixes #196 (test/code drift between the FlashInfer head-dim gate and the mocked-module fixtures; resolved in 86e7444 + d35b5b6).