A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 20:45 UTC
[LoRA] Lora shrink swizzle (#27694)
Signed-off-by: li2haipeng <44383182+li2haipeng@users.noreply.github.com> Signed-off-by: Haipeng Li <li2haipeng@gmail.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
li2haipeng committed
Nov 3, 2025 at 17:30 UTC
6ddae74054d4d9b55b367bfc9db82969f9d81930
2 files changed
+14
-2
vllm/lora/ops/triton_ops/lora_shrink_op.py
+13
-2
@@ -41,6 +41,7 @@ def _lora_shrink_kernel(
41
BLOCK_K: tl.constexpr,
42
EVEN_K: tl.constexpr,
43
SPLIT_K: tl.constexpr,
44
+ GROUP_SIZE_M: tl.constexpr,
45
SLICE_NUM: tl.constexpr,
46
):
47
cta_n_num = tl.cdiv(N, BLOCK_N)
@@ -48,8 +49,16 @@ def _lora_shrink_kernel(
49
50
pid_sk_m_n = tl.program_id(axis=0)
51
pid_sk = pid_sk_m_n % SPLIT_K
51
- pid_m = (pid_sk_m_n // SPLIT_K) % cta_m_num
52
- pid_n = pid_sk_m_n // (SPLIT_K * cta_m_num) % cta_n_num
52
+
53
+ pid_m_n = pid_sk_m_n // SPLIT_K
54
+ num_pid_in_group = GROUP_SIZE_M * cta_n_num
55
+ group_id = pid_m_n // num_pid_in_group
56
+ first_pid_m = group_id * GROUP_SIZE_M
57
+ group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M)
58
+
59
+ # Column-major ordering within groups for better cache reuse
60
+ pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
61
+ pid_n = (pid_m_n % num_pid_in_group) // group_size_m
62
63
slice_id = tl.program_id(axis=1)
64
lora_idx = tl.program_id(axis=2)
@@ -194,6 +203,7 @@ def _lora_shrink(
203
NUM_WARPS = kernel_config["num_warps"]
204
NUM_STAGES = kernel_config["num_stages"]
205
NUM_CTAS = kernel_config["num_ctas"]
206
+ GROUP_SIZE_M = kernel_config.get("group_size_m", 8)
207
EVEN_K = K % (BLOCK_K * SPLIT_K) == 0 # type: ignore
208
209
# TODO (varun): This grid formulation maximizes parallelization at the
@@ -233,6 +243,7 @@ def _lora_shrink(
243
BLOCK_K,
244
EVEN_K,
245
SPLIT_K,
246
+ GROUP_SIZE_M,
247
NUM_SLICES,
248
num_warps=NUM_WARPS,
249
num_ctas=NUM_CTAS,
vllm/lora/ops/triton_ops/utils.py
+1
@@ -199,6 +199,7 @@ def get_lora_op_configs(
199
"split_k": 64 if batch < 128 else 8,
200
"num_warps": 4,
201
"num_ctas": 1,
202
+ "group_size_m": 8,
203
"num_stages": 2,
204
"max_nreg": None,
205
}