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