diff --git a/python/sglang/kernels/aot/csrc/common_extension.cc b/python/sglang/kernels/aot/csrc/common_extension.cc index 39d9a6a2c..211bbdc05 100644 --- a/python/sglang/kernels/aot/csrc/common_extension.cc +++ b/python/sglang/kernels/aot/csrc/common_extension.cc @@ -438,6 +438,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { /* * From csrc/mamba */ + // Compatibility API: SGLang dispatches to the JIT implementation, but external + // sgl_kernel consumers still rely on these exported CUDA ops. m.def( "causal_conv1d_update(Tensor! x," "Tensor! conv_state," diff --git a/python/sglang/kernels/aot/tests/test_causal_conv1d.py b/python/sglang/kernels/aot/tests/test_causal_conv1d.py deleted file mode 100644 index 0a1be8b63..000000000 --- a/python/sglang/kernels/aot/tests/test_causal_conv1d.py +++ /dev/null @@ -1,492 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from https://github.com/vllm-project/vllm/blob/main/tests/kernels/mamba/test_causal_conv1d.py -import sys -from typing import Optional - -import torch -from sgl_kernel import causal_conv1d_fwd -from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel - -PAD_SLOT_ID = -1 - - -def causal_conv1d_fn( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - query_start_loc: Optional[torch.Tensor] = None, - cache_indices: Optional[torch.Tensor] = None, - has_initial_state: Optional[torch.Tensor] = None, - conv_states: Optional[torch.Tensor] = None, - activation: Optional[str] = "silu", - pad_slot_id: int = PAD_SLOT_ID, -): - """ - x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen - sequences are concatenated from left to right for varlen - weight: (dim, width) - bias: (dim,) - query_start_loc: (batch + 1) int32 - The cumulative sequence lengths of the sequences in - the batch, used to index into sequence. prepended by 0. - for example: query_start_loc = torch.Tensor([0,10,16,17]), - x.shape=(dim,17) - cache_indices: (batch) int32 - indicates the corresponding state index, - like so: conv_state = conv_states[cache_indices[batch_id]] - has_initial_state: (batch) bool - indicates whether should the kernel take the current state as initial - state for the calculations - conv_states: (...,dim,width - 1) itype - updated inplace if provided - activation: either None or "silu" or "swish" - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padded - entries that will not be processed, - for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] - in this case, the kernel will not process entries at - indices 0 and 3 - out: (batch, dim, seqlen) - """ - if activation not in [None, "silu", "swish"]: - raise NotImplementedError("activation must be None, silu, or swish") - if x.stride(-1) != 1: - x = x.contiguous() - bias = bias.contiguous() if bias is not None else None - - causal_conv1d_fwd( - x, - weight, - bias, - conv_states, - query_start_loc, - cache_indices, - has_initial_state, - activation in ["silu", "swish"], - pad_slot_id, - ) - return x - - -def causal_conv1d_update( - x: torch.Tensor, - conv_state: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - activation: Optional[str] = None, - cache_seqlens: Optional[torch.Tensor] = None, - conv_state_indices: Optional[torch.Tensor] = None, - pad_slot_id: int = PAD_SLOT_ID, -): - """ - x: (batch, dim) or (batch, dim, seqlen) - conv_state: (batch, dim, state_len), where state_len >= width - 1 - weight: (dim, width) - bias: (dim,) - cache_seqlens: (batch,), dtype int32. - If not None, the conv_state is treated as a circular buffer. - The conv_state will be updated by copying x to the conv_state - starting at the index - @cache_seqlens % state_len. - conv_state_indices: (batch,), dtype int32 - If not None, the conv_state is a larger tensor along the batch dim, - and we are selecting the batch coords specified by conv_state_indices. - Useful for a continuous batching scenario. - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padded - entries that will not be processed, - for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id] - in this case, the kernel will not process entries at - indices 0 and 3 - out: (batch, dim) or (batch, dim, seqlen) - """ - if activation not in [None, "silu", "swish"]: - raise NotImplementedError( - f"activation must be None, silu, or swish, actual: {activation}" - ) - activation_val = activation in ["silu", "swish"] - unsqueeze = x.dim() == 2 - if unsqueeze: - x = x.unsqueeze(-1) - causal_conv1d_update_kernel( - x, - conv_state, - weight, - bias, - activation_val, - cache_seqlens, - conv_state_indices, - pad_slot_id, - ) - if unsqueeze: - x = x.squeeze(-1) - return x - - -# SPDX-License-Identifier: Apache-2.0 - -from typing import Optional - -import pytest -import torch -import torch.nn.functional as F - - -def causal_conv1d_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - initial_states: Optional[torch.Tensor] = None, - return_final_states: bool = False, - final_states_out: Optional[torch.Tensor] = None, - activation: Optional[str] = "silu", -): - """ - x: (batch, dim, seqlen) - weight: (dim, width) - bias: (dim,) - initial_states: (batch, dim, width - 1) - final_states_out: (batch, dim, width - 1) - - out: (batch, dim, seqlen) - """ - if activation not in [None, "silu", "swish"]: - raise NotImplementedError("activation must be None, silu, or swish") - dtype_in = x.dtype - x = x.to(weight.dtype) - seqlen = x.shape[-1] - dim, width = weight.shape - if initial_states is None: - out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) - else: - x = torch.cat([initial_states, x], dim=-1) - out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) - out = out[..., :seqlen] - if return_final_states: - final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( - dtype_in - ) # (batch, dim, width - 1) - if final_states_out is not None: - final_states_out.copy_(final_states) - else: - final_states_out = final_states - out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) - return (out, None) if not return_final_states else (out, final_states_out) - - -def causal_conv1d_update_ref( - x, conv_state, weight, bias=None, activation=None, cache_seqlens=None -): - """ - x: (batch, dim) or (batch, dim, seqlen) - conv_state: (batch, dim, state_len), where state_len >= width - 1 - weight: (dim, width) - bias: (dim,) - cache_seqlens: (batch,), dtype int32. - If not None, the conv_state is treated as a circular buffer. - The conv_state will be updated by copying x to the - conv_state starting at the index - @cache_seqlens % state_len before performing the convolution. - - out: (batch, dim) or (batch, dim, seqlen) - """ - if activation not in [None, "silu", "swish"]: - raise NotImplementedError("activation must be None, silu, or swish") - dtype_in = x.dtype - unsqueeze = x.dim() == 2 - if unsqueeze: - x = x.unsqueeze(-1) - batch, dim, seqlen = x.shape - width = weight.shape[1] - state_len = conv_state.shape[-1] - assert conv_state.shape == (batch, dim, state_len) - assert weight.shape == (dim, width) - if cache_seqlens is None: - x_new = torch.cat([conv_state, x], dim=-1).to( - weight.dtype - ) # (batch, dim, state_len + seqlen) - conv_state.copy_(x_new[:, :, -state_len:]) - else: - width_idx = torch.arange( - -(width - 1), 0, dtype=torch.long, device=x.device - ).unsqueeze(0) + cache_seqlens.unsqueeze(1) - width_idx = ( - torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) - ) - x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) - copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze( - 0 - ) + cache_seqlens.unsqueeze(1) - copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) - conv_state.scatter_(2, copy_idx, x) - out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[ - :, :, -seqlen: - ] - if unsqueeze: - out = out.squeeze(-1) - return (out if activation is None else F.silu(out)).to(dtype=dtype_in) - - -@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float]) -@pytest.mark.parametrize("silu_activation", [True]) -@pytest.mark.parametrize("has_bias", [True]) -@pytest.mark.parametrize("has_initial_state", [True, False]) -@pytest.mark.parametrize("width", [4]) -@pytest.mark.parametrize( - "seqlen", [1, 8, 16, 32, 64, 128, 256, 512, 784, 1024, 1025, 2048, 4096] -) -@pytest.mark.parametrize("dim", [64]) -@pytest.mark.parametrize("batch", [1]) -def test_causal_conv1d( - batch, dim, seqlen, width, has_bias, silu_activation, has_initial_state, itype -): - device = "cuda" - rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) - if itype == torch.bfloat16: - rtol, atol = 1e-2, 5e-2 - x = torch.randn(batch, dim, seqlen, device=device, dtype=itype).contiguous() - - weight = torch.randn(dim, width, device=device, dtype=itype) - bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None - if has_initial_state: - initial_states = torch.randn(batch, dim, width - 1, device=device, dtype=itype) - has_initial_state_tensor = torch.ones(batch, dtype=torch.bool, device=x.device) - else: - initial_states = None - has_initial_state_tensor = None - x_ref = x.clone() - weight_ref = weight.clone() - bias_ref = bias.clone() if bias is not None else None - initial_states_ref = initial_states.clone() if initial_states is not None else None - activation = None if not silu_activation else "silu" - out = causal_conv1d_fn( - x, - weight, - bias, - activation=activation, - conv_states=initial_states, - has_initial_state=has_initial_state_tensor, - ) - out_ref, final_states_ref = causal_conv1d_ref( - x_ref, - weight_ref, - bias_ref, - initial_states=initial_states_ref, - return_final_states=True, - activation=activation, - ) - if has_initial_state: - assert initial_states is not None and final_states_ref is not None - assert torch.allclose(initial_states, final_states_ref, rtol=rtol, atol=atol) - assert torch.allclose(out, out_ref, rtol=rtol, atol=atol) - - -@pytest.mark.parametrize("itype", [torch.bfloat16]) -@pytest.mark.parametrize("silu_activation", [False, True]) -@pytest.mark.parametrize("has_bias", [False, True]) -@pytest.mark.parametrize("seqlen", [1]) -@pytest.mark.parametrize("width", [4]) -@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) -def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, itype): - device = "cuda" - rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) - if itype == torch.bfloat16: - rtol, atol = 1e-2, 5e-2 - - batch = 2 - x = torch.randn(batch, dim, seqlen, device=device, dtype=itype) - x_ref = x.clone() - conv_state = torch.randn(batch, dim, width - 1, device=device, dtype=itype) - - weight = torch.randn(dim, width, device=device, dtype=itype) - bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None - conv_state_ref = conv_state.detach().clone() - activation = None if not silu_activation else "silu" - out = causal_conv1d_update(x, conv_state, weight, bias, activation=activation) - out_ref = causal_conv1d_update_ref( - x_ref, conv_state_ref, weight, bias, activation=activation - ) - - assert torch.equal(conv_state, conv_state_ref) - assert torch.allclose(out, out_ref, rtol=rtol, atol=atol) - - -@pytest.mark.parametrize("itype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("silu_activation", [False, True]) -@pytest.mark.parametrize("has_bias", [False, True]) -@pytest.mark.parametrize("seqlen", [1, 4, 5]) -@pytest.mark.parametrize("width", [2, 3, 4]) -@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) -# tests correctness in case subset of the sequences are padded -@pytest.mark.parametrize("with_padding", [True, False]) -def test_causal_conv1d_update_with_batch_gather( - with_padding, dim, width, seqlen, has_bias, silu_activation, itype -): - device = "cuda" - rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) - if itype == torch.bfloat16: - rtol, atol = 1e-2, 5e-2 - - batch_size = 3 - padding = 5 if with_padding else 0 - padded_batch_size = batch_size + padding - total_entries = 10 * batch_size - - x = torch.randn(padded_batch_size, dim, 1, device=device, dtype=itype) - x_ref = x.clone() - - conv_state_indices = torch.randperm(total_entries)[:batch_size].to( - dtype=torch.int32, device=device - ) - unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device) - unused_states_bool[conv_state_indices] = False - padded_state_indices = torch.concat( - [ - conv_state_indices, - torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device), - ], - dim=0, - ) - conv_state = torch.randn(total_entries, dim, width - 1, device=device, dtype=itype) - conv_state_for_padding_test = conv_state.clone() - - weight = torch.randn(dim, width, device=device, dtype=itype) - bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None - conv_state_ref = conv_state[conv_state_indices, :].detach().clone() - activation = None if not silu_activation else "silu" - out = causal_conv1d_update( - x, - conv_state, - weight, - bias, - activation=activation, - conv_state_indices=padded_state_indices, - pad_slot_id=PAD_SLOT_ID, - ) - out_ref = causal_conv1d_update_ref( - x_ref[:batch_size], conv_state_ref, weight, bias, activation=activation - ) - - assert torch.equal(conv_state[conv_state_indices, :], conv_state_ref) - assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol) - assert torch.equal( - conv_state[unused_states_bool], conv_state_for_padding_test[unused_states_bool] - ) - - -@pytest.mark.parametrize("itype", [torch.bfloat16]) -@pytest.mark.parametrize("silu_activation", [True]) -@pytest.mark.parametrize("has_bias", [True]) -@pytest.mark.parametrize("width", [4]) -@pytest.mark.parametrize( - "seqlen", [8, 16, 32, 64, 128, 256, 512, 784, 1024, 2048, 2049, 4096] -) -@pytest.mark.parametrize("dim", [64, 4096]) -# tests correctness in case subset of the sequences are padded -@pytest.mark.parametrize("with_padding", [True, False]) -def test_causal_conv1d_varlen( - with_padding, dim, seqlen, width, has_bias, silu_activation, itype -): - device = "cuda" - torch.cuda.empty_cache() - rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) - if itype == torch.bfloat16: - rtol, atol = 1e-2, 5e-2 - - seqlens = [] - batch_size = 4 - if seqlen < 10: - batch_size = 1 - padding = 3 if with_padding else 0 - padded_batch_size = batch_size + padding - nsplits = padded_batch_size - 1 - - eos_pos = torch.randperm(seqlen - 1)[:nsplits].sort().values - seqlens.append( - torch.diff( - torch.cat([torch.tensor([-1]), eos_pos, torch.tensor([seqlen - 1])]) - ).tolist() - ) - assert sum(seqlens[-1]) == seqlen - assert all(s > 0 for s in seqlens[-1]) - - total_entries = batch_size * 10 - cumsum = torch.cumsum(torch.tensor(seqlens[0]), dim=0).to(torch.int32) - cumsum = torch.concat([torch.tensor([0], dtype=torch.int32), cumsum], dim=0) - x = torch.randn(1, 4096 + dim + 64, seqlen, device=device, dtype=itype)[ - :, 4096 : 4096 + dim, : - ] - weight = torch.randn(dim, width, device=device, dtype=itype) - bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None - x_ref = x.clone() - weight_ref = weight.clone() - bias_ref = bias.clone() if bias is not None else None - activation = None if not silu_activation else "silu" - final_states = torch.randn( - total_entries, dim, width - 1, device=x.device, dtype=x.dtype - ) - final_states_ref = final_states.clone() - has_initial_states = torch.randint( - 0, 2, (cumsum.shape[0] - 1,), dtype=torch.bool, device=x.device - ) - state_indices = torch.randperm(total_entries, dtype=torch.int32, device=x.device)[ - :batch_size - ] - padded_state_indices = torch.concat( - [ - state_indices, - torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device), - ], - dim=-1, - ) - - out = causal_conv1d_fn( - x.squeeze(0), - weight, - bias, - cumsum.cuda(), - padded_state_indices, - has_initial_states, - final_states, - activation, - PAD_SLOT_ID, - ) - out_ref = [] - out_ref_b = [] - - splits = [torch.split(var, seqlens[0], dim=-1) for var in (x_ref)] - for i in range(len(seqlens[0])): - x_s = [v[i].unsqueeze(0) for v in splits][0] - if padded_state_indices[i] == PAD_SLOT_ID: - continue - out_ref_b.append( - causal_conv1d_ref( - x_s, - weight_ref, - bias_ref, - activation=activation, - return_final_states=True, - final_states_out=final_states_ref[padded_state_indices[i]].unsqueeze(0), - initial_states=( - final_states_ref[padded_state_indices[i]].unsqueeze(0) - if has_initial_states[i] - else None - ), - ) - ) - out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=2)) - out_ref_tensor = torch.cat(out_ref, dim=0) - - unpadded_out = out[:, : out_ref_tensor.shape[-1]] - assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol) - assert torch.allclose( - final_states[state_indices], - final_states_ref[state_indices], - rtol=rtol, - atol=atol, - ) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__])) diff --git a/python/sglang/kernels/jit/csrc/mamba/causal_conv1d.cuh b/python/sglang/kernels/jit/csrc/mamba/causal_conv1d.cuh new file mode 100644 index 000000000..79893cf28 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/mamba/causal_conv1d.cuh @@ -0,0 +1,767 @@ +/// \file causal_conv1d.cuh +/// \brief Depthwise causal conv1d: prefill (`causal_conv1d_fwd`) and decode +/// (`causal_conv1d_update`). +/// +/// Adapted from +/// https://github.com/Dao-AILab/causal-conv1d/blob/main/csrc/causal_conv1d_fwd.cu +/// and +/// https://github.com/Dao-AILab/causal-conv1d/blob/main/csrc/causal_conv1d_update.cu +/// +/// The device kernels are carried over unchanged from the AOT implementation, so +/// results stay bit-identical; only the host launchers are adapted to the tvm-ffi +/// `TensorView` API. + +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For CHECK_HOST, div_ceil + +#include // For DTypeTrait, fp16_t / bf16_t / fp32_t +#include // For LaunchKernel, SGL_DEVICE + +#include +#include +#include +#include +#include + +#include +#include + +namespace sglang { + +/// Nested so these generic names cannot collide with the unrelated +/// `causal_conv1d` of the Inkling short-conv kernels. +namespace mamba_conv { + +/// \brief Runtime parameters shared by the prefill and decode kernels. +/// +/// The subset of the AOT `ConvParamsBase` these kernels read. Strides are in +/// elements, not bytes, and stay `uint32_t` as they were there. +struct MambaConvParams { + using index_t = uint32_t; + + int32_t batch; + int32_t dim; + int32_t seqlen; + int32_t width; + int64_t pad_slot_id; + bool silu_activation; + + index_t x_batch_stride; + index_t x_c_stride; + index_t x_l_stride; + index_t weight_c_stride; + index_t weight_width_stride; + index_t out_batch_stride; + index_t out_c_stride; + index_t out_l_stride; + + int32_t conv_state_len; + index_t conv_state_batch_stride; + index_t conv_state_c_stride; + index_t conv_state_l_stride; + + // Common data pointers. + void* __restrict__ x_ptr; + void* __restrict__ weight_ptr; + void* __restrict__ bias_ptr; + void* __restrict__ out_ptr; + + void* __restrict__ conv_state_ptr; + void* __restrict__ query_start_loc_ptr; + void* __restrict__ has_initial_state_ptr; + void* __restrict__ cache_indices_ptr; + const int32_t* __restrict__ cache_seqlens; + + // For the continuous batching case: the conv state for the current batch does + // not need to be a contiguous tensor. + const int32_t* __restrict__ conv_state_indices_ptr; + + void* conv_states_ptr; + index_t conv_states_batch_stride; + index_t conv_states_c_stride; + index_t conv_states_l_stride; +}; + +/// \brief The unsigned integer type occupying exactly `kBytes` bytes. +template +struct BytesToType {}; + +template <> +struct BytesToType<16> { + using Type = uint4; +}; +template <> +struct BytesToType<8> { + using Type = uint64_t; +}; +template <> +struct BytesToType<4> { + using Type = uint32_t; +}; +template <> +struct BytesToType<2> { + using Type = uint16_t; +}; + +/// \brief Convert a stored element to fp32 through the dtype-specific intrinsic. +template +SGL_DEVICE fp32_t to_float(const T& value) { + return DTypeTrait::from(value); +} + +/// \brief Convert an fp32 accumulator back to the stored element type. +template +SGL_DEVICE T from_float(fp32_t value) { + return DTypeTrait::from(value); +} + +/// \brief Zero every element of a stored-element array. +template +SGL_DEVICE void zero_fill(T (&values)[kN]) { +#pragma unroll + for (int i = 0; i < kN; ++i) { + values[i] = from_float(0.0f); + } +} + +//////////////////////////////////////////////////////////////////////////////// +// Prefill +//////////////////////////////////////////////////////////////////////////////// + +/// \brief Compile-time configuration of the prefill kernel. +/// +/// \tparam kNThreads_ Threads per CTA. +/// \tparam kWidth_ Convolution width (2..4). +/// \tparam kIsVecLoad_ Whether the chunk can be loaded/stored as whole vectors. +/// \tparam T Element type: fp16_t | bf16_t | fp32_t. +template +struct CausalConv1dFwdTraits { + using input_t = T; + static constexpr int kNThreads = kNThreads_; + static constexpr int kWidth = kWidth_; + static constexpr int kNBytes = sizeof(T); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = kNBytes == 4 ? 4 : 8; + static_assert(kWidth <= kNElts); + static constexpr bool kIsVecLoad = kIsVecLoad_; + using vec_t = typename BytesToType::Type; + using BlockLoadT = cub::BlockLoad; + using BlockLoadVecT = cub::BlockLoad; + using BlockStoreT = cub::BlockStore; + using BlockStoreVecT = cub::BlockStore; + static constexpr int kSmemIOSize = + kIsVecLoad ? 0 + : static_cast( + std::max(sizeof(typename BlockLoadT::TempStorage), sizeof(typename BlockStoreT::TempStorage))); + static constexpr int kSmemExchangeSize = kNThreads * kNBytes * kNElts; + static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize; +}; + +/// \brief One CTA per (sequence, channel): convolve a channel over its sequence +/// and write back the trailing `kWidth - 1` taps as the new conv state. +template +__global__ __launch_bounds__(Ktraits::kNThreads) void causal_conv1d_fwd_kernel(MambaConvParams params) { + constexpr int kWidth = Ktraits::kWidth; + constexpr int kNThreads = Ktraits::kNThreads; + constexpr int kNElts = Ktraits::kNElts; + constexpr bool kIsVecLoad = Ktraits::kIsVecLoad; + using input_t = typename Ktraits::input_t; + using vec_t = typename Ktraits::vec_t; + + // Shared memory. + extern __shared__ char smem_[]; + auto& smem_load = reinterpret_cast(smem_); + auto& smem_load_vec = reinterpret_cast(smem_); + auto& smem_store = reinterpret_cast(smem_); + auto& smem_store_vec = reinterpret_cast(smem_); + vec_t* smem_exchange = reinterpret_cast(smem_ + Ktraits::kSmemIOSize); + + const bool kVarlen = params.query_start_loc_ptr != nullptr; + const int32_t tidx = threadIdx.x; + const int32_t batch_id = blockIdx.x; + const int32_t channel_id = blockIdx.y; + const int32_t* query_start_loc = kVarlen ? reinterpret_cast(params.query_start_loc_ptr) : nullptr; + const int32_t sequence_start_index = kVarlen ? query_start_loc[batch_id] : batch_id; + const int32_t seqlen = kVarlen ? query_start_loc[batch_id + 1] - sequence_start_index : params.seqlen; + + input_t* x = reinterpret_cast(params.x_ptr) + sequence_start_index * params.x_batch_stride + + channel_id * params.x_c_stride; + const input_t* weight = reinterpret_cast(params.weight_ptr) + channel_id * params.weight_c_stride; + input_t* out = reinterpret_cast(params.out_ptr) + sequence_start_index * params.out_batch_stride + + channel_id * params.out_c_stride; + const float bias_val = + params.bias_ptr == nullptr ? 0.f : to_float(reinterpret_cast(params.bias_ptr)[channel_id]); + + const bool has_initial_state = params.has_initial_state_ptr == nullptr + ? false + : reinterpret_cast(params.has_initial_state_ptr)[batch_id]; + + const int32_t* cache_indices = + params.cache_indices_ptr == nullptr ? nullptr : reinterpret_cast(params.cache_indices_ptr); + const int32_t cache_index = cache_indices == nullptr ? batch_id : cache_indices[batch_id]; + // cache_index == params.pad_slot_id is defined as padding, so we exit early. + if (cache_index == params.pad_slot_id) { + return; + } + input_t* conv_states = params.conv_states_ptr == nullptr ? nullptr + : reinterpret_cast(params.conv_states_ptr) + + cache_index * params.conv_states_batch_stride + + channel_id * params.conv_states_c_stride; + + // Thread 0 will load the last elements of the previous chunk, so we initialize those to 0. + if (tidx == 0) { + input_t initial_state[kNElts]; + zero_fill(initial_state); + if (has_initial_state) { +#pragma unroll + for (int w = 0; w < kWidth - 1; ++w) { + initial_state[kNElts - 1 - (kWidth - 2) + w] = conv_states[w]; + } + } + smem_exchange[kNThreads - 1] = reinterpret_cast(initial_state)[0]; + } + + float weight_vals[kWidth]; +#pragma unroll + for (int i = 0; i < kWidth; ++i) { + weight_vals[i] = to_float(weight[i * params.weight_width_stride]); + } + + constexpr int kChunkSize = kNThreads * kNElts; + const int32_t n_chunks = (seqlen + kChunkSize - 1) / kChunkSize; + for (int32_t chunk = 0; chunk < n_chunks; ++chunk) { + input_t x_vals_load[2 * kNElts]; + zero_fill(x_vals_load); + if constexpr (kIsVecLoad) { + typename Ktraits::BlockLoadVecT(smem_load_vec) + .Load( + reinterpret_cast(x), + *reinterpret_cast(&x_vals_load[kNElts]), + (seqlen - chunk * kChunkSize) / kNElts); + } else { + __syncthreads(); + typename Ktraits::BlockLoadT(smem_load).Load( + x, *reinterpret_cast(&x_vals_load[kNElts]), seqlen - chunk * kChunkSize); + } + x += kChunkSize; + __syncthreads(); + // Thread kNThreads - 1 doesn't write yet, so that thread 0 can read + // the last elements of the previous chunk. + if (tidx < kNThreads - 1) { + smem_exchange[tidx] = reinterpret_cast(x_vals_load)[1]; + } + __syncthreads(); + reinterpret_cast(x_vals_load)[0] = smem_exchange[tidx > 0 ? tidx - 1 : kNThreads - 1]; + __syncthreads(); + // Now thread kNThreads - 1 can write the last elements of the current chunk. + if (tidx == kNThreads - 1) { + smem_exchange[tidx] = reinterpret_cast(x_vals_load)[1]; + } + + float x_vals[2 * kNElts]; +#pragma unroll + for (int i = 0; i < 2 * kNElts; ++i) { + x_vals[i] = to_float(x_vals_load[i]); + } + + float out_vals[kNElts]; +#pragma unroll + for (int i = 0; i < kNElts; ++i) { + out_vals[i] = bias_val; +#pragma unroll + for (int w = 0; w < kWidth; ++w) { + out_vals[i] += weight_vals[w] * x_vals[kNElts + i - (kWidth - w - 1)]; + } + } + + if (params.silu_activation) { +#pragma unroll + for (int i = 0; i < kNElts; ++i) { + out_vals[i] = out_vals[i] / (1 + expf(-out_vals[i])); + } + } + + input_t out_vals_store[kNElts]; +#pragma unroll + for (int i = 0; i < kNElts; ++i) { + out_vals_store[i] = from_float(out_vals[i]); + } + if constexpr (kIsVecLoad) { + typename Ktraits::BlockStoreVecT(smem_store_vec) + .Store( + reinterpret_cast(out), + reinterpret_cast(out_vals_store), + (seqlen - chunk * kChunkSize) / kNElts); + } else { + typename Ktraits::BlockStoreT(smem_store).Store(out, out_vals_store, seqlen - chunk * kChunkSize); + } + out += kChunkSize; + + const int32_t final_state_position = ((seqlen - (kWidth - 1)) - (n_chunks - 1) * kChunkSize); + // In case the final state is separated between the last "smem_exchange" and + // the one before it (chunk = n_chunks - 1 and chunk = n_chunks - 2), + // (which occurs when `final_state_position` is a non-positive index) + // we load the correct data from smem_exchange from both chunks, the last + // chunk iteration and the one before it. + if (conv_states != nullptr && final_state_position < 0 && seqlen > kWidth) { + input_t vals_load[kNElts]; + zero_fill(vals_load); + if ((chunk == n_chunks - 2) && (tidx == kNThreads - 1)) { + // chunk = n_chunks - 2, a segment of the final state sits in the last index + reinterpret_cast(vals_load)[0] = smem_exchange[kNThreads - 1]; +#pragma unroll + for (int w = 0; w < -final_state_position; ++w) { + conv_states[w] = vals_load[kNElts + final_state_position + w]; + } + } + if ((chunk == n_chunks - 1) && tidx == 0) { + // chunk = n_chunks - 1, the second segment of the final state first positions + reinterpret_cast(vals_load)[0] = smem_exchange[0]; + for (int w = -final_state_position; w < kWidth - 1; ++w) { + conv_states[w] = vals_load[w + final_state_position]; + } + return; + } + } + } + // Final state is stored in the smem_exchange last token slot, + // in case seqlen < kWidth, we would need to take the final state from the + // initial state which is stored in conv_states + // in case seqlen > kWidth, we would need to load the last kWidth - 1 data + // and load it into conv_state accordingly + const int32_t last_thread = ((seqlen - (kWidth - 1)) - (n_chunks - 1) * kChunkSize) / kNElts; + if (conv_states != nullptr && tidx == last_thread) { + input_t x_vals_load[kNElts * 2]; + zero_fill(x_vals_load); + // in case we are on the first kWidth tokens + if (last_thread == 0 && seqlen < kWidth) { + // Need to take the initial state + reinterpret_cast(x_vals_load)[0] = smem_exchange[0]; + const int32_t offset = seqlen - (kWidth - 1); +#pragma unroll + for (int w = 0; w < kWidth - 1; ++w) { + // pad the existing state + if ((w - seqlen) >= 0 && has_initial_state) { + conv_states[w - seqlen] = conv_states[w]; + } else if ((w - seqlen) >= 0 && !has_initial_state) { + conv_states[w - seqlen] = from_float(0.0f); + } + } +#pragma unroll + for (int w = 0; w < kWidth - 1; ++w) { + if (offset + w >= 0) { + conv_states[w] = x_vals_load[offset + w]; + } + } + } else { + // in case the final state is in between the threads data + const int32_t offset = ((seqlen - (kWidth - 1)) % (kNElts)); + if ((offset + kWidth - 2) >= kNElts && (last_thread + 1 < kNThreads)) { + // In case last_thread == kNThreads - 1, accessing last_thread + 1 will result in an + // illegal access error on H100. + // Therefore, we access last_thread + 1 only if the final state data sits there. + reinterpret_cast(x_vals_load)[1] = smem_exchange[last_thread + 1]; + } + reinterpret_cast(x_vals_load)[0] = smem_exchange[last_thread]; +#pragma unroll + for (int w = 0; w < kWidth - 1; ++w) { + conv_states[w] = x_vals_load[offset + w]; + } + } + } +} + +template +void causal_conv1d_fwd_launch(const MambaConvParams& params, DLDevice device) { + static constexpr int kNElts = sizeof(T) == 4 ? 4 : 8; + const bool is_varlen = params.query_start_loc_ptr != nullptr; + const bool is_vec_load = params.seqlen % kNElts == 0 && !is_varlen; + + const dim3 grid(params.batch, params.dim); + if (is_vec_load) { + using Ktraits = CausalConv1dFwdTraits; + // The AOT launcher raised the dynamic-smem cap past 48 KB; these traits stay + // near 4 KB, so the default limit is enough. + static_assert(Ktraits::kSmemSize < 48 * 1024); + host::LaunchKernel(grid, kNThreads, device, Ktraits::kSmemSize)(causal_conv1d_fwd_kernel, params); + } else { + using Ktraits = CausalConv1dFwdTraits; + static_assert(Ktraits::kSmemSize < 48 * 1024); + host::LaunchKernel(grid, kNThreads, device, Ktraits::kSmemSize)(causal_conv1d_fwd_kernel, params); + } +} + +template +void causal_conv1d_fwd_cuda(const MambaConvParams& params, DLDevice device) { + switch (params.width) { + case 2: + return causal_conv1d_fwd_launch<128, 2, T>(params, device); + case 3: + return causal_conv1d_fwd_launch<128, 3, T>(params, device); + case 4: + return causal_conv1d_fwd_launch<128, 4, T>(params, device); + default: + host::Panic("causal_conv1d_fwd: width must be between 2 and 4, got ", params.width); + } +} + +//////////////////////////////////////////////////////////////////////////////// +// Decode +//////////////////////////////////////////////////////////////////////////////// + +/// \brief One thread per (sequence, channel): advance the conv state by +/// `seqlen` tokens and emit the convolution over the sliding window. +template +__global__ __launch_bounds__(kNThreads) void causal_conv1d_update_kernel(MambaConvParams params) { + using input_t = T; + + const int32_t tidx = threadIdx.x; + const int32_t batch_id = blockIdx.x; + const int32_t channel_id = blockIdx.y * kNThreads + tidx; + if (channel_id >= params.dim) return; + + const input_t* x = reinterpret_cast(params.x_ptr) + batch_id * params.x_batch_stride + + channel_id * params.x_c_stride; + + // If params.conv_state_indices_ptr is set, the conv state is gathered from the conv state + // tensor along the batch axis. Otherwise, the conv state coordinate is the same as the batch id. + const int32_t conv_state_batch_coord = + params.conv_state_indices_ptr == nullptr ? batch_id : params.conv_state_indices_ptr[batch_id]; + // conv_state_batch_coord == params.pad_slot_id is defined as padding so we exit early. + if (conv_state_batch_coord == params.pad_slot_id) { + return; + } + input_t* conv_state = reinterpret_cast(params.conv_state_ptr) + + conv_state_batch_coord * params.conv_state_batch_stride + + channel_id * params.conv_state_c_stride; + + const input_t* weight = reinterpret_cast(params.weight_ptr) + channel_id * params.weight_c_stride; + input_t* out = reinterpret_cast(params.out_ptr) + batch_id * params.out_batch_stride + + channel_id * params.out_c_stride; + const float bias_val = + params.bias_ptr == nullptr ? 0.f : to_float(reinterpret_cast(params.bias_ptr)[channel_id]); + + const int32_t state_len = params.conv_state_len; + const int32_t advance_len = params.seqlen; + const int32_t cache_seqlen = kIsCircularBuffer ? params.cache_seqlens[batch_id] % state_len : 0; + int32_t update_idx = cache_seqlen - (kWidth - 1); + update_idx = update_idx < 0 ? update_idx + state_len : update_idx; + + float weight_vals[kWidth] = {0}; +#pragma unroll + for (int i = 0; i < kWidth; ++i) { + weight_vals[i] = to_float(weight[i * params.weight_width_stride]); + } + + float x_vals[kWidth] = {0}; + if constexpr (!kIsCircularBuffer) { +#pragma unroll 2 + for (int32_t i = 0; i < state_len - advance_len - (kWidth - 1); ++i) { + conv_state[i * params.conv_state_l_stride] = conv_state[(i + advance_len) * params.conv_state_l_stride]; + } +#pragma unroll + for (int i = 0; i < kWidth - 1; ++i) { + const input_t state_val = conv_state[(state_len - (kWidth - 1) + i) * params.conv_state_l_stride]; + if (i < advance_len + (kWidth - 1) && state_len - advance_len - (kWidth - 1) + i >= 0) { + conv_state[(state_len - advance_len - (kWidth - 1) + i) * params.conv_state_l_stride] = state_val; + } + x_vals[i] = to_float(state_val); + } + } else { +#pragma unroll + for (int i = 0; i < kWidth - 1; + ++i, update_idx = update_idx + 1 >= state_len ? update_idx + 1 - state_len : update_idx + 1) { + const input_t state_val = conv_state[update_idx * params.conv_state_l_stride]; + x_vals[i] = to_float(state_val); + } + } +#pragma unroll 2 + for (int32_t i = 0; i < params.seqlen; ++i) { + const input_t x_val = x[i * params.x_l_stride]; + if constexpr (!kIsCircularBuffer) { + if (i < advance_len && state_len - advance_len + i >= 0) { + conv_state[(state_len - advance_len + i) * params.conv_state_l_stride] = x_val; + } + } else { + conv_state[update_idx * params.conv_state_l_stride] = x_val; + ++update_idx; + update_idx = update_idx >= state_len ? update_idx - state_len : update_idx; + } + x_vals[kWidth - 1] = to_float(x_val); + float out_val = bias_val; +#pragma unroll + for (int j = 0; j < kWidth; ++j) { + out_val += weight_vals[j] * x_vals[j]; + } + if (params.silu_activation) { + out_val = out_val / (1 + expf(-out_val)); + } + out[i * params.out_l_stride] = from_float(out_val); + // Shift the input buffer by 1 +#pragma unroll + for (int k = 0; k < kWidth - 1; ++k) { + x_vals[k] = x_vals[k + 1]; + } + } +} + +template +void causal_conv1d_update_launch(const MambaConvParams& params, DLDevice device) { + const dim3 grid(params.batch, host::div_ceil(params.dim, kNThreads)); + if (params.cache_seqlens == nullptr) { + host::LaunchKernel(grid, kNThreads, device)(causal_conv1d_update_kernel, params); + } else { + host::LaunchKernel(grid, kNThreads, device)(causal_conv1d_update_kernel, params); + } +} + +template +void causal_conv1d_update_cuda(const MambaConvParams& params, DLDevice device) { + switch (params.width) { + case 2: + return causal_conv1d_update_launch<64, 2, T>(params, device); + case 3: + return causal_conv1d_update_launch<64, 3, T>(params, device); + case 4: + return causal_conv1d_update_launch<64, 4, T>(params, device); + default: + host::Panic("causal_conv1d_update: width must be between 2 and 4, got ", params.width); + } +} + +//////////////////////////////////////////////////////////////////////////////// +// Host entry points +//////////////////////////////////////////////////////////////////////////////// + +/// \brief Depthwise causal conv1d over whole sequences, in place on `x`. +/// +/// \tparam T Element type: fp16_t | bf16_t | fp32_t. +/// \param x `(batch, dim, seqlen)`, or `(dim, cu_seqlen)` when +/// `query_start_loc` is given. Overwritten with the output. +/// \param weight `(dim, width)`, `width` in 2..4. +/// \param bias Optional `(dim,)`. +/// \param conv_states Optional `(num_slots, dim, state_len)`; the trailing +/// `width - 1` taps of each sequence are written back. +/// \param query_start_loc Optional int32 `(batch + 1,)` varlen cumulative lengths. +/// \param cache_indices Optional int32 `(batch,)` conv-state slot per sequence. +/// \param has_initial_state Optional bool `(batch,)`; whether to seed from `conv_states`. +/// \param silu_activation Whether to apply SiLU to the output. +/// \param pad_slot_id Sequences whose cache index equals this are skipped. +template +void causal_conv1d_fwd( + tvm::ffi::TensorView x, + tvm::ffi::TensorView weight, + tvm::ffi::Optional bias, + tvm::ffi::Optional conv_states, + tvm::ffi::Optional query_start_loc, + tvm::ffi::Optional cache_indices, + tvm::ffi::Optional has_initial_state, + bool silu_activation, + int64_t pad_slot_id) { + using namespace host; + + const bool varlen = query_start_loc.has_value(); + auto batch_sym = SymbolicSize{"batch"}; + auto dim_sym = SymbolicSize{"dim"}; + auto seqlen_sym = SymbolicSize{"seqlen"}; + auto width_sym = SymbolicSize{"width"}; + auto device_sym = SymbolicDevice{}; + device_sym.set_options(); + + // Only the innermost stride is pinned: the Python wrapper makes `x` unit-stride + // in its last dimension, everything else is carried through as a stride. + if (varlen) { + TensorMatcher({dim_sym, seqlen_sym}).with_strides({-1, 1}).with_dtype().with_device(device_sym).verify(x); + TensorMatcher({-1}).with_dtype().with_device(device_sym).verify(query_start_loc.value()); + batch_sym.set_value(query_start_loc.value().size(0) - 1); + } else { + TensorMatcher({batch_sym, dim_sym, seqlen_sym}) + .with_strides({-1, -1, 1}) + .with_dtype() + .with_device(device_sym) + .verify(x); + } + TensorMatcher({dim_sym, width_sym}).with_strides({-1, -1}).with_dtype().with_device(device_sym).verify(weight); + + const int64_t batch = batch_sym.unwrap(); + const int64_t dim = dim_sym.unwrap(); + const int64_t seqlen = seqlen_sym.unwrap(); + const int64_t width = width_sym.unwrap(); + CHECK_HOST(batch > 0) << "causal_conv1d_fwd: batch must be positive, got " << batch; + CHECK_HOST(width >= 2 && width <= 4) << "causal_conv1d only supports width between 2 and 4, got " << width; + + if (bias.has_value()) { + TensorMatcher({dim_sym}).with_strides({1}).with_dtype().with_device(device_sym).verify(bias.value()); + } + if (cache_indices.has_value()) { + TensorMatcher({batch_sym}).with_dtype().with_device(device_sym).verify(cache_indices.value()); + } + if (has_initial_state.has_value()) { + const auto& initial_state_mask = has_initial_state.value(); + TensorMatcher({batch_sym}).with_device(device_sym).verify(initial_state_mask); + // Read as `const bool*` by the kernel. `kDLBool` has no C++ trait here, so + // `.with_dtype()` is unavailable -- check the code directly. + CHECK_HOST(initial_state_mask.dtype().code == kDLBool && initial_state_mask.dtype().bits == 8) + << "causal_conv1d_fwd: has_initial_state must be a bool tensor, got dtype code " + << static_cast(initial_state_mask.dtype().code) << " with " + << static_cast(initial_state_mask.dtype().bits) << " bits"; + } + + // `out` aliases `x`: this op is in-place and callers rely on that. + auto params = MambaConvParams{}; + params.batch = static_cast(batch); + params.dim = static_cast(dim); + params.seqlen = static_cast(seqlen); + params.width = static_cast(width); + params.pad_slot_id = pad_slot_id; + params.silu_activation = silu_activation; + params.x_ptr = x.data_ptr(); + params.weight_ptr = weight.data_ptr(); + params.bias_ptr = bias.has_value() ? bias.value().data_ptr() : nullptr; + params.out_ptr = x.data_ptr(); + params.query_start_loc_ptr = varlen ? query_start_loc.value().data_ptr() : nullptr; + params.cache_indices_ptr = cache_indices.has_value() ? cache_indices.value().data_ptr() : nullptr; + params.has_initial_state_ptr = has_initial_state.has_value() ? has_initial_state.value().data_ptr() : nullptr; + // In the varlen layout `x` is (dim, cu_seqlen): the "batch" axis is the token + // axis, so the token stride doubles as the batch stride. + params.x_batch_stride = static_cast(x.stride(varlen ? 1 : 0)); + params.x_c_stride = static_cast(x.stride(varlen ? 0 : 1)); + params.x_l_stride = static_cast(x.stride(varlen ? 1 : 2)); + params.weight_c_stride = static_cast(weight.stride(0)); + params.weight_width_stride = static_cast(weight.stride(1)); + params.out_batch_stride = params.x_batch_stride; + params.out_c_stride = params.x_c_stride; + params.out_l_stride = params.x_l_stride; + + if (conv_states.has_value()) { + const auto& states = conv_states.value(); + TensorMatcher({-1, dim_sym, -1}).with_strides({-1, -1, -1}).with_dtype().with_device(device_sym).verify(states); + params.conv_states_ptr = states.data_ptr(); + params.conv_states_batch_stride = static_cast(states.stride(0)); + params.conv_states_c_stride = static_cast(states.stride(1)); + params.conv_states_l_stride = static_cast(states.stride(2)); + } else { + params.conv_states_ptr = nullptr; + } + + causal_conv1d_fwd_cuda(params, device_sym.unwrap()); +} + +/// \brief Single-step (decode) depthwise causal conv1d, in place on `x`. +/// +/// \tparam T Element type: fp16_t | bf16_t | fp32_t. +/// \param x `(batch, dim, seqlen)`, overwritten with the output. +/// \param conv_state `(num_entries, dim, state_len)`, `state_len >= width - 1`, +/// advanced in place. +/// \param weight `(dim, width)`, `width` in 2..4. +/// \param bias Optional `(dim,)`. +/// \param silu_activation Whether to apply SiLU to the output. +/// \param cache_seqlens Optional int32 `(batch,)`; when given, `conv_state` is +/// treated as a circular buffer starting at +/// `cache_seqlens % state_len`. +/// \param conv_state_indices Optional int32 `(batch,)` conv-state slot per sequence. +/// \param pad_slot_id Sequences whose state index equals this are skipped. +template +void causal_conv1d_update( + tvm::ffi::TensorView x, + tvm::ffi::TensorView conv_state, + tvm::ffi::TensorView weight, + tvm::ffi::Optional bias, + bool silu_activation, + tvm::ffi::Optional cache_seqlens, + tvm::ffi::Optional conv_state_indices, + int64_t pad_slot_id) { + using namespace host; + + auto batch_sym = SymbolicSize{"batch"}; + auto dim_sym = SymbolicSize{"dim"}; + auto seqlen_sym = SymbolicSize{"seqlen"}; + auto width_sym = SymbolicSize{"width"}; + auto state_len_sym = SymbolicSize{"state_len"}; + auto entries_sym = SymbolicSize{"conv_state_entries"}; + auto device_sym = SymbolicDevice{}; + device_sym.set_options(); + + TensorMatcher({batch_sym, dim_sym, seqlen_sym}) + .with_strides({-1, -1, -1}) + .with_dtype() + .with_device(device_sym) + .verify(x); + TensorMatcher({dim_sym, width_sym}).with_strides({-1, -1}).with_dtype().with_device(device_sym).verify(weight); + // Gathered decode indexes `conv_state` by slot, so its leading dimension is + // the pool size rather than the batch size. + if (conv_state_indices.has_value()) { + TensorMatcher({batch_sym}) + .with_strides({1}) + .with_dtype() + .with_device(device_sym) + .verify(conv_state_indices.value()); + } else { + entries_sym.set_value(batch_sym.unwrap()); + } + TensorMatcher({entries_sym, dim_sym, state_len_sym}) + .with_strides({-1, -1, -1}) + .with_dtype() + .with_device(device_sym) + .verify(conv_state); + + const int64_t width = width_sym.unwrap(); + const int64_t state_len = state_len_sym.unwrap(); + CHECK_HOST(width >= 2 && width <= 4) << "causal_conv1d only supports width between 2 and 4, got " << width; + CHECK_HOST(state_len >= width - 1) << "causal_conv1d_update: conv_state length " << state_len + << " is shorter than width - 1 = " << width - 1; + + if (bias.has_value()) { + TensorMatcher({dim_sym}).with_strides({1}).with_dtype().with_device(device_sym).verify(bias.value()); + } + + auto params = MambaConvParams{}; + params.batch = static_cast(batch_sym.unwrap()); + params.dim = static_cast(dim_sym.unwrap()); + params.seqlen = static_cast(seqlen_sym.unwrap()); + params.width = static_cast(width); + params.pad_slot_id = pad_slot_id; + params.silu_activation = silu_activation; + params.x_ptr = x.data_ptr(); + params.weight_ptr = weight.data_ptr(); + params.bias_ptr = bias.has_value() ? bias.value().data_ptr() : nullptr; + params.out_ptr = x.data_ptr(); + params.x_batch_stride = static_cast(x.stride(0)); + params.x_c_stride = static_cast(x.stride(1)); + params.x_l_stride = static_cast(x.stride(2)); + params.weight_c_stride = static_cast(weight.stride(0)); + params.weight_width_stride = static_cast(weight.stride(1)); + params.out_batch_stride = params.x_batch_stride; + params.out_c_stride = params.x_c_stride; + params.out_l_stride = params.x_l_stride; + + params.conv_state_ptr = conv_state.data_ptr(); + params.conv_state_len = static_cast(state_len); + params.conv_state_batch_stride = static_cast(conv_state.stride(0)); + params.conv_state_c_stride = static_cast(conv_state.stride(1)); + params.conv_state_l_stride = static_cast(conv_state.stride(2)); + + if (cache_seqlens.has_value()) { + TensorMatcher({batch_sym}) + .with_strides({1}) + .with_dtype() + .with_device(device_sym) + .verify(cache_seqlens.value()); + params.cache_seqlens = static_cast(cache_seqlens.value().data_ptr()); + } else { + params.cache_seqlens = nullptr; + } + params.conv_state_indices_ptr = + conv_state_indices.has_value() ? static_cast(conv_state_indices.value().data_ptr()) : nullptr; + + causal_conv1d_update_cuda(params, device_sym.unwrap()); +} + +} // namespace mamba_conv + +using mamba_conv::causal_conv1d_fwd; +using mamba_conv::causal_conv1d_update; + +} // namespace sglang diff --git a/python/sglang/kernels/ops/mamba/__init__.py b/python/sglang/kernels/ops/mamba/__init__.py index bde9532a7..e4664256b 100644 --- a/python/sglang/kernels/ops/mamba/__init__.py +++ b/python/sglang/kernels/ops/mamba/__init__.py @@ -6,31 +6,44 @@ from typing import TYPE_CHECKING, Optional from sglang.kernels.registry import register_kernel from sglang.kernels.selector import get_kernel -from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, + KernelSpec, +) if TYPE_CHECKING: import torch +_CUDA = frozenset({CapabilityRequirement.CUDA}) + +# JIT is the only backend: the AOT kernel it replaced was built for CUDA alone, +# never by the ROCm / MUSA / Metal extensions. Non-CUDA resolves nothing here -- +# the Triton fallback is picked by the serving wrapper's `_HAS_CONV1D_KERNEL` +# branch, not by the registry. register_kernel( KernelSpec( op="mamba.causal_conv1d_fwd", - backend=KernelBackend.AOT, - target="sgl_kernel.mamba:causal_conv1d_fwd", + backend=KernelBackend.JIT, + target="sglang.kernels.ops.mamba.causal_conv1d:causal_conv1d_fwd", + capabilities=_CUDA, format_signature=FormatSignature( in_place=True, description="causal depthwise conv1d forward (prefill)" ), - description="Causal conv1d forward (sgl_kernel wheel).", + description="Causal conv1d forward (sglang.kernels.jit).", ) ) register_kernel( KernelSpec( op="mamba.causal_conv1d_update", - backend=KernelBackend.AOT, - target="sgl_kernel.mamba:causal_conv1d_update", + backend=KernelBackend.JIT, + target="sglang.kernels.ops.mamba.causal_conv1d:causal_conv1d_update", + capabilities=_CUDA, format_signature=FormatSignature( in_place=True, description="causal depthwise conv1d update (decode)" ), - description="Causal conv1d update (sgl_kernel wheel).", + description="Causal conv1d update (sglang.kernels.jit).", ) ) @@ -47,7 +60,7 @@ def causal_conv1d_fwd( pad_slot_id: int, ): """Causal depthwise conv1d forward (prefill).""" - return get_kernel("mamba.causal_conv1d_fwd", KernelBackend.AOT)( + return get_kernel("mamba.causal_conv1d_fwd", KernelBackend.JIT)( x, weight, bias_, @@ -71,7 +84,7 @@ def causal_conv1d_update( pad_slot_id: int, ): """Causal depthwise conv1d update (decode).""" - return get_kernel("mamba.causal_conv1d_update", KernelBackend.AOT)( + return get_kernel("mamba.causal_conv1d_update", KernelBackend.JIT)( x, conv_state, weight, diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d.py b/python/sglang/kernels/ops/mamba/causal_conv1d.py new file mode 100644 index 000000000..2558fb2e9 --- /dev/null +++ b/python/sglang/kernels/ops/mamba/causal_conv1d.py @@ -0,0 +1,106 @@ +"""JIT depthwise causal conv1d: prefill (``fwd``) and decode (``update``). + +In-place, like the AOT ops they replace on CUDA: the output is written back into +``x`` and the conv state is advanced in place. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + get_jit_cuda_arch, + load_jit, + make_cpp_args, +) +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_causal_conv1d_module(dtype: torch.dtype) -> Module: + if dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise RuntimeError( + f"Unsupported dtype {dtype}. Supported: float16, bfloat16, float32" + ) + # The AOT wheel ships an SM90 build compiled with `-use_fast_math` and a + # precise-math build for every other arch; match that split so the SiLU + # epilogue keeps producing the same bits as the op being replaced. + arch = get_jit_cuda_arch() + use_fast_math = (arch.major, arch.minor) == (9, 0) + math_mode = "fast_math" if use_fast_math else "precise_math" + args = make_cpp_args(dtype) + return load_jit( + "causal_conv1d", + math_mode, + *args, + cuda_files=["mamba/causal_conv1d.cuh"], + cuda_wrappers=[ + ("causal_conv1d_fwd", f"causal_conv1d_fwd<{args}>"), + ("causal_conv1d_update", f"causal_conv1d_update<{args}>"), + ], + extra_cuda_cflags=["--use_fast_math"] if use_fast_math else [], + ) + + +@register_custom_op( + op_name="mamba_causal_conv1d_fwd", + mutates_args=["x", "conv_states"], +) +def causal_conv1d_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias_: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool, + pad_slot_id: int, +) -> None: + """Causal depthwise conv1d forward (prefill), written back into ``x``.""" + module = _jit_causal_conv1d_module(x.dtype) + module.causal_conv1d_fwd( + x, + weight, + bias_, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + silu_activation, + pad_slot_id, + ) + + +@register_custom_op( + op_name="mamba_causal_conv1d_update", + mutates_args=["x", "conv_state"], +) +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias_: Optional[torch.Tensor], + silu_activation: bool, + cache_seqlens: Optional[torch.Tensor], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, +) -> None: + """Causal depthwise conv1d update (decode), written back into ``x``.""" + module = _jit_causal_conv1d_module(x.dtype) + module.causal_conv1d_update( + x, + conv_state, + weight, + bias_, + silu_activation, + cache_seqlens, + conv_state_indices, + pad_slot_id, + ) diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py index 4601f1df5..31dda54d5 100644 --- a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py +++ b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py @@ -18,15 +18,20 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import ( from sglang.kernels.ops.mamba.causal_conv1d_triton import ( causal_conv1d_update as _causal_conv1d_update_triton, ) +from sglang.srt.utils import is_cuda -try: - from sgl_kernel import causal_conv1d_fwd - from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel +# The compiled causal conv1d is CUDA-only -- the sgl_kernel wheel never built it +# for ROCm / MUSA either, so the old import probe always fell through to Triton +# there. Select that fallback directly instead of via a failed import. +_HAS_CONV1D_KERNEL = is_cuda() - torch.ops.sgl_kernel.causal_conv1d_update - _HAS_SGL_KERNEL = True -except (ImportError, AttributeError): - _HAS_SGL_KERNEL = False +if _HAS_CONV1D_KERNEL: + from sglang.kernels.ops.mamba import ( + causal_conv1d_fwd, + ) + from sglang.kernels.ops.mamba import ( + causal_conv1d_update as causal_conv1d_update_kernel, + ) def _get_seq_lens_cpu(query_start_loc, x): @@ -76,11 +81,13 @@ def causal_conv1d_fn( out: (batch, dim, seqlen) """ - # Use Triton when: (1) sgl_kernel not available, or (2) input is - # non-contiguous and seq_lens_cpu is already pre-computed by caller. + # Use Triton when: (1) there is no compiled conv1d kernel for this device, + # or (2) input is non-contiguous and seq_lens_cpu is pre-computed by caller. # The Triton kernel accepts arbitrary strides, avoiding a .contiguous() # copy that can cost >0.6 ms/layer on large prefill batches. - use_triton = not _HAS_SGL_KERNEL or (x.stride(-1) != 1 and "seq_lens_cpu" in kwargs) + use_triton = not _HAS_CONV1D_KERNEL or ( + x.stride(-1) != 1 and "seq_lens_cpu" in kwargs + ) if use_triton: if "seq_lens_cpu" not in kwargs: kwargs["seq_lens_cpu"] = _get_seq_lens_cpu(query_start_loc, x) @@ -150,7 +157,7 @@ def causal_conv1d_update( indices 0 and 3 out: (batch, dim) or (batch, dim, seqlen) """ - use_triton = not _HAS_SGL_KERNEL + use_triton = not _HAS_CONV1D_KERNEL if use_triton: return _causal_conv1d_update_triton( x, diff --git a/test/registered/kernels/benchmark/mamba/bench_causal_conv1d.py b/test/registered/kernels/benchmark/mamba/bench_causal_conv1d.py new file mode 100644 index 000000000..9316b93de --- /dev/null +++ b/test/registered/kernels/benchmark/mamba/bench_causal_conv1d.py @@ -0,0 +1,104 @@ +"""AOT vs. JIT benchmark for the depthwise causal conv1d prefill/decode kernels.""" + +import torch +from sgl_kernel import causal_conv1d_fwd as aot_causal_conv1d_fwd +from sgl_kernel import causal_conv1d_update as aot_causal_conv1d_update + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE, create_random +from sglang.kernels.ops.mamba.causal_conv1d import ( + causal_conv1d_fwd as jit_causal_conv1d_fwd, +) +from sglang.kernels.ops.mamba.causal_conv1d import ( + causal_conv1d_update as jit_causal_conv1d_update, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +PAD_SLOT_ID = -1 +WIDTH = 4 + +FWD_FN_MAP = {"jit": jit_causal_conv1d_fwd, "aot": aot_causal_conv1d_fwd} +UPDATE_FN_MAP = {"jit": jit_causal_conv1d_update, "aot": aot_causal_conv1d_update} + + +@marker.parametrize("seqlen", [128, 512, 2048, 8192], [512]) +@marker.parametrize("dim", [2048, 4096, 8192], [4096]) +@marker.parametrize("dtype", [torch.float16, torch.bfloat16]) +@marker.benchmark("impl", ["jit", "aot"]) +def benchmark_fwd(seqlen: int, dim: int, dtype: torch.dtype, impl: str): + """Prefill: one varlen batch of four sequences, conv state written back.""" + batch = 4 + x = create_random(dim, seqlen, dtype=dtype) + weight = create_random(dim, WIDTH, dtype=dtype) + bias = create_random(dim, dtype=dtype) + conv_states = create_random(batch, dim, WIDTH - 1, dtype=dtype) + lengths = [seqlen // batch] * batch + lengths[-1] += seqlen - sum(lengths) + query_start_loc = torch.tensor( + [0] + torch.cumsum(torch.tensor(lengths), 0).tolist(), + dtype=torch.int32, + device=DEFAULT_DEVICE, + ) + cache_indices = torch.arange(batch, dtype=torch.int32, device=DEFAULT_DEVICE) + has_initial_state = torch.ones(batch, dtype=torch.bool, device=DEFAULT_DEVICE) + return marker.do_bench( + FWD_FN_MAP[impl], + input_args=( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + True, + PAD_SLOT_ID, + ), + # x and conv_states are read-modify-write, so both need cloning. + graph_clone_args=(0, 1, 2, 3), + memory_args=(x, weight, bias, conv_states), + memory_output=(x, conv_states), + ) + + +@marker.parametrize("batch", [1, 8, 64, 256], [64]) +@marker.parametrize("dim", [2048, 4096, 8192], [4096]) +@marker.parametrize("dtype", [torch.float16, torch.bfloat16]) +@marker.benchmark("impl", ["jit", "aot"]) +def benchmark_update(batch: int, dim: int, dtype: torch.dtype, impl: str): + """Decode: one token per sequence, conv state gathered by slot index.""" + entries = max(batch * 4, 64) + x = create_random(batch, dim, 1, dtype=dtype) + conv_state = create_random(entries, dim, WIDTH - 1, dtype=dtype) + weight = create_random(dim, WIDTH, dtype=dtype) + bias = create_random(dim, dtype=dtype) + conv_state_indices = torch.randperm(entries, device=DEFAULT_DEVICE)[:batch].to( + torch.int32 + ) + return marker.do_bench( + UPDATE_FN_MAP[impl], + input_args=( + x, + conv_state, + weight, + bias, + True, + None, + conv_state_indices, + PAD_SLOT_ID, + ), + # conv_state is the large pool; only the gathered rows are touched, so + # leave it out of the rotation and count just those rows as traffic. + graph_clone_args=(0, 2, 3, 6), + memory_args=(x, weight, bias, conv_state_indices), + memory_output=(x,), + ) + + +if __name__ == "__main__": + benchmark_fwd.run() + benchmark_update.run() diff --git a/test/registered/kernels/ops/mamba/test_causal_conv1d.py b/test/registered/kernels/ops/mamba/test_causal_conv1d.py new file mode 100644 index 000000000..38a916aa8 --- /dev/null +++ b/test/registered/kernels/ops/mamba/test_causal_conv1d.py @@ -0,0 +1,510 @@ +"""Correctness coverage for the JIT depthwise causal conv1d kernels. + +Two layers. The reference tests are the broad gate: they compare against an +``F.conv1d`` formulation across the full dispatch grid, and depend only on the +kernel this repo builds. The differential tests are a small smoke set proving +the migration is bit-faithful to the AOT ops -- narrow on purpose, since they +compare two independently built binaries whose toolchains nothing pins together +(JIT: c++20 / sm_90a; wheel: c++17 / sm_90 / -DNDEBUG, and the pinned PyPI +release on scheduled runs). A bitwise failure with the reference cases green +points at the build environment, not a numerics regression. +""" + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Reference implementations adapted from +# https://github.com/vllm-project/vllm/blob/main/tests/kernels/mamba/test_causal_conv1d.py + +import sys +from typing import Optional + +import pytest +import torch +import torch.nn.functional as F +from sgl_kernel import causal_conv1d_fwd as aot_causal_conv1d_fwd +from sgl_kernel import causal_conv1d_update as aot_causal_conv1d_update + +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.mamba.causal_conv1d import ( + causal_conv1d_fwd as jit_causal_conv1d_fwd, +) +from sglang.kernels.ops.mamba.causal_conv1d import ( + causal_conv1d_update as jit_causal_conv1d_update, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="4-gpu-b200") +register_cuda_ci(est_time=180, stage="nightly", runner_config="1-gpu-large") + +PAD_SLOT_ID = -1 +DTYPES = [torch.float32, torch.float16, torch.bfloat16] +WIDTHS = [2, 3, 4] +# 8/128 hit the vectorized chunk load (divisible by both vector widths), 15/1025 +# the scalar path, 1025/4096 the multi-chunk conv-state stitching, and 1/3 the +# seqlen < width padding branch. +FWD_SEQLENS = get_ci_test_range([1, 3, 8, 15, 128, 1025, 4096], [3, 15, 1025]) +# One seqlen per prefill dispatch path, for the bitwise smoke set. +SMOKE_SEQLENS = [15, 128, 1025] +UPDATE_SEQLENS = [1, 2, 5] +# Larger than the `width - 1` every in-tree caller allocates, to reach the shift +# loop that only runs when state_len > width - 1. +UPDATE_STATE_LENS = [8, 16] + + +def _assert_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor) -> None: + assert actual.dtype == expected.dtype + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), ( + "JIT and AOT outputs differ bit-for-bit. If the reference tests in this " + "file pass, suspect a JIT-vs-wheel toolchain divergence before a " + "numerics regression -- see the module docstring." + ) + + +def _tolerance(dtype: torch.dtype) -> tuple[float, float]: + if dtype == torch.float32: + return 3e-4, 1e-3 + if dtype == torch.float16: + return 3e-3, 5e-3 + return 1e-2, 5e-2 + + +def causal_conv1d_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + initial_states: Optional[torch.Tensor] = None, + final_states_out: Optional[torch.Tensor] = None, + activation: Optional[str] = "silu", +): + """x: (batch, dim, seqlen); initial/final states: (batch, dim, width - 1).""" + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_states is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_states, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to(dtype_in) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out, final_states_out + + +def causal_conv1d_update_ref( + x, conv_state, weight, bias=None, activation=None, cache_seqlens=None +): + """x: (batch, dim, seqlen); conv_state: (batch, dim, state_len).""" + dtype_in = x.dtype + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange( + -(width - 1), 0, dtype=torch.long, device=x.device + ).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = ( + torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + ) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze( + 0 + ) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[ + :, :, -seqlen: + ] + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) + + +def _make_fwd_inputs(dtype, batch, dim, seqlen, width, varlen, seed=0): + device = "cuda" + gen = torch.Generator(device=device).manual_seed(seed) + + def randn(*shape): + return torch.randn(*shape, device=device, dtype=dtype, generator=gen) + + if varlen: + x = randn(dim, seqlen) + lengths = [seqlen // batch] * batch + lengths[-1] += seqlen - sum(lengths) + query_start_loc = torch.tensor( + [0] + torch.cumsum(torch.tensor(lengths), 0).tolist(), + dtype=torch.int32, + device=device, + ) + cache_indices = torch.arange(batch, dtype=torch.int32, device=device) + else: + x = randn(batch, dim, seqlen) + query_start_loc = None + cache_indices = None + return { + "x": x, + "weight": randn(dim, width), + "bias": randn(dim), + "conv_states": randn(batch, dim, width - 1), + "query_start_loc": query_start_loc, + "cache_indices": cache_indices, + "has_initial_state": torch.randint( + 0, 2, (batch,), dtype=torch.bool, device=device, generator=gen + ), + } + + +def _run_fwd(impl, inputs, silu_activation): + x = inputs["x"].clone() + conv_states = inputs["conv_states"].clone() + impl( + x, + inputs["weight"], + inputs["bias"], + conv_states, + inputs["query_start_loc"], + inputs["cache_indices"], + inputs["has_initial_state"], + silu_activation, + PAD_SLOT_ID, + ) + return x, conv_states + + +def _make_update_inputs( + dtype, batch, dim, seqlen, width, state_len, circular, gather, seed=0 +): + device = "cuda" + gen = torch.Generator(device=device).manual_seed(seed) + + def randn(*shape): + return torch.randn(*shape, device=device, dtype=dtype, generator=gen) + + entries = batch * 4 if gather else batch + return { + "x": randn(batch, dim, seqlen), + "conv_state": randn(entries, dim, state_len), + "weight": randn(dim, width), + "bias": randn(dim), + "cache_seqlens": ( + torch.randint( + 0, state_len, (batch,), dtype=torch.int32, device=device, generator=gen + ) + if circular + else None + ), + "conv_state_indices": ( + torch.randperm(entries, device=device, generator=gen)[:batch].to( + torch.int32 + ) + if gather + else None + ), + } + + +def _run_update(impl, inputs, silu_activation): + x = inputs["x"].clone() + conv_state = inputs["conv_state"].clone() + impl( + x, + conv_state, + inputs["weight"], + inputs["bias"], + silu_activation, + inputs["cache_seqlens"], + inputs["conv_state_indices"], + PAD_SLOT_ID, + ) + return x, conv_state + + +############################################################################### +# Reference coverage -- the broad gate. +############################################################################### + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +@pytest.mark.parametrize("seqlen", FWD_SEQLENS) +@pytest.mark.parametrize("has_initial_state", [True, False]) +def test_causal_conv1d_fwd_matches_reference(dtype, width, seqlen, has_initial_state): + """Output and final conv state must match an F.conv1d reference.""" + rtol, atol = _tolerance(dtype) + inputs = _make_fwd_inputs(dtype, 1, 64, seqlen, width, varlen=False, seed=7) + inputs["has_initial_state"] = torch.full( + (1,), has_initial_state, dtype=torch.bool, device="cuda" + ) + + out, conv_states = _run_fwd(jit_causal_conv1d_fwd, inputs, silu_activation=True) + out_ref, final_states_ref = causal_conv1d_ref( + inputs["x"].clone(), + inputs["weight"], + inputs["bias"], + initial_states=inputs["conv_states"].clone() if has_initial_state else None, + activation="silu", + ) + + torch.testing.assert_close(out, out_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(conv_states, final_states_ref, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +@pytest.mark.parametrize("seqlen", [s for s in FWD_SEQLENS if s >= 8]) +@pytest.mark.parametrize("silu_activation", [True, False]) +def test_causal_conv1d_fwd_varlen_matches_reference( + dtype, width, seqlen, silu_activation +): + """Varlen prefill, per sequence. The layout serving uses; always scalar load.""" + rtol, atol = _tolerance(dtype) + batch = 4 + inputs = _make_fwd_inputs(dtype, batch, 64, seqlen, width, varlen=True, seed=11) + activation = "silu" if silu_activation else None + + out, conv_states = _run_fwd(jit_causal_conv1d_fwd, inputs, silu_activation) + + conv_states_ref = inputs["conv_states"].clone() + starts = inputs["query_start_loc"].tolist() + for i in range(batch): + slot = int(inputs["cache_indices"][i]) + x_s = inputs["x"][:, starts[i] : starts[i + 1]].unsqueeze(0) + out_ref, _ = causal_conv1d_ref( + x_s.clone(), + inputs["weight"], + inputs["bias"], + initial_states=( + conv_states_ref[slot].unsqueeze(0).clone() + if inputs["has_initial_state"][i] + else None + ), + final_states_out=conv_states_ref[slot].unsqueeze(0), + activation=activation, + ) + torch.testing.assert_close( + out[:, starts[i] : starts[i + 1]].unsqueeze(0), + out_ref, + rtol=rtol, + atol=atol, + ) + torch.testing.assert_close(conv_states, conv_states_ref, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +@pytest.mark.parametrize("seqlen", UPDATE_SEQLENS) +@pytest.mark.parametrize("state_len_kind", ["exact", *UPDATE_STATE_LENS]) +@pytest.mark.parametrize("circular", [True, False]) +def test_causal_conv1d_update_matches_reference( + dtype, width, seqlen, state_len_kind, circular +): + """Both conv-state layouts: tail-anchored shift buffer and circular buffer.""" + state_len = width - 1 if state_len_kind == "exact" else state_len_kind + if circular and seqlen > state_len: + # The reference advances the ring with `scatter_`, whose behavior for the + # duplicate indices this produces is unspecified -- it cannot arbitrate. + pytest.skip("circular reference is ambiguous when seqlen > state_len") + rtol, atol = _tolerance(dtype) + inputs = _make_update_inputs( + dtype, 3, 2048, seqlen, width, state_len, circular, gather=False, seed=7 + ) + out, conv_state = _run_update(jit_causal_conv1d_update, inputs, True) + conv_state_ref = inputs["conv_state"].clone() + out_ref = causal_conv1d_update_ref( + inputs["x"].clone(), + conv_state_ref, + inputs["weight"], + inputs["bias"], + activation="silu", + cache_seqlens=inputs["cache_seqlens"], + ) + torch.testing.assert_close(out, out_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(conv_state, conv_state_ref, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +@pytest.mark.parametrize("seqlen", UPDATE_SEQLENS) +def test_causal_conv1d_update_gather_matches_reference(dtype, width, seqlen): + """Gathered decode: only the indexed slots advance, and they match the ref.""" + rtol, atol = _tolerance(dtype) + inputs = _make_update_inputs( + dtype, 3, 2048, seqlen, width, width - 1, circular=False, gather=True, seed=13 + ) + indices = inputs["conv_state_indices"] + out, conv_state = _run_update(jit_causal_conv1d_update, inputs, True) + + conv_state_ref = inputs["conv_state"][indices].clone() + out_ref = causal_conv1d_update_ref( + inputs["x"].clone(), + conv_state_ref, + inputs["weight"], + inputs["bias"], + activation="silu", + ) + torch.testing.assert_close(out, out_ref, rtol=rtol, atol=atol) + torch.testing.assert_close( + conv_state[indices], conv_state_ref, rtol=rtol, atol=atol + ) + + untouched = torch.ones(conv_state.shape[0], dtype=torch.bool, device="cuda") + untouched[indices] = False + assert torch.equal(conv_state[untouched], inputs["conv_state"][untouched]) + + +############################################################################### +# Differential smoke set -- the migration proof. +############################################################################### + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +@pytest.mark.parametrize("seqlen", SMOKE_SEQLENS) +def test_causal_conv1d_fwd_is_bit_exact(dtype, width, seqlen): + """One seqlen per prefill dispatch path: vectorized, scalar, multi-chunk.""" + inputs = _make_fwd_inputs(dtype, 1, 64, seqlen, width, varlen=False) + actual = _run_fwd(jit_causal_conv1d_fwd, inputs, True) + expected = _run_fwd(aot_causal_conv1d_fwd, inputs, True) + _assert_bitwise_equal(actual[0], expected[0]) + _assert_bitwise_equal(actual[1], expected[1]) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +def test_causal_conv1d_fwd_varlen_is_bit_exact(dtype, width): + inputs = _make_fwd_inputs(dtype, 4, 64, 1025, width, varlen=True) + actual = _run_fwd(jit_causal_conv1d_fwd, inputs, True) + expected = _run_fwd(aot_causal_conv1d_fwd, inputs, True) + _assert_bitwise_equal(actual[0], expected[0]) + _assert_bitwise_equal(actual[1], expected[1]) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("width", WIDTHS) +@pytest.mark.parametrize("circular", [True, False]) +@pytest.mark.parametrize("gather", [True, False]) +def test_causal_conv1d_update_is_bit_exact(dtype, width, circular, gather): + """Both state layouts crossed with both slot-addressing modes.""" + inputs = _make_update_inputs(dtype, 3, 2048 + 16, 2, width, 8, circular, gather) + actual = _run_update(jit_causal_conv1d_update, inputs, True) + expected = _run_update(aot_causal_conv1d_update, inputs, True) + _assert_bitwise_equal(actual[0], expected[0]) + _assert_bitwise_equal(actual[1], expected[1]) + + +############################################################################### +# Padding and argument validation. +############################################################################### + + +@pytest.mark.parametrize("width", [2, 4]) +def test_causal_conv1d_fwd_skips_padded_slots(width): + """Varlen sequences whose cache index is pad_slot_id are not processed.""" + dtype = torch.bfloat16 + batch, padding, dim, seqlen, entries = 4, 3, 64, 512, 40 + device = "cuda" + x = torch.randn(dim, seqlen, device=device, dtype=dtype) + x_before = x.clone() + weight = torch.randn(dim, width, device=device, dtype=dtype) + conv_states = torch.randn(entries, dim, width - 1, device=device, dtype=dtype) + conv_states_before = conv_states.clone() + # The trailing `padding` sequences are empty, so no output token belongs to + # them; only their conv-state slots would be touched without the pad check. + lengths = [seqlen // batch] * batch + [0] * padding + lengths[batch - 1] += seqlen - sum(lengths) + query_start_loc = torch.tensor( + [0] + torch.cumsum(torch.tensor(lengths), 0).tolist(), + dtype=torch.int32, + device=device, + ) + indices = torch.randperm(entries, device=device)[:batch].to(torch.int32) + padded_indices = torch.cat( + [ + indices, + torch.full((padding,), PAD_SLOT_ID, dtype=torch.int32, device=device), + ] + ) + has_initial_state = torch.zeros(batch + padding, dtype=torch.bool, device=device) + + jit_causal_conv1d_fwd( + x, + weight, + None, + conv_states, + query_start_loc, + padded_indices, + has_initial_state, + True, + PAD_SLOT_ID, + ) + + untouched = torch.ones(entries, dtype=torch.bool, device=device) + untouched[indices] = False + assert torch.equal(conv_states[untouched], conv_states_before[untouched]) + assert not torch.equal(x, x_before) + + +@pytest.mark.parametrize("width", [2, 4]) +def test_causal_conv1d_update_skips_padded_slots(width): + """Slots marked with pad_slot_id must be left untouched.""" + dtype = torch.bfloat16 + batch, padding, dim, entries = 3, 5, 128, 30 + device = "cuda" + x = torch.randn(batch + padding, dim, 1, device=device, dtype=dtype) + conv_state = torch.randn(entries, dim, width - 1, device=device, dtype=dtype) + conv_state_before = conv_state.clone() + weight = torch.randn(dim, width, device=device, dtype=dtype) + indices = torch.randperm(entries, device=device)[:batch].to(torch.int32) + padded_indices = torch.cat( + [ + indices, + torch.full((padding,), PAD_SLOT_ID, dtype=torch.int32, device=device), + ] + ) + + jit_causal_conv1d_update( + x, conv_state, weight, None, True, None, padded_indices, PAD_SLOT_ID + ) + + untouched = torch.ones(entries, dtype=torch.bool, device=device) + untouched[indices] = False + assert torch.equal(conv_state[untouched], conv_state_before[untouched]) + + +def test_causal_conv1d_rejects_unsupported_dtype(): + x = torch.ones((1, 8, 4), dtype=torch.int32, device="cuda") + weight = torch.ones((8, 4), dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="Unsupported dtype"): + jit_causal_conv1d_fwd(x, weight, None, None, None, None, None, True, -1) + + +def test_causal_conv1d_rejects_unsupported_width(): + x = torch.randn((1, 8, 4), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((8, 5), dtype=torch.bfloat16, device="cuda") + with pytest.raises(Exception, match="width between 2 and 4"): + jit_causal_conv1d_fwd(x, weight, None, None, None, None, None, True, -1) + + +@pytest.mark.parametrize("bad_dtype", [torch.uint8, torch.int32]) +def test_causal_conv1d_rejects_non_bool_has_initial_state(bad_dtype): + """Nothing normalizes this mask on the way in (unlike `cache_indices`), so a + wider dtype would silently read the wrong byte per sequence.""" + x = torch.randn((1, 8, 4), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((8, 4), dtype=torch.bfloat16, device="cuda") + conv_states = torch.zeros((1, 8, 3), dtype=torch.bfloat16, device="cuda") + has_initial_state = torch.ones((1,), dtype=bad_dtype, device="cuda") + with pytest.raises(Exception, match="has_initial_state must be a bool tensor"): + jit_causal_conv1d_fwd( + x, weight, None, conv_states, None, None, has_initial_state, True, -1 + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))