Preallocate the static KV cache with config.head_dim - #8389
Conversation
DeepSpeedStaticCache sized its buffers with hidden_size // num_attention_heads, which is wrong for any model that sets head_dim explicitly. The prefill copy in _generate_graph then fails on a shape mismatch. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c6bea0c8c4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| num_heads = getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1)) | ||
| head_dim = getattr(text_config, "hidden_size", 1) // getattr(text_config, "num_attention_heads", 1) | ||
| # head_dim is not always hidden_size // num_attention_heads, so prefer the config value | ||
| head_dim = getattr(text_config, "head_dim", None) or (getattr(text_config, "hidden_size", 1) // |
There was a problem hiding this comment.
This is a non-merge commit, but the commit message for e7c1229c51d6ba3c4896f6681196c8e2a2e11e1c has no Signed-off-by trailer, so it does not satisfy the repository's commit requirements; recreate the commit with --signoff before merging.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The commit on this PR is c6bea0c and its message ends with the Signed-off-by trailer, and the DCO check is green. The sha in the comment above, e7c1229, is not in this repo at all (the API returns 422 for it).
2 participants
DeepSpeedStaticCachepreallocates its KV buffers withhidden_size // num_attention_heads. Models that sethead_dimexplicitly (Qwen3, Gemma, Llama 4) break that identity, so the buffers get the wrong last dimension and the prefill copy in_generate_graphdies before the graph is ever captured:Now it reads
head_dimoff the config and falls back to the division when the attribute is missing or None (older configs set it to None).To reproduce on CPU: a tiny Qwen3 with
hidden_size=64,num_attention_heads=4,head_dim=32, prefilled through the HFStaticCachethe way_generate_graphdoes, then the copy loop athybrid_engine_rollout.py:346verbatim. The prefill cache comes back(1, 2, 8, 32)andDeepSpeedStaticCacheallocates(1, 2, 8, 16). Withhead_dim=16the same script is fine on both sides, which is why nothing has caught this so far.Two tests in
tests/unit/runtime/rollout/. Thehead_dimone fails on master. No GPU on this machine, so the capture and replay path itself is not exercised, only the allocation and the copy that feeds it.