diff --git a/python/sglang/srt/configs/linear_attn_model_registry.py b/python/sglang/srt/configs/linear_attn_model_registry.py index 33fdae8f0..ba30fcf43 100644 --- a/python/sglang/srt/configs/linear_attn_model_registry.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 594646a44..603ea0c99 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/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. " diff --git a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py b/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py index d9f56c1b7..4d03f9cdb 100644 --- a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py +++ b/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 diff --git a/python/sglang/srt/models/inkling_common/sconv.py b/python/sglang/srt/models/inkling_common/sconv.py index 33e726658..167fc7dad 100644 --- a/python/sglang/srt/models/inkling_common/sconv.py +++ b/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 diff --git a/test/registered/kernels/ops/mamba/test_sconv_cache.py b/test/registered/kernels/ops/mamba/test_sconv_cache.py new file mode 100644 index 000000000..d4ccbbd67 --- /dev/null +++ b/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"]))