minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(mamba): widen causal_conv1d token offsets to int64 (#33665)

Browse Source
This commit is contained in:
YAMY
2026-08-07 14:31:48 -07:00
committed by GitHub
parent b3ee679467
commit 9e3f6b746b
1 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/mamba/causal_conv1d_triton.py
+11 -5
View File
@@ -91,7 +91,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching
# base of the sequence # base of the sequence
x_base = ( 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,] ) # [BLOCK_N,]
if IS_CONTINUOUS_BATCHING: if IS_CONTINUOUS_BATCHING:
@@ -173,7 +175,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching
) # [BLOCK_M] ) # [BLOCK_M]
x_ptrs = ( x_ptrs = (
x_ptr 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, :] + (idx_feats * stride_x_dim)[None, :]
) # [BLOCK_M,BLOCK_N,] ) # [BLOCK_M,BLOCK_N,]
mask_x = ( mask_x = (
@@ -268,7 +273,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
else: # chunk_offset > 0 else: # chunk_offset > 0
# read prior-token data from `x` # read prior-token data from `x`
load_init_state = True 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 mask_w = idx_feats < dim
if KERNEL_WIDTH == 2: if KERNEL_WIDTH == 2:
conv_states_ptrs = prior_tokens # [BLOCK_N] conv_states_ptrs = prior_tokens # [BLOCK_N]
@@ -305,7 +310,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
else: else:
acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) 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 # PRE-LOAD WEIGHTS
mask_w = idx_feats < dim mask_w = idx_feats < dim
@@ -372,7 +377,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching
) # token-index # feature-index ) # token-index # feature-index
o_ptrs = ( o_ptrs = (
o_ptr 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) + (idx_feats * stride_o_dim)
) )