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