fix: speculative draft worker clobbering target attention backend (#28559)

This commit is contained in:
cctry
2026-06-18 01:21:14 -07:00
committed by GitHub
parent 59001267c3
commit fa71064147
3 changed files with 30 additions and 10 deletions
@@ -2415,6 +2415,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
else:
self.attn_backend = self._get_attention_backend()
# Record resolved per-mode backends on the backend for model dispatch.
self.attn_backend.prefill_attention_backend_str = (
self.prefill_attention_backend_str
)
self.attn_backend.decode_attention_backend_str = (
self.decode_attention_backend_str
)
def _get_attention_backend(self, init_new_workspace: bool = False):
"""Init attention kernel backend."""
draft_attn_backend = self.server_args.speculative_draft_attention_backend
@@ -2422,6 +2430,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
logger.warning(
f"Overriding draft attention backend to {draft_attn_backend}."
)
# Single backend for all draft modes (no prefill/decode split).
self.prefill_attention_backend_str = draft_attn_backend
self.decode_attention_backend_str = draft_attn_backend
return self._get_attention_backend_from_str(
draft_attn_backend,
init_new_workspace=init_new_workspace,
@@ -2463,10 +2474,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
init_new_workspace=init_new_workspace,
)
(
get_global_server_args().prefill_attention_backend,
get_global_server_args().decode_attention_backend,
) = (self.prefill_attention_backend_str, self.decode_attention_backend_str)
return attn_backend
def _get_attention_backend_from_str(