Fix hybrid_linear_attn_backend crash with ngram speculation (#20739)

Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yufeng He
2026-04-08 12:52:07 -07:00
committed by GitHub
co-authored by kpham-sgl Claude Opus 4.6
parent c26b8b4a4b
commit c89afaea7c
@@ -139,6 +139,7 @@ class MambaAttnBackendBase(AttentionBackend):
super().__init__()
self.pad_slot_id = PAD_SLOT_ID
self.device = model_runner.device
self.topk = model_runner.server_args.speculative_eagle_topk or 0
self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool
self.forward_metadata: ForwardMetadata = None
self.state_indices_list = []
@@ -185,7 +186,7 @@ class MambaAttnBackendBase(AttentionBackend):
device=forward_batch.input_ids.device,
)
if forward_batch.spec_info.topk > 1:
if self.topk > 1:
retrieve_next_token = forward_batch.spec_info.retrive_next_token
retrieve_next_sibling = forward_batch.spec_info.retrive_next_sibling
# retrieve_next_token is None during dummy run so skip tensor creation
@@ -482,7 +483,7 @@ class MambaAttnBackendBase(AttentionBackend):
self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
if forward_mode.is_target_verify() and spec_info.topk > 1:
if forward_mode.is_target_verify() and self.topk > 1:
# They are None during cuda graph capture so skip the copy_...
# self.retrieve_next_token_list[bs - 1].copy_(spec_info.retrive_next_token)
# self.retrieve_next_sibling_list[bs - 1].copy_(spec_info.retrive_next_sibling)
@@ -543,7 +544,7 @@ class MambaAttnBackendBase(AttentionBackend):
raise ValueError(f"Invalid forward mode: {forward_mode=}")
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
if forward_mode.is_target_verify() and spec_info.topk > 1:
if forward_mode.is_target_verify() and self.topk > 1:
bs_without_pad = spec_info.retrive_next_token.shape[0]
self.retrieve_next_token_list[bs - 1][:bs_without_pad].copy_(
spec_info.retrive_next_token