From 87549f8f0bf9e3514c1f6d30ba79844163b9221d Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Thu, 19 Mar 2026 19:38:46 -0700 Subject: [PATCH] perf(mamba): use Triton conv1d for non-contiguous input to avoid .contiguous() copy (#20469) --- .../layers/attention/mamba/causal_conv1d.py | 33 +++++++++++-------- python/sglang/test/run_eval.py | 4 +-- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py index e7beb75b0..aa9073a59 100644 --- a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py +++ b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py @@ -9,20 +9,23 @@ from typing import Optional import torch from .causal_conv1d_triton import PAD_SLOT_ID +from .causal_conv1d_triton import causal_conv1d_fn as _causal_conv1d_fn_triton +from .causal_conv1d_triton import causal_conv1d_update as _causal_conv1d_update_triton try: from sgl_kernel import causal_conv1d_fwd from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel torch.ops.sgl_kernel.causal_conv1d_update - _USE_TRITON = False + _HAS_SGL_KERNEL = True except (ImportError, AttributeError): - from .causal_conv1d_triton import causal_conv1d_fn as _causal_conv1d_fn_triton - from .causal_conv1d_triton import ( - causal_conv1d_update as _causal_conv1d_update_triton, - ) + _HAS_SGL_KERNEL = False - _USE_TRITON = True + +def _get_seq_lens_cpu(query_start_loc, x): + if query_start_loc is not None: + return (query_start_loc[1:] - query_start_loc[:-1]).cpu().tolist() + return [x.shape[-1]] def causal_conv1d_fn( @@ -66,19 +69,20 @@ def causal_conv1d_fn( out: (batch, dim, seqlen) """ - if _USE_TRITON: - seq_lens_cpu = ( - (query_start_loc[1:] - query_start_loc[:-1]).cpu().tolist() - if query_start_loc is not None - else [x.shape[-1]] - ) + # Use Triton when: (1) sgl_kernel not available, or (2) input is + # non-contiguous and seq_lens_cpu is already 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) + if use_triton: + if "seq_lens_cpu" not in kwargs: + kwargs["seq_lens_cpu"] = _get_seq_lens_cpu(query_start_loc, x) return _causal_conv1d_fn_triton( x, weight, bias, conv_states=conv_states, query_start_loc=query_start_loc, - seq_lens_cpu=seq_lens_cpu, cache_indices=cache_indices, has_initial_state=has_initial_state, activation=activation, @@ -137,7 +141,8 @@ def causal_conv1d_update( indices 0 and 3 out: (batch, dim) or (batch, dim, seqlen) """ - if _USE_TRITON: + use_triton = not _HAS_SGL_KERNEL + if use_triton: return _causal_conv1d_update_triton( x, conv_state, diff --git a/python/sglang/test/run_eval.py b/python/sglang/test/run_eval.py index 11ba1dc5c..51f650d4e 100644 --- a/python/sglang/test/run_eval.py +++ b/python/sglang/test/run_eval.py @@ -118,7 +118,7 @@ def run_eval(args): elif args.eval_name == "gpqa": from sglang.test.simple_eval_gpqa import GPQAEval - filename = ( + filename = getattr(args, "dataset_path", None) or ( "https://openaipublic.blob.core.windows.net/simple-evals/gpqa_diamond.csv" ) eval_obj = GPQAEval(filename, args.num_examples, args.num_threads) @@ -130,7 +130,7 @@ def run_eval(args): from sglang.test.simple_eval_longbench_v2 import LongBenchV2Eval # Default to HuggingFace dataset, can be overridden with --dataset-path - data_source = args.dataset_path + data_source = getattr(args, "dataset_path", None) categories = args.categories.split(",") if args.categories else None eval_obj = LongBenchV2Eval(