From 03e4c06589e35191b740b05d16f25aaaacb086ce Mon Sep 17 00:00:00 2001 From: DAI0818 <53027952+daii-0818@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:30:39 +0800 Subject: [PATCH] Fix Qwen3.5 GDN multi-item scoring (#33922) Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com> --- .../sglang/kernels/ops/attention/fla/chunk.py | 8 + .../ops/attention/fla/chunk_delta_h.py | 3 +- .../xpu/kernels/fla/chunk_delta_h.py | 5 + .../attention/hybrid_linear_attn_backend.py | 12 + .../layers/attention/linear/gdn_backend.py | 265 ++++++++++++++++++ .../attention/linear/kernels/gdn_triton.py | 18 +- .../attention_methods/gdn_attention.py | 137 +++++++-- .../attention/test_chunk_gated_delta_rule.py | 83 ++++++ .../attention/unittests/gdn/test_triton.py | 30 ++ .../gdn/test_qwen35_gdn_multi_item_scoring.py | 133 +++++++++ .../layers/mamba/test_causal_conv1d.py | 57 ++++ .../test_gdn_prefill_backend_policy.py | 19 ++ .../unit/layers/test_gdn_mis_metadata.py | 99 +++++++ 13 files changed, 842 insertions(+), 27 deletions(-) create mode 100644 test/registered/e2e/gdn/test_qwen35_gdn_multi_item_scoring.py create mode 100644 test/registered/unit/layers/test_gdn_mis_metadata.py diff --git a/python/sglang/kernels/ops/attention/fla/chunk.py b/python/sglang/kernels/ops/attention/fla/chunk.py index 80cfb4e1f..f4b892f68 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk.py +++ b/python/sglang/kernels/ops/attention/fla/chunk.py @@ -44,6 +44,7 @@ def chunk_gated_delta_rule_fwd( initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: torch.LongTensor | None = None, + inplace_update: bool = True, ): g = chunk_local_cumsum( g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices @@ -68,6 +69,7 @@ def chunk_gated_delta_rule_fwd( initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + inplace_update=inplace_update, ) o = chunk_fwd_o( q=q, @@ -100,6 +102,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False, + inplace_update: bool = True, ): q_orig = q k_orig = k @@ -124,6 +127,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + inplace_update=inplace_update, ) return o.to(q.dtype), h @@ -141,6 +145,7 @@ def chunk_gated_delta_rule( cu_seqlens: Optional[torch.LongTensor] = None, head_first: bool = False, use_qk_l2norm_in_kernel: bool = False, + inplace_update: bool = True, ): r""" Args: @@ -169,6 +174,8 @@ def chunk_gated_delta_rule( head_first (Optional[bool]): Whether the inputs are in the head-first format, which is not supported for variable-length inputs. Default: `False`. + inplace_update (Optional[bool]): + Whether to write final states back to `initial_state`. Default: `True`. Returns: o (torch.Tensor): @@ -255,6 +262,7 @@ def chunk_gated_delta_rule( initial_state_indices, cu_seqlens, use_qk_l2norm_in_kernel, + inplace_update, ) if head_first: o = rearrange(o, "b t h ... -> b h t ...") diff --git a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py index 7241e145c..84ae081f3 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py @@ -360,6 +360,7 @@ def chunk_gated_delta_rule_fwd_h( cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: Optional[torch.LongTensor] = None, use_exp2: bool = False, + inplace_update: bool = True, track_state: Optional[torch.Tensor] = None, track_chunk_idx: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -426,7 +427,7 @@ def chunk_gated_delta_rule_fwd_h( USE_G=g is not None, USE_GK=gk is not None, USE_INITIAL_STATE=initial_state is not None, - INPLACE_UPDATE=True, + INPLACE_UPDATE=inplace_update, SAVE_NEW_VALUE=v_new is not None, IS_VARLEN=cu_seqlens is not None, NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)), diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py index 2f6a354bc..e3b7da63a 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py @@ -242,7 +242,12 @@ def chunk_gated_delta_rule_fwd_h( cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: Optional[torch.LongTensor] = None, use_exp2: bool = False, + inplace_update: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if not inplace_update: + raise NotImplementedError( + "GDN multi-item scoring is not supported by the XPU chunk kernel" + ) assert not (use_exp2 and g is not None), ( "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" ) 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 7587839f7..b1e4ac897 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -47,12 +47,24 @@ _validate_mamba_replay_state_indices = ( class MambaAttnBackendBase(AttentionBackend): + supports_mis: bool = False + + @classmethod + def validate_mis_support(cls, server_args) -> None: + if server_args.enable_mis and not cls.supports_mis: + raise ValueError( + f"{cls.__name__} does not support multi-item scoring. " + "Hybrid models require a linear-attention backend that explicitly " + "declares MIS support." + ) + # Per-slot accept lengths for the KDA fused-accept spec path; allocated only # by KDAAttnBackend where `_can_fuse_accept_state` holds. None everywhere # else — update_mamba_state_after_mtp_verify keys the fused branch on it. accept_lens_pool: Optional[torch.Tensor] = None def __init__(self, model_runner: ModelRunner): + self.validate_mis_support(model_runner.server_args) super().__init__() self.pad_slot_id = PAD_SLOT_ID self.device = model_runner.device diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index faca9e1d0..dfdfee790 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -1,5 +1,6 @@ from typing import Optional, Tuple, Union +import msgspec import torch from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating @@ -49,6 +50,117 @@ _fused_decode_verify_real_tensors = ( envs.SGLANG_GDN_DECODE_FUSION_VERIFY_REAL_TENSORS.get() ) + +class GDNMISMetadata(msgspec.Struct, frozen=True): + query_token_indices: torch.Tensor + query_cu_seqlens: torch.Tensor + query_seq_lens_cpu: list[int] + query_request_indices: torch.Tensor + item_token_indices: torch.Tensor + item_cu_seqlens: torch.Tensor + item_seq_lens_cpu: list[int] + item_request_indices: torch.Tensor + + +def build_gdn_mis_metadata(forward_batch: ForwardBatch) -> GDNMISMetadata: + """Build compact query/item segments from request-local MIS delimiters.""" + if not forward_batch.is_prefill_only: + raise ValueError("GDN MIS is only supported for prefill-only requests") + + prefix_lens = forward_batch.extend_prefix_lens_cpu + if isinstance(prefix_lens, torch.Tensor): + prefix_lens = prefix_lens.tolist() + if any(int(prefix_len) != 0 for prefix_len in prefix_lens): + raise ValueError("GDN MIS does not support cached prefixes") + + seq_lens = forward_batch.extend_seq_lens_cpu + if isinstance(seq_lens, torch.Tensor): + seq_lens = seq_lens.tolist() + seq_lens = [int(seq_len) for seq_len in seq_lens] + delimiter_indices = forward_batch.multi_item_delimiter_indices + if delimiter_indices is None or len(delimiter_indices) != len(seq_lens): + raise ValueError("GDN MIS requires delimiter indices for every request") + if sum(seq_lens) > forward_batch.input_ids.numel(): + raise ValueError("GDN MIS sequence lengths exceed the input tokens") + + query_token_indices: list[int] = [] + query_seq_lens_cpu: list[int] = [] + query_request_indices: list[int] = [] + item_token_indices: list[int] = [] + item_seq_lens_cpu: list[int] = [] + item_request_indices: list[int] = [] + + request_start = 0 + for request_idx, (seq_len, request_delimiters) in enumerate( + zip(seq_lens, delimiter_indices) + ): + delimiters = [int(index) for index in request_delimiters.tolist()] + if len(delimiters) < 2: + raise ValueError("GDN MIS requires at least two delimiters per request") + if any( + current >= following + for current, following in zip(delimiters, delimiters[1:]) + ): + raise ValueError("GDN MIS delimiter indices must be strictly increasing") + if delimiters[0] < 0 or delimiters[-1] >= seq_len: + raise ValueError("GDN MIS delimiter index is outside the request") + if delimiters[-1] != seq_len - 1: + raise ValueError("GDN MIS final delimiter must be the last request token") + + query_len = delimiters[0] + if query_len > 0: + query_token_indices.extend(range(request_start, request_start + query_len)) + query_seq_lens_cpu.append(query_len) + query_request_indices.append(request_idx) + + branch_ends = delimiters[1:] + [seq_len] + for branch_start, branch_end in zip(delimiters, branch_ends): + branch_len = branch_end - branch_start + item_token_indices.extend( + range(request_start + branch_start, request_start + branch_end) + ) + item_seq_lens_cpu.append(branch_len) + item_request_indices.append(request_idx) + + request_start += seq_len + + device = forward_batch.input_ids.device + + def _indices(values: list[int], dtype: torch.dtype) -> torch.Tensor: + return torch.tensor(values, dtype=dtype, device=device) + + def _cu_seqlens(lengths: list[int]) -> torch.Tensor: + result = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=device) + if lengths: + result[1:] = torch.tensor(lengths, dtype=torch.int32, device=device).cumsum( + dim=0 + ) + return result + + return GDNMISMetadata( + query_token_indices=_indices(query_token_indices, torch.int64), + query_cu_seqlens=_cu_seqlens(query_seq_lens_cpu), + query_seq_lens_cpu=query_seq_lens_cpu, + query_request_indices=_indices(query_request_indices, torch.int64), + item_token_indices=_indices(item_token_indices, torch.int64), + item_cu_seqlens=_cu_seqlens(item_seq_lens_cpu), + item_seq_lens_cpu=item_seq_lens_cpu, + item_request_indices=_indices(item_request_indices, torch.int64), + ) + + +def validate_gdn_mis_backend(prefill_backend: LinearAttnKernelBackend) -> None: + if not get_exec().features.enable_mis: + return + if not prefill_backend.is_triton(): + raise ValueError( + "GDN multi-item scoring requires the Triton linear-attention prefill " + "backend. Set --linear-attn-prefill-backend triton." + ) + if get_memory().enable_page_major_kv_layout: + raise ValueError("GDN multi-item scoring does not support page-major layout") + + if is_cuda(): from sglang.srt.layers.attention.mamba.causal_conv1d import ( causal_conv1d_fn as causal_conv1d_fn_cuda, @@ -391,10 +503,13 @@ class GDNAttnBackend(MambaAttnBackendBase): """Attention backend for GDN (Gated Delta Network) linear attention.""" needs_cpu_seq_lens: bool = False + supports_mis: bool = True def __init__(self, model_runner: ModelRunner): _validate_gdn_linear_attn_backends(model_runner.linear_attn_backends) super().__init__(model_runner) + self.enable_mis = get_exec().features.enable_mis + self.mis_metadata: Optional[GDNMISMetadata] = None self.conv_states_shape = ( model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape ) @@ -404,6 +519,7 @@ class GDNAttnBackend(MambaAttnBackendBase): ) backends = model_runner.linear_attn_backends + validate_gdn_mis_backend(backends.prefill) self.linear_attn_backends = backends self.kernel_dispatcher = GDNKernelDispatcher( backends.decode, backends.prefill, backends.verify @@ -418,6 +534,11 @@ class GDNAttnBackend(MambaAttnBackendBase): def init_forward_metadata(self, forward_batch: ForwardBatch): super().init_forward_metadata(forward_batch) + self.mis_metadata = None + if forward_batch.multi_item_delimiter_indices is not None: + if not self.enable_mis: + 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] @@ -710,6 +831,18 @@ class GDNAttnBackend(MambaAttnBackendBase): mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) conv_states = mamba_cache_params.conv[0] ssm_states = mamba_cache_params.temporal + if self.mis_metadata is not None: + if is_target_verify: + raise ValueError("GDN MIS does not support target verify") + return self._forward_extend_mis( + layer=layer, + mixed_qkv=mixed_qkv, + a=a, + b=b, + conv_states=conv_states, + ssm_states=ssm_states, + cache_indices=cache_indices, + ) if is_target_verify: assert isinstance(mamba_cache_params, MambaPool.SpeculativeState) intermediate_state_cache = mamba_cache_params.intermediate_ssm @@ -923,6 +1056,138 @@ class GDNAttnBackend(MambaAttnBackendBase): return core_attn_out + def _forward_extend_mis( + self, + *, + layer: RadixLinearAttention, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + ) -> torch.Tensor: + metadata = self.mis_metadata + assert metadata is not None + + conv_states[cache_indices] = 0 + ssm_states[cache_indices] = 0 + output = mixed_qkv.new_zeros( + 1, mixed_qkv.shape[0], layer.num_v_heads, layer.head_v_dim + ) + + if metadata.query_token_indices.numel() > 0: + query_cache_indices = cache_indices[metadata.query_request_indices] + query_output = self._forward_mis_segments( + layer=layer, + mixed_qkv=mixed_qkv[metadata.query_token_indices], + a=a[metadata.query_token_indices], + b=b[metadata.query_token_indices], + conv_states=conv_states, + conv_cache_indices=query_cache_indices, + has_initial_states=torch.zeros( + len(metadata.query_seq_lens_cpu), + dtype=torch.bool, + device=mixed_qkv.device, + ), + query_start_loc=metadata.query_cu_seqlens, + seq_lens_cpu=metadata.query_seq_lens_cpu, + ssm_states=ssm_states, + ssm_cache_indices=query_cache_indices, + inplace_update=True, + ) + output[:, metadata.query_token_indices] = query_output + + item_ssm_indices = cache_indices[metadata.item_request_indices] + item_conv_states = conv_states[item_ssm_indices] + item_conv_indices = torch.arange( + item_ssm_indices.shape[0], + dtype=cache_indices.dtype, + device=cache_indices.device, + ) + item_output = self._forward_mis_segments( + layer=layer, + mixed_qkv=mixed_qkv[metadata.item_token_indices], + a=a[metadata.item_token_indices], + b=b[metadata.item_token_indices], + conv_states=item_conv_states, + conv_cache_indices=item_conv_indices, + has_initial_states=torch.ones( + len(metadata.item_seq_lens_cpu), + dtype=torch.bool, + device=mixed_qkv.device, + ), + query_start_loc=metadata.item_cu_seqlens, + seq_lens_cpu=metadata.item_seq_lens_cpu, + ssm_states=ssm_states, + ssm_cache_indices=item_ssm_indices, + inplace_update=False, + ) + output[:, metadata.item_token_indices] = item_output + return output + + def _forward_mis_segments( + self, + *, + layer: RadixLinearAttention, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_states: torch.Tensor, + conv_cache_indices: torch.Tensor, + has_initial_states: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens_cpu: list[int], + ssm_states: torch.Tensor, + ssm_cache_indices: torch.Tensor, + inplace_update: bool, + ) -> torch.Tensor: + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer.conv_weights, + layer.bias, + activation=layer.activation, + conv_states=conv_states, + has_initial_state=has_initial_states, + cache_indices=conv_cache_indices, + query_start_loc=query_start_loc, + seq_lens_cpu=seq_lens_cpu, + ).transpose(0, 1) + + qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim + if (is_cuda() or is_hip()) and qkv_dim <= MAX_FUSED_QKV_SPLIT_DIM: + query, key, value = fused_qkv_split_gdn_prefill( + mixed_qkv, + layer.num_q_heads, + layer.num_k_heads, + layer.num_v_heads, + layer.head_q_dim, + layer.head_k_dim, + layer.head_v_dim, + ) + else: + query, key, value = torch.split( + mixed_qkv, [layer.q_dim, layer.k_dim, layer.v_dim], dim=-1 + ) + num_tokens = mixed_qkv.shape[0] + query = query.view(1, num_tokens, layer.num_q_heads, layer.head_q_dim) + key = key.view(1, num_tokens, layer.num_k_heads, layer.head_k_dim) + value = value.view(1, num_tokens, layer.num_v_heads, layer.head_v_dim) + + g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) + output, _, _ = self.kernel_dispatcher.extend( + q=query, + k=key, + v=value, + g=g, + beta=beta, + ssm_states=ssm_states, + cache_indices=ssm_cache_indices, + query_start_loc=query_start_loc, + inplace_update=inplace_update, + ) + return output + def _replayssm_fold_target_verify( self, *, diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py index c6942226b..b16c41ca8 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py @@ -177,13 +177,28 @@ class TritonGDNKernel(LinearAttnKernelBase): ssm_states: torch.Tensor, cache_indices: torch.Tensor, query_start_loc: torch.Tensor, + inplace_update: bool = True, **kwargs, ) -> tuple: recurrent_state = ssm_states recurrent_state_indices_args = {"initial_state_indices": cache_indices} - if is_npu(): + inplace_update_args = {"inplace_update": inplace_update} + if is_cpu(): + if not inplace_update: + raise NotImplementedError( + "GDN multi-item scoring is not supported by the CPU chunk kernel" + ) + inplace_update_args = {} + elif is_npu(): + if not inplace_update: + raise NotImplementedError( + "GDN multi-item scoring is not supported by the NPU chunk kernel" + ) recurrent_state = ssm_states[cache_indices] recurrent_state_indices_args = {} + # The external NPU kernel does not expose the optional write-back + # control. Its existing behavior is equivalent to True. + inplace_update_args = {} return chunk_gated_delta_rule( q=q, @@ -196,6 +211,7 @@ class TritonGDNKernel(LinearAttnKernelBase): head_first=False, use_qk_l2norm_in_kernel=True, **recurrent_state_indices_args, + **inplace_update_args, ) def target_verify( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index 556bc9293..83b07097d 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py @@ -30,7 +30,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner -from sglang.srt.runtime_context import get_context, get_parallel +from sglang.srt.runtime_context import get_context, get_parallel, get_server_args _parallel_override = get_parallel().override(attn_tp_size=1) _parallel_override.__enter__() @@ -56,6 +56,8 @@ class GDNAttentionCase: prefix_lens: tuple[int, ...] extend_lens: tuple[int, ...] = () linear_attn_prefill_backend: str | None = None + mis_delimiter_indices: tuple[tuple[int, ...], ...] = () + conv_history_weight: float = 0.0 @property def batch_size(self) -> int: @@ -252,7 +254,7 @@ class MockGDNModelRunner(ModelRunner): dllm_algorithm=None, dllm_algorithm_config=None, enable_deterministic_inference=False, - enable_mis=False, + enable_mis=bool(case.mis_delimiter_indices), linear_attn_backend="triton", linear_attn_decode_backend=None, linear_attn_prefill_backend=case.linear_attn_prefill_backend, @@ -268,7 +270,8 @@ class MockGDNModelRunner(ModelRunner): # derives it from hf_config + page_size, which needs a real model. _mamba_cache_chunk_size=64, ) - self.server_args = self._server_args_override.install() + self._server_args_override.install() + self.server_args = get_server_args() cache_shape = Mamba2StateShape.create( tp_world_size=1, intermediate_size=case.num_v_heads * head_v_dim, @@ -364,6 +367,7 @@ class ProjectedGDNAttention(nn.Module): head_v_dim: int, dtype: torch.dtype, device: str, + conv_history_weight: float = 0.0, ): super().__init__() self.num_k_heads = num_k_heads @@ -372,6 +376,7 @@ class ProjectedGDNAttention(nn.Module): self.head_v_dim = head_v_dim mixed_qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim conv_weights = torch.zeros(mixed_qkv_dim, 2, dtype=dtype, device=device) + conv_weights[:, 0] = conv_history_weight conv_weights[:, 1] = 1 self.A_log = nn.Parameter( torch.randn(num_v_heads, dtype=torch.float32, device=device) * 0.1 @@ -547,6 +552,15 @@ def _make_forward_batch( out_cache_loc=torch.tensor(out_cache_locs, dtype=torch.int64, device=device), seq_lens_sum=sum(seq_lens), positions=torch.tensor(positions, dtype=torch.int64, device=device), + is_prefill_only=bool(case.mis_delimiter_indices), + multi_item_delimiter_indices=( + [ + torch.tensor(indices, dtype=torch.int64) + for indices in case.mis_delimiter_indices + ] + if case.mis_delimiter_indices + else None + ), ) if case.forward_mode.is_extend(include_draft_extend_v2=True): @@ -626,6 +640,7 @@ def build_gdn_attention_fixture( head_v_dim=head_v_dim, dtype=dtype, device=device, + conv_history_weight=case.conv_history_weight, ) reference_module = ReferenceGDNAttention( num_k_heads=case.num_k_heads, @@ -636,6 +651,13 @@ def build_gdn_attention_fixture( device=device, ) _copy_gdn_parameters(actual_module, reference_module) + if case.mis_delimiter_indices: + # Keep recurrent history strong so cross-item state leakage cannot hide + # behind the normal random gate decay in short focused test sequences. + with torch.no_grad(): + actual_module.A_log.fill_(-4.0) + actual_module.dt_bias.fill_(-4.0) + _copy_gdn_parameters(actual_module, reference_module) from .dense_attention import make_loc_fn as _dense_make_loc_fn loc_fn = _dense_make_loc_fn( @@ -785,7 +807,8 @@ def _pure_torch_gdn_reference( initial_ssm_states: torch.Tensor, ) -> GDNReferenceOutput: module = fixture.reference_module - q, k, v = module.split_qkv(fixture.mixed_qkv) + mixed_qkv = _pure_torch_mis_conv_reference(fixture) + q, k, v = module.split_qkv(mixed_qkv) cache_indices = _cache_indices(fixture) g, beta = _pure_torch_gdn_gating(module, fixture.a, fixture.b) q = q.float() @@ -808,29 +831,52 @@ def _pure_torch_gdn_reference( state_idx = cache_indices[req_idx] state = initial_ssm_states[state_idx].float().clone() - for offset in range(input_len): - token_idx = start + offset - for v_head in range(fixture.case.num_v_heads): - k_head = v_head // q_head_ratio - q_vec = q[0, token_idx, k_head] - k_vec = k[0, token_idx, k_head] - v_vec = v[0, token_idx, v_head] + if fixture.case.mis_delimiter_indices: + delimiters = fixture.case.mis_delimiter_indices[req_idx] + segments = [(0, delimiters[0])] + segments.extend(zip(delimiters, tuple(delimiters[1:]) + (input_len,))) + else: + segments = [(0, input_len)] - q_norm = q_vec / torch.sqrt(torch.sum(q_vec * q_vec) + 1e-6) - k_norm = k_vec / torch.sqrt(torch.sum(k_vec * k_vec) + 1e-6) - q_norm = q_norm * (module.head_k_dim**-0.5) + query_final_state = state.clone() + for segment_idx, (segment_start, segment_end) in enumerate(segments): + state = ( + query_final_state.clone() + if segment_idx > 0 + else initial_ssm_states[state_idx].float().clone() + ) + for offset in range(segment_start, segment_end): + token_idx = start + offset + for v_head in range(fixture.case.num_v_heads): + k_head = v_head // q_head_ratio + q_vec = q[0, token_idx, k_head] + k_vec = k[0, token_idx, k_head] + v_vec = v[0, token_idx, v_head] - head_state = state[v_head] - head_state = head_state * torch.exp(g[token_idx, v_head]) - residual_v = v_vec - torch.sum(head_state * k_norm.unsqueeze(0), dim=1) - residual_v = residual_v * beta[token_idx, v_head] - head_state = head_state + residual_v.unsqueeze(1) * k_norm.unsqueeze(0) - state[v_head] = head_state - outputs[0, token_idx, v_head] = torch.sum( - head_state * q_norm.unsqueeze(0), dim=1 - ) + q_norm = q_vec / torch.sqrt(torch.sum(q_vec * q_vec) + 1e-6) + k_norm = k_vec / torch.sqrt(torch.sum(k_vec * k_vec) + 1e-6) + q_norm = q_norm * (module.head_k_dim**-0.5) - final_states[state_idx] = state.to(final_states.dtype) + head_state = state[v_head] + head_state = head_state * torch.exp(g[token_idx, v_head]) + residual_v = v_vec - torch.sum( + head_state * k_norm.unsqueeze(0), dim=1 + ) + residual_v = residual_v * beta[token_idx, v_head] + head_state = head_state + residual_v.unsqueeze( + 1 + ) * k_norm.unsqueeze(0) + state[v_head] = head_state + outputs[0, token_idx, v_head] = torch.sum( + head_state * q_norm.unsqueeze(0), dim=1 + ) + + if segment_idx == 0: + query_final_state = state.clone() + + final_states[state_idx] = ( + query_final_state if fixture.case.mis_delimiter_indices else state + ).to(final_states.dtype) start += input_len return GDNReferenceOutput( @@ -839,6 +885,45 @@ def _pure_torch_gdn_reference( ) +def _pure_torch_mis_conv_reference(fixture: GDNAttentionFixture) -> torch.Tensor: + """Token-by-token causal-conv reference with MIS query-state branching.""" + if fixture.case.conv_history_weight == 0: + return fixture.mixed_qkv + + weights = fixture.actual_module.attn.conv_weights.float() + width = weights.shape[1] + result = torch.empty_like(fixture.mixed_qkv) + request_start = 0 + for request_idx, input_len in enumerate(fixture.case.input_lens): + request_input = fixture.mixed_qkv[ + request_start : request_start + input_len + ].float() + delimiters = fixture.case.mis_delimiter_indices[request_idx] + segments = [(0, delimiters[0])] + segments.extend(zip(delimiters, tuple(delimiters[1:]) + (input_len,))) + + query_state = torch.zeros( + width - 1, + request_input.shape[1], + dtype=torch.float32, + device=request_input.device, + ) + for segment_idx, (segment_start, segment_end) in enumerate(segments): + state = query_state.clone() + for token_offset in range(segment_start, segment_end): + window = torch.cat( + [state, request_input[token_offset : token_offset + 1]] + ) + result[request_start + token_offset] = torch.sum( + window.transpose(0, 1) * weights, dim=1 + ).to(result.dtype) + state = window[1:] + if segment_idx == 0: + query_state = state + request_start += input_len + return result + + def make_gdn_case_with_prefix_lens( case: GDNAttentionCase, name: str, @@ -864,6 +949,8 @@ def make_gdn_case_with_prefix_lens( page_size=case.page_size, prefix_lens=prefix_lens, extend_lens=extend_lens, + mis_delimiter_indices=case.mis_delimiter_indices, + conv_history_weight=case.conv_history_weight, ) @@ -1117,7 +1204,7 @@ def run_gdn_attention_case( expected = _pure_torch_gdn_reference(fixture, initial_ssm_states) torch.testing.assert_close(actual, expected.output, atol=GDN_ATOL, rtol=GDN_RTOL) - if case.forward_mode.is_decode(): + if case.forward_mode.is_decode() or case.mis_delimiter_indices: torch.testing.assert_close( _ssm_states(fixture)[_cache_indices(fixture)], expected.final_states[_cache_indices(fixture)], diff --git a/test/registered/attention/test_chunk_gated_delta_rule.py b/test/registered/attention/test_chunk_gated_delta_rule.py index 70ef9a28c..83cff71ac 100644 --- a/test/registered/attention/test_chunk_gated_delta_rule.py +++ b/test/registered/attention/test_chunk_gated_delta_rule.py @@ -161,6 +161,89 @@ class TestChunkGatedDeltaRule(unittest.TestCase): def test_batch_32(self): self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256) + @unittest.skipUnless( + torch.cuda.is_available(), + "The read-only initial-state path is not implemented by the XPU chunk kernel", + ) + def test_read_only_initial_state_supports_duplicate_indices(self): + """MIS item branches may share one query-end state without updating it.""" + device = get_device() + dtype = torch.bfloat16 + batch_size, tokens_per_item = 3, 65 + num_heads, key_dim, value_dim = 4, 32, 32 + total_tokens = batch_size * tokens_per_item + + torch.manual_seed(1234) + pool_init = torch.randn( + 4, + num_heads, + value_dim, + key_dim, + dtype=torch.float32, + device=device, + ) + duplicate_indices = torch.tensor([2, 2, 2], dtype=torch.int32, device=device) + cu_seqlens = torch.arange( + 0, + total_tokens + 1, + tokens_per_item, + dtype=torch.int32, + device=device, + ) + q = torch.randn(1, total_tokens, num_heads, key_dim, dtype=dtype, device=device) + k = torch.randn_like(q) + v = torch.randn( + 1, total_tokens, num_heads, value_dim, dtype=dtype, device=device + ) + g = torch.nn.functional.logsigmoid( + torch.randn(1, total_tokens, num_heads, dtype=dtype, device=device) + ) + beta = torch.sigmoid( + torch.randn(1, total_tokens, num_heads, dtype=dtype, device=device) + ) + + expected, _ = self._run_reference( + pool_init, duplicate_indices, q, k, v, g, beta + ) + actual_pool = pool_init.clone() + actual, _, _ = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=actual_pool, + initial_state_indices=duplicate_indices, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True, + inplace_update=False, + ) + + torch.testing.assert_close( + actual.float(), expected.float(), atol=self.ATOL, rtol=self.RTOL + ) + self.assertTrue(torch.equal(actual_pool, pool_init)) + + updating_pool = pool_init.clone() + chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=updating_pool, + initial_state_indices=torch.arange( + batch_size, dtype=torch.int32, device=device + ), + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + self.assertFalse( + torch.equal(updating_pool[:batch_size], pool_init[:batch_size]) + ) + # ------------------------------------------------------------------ # Head count sweep # ------------------------------------------------------------------ diff --git a/test/registered/attention/unittests/gdn/test_triton.py b/test/registered/attention/unittests/gdn/test_triton.py index ab00d553b..24be44741 100644 --- a/test/registered/attention/unittests/gdn/test_triton.py +++ b/test/registered/attention/unittests/gdn/test_triton.py @@ -216,6 +216,36 @@ class TestTritonGDNBackendCorrectness(CustomTestCase): with self.subTest(case=case.name, backend=case.backend): run_gdn_attention_case(self, case) + def test_multi_item_scoring_mixed_batch_with_empty_query(self): + case = GDNAttentionCase( + name="gdn_mis_mixed_batch_empty_query", + backend="triton", + forward_mode=ForwardMode.EXTEND, + num_k_heads=2, + num_v_heads=2, + page_size=1, + prefix_lens=(0, 0), + extend_lens=(9, 7), + mis_delimiter_indices=((0, 3, 8), (4, 6)), + conv_history_weight=0.25, + ) + run_gdn_attention_case(self, case) + + def test_multi_item_scoring_crosses_chunk_boundaries(self): + case = GDNAttentionCase( + name="gdn_mis_chunk_boundaries", + backend="triton", + forward_mode=ForwardMode.EXTEND, + num_k_heads=2, + num_v_heads=2, + page_size=1, + prefix_lens=(0,), + extend_lens=(198,), + mis_delimiter_indices=((5, 68, 132, 197),), + conv_history_weight=0.25, + ) + run_gdn_attention_case(self, case, max_context_len=256) + # Layout-robustness. See dense/test_triton.py for the rationale. # shuffled_pages is the default for all tests; this method opts # into the more aggressive interleaved_pages + non_monotonic_extend. diff --git a/test/registered/e2e/gdn/test_qwen35_gdn_multi_item_scoring.py b/test/registered/e2e/gdn/test_qwen35_gdn_multi_item_scoring.py new file mode 100644 index 000000000..fe9737147 --- /dev/null +++ b/test/registered/e2e/gdn/test_qwen35_gdn_multi_item_scoring.py @@ -0,0 +1,133 @@ +"""End-to-end MIS coverage for the hybrid Qwen3.5 GDN architecture.""" + +import asyncio +import os +import unittest + +import torch +from transformers import AutoTokenizer + +from sglang.srt.entrypoints.engine import Engine +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST, + CustomTestCase, +) + +register_cuda_ci(est_time=240, stage="extra-a", runner_config="1-gpu-large") + + +class TestQwen35GDNMultiItemScoring(CustomTestCase): + model = os.environ.get( + "QWEN35_GDN_TEST_MODEL", DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST + ) + atol = 2e-2 + rtol = 2e-2 + + @classmethod + def setUpClass(cls): + cls.engine = Engine( + model_path=cls.model, + trust_remote_code=True, + dtype="bfloat16", + enable_mis=True, + attention_backend="flashinfer", + linear_attn_prefill_backend="triton", + disable_radix_cache=True, + chunked_prefill_size=-1, + mem_fraction_static=0.8, + log_level="error", + ) + tokenizer = AutoTokenizer.from_pretrained(cls.model, trust_remote_code=True) + cls.label_token_ids = [ + tokenizer.encode(label, add_special_tokens=False)[0] + for label in (" yes", " no") + ] + + @classmethod + def tearDownClass(cls): + if getattr(cls, "engine", None) is not None: + cls.engine.shutdown() + torch.cuda.empty_cache() + + def _score(self, query, items): + return self.engine.score( + query=query, + items=items, + label_token_ids=self.label_token_ids, + apply_softmax=False, + ).scores + + async def _async_score(self, query, items): + result = await self.engine.async_score( + query=query, + items=items, + label_token_ids=self.label_token_ids, + apply_softmax=False, + ) + return result.scores + + def _pointwise(self, query, items): + return [self._score(query, [item])[0] for item in items] + + def _assert_scores_close(self, actual, expected): + torch.testing.assert_close( + torch.tensor(actual), + torch.tensor(expected), + atol=self.atol, + rtol=self.rtol, + ) + + def test_batched_matches_pointwise_for_varied_requests(self): + cases = [ + ( + "Decide whether each statement is true:", + ["The sky is blue.", "Two plus two is five.", "Water freezes."], + ), + ("", ["empty query short", "empty query with a much longer item " * 8]), + ("Classify:", ["one"]), + ( + "Judge each passage:", + ["tiny", "medium length passage " * 5, "long passage " * 24], + ), + ] + + for query, items in cases: + with self.subTest(query=query, item_count=len(items)): + self._assert_scores_close( + self._score(query, items), self._pointwise(query, items) + ) + + def test_sibling_changes_and_reordering_do_not_change_target(self): + query = "Rate each answer:" + target = "The target answer remains unchanged." + baseline = self._score(query, [target, "sibling A", "sibling B"])[0] + changed = self._score( + query, [target, "a completely different sibling " * 8, "sibling B"] + )[0] + reordered = self._score(query, ["sibling B", "sibling A", target])[2] + + self._assert_scores_close(changed, baseline) + self._assert_scores_close(reordered, baseline) + + def test_concurrent_requests_match_pointwise(self): + cases = [ + ("Is it an animal?", ["cat", "table", "blue whale"]), + ("", ["alpha", "beta"]), + ("Choose:", ["first"]), + ("Is it a city?", ["Paris", "bread", "Tokyo", "chair"]), + ] + expected = [self._pointwise(query, items) for query, items in cases] + + async def gather_scores(): + return await asyncio.gather( + *(self._async_score(query, items) for query, items in cases) + ) + + actual = self.engine.loop.run_until_complete(gather_scores()) + for actual_scores, expected_scores in zip(actual, expected): + self._assert_scores_close(actual_scores, expected_scores) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/layers/mamba/test_causal_conv1d.py b/test/registered/layers/mamba/test_causal_conv1d.py index 9148f33b8..638da5dd0 100644 --- a/test/registered/layers/mamba/test_causal_conv1d.py +++ b/test/registered/layers/mamba/test_causal_conv1d.py @@ -119,6 +119,63 @@ def causal_conv1d_update_ref( return (out if activation is None else F.silu(out)).to(dtype=dtype_in) +def test_causal_conv1d_branches_from_cloned_query_state(): + device = get_device() + dtype = torch.bfloat16 + torch.manual_seed(7) + dim, width = 96, 4 + item_lens = [2, 5, 7] + total_tokens = sum(item_lens) + + source_state = torch.randn(1, dim, width - 1, device=device, dtype=dtype) + source_state_before = source_state.clone() + branch_states = source_state.expand(len(item_lens), -1, -1).contiguous().clone() + x = torch.randn(dim, total_tokens, device=device, dtype=dtype) + weight = torch.randn(dim, width, device=device, dtype=dtype) + bias = torch.randn(dim, device=device, dtype=dtype) + query_start_loc = torch.tensor( + [0, *torch.tensor(item_lens).cumsum(0).tolist()], + dtype=torch.int32, + device=device, + ) + cache_indices = torch.arange(len(item_lens), dtype=torch.int32, device=device) + + actual = causal_conv1d_fn( + x, + weight, + bias=bias, + conv_states=branch_states, + query_start_loc=query_start_loc, + seq_lens_cpu=torch.tensor(item_lens), + cache_indices=cache_indices, + has_initial_state=torch.ones(len(item_lens), dtype=torch.bool, device=device), + activation="silu", + pad_slot_id=PAD_SLOT_ID, + ) + + expected_outputs = [] + expected_states = [] + for item in torch.split(x, item_lens, dim=-1): + item_output, item_state = causal_conv1d_ref( + item.unsqueeze(0), + weight, + bias, + initial_states=source_state, + return_final_states=True, + activation="silu", + ) + expected_outputs.append(item_output.squeeze(0)) + expected_states.append(item_state.squeeze(0)) + + torch.testing.assert_close( + actual, torch.cat(expected_outputs, dim=-1), atol=5e-2, rtol=1e-2 + ) + torch.testing.assert_close( + branch_states, torch.stack(expected_states), atol=5e-2, rtol=1e-2 + ) + assert torch.equal(source_state, source_state_before) + + @pytest.mark.parametrize("itype", [torch.bfloat16, torch.float]) @pytest.mark.parametrize("silu_activation", [True]) @pytest.mark.parametrize("has_bias", [True]) diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py index 1e2b66f44..012bc8efe 100644 --- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py +++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py @@ -13,6 +13,7 @@ from sglang.srt.layers.attention.linear.gdn_backend import ( GDNKernelDispatcher, _validate_gdn_linear_attn_backends, flashinfer_gdn_prefill_default, + validate_gdn_mis_backend, ) from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( maybe_build_flashinfer_checkpoint_plan, @@ -59,6 +60,7 @@ def make_runner( mamba_radix_cache_strategy="no_buffer", enable_dynamic_chunking=False, chunked_prefill_size=8192, + enable_mis=False, ) fields.update(arg_overrides) args = _publish(testcase, **fields) @@ -79,6 +81,22 @@ def make_runner( class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase): + def test_mis_requires_triton_prefill_backend(self): + runner = make_runner(self, enable_mis=True) + with self.assertRaisesRegex(ValueError, "Triton linear-attention prefill"): + validate_gdn_mis_backend(LinearAttnKernelBackend.FLASHINFER) + + def test_mis_rejects_page_major_layout(self): + make_runner(self, enable_mis=True, enable_page_major_kv_layout=True) + with self.assertRaisesRegex(ValueError, "page-major"): + validate_gdn_mis_backend(LinearAttnKernelBackend.TRITON) + + def test_non_gdn_linear_backend_rejects_mis(self): + with self.assertRaisesRegex(ValueError, "does not support multi-item scoring"): + MambaAttnBackendBase.validate_mis_support(SimpleNamespace(enable_mis=True)) + + GDNAttnBackend.validate_mis_support(SimpleNamespace(enable_mis=True)) + def apply_policy( self, runner, @@ -232,6 +250,7 @@ class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase): backend.kernel_dispatcher = SimpleNamespace(extend_uses_state_checkpoints=True) metadata = SimpleNamespace(has_mamba_track_mask=True, track_ssm_h_src=None) forward_batch = SimpleNamespace( + multi_item_delimiter_indices=None, mamba_track_mask=torch.tensor([True]), mamba_track_indices=torch.tensor([7]), ) diff --git a/test/registered/unit/layers/test_gdn_mis_metadata.py b/test/registered/unit/layers/test_gdn_mis_metadata.py new file mode 100644 index 000000000..653de57b1 --- /dev/null +++ b/test/registered/unit/layers/test_gdn_mis_metadata.py @@ -0,0 +1,99 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.attention.linear.gdn_backend import build_gdn_mis_metadata +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestGDNMISMetadata(CustomTestCase): + def test_allows_trailing_attention_padding(self): + forward_batch = SimpleNamespace( + input_ids=torch.empty(8, dtype=torch.int64), + extend_seq_lens_cpu=[5], + extend_prefix_lens_cpu=[0], + multi_item_delimiter_indices=[torch.tensor([1, 4], dtype=torch.int64)], + is_prefill_only=True, + ) + + metadata = build_gdn_mis_metadata(forward_batch) + + torch.testing.assert_close( + torch.cat([metadata.query_token_indices, metadata.item_token_indices]) + .sort() + .values, + torch.arange(5, dtype=torch.int64), + ) + + def test_mixed_batch_with_empty_query(self): + forward_batch = SimpleNamespace( + input_ids=torch.empty(16, dtype=torch.int64), + extend_seq_lens_cpu=[9, 7], + extend_prefix_lens_cpu=[0, 0], + multi_item_delimiter_indices=[ + torch.tensor([0, 3, 8], dtype=torch.int64), + torch.tensor([4, 6], dtype=torch.int64), + ], + is_prefill_only=True, + ) + + metadata = build_gdn_mis_metadata(forward_batch) + + self.assertEqual(metadata.query_seq_lens_cpu, [4]) + self.assertEqual(metadata.item_seq_lens_cpu, [3, 5, 1, 2, 1]) + torch.testing.assert_close( + metadata.query_token_indices, + torch.tensor([9, 10, 11, 12], dtype=torch.int64), + ) + torch.testing.assert_close( + metadata.query_cu_seqlens, + torch.tensor([0, 4], dtype=torch.int32), + ) + torch.testing.assert_close( + metadata.query_request_indices, + torch.tensor([1], dtype=torch.int64), + ) + torch.testing.assert_close( + metadata.item_token_indices, + torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 13, 14, 15], dtype=torch.int64), + ) + torch.testing.assert_close( + metadata.item_cu_seqlens, + torch.tensor([0, 3, 8, 9, 11, 12], dtype=torch.int32), + ) + torch.testing.assert_close( + metadata.item_request_indices, + torch.tensor([0, 0, 0, 1, 1], dtype=torch.int64), + ) + + def test_rejects_sequence_lengths_beyond_input(self): + forward_batch = SimpleNamespace( + input_ids=torch.empty(4, dtype=torch.int64), + extend_seq_lens_cpu=[5], + extend_prefix_lens_cpu=[0], + multi_item_delimiter_indices=[torch.tensor([1, 4], dtype=torch.int64)], + is_prefill_only=True, + ) + + with self.assertRaisesRegex(ValueError, "exceed the input tokens"): + build_gdn_mis_metadata(forward_batch) + + def test_rejects_missing_final_delimiter(self): + forward_batch = SimpleNamespace( + input_ids=torch.empty(5, dtype=torch.int64), + extend_seq_lens_cpu=[5], + extend_prefix_lens_cpu=[0], + multi_item_delimiter_indices=[torch.tensor([2, 3], dtype=torch.int64)], + is_prefill_only=True, + ) + + with self.assertRaisesRegex(ValueError, "final delimiter"): + build_gdn_mis_metadata(forward_batch) + + +if __name__ == "__main__": + unittest.main()