From c8eb54c41da1d6b452e77469fda548547dd21b14 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang Date: Sun, 20 Sep 2026 14:46:33 +0800 Subject: [PATCH] Fuse GLM-5.3-Flash KDA projections and prefill metadata (#39688) Co-authored-by: Xinyuan Tong --- .../sglang/kernels/ops/attention/fla/kda.py | 67 +++- .../attention/hybrid_linear_attn_backend.py | 99 +++++- .../layers/attention/linear/gdn_backend.py | 7 +- .../layers/attention/linear/kda_backend.py | 11 +- .../layers/attention/mamba/mamba2_metadata.py | 1 + .../attention/mamba/prefill_track_metadata.py | 57 ++++ python/sglang/srt/managers/schedule_batch.py | 13 + .../srt/model_executor/forward_batch_info.py | 24 +- python/sglang/srt/models/glm5_next.py | 211 ++++++++---- python/sglang/srt/speculative/spec_utils.py | 2 + .../attention/test_kda_gate_beta_cumsum.py | 130 ++++++++ .../layers/test_mamba2_track_ssm_indices.py | 8 + .../test_mamba_prefill_track_metadata.py | 302 ++++++++++++++++++ .../unit/models/test_glm5_next_bfg_fusion.py | 182 +++++++++++ 14 files changed, 1027 insertions(+), 87 deletions(-) create mode 100644 python/sglang/srt/layers/attention/mamba/prefill_track_metadata.py create mode 100644 test/registered/kernel/attention/test_kda_gate_beta_cumsum.py create mode 100644 test/registered/unit/layers/test_mamba_prefill_track_metadata.py create mode 100644 test/registered/unit/models/test_glm5_next_bfg_fusion.py diff --git a/python/sglang/kernels/ops/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py index 9b91feb23..3b877ec67 100644 --- a/python/sglang/kernels/ops/attention/fla/kda.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index b1e4ac897..6249cb6b5 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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( - forward_batch.mamba_track_mask is not None - and forward_batch.mamba_track_mask.any() - ) + cpu_track_metadata = self._has_cpu_prefill_track_metadata(forward_batch) + if cpu_track_metadata: + 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, diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index dfdfee790..21148e6b0 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/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 = ( - forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] - ) + if getattr(self.forward_metadata, "mamba_track_mask_indices", None) is None: + 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 diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index b3b5cf8ab..54ffc91d5 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/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 = ( - forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] - ) + if self.forward_metadata.mamba_track_mask_indices is None: + 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] diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 369654de6..c0cba1d8b 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/mamba/prefill_track_metadata.py b/python/sglang/srt/layers/attention/mamba/prefill_track_metadata.py new file mode 100644 index 000000000..06ac00b80 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 777b27d68..81ea87a14 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 94b7b597f..40939377d 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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_seq_lens_cpu = extend_seq_lens + ret.extend_prefix_lens_cpu = list(extend_prefix_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: diff --git a/python/sglang/srt/models/glm5_next.py b/python/sglang/srt/models/glm5_next.py index 47d2b20c0..50aafeab7 100644 --- a/python/sglang/srt/models/glm5_next.py +++ b/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( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.f_a_proj", - ) + if self.fuse_bfg: + self.fused_bfg_a_proj = MergedColumnParallelRepeatedLinear( + self.hidden_size, + [self.num_heads], + [self.head_dim, self.head_dim], + 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.head_dim, - projection_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.f_b_proj", - tp_rank=head_shard_rank, - tp_size=head_shard_size, - ) + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_b_proj", + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) - self.b_proj = ColumnParallelLinear( - self.hidden_size, - self.num_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.b_proj", - tp_rank=head_shard_rank, - tp_size=head_shard_size, - ) + self.b_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.b_proj", + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) - self.g_a_proj = ReplicatedLinear( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.g_a_proj", - ) - self.g_b_proj = ColumnParallelLinear( - self.head_dim, - projection_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.g_b_proj", - tp_rank=head_shard_rank, - tp_size=head_shard_size, - ) + self.g_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_a_proj", + ) + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_b_proj", + tp_rank=head_shard_rank, + 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] - 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] + if self.fuse_bfg: + fused_states = self.fused_bfg_a_proj(hidden_states) + 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": [ - "q_proj", - "k_proj", - "v_proj", - "b_proj", - "f_a_proj", - "g_a_proj", - ], - "fused_fg_b_proj": ["f_b_proj", "g_b_proj"], + **Glm5NextLinearAttention._PACKED_MODULES_MAPPING, "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", diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 0dc7092a4..a94538c44 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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( diff --git a/test/registered/kernel/attention/test_kda_gate_beta_cumsum.py b/test/registered/kernel/attention/test_kda_gate_beta_cumsum.py new file mode 100644 index 000000000..60dc2b387 --- /dev/null +++ b/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() diff --git a/test/registered/unit/layers/test_mamba2_track_ssm_indices.py b/test/registered/unit/layers/test_mamba2_track_ssm_indices.py index 016ec28d7..48828418f 100644 --- a/test/registered/unit/layers/test_mamba2_track_ssm_indices.py +++ b/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, ) diff --git a/test/registered/unit/layers/test_mamba_prefill_track_metadata.py b/test/registered/unit/layers/test_mamba_prefill_track_metadata.py new file mode 100644 index 000000000..dee5d7d51 --- /dev/null +++ b/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() diff --git a/test/registered/unit/models/test_glm5_next_bfg_fusion.py b/test/registered/unit/models/test_glm5_next_bfg_fusion.py new file mode 100644 index 000000000..c7607ff6b --- /dev/null +++ b/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()