Fix hybrid linear attention dispatch by layer id with draft-worker awareness (#27120)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

python/sglang/srt/layers/attention/attention_registry.py

@@ -314,7 +314,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
                     "If this is a custom hybrid model, use register_linear_attn_model() "
                     "from sglang.srt.configs.linear_attn_model_registry."
                 )
-        full_attn_layers = cfg.full_attention_layer_ids
+        if runner.is_draft_worker:
+            # FIXME: we assume that MTP/NEXTN always use full-attention.
+            full_attn_layers = [0]
+        else:
+            full_attn_layers = cfg.full_attention_layer_ids
         return HybridLinearAttnBackend(
             full_attn_backend, linear_attn_backend, full_attn_layers
         )

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -16,7 +16,6 @@ from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
     fused_mamba_state_scatter_with_mask,
 )
 from sglang.srt.layers.radix_attention import RadixAttention
-from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
@@ -784,21 +783,6 @@ class HybridLinearAttnBackend(AttentionBackend):
     def _is_full_attn(
         self, layer: Optional[RadixAttention], layer_id: Optional[int] = None
     ) -> bool:
-        # Explicit linear-attention subclass → strong linear signal (KDA, GDN,
-        # Qwen3-Next, Qwen3.5 main linear layers).
-        if isinstance(layer, RadixLinearAttention):
-            return False
-        # Some hybrid models (Ling-2.5/2.6) wrap their linear layers in plain
-        # `RadixAttention` rather than `RadixLinearAttention`. Those wrappers
-        # set `_is_linear_attention=True` on the attn module so we can
-        # distinguish them from full-attention RadixAttention instances —
-        # including MTP/NEXTN draft layers, which are full and must default to
-        # the full-attn path.
-        if layer is not None and getattr(layer, "_is_linear_attention", False):
-            return False
-        if isinstance(layer, RadixAttention):
-            return True
-
         if layer is not None:
             layer_id = layer.layer_id
         assert layer_id is not None, "either layer or layer_id must be provided"

python/sglang/srt/models/bailing_moe_linear.py

@@ -508,12 +508,6 @@ class BailingMoELinearAttention(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.attn",
         )
-        # Marker for HybridLinearAttnBackend._is_full_attn: Bailing wraps
-        # linear-attention layers in a plain RadixAttention, so the
-        # dispatcher can't tell from the type alone that this is a linear
-        # layer (would otherwise default to the full-attn backend, e.g. the
-        # same way MTP/NEXTN draft layers are routed).
-        self.attn._is_linear_attention = True
 
         self.group_norm_size = getattr(config, "group_norm_size", 1)
         self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))