@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7320ca3942
Fork 0

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 9, 03:45 UTC

Add unpermute-aware fused MoE LoRA path (#32655)

Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu>

Runkai Tao committed Feb 1, 2026 at 20:46 UTC 7320ca3942d008a59c24da3305c81810cb586c9d
6 files changed +474 -119
benchmarks/kernels/benchmark_lora.py
+2
@@ -842,6 +842,7 @@ class BenchmarkTensors:
842 "sorted_token_ids": sorted_token_ids,
843 "expert_ids": expert_ids,
844 "num_tokens_post_padded": num_tokens_post_padded,
845 + "token_lora_mapping": self.lora_kernel_meta.token_lora_mapping,
846 "top_k_num": ctx.top_k_num,
847 "device": self.input.device,
848 "N": lora_rank,
@@ -915,6 +916,7 @@ class BenchmarkTensors:
916 "sorted_token_ids": sorted_token_ids,
917 "expert_ids": expert_ids,
918 "num_tokens_post_padded": num_tokens_post_padded,
919 + "token_lora_mapping": self.lora_kernel_meta.token_lora_mapping,
920 "top_k_num": ctx.top_k_num,
921 "device": self.input.device,
922 "N": lora_rank,
tests/lora/test_fused_moe_lora_kernel.py
+184
@@ -190,6 +190,7 @@ def use_fused_moe_lora_kernel(
190 sorted_token_ids,
191 expert_ids,
192 num_tokens_post_padded,
193 + token_lora_mapping,
194 max_lora_rank,
195 top_k_num,
196 lora_ids,
@@ -333,6 +334,189 @@ def test_fused_moe_lora_kernel(
334 torch.testing.assert_close(output, output2, atol=1e-1, rtol=1e-1)
335
336
337 +def use_fused_moe_lora_kernel_naive(
338 + topk_ids,
339 + topk_weights,
340 + token_lora_mapping,
341 + max_lora_rank,
342 + top_k_num,
343 + lora_a_stacked,
344 + lora_b_stacked,
345 + hidden_states,
346 + output,
347 + max_loras,
348 + block_size,
349 + fully_sharded=False,
350 + offset=0,
351 +):
352 + """
353 + Test helper for naive_block_assignment path.
354 + Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids.
355 + """
356 + config = {
357 + "BLOCK_SIZE_M": block_size,
358 + "BLOCK_SIZE_N": 32,
359 + "BLOCK_SIZE_K": 64,
360 + "GROUP_SIZE_M": 1,
361 + "NUM_WARPS": 4,
362 + "NUM_STAGES": 3,
363 + "SPLIT_K": 1,
364 + }
365 +
366 + mul_routed_weight = False
367 +
368 + # In naive mode:
369 + # - expert_ids = topk_ids.view(-1), shape: (num_tokens * top_k,)
370 + # - sorted_token_ids = None
371 + # - num_tokens_post_padded = None
372 + expert_ids = topk_ids.reshape(-1)
373 + sorted_token_ids = None
374 + num_tokens_post_padded = None
375 +
376 + adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
377 + lora_ids = torch.arange(max_loras + 2, dtype=torch.int32)
378 +
379 + fused_moe_lora(
380 + output,
381 + hidden_states,
382 + lora_a_stacked,
383 + lora_b_stacked,
384 + topk_weights,
385 + sorted_token_ids,
386 + expert_ids,
387 + num_tokens_post_padded,
388 + token_lora_mapping,
389 + max_lora_rank,
390 + top_k_num,
391 + lora_ids,
392 + adapter_enabled,
393 + config["BLOCK_SIZE_M"],
394 + config["BLOCK_SIZE_N"],
395 + config["BLOCK_SIZE_K"],
396 + config["GROUP_SIZE_M"],
397 + config["NUM_WARPS"],
398 + config["NUM_STAGES"],
399 + config["SPLIT_K"],
400 + config["BLOCK_SIZE_M"],
401 + config["BLOCK_SIZE_N"],
402 + config["BLOCK_SIZE_K"],
403 + config["GROUP_SIZE_M"],
404 + config["NUM_WARPS"],
405 + config["NUM_STAGES"],
406 + config["SPLIT_K"],
407 + mul_routed_weight=mul_routed_weight,
408 + fully_sharded=fully_sharded,
409 + offset=offset,
410 + )
411 +
412 +
413 +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8])
414 +@pytest.mark.parametrize("top_k_num", [1, 2])
415 +@pytest.mark.parametrize("num_experts", [64, 128])
416 +@pytest.mark.parametrize("max_loras", [4, 8])
417 +@pytest.mark.parametrize("N", [1408])
418 +@pytest.mark.parametrize("K", [2048])
419 +@pytest.mark.parametrize("max_lora_rank", [16, 32])
420 +@pytest.mark.parametrize("block_size", [16])
421 +@pytest.mark.parametrize("dtype", DTYPES)
422 +@pytest.mark.parametrize("device", DEVICES)
423 +@pytest.mark.parametrize("seed", SEED)
424 +def test_fused_moe_lora_kernel_naive_block_assignment(
425 + num_tokens,
426 + top_k_num,
427 + num_experts,
428 + max_loras,
429 + N,
430 + K,
431 + max_lora_rank,
432 + block_size,
433 + dtype,
434 + device,
435 + seed,
436 +):
437 + """
438 + Test the naive_block_assignment path of the fused_moe_lora kernel.
439 + This path is triggered when batch_size * top_k is much smaller than
440 + num_experts * max_loras, and skips the moe_lora_align_block_size kernel.
441 + """
442 + torch.set_default_device(device)
443 + set_random_seed(seed)
444 +
445 + # Verify this configuration would trigger naive_block_assignment
446 + # (num_tokens * top_k * SPARSITY_FACTOR <= num_experts * max_loras)
447 + SPARSITY_FACTOR = 8
448 + assert num_tokens * top_k_num * SPARSITY_FACTOR <= num_experts * max_loras, (
449 + f"Test configuration doesn't meet naive_block_assignment condition: "
450 + f"{num_tokens} * {top_k_num} * {SPARSITY_FACTOR} > {num_experts} * {max_loras}"
451 + )
452 +
453 + # the number of randomly generated sentences.
454 + num_sequences = min(num_tokens, 4)
455 + # generate data
456 + topk_ids, topk_weights, token_lora_mapping = sample_data(
457 + num_tokens, num_sequences, max_loras, num_experts, top_k_num
458 + )
459 +
460 + # init lora weights
461 + lora_a_stacked = [
462 + torch.rand(
463 + (
464 + max_loras,
465 + num_experts,
466 + max_lora_rank,
467 + K,
468 + ),
469 + dtype=dtype,
470 + )
471 + ]
472 + lora_b_stacked = [
473 + torch.rand(
474 + (
475 + max_loras,
476 + num_experts,
477 + N,
478 + max_lora_rank,
479 + ),
480 + dtype=dtype,
481 + )
482 + ]
483 + hidden_states = torch.rand(
484 + (
485 + num_tokens,
486 + K,
487 + ),
488 + dtype=dtype,
489 + )
490 +
491 + # fused_moe_lora_kernel output (naive path)
492 + output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
493 + use_fused_moe_lora_kernel_naive(
494 + topk_ids,
495 + topk_weights,
496 + token_lora_mapping,
497 + max_lora_rank,
498 + top_k_num,
499 + lora_a_stacked,
500 + lora_b_stacked,
501 + hidden_states,
502 + output,
503 + max_loras,
504 + block_size,
505 + )
506 +
507 + # pytorch reference output
508 + output_ref = use_torch(
509 + hidden_states,
510 + token_lora_mapping,
511 + topk_ids,
512 + lora_a_stacked,
513 + lora_b_stacked,
514 + top_k_num,
515 + )
516 +
517 + torch.testing.assert_close(output, output_ref, atol=1e-1, rtol=1e-1)
518 +
519 +
520 @multi_gpu_test(num_gpus=2)
521 @pytest.mark.parametrize("num_tokens", [100])
522 @pytest.mark.parametrize("top_k_num", [6])
vllm/lora/layers/fused_moe.py
+25 -4
@@ -190,8 +190,18 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
190 config_dtype=config_dtype,
191 )
192
193 + # SPARSITY_FACTOR is a heuristic margin ensuring tokens * top_k
194 + # activates only a small fraction of total experts * loras.
195 + SPARSITY_FACTOR = 8
196 + naive_block_assignment = (
197 + expert_map is None
198 + and num_tokens * top_k * SPARSITY_FACTOR
199 + <= self.base_layer.local_num_experts * self.max_loras
200 + )
201 +
202 # get the block size of m from customized config or default config
203 (
204 + token_lora_mapping,
205 sorted_token_ids_lora,
206 expert_ids_lora,
207 num_tokens_post_padded_lora,
@@ -203,6 +213,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
213 self.max_loras,
214 self.adapter_enabled,
215 expert_map,
216 + naive_block_assignment,
217 )
218
219 moe_state_dict["sorted_token_ids_lora"] = sorted_token_ids_lora
@@ -210,9 +221,13 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
221 moe_state_dict["num_tokens_post_padded_lora"] = (
222 num_tokens_post_padded_lora
223 )
224 + moe_state_dict["token_lora_mapping"] = token_lora_mapping
225
214 - expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
215 - sorted_token_ids_lora = sorted_token_ids_lora.view(self.max_loras, -1)
226 + if sorted_token_ids_lora is not None:
227 + expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
228 + sorted_token_ids_lora = sorted_token_ids_lora.view(
229 + self.max_loras, -1
230 + )
231 #
232
233 self.punica_wrapper.add_lora_fused_moe(
@@ -230,6 +245,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
245 expand_config, ## pass the expand config
246 self.adapter_enabled,
247 fully_sharded=self.fully_sharded,
248 + token_lora_mapping=token_lora_mapping,
249 )
250
251 result = func(*args, **kwargs)
@@ -270,9 +286,13 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
286 num_tokens_post_padded_lora = moe_state_dict[
287 "num_tokens_post_padded_lora"
288 ]
289 + token_lora_mapping = moe_state_dict.get("token_lora_mapping")
290
274 - expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
275 - sorted_token_ids_lora = sorted_token_ids_lora.view(self.max_loras, -1)
291 + if sorted_token_ids_lora is not None:
292 + expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
293 + sorted_token_ids_lora = sorted_token_ids_lora.view(
294 + self.max_loras, -1
295 + )
296 intermediate_cache2 = moe_state_dict["intermediate_cache2"]
297 intermediate_cache3 = args[0]
298
@@ -295,6 +315,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
315 True,
316 fully_sharded=self.fully_sharded,
317 offset=shard_size_w2 * self.tp_rank if self.fully_sharded else 0,
318 + token_lora_mapping=token_lora_mapping,
319 )
320
321 result = func(*args, **kwargs)
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+201 -70
@@ -12,6 +12,64 @@ from vllm.utils.torch_utils import direct_register_custom_op
12
13 from .utils import supports_pdl
14
15 +
16 +@triton.jit
17 +def _get_lora_id(
18 + lora_ids,
19 + token_lora_mapping_ptr,
20 + lora_idx,
21 + pid_m,
22 + top_k_num,
23 + naive_block_assignment: tl.constexpr,
24 +):
25 + """Returns lora_id"""
26 + if naive_block_assignment:
27 + token_idx = pid_m // top_k_num
28 + return tl.load(token_lora_mapping_ptr + token_idx)
29 + else:
30 + return tl.load(lora_ids + lora_idx)
31 +
32 +
33 +@triton.jit
34 +def _get_expert_id(
35 + expert_ids_ptr,
36 + lora_id,
37 + pid_m,
38 + stride_el,
39 + max_loras,
40 + naive_block_assignment: tl.constexpr,
41 +):
42 + """Returns expert_id"""
43 + if naive_block_assignment:
44 + return tl.load(expert_ids_ptr + pid_m)
45 + else:
46 + ind = lora_id * stride_el + pid_m
47 + return tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)
48 +
49 +
50 +@triton.jit
51 +def _get_token_offs(
52 + sorted_token_ids_ptr,
53 + lora_id,
54 + pid_m,
55 + offs,
56 + stride_tl,
57 + max_loras,
58 + num_valid_tokens,
59 + naive_block_assignment: tl.constexpr,
60 + BLOCK_SIZE_M: tl.constexpr,
61 +):
62 + """Returns token offsets"""
63 + if naive_block_assignment:
64 + return tl.where(offs == 0, pid_m, num_valid_tokens)
65 + else:
66 + offs_token_id = pid_m * BLOCK_SIZE_M + offs
67 + token_ind = stride_tl * lora_id + offs_token_id
68 + return tl.load(
69 + sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0
70 + )
71 +
72 +
73 _LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
74
75
@@ -36,6 +94,25 @@ def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
94 return _LORA_PTR_DICT.get(key)
95
96
97 +def _adjust_kernel_inputs(
98 + max_loras: int,
99 + sorted_token_ids: torch.Tensor | None,
100 + expert_ids: torch.Tensor,
101 +):
102 + """
103 + helper function to adjust kernel inputs when sorted_token_ids is None
104 + """
105 + if sorted_token_ids is None:
106 + stride_tl = 0
107 + stride_el = 0
108 + grid_lora_dim = 1
109 + else:
110 + stride_tl = sorted_token_ids.stride(0)
111 + stride_el = expert_ids.stride(0)
112 + grid_lora_dim = max_loras + 1
113 + return grid_lora_dim, stride_tl, stride_el
114 +
115 +
116 @triton.jit(
117 do_not_specialize=[
118 "num_valid_tokens",
@@ -54,12 +131,14 @@ def _fused_moe_lora_kernel(
131 sorted_token_ids_ptr,
132 expert_ids_ptr,
133 num_tokens_post_padded_ptr,
134 + token_lora_mapping_ptr,
135 # Matrix dimensions
136 N,
137 K,
138 EM,
139 num_valid_tokens,
140 num_experts,
141 + top_k_num,
142 lora_ids,
143 adapter_enabled,
144 max_loras, # <<< PR2: rename, used for masks when grid axis-2 != max_loras
@@ -82,7 +161,11 @@ def _fused_moe_lora_kernel(
161 # Meta-parameters
162 num_slice_a: tl.constexpr,
163 num_slice_c: tl.constexpr,
85 - top_k: tl.constexpr,
164 + # top_k_num or 1 depending on input token
165 + # is expanded by top_k or not
166 + token_mapping_factor: tl.constexpr,
167 + # whether use naive block assignment
168 + naive_block_assignment: tl.constexpr,
169 MUL_ROUTED_WEIGHT: tl.constexpr,
170 ADD_INPUTS: tl.constexpr,
171 USE_B_L2_CACHE: tl.constexpr, # new, enable .ca load for B
@@ -97,26 +180,10 @@ def _fused_moe_lora_kernel(
180 ):
181 pid = tl.program_id(axis=0)
182 slice_id = tl.program_id(axis=1)
100 - lora_idx = tl.program_id(axis=2)
101 - lora_id = tl.load(lora_ids + lora_idx)
102 -
103 - if lora_id == -1:
104 - # Early exit for the no-lora case.
105 - return
106 - moe_enabled = tl.load(adapter_enabled + lora_id)
107 - if moe_enabled == 0:
108 - # Early exit for the no moe lora case.
109 - return
110 - # The grid's axis-2 dimension is max_loras + 1 to accommodate the -1 sentinel.
111 - # This guard ensures we don't access sorted_token_ids / expert_ids /
112 - # num_tokens_post_padded beyond their allocated bounds if an invalid
113 - # lora_id somehow appears. Although the caller should pass correct
114 - # max_loras, defensive programming prevents accidental out-of-bounds.
115 - if lora_id >= max_loras:
116 - return
183 grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
184
185 # calculate pid_m,pid_n
186 + lora_idx = tl.program_id(axis=2)
187 pid_sk = pid % SPLIT_K
188 pid_m_n = pid // SPLIT_K
189 num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
@@ -129,14 +196,55 @@ def _fused_moe_lora_kernel(
196 pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
197 pid_n = (pid_m_n % num_pid_in_group) // group_size_m
198
132 - num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)
133 - if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
199 + offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
200 +
201 + # Get lora_id
202 + lora_id = _get_lora_id(
203 + lora_ids,
204 + token_lora_mapping_ptr,
205 + lora_idx,
206 + pid_m,
207 + top_k_num,
208 + naive_block_assignment,
209 + )
210 + if lora_id == -1:
211 + return
212 + moe_enabled = tl.load(adapter_enabled + lora_id)
213 + if moe_enabled == 0:
214 + return
215 + if lora_id >= max_loras:
216 return
135 - # get the expert_id to process curr shard
136 - ind = lora_id * stride_el + pid_m
137 - expert_id = tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)
217 +
218 + # Non-naive only: check num_tokens_post_padded
219 + if not naive_block_assignment:
220 + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)
221 + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
222 + return
223 +
224 + # Get expert_id
225 + expert_id = _get_expert_id(
226 + expert_ids_ptr,
227 + lora_id,
228 + pid_m,
229 + stride_el,
230 + max_loras,
231 + naive_block_assignment,
232 + )
233 if expert_id == -1:
234 return
235 +
236 + # Get token offsets
237 + offs_token = _get_token_offs(
238 + sorted_token_ids_ptr,
239 + lora_id,
240 + pid_m,
241 + offs,
242 + stride_tl,
243 + max_loras,
244 + num_valid_tokens,
245 + naive_block_assignment,
246 + BLOCK_SIZE_M,
247 + )
248 # get a_ptr,b_ptr,c_ptr
249 cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size
250 cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))
@@ -145,19 +253,12 @@ def _fused_moe_lora_kernel(
253 # remove modulo wrap-around
254 offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int32)
255 offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
148 -
149 - offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int32)
150 - token_ind = stride_tl * lora_id + offs_token_id
151 - offs_token = tl.load(
152 - sorted_token_ids_ptr + token_ind,
153 - mask=token_ind < max_loras * stride_tl,
154 - other=num_valid_tokens,
155 - )
256 token_mask = offs_token < num_valid_tokens
257
258 # get a_ptrs,b_ptrs
259 a_ptrs = cur_a_ptr + (
160 - offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
260 + offs_token[:, None] // token_mapping_factor * stride_am
261 + + offs_k[None, :] * stride_ak
262 )
263
264 b_ptrs = (
@@ -230,9 +331,10 @@ def _fused_moe_lora_shrink(
331 torch.Tensor
332 ], # [(max_loras, num_experts, max_lora_rank, K,),...]
333 topk_weights: torch.Tensor, # (num_tokens, top_k_num)
233 - sorted_token_ids: torch.Tensor, # (max_loras, _)
234 - expert_ids: torch.Tensor, # (max_loras, _ ,)
235 - num_tokens_post_padded: torch.Tensor, # (max_loras, )
334 + sorted_token_ids: torch.Tensor | None, # (max_loras, _)
335 + expert_ids: torch.Tensor, # (max_loras, _ ,) or (num_tokens * top_k,)
336 + num_tokens_post_padded: torch.Tensor | None, # (max_loras, )
337 + token_lora_mapping: torch.Tensor,
338 top_k_num: int,
339 lora_ids: torch.Tensor,
340 adapter_enabled: torch.Tensor,
@@ -270,13 +372,15 @@ def _fused_moe_lora_shrink(
372
373 b_ptr = _get_ptr(lora_a_stacked, device)
374
375 + grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(
376 + w1_lora_a_stacked.shape[0], sorted_token_ids, expert_ids
377 + )
378 grid = lambda META: (
379 split_k
380 * triton.cdiv(EM, META["BLOCK_SIZE_M"])
381 * triton.cdiv(N, META["BLOCK_SIZE_N"]),
382 len(lora_a_stacked),
278 - ## max_loras + 1 to handle the no-lora case (lora_id == -1)
279 - lora_a_stacked[0].shape[0] + 1,
383 + grid_lora_dim,
384 )
385 _fused_moe_lora_kernel[grid](
386 qcurr_hidden_states,
@@ -286,11 +390,13 @@ def _fused_moe_lora_shrink(
390 sorted_token_ids,
391 expert_ids,
392 num_tokens_post_padded,
393 + token_lora_mapping,
394 N,
395 K,
396 EM,
397 num_tokens,
398 num_experts,
399 + top_k_num,
400 lora_ids,
401 adapter_enabled,
402 lora_a_stacked[0].shape[0],
@@ -302,13 +408,14 @@ def _fused_moe_lora_shrink(
408 w1_lora_a_stacked.stride(2),
409 a_intermediate_cache1.stride(2),
410 a_intermediate_cache1.stride(3),
305 - sorted_token_ids.stride(0),
306 - expert_ids.stride(0),
411 + stride_tl,
412 + stride_el,
413 slice_a_size=qcurr_hidden_states.numel(),
414 slice_c_size=a_intermediate_cache1.numel() // num_slices,
415 num_slice_a=1,
416 num_slice_c=num_slices,
311 - top_k=1 if mul_routed_weight else top_k_num,
417 + token_mapping_factor=1 if mul_routed_weight else top_k_num,
418 + naive_block_assignment=sorted_token_ids is None,
419 MUL_ROUTED_WEIGHT=False,
420 ADD_INPUTS=False,
421 USE_B_L2_CACHE=True, # new
@@ -325,9 +432,10 @@ def _fused_moe_lora_expand(
432 torch.Tensor
433 ], # [(max_loras, num_experts, max_lora_rank, K,),...]
434 topk_weights: torch.Tensor, # (num_tokens, top_k_num)
328 - sorted_token_ids: torch.Tensor, # (max_loras, _)
329 - expert_ids: torch.Tensor, # (max_loras, _ ,)
330 - num_tokens_post_padded: torch.Tensor, # (max_loras, )
435 + sorted_token_ids: torch.Tensor | None, # (max_loras, _)
436 + expert_ids: torch.Tensor, # (max_loras, _ ,) or (num_tokens * top_k,)
437 + num_tokens_post_padded: torch.Tensor | None, # (max_loras, )
438 + token_lora_mapping: torch.Tensor,
439 top_k_num: int,
440 lora_ids: torch.Tensor,
441 adapter_enabled: torch.Tensor,
@@ -375,11 +483,14 @@ def _fused_moe_lora_expand(
483 "launch_pdl": use_gdc, # triton kernel metadata
484 }
485
486 + grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(
487 + w1_lora_b_stacked.shape[0], sorted_token_ids, expert_ids
488 + )
489 +
490 grid = lambda META: (
491 triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
492 len(lora_b_stacked),
381 - ## max_loras + 1 to handle the no-lora case (lora_id == -1)
382 - lora_b_stacked[0].shape[0] + 1,
493 + grid_lora_dim,
494 )
495
496 # Fast path: directly accumulate into the corresponding slice interval of output.
@@ -394,11 +505,13 @@ def _fused_moe_lora_expand(
505 sorted_token_ids,
506 expert_ids,
507 num_tokens_post_padded,
508 + token_lora_mapping,
509 N,
510 K,
511 EM,
512 num_tokens,
513 num_experts,
514 + top_k_num,
515 lora_ids,
516 adapter_enabled,
517 lora_b_stacked[0].shape[0],
@@ -410,13 +523,14 @@ def _fused_moe_lora_expand(
523 w1_lora_b_stacked.stride(2),
524 out_view.stride(1),
525 out_view.stride(2),
413 - sorted_token_ids.stride(0),
414 - expert_ids.stride(0),
526 + stride_tl,
527 + stride_el,
528 slice_a_size=a_intermediate_cache1.numel() // num_slices,
529 slice_c_size=slice_c_size,
530 num_slice_a=num_slices,
531 num_slice_c=num_slices,
419 - top_k=1,
532 + token_mapping_factor=1,
533 + naive_block_assignment=sorted_token_ids is None,
534 MUL_ROUTED_WEIGHT=mul_routed_weight,
535 ADD_INPUTS=True,
536 USE_B_L2_CACHE=True, # new
@@ -436,9 +550,10 @@ def _fused_moe_lora(
550 torch.Tensor
551 ], # [(max_loras, num_experts, N, max_lora_rank,),...]
552 topk_weights: torch.Tensor, # (num_tokens, top_k_num)
439 - sorted_token_ids: torch.Tensor, # (max_loras, _)
440 - expert_ids: torch.Tensor, # (max_loras, _ ,)
441 - num_tokens_post_padded: torch.Tensor, # (max_loras, )
553 + sorted_token_ids: torch.Tensor | None, # (max_loras, _)
554 + expert_ids: torch.Tensor, # (max_loras, _ ,) or (num_tokens * top_k,)
555 + num_tokens_post_padded: torch.Tensor | None, # (max_loras, )
556 + token_lora_mapping: torch.Tensor,
557 max_lora_rank: int,
558 top_k_num: int,
559 lora_ids: torch.Tensor,
@@ -462,18 +577,24 @@ def _fused_moe_lora(
577 offset: int = 0,
578 ) -> None:
579 assert len(lora_a_stacked) == len(lora_b_stacked) > 0
465 - assert (
466 - sorted_token_ids.dim()
467 - == expert_ids.dim()
468 - == topk_weights.dim()
469 - == qcurr_hidden_states.dim()
470 - == 2
471 - )
472 - assert (
473 - sorted_token_ids.shape[0]
474 - == expert_ids.shape[0]
475 - == num_tokens_post_padded.shape[0]
476 - )
580 + assert topk_weights.dim() == qcurr_hidden_states.dim() == 2
581 + if sorted_token_ids is None:
582 + assert expert_ids.dim() == 1
583 + else:
584 + assert sorted_token_ids is not None
585 + assert num_tokens_post_padded is not None
586 + assert (
587 + sorted_token_ids.dim()
588 + == expert_ids.dim()
589 + == topk_weights.dim()
590 + == qcurr_hidden_states.dim()
591 + == 2
592 + )
593 + assert (
594 + sorted_token_ids.shape[0]
595 + == expert_ids.shape[0]
596 + == num_tokens_post_padded.shape[0]
597 + )
598 assert output.shape[0] == topk_weights.shape[0]
599 assert top_k_num == topk_weights.shape[1]
600 device = qcurr_hidden_states.device
@@ -482,10 +603,15 @@ def _fused_moe_lora(
603 num_experts = lora_a_stacked[0].shape[1]
604 N = max_lora_rank
605 M = topk_weights.shape[0]
485 - EM = sorted_token_ids.shape[1]
606 K = qcurr_hidden_states.shape[1]
607 num_tokens = M * top_k_num
608 w1_output_dim_size = w1_lora_b_stacked.shape[2]
609 + assert shrink_block_size_m == expand_block_size_m
610 + EM = (
611 + sorted_token_ids.shape[1]
612 + if sorted_token_ids is not None
613 + else num_tokens * shrink_block_size_m
614 + )
615
616 a_intermediate_cache1 = torch.zeros(
617 (num_slices, M, top_k_num, max_lora_rank),
@@ -502,6 +628,7 @@ def _fused_moe_lora(
628 sorted_token_ids,
629 expert_ids,
630 num_tokens_post_padded,
631 + token_lora_mapping,
632 top_k_num,
633 lora_ids,
634 adapter_enabled,
@@ -546,6 +673,7 @@ def _fused_moe_lora(
673 sorted_token_ids,
674 expert_ids,
675 num_tokens_post_padded,
676 + token_lora_mapping,
677 top_k_num,
678 lora_ids,
679 adapter_enabled,
@@ -579,9 +707,10 @@ def _fused_moe_lora_fake(
707 lora_a_stacked: list[torch.Tensor],
708 lora_b_stacked: list[torch.Tensor],
709 topk_weights: torch.Tensor,
582 - sorted_token_ids: torch.Tensor,
710 + sorted_token_ids: torch.Tensor | None,
711 expert_ids: torch.Tensor,
584 - num_tokens_post_padded: torch.Tensor,
712 + num_tokens_post_padded: torch.Tensor | None,
713 + token_lora_mapping: torch.Tensor,
714 max_lora_rank: int,
715 top_k_num: int,
716 lora_ids: torch.Tensor,
@@ -610,9 +739,10 @@ def _fused_moe_lora_shrink_fake(
739 qcurr_hidden_states: torch.Tensor,
740 lora_a_stacked: list[torch.Tensor],
741 topk_weights: torch.Tensor,
613 - sorted_token_ids: torch.Tensor,
742 + sorted_token_ids: torch.Tensor | None,
743 expert_ids: torch.Tensor,
615 - num_tokens_post_padded: torch.Tensor,
744 + num_tokens_post_padded: torch.Tensor | None,
745 + token_lora_mapping: torch.Tensor,
746 top_k_num: int,
747 lora_ids: torch.Tensor,
748 adapter_enabled: torch.Tensor,
@@ -642,9 +772,10 @@ def _fused_moe_lora_expand_fake(
772 a_intermediate_cache1: torch.Tensor,
773 lora_b_stacked: list[torch.Tensor],
774 topk_weights: torch.Tensor,
645 - sorted_token_ids: torch.Tensor,
775 + sorted_token_ids: torch.Tensor | None,
776 expert_ids: torch.Tensor,
647 - num_tokens_post_padded: torch.Tensor,
777 + num_tokens_post_padded: torch.Tensor | None,
778 + token_lora_mapping: torch.Tensor,
779 top_k_num: int,
780 lora_ids: torch.Tensor,
781 adapter_enabled: torch.Tensor,
vllm/lora/punica_wrapper/punica_base.py
+4 -3
@@ -458,7 +458,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
458 adapter_enabled: torch.Tensor,
459 expert_map: torch.Tensor | None = None,
460 pad_sorted_ids: bool = False,
461 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
461 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
462 """
463 Aligns tokens and experts into block-sized chunks for LoRA-based
464 mixture-of-experts (MoE) execution.
@@ -473,9 +473,9 @@ class PunicaWrapperBase(PunicaWrapperABC):
473 lora_a_stacked: tuple[torch.Tensor, ...],
474 lora_b_stacked: tuple[torch.Tensor, ...],
475 topk_weights: torch.Tensor,
476 - sorted_token_ids: torch.Tensor,
476 + sorted_token_ids: torch.Tensor | None,
477 expert_ids: torch.Tensor,
478 - num_tokens_post_padded: torch.Tensor,
478 + num_tokens_post_padded: torch.Tensor | None,
479 max_lora_rank: int,
480 top_k_num: int,
481 shrink_config,
@@ -484,6 +484,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
484 mul_routed_weight=False,
485 fully_sharded: bool = False,
486 offset: int = 0,
487 + token_lora_mapping: torch.Tensor | None = None,
488 ):
489 """
490 Performs a fused forward computation for LoRA of
vllm/lora/punica_wrapper/punica_gpu.py
+58 -42
@@ -310,52 +310,57 @@ class PunicaWrapperGPU(PunicaWrapperBase):
310 adapter_enabled: torch.Tensor,
311 expert_map: torch.Tensor | None = None,
312 pad_sorted_ids: bool = False,
313 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
313 + naive_block_assignment: bool = False,
314 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
315 """
316 Aligns tokens and experts into block-sized chunks for LoRA-based
317 mixture-of-experts (MoE) execution.
318 """
318 - max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
319 - if pad_sorted_ids:
320 - max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
321 - sorted_ids = torch.empty(
322 - (max_loras * max_num_tokens_padded,),
323 - dtype=torch.int32,
324 - device=topk_ids.device,
325 - )
326 - max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
327 - # Expert ids must be set default to -1 to prevent a blank block
328 - expert_ids = torch.empty(
329 - (max_loras * max_num_m_blocks,),
330 - dtype=torch.int32,
331 - device=topk_ids.device,
332 - )
333 - num_tokens_post_pad = torch.empty(
334 - (max_loras), dtype=torch.int32, device=topk_ids.device
335 - )
336 -
319 (token_lora_mapping, _, _, _, lora_ids, _) = self.token_mapping_meta.meta_args(
320 num_tokens
321 )
340 -
341 - ops.moe_lora_align_block_size(
342 - topk_ids,
343 - token_lora_mapping,
344 - num_experts,
345 - block_size,
346 - max_loras,
347 - max_num_tokens_padded,
348 - max_num_m_blocks,
349 - sorted_ids,
350 - expert_ids,
351 - num_tokens_post_pad,
352 - adapter_enabled,
353 - lora_ids,
354 - )
355 - if expert_map is not None:
356 - expert_ids = expert_map[expert_ids]
357 -
358 - return sorted_ids, expert_ids, num_tokens_post_pad
322 + if naive_block_assignment:
323 + expert_ids = topk_ids.reshape(-1)
324 + sorted_ids = None
325 + num_tokens_post_pad = None
326 + else:
327 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
328 + if pad_sorted_ids:
329 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
330 + sorted_ids = torch.empty(
331 + (max_loras * max_num_tokens_padded,),
332 + dtype=torch.int32,
333 + device=topk_ids.device,
334 + )
335 + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
336 + # Expert ids must be set default to -1 to prevent a blank block
337 + expert_ids = torch.empty(
338 + (max_loras * max_num_m_blocks,),
339 + dtype=torch.int32,
340 + device=topk_ids.device,
341 + )
342 + num_tokens_post_pad = torch.empty(
343 + (max_loras), dtype=torch.int32, device=topk_ids.device
344 + )
345 +
346 + ops.moe_lora_align_block_size(
347 + topk_ids,
348 + token_lora_mapping,
349 + num_experts,
350 + block_size,
351 + max_loras,
352 + max_num_tokens_padded,
353 + max_num_m_blocks,
354 + sorted_ids,
355 + expert_ids,
356 + num_tokens_post_pad,
357 + adapter_enabled,
358 + lora_ids,
359 + )
360 + if expert_map is not None:
361 + expert_ids = expert_map[expert_ids]
362 +
363 + return None, sorted_ids, expert_ids, num_tokens_post_pad
364
365 def add_lora_fused_moe(
366 self,
@@ -364,9 +369,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):
369 lora_a_stacked: tuple[torch.Tensor, ...],
370 lora_b_stacked: tuple[torch.Tensor, ...],
371 topk_weights: torch.Tensor,
367 - sorted_token_ids: torch.Tensor,
372 + sorted_token_ids: torch.Tensor | None,
373 expert_ids: torch.Tensor,
369 - num_tokens_post_padded: torch.Tensor,
374 + num_tokens_post_padded: torch.Tensor | None,
375 max_lora_rank: int,
376 top_k_num: int,
377 shrink_config,
@@ -375,11 +380,21 @@ class PunicaWrapperGPU(PunicaWrapperBase):
380 mul_routed_weight=False,
381 fully_sharded: bool = False,
382 offset: int = 0,
383 + token_lora_mapping: torch.Tensor | None = None,
384 ):
385 """
386 Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
387 """
382 - (_, _, _, _, lora_ids, _) = self.token_mapping_meta.meta_args(x.size(0))
388 + (
389 + token_lora_mapping_meta,
390 + _,
391 + _,
392 + _,
393 + lora_ids,
394 + _,
395 + ) = self.token_mapping_meta.meta_args(x.size(0))
396 + if token_lora_mapping is None:
397 + token_lora_mapping = token_lora_mapping_meta
398 fused_moe_lora(
399 y,
400 x,
@@ -389,6 +404,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
404 sorted_token_ids,
405 expert_ids,
406 num_tokens_post_padded,
407 + token_lora_mapping,
408 max_lora_rank,
409 top_k_num,
410 lora_ids,