fix(attention): read per-runner kv cache dtype off model_runner - #32251
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![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: 6597fcfc65
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| from sglang.srt.runtime_context import get_model | ||
|
|
||
| self.kv_cache_dtype_str = get_model().kv_cache_dtype | ||
| self.kv_cache_dtype_str = model_runner.kv_cache_dtype_str |
There was a problem hiding this comment.
When GPU attention tests construct this backend with their mock runners, this direct access raises AttributeError: only test/manual/attention/test_flashattn_backend.py was updated, while test_flashattn_mla_backend.MockModelRunner and the dense-attention MockModelRunner still omit kv_cache_dtype_str. The same omission exists in MockLightningModelRunner and DualChunkMockModelRunner for the other direct accesses introduced here, so the corresponding attention-unittest fixtures also fail during backend construction rather than exercising kernels. Please update every in-repo runner fixture that instantiates these backends, or retain a compatible per-runner fallback.
Useful? React with 👍 / 👎.
|
/tag-and-rerun-ci |
The FlashAttention and XPU attention backends resolved `kv_cache_dtype_str` from the process-global published model config via `get_model().kv_cache_dtype`. That global reflects only the target runner: a draft (speculative) worker configures its own kv cache dtype but does not publish it, so a draft backend would inherit the target's dtype string and mis-select its fp8 code paths whenever the draft's kv cache dtype differs from the target's. Read the dtype string straight from `model_runner.kv_cache_dtype_str` instead, which is set per runner during kv-cache configuration before any attention backend is initialized, so each backend — target or draft — uses its own resolved dtype. Also drop the `getattr(..., server_args.kv_cache_dtype)` fallback in the lightning and dual-chunk backends: `kv_cache_dtype_str` is always set on a real runner before backend init, so the fallback only masked setup bugs. Because the backends now read the attribute directly, the attention unit-test mock model runners (the shared attention_unittest kit and the manual FlashAttention backend tests) set `kv_cache_dtype_str` to match a real runner. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
b86ddc7 to
987d44a
Compare
1 participant
Summary
The FlashAttention and XPU attention backends resolved
kv_cache_dtype_strfrom theprocess-global published model config (
get_model().kv_cache_dtype). That globalreflects only the target runner: a draft (speculative) worker configures its own kv cache
dtype but does not publish it, so a draft backend would inherit the target's dtype string
and mis-select its fp8 code paths whenever the draft's kv cache dtype differs from the
target's.
Read the dtype string straight from
model_runner.kv_cache_dtype_strinstead, which isset per runner during kv-cache configuration before any attention backend is initialized,
so each backend — target or draft — uses its own resolved dtype.
Also drop the
getattr(..., server_args.kv_cache_dtype)fallback in the lightning anddual-chunk backends:
kv_cache_dtype_stris always set on a real runner before backendinit, so the fallback only masked setup bugs.
Test
test/manual/attention/test_flashattn_backend.pymock updated to set the attributedirectly; the non-CP tests pass.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ✅ Run #30054771755
Latest PR Test (Extra): ❌ Run #30054771718