Add registration API for external linear attention backend (#21983)

Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com>

python/sglang/srt/configs/linear_attn_model_registry.py

@@ -0,0 +1,72 @@
+"""Registry for linear attention hybrid models (softmax + linear attention).
+
+External models can register themselves without modifying SGLang core files:
+
+    from sglang.srt.configs.linear_attn_model_registry import (
+        register_linear_attn_model, LinearAttnModelSpec,
+    )
+
+    register_linear_attn_model(LinearAttnModelSpec(
+        config_class=MyLinearAttnConfig,
+        backend_class_name="sglang.srt.layers.attention.linear.kda_backend.KDAAttnBackend",
+        arch_names=["MyLinearAttnForCausalLM"],
+        uses_mamba_radix_cache=True,
+        support_mamba_cache=True,
+    ))
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LinearAttnModelSpec:
+    """Specification for a hybrid (softmax + linear attention) model."""
+
+    config_class: type
+    backend_class_name: str  # fully-qualified class name, lazily imported
+    arch_names: list[str] = field(default_factory=list)
+    uses_mamba_radix_cache: bool = True
+    support_mamba_cache: bool = True
+    support_mamba_cache_extra_buffer: bool = False
+    unwrap_text_config: bool = False  # call get_text_config() before isinstance check
+
+
+_LINEAR_ATTN_MODEL_REGISTRY: list[LinearAttnModelSpec] = []
+
+
+def register_linear_attn_model(spec: LinearAttnModelSpec) -> None:
+    _LINEAR_ATTN_MODEL_REGISTRY.append(spec)
+    logger.info(
+        "Registered linear attn model: config=%s, backend=%s, archs=%s",
+        spec.config_class.__name__,
+        spec.backend_class_name.rsplit(".", 1)[-1],
+        spec.arch_names,
+    )
+
+
+def get_linear_attn_config(hf_config: Any) -> Optional[tuple[LinearAttnModelSpec, Any]]:
+    for spec in _LINEAR_ATTN_MODEL_REGISTRY:
+        config = hf_config.get_text_config() if spec.unwrap_text_config else hf_config
+        if isinstance(config, spec.config_class):
+            return spec, config
+    return None
+
+
+def get_linear_attn_spec_by_arch(arch_name: str) -> Optional[LinearAttnModelSpec]:
+    for spec in _LINEAR_ATTN_MODEL_REGISTRY:
+        if arch_name in spec.arch_names:
+            return spec
+    return None
+
+
+def import_backend_class(dotted_name: str) -> type:
+    module_path, class_name = dotted_name.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)

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

@@ -1,6 +1,11 @@
 import logging
 from typing import TYPE_CHECKING
 
+from sglang.srt.configs.linear_attn_model_registry import (
+    get_linear_attn_config,
+    import_backend_class,
+)
+
 logger = logging.getLogger(__name__)
 
 
@@ -225,9 +230,17 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
         elif runner.hybrid_lightning_config is not None:
             linear_attn_backend = LightningAttentionBackend(runner)
         else:
-            raise ValueError(
-                "Expected hybrid GDN or NemotronH models, but got unknown model."
-            )
+            spec_result = get_linear_attn_config(runner.model_config.hf_config)
+            if spec_result is not None:
+                spec, _ = spec_result
+                BackendClass = import_backend_class(spec.backend_class_name)
+                linear_attn_backend = BackendClass(runner)
+            else:
+                raise ValueError(
+                    "Expected hybrid GDN or NemotronH models, but got unknown model. "
+                    "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
         return HybridLinearAttnBackend(
             full_attn_backend, linear_attn_backend, full_attn_layers

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

@@ -111,6 +111,7 @@ class TritonAttnBackend(AttentionBackend):
         elif (
             model_runner.hybrid_gdn_config is not None
             or model_runner.kimi_linear_config is not None
+            or model_runner.linear_attn_model_spec is not None
         ):
             # For hybrid linear models, layer_id = 0 may not be full attention
             self.v_head_dim = model_runner.token_to_kv_pool.get_v_head_dim()

python/sglang/srt/managers/scheduler.py

@@ -725,9 +725,14 @@ class Scheduler(
 
         # Hybrid memory pool
         self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
+        _spec = self.tp_worker.model_runner.linear_attn_model_spec
+        _registry_needs_mamba = (
+            _spec.uses_mamba_radix_cache if _spec is not None else False
+        )
         self.is_hybrid_ssm = (
             self.tp_worker.model_runner.hybrid_gdn_config is not None
             or self.tp_worker.model_runner.mamba2_config is not None
+            or _registry_needs_mamba
         )
 
         self.sliding_window_size = None

python/sglang/srt/model_executor/model_runner.py

@@ -50,6 +50,7 @@ from sglang.srt.configs import (
     Qwen3NextConfig,
 )
 from sglang.srt.configs.device_config import DeviceConfig
+from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config
 from sglang.srt.configs.load_config import LoadConfig, LoadFormat
 from sglang.srt.configs.model_config import AttentionArch, ModelConfig, ModelImpl
 from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp
@@ -1890,14 +1891,30 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             return config
         return None
 
+    def _get_linear_attn_registry_result(self):
+        if not hasattr(self, "_linear_attn_registry_cache"):
+            self._linear_attn_registry_cache = get_linear_attn_config(
+                self.model_config.hf_config
+            )
+        return self._linear_attn_registry_cache
+
+    @property
+    def linear_attn_model_spec(self):
+        result = self._get_linear_attn_registry_result()
+        return result[0] if result else None
+
     @property
     def mambaish_config(self):
-        return (
+        existing = (
             self.mamba2_config
             or self.hybrid_gdn_config
             or self.kimi_linear_config
             or self.hybrid_lightning_config
         )
+        if existing:
+            return existing
+        result = self._get_linear_attn_registry_result()
+        return result[1] if result else None
 
     def configure_kv_cache_dtype(self):
         if self.server_args.kv_cache_dtype == "auto":

python/sglang/srt/server_args.py

@@ -26,6 +26,7 @@ import random
 import tempfile
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
+from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_spec_by_arch
 from sglang.srt.connector import ConnectorType
 from sglang.srt.environ import envs
 from sglang.srt.function_call.function_call_parser import FunctionCallParser
@@ -1501,6 +1502,14 @@ class ServerArgs:
         hf_config = self.get_model_config().hf_config
         model_arch = hf_config.architectures[0]
 
+        _hybrid_spec = get_linear_attn_spec_by_arch(model_arch)
+        if _hybrid_spec is not None:
+            self._handle_mamba_radix_cache(
+                model_arch=model_arch,
+                support_mamba_cache=_hybrid_spec.support_mamba_cache,
+                support_mamba_cache_extra_buffer=_hybrid_spec.support_mamba_cache_extra_buffer,
+            )
+
         if model_arch in [
             "MistralLarge3ForCausalLM",
             "PixtralForConditionalGeneration",