perf(mamba): use Triton conv1d for non-contiguous input to avoid .contiguous() copy (#20469)

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,

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(