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] Enable moe LoRA kernel support FP16 (#27468)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Oct 27, 2025 at 19:48 UTC
f4e81540769848bbac8f22c356a121d3b3eb5038
2 files changed
+26
-16
tests/lora/test_fused_moe_lora_kernel.py
+17
-6
@@ -204,6 +204,11 @@ def use_torch(
204
return torch.stack(outputs, dim=0)
205
206
207
+DTYPES = [torch.float16, torch.bfloat16]
208
+DEVICES = [f"cuda:{0}"]
209
+SEED = [42]
210
+
211
+
212
@pytest.mark.parametrize("num_tokens", [100])
213
@pytest.mark.parametrize("top_k_num", [6, 12])
214
@pytest.mark.parametrize("num_experts", [64])
@@ -212,6 +217,9 @@ def use_torch(
217
@pytest.mark.parametrize("K", [2048])
218
@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
219
@pytest.mark.parametrize("block_size", [16])
220
+@pytest.mark.parametrize("dtype", DTYPES)
221
+@pytest.mark.parametrize("device", DEVICES)
222
+@pytest.mark.parametrize("seed", SEED)
223
def test_fused_moe_lora_kernel(
224
num_tokens,
225
top_k_num,
@@ -221,9 +229,12 @@ def test_fused_moe_lora_kernel(
229
K,
230
max_lora_rank,
231
block_size,
232
+ dtype,
233
+ device,
234
+ seed,
235
):
225
- torch.set_default_device("cuda:0")
226
- current_platform.seed_everything(42)
236
+ torch.set_default_device(device)
237
+ current_platform.seed_everything(seed)
238
# the number of randomly generated sentences.
239
num_sequences = 10
240
# generate data
@@ -240,7 +251,7 @@ def test_fused_moe_lora_kernel(
251
max_lora_rank,
252
K,
253
),
243
- dtype=torch.bfloat16,
254
+ dtype=dtype,
255
)
256
]
257
lora_b_stacked = [
@@ -251,7 +262,7 @@ def test_fused_moe_lora_kernel(
262
N,
263
max_lora_rank,
264
),
254
- dtype=torch.bfloat16,
265
+ dtype=dtype,
266
)
267
]
268
hidden_states = torch.rand(
@@ -259,11 +270,11 @@ def test_fused_moe_lora_kernel(
270
num_tokens,
271
K,
272
),
262
- dtype=torch.bfloat16,
273
+ dtype=dtype,
274
)
275
276
# fused_moe_lora_kernel output
266
- output = torch.zeros((num_tokens, top_k_num, N), dtype=torch.bfloat16)
277
+ output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
278
use_fused_moe_lora_kernel(
279
topk_ids,
280
topk_weights,
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+9
-10
@@ -2,9 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import torch
5
-import triton
6
-import triton.language as tl
5
6
+from vllm.triton_utils import tl, triton
7
from vllm.utils.torch_utils import direct_register_custom_op
8
9
_LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
@@ -110,7 +109,7 @@ def _fused_moe_lora_kernel(
109
110
# get a_ptr,b_ptr,c_ptr
111
cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size
113
- cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(tl.bfloat16))
112
+ cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))
113
cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size
114
115
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
@@ -154,7 +153,7 @@ def _fused_moe_lora_kernel(
153
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
154
accumulator = accumulator * moe_weight[:, None]
155
157
- accumulator = accumulator.to(tl.bfloat16)
156
+ accumulator = accumulator.to(c_ptr.dtype.element_ty)
157
# Write back the block of the output
158
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
159
c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
@@ -205,6 +204,10 @@ def _fused_moe_lora(
204
assert output.shape[0] == topk_weights.shape[0]
205
assert top_k_num == topk_weights.shape[1]
206
207
+ for lora_a, lora_b in zip(lora_a_stacked, lora_b_stacked):
208
+ assert lora_a.dtype == lora_b.dtype == output.dtype == qcurr_hidden_states.dtype
209
+ assert lora_a.dtype in [torch.float16, torch.bfloat16]
210
+
211
device = qcurr_hidden_states.device
212
num_slices = len(lora_a_stacked)
213
@@ -227,9 +230,9 @@ def _fused_moe_lora(
230
num_tokens = M * top_k_num
231
w1_output_dim_size = w1_lora_b_stacked.shape[2]
232
230
- lora_intermediate_cache1 = torch.zeros(
233
+ lora_intermediate_cache1 = torch.empty(
234
(num_slices * M * top_k_num * (max_lora_rank + w1_output_dim_size)),
232
- dtype=torch.bfloat16,
235
+ dtype=output.dtype,
236
device=device,
237
)
238
@@ -288,10 +291,6 @@ def _fused_moe_lora(
291
K = max_lora_rank
292
N = w1_output_dim_size
293
291
- # a_intermediate_cache1 = a_intermediate_cache1.view(
292
- # M, -1, a_intermediate_cache1.shape[3]
293
- # )
294
-
294
a_intermediate_cache1 = a_intermediate_cache1.view(
295
-1, a_intermediate_cache1.shape[3]
296
)