Add registered short-conv tests and backend extensions (#34045)

python/sglang/srt/configs/linear_attn_model_registry.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 
 import importlib
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
@@ -36,6 +37,8 @@ class LinearAttnModelSpec:
     support_mamba_cache: bool = True
     support_mamba_cache_extra_buffer: bool = False
     unwrap_text_config: bool = False  # call get_text_config() before isinstance check
+    hybrid_backend_class_name: str | None = None
+    config_predicate: Callable[[Any], bool] | None = None
 
 
 _LINEAR_ATTN_MODEL_REGISTRY: list[LinearAttnModelSpec] = []
@@ -54,7 +57,9 @@ def register_linear_attn_model(spec: LinearAttnModelSpec) -> None:
 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):
+        if isinstance(config, spec.config_class) and (
+            spec.config_predicate is None or spec.config_predicate(config)
+        ):
             return spec, config
     return None
 

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

@@ -460,8 +460,13 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
             spec_result = get_linear_attn_config(runner.model_config.hf_config)
             if spec_result is not None:
                 spec, _ = spec_result
+                cfg = runner.model_config
                 BackendClass = import_backend_class(spec.backend_class_name)
                 linear_attn_backend = BackendClass(runner)
+                if spec.hybrid_backend_class_name is not None:
+                    hybrid_backend_cls = import_backend_class(
+                        spec.hybrid_backend_class_name
+                    )
             else:
                 raise ValueError(
                     "Expected hybrid GDN or NemotronH models, but got unknown model. "

python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py

@@ -564,6 +564,10 @@ class InklingShortConvHybridAttnBackend(ShortConvHybridAttnBackend):
         # one (KV write locs, the SWA loc translate).
         return self.full_attn_backend.forward_metadata
 
+    @forward_metadata.setter
+    def forward_metadata(self, value):
+        self.full_attn_backend.forward_metadata = value
+
     @property
     def supports_ragged_verify_graph(self) -> bool:
         return self.full_attn_backend.supports_ragged_verify_graph

python/sglang/srt/models/inkling_common/sconv.py

@@ -143,22 +143,16 @@ class ShortConvolution(nn.Module):
     def _apply_training_sconv_kernel(
         self,
         hidden_states: torch.Tensor,
-        weight: torch.Tensor,
         sconv_cache: torch.Tensor,
         cache_indices: torch.Tensor,
         query_start_loc: torch.Tensor,
         has_initial_state: torch.Tensor,
         precomputed: SconvDecodeMetadata | SconvExtendMetadata,
-        is_decode: bool = False,
     ) -> torch.Tensor:
-        y = causal_conv1d(
-            x=hidden_states,
-            weight=weight,
+        y = self._apply_causal_sconv_kernel(
+            hidden_states=hidden_states,
             sconv_cache=sconv_cache,
-            activation=self.activation,
-            use_residual=self.use_residual,
-            is_decode=is_decode,
-            **precomputed,
+            precomputed=precomputed,
         )
         update_sconv_cache(
             x=hidden_states,
@@ -169,6 +163,41 @@ class ShortConvolution(nn.Module):
         )
         return y
 
+    def _apply_causal_sconv_kernel(
+        self,
+        hidden_states: torch.Tensor,
+        sconv_cache: torch.Tensor,
+        precomputed: SconvDecodeMetadata | SconvExtendMetadata,
+    ) -> torch.Tensor:
+        return causal_conv1d(
+            x=hidden_states,
+            weight=self._weight_2d(),
+            sconv_cache=sconv_cache,
+            activation=self.activation,
+            use_residual=self.use_residual,
+            **precomputed,
+        )
+
+    def _apply_decode_sconv_kernel(
+        self,
+        hidden_states: torch.Tensor,
+        sconv_cache: torch.Tensor,
+        cache_indices: torch.Tensor,
+        precomputed: SconvDecodeMetadata | SconvExtendMetadata,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        return fused_causal_conv1d_update_decode(
+            x=hidden_states,
+            weight=self._weight_2d(),
+            sconv_cache=sconv_cache,
+            cache_indices=cache_indices,
+            cache_mask=precomputed["cache_mask"],
+            activation=self.activation,
+            use_residual=self.use_residual,
+            track_mask=forward_batch.mamba_track_mask,
+            track_indices=forward_batch.mamba_track_indices,
+        )
+
     def _prepare_extend_sconv_cache(
         self,
         forward_batch: ForwardBatch,
@@ -378,17 +407,12 @@ class ShortConvolution(nn.Module):
         cache_indices = meta.cache_indices
         sconv_cache = self._sconv_cache()
         precomputed = meta.precomputed
-        weight = self._weight_2d()
 
         if forward_batch.forward_mode.is_target_verify():
-            y = causal_conv1d(
-                x=hidden_states,
-                weight=weight,
+            y = self._apply_causal_sconv_kernel(
+                hidden_states=hidden_states,
                 sconv_cache=sconv_cache,
-                activation=self.activation,
-                use_residual=self.use_residual,
-                is_decode=False,
-                **precomputed,
+                precomputed=precomputed,
             )
             self._save_intermediate_conv_windows(
                 forward_batch=forward_batch,
@@ -403,14 +427,10 @@ class ShortConvolution(nn.Module):
             )
 
             if forward_batch.forward_mode.is_draft_extend_v2():
-                y = causal_conv1d(
-                    x=hidden_states,
-                    weight=weight,
+                y = self._apply_causal_sconv_kernel(
+                    hidden_states=hidden_states,
                     sconv_cache=sconv_cache,
-                    activation=self.activation,
-                    use_residual=self.use_residual,
-                    is_decode=False,
-                    **precomputed,
+                    precomputed=precomputed,
                 )
                 self._update_sconv_cache_for_draft_extend(
                     forward_batch,
@@ -421,13 +441,11 @@ class ShortConvolution(nn.Module):
             else:
                 y = self._apply_training_sconv_kernel(
                     hidden_states=hidden_states,
-                    weight=weight,
                     sconv_cache=sconv_cache,
                     cache_indices=cache_indices,
                     query_start_loc=meta.query_start_loc,
                     has_initial_state=meta.has_initial_state,
                     precomputed=precomputed,
-                    is_decode=False,
                 )
         else:
             # Fused decode: prefix construction + conv + cache update + prefix-cache
@@ -436,16 +454,12 @@ class ShortConvolution(nn.Module):
             # into the persistent ping-pong slot in-register (no separate
             # copy_if_needed launch). track_mask is None when prefix caching with the
             # mamba extra buffer is disabled, which disables the track-copy path.
-            y = fused_causal_conv1d_update_decode(
-                x=hidden_states,
-                weight=weight,
+            y = self._apply_decode_sconv_kernel(
+                hidden_states=hidden_states,
                 sconv_cache=sconv_cache,
                 cache_indices=cache_indices,
-                cache_mask=precomputed["cache_mask"],
-                activation=self.activation,
-                use_residual=self.use_residual,
-                track_mask=forward_batch.mamba_track_mask,
-                track_indices=forward_batch.mamba_track_indices,
+                precomputed=precomputed,
+                forward_batch=forward_batch,
             )
 
         return y