Add offline auto-tuning for LoRA CSGMV kernel (#20391)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

benchmark/kernels/lora_csgmv/tune_lora_csgmv.py

@@ -0,0 +1,747 @@
+"""
+Auto-tuning script for LoRA CSGMV (Chunked Segmented Matrix-Vector) kernels.
+
+LoRA adds low-rank adapters to linear layers. The two kernels are:
+  - Shrink (lora_a): x @ A^T, projecting from input_dim down to rank
+  - Expand (lora_b): (x @ A^T) @ B^T, projecting from rank back up to output_dim
+
+Terminology / dimensions:
+  K          For shrink: input_dim (the large dimension, e.g. hidden_size).
+             For expand: output_dim (e.g. hidden_size or qkv_output_dim).
+  R          Max LoRA rank (e.g. 16, 32, 64). The small dimension.
+  S          num_slices — how many weight slices a layer fuses together:
+               qkv_proj → 3 (q, k, v), gate_up_proj → 2, others → 1.
+             Affects the Triton grid (N = S * R for shrink, grid dim for expand).
+  chunk_size BLOCK_M — the max segment length in the chunked batch. Sequences
+             are split into fixed-size chunks for load-balanced GPU scheduling.
+             Typical values: 16, 32, 64, 128.
+
+Tuned parameters (per kernel, K, R, S, chunk_size):
+  BLOCK_N    Tile size along the N (output) dimension.
+  BLOCK_K    Tile size along the K (reduction) dimension.
+  num_warps  Number of warps per Triton program instance.
+  num_stages Number of software pipelining stages.
+  maxnreg    (expand only) Register cap to improve occupancy.
+
+Config files are saved as JSON keyed by chunk_size, e.g.:
+  lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H100.json
+
+The server loads these at startup via lora_tuning_config.py. If no tuned
+config exists, hardcoded defaults are used.
+
+Usage:
+    # Tune from model name (auto-derives hidden_size, QKV dims)
+    python benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
+        --model Qwen/Qwen3-0.6B --rank 64
+
+    # Tune with explicit dimensions
+    python benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
+        --hidden-size 1024 --rank 64
+
+    # Tune for specific chunk sizes
+    python benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
+        --model Qwen/Qwen3-0.6B --rank 64 --chunk-sizes 32 64 128
+
+    # Another model
+    python benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
+        --model meta-llama/Llama-2-7b-hf --rank 32
+"""
+
+import argparse
+import json
+import math
+import os
+import statistics
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import torch
+import triton
+
+from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel
+from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel
+from sglang.srt.lora.triton_ops.lora_tuning_config import (
+    DEFAULT_EXPAND_CONFIG,
+    DEFAULT_SHRINK_CONFIG,
+    get_lora_config_file_name,
+)
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+def _get_raw_kernel(cached_kernel):
+    """Get the underlying triton.jit function, bypassing cached_triton_kernel."""
+    return getattr(cached_kernel, "fn", cached_kernel)
+
+
+def build_batch_info(
+    total_tokens: int,
+    chunk_size: int,
+    rank: int,
+    device: torch.device,
+) -> LoRABatchInfo:
+    """Build a LoRABatchInfo for benchmarking with a single LoRA adapter."""
+    num_segments = math.ceil(total_tokens / chunk_size)
+
+    seg_indptr = []
+    for i in range(num_segments):
+        seg_indptr.append(i * chunk_size)
+    seg_indptr.append(total_tokens)
+    seg_indptr = torch.tensor(seg_indptr, dtype=torch.int32, device=device)
+
+    weight_indices = torch.ones(num_segments, dtype=torch.int32, device=device)
+    lora_ranks = torch.tensor([0, rank], dtype=torch.int32, device=device)
+    scalings = torch.ones(2, dtype=torch.float32, device=device)
+    permutation = torch.arange(total_tokens, dtype=torch.int32, device=device)
+
+    return LoRABatchInfo(
+        use_cuda_graph=False,
+        bs=1,
+        num_segments=num_segments,
+        max_len=chunk_size,
+        seg_indptr=seg_indptr,
+        weight_indices=weight_indices,
+        lora_ranks=lora_ranks,
+        scalings=scalings,
+        seg_lens=None,
+        permutation=permutation,
+    )
+
+
+def timed_cuda_ms(fn, warmup: int = 10, trials: int = 50) -> float:
+    """Time a GPU function using CUDA events. Returns median time in ms."""
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+
+    times = []
+    for _ in range(trials):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return statistics.median(times)
+
+
+# ---------------------------------------------------------------------------
+# Search spaces
+# ---------------------------------------------------------------------------
+
+
+def get_shrink_search_space() -> List[Dict[str, Any]]:
+    """Generate candidate configs for the shrink kernel."""
+    configs = []
+    for block_n in [16, 32, 64]:
+        for block_k in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [2, 3, 4]:
+                    configs.append(
+                        {
+                            "BLOCK_N": block_n,
+                            "BLOCK_K": block_k,
+                            "num_warps": num_warps,
+                            "num_stages": num_stages,
+                        }
+                    )
+    return configs
+
+
+def get_expand_search_space() -> List[Dict[str, Any]]:
+    """Generate candidate configs for the expand kernel."""
+    configs = []
+    for block_n in [32, 64]:
+        for block_k in [16, 32]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    # Without maxnreg
+                    configs.append(
+                        {
+                            "BLOCK_N": block_n,
+                            "BLOCK_K": block_k,
+                            "num_warps": num_warps,
+                            "num_stages": num_stages,
+                        }
+                    )
+                    # With maxnreg (register capping for occupancy)
+                    for maxnreg in [96, 112, 128, 160]:
+                        configs.append(
+                            {
+                                "BLOCK_N": block_n,
+                                "BLOCK_K": block_k,
+                                "num_warps": num_warps,
+                                "num_stages": num_stages,
+                                "maxnreg": maxnreg,
+                            }
+                        )
+    return configs
+
+
+# ---------------------------------------------------------------------------
+# Benchmark functions
+# ---------------------------------------------------------------------------
+
+
+def benchmark_shrink_config(
+    config: Dict[str, Any],
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    num_slices: int,
+    N: int,
+    K: int,
+) -> Optional[float]:
+    """Benchmark a single shrink config. Returns median ms or None on failure."""
+    kernel = _get_raw_kernel(_chunked_lora_shrink_kernel)
+    S = x.shape[0]
+    num_segments = batch_info.num_segments
+
+    grid = (triton.cdiv(N, config["BLOCK_N"]), num_segments)
+    output = torch.empty((S, N), device=x.device, dtype=x.dtype)
+
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+
+    try:
+        kernel[grid](
+            x=x,
+            weights=weights,
+            output=output,
+            seg_indptr=batch_info.seg_indptr,
+            weight_indices=batch_info.weight_indices,
+            lora_ranks=batch_info.lora_ranks,
+            permutation=batch_info.permutation,
+            num_segs=num_segments,
+            N=N,
+            K=K,
+            NUM_SLICES=num_slices,
+            BLOCK_M=batch_info.max_len,
+            BLOCK_N=config["BLOCK_N"],
+            BLOCK_K=config["BLOCK_K"],
+            **extra_kwargs,
+        )
+        torch.cuda.synchronize()
+    except Exception:
+        return None
+
+    def run():
+        kernel[grid](
+            x=x,
+            weights=weights,
+            output=output,
+            seg_indptr=batch_info.seg_indptr,
+            weight_indices=batch_info.weight_indices,
+            lora_ranks=batch_info.lora_ranks,
+            permutation=batch_info.permutation,
+            num_segs=num_segments,
+            N=N,
+            K=K,
+            NUM_SLICES=num_slices,
+            BLOCK_M=batch_info.max_len,
+            BLOCK_N=config["BLOCK_N"],
+            BLOCK_K=config["BLOCK_K"],
+            **extra_kwargs,
+        )
+
+    return timed_cuda_ms(run, warmup=10, trials=50)
+
+
+def benchmark_expand_config(
+    config: Dict[str, Any],
+    x: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    slice_offsets: torch.Tensor,
+    max_slice_size: int,
+    output_dim: int,
+    num_slices: int,
+    max_rank: int,
+) -> Optional[float]:
+    """Benchmark a single expand config. Returns median ms or None on failure."""
+    kernel = _get_raw_kernel(_chunked_lora_expand_kernel)
+    M = x.shape[0]
+    num_segments = batch_info.num_segments
+
+    grid = (
+        triton.cdiv(max_slice_size, config["BLOCK_N"]),
+        num_slices,
+        num_segments,
+    )
+    output = torch.zeros((M, output_dim), device=x.device, dtype=x.dtype)
+
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+    if "maxnreg" in config:
+        extra_kwargs["maxnreg"] = config["maxnreg"]
+
+    try:
+        kernel[grid](
+            x=x,
+            weights=weights,
+            output=output,
+            seg_indptr=batch_info.seg_indptr,
+            weight_indices=batch_info.weight_indices,
+            lora_ranks=batch_info.lora_ranks,
+            permutation=batch_info.permutation,
+            num_segs=num_segments,
+            scalings=batch_info.scalings,
+            slice_offsets=slice_offsets,
+            NUM_SLICES=num_slices,
+            OUTPUT_DIM=output_dim,
+            MAX_RANK=max_rank,
+            BLOCK_M=batch_info.max_len,
+            BLOCK_N=config["BLOCK_N"],
+            BLOCK_K=config["BLOCK_K"],
+            **extra_kwargs,
+        )
+        torch.cuda.synchronize()
+    except Exception:
+        return None
+
+    def run():
+        output.zero_()
+        kernel[grid](
+            x=x,
+            weights=weights,
+            output=output,
+            seg_indptr=batch_info.seg_indptr,
+            weight_indices=batch_info.weight_indices,
+            lora_ranks=batch_info.lora_ranks,
+            permutation=batch_info.permutation,
+            num_segs=num_segments,
+            scalings=batch_info.scalings,
+            slice_offsets=slice_offsets,
+            NUM_SLICES=num_slices,
+            OUTPUT_DIM=output_dim,
+            MAX_RANK=max_rank,
+            BLOCK_M=batch_info.max_len,
+            BLOCK_N=config["BLOCK_N"],
+            BLOCK_K=config["BLOCK_K"],
+            **extra_kwargs,
+        )
+
+    return timed_cuda_ms(run, warmup=10, trials=50)
+
+
+# ---------------------------------------------------------------------------
+# Config saving
+# ---------------------------------------------------------------------------
+
+
+def save_config(
+    configs: Dict[int, Dict[str, Any]],
+    kernel: str,
+    major_dim: int,
+    max_rank: int,
+    num_slices: int,
+) -> str:
+    """Save tuned configs to the standard config directory. Returns filepath.
+
+    Args:
+        configs: Dict mapping chunk_size -> best block config.
+        kernel: "shrink" or "expand".
+        major_dim: The large dimension (input_dim for shrink, output_dim for expand).
+        max_rank: The max LoRA rank.
+        num_slices: Number of fused weight slices (qkv=3, gate_up=2, others=1).
+    """
+    filename = get_lora_config_file_name(kernel, major_dim, max_rank, num_slices)
+
+    triton_version = triton.__version__
+    version_dir = f"triton_{triton_version.replace('.', '_')}"
+    config_dir = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "..",
+        "..",
+        "..",
+        "python",
+        "sglang",
+        "srt",
+        "lora",
+        "triton_ops",
+        "csgmv_configs",
+        version_dir,
+    )
+    config_dir = os.path.normpath(config_dir)
+    os.makedirs(config_dir, exist_ok=True)
+
+    filepath = os.path.join(config_dir, filename)
+    with open(filepath, "w") as f:
+        json.dump(configs, f, indent=4)
+        f.write("\n")
+    return filepath
+
+
+def sort_config(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Sort config keys for consistent JSON output."""
+    ordered = {}
+    for key in ["BLOCK_N", "BLOCK_K", "num_warps", "num_stages", "maxnreg"]:
+        if key in config:
+            ordered[key] = config[key]
+    return ordered
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def get_model_dims(args: argparse.Namespace):
+    """Extract all LoRA layer dimensions from model config or CLI args.
+
+    Returns a list of (label, shrink_K, expand_output_dim, num_slices,
+    slice_offsets_list) tuples for each LoRA layer type.
+    """
+    if args.model:
+        from transformers import AutoConfig
+
+        config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
+        hidden_size = config.hidden_size
+
+        num_heads = config.num_attention_heads
+        num_kv_heads = getattr(config, "num_key_value_heads", num_heads)
+        head_dim = getattr(config, "head_dim", hidden_size // num_heads)
+        intermediate_size = config.intermediate_size
+
+        q_dim = num_heads * head_dim
+        kv_dim = num_kv_heads * head_dim
+        qkv_output_dim = q_dim + 2 * kv_dim
+
+        print(f"Model: {args.model}")
+        print(
+            f"  hidden_size={hidden_size}, num_heads={num_heads}, "
+            f"num_kv_heads={num_kv_heads}, head_dim={head_dim}"
+        )
+        print(f"  intermediate_size={intermediate_size}")
+    else:
+        hidden_size = args.hidden_size
+        intermediate_size = getattr(args, "intermediate_size", None) or hidden_size * 3
+        if args.qkv_output_dim:
+            qkv_output_dim = args.qkv_output_dim
+            q_dim = qkv_output_dim // 2
+            kv_dim = (qkv_output_dim - q_dim) // 2
+        else:
+            q_dim = hidden_size * 2
+            kv_dim = hidden_size
+            qkv_output_dim = q_dim + 2 * kv_dim
+
+    # All LoRA layer types with their dimensions:
+    #   (label, shrink_K, expand_output_dim, num_slices, slice_offsets)
+    layers = [
+        (
+            "qkv",
+            hidden_size,
+            qkv_output_dim,
+            3,
+            [0, q_dim, q_dim + kv_dim, qkv_output_dim],
+        ),
+        ("o_proj", q_dim, hidden_size, 1, [0, hidden_size]),
+        (
+            "gate_up",
+            hidden_size,
+            2 * intermediate_size,
+            2,
+            [0, intermediate_size, 2 * intermediate_size],
+        ),
+        ("down_proj", intermediate_size, hidden_size, 1, [0, hidden_size]),
+    ]
+
+    print(f"\nLoRA layer dimensions:")
+    for label, sk, eo, ns, so in layers:
+        print(f"  {label:>10}: shrink K={sk}, expand output_dim={eo}, num_slices={ns}")
+
+    return layers
+
+
+def _tune_shrink(
+    label: str,
+    K: int,
+    N: int,
+    num_slices: int,
+    rank: int,
+    chunk_sizes: List[int],
+    total_tokens: int,
+    device: torch.device,
+) -> tuple:
+    """Tune shrink kernel for one layer type. Returns (best_configs, results)."""
+    print(f"\n{'='*80}")
+    print(f"Tuning SHRINK — {label} (K={K}, N={N}, slices={num_slices})")
+    print(f"{'='*80}")
+
+    search = get_shrink_search_space()
+    print(f"Search space: {len(search)} configs")
+
+    best_configs = {}
+    results = {}
+
+    for chunk_size in chunk_sizes:
+        batch_info = build_batch_info(total_tokens, chunk_size, rank, device)
+        x = torch.randn(total_tokens, K, device=device, dtype=torch.float16)
+        weights = torch.randn(2, N, K, device=device, dtype=torch.float16)
+
+        baseline_time = benchmark_shrink_config(
+            DEFAULT_SHRINK_CONFIG,
+            x,
+            weights,
+            batch_info,
+            num_slices,
+            N,
+            K,
+        )
+        print(f"  chunk={chunk_size}: baseline={baseline_time:.3f}ms")
+
+        best_config = None
+        best_time = float("inf")
+
+        for i, config in enumerate(search):
+            t = benchmark_shrink_config(
+                config, x, weights, batch_info, num_slices, N, K
+            )
+            if t is not None and t < best_time:
+                best_time = t
+                best_config = config
+            if (i + 1) % 20 == 0:
+                print(
+                    f"  chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms"
+                )
+
+        best_configs[chunk_size] = sort_config(best_config)
+        results[chunk_size] = (baseline_time, best_time, best_configs[chunk_size])
+        speedup = baseline_time / best_time if best_time > 0 else 0
+        print(
+            f"  chunk={chunk_size}: best={best_time:.3f}ms ({speedup:.2f}x), config={best_configs[chunk_size]}"
+        )
+
+    return best_configs, results
+
+
+def _tune_expand(
+    label: str,
+    output_dim: int,
+    num_slices: int,
+    slice_offsets_list: List[int],
+    max_slice_size: int,
+    rank: int,
+    chunk_sizes: List[int],
+    total_tokens: int,
+    device: torch.device,
+) -> tuple:
+    """Tune expand kernel for one layer type. Returns (best_configs, results)."""
+    print(f"\n{'='*80}")
+    print(f"Tuning EXPAND — {label} (output_dim={output_dim}, slices={num_slices})")
+    print(f"{'='*80}")
+
+    search = get_expand_search_space()
+    print(f"Search space: {len(search)} configs")
+
+    slice_offsets = torch.tensor(slice_offsets_list, dtype=torch.int64, device=device)
+    best_configs = {}
+    results = {}
+
+    for chunk_size in chunk_sizes:
+        batch_info = build_batch_info(total_tokens, chunk_size, rank, device)
+        x = torch.randn(
+            total_tokens, num_slices * rank, device=device, dtype=torch.float16
+        )
+        weights = torch.randn(2, output_dim, rank, device=device, dtype=torch.float16)
+
+        baseline_time = benchmark_expand_config(
+            DEFAULT_EXPAND_CONFIG,
+            x,
+            weights,
+            batch_info,
+            slice_offsets,
+            max_slice_size,
+            output_dim,
+            num_slices,
+            rank,
+        )
+        print(f"  chunk={chunk_size}: baseline={baseline_time:.3f}ms")
+
+        best_config = None
+        best_time = float("inf")
+
+        for i, config in enumerate(search):
+            t = benchmark_expand_config(
+                config,
+                x,
+                weights,
+                batch_info,
+                slice_offsets,
+                max_slice_size,
+                output_dim,
+                num_slices,
+                rank,
+            )
+            if t is not None and t < best_time:
+                best_time = t
+                best_config = config
+            if (i + 1) % 50 == 0:
+                print(
+                    f"  chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms"
+                )
+
+        best_configs[chunk_size] = sort_config(best_config)
+        results[chunk_size] = (baseline_time, best_time, best_configs[chunk_size])
+        speedup = baseline_time / best_time if best_time > 0 else 0
+        print(
+            f"  chunk={chunk_size}: best={best_time:.3f}ms ({speedup:.2f}x), config={best_configs[chunk_size]}"
+        )
+
+    return best_configs, results
+
+
+def main(args: argparse.Namespace):
+    device = torch.device("cuda:0")
+    rank = args.rank
+    chunk_sizes = args.chunk_sizes
+    total_tokens = args.total_tokens
+
+    layers = get_model_dims(args)
+
+    print(f"\nLoRA CSGMV Tuning")
+    print(f"  rank={rank}, total_tokens={total_tokens}, chunk_sizes={chunk_sizes}")
+
+    # Collect all results for summary
+    all_results = []  # (label, kernel, K_or_outdim, results_dict)
+
+    # Deduplicate: multiple layers can share the same (shrink_K, num_slices) or
+    # (expand_output_dim, num_slices). No need to tune the same config twice.
+    tuned_shrink = {}  # (shrink_K, num_slices) -> best_configs
+    tuned_expand = {}  # (expand_output_dim, num_slices) -> best_configs
+
+    for label, shrink_K, expand_output_dim, num_slices, slice_offsets_list in layers:
+        # --- Shrink ---
+        shrink_key = (shrink_K, num_slices)
+        if shrink_key not in tuned_shrink:
+            N_shrink = num_slices * rank
+            best_configs, results = _tune_shrink(
+                label,
+                shrink_K,
+                N_shrink,
+                num_slices,
+                rank,
+                chunk_sizes,
+                total_tokens,
+                device,
+            )
+            filepath = save_config(best_configs, "shrink", shrink_K, rank, num_slices)
+            print(f"  Saved to: {filepath}")
+            tuned_shrink[shrink_key] = best_configs
+            all_results.append((label, "shrink", shrink_K, results))
+        else:
+            print(
+                f"\n  Skipping shrink {label} (K={shrink_K}, S={num_slices}) — already tuned"
+            )
+
+        # --- Expand ---
+        expand_key = (expand_output_dim, num_slices)
+        if expand_key not in tuned_expand:
+            # max_slice_size = largest slice width
+            slice_widths = [
+                slice_offsets_list[i + 1] - slice_offsets_list[i]
+                for i in range(num_slices)
+            ]
+            max_slice_size = max(slice_widths)
+
+            best_configs, results = _tune_expand(
+                label,
+                expand_output_dim,
+                num_slices,
+                slice_offsets_list,
+                max_slice_size,
+                rank,
+                chunk_sizes,
+                total_tokens,
+                device,
+            )
+            filepath = save_config(
+                best_configs, "expand", expand_output_dim, rank, num_slices
+            )
+            print(f"  Saved to: {filepath}")
+            tuned_expand[expand_key] = best_configs
+            all_results.append((label, "expand", expand_output_dim, results))
+        else:
+            print(
+                f"\n  Skipping expand {label} (output_dim={expand_output_dim}, S={num_slices}) — already tuned"
+            )
+
+    # --- Summary ---
+    print(f"\n{'='*80}")
+    print(f"SUMMARY")
+    print(f"{'='*80}")
+    print(
+        f"\n{'layer':<10} {'kernel':<8} {'K/dim':>6} {'chunk':>6}"
+        f" {'baseline':>10} {'tuned':>10} {'speedup':>8}  config"
+    )
+    print("-" * 100)
+    for label, kernel, dim, results in all_results:
+        for chunk_size in chunk_sizes:
+            if chunk_size in results:
+                base, best, cfg = results[chunk_size]
+                spd = base / best if best > 0 else 0
+                print(
+                    f"{label:<10} {kernel:<8} {dim:>6} {chunk_size:>6}"
+                    f" {base:>9.3f}ms {best:>9.3f}ms {spd:>7.2f}x  {cfg}"
+                )
+
+    now = datetime.now()
+    print(f"\nTuning completed at {now.ctime()}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Auto-tune LoRA CSGMV kernel block dimensions"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="HuggingFace model name to auto-derive dimensions "
+        "(e.g., Qwen/Qwen3-0.6B, meta-llama/Llama-2-7b-hf)",
+    )
+    parser.add_argument(
+        "--hidden-size",
+        type=int,
+        default=None,
+        help="Model hidden size (e.g., 1024 for Qwen3-0.6B). "
+        "Required if --model is not specified.",
+    )
+    parser.add_argument(
+        "--rank",
+        type=int,
+        required=True,
+        help="LoRA rank (e.g., 16, 32, 64)",
+    )
+    parser.add_argument(
+        "--qkv-output-dim",
+        type=int,
+        default=None,
+        help="QKV output dimension. Only used with --hidden-size. "
+        "Default: 4 * hidden_size",
+    )
+    parser.add_argument(
+        "--chunk-sizes",
+        type=int,
+        nargs="+",
+        default=[16, 32, 64, 128],
+        help="Chunk sizes to tune (default: 16 32 64 128)",
+    )
+    parser.add_argument(
+        "--total-tokens",
+        type=int,
+        default=30720,
+        help="Total tokens for benchmarking (default: 30720 = 2 reqs x 15360)",
+    )
+    args = parser.parse_args()
+
+    if not args.model and not args.hidden_size:
+        parser.error("Either --model or --hidden-size is required")
+
+    main(args)

python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py

@@ -4,6 +4,7 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.lora.triton_ops.lora_tuning_config import get_lora_expand_config
 from sglang.srt.lora.utils import LoRABatchInfo
 from sglang.srt.utils import cached_triton_kernel
 
@@ -175,10 +176,13 @@ def chunked_sgmv_lora_expand_forward(
     num_slices = len(slice_offsets) - 1
     assert input_dim == num_slices * MAX_RANK
 
-    # TODO (lifuhuang): fine-tune per operation
+    # Block shapes — use auto-tuned config if available, else defaults
     BLOCK_M = batch_info.max_len
-    BLOCK_K = 16
-    BLOCK_N = 64
+    config = get_lora_expand_config(
+        K=OUTPUT_DIM, R=MAX_RANK, num_slices=num_slices, chunk_size=BLOCK_M
+    )
+    BLOCK_K = config["BLOCK_K"]
+    BLOCK_N = config["BLOCK_N"]
 
     num_segments = batch_info.num_segments
 
@@ -193,6 +197,15 @@ def chunked_sgmv_lora_expand_forward(
     else:
         output = base_output
 
+    # Optional launch params from tuned config
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+    if "maxnreg" in config:
+        extra_kwargs["maxnreg"] = config["maxnreg"]
+
     _chunked_lora_expand_kernel[grid](
         x=x,
         weights=weights,
@@ -211,6 +224,7 @@ def chunked_sgmv_lora_expand_forward(
         BLOCK_M=BLOCK_M,
         BLOCK_N=BLOCK_N,
         BLOCK_K=BLOCK_K,
+        **extra_kwargs,
     )
 
     return output

python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py

@@ -2,6 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.lora.triton_ops.lora_tuning_config import get_lora_shrink_config
 from sglang.srt.lora.utils import LoRABatchInfo
 from sglang.srt.utils import cached_triton_kernel
 
@@ -137,11 +138,15 @@ def chunked_sgmv_lora_shrink_forward(
     assert len(x.shape) == 2
     assert len(weights.shape) == 3
 
-    # Block shapes
-    # TODO (lifuhuang): experiment with split-k
+    # Block shapes — use auto-tuned config if available, else defaults
     BLOCK_M = batch_info.max_len
-    BLOCK_N = 16
-    BLOCK_K = 256
+    # weights shape is (num_lora, num_slices * rank, input_dim)
+    MAX_RANK = weights.shape[1] // num_slices
+    config = get_lora_shrink_config(
+        K=weights.shape[2], R=MAX_RANK, num_slices=num_slices, chunk_size=BLOCK_M
+    )
+    BLOCK_N = config["BLOCK_N"]
+    BLOCK_K = config["BLOCK_K"]
 
     S = x.shape[0]
     N = weights.shape[1]
@@ -154,6 +159,13 @@ def chunked_sgmv_lora_shrink_forward(
         batch_info.bs if batch_info.use_cuda_graph else num_segments,
     )
 
+    # Optional launch params from tuned config
+    extra_kwargs = {}
+    if "num_warps" in config:
+        extra_kwargs["num_warps"] = config["num_warps"]
+    if "num_stages" in config:
+        extra_kwargs["num_stages"] = config["num_stages"]
+
     output = torch.empty((S, N), device=x.device, dtype=x.dtype)
     _chunked_lora_shrink_kernel[grid](
         x=x,
@@ -171,6 +183,7 @@ def chunked_sgmv_lora_shrink_forward(
         BLOCK_M=BLOCK_M,
         BLOCK_N=BLOCK_N,
         BLOCK_K=BLOCK_K,
+        **extra_kwargs,
     )
 
     return output

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 1,
+        "maxnreg": 128
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 2,
+        "maxnreg": 128
+    }
+}

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3,
+        "maxnreg": 160
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 1,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 2,
+        "maxnreg": 128
+    }
+}

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json

@@ -0,0 +1,29 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2,
+        "maxnreg": 112
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 32,
+        "num_warps": 4,
+        "num_stages": 3,
+        "maxnreg": 160
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 16,
+        "num_warps": 8,
+        "num_stages": 3,
+        "maxnreg": 128
+    }
+}

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    }
+}

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 128,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 4
+    }
+}

python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json

@@ -0,0 +1,26 @@
+{
+    "16": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 128,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "64": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_N": 64,
+        "BLOCK_K": 64,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

python/sglang/srt/lora/triton_ops/lora_tuning_config.py

@@ -0,0 +1,201 @@
+"""
+Configuration loader for auto-tuned LoRA CSGMV kernel block sizes.
+
+Follows the same pattern as fused_moe_triton_config.py:
+- Offline tuning script writes JSON files keyed by chunk_size (BLOCK_M)
+- At server startup, the config loader reads the best block sizes for each kernel
+- Kernels use these instead of hardcoded defaults
+
+Config file naming: lora_{kernel},K={K},R={R},S={S},device={device}.json
+Where kernel is "shrink" or "expand", K is input_dim, R is max_rank, S is num_slices.
+
+Config file format (keyed by chunk_size):
+{
+    "16": {"BLOCK_N": 16, "BLOCK_K": 256, "num_warps": 4, "num_stages": 3},
+    "32": {"BLOCK_N": 32, "BLOCK_K": 128, "num_warps": 4, "num_stages": 4},
+    "128": {"BLOCK_N": 64, "BLOCK_K": 256, "num_warps": 8, "num_stages": 3}
+}
+
+Usage:
+    python3 benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
+        --model Qwen/Qwen3-Embedding-0.6B --max-lora-rank 64
+
+    # Configs saved to python/sglang/srt/lora/triton_ops/configs/
+
+    # Server automatically picks them up:
+    python3 -m sglang.launch_server --model ... --enable-lora --lora-backend csgmv
+"""
+
+from __future__ import annotations
+
+import functools
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+import triton
+
+from sglang.srt.utils import get_device_name
+
+logger = logging.getLogger(__name__)
+
+
+def get_lora_config_file_name(
+    kernel: str,
+    K: int,
+    R: int,
+    S: int,
+) -> str:
+    """Generate config filename for a LoRA kernel configuration.
+
+    Args:
+        kernel: "shrink" or "expand"
+        K: The large dimension (input_dim for shrink, output_dim for expand)
+        R: The max LoRA rank
+        S: num_slices (qkv=3, gate_up=2, others=1)
+    """
+    device_name = get_device_name().replace(" ", "_")
+    return f"lora_{kernel},K={K},R={R},S={S},device={device_name}.json"
+
+
+@functools.lru_cache
+def get_lora_configs(
+    kernel: str,
+    K: int,
+    R: int,
+    S: int,
+) -> Optional[Dict[int, Dict[str, Any]]]:
+    """Load pre-tuned LoRA kernel configs from JSON files.
+
+    Returns a dict mapping chunk_size (BLOCK_M) to block size configs,
+    or None if no config file is found.
+    """
+    json_file_name = get_lora_config_file_name(kernel, K, R, S)
+
+    config_dir = os.environ.get(
+        "SGLANG_LORA_CONFIG_DIR", os.path.dirname(os.path.realpath(__file__))
+    )
+    configs_root = os.path.join(config_dir, "csgmv_configs")
+
+    triton_version = triton.__version__
+    version_dir = f"triton_{triton_version.replace('.', '_')}"
+
+    # Try exact triton version first
+    config_file_path = os.path.join(configs_root, version_dir, json_file_name)
+    if os.path.exists(config_file_path):
+        with open(config_file_path) as f:
+            logger.info(f"Using LoRA {kernel} config from {config_file_path}.")
+            return {int(key): val for key, val in json.load(f).items()}
+
+    # Scan existing version directories as fallback (newest first)
+    if os.path.isdir(configs_root):
+        version_dirs = sorted(
+            (d for d in os.listdir(configs_root) if d.startswith("triton_")),
+            reverse=True,
+        )
+        for vdir in version_dirs:
+            if vdir == version_dir:
+                continue
+            try_path = os.path.join(configs_root, vdir, json_file_name)
+            if os.path.exists(try_path):
+                with open(try_path) as f:
+                    logger.warning(
+                        f"LoRA {kernel} config not found for Triton {triton_version}. "
+                        f"Falling back to {try_path}."
+                    )
+                    return {int(key): val for key, val in json.load(f).items()}
+
+    return None
+
+
+# Default block sizes (current hardcoded values)
+DEFAULT_SHRINK_CONFIG = {"BLOCK_N": 16, "BLOCK_K": 256}
+DEFAULT_EXPAND_CONFIG = {"BLOCK_N": 64, "BLOCK_K": 16}
+
+# Track which configs have been logged to avoid spamming on every forward pass
+_logged_configs: set = set()
+
+
+def get_lora_shrink_config(
+    K: int,
+    R: int,
+    num_slices: int,
+    chunk_size: int,
+) -> Dict[str, int]:
+    """Get block sizes for the CSGMV shrink (lora_a) kernel.
+
+    Args:
+        K: input_dim
+        R: max_rank
+        num_slices: number of slices (qkv=3, gate_up=2, others=1)
+        chunk_size: BLOCK_M value (= batch_info.max_len)
+    """
+    log_key = ("shrink", K, R, num_slices, chunk_size)
+    configs = get_lora_configs("shrink", K, R, num_slices)
+    if configs is not None:
+        config = configs.get(chunk_size)
+        if config is None:
+            closest = min(configs.keys(), key=lambda x: abs(x - chunk_size))
+            config = configs[closest]
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA shrink (K={K}, R={R}): no config for chunk_size={chunk_size}, "
+                    f"using closest={closest}: {config}"
+                )
+        else:
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA shrink (K={K}, R={R}, chunk_size={chunk_size}): tuned config {config}"
+                )
+        return config
+    if log_key not in _logged_configs:
+        _logged_configs.add(log_key)
+        logger.info(
+            f"LoRA shrink (K={K}, R={R}): no tuned config, using defaults {DEFAULT_SHRINK_CONFIG}"
+        )
+    return dict(DEFAULT_SHRINK_CONFIG)
+
+
+def get_lora_expand_config(
+    K: int,
+    R: int,
+    num_slices: int,
+    chunk_size: int,
+) -> Dict[str, int]:
+    """Get block sizes for the CSGMV expand (lora_b) kernel.
+
+    Args:
+        K: output_dim
+        R: max_rank
+        num_slices: number of slices (qkv=3, gate_up=2, others=1)
+        chunk_size: BLOCK_M value (= batch_info.max_len)
+    """
+    log_key = ("expand", K, R, num_slices, chunk_size)
+    configs = get_lora_configs("expand", K, R, num_slices)
+    if configs is not None:
+        config = configs.get(chunk_size)
+        if config is None:
+            closest = min(configs.keys(), key=lambda x: abs(x - chunk_size))
+            config = configs[closest]
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA expand (K={K}, R={R}): no config for chunk_size={chunk_size}, "
+                    f"using closest={closest}: {config}"
+                )
+        else:
+            if log_key not in _logged_configs:
+                _logged_configs.add(log_key)
+                logger.info(
+                    f"LoRA expand (K={K}, R={R}, chunk_size={chunk_size}): tuned config {config}"
+                )
+        return config
+    if log_key not in _logged_configs:
+        _logged_configs.add(log_key)
+        logger.info(
+            f"LoRA expand (K={K}, R={R}): no tuned config, using defaults {DEFAULT_EXPAND_CONFIG}"
+        )
+    return dict(DEFAULT_EXPAND_CONFIG)

test/manual/lora/test_lora_tuning_config.py

@@ -0,0 +1,118 @@
+"""Unit tests for LoRA CSGMV tuning config loading."""
+
+import json
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from sglang.srt.lora.triton_ops.lora_tuning_config import (
+    DEFAULT_EXPAND_CONFIG,
+    DEFAULT_SHRINK_CONFIG,
+    get_lora_config_file_name,
+    get_lora_configs,
+    get_lora_expand_config,
+    get_lora_shrink_config,
+)
+
+_MODULE = "sglang.srt.lora.triton_ops.lora_tuning_config"
+
+# Shared fixture
+_TUNED_CONFIGS = {
+    32: {"BLOCK_N": 32, "BLOCK_K": 128, "num_warps": 4, "num_stages": 3},
+    128: {"BLOCK_N": 64, "BLOCK_K": 256, "num_warps": 8, "num_stages": 2},
+}
+
+
+class TestLoraConfigFileName(unittest.TestCase):
+    @patch(f"{_MODULE}.get_device_name", return_value="NVIDIA H100")
+    def test_includes_all_params(self, _):
+        name = get_lora_config_file_name("shrink", K=1024, R=64, S=3)
+        self.assertEqual(name, "lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H100.json")
+
+    @patch(f"{_MODULE}.get_device_name", return_value="GPU")
+    def test_different_slices_different_filenames(self, _):
+        s1 = get_lora_config_file_name("shrink", 1024, 64, S=1)
+        s3 = get_lora_config_file_name("shrink", 1024, 64, S=3)
+        self.assertNotEqual(s1, s3)
+
+
+class TestLoraConfigLoading(unittest.TestCase):
+    def setUp(self):
+        get_lora_configs.cache_clear()
+        self.tmpdir = tempfile.mkdtemp()
+
+    def _write_config(self, triton_ver_dir, filename, data):
+        d = os.path.join(self.tmpdir, "csgmv_configs", triton_ver_dir)
+        os.makedirs(d, exist_ok=True)
+        with open(os.path.join(d, filename), "w") as f:
+            json.dump(data, f)
+
+    @patch(f"{_MODULE}.get_device_name", return_value="TestGPU")
+    @patch(f"{_MODULE}.triton")
+    def test_load_and_fallback(self, mock_triton, _):
+        """Loads exact version, falls back to other version, returns None if missing."""
+        config_data = {"32": {"BLOCK_N": 32, "BLOCK_K": 128}}
+        self._write_config(
+            "triton_3_5_1",
+            "lora_shrink,K=1024,R=64,S=3,device=TestGPU.json",
+            config_data,
+        )
+
+        # Exact match
+        mock_triton.__version__ = "3.5.1"
+        with patch.dict(os.environ, {"SGLANG_LORA_CONFIG_DIR": self.tmpdir}):
+            result = get_lora_configs("shrink", 1024, 64, 3)
+        self.assertEqual(result[32]["BLOCK_N"], 32)
+
+        # Fallback from newer version
+        get_lora_configs.cache_clear()
+        mock_triton.__version__ = "3.6.0"
+        with patch.dict(os.environ, {"SGLANG_LORA_CONFIG_DIR": self.tmpdir}):
+            result = get_lora_configs("shrink", 1024, 64, 3)
+        self.assertIsNotNone(result)
+
+        # Missing config returns None
+        get_lora_configs.cache_clear()
+        with patch.dict(os.environ, {"SGLANG_LORA_CONFIG_DIR": self.tmpdir}):
+            self.assertIsNone(get_lora_configs("shrink", 9999, 64, 1))
+
+
+class TestConfigSelection(unittest.TestCase):
+    """Test exact match, nearest-neighbor, and default fallback for both kernels."""
+
+    KERNELS = [
+        (get_lora_shrink_config, DEFAULT_SHRINK_CONFIG),
+        (get_lora_expand_config, DEFAULT_EXPAND_CONFIG),
+    ]
+
+    def setUp(self):
+        get_lora_configs.cache_clear()
+        from sglang.srt.lora.triton_ops import lora_tuning_config
+
+        lora_tuning_config._logged_configs.clear()
+
+    def test_defaults_when_no_config(self):
+        for get_fn, default in self.KERNELS:
+            with self.subTest(fn=get_fn.__name__):
+                with patch(f"{_MODULE}.get_lora_configs", return_value=None):
+                    config = get_fn(K=1024, R=64, num_slices=1, chunk_size=32)
+                self.assertEqual(config, default)
+
+    def test_exact_and_nearest_neighbor(self):
+        for get_fn, _ in self.KERNELS:
+            with self.subTest(fn=get_fn.__name__):
+                with patch(f"{_MODULE}.get_lora_configs", return_value=_TUNED_CONFIGS):
+                    # Exact match for chunk_size=32
+                    self.assertEqual(
+                        get_fn(K=1024, R=64, num_slices=1, chunk_size=32)["BLOCK_N"], 32
+                    )
+                    # Nearest neighbor: 100 is closer to 128
+                    self.assertEqual(
+                        get_fn(K=1024, R=64, num_slices=1, chunk_size=100)["BLOCK_N"],
+                        64,
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()