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

feat(short-conv): shared ShortConvAttnBackend for ZAYA1 CCA + LFM2 short conv (#29867)

Browse Source
This commit is contained in:
Cheng Wan
2026-07-01 20:08:00 -07:00
committed by GitHub
parent 70df09b833
commit b558cc1abb
8 changed files with 652 additions and 338 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/attention_registry.py
+42 -2
View File
@@ -288,6 +288,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
check_environments()
initialize_linear_attn_config(runner.server_args)
hybrid_backend_cls = HybridLinearAttnBackend
if runner.hybrid_gdn_config is not None:
if is_blackwell():
assert (
@@ -303,7 +304,46 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
logger.info(f"Using hybrid linear attention backend for hybrid GDN models.")
linear_attn_backend = GDNAttnBackend(runner)
elif runner.mamba2_config is not None:
linear_attn_backend = Mamba2AttnBackend(runner)
from sglang.srt.configs.lfm2 import Lfm2Config
from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
from sglang.srt.configs.lfm2_vl import Lfm2VlConfig
from sglang.srt.configs.zaya import ZayaConfig
# Short-conv hybrids (ZAYA1 CCA, LFM2 short conv) share a conv-state
# sidecar that owns the per-request state plumbing and is invoked by
# the model via conv_state_metadata (never as a full-vs-linear
# alternative). Other mamba2 models keep the full Mamba2 SSM backend.
short_conv_cfgs = (
ZayaConfig,
Lfm2Config,
Lfm2MoeConfig,
Lfm2VlConfig,
)
if isinstance(runner.mamba2_config, short_conv_cfgs):
if is_npu():
# The model conv layers call
# get_attn_backend().conv_state_metadata() unconditionally,
# but the Ascend hybrid/mamba backend has no such method.
# Fail here (before model execution) with a clear message
# rather than an AttributeError deep in the first conv layer.
raise NotImplementedError(
"Short-conv hybrid models (ZAYA1 CCA, LFM2 / LFM2-MoE) "
"are not yet supported on NPU: the conv-state sidecar "
"(ShortConvAttnBackend.conv_state_metadata) has no Ascend "
"implementation. Add an Ascend conv-state backend before "
"serving these models on NPU."
)
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
ShortConvHybridAttnBackend,
)
from sglang.srt.layers.attention.linear.short_conv_backend import (
ShortConvAttnBackend,
)
linear_attn_backend = ShortConvAttnBackend(runner)
hybrid_backend_cls = ShortConvHybridAttnBackend
else:
linear_attn_backend = Mamba2AttnBackend(runner)
elif runner.kimi_linear_config is not None:
linear_attn_backend = KDAAttnBackend(runner)
elif runner.hybrid_lightning_config is not None:
@@ -325,7 +365,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
full_attn_layers = [0]
else:
full_attn_layers = cfg.full_attention_layer_ids
return HybridLinearAttnBackend(
return hybrid_backend_cls(
full_attn_backend, linear_attn_backend, full_attn_layers
)
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+28
View File
@@ -1096,3 +1096,31 @@ class HybridLinearAttnBackend(AttentionBackend):
mamba_track_indices,
mamba_steps_to_track,
)
class ShortConvHybridAttnBackend(HybridLinearAttnBackend):
"""HybridLinearAttnBackend variant for short-conv hybrid models (ZAYA1 CCA,
LFM2 short conv).
The linear sidecar is a :class:`ShortConvAttnBackend
<sglang.srt.layers.attention.linear.short_conv_backend.ShortConvAttnBackend>`
that owns the per-request conv-state plumbing. The model's conv module
reaches it via :meth:`conv_state_metadata` (``get_attn_backend()`` returns
this wrapper) and runs its own conv kernel against the returned handle, so
the model definition holds no pool access. The sidecar is never reached
through the full-vs-linear ``forward_decode`` / ``forward_extend`` dispatch.
"""
def __init__(
self,
full_attn_backend: AttentionBackend,
short_conv_backend: MambaAttnBackendBase,
full_attn_layers: list,
):
# Register short_conv_backend as the linear sidecar so it rides in
# attn_backend_list and inherits the metadata / cuda-graph fan-out.
super().__init__(full_attn_backend, short_conv_backend, full_attn_layers)
self.short_conv_backend = short_conv_backend
def conv_state_metadata(self, layer_id: int, forward_batch: ForwardBatch):
return self.short_conv_backend.conv_state_metadata(layer_id, forward_batch)
python/sglang/srt/layers/attention/linear/short_conv_backend.py
+219
View File
@@ -0,0 +1,219 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Short-convolution attention backend.
Several hybrid models interleave a *causal short conv with per-request conv
state* (stored in the centralized ``MambaPool``) with softmax attention layers:
* **LFM2** (:class:`Lfm2ShortConv <sglang.srt.models.lfm2.Lfm2ShortConv>`) --
a depthwise gated short conv (``causal_conv1d_fn`` / ``causal_conv1d_update``)
as a standalone token mixer on its own conv layers.
* **ZAYA1** (:class:`CCA <sglang.srt.models.zaya.CCA>`) -- a two-stage grouped
conv plus a one-token ``prev_hs`` lag, preprocessing q/k for the layer's
softmax attention.
These share the *state plumbing* -- resolving the per-request slot indices, the
``has_initial_state`` prefix mask, the ``query_start_loc`` cu-seqlens, and the
cuda-graph static index buffers, all once per forward step -- but NOT the conv
kernel itself. ``ShortConvAttnBackend`` owns only the plumbing and hands it out
via :meth:`conv_state_metadata` as a :class:`ShortConvMetadata`; each model runs
its own conv kernel against that handle, so the model definition holds no pool
access.
The backend is a *sidecar*: it is invoked directly by the model (through
:class:`ShortConvHybridAttnBackend
<sglang.srt.layers.attention.hybrid_linear_attn_backend.ShortConvHybridAttnBackend>`),
never through the full-vs-linear ``forward_decode`` / ``forward_extend``
dispatch. Metadata + cuda-graph capture/replay come from
:class:`MambaAttnBackendBase`.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, List, NamedTuple, Optional
import torch
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
MambaAttnBackendBase,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
class ShortConvMetadata(NamedTuple):
"""Per-(layer, step) conv-state handle handed to a model's conv kernel.
``layer_cache`` exposes the per-layer pool views (``conv[0]`` = conv state,
``conv[1]`` = an optional second state such as ZAYA1's ``prev_hs``,
``temporal`` = SSM state, unused by pure short convs). The device tensors are
cuda-graph-static on the decode/replay path; the ``*_cpu`` host mirrors are
built once per step only for models whose extend path runs a host loop
(e.g. ZAYA1 v1) and are ``None`` on decode.
"""
layer_cache: Any
cache_indices: torch.Tensor
# cu-seqlens for the varlen prefill conv (device, int32). None on decode.
query_start_loc: Optional[torch.Tensor] = None
# Per-request "resumes a cached prefix" mask (device bool). None on decode.
has_initial_state: Optional[torch.Tensor] = None
# Host mirror of cache_indices for extend host loops. None on decode.
slot_ids_cpu: Optional[List[int]] = None
# Host mirror of has_initial_state for extend host loops. None on decode.
has_prefix_cpu: Optional[List[bool]] = None
class ShortConvAttnBackend(MambaAttnBackendBase):
"""Owns the short-conv per-request state plumbing (see module docstring)."""
# State IO is index-driven; no host seq-lens plumbing required from the
# runner. (The extend path reads ``extend_*_cpu`` off the batch, which is
# always populated for extend regardless of this flag.)
needs_cpu_seq_lens: bool = False
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
mamba_cache = self.req_to_token_pool.mamba_pool.mamba_cache
# conv[0] == conv_state: [n_layers, n_slots, conv_dim, conv_kernel - 1]
self.conv_states_shape = mamba_cache.conv[0].shape
# Per-step state, resolved ONCE per step in init_forward_metadata /
# init_forward_metadata_out_graph (never per conv layer). The extend host
# mirrors drive the extend loop; ``_cache_indices`` is the int64 slot
# index view shared by all conv layers within the step.
self._has_initial_state: Optional[torch.Tensor] = None
self._slot_ids_cpu: Optional[List[int]] = None
self._has_prefix_cpu: Optional[List[bool]] = None
self._cache_indices: Optional[torch.Tensor] = None
self._cache_indices_buf: Optional[torch.Tensor] = None
def _reset_step_state(self):
self._has_initial_state = None
self._slot_ids_cpu = None
self._has_prefix_cpu = None
def _alloc_cache_indices_buf(self, max_bs: int):
# Persistent int64 index buffer, refilled in place per step so the
# captured (cuda or cpu) graph reads a stable address.
self._cache_indices_buf = torch.empty(
max_bs, dtype=torch.int64, device=self.device
)
def _refresh_cache_indices(self):
# Resolve the int64 slot-index view ONCE per step, shared by every conv
# layer. When a graph index buffer is allocated and large enough, refill
# it IN PLACE and hand out a view -- the captured graph then reads a
# stable address that this (pre-replay) hook keeps current, so it is
# cuda- and cpu-graph safe. Otherwise (eager, or bs beyond the buffer)
# a fresh cast is fine.
md = self.forward_metadata
idx = md.mamba_cache_indices if md is not None else None
buf = self._cache_indices_buf
if idx is None:
self._cache_indices = None
elif buf is not None and idx.shape[0] <= buf.shape[0]:
n = idx.shape[0]
buf[:n].copy_(idx)
self._cache_indices = buf[:n]
else:
self._cache_indices = idx.to(torch.long)
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
super().init_cuda_graph_state(max_bs, max_num_tokens)
self._alloc_cache_indices_buf(max_bs)
def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int):
super().init_cpu_graph_state(max_bs, max_num_tokens)
self._alloc_cache_indices_buf(max_bs)
def init_forward_metadata(self, forward_batch: ForwardBatch):
# Eager path (also the CPU-graph replay path). Builds
# self.forward_metadata and runs the deferred mamba clear/COW ops.
super().init_forward_metadata(forward_batch)
self._reset_step_state()
self._refresh_cache_indices()
mode = forward_batch.forward_mode
if (
mode.is_extend()
and not mode.is_target_verify()
and not mode.is_draft_extend_v2()
):
self._has_initial_state = forward_batch.extend_prefix_lens > 0
if self._cache_indices is not None:
self._slot_ids_cpu = self._cache_indices.tolist()
self._has_prefix_cpu = [
int(p) > 0 for p in forward_batch.extend_prefix_lens_cpu
]
def init_forward_metadata_out_graph(
self, forward_batch: ForwardBatch, in_capture: bool = False
):
# Decode cuda-graph capture + replay path -- no extend prefix state.
super().init_forward_metadata_out_graph(forward_batch, in_capture)
self._reset_step_state()
self._refresh_cache_indices()
def init_forward_metadata_capture_cpu_graph(self, *args, **kwargs):
# Decode CPU-graph capture path. The base fills forward_metadata but not
# the int64 view; without this the conv layers would capture a ``None``
# index (crash / corrupt state). Replay goes through init_forward_metadata
# and refills the SAME buffer, so the captured cpu graph reads a stable
# address kept current at replay.
super().init_forward_metadata_capture_cpu_graph(*args, **kwargs)
self._reset_step_state()
self._refresh_cache_indices()
def conv_state_metadata(
self, layer_id: int, forward_batch: ForwardBatch
) -> ShortConvMetadata:
"""Return the conv-state handle for ``layer_id`` at the current step.
The per-step fields are already resolved on ``self.forward_metadata`` /
``self._*`` (in ``init_forward_metadata`` / ``_out_graph``);
``forward_batch`` is accepted for interface parity with the unit-test
mock and is not otherwise required here.
"""
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id)
md = self.forward_metadata
# Slot indices are cached ONCE per step in init_forward_metadata /
# init_forward_metadata_out_graph (int64). Hand back the cached view -- no
# per-layer recompute. Decode is cuda-graph-safe because that view is a
# persistent buffer refilled in place before each replay.
return ShortConvMetadata(
layer_cache=layer_cache,
cache_indices=self._cache_indices,
query_start_loc=md.query_start_loc,
has_initial_state=self._has_initial_state,
slot_ids_cpu=self._slot_ids_cpu,
has_prefix_cpu=self._has_prefix_cpu,
)
# The short-conv layers are invoked via conv_state_metadata + the model's own
# conv kernel, never through the HybridLinearAttnBackend full-vs-linear
# dispatch. Mirror Mamba2AttnBackend and guard the routed entrypoints.
def forward_decode(self, *args, **kwargs):
raise NotImplementedError(
"ShortConvAttnBackend is invoked via conv_state_metadata; "
"it does not run through forward_decode."
)
def forward_extend(self, *args, **kwargs):
raise NotImplementedError(
"ShortConvAttnBackend is invoked via conv_state_metadata; "
"it does not run through forward_extend."
)
python/sglang/srt/layers/attention/mamba/causal_conv1d.py
+4
View File
@@ -96,6 +96,8 @@ def causal_conv1d_fn(
x = x.contiguous()
bias = bias.contiguous() if bias is not None else None
if cache_indices is not None and cache_indices.dtype != torch.int32:
cache_indices = cache_indices.to(torch.int32)
causal_conv1d_fwd(
x,
weight,
@@ -162,6 +164,8 @@ def causal_conv1d_update(
unsqueeze = x.dim() == 2
if unsqueeze:
x = x.unsqueeze(-1)
if conv_state_indices is not None and conv_state_indices.dtype != torch.int32:
conv_state_indices = conv_state_indices.to(torch.int32)
causal_conv1d_update_kernel(
x,
conv_state,
python/sglang/srt/models/lfm2.py
+10 -31
View File
@@ -40,7 +40,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_req_to_token_pool
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
sharded_weight_loader,
@@ -265,10 +265,11 @@ class Lfm2ShortConv(nn.Module):
if forward_batch.forward_mode.is_idle():
return hidden_states
layer_cache = get_req_to_token_pool().mamba2_layer_cache(self.layer_idx)
conv_state = layer_cache.conv[0]
req_pool_indices = forward_batch.req_pool_indices
mamba_indices = get_req_to_token_pool().get_mamba_indices(req_pool_indices)
# The backend owns the per-request conv-state plumbing (slot indices,
# prefix mask, cu-seqlens, cuda-graph buffers); this layer just runs its
# depthwise conv against the returned handle.
meta = get_attn_backend().conv_state_metadata(self.layer_idx, forward_batch)
conv_state = meta.layer_cache.conv[0]
# Project and split into gates: B (pre-conv), C (post-conv), x (input)
proj, _ = self.in_proj(hidden_states)
@@ -283,40 +284,18 @@ class Lfm2ShortConv(nn.Module):
self.conv_weight,
self.conv_bias,
activation=None,
conv_state_indices=mamba_indices.to(torch.int32),
conv_state_indices=meta.cache_indices,
)
else:
# Prefill: multiple tokens, use varlen kernel
T = hidden_states.shape[0]
Bx_t = Bx.transpose(0, 1).contiguous()
# Build query_start_loc: [0, cumsum(seq_lens)...]
extend_start_loc = forward_batch.extend_start_loc
if extend_start_loc is not None and len(extend_start_loc) > 1:
query_start_loc = torch.cat(
[
extend_start_loc,
torch.tensor(
[T], dtype=torch.int32, device=hidden_states.device
),
]
)
cache_indices = mamba_indices.to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens > 0
else:
query_start_loc = torch.tensor(
[0, T], dtype=torch.int32, device=hidden_states.device
)
cache_indices = mamba_indices[:1].to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens[:1] > 0
conv_out = causal_conv1d_fn(
Bx_t,
self.conv_weight,
self.conv_bias,
query_start_loc=query_start_loc,
cache_indices=cache_indices,
has_initial_state=has_initial_state,
query_start_loc=meta.query_start_loc,
cache_indices=meta.cache_indices,
has_initial_state=meta.has_initial_state,
conv_states=conv_state,
activation=None,
).transpose(0, 1)
python/sglang/srt/models/lfm2_moe.py
+10 -28
View File
@@ -42,7 +42,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_req_to_token_pool
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
sharded_weight_loader,
@@ -328,10 +328,11 @@ class Lfm2MoeShortConv(nn.Module):
if forward_batch.forward_mode.is_idle():
return hidden_states
layer_cache = get_req_to_token_pool().mamba2_layer_cache(self.layer_idx)
conv_state = layer_cache.conv[0]
req_pool_indices = forward_batch.req_pool_indices
mamba_indices = get_req_to_token_pool().get_mamba_indices(req_pool_indices)
# The backend owns the per-request conv-state plumbing (slot indices,
# prefix mask, cu-seqlens, cuda-graph buffers); this layer just runs its
# depthwise conv against the returned handle.
meta = get_attn_backend().conv_state_metadata(self.layer_idx, forward_batch)
conv_state = meta.layer_cache.conv[0]
proj, _ = self.in_proj(hidden_states)
B_gate, C_gate, x = proj.chunk(3, dim=-1)
@@ -344,36 +345,17 @@ class Lfm2MoeShortConv(nn.Module):
self.conv_weight,
self.conv_bias,
activation=None,
conv_state_indices=mamba_indices.to(torch.int32),
conv_state_indices=meta.cache_indices,
)
else:
T = hidden_states.shape[0]
Bx_t = Bx.transpose(0, 1).contiguous()
# Build query_start_loc for variable-length sequences
# causal_conv1d_fn expects [start0, start1, ..., startN, T]
extend_start_loc = forward_batch.extend_start_loc
if extend_start_loc is not None and len(extend_start_loc) > 1:
# Multiple sequences: append T to extend_start_loc
# Allocate and fill to avoid torch.cat overhead
query_start_loc = extend_start_loc.new_empty(len(extend_start_loc) + 1)
query_start_loc[:-1] = extend_start_loc
query_start_loc[-1] = T
cache_indices = mamba_indices.to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens > 0
else:
# Single sequence: [0, T]
query_start_loc = hidden_states.new_tensor([0, T], dtype=torch.int32)
cache_indices = mamba_indices[:1].to(torch.int32)
has_initial_state = forward_batch.extend_prefix_lens[:1] > 0
conv_out = causal_conv1d_fn(
Bx_t,
self.conv_weight,
self.conv_bias,
query_start_loc=query_start_loc,
cache_indices=cache_indices,
has_initial_state=has_initial_state,
query_start_loc=meta.query_start_loc,
cache_indices=meta.cache_indices,
has_initial_state=meta.has_initial_state,
conv_states=conv_state,
activation=None,
).transpose(0, 1)
python/sglang/srt/models/zaya.py
+256 -257
View File
@@ -31,9 +31,13 @@ for the full design notes):
averaging across MoE layers) and MOD (mixture-of-depths skip expert).
- Per-layer :class:`ResidualScaling` keeps the residual stream in fp32 with
affine scale/bias both on the residual and on the post-mixer hidden states.
- Per-request CCA state (``conv_state`` + ``prev_hs``) is managed by
SGLang's centralized ``MambaPool`` inside ``HybridReqToTokenPool``,
accessed via ``get_req_to_token_pool().mamba2_layer_cache()``.
- Per-request CCA state (``conv_state`` + ``prev_hs``) lives in SGLang's
centralized ``MambaPool`` inside ``HybridReqToTokenPool``. The per-request
state plumbing (slot indices, prefix mask, cuda-graph buffers) is owned by
``ShortConvAttnBackend`` and reached via
``get_attn_backend().conv_state_metadata()``, so the model holds no pool
access; CCA runs its own conv (:func:`cca_extend` / :func:`cca_decode`)
against the returned handle.
"""
from __future__ import annotations
@@ -41,7 +45,7 @@ from __future__ import annotations
import logging
import re
from collections.abc import Iterable
from typing import Optional
from typing import List, Optional, Tuple
import torch
import torch.nn.functional as F
@@ -70,7 +74,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.forward_context import get_req_to_token_pool
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs
@@ -78,14 +82,6 @@ from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs
logger = logging.getLogger(__name__)
# Attribute names used to memoize the per-request MambaPool slot indices on the
# ForwardBatch. The req -> slot mapping is identical for every CCA layer in a
# step, so caching it here makes the lookup (and its GPU->CPU sync) run once per
# forward step instead of once per attention layer.
_MAMBA_INDICES_ATTR = "_zaya_mamba_indices"
_MAMBA_INDICES_CPU_ATTR = "_zaya_mamba_indices_cpu"
# ---------------------------------------------------------------------------
# Residual scaling
# ---------------------------------------------------------------------------
@@ -141,6 +137,187 @@ def _apply_norm_with_fp32_residual(
return norm(residual.to(target_dtype))
# ---------------------------------------------------------------------------
# CCA conv-state kernels (v1 torch)
#
# ZAYA1-specific conv step: the CCA conv is a causal two-stage conv over
# ``qk = [W_q hs || W_k hs]`` plus a one-token ``prev_hs`` lag for val_proj2.
# The per-request conv state lives in the centralized MambaPool; the backend
# (ShortConvAttnBackend) hands out the slot indices + prefix flags and CCA runs
# these functions against them. ``conv_qk`` is the module's two-stage conv;
# both functions mutate ``conv_state`` / ``prev_hs_state`` in place and return
# ``(qk_out, v2_input)`` -- the conv output ``[T, in_out_ch]`` and the (shifted)
# ``val_proj2`` input ``[T, hidden_size]``.
# ---------------------------------------------------------------------------
def cca_extend(
qk: torch.Tensor,
hidden_states: torch.Tensor,
conv_qk: nn.Module,
conv_state: torch.Tensor,
prev_hs_state: torch.Tensor,
slot_ids: List[int],
has_prefix: List[bool],
extend_seq_lens_cpu: List[int],
total_padding: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Prefill / extend conv-state step (v1, pure torch).
Walks each request in the batch, applies ``conv_qk`` with the request's own
initial state (zeros on a fresh first chunk, the cached ``conv_state`` slot
otherwise), writes the updated ``conv_state`` / ``prev_hs_state`` back, and
returns the concatenated ``(qk_out, v2_input)`` in the original token layout.
``slot_ids`` is the host mirror of the per-request MambaPool slot indices and
``has_prefix[i]`` is ``True`` when request ``i`` resumes a cached prefix.
The Triton swap (:func:`cca_conv1d_fn`) removes this per-request loop.
"""
dtype = hidden_states.dtype
if total_padding is None:
total_padding = conv_state.shape[-1]
in_out_ch = qk.shape[-1]
hidden_size = hidden_states.shape[-1]
qk_out = torch.empty_like(qk)
v2_input = torch.empty_like(hidden_states)
# Fresh-prefill fast path: when no request has a cached prefix the per-request
# convs can be coalesced into a single packed convolution. Each request's
# segment is laid out as ``[total_padding zeros, S_i tokens]``.
all_fresh = bool(extend_seq_lens_cpu) and not any(has_prefix)
if all_fresh:
seq_lens = [int(s) for s in extend_seq_lens_cpu]
pad = total_padding
offsets_in = [0]
for s in seq_lens:
offsets_in.append(offsets_in[-1] + s + pad)
packed = qk.new_zeros((1, in_out_ch, offsets_in[-1]))
start = 0
for i, s in enumerate(seq_lens):
end = start + s
packed[0, :, offsets_in[i] + pad : offsets_in[i + 1]] = qk[
start:end
].transpose(0, 1)
start = end
packed_out = conv_qk(packed) # [1, C, offsets_in[-1] - pad]
start = 0
for i, s in enumerate(seq_lens):
end = start + s
a_i = offsets_in[i]
qk_out[start:end] = packed_out[0, :, a_i : a_i + s].transpose(0, 1)
new_state = packed[0, :, a_i + s : a_i + s + pad]
conv_state[slot_ids[i]] = new_state.to(conv_state.dtype)
hs_cur = hidden_states[start:end]
first = hidden_states.new_zeros((1, hidden_size))
v2_input[start:end] = torch.cat([first, hs_cur[:-1]], dim=0)
prev_hs_state[slot_ids[i]] = (
hs_cur[-1].unsqueeze(-1).to(prev_hs_state.dtype)
)
start = end
else:
start = 0
for i, seq_len in enumerate(extend_seq_lens_cpu):
end = start + int(seq_len)
slot = slot_ids[i]
prefix = bool(has_prefix[i])
qk_cur = qk[start:end].transpose(0, 1).unsqueeze(0) # [1, C, S_cur]
if prefix:
left_pad = conv_state[slot].unsqueeze(0).to(dtype)
else:
left_pad = qk_cur.new_zeros((1, in_out_ch, total_padding))
padded = torch.cat([left_pad, qk_cur], dim=-1)
out = conv_qk(padded) # [1, C, S_cur]
qk_out[start:end] = out.squeeze(0).transpose(0, 1)
new_state = padded[..., -total_padding:]
conv_state[slot] = new_state.squeeze(0).to(conv_state.dtype)
hs_cur = hidden_states[start:end]
if prefix:
first = prev_hs_state[slot].squeeze(-1).to(dtype).unsqueeze(0)
else:
first = hidden_states.new_zeros((1, hidden_size))
v2_input[start:end] = torch.cat([first, hs_cur[:-1]], dim=0)
prev_hs_state[slot] = hs_cur[-1].unsqueeze(-1).to(prev_hs_state.dtype)
start = end
return qk_out, v2_input
def cca_decode(
qk: torch.Tensor,
hidden_states: torch.Tensor,
conv_qk: nn.Module,
conv_state: torch.Tensor,
prev_hs_state: torch.Tensor,
mamba_indices: torch.Tensor,
total_padding: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Single-token decode conv-state step (v1, pure torch).
Gathers each request's cached ``conv_state`` / ``prev_hs_state`` via
``index_select``, runs ``conv_qk`` on the ``[T, C, total_padding + 1]``
window, and scatters the updated state back with ``index_copy_``. All ops are
on-device (``mamba_indices`` is a device ``long`` tensor), so this stays
CUDA-graph capturable. Returns ``(qk_out, prev_hs)`` where ``prev_hs`` is the
previous hidden state feeding ``val_proj2``.
The Triton swap is :func:`cca_conv1d_update`.
"""
dtype = hidden_states.dtype
if total_padding is None:
total_padding = conv_state.shape[-1]
left_pad = conv_state.index_select(0, mamba_indices).to(dtype)
cur = qk.unsqueeze(-1) # [T, C, 1]
padded = torch.cat([left_pad, cur], dim=-1) # [T, C, total_padding + 1]
out = conv_qk(padded) # [T, C, 1]
qk_out = out.squeeze(-1) # [T, C]
new_state = padded[..., -total_padding:]
conv_state.index_copy_(0, mamba_indices, new_state.to(conv_state.dtype))
# Read the previous hidden state (val_proj2 input) BEFORE overwriting the
# slot with the current token.
prev_hs = prev_hs_state.index_select(0, mamba_indices).squeeze(-1).to(dtype)
prev_hs_state.index_copy_(
0, mamba_indices, hidden_states.unsqueeze(-1).to(prev_hs_state.dtype)
)
return qk_out, prev_hs
# Fused kernel seam (TODO) -- perf swap for the v1 torch paths above. These
# mirror the ``causal_conv1d_fn`` / ``causal_conv1d_update`` contract but for
# CCA's two-stage *grouped* conv (conv_qk[0] depthwise + conv_qk[1] grouped
# per-head), which the stock depthwise ``causal_conv1d`` cannot express. Once
# implemented they replace the per-request loop in ``cca_extend`` and the
# separate gather/conv/scatter launches in ``cca_decode`` with a single
# index-driven kernel. Same ``(qk_out, v2_input)`` return contract.
def cca_conv1d_fn(*args, **kwargs):
raise NotImplementedError(
"Fused CCA prefill conv-with-state kernel not implemented yet; "
"the model uses cca_extend (v1 torch) in the meantime."
)
def cca_conv1d_update(*args, **kwargs):
raise NotImplementedError(
"Fused CCA decode conv-with-state kernel not implemented yet; "
"the model uses cca_decode (v1 torch) in the meantime."
)
# ---------------------------------------------------------------------------
# CCA: Compressed Convolutional Attention QKV projection
# ---------------------------------------------------------------------------
@@ -401,58 +578,6 @@ class CCA(nn.Module):
# ----- helpers ---------------------------------------------------------
@staticmethod
def _get_mamba_indices(forward_batch: ForwardBatch) -> torch.Tensor:
"""Per-request MambaPool slot indices as an int64 device tensor.
The req -> slot mapping depends only on ``forward_batch.req_pool_indices``,
which is constant for every CCA layer within one forward step. Computing
it inside each of the ~60 attention layers would issue one redundant
gather per layer, so it is computed once and memoized on the ForwardBatch
(whose lifetime is exactly one forward step). The lookup is pure on-device
work, so this stays compatible with CUDA graph capture on the decode path.
"""
cached = getattr(forward_batch, _MAMBA_INDICES_ATTR, None)
if cached is None:
cached = (
get_req_to_token_pool()
.get_mamba_indices(forward_batch.req_pool_indices)
.to(torch.long)
)
setattr(forward_batch, _MAMBA_INDICES_ATTR, cached)
return cached
@staticmethod
def _get_mamba_indices_cpu(
forward_batch: ForwardBatch, mamba_indices: torch.Tensor
) -> list[int]:
"""Host mirror of :meth:`_get_mamba_indices`, memoized per forward step.
Only the extend/prefill path needs the indices on the host to drive its
per-request Python loop; the decode path indexes the pool entirely
on-device. Memoizing turns the previous one-``.tolist()``-sync-per-layer
behavior into a single GPU->CPU sync per forward step. This helper is
never reached on the decode path that CUDA graphs capture.
"""
cached = getattr(forward_batch, _MAMBA_INDICES_CPU_ATTR, None)
if cached is None:
cached = mamba_indices.tolist()
setattr(forward_batch, _MAMBA_INDICES_CPU_ATTR, cached)
return cached
def _get_pool_state(self, forward_batch: ForwardBatch):
"""Retrieve per-request CCA state from the centralized MambaPool.
``conv_state`` / ``prev_hs_state`` are layer-local pool views, but the
``mamba_indices`` req -> slot mapping is shared across layers and so is
memoized on the ForwardBatch (see :meth:`_get_mamba_indices`).
"""
layer_cache = get_req_to_token_pool().mamba2_layer_cache(self.layer_id)
conv_state = layer_cache.conv[0]
prev_hs_state = layer_cache.conv[1]
mamba_indices = self._get_mamba_indices(forward_batch)
return conv_state, prev_hs_state, mamba_indices
def _normalize_qk(
self, query: torch.Tensor, key: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -583,195 +708,6 @@ class CCA(nn.Module):
value = self._slice_v_per_rank(value_full)
return query, key, value
def _forward_extend(
self,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Prefill / extend path.
Walks every request in the batch, applies the conv with each request's
own initial state (zero on first chunk, cached otherwise), writes the
updated state and ``prev_hs`` back into the centralized MambaPool, and
returns the concatenated q/k/v in the original token layout.
"""
dtype = hidden_states.dtype
T = hidden_states.shape[0]
q_raw, _ = self.linear_q(hidden_states) # [T, latent_q]
k_raw, _ = self.linear_k(hidden_states)
qk = torch.cat([q_raw, k_raw], dim=-1) # [T, in_out_ch]
query_pre = q_raw.view(T, self.num_q_heads, self.head_dim)
key_base = k_raw.view(T, self.num_k_heads, self.head_dim)
qk_out = torch.empty_like(qk)
v2_input = torch.empty_like(hidden_states)
conv_state, prev_hs_state, mamba_indices = self._get_pool_state(forward_batch)
# Host view of the slot indices to drive the per-request loop below.
# Memoized on the ForwardBatch, so the GPU->CPU sync runs once per forward
# step rather than once per attention layer (~60 syncs/step otherwise).
mamba_idx_cpu = self._get_mamba_indices_cpu(forward_batch, mamba_indices)
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
extend_prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu
# Fresh-prefill fast path: when no request has a cached prefix the
# per-request convs (one launch each, ×60 attention layers ×B requests)
# can be coalesced into a single packed convolution. The conv chain is
# two ``kernel_size=2`` convs (effective receptive field = 3), so each
# request's S valid outputs are produced from the packed positions
# ``[a_i, a_i + S_i - 1]`` where the input segment for request i is
# ``[pad, pad, x_0, ..., x_{S-1}]`` of length ``S_i + total_padding``.
all_fresh = bool(extend_seq_lens_cpu) and not any(
int(p) > 0 for p in extend_prefix_lens_cpu
)
if all_fresh:
seq_lens = [int(s) for s in extend_seq_lens_cpu]
pad = self.total_padding
# Build packed buffer: per request -> [pad zeros, S_i tokens].
offsets_in = [0]
for s in seq_lens:
offsets_in.append(offsets_in[-1] + s + pad)
packed = qk.new_zeros((1, self.in_out_ch, offsets_in[-1]))
start = 0
for i, s in enumerate(seq_lens):
end = start + s
packed[0, :, offsets_in[i] + pad : offsets_in[i + 1]] = qk[
start:end
].transpose(0, 1)
start = end
packed_out = self._conv_qk_run(packed) # [1, C, offsets_in[-1] - pad]
start = 0
for i, s in enumerate(seq_lens):
end = start + s
a_i = offsets_in[i]
qk_out[start:end] = packed_out[0, :, a_i : a_i + s].transpose(0, 1)
new_state = packed[0, :, a_i + s : a_i + s + pad]
conv_state[mamba_idx_cpu[i]] = new_state.to(conv_state.dtype)
hs_cur = hidden_states[start:end]
first = hidden_states.new_zeros((1, self.hidden_size))
v2_input[start:end] = torch.cat([first, hs_cur[:-1]], dim=0)
prev_hs_state[mamba_idx_cpu[i]] = (
hs_cur[-1].unsqueeze(-1).to(prev_hs_state.dtype)
)
start = end
else:
start = 0
for i, seq_len in enumerate(extend_seq_lens_cpu):
end = start + int(seq_len)
mamba_idx = mamba_idx_cpu[i]
has_prefix = int(extend_prefix_lens_cpu[i]) > 0
qk_cur = qk[start:end].transpose(0, 1).unsqueeze(0) # [1, C, S_cur]
if has_prefix:
left_pad = conv_state[mamba_idx].unsqueeze(0).to(dtype)
else:
left_pad = qk_cur.new_zeros((1, self.in_out_ch, self.total_padding))
padded = torch.cat([left_pad, qk_cur], dim=-1)
out = self._conv_qk_run(padded) # [1, C, S_cur]
qk_out[start:end] = out.squeeze(0).transpose(0, 1)
new_state = padded[..., -self.total_padding :]
conv_state[mamba_idx] = new_state.squeeze(0).to(conv_state.dtype)
hs_cur = hidden_states[start:end]
if has_prefix:
first = prev_hs_state[mamba_idx].squeeze(-1).to(dtype).unsqueeze(0)
else:
first = hidden_states.new_zeros((1, self.hidden_size))
shifted = torch.cat([first, hs_cur[:-1]], dim=0)
v2_input[start:end] = shifted
prev_hs_state[mamba_idx] = (
hs_cur[-1].unsqueeze(-1).to(prev_hs_state.dtype)
)
start = end
query_conv = qk_out[:, : self.latent_q_dim].view(
T, self.num_q_heads, self.head_dim
)
key_conv = qk_out[:, self.latent_q_dim :].view(
T, self.num_k_heads, self.head_dim
)
query, key = self._add_grouped_qk_means(
query_conv, key_conv, query_pre, key_base
)
query, key = self._normalize_qk(query, key)
v1, _ = self.val_proj1(hidden_states)
v2, _ = self.val_proj2(v2_input)
value_full = torch.cat([v1, v2], dim=-1).view(
T, self.num_k_heads_full, self.head_dim
)
value = self._slice_v_per_rank(value_full)
return query, key, value
def _forward_decode(
self,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Single-token decode path for a whole batch.
Reads each request's cached conv state and ``prev_hs`` from the
centralized MambaPool via ``index_select``, runs the conv on the
small ``[T, C, total_padding+1]`` window, and writes back via
``index_copy_``.
"""
T = hidden_states.shape[0]
dtype = hidden_states.dtype
conv_state, prev_hs_state, mamba_indices = self._get_pool_state(forward_batch)
q_raw, _ = self.linear_q(hidden_states)
k_raw, _ = self.linear_k(hidden_states)
qk = torch.cat([q_raw, k_raw], dim=-1) # [T, C]
query_pre = q_raw.view(T, self.num_q_heads, self.head_dim)
key_base = k_raw.view(T, self.num_k_heads, self.head_dim)
left_pad = conv_state.index_select(0, mamba_indices).to(dtype)
cur = qk.unsqueeze(-1) # [T, C, 1]
padded = torch.cat([left_pad, cur], dim=-1) # [T, C, total_padding+1]
out = self._conv_qk_run(padded) # [T, C, 1]
qk_out = out.squeeze(-1) # [T, C]
new_state = padded[..., -self.total_padding :]
conv_state.index_copy_(0, mamba_indices, new_state.to(conv_state.dtype))
query_conv = qk_out[:, : self.latent_q_dim].view(
T, self.num_q_heads, self.head_dim
)
key_conv = qk_out[:, self.latent_q_dim :].view(
T, self.num_k_heads, self.head_dim
)
query, key = self._add_grouped_qk_means(
query_conv, key_conv, query_pre, key_base
)
query, key = self._normalize_qk(query, key)
prev_hs = prev_hs_state.index_select(0, mamba_indices).squeeze(-1).to(dtype)
v1, _ = self.val_proj1(hidden_states)
v2, _ = self.val_proj2(prev_hs)
value_full = torch.cat([v1, v2], dim=-1).view(
T, self.num_k_heads_full, self.head_dim
)
value = self._slice_v_per_rank(value_full)
prev_hs_state.index_copy_(
0, mamba_indices, hidden_states.unsqueeze(-1).to(prev_hs_state.dtype)
)
return query, key, value
def forward(
self,
hidden_states: torch.Tensor,
@@ -779,6 +715,16 @@ class CCA(nn.Module):
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Project ``hidden_states`` into ``(q, k, v)`` honoring per-request state.
The per-request conv-state plumbing (slot gather/scatter, prefix mask,
cuda-graph buffers) is owned by :class:`ShortConvAttnBackend
<sglang.srt.layers.attention.linear.short_conv_backend.ShortConvAttnBackend>`,
reached via ``get_attn_backend().conv_state_metadata``; CCA runs its own
two-stage grouped conv (:func:`cca_extend` / :func:`cca_decode`) against
that handle, so this module holds no pool access. Those functions return
the conv output ``qk_out`` and the ``val_proj2`` input ``v2_input`` (the
shifted / previous hidden state), updating the ``conv_state`` /
``prev_hs`` pool slots in place.
``q`` / ``k`` are returned in fp32 (the normalize step keeps fp32 for
stability); ``v`` is returned in the input dtype since the caller
casts everything back to ``hidden_states.dtype`` before rotary +
@@ -798,10 +744,63 @@ class CCA(nn.Module):
zero.view(0, self.num_k_heads, self.head_dim),
)
T = hidden_states.shape[0]
q_raw, _ = self.linear_q(hidden_states) # [T, latent_q]
k_raw, _ = self.linear_k(hidden_states)
qk = torch.cat([q_raw, k_raw], dim=-1) # [T, in_out_ch]
query_pre = q_raw.view(T, self.num_q_heads, self.head_dim)
key_base = k_raw.view(T, self.num_k_heads, self.head_dim)
# The backend hands out the per-request conv-state handle (slot indices,
# prefix mask, cuda-graph buffers); CCA runs its own two-stage grouped
# conv against it and gets back the conv output + val_proj2 input, with
# the conv_state / prev_hs pool slots updated in place.
meta = get_attn_backend().conv_state_metadata(self.layer_id, forward_batch)
conv_state = meta.layer_cache.conv[0]
prev_hs_state = meta.layer_cache.conv[1]
if forward_batch.forward_mode.is_decode_or_idle():
return self._forward_decode(hidden_states, forward_batch)
# EXTEND / MIXED / DLLM_EXTEND all share the prefill loop.
return self._forward_extend(hidden_states, forward_batch)
qk_out, v2_input = cca_decode(
qk,
hidden_states,
self.conv_qk,
conv_state,
prev_hs_state,
meta.cache_indices,
self.total_padding,
)
else:
qk_out, v2_input = cca_extend(
qk,
hidden_states,
self.conv_qk,
conv_state,
prev_hs_state,
meta.slot_ids_cpu,
meta.has_prefix_cpu,
forward_batch.extend_seq_lens_cpu,
self.total_padding,
)
query_conv = qk_out[:, : self.latent_q_dim].view(
T, self.num_q_heads, self.head_dim
)
key_conv = qk_out[:, self.latent_q_dim :].view(
T, self.num_k_heads, self.head_dim
)
query, key = self._add_grouped_qk_means(
query_conv, key_conv, query_pre, key_base
)
query, key = self._normalize_qk(query, key)
v1, _ = self.val_proj1(hidden_states)
v2, _ = self.val_proj2(v2_input)
value_full = torch.cat([v1, v2], dim=-1).view(
T, self.num_k_heads_full, self.head_dim
)
value = self._slice_v_per_rank(value_full)
return query, key, value
# ---------------------------------------------------------------------------