test: update hybrid attention runner fixtures (#37345)
This commit is contained in:
@@ -162,6 +162,7 @@ def test_hybrid_wrappers_forward_in_graph_hook():
|
|||||||
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,
|
||||||
|
kv_index_translator=None,
|
||||||
server_args=SimpleNamespace(speculative_attention_mode="decode"),
|
server_args=SimpleNamespace(speculative_attention_mode="decode"),
|
||||||
model_config=SimpleNamespace(context_len=2048),
|
model_config=SimpleNamespace(context_len=2048),
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -148,6 +148,7 @@ def _make_hybrid_backend(speculative_attention_mode, prefill_mask, decode_mask):
|
|||||||
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(),
|
||||||
|
kv_index_translator=None,
|
||||||
server_args=SimpleNamespace(
|
server_args=SimpleNamespace(
|
||||||
speculative_attention_mode=speculative_attention_mode
|
speculative_attention_mode=speculative_attention_mode
|
||||||
),
|
),
|
||||||
|
|||||||
+1
@@ -37,6 +37,7 @@ def test_split_full_attention_applies_model_wrapper_once():
|
|||||||
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(),
|
||||||
|
kv_index_translator=None,
|
||||||
init_new_workspace=None,
|
init_new_workspace=None,
|
||||||
)
|
)
|
||||||
wrapper_inputs = []
|
wrapper_inputs = []
|
||||||
|
|||||||
@@ -235,6 +235,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,
|
||||||
|
kv_index_translator=None,
|
||||||
model_config=SimpleNamespace(context_len=2048),
|
model_config=SimpleNamespace(context_len=2048),
|
||||||
)
|
)
|
||||||
# The backend takes the mode from the published configuration, not from
|
# The backend takes the mode from the published configuration, not from
|
||||||
|
|||||||
Reference in New Issue
Block a user