diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 5ff13ebbb..34b839d6c 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -3,7 +3,7 @@ from __future__ import annotations from typing import TYPE_CHECKING import torch - +from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged from sglang.srt.layers.attention.tbo_backend import TboAttnBackend @@ -466,6 +466,8 @@ class DeepseekMHAForwardMixin: backend = forward_batch.attn_backend if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend backend = backend.primary + if isinstance(backend, HybridAttnBackend): + backend = backend._select_backend(forward_batch.forward_mode) kv_indices = backend.forward_metadata.page_table_1_flattened assert ( kv_indices is not None