[Perf] mamba: write the deduplicated conv verify window as one tile

python/sglang/kernels/ops/mamba/causal_conv1d_triton.py

@@ -664,6 +664,8 @@ def _causal_conv1d_update_kernel(
     SAVE_INTERMEDIATE: tl.constexpr,
     HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr,
     USE_GDC: tl.constexpr = False,
+    DEDUP_CONV_WINDOW: tl.constexpr = False,
+    NP2_WINDOW_COLS: tl.constexpr = 1,
 ):
     # ruff: noqa: E501
     if USE_GDC:
@@ -832,6 +834,34 @@ def _causal_conv1d_update_kernel(
     x_base_1d = x_base  # starting of chunk [BLOCK_N]
     mask_x_1d = idx_feats < dim
 
+    if SAVE_INTERMEDIATE and DEDUP_CONV_WINDOW:
+        tl.static_assert(KERNEL_WIDTH == 4)
+        tl.static_assert(not HAS_EAGLE_TREE_CUSTOM_ATTN_MASK)
+        # Deduplicated window (KERNEL_WIDTH == 4, linear chain): step t's window is
+        # columns t..t+2 of one shared per-channel row
+        # [state[1], state[2], x_0, ..., x_{seqlen-1}]. Write those seqlen + 2
+        # columns as one [BLOCK_N, NP2_WINDOW_COLS] tile instead of three
+        # lane-strided stores per token; the bytes are the same.
+        idx_c = tl.arange(0, NP2_WINDOW_COLS)
+        x_win = tl.load(
+            x_base[:, None] + ((idx_c - 2) * stride_x_token)[None, :],
+            mask_w[:, None] & ((idx_c >= 2) & (idx_c < seqlen + 2))[None, :],
+            other=0.0,
+        )
+        win = tl.where(
+            (idx_c == 0)[None, :],
+            col1[:, None],
+            tl.where((idx_c == 1)[None, :], col2[:, None], x_win),
+        )
+        tl.store(
+            intermediate_conv_window_ptr
+            + intermediate_state_batch_coord * stride_inter_seq
+            + (idx_feats * stride_inter_dim)[:, None]
+            + (idx_c * stride_inter_win)[None, :],
+            win,
+            mask=mask_w[:, None] & (idx_c < seqlen + 2)[None, :],
+        )
+
     # STEP 5: compute each token
     for idx_token in tl.static_range(seqlen):
         acc = acc_preload
@@ -977,7 +1007,7 @@ def _causal_conv1d_update_kernel(
                 col1 = col2
                 col2 = matrix_x
 
-            if SAVE_INTERMEDIATE:
+            if SAVE_INTERMEDIATE and not DEDUP_CONV_WINDOW:
                 # Save the window state after consuming this token
                 # Layout: [seq(cache line), step, dim, win(K-1)]
                 base_ptr = (
@@ -1135,6 +1165,19 @@ def grid(META):
         )
     else:
         stride_inter_seq = stride_inter_step = stride_inter_dim = stride_inter_win = 0
+    # In a width-4 linear chain, step t's window is columns t..t+2 of the row
+    # [state[1], state[2], x_0, ..., x_{seqlen-1}]. When the window view's step
+    # stride equals its window stride s, the per-token stores write column t + w
+    # of that row at (t + w) * s, and one store of the whole row writes column c
+    # at c * s: the same addresses and values for any layout with these strides
+    # (MambaPool's deduplicated view is one). The kernel then writes the row as
+    # one tile. Tree verify and other widths keep the per-token stores.
+    dedup_conv_window = (
+        intermediate_conv_window is not None
+        and retrieve_next_token is None
+        and width == 4
+        and stride_inter_step == stride_inter_win
+    )
 
     # prepare retrieve next token buffer strides if provided
     if retrieve_next_token is not None:
@@ -1224,6 +1267,10 @@ def grid(META):
         BLOCK_N=256,
         SAVE_INTERMEDIATE=intermediate_conv_window is not None,
         HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_next_token is not None,
+        DEDUP_CONV_WINDOW=dedup_conv_window,
+        NP2_WINDOW_COLS=(
+            triton.next_power_of_2(seqlen + width - 2) if dedup_conv_window else 1
+        ),
         **pdl_kwargs,
     )
     if unsqueeze:

python/sglang/srt/mem_cache/memory_pool.py

@@ -902,9 +902,10 @@ def __init__(
                 # logical shape [num_layers, size+1, draft_tokens, dim, K-1] where
                 # step `t`'s window is the slice shared[..., :, t:t+K-1]. This
                 # halves the conv-intermediate footprint (D*(K-1) -> D+K-2 columns)
-                # with no numerical change: both the conv kernel write (idempotent
-                # overlapping stores) and `fused_conv_window_scatter_with_mask`
-                # consume the view through its strides.
+                # with no numerical change: both the conv kernel write (one tile
+                # store per row for width-4 chains, idempotent overlapping stores
+                # otherwise) and `fused_conv_window_scatter_with_mask` consume the
+                # view through its strides.
                 #
                 # Dedup the sliding-window conv-intermediate only when it is safe:
                 # CUDA + a linear draft chain (topk <= 1). NPU/CPU and EAGLE tree

test/registered/layers/mamba/test_causal_conv1d.py

@@ -491,6 +491,111 @@ def test_causal_conv1d_varlen_mixed_input_and_state_dtype():
     torch.testing.assert_close(conv_states, conv_states_ref, rtol=1e-2, atol=5e-2)
 
 
+@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(
+    "seqlen, case",
+    [
+        # T + 2 window columns pad to tiles of 4, 8, 16 and 32; T = 3 is odd.
+        (2, "plain"),
+        (3, "plain"),
+        (8, "plain"),
+        (16, "plain"),
+        # Row i's prior state starts num_accept_tokens[i] - 1 columns in.
+        (4, "accept_offset"),
+        # MambaPool's window axis 0: a [T + K - 2, dim] row, seen transposed.
+        (4, "window_axis_0"),
+    ],
+)
+@pytest.mark.parametrize("with_padding", [True, False])
+def test_causal_conv1d_update_dedup_window_matches_per_token(
+    with_padding, case, seqlen, itype
+):
+    """The deduplicated verify window is written as one tile; it must hold the
+    same bytes as the per-token stores into a dense window and leave the slots of
+    padded rows untouched."""
+    device = get_device()
+    torch.manual_seed(0)
+    # dim is not a multiple of the 256-channel block, so the last block is partial.
+    batch, dim, width, num_slots = 6, 1000, 4, 16
+    # Verify tokens are token-major, as in the GDN and Mamba verify reshape.
+    x = torch.randn(batch, seqlen, dim, device=device, dtype=itype).transpose(1, 2)
+    weight = torch.randn(dim, width, device=device, dtype=itype)
+    bias = torch.randn(dim, device=device, dtype=itype)
+    num_accept_tokens = None
+    state_len = width - 1
+    if case == "accept_offset":
+        num_accept_tokens = torch.randint(
+            1, seqlen + 1, (batch,), device=device, dtype=torch.int32
+        )
+        state_len += seqlen - 1
+    conv_state = torch.randn(num_slots, dim, state_len, device=device, dtype=itype)
+    slots = torch.randperm(num_slots, device=device)[: 2 * batch].to(torch.int32)
+    conv_state_indices, window_indices = slots[:batch].clone(), slots[batch:]
+    if with_padding:
+        conv_state_indices[0] = conv_state_indices[-1] = PAD_SLOT_ID
+    real = conv_state_indices != PAD_SLOT_ID
+
+    # MambaPool._allocate_deduplicated_conv_window: one row of T + K - 2 columns
+    # per slot and channel, viewed as [T, dim, K - 1] windows with step stride ==
+    # window stride.
+    num_cols = seqlen + width - 2
+    if case == "window_axis_0":
+        shared = torch.randn(num_slots, num_cols, dim, device=device, dtype=itype)
+        col_stride, dim_stride = shared.stride(1), shared.stride(2)
+    else:
+        shared = torch.randn(num_slots, dim, num_cols, device=device, dtype=itype)
+        dim_stride, col_stride = shared.stride(1), shared.stride(2)
+    dedup = shared.as_strided(
+        (num_slots, seqlen, dim, width - 1),
+        (shared.stride(0), col_stride, dim_stride, col_stride),
+    )
+    dense = torch.randn(num_slots, seqlen, dim, width - 1, device=device, dtype=itype)
+    shared_before, dense_before = shared.clone(), dense.clone()
+
+    def run(window):
+        state = conv_state.clone()
+        out = causal_conv1d_update(
+            x,
+            state,
+            weight,
+            bias,
+            activation="silu",
+            conv_state_indices=conv_state_indices,
+            num_accept_tokens=num_accept_tokens,
+            intermediate_conv_window=window,
+            intermediate_state_indices=window_indices,
+            pad_slot_id=PAD_SLOT_ID,
+        )
+        return out, state
+
+    out_tile, state_tile = run(dedup)
+    out_ref, state_ref = run(dense)
+
+    # Each real row's shared window row is [state[o + 1], state[o + 2], x_0, ...,
+    # x_{T-1}] with o = num_accept_tokens - 1 (0 without num_accept_tokens).
+    offset = torch.zeros(batch, dtype=torch.long, device=device)
+    if num_accept_tokens is not None:
+        offset = num_accept_tokens.long() - 1
+    prior_cols = offset[real, None] + torch.arange(1, width - 1, device=device)
+    prior = conv_state[conv_state_indices[real].long()].gather(
+        2, prior_cols[:, None, :].expand(-1, dim, -1)
+    )
+    rows = torch.cat([prior, x[real]], dim=-1)
+    written = window_indices[real].long()
+    expected_shared = shared_before.clone()
+    expected_shared[written] = rows.transpose(1, 2) if case == "window_axis_0" else rows
+    expected_dense = dense_before.clone()
+    expected_dense[written] = rows.unfold(-1, width - 1, 1).transpose(1, 2)
+
+    # Windows and conv states are copies, so they must match bitwise.
+    assert torch.equal(shared, expected_shared)
+    assert torch.equal(dense, expected_dense)
+    assert torch.equal(dedup[written], dense[written])
+    assert torch.equal(state_tile, state_ref)
+    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (1e-2, 5e-2)
+    torch.testing.assert_close(out_tile[real], out_ref[real], rtol=rtol, atol=atol)
+
+
 if __name__ == "__main__":
     import sys