diff --git a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py new file mode 100755 index 000000000..1c162beca --- /dev/null +++ b/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) diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py b/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py index a24abbe37..52e3d6038 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +++ b/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 diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py b/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py index b0ffdb763..e625588aa 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py +++ b/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 diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=1024,R=64,S=1,device=NVIDIA_H200.json new file mode 100644 index 000000000..b019801e4 --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=4096,R=64,S=3,device=NVIDIA_H200.json new file mode 100644 index 000000000..c2dc948e0 --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_expand,K=6144,R=64,S=2,device=NVIDIA_H200.json new file mode 100644 index 000000000..930bb3214 --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=2,device=NVIDIA_H200.json new file mode 100644 index 000000000..b7054db7e --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=1024,R=64,S=3,device=NVIDIA_H200.json new file mode 100644 index 000000000..f1eb38c7f --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=2048,R=64,S=1,device=NVIDIA_H200.json new file mode 100644 index 000000000..4741c164f --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json b/python/sglang/srt/lora/triton_ops/csgmv_configs/triton_3_5_1/lora_shrink,K=3072,R=64,S=1,device=NVIDIA_H200.json new file mode 100644 index 000000000..1715cfe9a --- /dev/null +++ b/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 + } +} diff --git a/python/sglang/srt/lora/triton_ops/lora_tuning_config.py b/python/sglang/srt/lora/triton_ops/lora_tuning_config.py new file mode 100644 index 000000000..33e9e72ed --- /dev/null +++ b/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) diff --git a/test/manual/lora/test_lora_tuning_config.py b/test/manual/lora/test_lora_tuning_config.py new file mode 100644 index 000000000..6601a7bde --- /dev/null +++ b/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()