fix: fix MockModelRunner in attention tests (#18240)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user