llm: patch 04b v2 - plan-time signal is the spec-config METHOD

use_aux_hidden_state_outputs lives on the gpu_model_runner instance
(gpu_model_runner.py:623/678), not on the spec config, so the first
registry check never fired and the reserve stayed at 13,600 B/token.
Switch to method in {dflash, dspark, extract_hidden_states}.
This commit is contained in:
Frank
2026-09-30 19:29:18 -07:00
parent 9198da375f
commit caee26e3c3
@@ -1,6 +1,6 @@
--- a/vllm/model_executor/offloader/expert_activation_registry.py
+++ b/vllm/model_executor/offloader/expert_activation_registry.py
@@ -316,7 +316,19 @@
@@ -316,7 +316,27 @@
def _glm53_kv_bytes_per_token(cfg: "VllmConfig") -> int:
@@ -13,8 +13,16 @@
+ # cannot absorb, so the reserve must include them. The drafter is a 5-layer
+ # qwen3 (8 kv heads x 128 x K+V x bf16 = 4096 B/token); the aux count is
+ # the dflash target_layer_ids count (5).
+ # NOTE: use_aux_hidden_state_outputs lives on the gpu_model_runner
+ # instance (gpu_model_runner.py:623/678), not on the spec config -- it is
+ # not set yet when the expert plan runs. The spec-config METHOD is the
+ # plan-time signal for the same fact.
+ sc = getattr(cfg, "speculative_config", None)
+ if sc is not None and getattr(sc, "use_aux_hidden_state_outputs", False):
+ if sc is not None and getattr(sc, "method", "") in (
+ "dflash",
+ "dspark",
+ "extract_hidden_states",
+ ):
+ hidden = int(cfg.model_config.hf_config.hidden_size)
+ per_tok += 5 * hidden * 2 + 4096
+ return per_tok