[Fix] missing max_context_len on HybridAttnBackend (#32690)
This commit is contained in:
@@ -42,6 +42,7 @@ class HybridAttnBackend(AttentionBackend):
|
|||||||
self.needs_cpu_seq_lens = (
|
self.needs_cpu_seq_lens = (
|
||||||
prefill_backend.needs_cpu_seq_lens or decode_backend.needs_cpu_seq_lens
|
prefill_backend.needs_cpu_seq_lens or decode_backend.needs_cpu_seq_lens
|
||||||
)
|
)
|
||||||
|
self.max_context_len = model_runner.model_config.context_len
|
||||||
|
|
||||||
def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend:
|
def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -146,6 +146,7 @@ def test_hybrid_wrappers_forward_in_graph_hook():
|
|||||||
token_to_kv_pool=None,
|
token_to_kv_pool=None,
|
||||||
req_to_token_pool=None,
|
req_to_token_pool=None,
|
||||||
server_args=SimpleNamespace(speculative_attention_mode="decode"),
|
server_args=SimpleNamespace(speculative_attention_mode="decode"),
|
||||||
|
model_config=SimpleNamespace(context_len=2048),
|
||||||
),
|
),
|
||||||
prefill_backend=make_fake("prefill", calls),
|
prefill_backend=make_fake("prefill", calls),
|
||||||
decode_backend=make_fake("decode", calls),
|
decode_backend=make_fake("decode", calls),
|
||||||
|
|||||||
+1
@@ -26,6 +26,7 @@ class _FakeBackend:
|
|||||||
def test_split_full_attention_applies_model_wrapper_once():
|
def test_split_full_attention_applies_model_wrapper_once():
|
||||||
runner = SimpleNamespace(
|
runner = SimpleNamespace(
|
||||||
server_args=SimpleNamespace(speculative_attention_mode="prefill"),
|
server_args=SimpleNamespace(speculative_attention_mode="prefill"),
|
||||||
|
model_config=SimpleNamespace(context_len=2048),
|
||||||
kv_cache_dtype=None,
|
kv_cache_dtype=None,
|
||||||
token_to_kv_pool=object(),
|
token_to_kv_pool=object(),
|
||||||
req_to_token_pool=object(),
|
req_to_token_pool=object(),
|
||||||
|
|||||||
@@ -231,6 +231,7 @@ class TestHybridNeedsCpuSeqLens(CustomTestCase):
|
|||||||
kv_cache_dtype=torch.bfloat16,
|
kv_cache_dtype=torch.bfloat16,
|
||||||
token_to_kv_pool=None,
|
token_to_kv_pool=None,
|
||||||
req_to_token_pool=None,
|
req_to_token_pool=None,
|
||||||
|
model_config=SimpleNamespace(context_len=2048),
|
||||||
)
|
)
|
||||||
return HybridAttnBackend(runner, backend(prefill_flag), backend(decode_flag))
|
return HybridAttnBackend(runner, backend(prefill_flag), backend(decode_flag))
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user