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

Fix issues/16714: Revert comment out of tl.debug_barrier() in causal_conv1d_triton (#16899)

Browse Source 
Co-authored-by: Yi Zhang <1109276519@qq.com>
This commit is contained in:
Netanel Haber
2026-01-14 17:26:48 +08:00
committed by GitHub
co-authored by Yi Zhang
parent 72bacc88c8
commit e75299a111
2 changed files with 3 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py
+3 -3
View File
@@ -186,7 +186,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
) )
mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
# tl.debug_barrier() # NOTE: use this due to bug in Triton compiler tl.debug_barrier() # NOTE: use this due to bug in Triton compiler
tl.store(conv_states_ptrs_target, new_conv_state, mask) tl.store(conv_states_ptrs_target, new_conv_state, mask)
else: else:
@@ -221,7 +221,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
) # token-index # token-index # feature-index ) # token-index # token-index # feature-index
loaded_x = tl.load(x_ptrs, mask_x, 0.0) loaded_x = tl.load(x_ptrs, mask_x, 0.0)
# tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load
new_conv_state = tl.where( new_conv_state = tl.where(
mask, conv_state, loaded_x mask, conv_state, loaded_x
) # BUG in 'tl.where' which requires a barrier before this ) # BUG in 'tl.where' which requires a barrier before this
@@ -728,7 +728,7 @@ def _causal_conv1d_update_kernel(
& (idx_feats < dim)[None, :] & (idx_feats < dim)[None, :]
) # token-index # token-index # feature-index ) # token-index # token-index # feature-index
loaded_x = tl.load(x_ptrs, mask_x, 0.0) loaded_x = tl.load(x_ptrs, mask_x, 0.0)
# tl.debug_barrier() tl.debug_barrier()
new_conv_state = tl.where(mask, conv_state, loaded_x) new_conv_state = tl.where(mask, conv_state, loaded_x)
test/registered/layers/mamba/test_causal_conv1d.py
-6
View File
@@ -196,9 +196,6 @@ def test_causal_conv1d_update_with_batch_gather(
if not torch.cuda.is_available(): if not torch.cuda.is_available():
pytest.skip("CUDA device not available") pytest.skip("CUDA device not available")
# TODO: Fix Triton kernel bug - see issue #16714
pytest.skip("Triton kernel bug - see issue #16714")
device = "cuda" device = "cuda"
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
if itype == torch.bfloat16: if itype == torch.bfloat16:
@@ -279,9 +276,6 @@ def test_causal_conv1d_varlen(
if not torch.cuda.is_available(): if not torch.cuda.is_available():
pytest.skip("CUDA device not available") pytest.skip("CUDA device not available")
# TODO: Fix Triton kernel bug - see issue #16714
pytest.skip("Triton kernel bug - see issue #16714")
device = "cuda" device = "cuda"
torch.cuda.empty_cache() torch.cuda.empty_cache()
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)