From d5ec233834ab58eb5f551aa75dc9f65ea0ae2234 Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Wed, 25 Mar 2026 21:20:47 +0800
Subject: [PATCH] feat(mtp): add tbo mtp support

---
 .../deepseek_common/attention_forward_methods/forward_mha.py  | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

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