minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Aurick Qiao
2026-08-08 14:38:03 +08:00
committed by GitHub
parent 6679d9b60c
commit 6185ed8011
5 changed files with 234 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/linear_attn_model_registry.py
+6 -1
View File
@@ -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
+5
View File
@@ -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
+4
View File
@@ -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
+48 -34
View File
@@ -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
test/registered/kernels/ops/mamba/test_sconv_cache.py
+171
View File
@@ -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"]))