diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index 6080d4001..f677e1518 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py @@ -91,7 +91,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching # base of the sequence x_base = ( - x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim + x_ptr + + sequence_start_index.to(tl.int64) * stride_x_token + + idx_feats * stride_x_dim ) # [BLOCK_N,] if IS_CONTINUOUS_BATCHING: @@ -173,7 +175,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching ) # [BLOCK_M] x_ptrs = ( x_ptr - + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] + + ( + (sequence_start_index + idx_tokens_last).to(tl.int64) + * stride_x_token + )[:, None] + (idx_feats * stride_x_dim)[None, :] ) # [BLOCK_M,BLOCK_N,] mask_x = ( @@ -268,7 +273,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching else: # chunk_offset > 0 # read prior-token data from `x` load_init_state = True - prior_tokens = x_base + (token_offset - 1) * stride_x_token + prior_tokens = x_base + (token_offset - 1).to(tl.int64) * stride_x_token mask_w = idx_feats < dim if KERNEL_WIDTH == 2: conv_states_ptrs = prior_tokens # [BLOCK_N] @@ -305,7 +310,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching else: acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) - x_base_1d = x_base + token_offset * stride_x_token # starting of chunk + x_base_1d = x_base + token_offset.to(tl.int64) * stride_x_token # starting of chunk # PRE-LOAD WEIGHTS mask_w = idx_feats < dim @@ -372,7 +377,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching ) # token-index # feature-index o_ptrs = ( o_ptr - + (sequence_start_index + token_offset + idx_token) * stride_o_token + + (sequence_start_index + token_offset + idx_token).to(tl.int64) + * stride_o_token + (idx_feats * stride_o_dim) )