diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index ce8c4f1bc92b..d9add5e3fdc1 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 53b271a1c884..4cc65791ed49 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 diff --git a/test/registered/layers/mamba/test_causal_conv1d.py b/test/registered/layers/mamba/test_causal_conv1d.py index 638da5dd0d04..7f0ba08fa996 100644 --- a/test/registered/layers/mamba/test_causal_conv1d.py +++ b/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