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(
+        y = self._apply_causal_sconv_kernel(
-            x=hidden_states,
+            hidden_states=hidden_states,
-            weight=weight,
             sconv_cache=sconv_cache,
-            activation=self.activation,
+            precomputed=precomputed,
-            use_residual=self.use_residual,
-            is_decode=is_decode,
-            **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(
+            y = self._apply_causal_sconv_kernel(
-                x=hidden_states,
+                hidden_states=hidden_states,
-                weight=weight,
                 sconv_cache=sconv_cache,
-                activation=self.activation,
+                precomputed=precomputed,
-                use_residual=self.use_residual,
-                is_decode=False,
-                **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(
+                y = self._apply_causal_sconv_kernel(
-                    x=hidden_states,
+                    hidden_states=hidden_states,
-                    weight=weight,
                     sconv_cache=sconv_cache,
-                    activation=self.activation,
+                    precomputed=precomputed,
-                    use_residual=self.use_residual,
-                    is_decode=False,
-                    **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(
+            y = self._apply_decode_sconv_kernel(
-                x=hidden_states,
+                hidden_states=hidden_states,
-                weight=weight,
                 sconv_cache=sconv_cache,
                 cache_indices=cache_indices,
-                cache_mask=precomputed["cache_mask"],
+                precomputed=precomputed,
-                activation=self.activation,
+                forward_batch=forward_batch,
-                use_residual=self.use_residual,
-                track_mask=forward_batch.mamba_track_mask,
-                track_indices=forward_batch.mamba_track_indices,
             )
 
         return y

test/registered/kernels/ops/mamba/test_sconv_cache.py

@@ -0,0 +1,171 @@
+import pytest
+import torch
+
+from sglang.srt.models.inkling_common.kernels.sconv import (
+    HIS_PREFIX,
+    HIS_ZEROS,
+    PAD_SLOT_ID,
+    causal_conv1d,
+    fused_decode_sconv_metadata,
+    fused_extend_sconv_metadata,
+    update_sconv_cache,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+
+@requires_cuda
+def test_update_sconv_cache_matches_reference():
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    dim, width = 128, 4
+    query_lens = torch.tensor([0, 1, 2, 3, 4, 5, 7], device="cuda")
+    query_start_loc = torch.cat(
+        [
+            torch.zeros(1, dtype=torch.int32, device="cuda"),
+            query_lens.cumsum(0).to(torch.int32),
+        ]
+    )
+    cache_indices = torch.tensor(
+        [0, 1, PAD_SLOT_ID, 3, 4, 5, 6], dtype=torch.int32, device="cuda"
+    )
+    has_initial_state = torch.tensor(
+        [True, False, True, True, False, True, False],
+        dtype=torch.bool,
+        device="cuda",
+    )
+    hidden_states = torch.randn(int(query_lens.sum()), dim, dtype=dtype, device="cuda")
+    initial_cache = torch.randn(8, width - 1, dim, dtype=dtype, device="cuda")
+    cache = initial_cache.clone()
+
+    update_sconv_cache(
+        x=hidden_states,
+        sconv_cache=cache,
+        cache_indices=cache_indices,
+        has_initial_state=has_initial_state,
+        query_start_loc=query_start_loc,
+    )
+
+    expected = initial_cache.clone()
+    for batch_idx, slot in enumerate(cache_indices.tolist()):
+        start = int(query_start_loc[batch_idx])
+        end = int(query_start_loc[batch_idx + 1])
+        if slot == PAD_SLOT_ID or start == end:
+            continue
+        prior = (
+            initial_cache[slot]
+            if has_initial_state[batch_idx]
+            else torch.zeros_like(initial_cache[slot])
+        )
+        expected[slot] = torch.cat([prior, hidden_states[start:end]])[-(width - 1) :]
+
+    torch.testing.assert_close(cache, expected, rtol=0, atol=0)
+
+
+def _extend_metadata(
+    cache_indices: torch.Tensor,
+    length: int,
+    *,
+    has_prefix: bool,
+):
+    lens = torch.tensor([length], dtype=torch.int64, device="cuda")
+    result = fused_extend_sconv_metadata(
+        B=1,
+        T=length,
+        cache_indices=cache_indices,
+        his_mode=HIS_PREFIX if has_prefix else HIS_ZEROS,
+        extend_seq_lens=lens,
+        his_src=lens if has_prefix else None,
+    )
+    assert result is not None
+    return result
+
+
+@requires_cuda
+def test_cached_continuations_match_full_prefill():
+    torch.manual_seed(1)
+    dtype = torch.bfloat16
+    length, prefix_len, dim, width = 24, 10, 128, 4
+    hidden_states = torch.randn(length, dim, dtype=dtype, device="cuda")
+    weight = torch.randn(dim, width, dtype=dtype, device="cuda")
+    cache_indices = torch.zeros(1, dtype=torch.int32, device="cuda")
+
+    full_cache = torch.zeros(8, width - 1, dim, dtype=dtype, device="cuda")
+    _, _, full_meta = _extend_metadata(cache_indices, length, has_prefix=False)
+    full_output = causal_conv1d(
+        x=hidden_states,
+        weight=weight,
+        sconv_cache=full_cache,
+        activation="silu",
+        use_residual=True,
+        **full_meta,
+    )
+
+    decode_cache = torch.zeros_like(full_cache)
+    query_start_loc, has_initial_state, decode_meta = fused_decode_sconv_metadata(
+        B=1, cache_indices=cache_indices
+    )
+    decode_outputs = []
+    for token in hidden_states.split(1):
+        decode_outputs.append(
+            causal_conv1d(
+                x=token,
+                weight=weight,
+                sconv_cache=decode_cache,
+                activation="silu",
+                use_residual=True,
+                is_decode=True,
+                **decode_meta,
+            )
+        )
+        update_sconv_cache(
+            x=token,
+            sconv_cache=decode_cache,
+            cache_indices=cache_indices,
+            has_initial_state=has_initial_state,
+            query_start_loc=query_start_loc,
+        )
+    decode_output = torch.cat(decode_outputs)
+
+    extend_cache = torch.zeros_like(full_cache)
+    prefix = hidden_states[:prefix_len]
+    prefix_qsl, prefix_his, prefix_meta = _extend_metadata(
+        cache_indices, prefix_len, has_prefix=False
+    )
+    causal_conv1d(
+        x=prefix,
+        weight=weight,
+        sconv_cache=extend_cache,
+        activation="silu",
+        use_residual=True,
+        **prefix_meta,
+    )
+    update_sconv_cache(
+        x=prefix,
+        sconv_cache=extend_cache,
+        cache_indices=cache_indices,
+        has_initial_state=prefix_his,
+        query_start_loc=prefix_qsl,
+    )
+    suffix = hidden_states[prefix_len:]
+    _, _, suffix_meta = _extend_metadata(cache_indices, len(suffix), has_prefix=True)
+    suffix_output = causal_conv1d(
+        x=suffix,
+        weight=weight,
+        sconv_cache=extend_cache,
+        activation="silu",
+        use_residual=True,
+        **suffix_meta,
+    )
+
+    torch.testing.assert_close(decode_output, full_output, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(
+        suffix_output, full_output[prefix_len:], rtol=2e-2, atol=2e-2
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v", "-x"]))