[Perf] Remove two operations in gdn_backend extend verify path (#22444)
This commit is contained in:
@@ -255,6 +255,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
||||
decode_backend = get_linear_attn_decode_backend()
|
||||
prefill_backend = get_linear_attn_prefill_backend()
|
||||
self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)
|
||||
self.verify_intermediate_state_indices = torch.arange(
|
||||
self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device
|
||||
)
|
||||
|
||||
def init_forward_metadata(self, forward_batch: ForwardBatch):
|
||||
super().init_forward_metadata(forward_batch)
|
||||
@@ -373,14 +376,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
||||
intermediate_conv_window_cache = (
|
||||
mamba_cache_params.intermediate_conv_window[0]
|
||||
)
|
||||
has_initial_states = torch.ones(
|
||||
seq_len // forward_batch.spec_info.draft_token_num,
|
||||
dtype=torch.bool,
|
||||
device=forward_batch.input_ids.device,
|
||||
)
|
||||
intermediate_state_indices = torch.arange(
|
||||
cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device
|
||||
)
|
||||
intermediate_state_indices = self.verify_intermediate_state_indices
|
||||
else:
|
||||
has_initial_states = forward_batch.extend_prefix_lens > 0
|
||||
|
||||
|
||||
Reference in New Issue
Block a user