Add offline auto-tuning for LoRA CSGMV kernel (#20391)
Co-authored-by: Satyam Kumar <satyamk@linkedin.com>
This commit is contained in:
satyamk7054
co-authored by
Satyam Kumar
Satyam Kumar
parent
d8831355a3
commit
059b287e25
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+29
@@ -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
|
||||
}
|
||||
}
|
||||
+29
@@ -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
|
||||
}
|
||||
}
|
||||
+29
@@ -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
|
||||
}
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user