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 14, 05:30 UTC
[Kernel] Adding split_K implementation for fused_moe_lora (#27291)
Signed-off-by: Danielle Robinson <dmmaddix@amazon.com> Signed-off-by: Danielle Robinson <dcmaddix@gmail.com> Co-authored-by: Danielle Robinson <dmmaddix@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Danielle Robinson committed
Oct 27, 2025 at 02:05 UTC
9932ed6a8352fcc2230990c9631bcb6e77a93f70
4 files changed
+35
-14
tests/lora/test_fused_moe_lora_kernel.py
+2
@@ -154,6 +154,7 @@ def use_fused_moe_lora_kernel(
154
"BLOCK_SIZE_N": 32,
155
"BLOCK_SIZE_K": 64,
156
"GROUP_SIZE_M": 1,
157
+ "SPLIT_K": 1,
158
}
159
160
mul_routed_weight = False
@@ -175,6 +176,7 @@ def use_fused_moe_lora_kernel(
176
config["BLOCK_SIZE_N"],
177
config["BLOCK_SIZE_K"],
178
config["GROUP_SIZE_M"],
179
+ config["SPLIT_K"],
180
mul_routed_weight,
181
)
182
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+21
-9
@@ -80,11 +80,13 @@ def _fused_moe_lora_kernel(
80
BLOCK_SIZE_N: tl.constexpr,
81
BLOCK_SIZE_K: tl.constexpr,
82
GROUP_SIZE_M: tl.constexpr,
83
+ SPLIT_K: tl.constexpr,
84
):
85
pid = tl.program_id(axis=0)
86
slice_id = tl.program_id(axis=1)
87
lora_idx = tl.program_id(axis=2)
88
max_loras = tl.num_programs(axis=2)
89
+ grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
90
91
# calculate pid_m,pid_n
92
num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
@@ -102,7 +104,7 @@ def _fused_moe_lora_kernel(
104
105
# get the expert_id to process curr shard
106
ind = lora_idx * stride_el + pid_m
105
- expert_id = tl.load(expert_ids_ptr + ind)
107
+ expert_id = tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)
108
if expert_id == -1:
109
return
110
@@ -117,7 +119,7 @@ def _fused_moe_lora_kernel(
119
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
120
token_ind = stride_tl * lora_idx + offs_token_id
121
offs_token = tl.load(
120
- sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0.0
122
+ sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0
123
)
124
token_mask = offs_token < num_valid_tokens
125
@@ -135,17 +137,18 @@ def _fused_moe_lora_kernel(
137
138
# accumulator
139
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
138
- for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
140
+ for k in range(0, grid_k):
141
+ k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
142
a = tl.load(
143
a_ptrs,
141
- mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
144
+ mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
145
other=0.0,
146
)
144
- b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
147
+ b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
148
accumulator += tl.dot(a, b)
149
# Advance the ptrs to the next K block.
147
- a_ptrs += BLOCK_SIZE_K * stride_ak
148
- b_ptrs += BLOCK_SIZE_K * stride_bk
150
+ a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
151
+ b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
152
153
if MUL_ROUTED_WEIGHT:
154
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
@@ -156,7 +159,10 @@ def _fused_moe_lora_kernel(
159
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
160
c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
161
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
159
- tl.store(c_ptrs, accumulator, mask=c_mask)
162
+ if SPLIT_K == 1:
163
+ tl.store(c_ptrs, accumulator, mask=c_mask)
164
+ else:
165
+ tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
166
167
168
@torch.inference_mode()
@@ -179,6 +185,7 @@ def _fused_moe_lora(
185
block_size_n: int,
186
block_size_k: int,
187
group_size_m: int,
188
+ split_k: int,
189
mul_routed_weight: bool = False,
190
) -> None:
191
assert len(lora_a_stacked) == len(lora_b_stacked) > 0
@@ -206,6 +213,7 @@ def _fused_moe_lora(
213
"BLOCK_SIZE_N": block_size_n,
214
"BLOCK_SIZE_K": block_size_k,
215
"GROUP_SIZE_M": group_size_m,
216
+ "SPLIT_K": split_k,
217
}
218
219
w1_lora_a_stacked = lora_a_stacked[0]
@@ -237,7 +245,9 @@ def _fused_moe_lora(
245
b_ptr = _get_ptr(lora_a_stacked, device)
246
247
grid = lambda META: (
240
- triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
248
+ split_k
249
+ * triton.cdiv(EM, META["BLOCK_SIZE_M"])
250
+ * triton.cdiv(N, META["BLOCK_SIZE_N"]),
251
len(lora_a_stacked),
252
lora_a_stacked[0].shape[0],
253
)
@@ -286,6 +296,8 @@ def _fused_moe_lora(
296
-1, a_intermediate_cache1.shape[3]
297
)
298
299
+ # Set split_k = 1 for expand calls
300
+ config["SPLIT_K"] = 1
301
grid = lambda META: (
302
triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
303
len(lora_b_stacked),
vllm/lora/punica_wrapper/punica_gpu.py
+1
@@ -385,5 +385,6 @@ class PunicaWrapperGPU(PunicaWrapperBase):
385
config["BLOCK_SIZE_N"],
386
config["BLOCK_SIZE_K"],
387
config["GROUP_SIZE_M"],
388
+ config.get("SPLIT_K", 1),
389
mul_routed_weight,
390
)
vllm/model_executor/layers/fused_moe/fused_moe.py
+11
-5
@@ -121,6 +121,7 @@ def fused_moe_kernel_gptq_awq(
121
BLOCK_SIZE_N: tl.constexpr,
122
BLOCK_SIZE_K: tl.constexpr,
123
GROUP_SIZE_M: tl.constexpr,
124
+ SPLIT_K: tl.constexpr,
125
MUL_ROUTED_WEIGHT: tl.constexpr,
126
top_k: tl.constexpr,
127
compute_type: tl.constexpr,
@@ -356,6 +357,7 @@ def fused_moe_kernel(
357
BLOCK_SIZE_N: tl.constexpr,
358
BLOCK_SIZE_K: tl.constexpr,
359
GROUP_SIZE_M: tl.constexpr,
360
+ SPLIT_K: tl.constexpr,
361
MUL_ROUTED_WEIGHT: tl.constexpr,
362
top_k: tl.constexpr,
363
compute_type: tl.constexpr,
@@ -646,7 +648,6 @@ def invoke_fused_moe_kernel(
648
bit,
649
)
650
return
649
-
651
fused_moe_kernel_gptq_awq[grid](
652
A,
653
B,
@@ -686,6 +687,7 @@ def invoke_fused_moe_kernel(
687
)
688
else:
689
config = config.copy()
690
+ config["SPLIT_K"] = 1
691
BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
692
if block_shape is not None:
693
BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
@@ -983,6 +985,7 @@ def get_default_config(
985
"BLOCK_SIZE_N": 64,
986
"BLOCK_SIZE_K": 32,
987
"GROUP_SIZE_M": 8,
988
+ "SPLIT_K": 1,
989
}
990
return config
991
@@ -996,6 +999,7 @@ def get_default_config(
999
"BLOCK_SIZE_N": block_shape[0],
1000
"BLOCK_SIZE_K": block_shape[1],
1001
"GROUP_SIZE_M": 32,
1002
+ "SPLIT_K": 1,
1003
"num_warps": 4,
1004
"num_stages": 3 if not current_platform.is_rocm() else 2,
1005
}
@@ -1006,19 +1010,20 @@ def get_default_config(
1010
bit = 4 if dtype == "int4_w4a16" else 8
1011
use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit)
1012
if use_moe_wna16_cuda:
1009
- config = {"BLOCK_SIZE_M": min(16, M)}
1013
+ config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1}
1014
elif M <= 20:
1011
- config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1}
1015
+ config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
1016
elif M <= 40:
1013
- config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
1017
+ config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
1018
else:
1015
- config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
1019
+ config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
1020
elif M <= E:
1021
config = {
1022
"BLOCK_SIZE_M": 16,
1023
"BLOCK_SIZE_N": 32,
1024
"BLOCK_SIZE_K": 64,
1025
"GROUP_SIZE_M": 1,
1026
+ "SPLIT_K": 1,
1027
}
1028
else:
1029
config = {
@@ -1026,6 +1031,7 @@ def get_default_config(
1031
"BLOCK_SIZE_N": 64,
1032
"BLOCK_SIZE_K": 32,
1033
"GROUP_SIZE_M": 8,
1034
+ "SPLIT_K": 1,
1035
}
1036
return config
1037