Fuse GLM-5.3-Flash KDA projections and prefill metadata (#39688)

Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>

python/sglang/kernels/ops/attention/fla/kda.py

@@ -917,6 +917,7 @@ def softplus_fwd(x):
 @triton.heuristics(
     {
         "HAS_BIAS": lambda args: args["dt_bias"] is not None,
+        "HAS_BETA": lambda args: args["beta"] is not None,
         "HAS_SCALE": lambda args: args["scale"] is not None,
         "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
         "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
@@ -928,7 +929,7 @@ def softplus_fwd(x):
         for BS in BS_LIST
         for num_warps in [2, 4, 8]
     ],
-    key=["H", "S", "BT", "IS_VARLEN"],
+    key=["H", "S", "BT", "IS_VARLEN", "HAS_BETA"],
 )
 @triton.jit(do_not_specialize=["T"])
 def kda_gate_chunk_cumsum_vector_kernel(
@@ -940,12 +941,18 @@ def kda_gate_chunk_cumsum_vector_kernel(
     cu_seqlens,
     chunk_indices,
     lower_bound,
+    beta,
+    beta_out,
+    beta_stride_b: tl.constexpr,
+    beta_stride_t: tl.constexpr,
+    beta_stride_h: tl.constexpr,
     T,
     H: tl.constexpr,
     S: tl.constexpr,
     BT: tl.constexpr,
     BS: tl.constexpr,
     HAS_BIAS: tl.constexpr,
+    HAS_BETA: tl.constexpr,
     HAS_SCALE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
     USE_LOWER_BOUND: tl.constexpr,
@@ -1011,6 +1018,24 @@ def kda_gate_chunk_cumsum_vector_kernel(
         b_o *= scale
     tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
 
+    if HAS_BETA:
+        if i_s == 0:
+            offsets_t = i_t * BT + tl.arange(0, BT)
+            if IS_VARLEN:
+                beta_offsets = (bos + offsets_t) * beta_stride_t
+            else:
+                beta_offsets = i_b * beta_stride_b + offsets_t * beta_stride_t
+            b_beta = tl.load(
+                beta + beta_offsets + i_h * beta_stride_h,
+                mask=offsets_t < T,
+                other=0.0,
+            ).to(tl.float32)
+            tl.store(
+                beta_out + (bos + offsets_t) * H + i_h,
+                tl.sigmoid(b_beta),
+                mask=offsets_t < T,
+            )
+
 
 def kda_gate_chunk_cumsum(
     g: torch.Tensor,
@@ -1022,9 +1047,10 @@ def kda_gate_chunk_cumsum(
     output_dtype: Optional[torch.dtype] = torch.float,
     chunk_indices: Optional[torch.LongTensor] = None,
     lower_bound: Optional[float] = None,
-) -> torch.Tensor:
+    beta: Optional[torch.Tensor] = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """
-    Fused KDA gate activation + chunk-local cumulative sum.
+    Fused KDA gate activation + chunk-local cumulative sum, with optional beta.
 
     Combines two memory-bound kernels into one:
       1. Gate activation: g = -exp(A_log) * softplus(raw_g + dt_bias)
@@ -1040,9 +1066,11 @@ def kda_gate_chunk_cumsum(
         output_dtype: Output dtype (default float32).
         chunk_indices: Pre-computed chunk indices for varlen mode.
         lower_bound: If set, use safe gate: lower_bound * sigmoid(exp(A_log) * g).
+        beta: Optional raw beta of shape [B, T, H], including strided projections.
 
     Returns:
-        Cumulative-summed gated tensor of shape [B, T, H, K].
+        Cumulative-summed gated tensor of shape [B, T, H, K]. If beta is
+        supplied, also return its sigmoid in a contiguous float32 [B, T, H] tensor.
     """
     if cu_seqlens is not None:
         assert g.shape[0] == 1, (
@@ -1059,6 +1087,18 @@ def kda_gate_chunk_cumsum(
     )
 
     g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    if beta is not None:
+        assert beta.shape == (B, T, H)
+        assert beta.device == g.device
+        beta_out = torch.empty((B, T, H), dtype=torch.float32, device=beta.device)
+        beta_strides = beta.stride()
+        if cu_seqlens is not None:
+            beta_strides = (0, beta_strides[1], beta_strides[2])
+    else:
+        beta_out = None
+        beta_strides = (0, 0, 0)
+    if B * T == 0:
+        return g if beta is None else (g, beta_out)
 
     def grid(meta):
         return (cdiv(meta["S"], meta["BS"]), NT, B * H)
@@ -1072,12 +1112,17 @@ def kda_gate_chunk_cumsum(
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
         lower_bound=lower_bound,
+        beta=beta,
+        beta_out=beta_out,
+        beta_stride_b=beta_strides[0],
+        beta_stride_t=beta_strides[1],
+        beta_stride_h=beta_strides[2],
         T=T,
         H=H,
         S=S,
         BT=BT,
     )
-    return g
+    return g if beta is None else (g, beta_out)
 
 
 def chunk_kda_fwd(
@@ -1096,6 +1141,7 @@ def chunk_kda_fwd(
     output_intermediate_states: bool = False,
     track_state: Optional[torch.Tensor] = None,
     track_chunk_idx: Optional[torch.Tensor] = None,
+    beta_is_raw: bool = False,
 ):
     chunk_size = 64
     # Pre-compute chunk indices once and thread through all downstream kernels.
@@ -1118,9 +1164,14 @@ def chunk_kda_fwd(
             cu_seqlens=cu_seqlens,
             chunk_indices=chunk_indices,
             lower_bound=lower_bound,
+            beta=beta if beta_is_raw else None,
         )
+        if beta_is_raw:
+            g, beta = g
     else:
         # g is already gate-activated by caller; just do cumsum.
+        if beta_is_raw:
+            beta = beta.float().sigmoid().contiguous()
         g = chunk_local_cumsum(
             g,
             chunk_size=chunk_size,
@@ -1226,16 +1277,13 @@ def chunk_kda(
         q = l2norm_fwd(q.contiguous())
         k = l2norm_fwd(k.contiguous())
 
-    if beta_is_raw:
-        beta = beta.float().sigmoid()
-
     # Returns o [B, T, H, V] when output_intermediate_states=False, or (o, h [B, NT, H, V, K]) when output_intermediate_states=True.
     return chunk_kda_fwd(
         q=q,
         k=k,
         v=v.contiguous(),
         g=g.contiguous(),
-        beta=beta.contiguous(),
+        beta=beta if beta_is_raw else beta.contiguous(),
         scale=scale,
         initial_state=initial_state,
         initial_state_indices=initial_state_indices,
@@ -1246,4 +1294,5 @@ def chunk_kda(
         output_intermediate_states=output_intermediate_states,
         track_state=track_state,
         track_chunk_idx=track_chunk_idx,
+        beta_is_raw=beta_is_raw,
     )

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

@@ -26,6 +26,9 @@ from sglang.srt.layers.attention.mamba.mamba2_metadata import (
     ForwardMetadata,
     Mamba2Metadata,
 )
+from sglang.srt.layers.attention.mamba.prefill_track_metadata import (
+    build_prefill_track_plan,
+)
 from sglang.srt.layers.attention.mamba.replay_state_indices_validator import (
     validate_replay_state_indices_cpu,
 )
@@ -36,6 +39,7 @@ from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.runtime_context import get_exec, get_memory, get_spec
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.spec_info import SpecInput
+from sglang.srt.utils import is_pin_memory_available
 
 if TYPE_CHECKING:
     from sglang.srt.layers.attention.verify_mask import VerifyMask
@@ -118,6 +122,22 @@ class MambaAttnBackendBase(AttentionBackend):
         state ops, incl. the cuda-graph replay-prep copy into ``state_indices_list``."""
         return self.req_to_token_pool.translate_mamba_indices(mamba_indices)
 
+    @staticmethod
+    def _has_cpu_prefill_track_metadata(forward_batch: ForwardBatch) -> bool:
+        return (
+            forward_batch.forward_mode.is_extend()
+            and not forward_batch.forward_mode.is_target_verify()
+            and all(
+                values is not None and len(values) == forward_batch.batch_size
+                for values in (
+                    forward_batch.mamba_prefill_track_mask_cpu,
+                    forward_batch.mamba_track_seqlens_cpu,
+                    forward_batch.extend_seq_lens_cpu,
+                    forward_batch.extend_prefix_lens_cpu,
+                )
+            )
+        )
+
     def _forward_metadata(self, forward_batch: ForwardBatch):
         bs = forward_batch.batch_size
 
@@ -134,6 +154,8 @@ class MambaAttnBackendBase(AttentionBackend):
         track_ssm_seq_idx = None
         track_ssm_end_locs = None
         track_ssm_recompute_dst = None
+        logical_num_tokens = None
+        track_mask_indices = None
 
         mamba_cache_indices = self.req_to_token_pool.get_mamba_indices(
             forward_batch.req_pool_indices
@@ -146,10 +168,20 @@ class MambaAttnBackendBase(AttentionBackend):
                 forward_batch.mamba_track_indices
             )
         # Resolve the tracked-row selection once per forward
-        has_mamba_track_mask = bool(
+        cpu_track_metadata = self._has_cpu_prefill_track_metadata(forward_batch)
-            forward_batch.mamba_track_mask is not None
+        if cpu_track_metadata:
-            and forward_batch.mamba_track_mask.any()
+            rows = [
-        )
+                i
+                for i, track in enumerate(forward_batch.mamba_prefill_track_mask_cpu)
+                if track
+            ]
+            has_mamba_track_mask = bool(rows)
+            track_mask_indices = self._track_indices_to_device(rows) if rows else None
+        else:
+            has_mamba_track_mask = bool(
+                forward_batch.mamba_track_mask is not None
+                and forward_batch.mamba_track_mask.any()
+            )
         _real_bs = forward_batch._original_batch_size
         if _real_bs is not None and _real_bs < mamba_cache_indices.shape[0]:
             mamba_cache_indices = mamba_cache_indices.clone()
@@ -258,9 +290,17 @@ class MambaAttnBackendBase(AttentionBackend):
                     forward_batch.extend_start_loc[-1]
                     + forward_batch.extend_seq_lens[-1]
                 )
+                if (
+                    forward_batch.extend_seq_lens_cpu is not None
+                    and len(forward_batch.extend_seq_lens_cpu) == bs
+                    and forward_batch.tbo_parent_token_range is None
+                ):
+                    logical_num_tokens = sum(forward_batch.extend_seq_lens_cpu)
+                else:
+                    logical_num_tokens = int(query_start_loc[-1])
                 if has_mamba_track_mask:
                     track_conv_indices = self._init_track_conv_indices(
-                        query_start_loc, forward_batch
+                        query_start_loc, forward_batch, track_mask_indices
                     )
 
                     (
@@ -279,6 +319,8 @@ class MambaAttnBackendBase(AttentionBackend):
 
         return ForwardMetadata(
             query_start_loc=query_start_loc,
+            logical_num_tokens=logical_num_tokens,
+            mamba_track_mask_indices=track_mask_indices,
             mamba_cache_indices=mamba_cache_indices,
             # Physical track destinations (None when tracking off); cuda-graph
             # supplies this via the static backend buffer in _replay_metadata.
@@ -347,7 +389,10 @@ class MambaAttnBackendBase(AttentionBackend):
         )
 
     def _init_track_conv_indices(
-        self, query_start_loc: torch.Tensor, forward_batch: ForwardBatch
+        self,
+        query_start_loc: torch.Tensor,
+        forward_batch: ForwardBatch,
+        track_mask_indices: Optional[torch.Tensor] = None,
     ):
         """Flattened input positions of conv states to track during extend (up to
         the last complete chunk boundary, mamba_track_mask rows only)."""
@@ -361,7 +406,11 @@ class MambaAttnBackendBase(AttentionBackend):
             "this path should only run when the track mask is set on an extend batch"
         )
         start_indices = query_start_loc[:-1] + aligned_len - conv_state_len
-        start_indices = start_indices[forward_batch.mamba_track_mask]
+        start_indices = (
+            start_indices.index_select(0, track_mask_indices)
+            if track_mask_indices is not None
+            else start_indices[forward_batch.mamba_track_mask]
+        )
 
         indices = start_indices.unsqueeze(-1) + torch.arange(
             conv_state_len,
@@ -379,6 +428,10 @@ class MambaAttnBackendBase(AttentionBackend):
         chunk boundary. Also returns ``track_ssm_h_batch_src``: the batch rows of
         the unaligned tracked seqs, used to integer-index the fp32 snapshot
         buffer on the KDA path so the copy stays free of GPU syncs."""
+        if self._has_cpu_prefill_track_metadata(forward_batch):
+            return self._init_track_ssm_indices_from_cpu(
+                mamba_cache_indices, forward_batch
+            )
         state_chunk_size = self.mamba_chunk_size
         # CPU to avoid kernel launches for the masking ops
         mamba_track_mask = forward_batch.mamba_track_mask.cpu()
@@ -454,6 +507,38 @@ class MambaAttnBackendBase(AttentionBackend):
             to_device(track_ssm_recompute_dst),
         )
 
+    def _track_indices_to_device(self, values, dtype=torch.int64):
+        return torch.tensor(
+            values, dtype=dtype, pin_memory=is_pin_memory_available(self.device)
+        ).to(self.device, non_blocking=True)
+
+    def _init_track_ssm_indices_from_cpu(self, mamba_cache_indices, forward_batch):
+        is_mamba2 = isinstance(self, Mamba2AttnBackend)
+        plan = build_prefill_track_plan(
+            forward_batch.mamba_prefill_track_mask_cpu,
+            forward_batch.mamba_track_seqlens_cpu,
+            forward_batch.extend_seq_lens_cpu,
+            forward_batch.extend_prefix_lens_cpu,
+            self.mamba_chunk_size,
+            mamba2=is_mamba2,
+        )
+        to_device = self._track_indices_to_device
+        final_rows = to_device(plan.final_rows)
+        h_rows = to_device(plan.h_rows)
+        recompute_rows = to_device(plan.recompute_rows) if is_mamba2 else None
+        destinations = forward_batch.mamba_track_indices
+        return (
+            to_device(plan.chunk_indices, torch.int32),
+            to_device(plan.h_src),
+            destinations.index_select(0, h_rows),
+            to_device(plan.unaligned_rows),
+            mamba_cache_indices.index_select(0, final_rows),
+            destinations.index_select(0, final_rows),
+            recompute_rows,
+            to_device(plan.recompute_end_locs) if is_mamba2 else None,
+            destinations.index_select(0, recompute_rows) if is_mamba2 else None,
+        )
+
     def init_forward_metadata_capture_cpu_graph(
         self,
         bs: int,

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

@@ -540,9 +540,10 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 raise ValueError("GDN MIS metadata requires --enable-mis")
             self.mis_metadata = build_gdn_mis_metadata(forward_batch)
         if self.forward_metadata.has_mamba_track_mask:
-            self.forward_metadata.mamba_track_mask_indices = (
+            if getattr(self.forward_metadata, "mamba_track_mask_indices", None) is None:
-                forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
+                self.forward_metadata.mamba_track_mask_indices = (
-            )
+                    forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
+                )
             self.forward_metadata.conv_states_mask_indices = (
                 forward_batch.mamba_track_indices[
                     self.forward_metadata.mamba_track_mask_indices

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

@@ -533,9 +533,10 @@ class KDAAttnBackend(MambaAttnBackendBase):
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         super().init_forward_metadata(forward_batch)
         if self.forward_metadata.has_mamba_track_mask:
-            self.forward_metadata.mamba_track_mask_indices = (
+            if self.forward_metadata.mamba_track_mask_indices is None:
-                forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
+                self.forward_metadata.mamba_track_mask_indices = (
-            )
+                    forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
+                )
             self.forward_metadata.conv_states_mask_indices = (
                 forward_batch.mamba_track_indices[
                     self.forward_metadata.mamba_track_mask_indices
@@ -822,7 +823,9 @@ class KDAAttnBackend(MambaAttnBackendBase):
         has_initial_state = forward_batch.extend_prefix_lens > 0
 
         physical_num_tokens = mixed_qkv.shape[0]
-        logical_num_tokens = int(query_start_loc[-1])
+        logical_num_tokens = self.forward_metadata.logical_num_tokens
+        if logical_num_tokens is None:
+            logical_num_tokens = int(query_start_loc[-1])
         if logical_num_tokens < physical_num_tokens:
             mixed_qkv = mixed_qkv[:logical_num_tokens]
             a = a[:, :logical_num_tokens]

python/sglang/srt/layers/attention/mamba/mamba2_metadata.py

@@ -29,6 +29,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 class ForwardMetadata:
     query_start_loc: torch.Tensor
     mamba_cache_indices: torch.Tensor
+    logical_num_tokens: Optional[int] = None
     mamba_cache_indices_gdn: Optional[torch.Tensor] = None
     # Mamba track DESTINATION slots (PHYSICAL, length == batch). Like
     # mamba_cache_indices: a backend-owned static buffer under cuda-graph (translated

python/sglang/srt/layers/attention/mamba/prefill_track_metadata.py

@@ -0,0 +1,57 @@
+"""Host-side row selection for prefill state snapshots.
+
+Slot IDs deliberately stay on the device: unified pools translate virtual IDs
+before these row indices gather physical source and destination slots.
+"""
+
+from dataclasses import dataclass
+from itertools import accumulate
+
+
+@dataclass
+class PrefillTrackPlan:
+    tracked_rows: list[int]
+    final_rows: list[int]
+    unaligned_rows: list[int]
+    h_rows: list[int]
+    h_src: list[int]
+    recompute_rows: list[int]
+    recompute_end_locs: list[int]
+    chunk_indices: list[int]
+
+
+def build_prefill_track_plan(
+    mask: list[bool],
+    track_lens: list[int],
+    extend_lens: list[int],
+    prefix_lens: list[int],
+    chunk_size: int,
+    *,
+    mamba2: bool,
+) -> PrefillTrackPlan:
+    """Use the backend's actual chunk size, including Mamba2's flat grid."""
+    assert len(mask) == len(track_lens) == len(extend_lens) == len(prefix_lens)
+    starts = list(accumulate(extend_lens, initial=0))
+    h_offsets = list(
+        accumulate(((n + chunk_size - 1) // chunk_size for n in extend_lens), initial=0)
+    )
+    plan = PrefillTrackPlan([], [], [], [], [], [], [], [-1] * len(mask))
+    for row, track in enumerate(mask):
+        if not track:
+            continue
+        plan.tracked_rows.append(row)
+        length = track_lens[row] - prefix_lens[row]
+        if length % chunk_size == 0:
+            plan.final_rows.append(row)
+            continue
+        chunk = length // chunk_size
+        plan.unaligned_rows.append(row)
+        plan.chunk_indices[row] = chunk
+        end = starts[row] + chunk * chunk_size
+        if mamba2 and end % chunk_size:
+            plan.recompute_rows.append(row)
+            plan.recompute_end_locs.append(end)
+        else:
+            plan.h_rows.append(row)
+            plan.h_src.append(end // chunk_size if mamba2 else h_offsets[row] + chunk)
+    return plan

python/sglang/srt/managers/schedule_batch.py

@@ -2411,6 +2411,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     mamba_track_buffer_indices: Optional[List[int]] = None  # shape: [b], 0 or 1
     mamba_track_mask: torch.Tensor = None  # shape: [b], bool
     mamba_track_seqlens: torch.Tensor = None  # shape: [b], int64
+    # TBO rejects Mamba tracking; enabling it must also slice these CPU lists.
+    mamba_track_seqlens_cpu: Optional[List[int]] = None
+    mamba_prefill_track_mask_cpu: Optional[List[bool]] = None
     mamba_track_mask_cpu: Optional[List[bool]] = None  # shape: [b]
     mamba_track_mask_next_cpu: Optional[List[bool]] = None  # shape: [b]
     mamba_decode_batch_idx_cpu: Optional[List[int]] = None  # shape: [b]
@@ -2934,6 +2937,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
 
         if get_exec().mamba.enable_mamba_extra_buffer:
+            self.mamba_prefill_track_mask_cpu = mamba_track_mask_cpu
+            self.mamba_track_seqlens_cpu = mamba_track_seqlens_cpu
             self.mamba_track_indices = torch.tensor(
                 mamba_track_indices_cpu,
                 dtype=torch.int64,
@@ -3500,6 +3505,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     def prepare_for_decode(self):
         self.forward_mode = ForwardMode.DECODE
+        self.mamba_track_seqlens_cpu = None
+        self.mamba_prefill_track_mask_cpu = None
         # Decode embeds the last output token via embed_tokens; clear the stale
         # prefill-time tensor so it doesn't leak into ForwardBatch.
         self.input_embeds = None
@@ -3653,6 +3660,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         self.mamba_track_buffer_indices = None
         self.mamba_track_mask = None
         self.mamba_track_seqlens = None
+        self.mamba_track_seqlens_cpu = None
+        self.mamba_prefill_track_mask_cpu = None
         self.mamba_track_mask_cpu = None
         self.mamba_track_mask_next_cpu = None
         self.mamba_decode_batch_idx_cpu = None
@@ -3719,6 +3728,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         self.mamba_track_buffer_indices = None
         self.mamba_track_mask = None
         self.mamba_track_seqlens = None
+        self.mamba_track_seqlens_cpu = None
+        self.mamba_prefill_track_mask_cpu = None
         self.mamba_track_mask_cpu = None
         self.mamba_track_mask_next_cpu = None
         self.mamba_decode_batch_idx_cpu = None
@@ -3782,6 +3793,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             mamba_track_buffer_indices=self.mamba_track_buffer_indices,
             mamba_track_mask=self.mamba_track_mask,
             mamba_track_seqlens=self.mamba_track_seqlens,
+            mamba_track_seqlens_cpu=self.mamba_track_seqlens_cpu,
+            mamba_prefill_track_mask_cpu=self.mamba_prefill_track_mask_cpu,
             mamba_track_mask_cpu=self.mamba_track_mask_cpu,
             mamba_track_mask_next_cpu=self.mamba_track_mask_next_cpu,
             mamba_decode_batch_idx_cpu=self.mamba_decode_batch_idx_cpu,

python/sglang/srt/model_executor/forward_batch_info.py

@@ -503,6 +503,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     mamba_track_mask: Optional[torch.Tensor] = None  # shape: [b], bool
     # The seqlens to track mamba state if masked, prefill only.
     mamba_track_seqlens: Optional[torch.Tensor] = None  # shape: [b], int64
+    mamba_prefill_track_mask_cpu: Optional[List[bool]] = None
+    mamba_track_seqlens_cpu: Optional[List[int]] = None
     # Deferred mamba init ops: COW pairs and clear indices (performed on forward stream)
     mamba_cow_src_indices: Optional[torch.Tensor] = None
     mamba_cow_dst_indices: Optional[torch.Tensor] = None
@@ -912,6 +914,16 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             mamba_track_indices=batch.mamba_track_indices,
             mamba_track_mask=batch.mamba_track_mask,
             mamba_track_seqlens=batch.mamba_track_seqlens,
+            mamba_prefill_track_mask_cpu=(
+                list(batch.mamba_prefill_track_mask_cpu)
+                if batch.mamba_prefill_track_mask_cpu is not None
+                else None
+            ),
+            mamba_track_seqlens_cpu=(
+                list(batch.mamba_track_seqlens_cpu)
+                if batch.mamba_track_seqlens_cpu is not None
+                else None
+            ),
             mamba_cow_src_indices=batch.mamba_cow_src_indices,
             mamba_cow_dst_indices=batch.mamba_cow_dst_indices,
             mamba_clear_indices=batch.mamba_clear_indices,
@@ -1035,8 +1047,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                 ret.extend_prefix_lens = torch.tensor(
                     extend_prefix_lens, dtype=torch.int32, pin_memory=pin_memory
                 ).to(device, non_blocking=True)
-                ret.extend_prefix_lens_cpu = extend_prefix_lens
+                ret.extend_prefix_lens_cpu = list(extend_prefix_lens)
-                ret.extend_seq_lens_cpu = extend_seq_lens
+                ret.extend_seq_lens_cpu = list(extend_seq_lens)
             else:
                 # gpu_only: device tensors handed in directly; leave *_cpu unset.
                 assert isinstance(extend_seq_lens, torch.Tensor)
@@ -1689,6 +1701,14 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             self.mamba_track_indices = self._pad_tensor_to_size(
                 self.mamba_track_indices, bs
             )
+        if self.mamba_prefill_track_mask_cpu is not None:
+            self.mamba_prefill_track_mask_cpu = self.mamba_prefill_track_mask_cpu + [
+                False
+            ] * (bs - len(self.mamba_prefill_track_mask_cpu))
+        if self.mamba_track_seqlens_cpu is not None:
+            self.mamba_track_seqlens_cpu = self.mamba_track_seqlens_cpu + [0] * (
+                bs - len(self.mamba_track_seqlens_cpu)
+            )
         if self.mamba_track_mask is not None:
             self.mamba_track_mask = self._pad_tensor_to_size(self.mamba_track_mask, bs)
         if self.mamba_track_seqlens is not None:

python/sglang/srt/models/glm5_next.py

@@ -35,6 +35,7 @@ from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelBatchedLinear,
     ColumnParallelLinear,
+    LinearBase,
     MergedColumnParallelLinear,
     MergedColumnParallelRepeatedLinear,
     QKVParallelLinear,
@@ -48,6 +49,7 @@ from sglang.srt.layers.moe.utils import (
     is_shared_experts_fusion_disabled,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.utils.common import PPMissingLayer
@@ -98,7 +100,13 @@ from sglang.srt.multimodal.mm_utils import (
     run_dp_presharded_mrope_vision_model,
     run_dp_sharded_mrope_vision_model,
 )
-from sglang.srt.runtime_context import get_forward, get_mm, get_parallel, get_spec
+from sglang.srt.runtime_context import (
+    get_forward,
+    get_lora,
+    get_mm,
+    get_parallel,
+    get_spec,
+)
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils.common import (
     BumpAllocator,
@@ -303,6 +311,55 @@ class Glm5NextVisionModel(GlmOcrVisionModel):
 
 
 class Glm5NextLinearAttention(nn.Module):
+    _PACKED_MODULES_MAPPING = {
+        "fused_qkvbfg_a_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "b_proj",
+            "f_a_proj",
+            "g_a_proj",
+        ],
+        "fused_bfg_a_proj": ["b_proj", "f_a_proj", "g_a_proj"],
+        "fused_fg_b_proj": ["f_b_proj", "g_b_proj"],
+    }
+
+    @classmethod
+    def _can_fuse_proj(
+        cls,
+        quant_config: Optional[QuantizationConfig],
+        prefix: str,
+        *fused_projs: str,
+    ) -> bool:
+        if get_lora().enable_lora or get_lora().lora_paths:
+            return False
+        if quant_config is None:
+            return True
+        if quant_config.get_name() not in {
+            "fp8",
+            "mxfp8",
+            "modelopt_fp8",
+            "modelopt_fp4",
+            "modelopt_mixed",
+        }:
+            return False
+
+        probe = LinearBase(1, 1)
+        source_projs = [
+            proj
+            for fused_proj in fused_projs
+            for proj in cls._PACKED_MODULES_MAPPING[fused_proj]
+        ]
+        if "fused_qkvbfg_a_proj" in fused_projs:
+            source_projs.append("qkv_proj")
+        return all(
+            isinstance(
+                quant_config.get_quant_method(probe, prefix=f"{prefix}.{proj}"),
+                UnquantizedLinearMethod,
+            )
+            for proj in source_projs
+        )
+
     def __init__(
         self,
         layer_idx: int,
@@ -336,7 +393,12 @@ class Glm5NextLinearAttention(nn.Module):
         projection_size = self.head_dim * self.num_heads
         self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
 
-        self.do_fuse_qkvbfg = quant_config is None and head_shard_size == self.tp_size
+        self.do_fuse_qkvbfg = self._can_fuse_proj(
+            quant_config, prefix, "fused_qkvbfg_a_proj", "fused_fg_b_proj"
+        )
+        self.fuse_bfg = not self.do_fuse_qkvbfg and self._can_fuse_proj(
+            quant_config, prefix, "fused_bfg_a_proj", "fused_fg_b_proj"
+        )
         if self.do_fuse_qkvbfg:
             self.qkvb_sizes = [
                 projection_size,
@@ -350,21 +412,23 @@ class Glm5NextLinearAttention(nn.Module):
                 self.hidden_size,
                 self.qkvb_sizes,
                 self.fg_sizes,
-                quant_config=quant_config,
+                quant_config=None,
                 prefix=f"{prefix}.fused_qkvbfg_a_proj",
+                tp_rank=head_shard_rank,
+                tp_size=head_shard_size,
             )
             self.split_sizes = [
                 3 * projection_size // head_shard_size,
                 self.num_heads // head_shard_size,
                 2 * self.head_dim,
             ]
-            fused_dtype = (
-                getattr(config, "dtype", None)
-                or getattr(config, "torch_dtype", None)
-                or torch.get_default_dtype()
-            )
             self.fused_fg_b_proj = ColumnParallelBatchedLinear(
-                2, self.head_dim, projection_size, dtype=fused_dtype
+                2,
+                self.head_dim,
+                projection_size,
+                dtype=self.fused_qkvbfg_a_proj.params_dtype,
+                tp_rank=head_shard_rank,
+                tp_size=head_shard_size,
             )
         else:
             self.qkv_proj = QKVParallelLinear(
@@ -379,50 +443,70 @@ class Glm5NextLinearAttention(nn.Module):
                 prefix=f"{prefix}.qkv_proj",
             )
 
-            self.f_a_proj = ReplicatedLinear(
+            if self.fuse_bfg:
-                self.hidden_size,
+                self.fused_bfg_a_proj = MergedColumnParallelRepeatedLinear(
-                self.head_dim,
+                    self.hidden_size,
-                bias=False,
+                    [self.num_heads],
-                quant_config=quant_config,
+                    [self.head_dim, self.head_dim],
-                prefix=f"{prefix}.f_a_proj",
+                    quant_config=None,
-            )
+                    prefix=f"{prefix}.fused_bfg_a_proj",
+                    tp_rank=head_shard_rank,
+                    tp_size=head_shard_size,
+                )
+                self.bfg_split_sizes = [self.local_num_heads, 2 * self.head_dim]
+                self.fused_fg_b_proj = ColumnParallelBatchedLinear(
+                    2,
+                    self.head_dim,
+                    projection_size,
+                    dtype=self.fused_bfg_a_proj.params_dtype,
+                    tp_rank=head_shard_rank,
+                    tp_size=head_shard_size,
+                )
+            else:
+                self.f_a_proj = ReplicatedLinear(
+                    self.hidden_size,
+                    self.head_dim,
+                    bias=False,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.f_a_proj",
+                )
 
-            self.f_b_proj = ColumnParallelLinear(
+                self.f_b_proj = ColumnParallelLinear(
-                self.head_dim,
+                    self.head_dim,
-                projection_size,
+                    projection_size,
-                bias=False,
+                    bias=False,
-                quant_config=quant_config,
+                    quant_config=quant_config,
-                prefix=f"{prefix}.f_b_proj",
+                    prefix=f"{prefix}.f_b_proj",
-                tp_rank=head_shard_rank,
+                    tp_rank=head_shard_rank,
-                tp_size=head_shard_size,
+                    tp_size=head_shard_size,
-            )
+                )
 
-            self.b_proj = ColumnParallelLinear(
+                self.b_proj = ColumnParallelLinear(
-                self.hidden_size,
+                    self.hidden_size,
-                self.num_heads,
+                    self.num_heads,
-                bias=False,
+                    bias=False,
-                quant_config=quant_config,
+                    quant_config=quant_config,
-                prefix=f"{prefix}.b_proj",
+                    prefix=f"{prefix}.b_proj",
-                tp_rank=head_shard_rank,
+                    tp_rank=head_shard_rank,
-                tp_size=head_shard_size,
+                    tp_size=head_shard_size,
-            )
+                )
 
-            self.g_a_proj = ReplicatedLinear(
+                self.g_a_proj = ReplicatedLinear(
-                self.hidden_size,
+                    self.hidden_size,
-                self.head_dim,
+                    self.head_dim,
-                bias=False,
+                    bias=False,
-                quant_config=quant_config,
+                    quant_config=quant_config,
-                prefix=f"{prefix}.g_a_proj",
+                    prefix=f"{prefix}.g_a_proj",
-            )
+                )
-            self.g_b_proj = ColumnParallelLinear(
+                self.g_b_proj = ColumnParallelLinear(
-                self.head_dim,
+                    self.head_dim,
-                projection_size,
+                    projection_size,
-                bias=False,
+                    bias=False,
-                quant_config=quant_config,
+                    quant_config=quant_config,
-                prefix=f"{prefix}.g_b_proj",
+                    prefix=f"{prefix}.g_b_proj",
-                tp_rank=head_shard_rank,
+                    tp_rank=head_shard_rank,
-                tp_size=head_shard_size,
+                    tp_size=head_shard_size,
-            )
+                )
 
         self.dt_bias = nn.Parameter(
             torch.empty(divide(projection_size, head_shard_size), dtype=torch.float32)
@@ -490,9 +574,16 @@ class Glm5NextLinearAttention(nn.Module):
     def forward_qkvbfg(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch):
         qkv, _ = self.qkv_proj(hidden_states)
 
-        beta = self.b_proj(hidden_states)[0]
+        if self.fuse_bfg:
-        forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
+            fused_states = self.fused_bfg_a_proj(hidden_states)
-        g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
+            beta, fg_a_states = torch.split(fused_states, self.bfg_split_sizes, dim=-1)
+            forget_gate, g_proj_states = self.fused_fg_b_proj(
+                fg_a_states.view(-1, 2, self.head_dim).transpose(0, 1)
+            )
+        else:
+            beta = self.b_proj(hidden_states)[0]
+            forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
+            g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
 
         return (
             qkv,
@@ -1089,15 +1180,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
 
     packed_modules_mapping = {
         "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"],
-        "fused_qkvbfg_a_proj": [
+        **Glm5NextLinearAttention._PACKED_MODULES_MAPPING,
-            "q_proj",
-            "k_proj",
-            "v_proj",
-            "b_proj",
-            "f_a_proj",
-            "g_a_proj",
-        ],
-        "fused_fg_b_proj": ["f_b_proj", "g_b_proj"],
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
         "qkv_conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"],
         "gate_up_proj": ["gate_proj", "up_proj"],
@@ -1391,6 +1474,9 @@ class Glm5NextForConditionalGeneration(nn.Module):
             (".fused_qkvbfg_a_proj", ".g_a_proj", 5),
             (".fused_fg_b_proj", ".f_b_proj", 0),
             (".fused_fg_b_proj", ".g_b_proj", 1),
+            (".fused_bfg_a_proj", ".b_proj", 0),
+            (".fused_bfg_a_proj", ".f_a_proj", 1),
+            (".fused_bfg_a_proj", ".g_a_proj", 2),
             (".qkv_proj", ".q_proj", "q"),
             (".qkv_proj", ".k_proj", "k"),
             (".qkv_proj", ".v_proj", "v"),
@@ -1491,6 +1577,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
                     param_name
                     in {
                         ".fused_qkvbfg_a_proj",
+                        ".fused_bfg_a_proj",
                         ".fused_fg_b_proj",
                         ".qkv_proj",
                         ".qkv_conv1d",

python/sglang/srt/speculative/spec_utils.py

@@ -816,6 +816,8 @@ def prepare_mamba_track_for_verify(batch: ScheduleBatch) -> None:
     set_mamba_track_indices_from_reqs(batch, track_positions)
     batch.mamba_track_mask = None
     batch.mamba_track_seqlens = None
+    batch.mamba_prefill_track_mask_cpu = None
+    batch.mamba_track_seqlens_cpu = None
 
 
 def _verify_commit_step_indices(

test/registered/kernel/attention/test_kda_gate_beta_cumsum.py

@@ -0,0 +1,130 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.fla.kda import chunk_kda, kda_gate_chunk_cumsum
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestKDAGateBetaCumsum(unittest.TestCase):
+    @torch.inference_mode()
+    def test_gate_cumsum_beta_matches_separate_sigmoid(self):
+        torch.manual_seed(42)
+        for cu_seqlens, chunks, batch, tokens in (
+            (None, None, 2, 65),
+            (
+                torch.tensor([0, 0, 1, 65, 130], device="cuda", dtype=torch.int32),
+                torch.tensor(
+                    [[1, 0], [2, 0], [3, 0], [3, 1]], device="cuda", dtype=torch.int32
+                ),
+                1,
+                130,
+            ),
+        ):
+            heads, dim = 3, 128
+            gate = torch.randn(
+                batch, tokens, heads, dim, device="cuda", dtype=torch.bfloat16
+            )
+            a_log = torch.randn(heads, device="cuda")
+            bias = torch.randn(heads * dim, device="cuda")
+            packed = torch.randn(
+                batch, tokens, 4 * heads + 7, device="cuda", dtype=torch.bfloat16
+            )
+            beta = packed[..., 2 : 2 + heads]
+            for lower_bound in (None, -5.0):
+                with self.subTest(
+                    varlen=cu_seqlens is not None, lower_bound=lower_bound
+                ):
+                    kwargs = dict(
+                        A_log=a_log,
+                        chunk_size=64,
+                        dt_bias=bias,
+                        cu_seqlens=cu_seqlens,
+                        chunk_indices=chunks,
+                        lower_bound=lower_bound,
+                    )
+                    expected_gate = kda_gate_chunk_cumsum(gate, **kwargs)
+                    actual_gate, actual_beta = kda_gate_chunk_cumsum(
+                        gate, beta=beta, **kwargs
+                    )
+                    torch.testing.assert_close(
+                        actual_gate, expected_gate, atol=1e-4, rtol=1e-6
+                    )
+                    torch.testing.assert_close(
+                        actual_beta, beta.float().sigmoid(), atol=1.2e-7, rtol=1e-6
+                    )
+                    self.assertEqual(actual_beta.dtype, torch.float32)
+
+    @torch.inference_mode()
+    def test_chunk_raw_beta_matches_activated_beta_and_final_state(self):
+        torch.manual_seed(17)
+        tokens, heads, dim = 68, 2, 64
+        shape = (1, tokens, heads, dim)
+        q, k, v, gate = [
+            torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(4)
+        ]
+        packed = torch.randn(
+            1,
+            tokens,
+            3 * heads * dim + heads + 2 * dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        layouts = (
+            torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16),
+            torch.randn(
+                1, heads, tokens, device="cuda", dtype=torch.bfloat16
+            ).transpose(1, 2),
+            packed[..., 3 * heads * dim : 3 * heads * dim + heads],
+        )
+        a_log = torch.zeros(heads, device="cuda")
+        bias = torch.randn(heads * dim, device="cuda")
+        cu_seqlens = torch.tensor([0, 3, tokens], device="cuda", dtype=torch.int32)
+        state = torch.randn(2, heads, dim, dim, device="cuda") * 0.01
+        indices = torch.arange(2, device="cuda", dtype=torch.int32)
+        for beta in layouts:
+            for fused_gate in (False, True):
+                with self.subTest(stride=beta.stride(), fused_gate=fused_gate):
+                    kwargs = dict(
+                        q=q,
+                        k=k,
+                        scale=dim**-0.5,
+                        use_qk_l2norm_in_kernel=True,
+                        initial_state_indices=indices,
+                        cu_seqlens=cu_seqlens,
+                        A_log=a_log if fused_gate else None,
+                        dt_bias=bias if fused_gate else None,
+                        lower_bound=-5.0 if fused_gate else None,
+                    )
+                    g = (
+                        gate
+                        if fused_gate
+                        else -torch.nn.functional.softplus(gate.float())
+                    )
+                    expected_state, actual_state = state.clone(), state.clone()
+                    expected = chunk_kda(
+                        v=v.clone(),
+                        g=g.clone(),
+                        beta=beta.float().sigmoid(),
+                        initial_state=expected_state,
+                        **kwargs,
+                    )
+                    actual = chunk_kda(
+                        v=v.clone(),
+                        g=g.clone(),
+                        beta=beta,
+                        beta_is_raw=True,
+                        initial_state=actual_state,
+                        **kwargs,
+                    )
+                    torch.testing.assert_close(actual, expected, atol=2e-3, rtol=2e-3)
+                    torch.testing.assert_close(
+                        actual_state, expected_state, atol=2e-4, rtol=2e-3
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_mamba2_track_ssm_indices.py

@@ -25,11 +25,19 @@ def _backend():
 
 def _forward_batch(extend_lens, prefix_lens, track_seqlens, track_mask):
     return SimpleNamespace(
+        forward_mode=SimpleNamespace(
+            is_extend=lambda: True, is_target_verify=lambda: False
+        ),
         extend_seq_lens=torch.tensor(extend_lens),
         extend_prefix_lens=torch.tensor(prefix_lens),
         mamba_track_seqlens=torch.tensor(track_seqlens),
         mamba_track_mask=torch.tensor(track_mask),
         mamba_track_indices=torch.arange(100, 100 + len(extend_lens)),
+        # Exercise the legacy GPU planner, not the CPU-metadata fast path.
+        mamba_prefill_track_mask_cpu=None,
+        mamba_track_seqlens_cpu=None,
+        extend_seq_lens_cpu=None,
+        extend_prefix_lens_cpu=None,
     )
 
 

test/registered/unit/layers/test_mamba_prefill_track_metadata.py

@@ -0,0 +1,302 @@
+import random
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    Mamba2AttnBackend,
+    MambaAttnBackendBase,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.runtime_context import get_context
+from sglang.srt.speculative import spec_utils
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class NoHostRead(torch.Tensor):
+    def cpu(self, *args, **kwargs):
+        raise AssertionError("CPU tracking must not copy device metadata to the host")
+
+
+def make_batch(lengths, prefix, track_lens, mask, mirrored):
+    def tensor(values):
+        return torch.tensor(values).as_subclass(
+            NoHostRead if mirrored else torch.Tensor
+        )
+
+    return SimpleNamespace(
+        batch_size=len(lengths),
+        forward_mode=ForwardMode.EXTEND,
+        extend_seq_lens=tensor(lengths),
+        extend_prefix_lens=tensor(prefix),
+        mamba_track_seqlens=tensor(track_lens),
+        mamba_track_mask=tensor(mask),
+        mamba_track_indices=tensor([37 + i * 17 for i in range(len(lengths))]),
+        extend_seq_lens_cpu=lengths,
+        extend_prefix_lens_cpu=prefix,
+        mamba_track_seqlens_cpu=track_lens if mirrored else None,
+        mamba_prefill_track_mask_cpu=mask if mirrored else None,
+    )
+
+
+def make_forward_batch(lengths, starts, cpu_lengths, mode=ForwardMode.EXTEND):
+    return ForwardBatch(
+        forward_mode=mode,
+        batch_size=len(lengths),
+        input_ids=torch.zeros(sum(lengths), dtype=torch.int64),
+        req_pool_indices=torch.arange(len(lengths)),
+        seq_lens=torch.tensor(lengths),
+        seq_lens_sum=sum(lengths),
+        out_cache_loc=torch.zeros(sum(lengths), dtype=torch.int64),
+        extend_start_loc=torch.tensor(starts, dtype=torch.int32),
+        extend_seq_lens=torch.tensor(lengths, dtype=torch.int32),
+        extend_seq_lens_cpu=cpu_lengths,
+    )
+
+
+def make_metadata_backend():
+    backend = object.__new__(MambaAttnBackendBase)
+    backend.device = "cpu"
+    backend.topk = 1
+    backend.req_to_token_pool = SimpleNamespace(
+        get_mamba_indices=lambda rows: rows,
+        translate_mamba_indices=lambda slots: slots,
+    )
+    return backend
+
+
+class TestMambaPrefillTrackMetadata(unittest.TestCase):
+    def test_cpu_plan_matches_existing_tensor_planner(self):
+        rng = random.Random(2026)
+        for backend_type in (MambaAttnBackendBase, Mamba2AttnBackend):
+            for chunk in (16, 64, 128):
+                backend = object.__new__(backend_type)
+                backend.device = "cpu"
+                backend._mamba_chunk_size = chunk
+                cases = [
+                    (
+                        [chunk + 6, 2 * chunk + 1, chunk],
+                        [0, 2 * chunk, 0],
+                        [chunk + 1, 3 * chunk + 1, chunk],
+                        [True, True, True],
+                    ),
+                    (
+                        [2 * chunk + 1, 1, 1],
+                        [chunk, 0, 0],
+                        [2 * chunk + 1, 0, 0],
+                        [True, False, False],
+                    ),
+                    ([1, chunk], [0, 0], [0, 0], [False, False]),
+                ]
+                for _ in range(10):
+                    lengths = [rng.randrange(1, chunk * 6) for _ in range(5)]
+                    prefix = [rng.randrange(4) * chunk for _ in lengths]
+                    cases.append(
+                        (
+                            lengths,
+                            prefix,
+                            [
+                                p + rng.randrange(1, n + 1)
+                                for p, n in zip(prefix, lengths)
+                            ],
+                            [bool(rng.randrange(2)) for _ in lengths],
+                        )
+                    )
+                for lengths, prefix, track, mask in cases:
+                    slots = torch.tensor([111 - i * 5 for i in range(len(lengths))])
+                    with self.subTest(
+                        backend=backend_type.__name__, chunk=chunk, mask=mask
+                    ):
+                        expected = backend._init_track_ssm_indices(
+                            slots, make_batch(lengths, prefix, track, mask, False)
+                        )
+                        actual = backend._init_track_ssm_indices(
+                            slots.as_subclass(NoHostRead),
+                            make_batch(lengths, prefix, track, mask, True),
+                        )
+                        for result, reference in zip(actual, expected):
+                            if reference is None:
+                                self.assertIsNone(result)
+                            else:
+                                torch.testing.assert_close(result, reference)
+
+    def test_verify_and_incomplete_mirrors_use_existing_planner(self):
+        batch = make_batch([64], [0], [64], [True], True)
+        eligible = MambaAttnBackendBase._has_cpu_prefill_track_metadata
+        self.assertTrue(eligible(batch))
+        batch.forward_mode = ForwardMode.TARGET_VERIFY
+        self.assertFalse(eligible(batch))
+        batch.forward_mode = ForwardMode.EXTEND
+        batch.mamba_track_seqlens_cpu = None
+        self.assertFalse(eligible(batch))
+        batch.mamba_track_seqlens_cpu = [64, 0]
+        self.assertFalse(eligible(batch))
+
+    def test_logical_token_extent_avoids_scalar_reads_with_valid_cpu_lengths(self):
+        backend = make_metadata_backend()
+        original_int = torch.Tensor.__int__
+        for lengths, starts, cpu_lengths, tbo_range, expected, scalar_reads in (
+            ([3, 5], [0, 3], [3, 5], None, 8, 0),
+            ([3, 5, 0], [0, 3, 8], [3, 5, 0], None, 8, 0),
+            ([3, 5], [4, 7], None, None, 12, 1),
+            ([3, 5], [4, 7], [8], None, 12, 1),
+            ([3, 5], [4, 7], [3, 5], (4, 12), 12, 1),
+        ):
+            with self.subTest(cpu_lengths=cpu_lengths, tbo_range=tbo_range):
+                batch = make_forward_batch(lengths, starts, cpu_lengths)
+                batch.tbo_parent_token_range = tbo_range
+                reads = []
+
+                def read_scalar(tensor):
+                    if scalar_reads == 0:
+                        raise AssertionError(
+                            "Valid CPU lengths must avoid scalar reads"
+                        )
+                    reads.append(tensor.clone())
+                    return original_int(tensor)
+
+                with patch.object(torch.Tensor, "__int__", read_scalar):
+                    metadata = backend._forward_metadata(batch)
+                self.assertEqual(metadata.logical_num_tokens, expected)
+                self.assertEqual(len(reads), scalar_reads)
+                self.assertEqual(metadata.query_start_loc[-1].item(), expected)
+
+    def test_verify_decode_and_idle_ignore_stale_cpu_token_lengths(self):
+        backend = make_metadata_backend()
+        for mode, lengths, expected_starts in (
+            (ForwardMode.TARGET_VERIFY, [3, 3], [0, 3, 6]),
+            (ForwardMode.DECODE, [1, 1], [0, 1, 2]),
+            (ForwardMode.IDLE, [], [0]),
+        ):
+            with self.subTest(mode=mode):
+                batch = make_forward_batch(
+                    lengths, [0, 3][: len(lengths)], [100, 200], mode
+                )
+                if mode == ForwardMode.TARGET_VERIFY:
+                    batch.spec_info = SimpleNamespace(
+                        ragged_verify_layout=None, draft_token_num=3
+                    )
+                with patch.object(
+                    torch.Tensor,
+                    "__int__",
+                    side_effect=AssertionError("This mode must not read token scalars"),
+                ):
+                    metadata = backend._forward_metadata(batch)
+                self.assertIsNone(metadata.logical_num_tokens)
+                torch.testing.assert_close(
+                    metadata.query_start_loc,
+                    torch.tensor(expected_starts, dtype=torch.int32),
+                )
+
+    def test_forward_snapshot_and_padding_do_not_mutate_scheduler_lists(self):
+        override = get_context().override_server_args(device="cpu")
+        override.install()
+        self.addCleanup(override.restore)
+        batch = ScheduleBatch(
+            reqs=[SimpleNamespace(rid="one", lora_id=None, token_type_ids=None)],
+            device="cpu",
+            forward_mode=ForwardMode.EXTEND,
+            input_ids=torch.tensor([3]),
+            req_pool_indices=torch.tensor([2]),
+            seq_lens=torch.tensor([65]),
+            seq_lens_cpu=torch.tensor([65]),
+            seq_lens_sum=65,
+            out_cache_loc=torch.tensor([1]),
+            extend_lens=[1],
+            prefix_lens=[64],
+            extend_num_tokens=1,
+            mamba_track_mask=torch.tensor([True]),
+            mamba_track_seqlens=torch.tensor([65]),
+            mamba_prefill_track_mask_cpu=[True],
+            mamba_track_seqlens_cpu=[65],
+        )
+        runner = SimpleNamespace(
+            device="cpu",
+            model_config=SimpleNamespace(
+                requires_mm_token_modalities=False, model_is_mrope=False
+            ),
+            kv_index_translator=SimpleNamespace(rebind_write_loc=lambda forward: None),
+            prefill_attention_backend_str="torch_native",
+            ngram_embedding_manager=SimpleNamespace(enabled=False),
+            lora_manager=None,
+            ps=SimpleNamespace(attn_dcp_size=1),
+            attn_backend=SimpleNamespace(
+                get_cpu_graph_seq_len_fill_value=lambda: 1,
+                get_cuda_graph_seq_len_fill_value=lambda: 1,
+            ),
+        )
+        forward = ForwardBatch.init_new(
+            batch,
+            runner,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+            return_hidden_states_before_norm=False,
+        )
+        for target, source in (
+            ("mamba_prefill_track_mask_cpu", "mamba_prefill_track_mask_cpu"),
+            ("mamba_track_seqlens_cpu", "mamba_track_seqlens_cpu"),
+            ("extend_seq_lens_cpu", "extend_lens"),
+            ("extend_prefix_lens_cpu", "prefix_lens"),
+        ):
+            self.assertEqual(getattr(forward, target), getattr(batch, source))
+            self.assertIsNot(getattr(forward, target), getattr(batch, source))
+        forward._pad_inputs_to_size(runner, num_tokens=3, bs=3)
+        self.assertEqual(batch.mamba_prefill_track_mask_cpu, [True])
+        self.assertEqual(batch.mamba_track_seqlens_cpu, [65])
+        self.assertEqual(batch.extend_lens, [1])
+        self.assertEqual(batch.prefix_lens, [64])
+        for host, device, expected in (
+            ("mamba_prefill_track_mask_cpu", "mamba_track_mask", [True, False, False]),
+            ("mamba_track_seqlens_cpu", "mamba_track_seqlens", [65, 0, 0]),
+            ("extend_seq_lens_cpu", "extend_seq_lens", [1, 0, 0]),
+            ("extend_prefix_lens_cpu", "extend_prefix_lens", [64, 0, 0]),
+        ):
+            self.assertEqual(getattr(forward, host), expected)
+            self.assertEqual(getattr(forward, device).tolist(), expected)
+
+    def test_decode_and_verify_clear_prefill_lists_without_losing_snapshot(self):
+        for verify in (False, True):
+            with self.subTest(verify=verify):
+                batch = ScheduleBatch(
+                    reqs=[],
+                    spec_algorithm=SimpleNamespace(is_none=lambda: False),
+                    mamba_track_mask=torch.tensor([True]),
+                    mamba_track_seqlens=torch.tensor([65]),
+                    mamba_prefill_track_mask_cpu=[True],
+                    mamba_track_seqlens_cpu=[65],
+                )
+                snapshot = batch.copy()
+                if verify:
+                    settings = SimpleNamespace(
+                        mamba=SimpleNamespace(
+                            enable_mamba_extra_buffer=True,
+                            enable_mamba_extra_buffer_lazy=False,
+                        )
+                    )
+                    with (
+                        patch.object(spec_utils, "get_exec", return_value=settings),
+                        patch.object(spec_utils, "set_mamba_track_indices_from_reqs"),
+                    ):
+                        spec_utils.prepare_mamba_track_for_verify(batch)
+                    self.assertIsNone(batch.mamba_track_mask)
+                    self.assertIsNone(batch.mamba_track_seqlens)
+                else:
+                    with patch.object(spec_utils, "spec_prepare_for_decode"):
+                        batch.prepare_for_decode()
+                self.assertIsNone(batch.mamba_prefill_track_mask_cpu)
+                self.assertIsNone(batch.mamba_track_seqlens_cpu)
+                self.assertEqual(snapshot.mamba_prefill_track_mask_cpu, [True])
+                self.assertEqual(snapshot.mamba_track_seqlens_cpu, [65])
+                self.assertIsNone(snapshot.mamba_track_mask_cpu)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/models/test_glm5_next_bfg_fusion.py

@@ -0,0 +1,182 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.models import glm5_next
+from sglang.srt.runtime_context import get_context, get_parallel
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+PREFIX = "model.layers.0.self_attn"
+QKV = ("q_proj", "k_proj", "v_proj")
+BFG = ("b_proj", "f_a_proj", "g_a_proj", "f_b_proj", "g_b_proj")
+
+
+class MockQuantizedLinearMethod:
+    """Keep dense storage so the test isolates routing and checkpoint loading."""
+
+    create_weights = UnquantizedLinearMethod.create_weights
+
+    def apply(self, layer, x, bias=None):
+        return F.linear(x, layer.weight, bias)
+
+
+class MockFp8Config:
+    def __init__(self, ignored):
+        self.ignored_layers = {f"{PREFIX}.{name}" for name in ignored}
+
+    def get_name(self):
+        return "fp8"
+
+    def get_quant_method(self, layer, prefix):
+        names = (
+            [prefix.replace("qkv_proj", name) for name in QKV]
+            if prefix.endswith(".qkv_proj")
+            else [prefix]
+        )
+        if all(name in self.ignored_layers for name in names):
+            return UnquantizedLinearMethod()
+        return MockQuantizedLinearMethod()
+
+
+class TestGlm5NextBfgFusion(unittest.TestCase):
+    def setUp(self):
+        self.addCleanup(torch.set_default_dtype, torch.get_default_dtype())
+        torch.set_default_dtype(torch.float32)
+        override = get_context().override_server_args(
+            device="cpu", enable_lora=False, lora_paths=None
+        )
+        override.install()
+        self.addCleanup(override.restore)
+        patcher = patch.object(
+            UnquantizedLinearMethod,
+            "apply",
+            MockQuantizedLinearMethod.apply,
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+
+    @torch.no_grad()
+    def test_projection_loading_matches_unfused_reference(self):
+        torch.manual_seed(42)
+        hidden, heads, dim = 16, 4, 8
+        shapes = {name: (heads * dim, hidden) for name in QKV}
+        shapes.update(
+            b_proj=(heads, hidden),
+            f_a_proj=(dim, hidden),
+            g_a_proj=(dim, hidden),
+            f_b_proj=(heads * dim, dim),
+            g_b_proj=(heads * dim, dim),
+        )
+        weights = {name: torch.randn(shape) for name, shape in shapes.items()}
+        x = torch.randn(7, hidden)
+        for ignored, expected_route in (
+            (QKV + BFG, (True, False)),
+            (BFG, (False, True)),
+            ((), (False, False)),
+        ):
+            for attn_tp, rank in ((1, 0), (2, 0), (2, 1)):
+                with (
+                    self.subTest(route=expected_route, attn_tp=attn_tp, rank=rank),
+                    get_parallel().override(
+                        tp_size=4, tp_rank=3, attn_tp_size=attn_tp, attn_tp_rank=rank
+                    ),
+                ):
+                    quant = MockFp8Config(ignored)
+                    attention = glm5_next.Glm5NextLinearAttention(
+                        layer_idx=0,
+                        hidden_size=hidden,
+                        config=SimpleNamespace(
+                            linear_attn_config={
+                                "head_dim": dim,
+                                "num_heads": heads,
+                                "short_conv_kernel_size": 4,
+                            }
+                        ),
+                        quant_config=quant,
+                        prefix=PREFIX,
+                    )
+                    self.assertEqual(
+                        (attention.do_fuse_qkvbfg, attention.fuse_bfg), expected_route
+                    )
+                    for parameter in attention.parameters():
+                        parameter.fill_(torch.nan)
+                    model = SimpleNamespace(
+                        config=SimpleNamespace(n_routed_experts=0),
+                        num_fused_shared_experts=0,
+                        quant_config=quant,
+                        named_parameters=lambda: (
+                            (f"{PREFIX}.{name}", param)
+                            for name, param in attention.named_parameters()
+                        ),
+                    )
+                    with patch.object(
+                        glm5_next.DeepseekV2WeightLoaderMixin, "post_load_weights"
+                    ):
+                        glm5_next.Glm5NextForConditionalGeneration.load_weights(
+                            model,
+                            [
+                                (f"{PREFIX}.{name}.weight", w)
+                                for name, w in weights.items()
+                            ],
+                        )
+
+                    def linear(value, name):
+                        weight = weights[name]
+                        if name not in ("f_a_proj", "g_a_proj"):
+                            weight = weight.chunk(attn_tp, dim=0)[rank]
+                        return F.linear(value, weight)
+
+                    expected = (
+                        torch.cat([linear(x, name) for name in QKV], dim=-1),
+                        linear(x, "b_proj"),
+                        linear(linear(x, "f_a_proj"), "f_b_proj"),
+                        linear(linear(x, "g_a_proj"), "g_b_proj"),
+                    )
+                    forward = (
+                        attention.forward_qkvbfg_fused
+                        if attention.do_fuse_qkvbfg
+                        else attention.forward_qkvbfg
+                    )
+                    for actual, reference in zip(forward(x, None), expected):
+                        torch.testing.assert_close(
+                            actual, reference, atol=1e-5, rtol=1e-5
+                        )
+
+    def test_each_quantized_gate_projection_disables_fusion(self):
+        for quantized in BFG:
+            quant = MockFp8Config(name for name in QKV + BFG if name != quantized)
+            for packed in ("fused_qkvbfg_a_proj", "fused_bfg_a_proj"):
+                with self.subTest(quantized=quantized, packed=packed):
+                    self.assertFalse(
+                        glm5_next.Glm5NextLinearAttention._can_fuse_proj(
+                            quant, PREFIX, packed, "fused_fg_b_proj"
+                        )
+                    )
+
+    def test_lora_disables_full_and_bfg_fusion(self):
+        for enable_lora, paths in ((True, None), (False, ["adapter"])):
+            with patch.object(
+                glm5_next,
+                "get_lora",
+                return_value=SimpleNamespace(enable_lora=enable_lora, lora_paths=paths),
+            ):
+                for quant in (None, MockFp8Config(QKV + BFG)):
+                    for packed in ("fused_qkvbfg_a_proj", "fused_bfg_a_proj"):
+                        with self.subTest(
+                            enabled=enable_lora, paths=paths, packed=packed
+                        ):
+                            self.assertFalse(
+                                glm5_next.Glm5NextLinearAttention._can_fuse_proj(
+                                    quant, PREFIX, packed, "fused_fg_b_proj"
+                                )
+                            )
+
+
+if __name__ == "__main__":
+    unittest.main()