Skip to content

fix: MLA and sliding window KV cache accuracy#10

Open
stmatengss wants to merge 1 commit into
mainfrom
feat/mla-accuracy
Open

fix: MLA and sliding window KV cache accuracy#10
stmatengss wants to merge 1 commit into
mainfrom
feat/mla-accuracy

Conversation

@stmatengss

Copy link
Copy Markdown
Collaborator

Summary

  • Fix 15x overestimate for MLA models (DeepSeek V2/V3/V3.1, Kimi K2): Replace the fake n_kv_heads=16, head_dim=512 hack with proper MLA compressed KV parameters (mla: true, kv_lora_rank: 512, qk_rope_head_dim: 64). The kvBytesPerToken() function now uses the correct formula: (kv_lora_rank + qk_rope_head_dim + kv_lora_rank) * dtype_bytes per layer = 1088 values instead of 16384 (a ~15x correction).
  • Add sliding window metadata to Gemma 3 models (27B, 12B, 4B): Add sliding_window: 1024, sliding_window_pattern: 6 params to expose hybrid attention properties in the model profile.
  • Update DeepSeek V3 deployment preset to use corrected MLA parameters.

Technical details

MLA (Multi-head Latent Attention) compresses KV cache using low-rank projection:

  • K cache per layer = (kv_lora_rank + qk_rope_head_dim) * dtype_bytes = (512 + 64) * 2 = 1152 bytes
  • V cache per layer = kv_lora_rank * dtype_bytes = 512 * 2 = 1024 bytes
  • Total per layer = 2176 bytes vs old formula's 2 * 16 * 512 * 2 = 32768 bytes

Test plan

  • All 178 JS parity validation checks pass (node tests/validate_js_parity.mjs)
  • Verify DeepSeek V3 671B KV cache estimate in the web calculator shows ~0.12 GB/1K tokens instead of ~1.9 GB/1K tokens
  • Verify Gemma 3 models display sliding window metadata in model notes

Fix 15x overestimate for MLA models (DeepSeek V2/V3/V3.1, Kimi K2) by
replacing the fake n_kv_heads=16/head_dim=512 hack with proper MLA
compressed KV formula: (kv_lora_rank + qk_rope_head_dim) for K +
kv_lora_rank for V = 1088 values per layer instead of 16384.

Also add sliding_window and sliding_window_pattern params to Gemma 3
models (27B, 12B, 4B) to expose hybrid attention metadata.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant