fix: fix MockModelRunner in attention tests (#18240)

This commit is contained in:
Zack Yu
2026-02-04 13:18:02 -08:00
committed by GitHub
parent 6fd878b41d
commit 2e87c2bd5e
2 changed files with 9 additions and 3 deletions

View File

@@ -24,6 +24,9 @@ class MockModelRunner:
):
self.device = "cuda"
self.dtype = torch.float16
self.kv_cache_dtype = torch.float16
self.is_hybrid_swa = False
self.attention_chunk_size = None
attention_arch = AttentionArch.MHA
# Max batch size for the test.
max_batch_size = 160
@@ -36,10 +39,11 @@ class MockModelRunner:
"context_len": max_context_len,
"is_multimodal": False,
"attention_arch": attention_arch,
"is_encoder_decoder": False,
"is_local_attention_model": False,
},
)
)()
self.sliding_window_size = None
self.device = self.device
# Create a large enough req_to_token_pool to fit the test usage.
self.req_to_token_pool = type(
"TokenPool",
@@ -55,7 +59,7 @@ class MockModelRunner:
device=self.device,
),
},
)
)()
self.page_size = page_size
max_total_num_tokens = max_batch_size * max_context_len
self.token_to_kv_pool = MHATokenToKVPool(

View File

@@ -28,6 +28,8 @@ class MockModelRunner:
{
"context_len": context_len,
"attention_arch": attention_arch,
"is_encoder_decoder": False,
"is_local_attention_model": False,
},
)
self.sliding_window_size = None