Fix Qwen3.5 GDN multi-item scoring (#33922)

Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

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 ...")

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)),

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"
     )

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

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,
         *,

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(

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)],

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
     # ------------------------------------------------------------------

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.

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()

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])

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]),
         )

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()