@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2ec401bc39
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 13, 20:45 UTC

Load tuned fused_moe_lora shrink and expand kernel configs separately (#27435)

Signed-off-by: Yu Gong <yu3.gong@gmail.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

yugong333 committed Nov 4, 2025 at 02:27 UTC 2ec401bc39daf0c8daa7f7c6bffe4f5e15cb7c79
9 files changed +910 -124
benchmarks/kernels/benchmark_lora.py
+446 -30
@@ -19,13 +19,24 @@ from torch.utils.benchmark import Measurement as TMeasurement
19 from utils import ArgPool, Bench, CudaGraphBenchParams
20 from weight_shapes import WEIGHT_SHAPES
21
22 -from vllm.triton_utils import HAS_TRITON
22 +from vllm.lora.ops.triton_ops.utils import get_lora_op_configs
23 +from vllm.triton_utils import HAS_TRITON, triton
24
25 if HAS_TRITON:
25 - from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_expand, lora_shrink
26 + from vllm.lora.ops.triton_ops import ( ## added fused_moe_lora
27 + LoRAKernelMeta,
28 + fused_moe_lora_expand,
29 + fused_moe_lora_shrink,
30 + lora_expand,
31 + lora_shrink,
32 + )
33 + from vllm.lora.ops.triton_ops.fused_moe_lora_op import (
34 + _LORA_PTR_DICT, ## added _LORA_PTR_DICT for fused_moe_lora
35 + )
36 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
27 -
37 +from vllm import _custom_ops as ops
38 from vllm.utils.argparse_utils import FlexibleArgumentParser
39 +from vllm.utils.math_utils import round_up
40
41 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
42 DEFAULT_TP_SIZES = [1]
@@ -59,6 +70,8 @@ DEFAULT_NUM_LORAS = [1, 2, 3, 4]
70 DEFAULT_SORT_BY_LORA_IDS = [False, True]
71 DEFAULT_SEQ_LENGTHS = [1]
72 DEFAULT_EXPAND_FN_ADD_INPUTS = [True, False]
73 +DEFAULT_TOP_K_NUMS = [1] # Added for MoE LoRA top_k
74 +DEFAULT_NUM_EXPERTS = [8] # Added for MoE LoRA num_experts
75
76
77 # Utilities
@@ -191,6 +204,11 @@ class OpType(Enum):
204
205 LORA_SHRINK = auto()
206 LORA_EXPAND = auto()
207 + ## Adding support for fused moe lora
208 + FUSED_MOE_LORA_GATE_UP_SHRINK = auto() ## Gate/Up projection variant with shrink
209 + FUSED_MOE_LORA_GATE_UP_EXPAND = auto() ## Gate/Up projection variant with expand
210 + FUSED_MOE_LORA_DOWN_SHRINK = auto() ## Down projection variant with shrink
211 + FUSED_MOE_LORA_DOWN_EXPAND = auto() ## Down projection variant with expand
212
213 @staticmethod
214 def from_str(s: str) -> "OpType":
@@ -198,6 +216,15 @@ class OpType(Enum):
216 return OpType.LORA_SHRINK
217 if s.lower() == "lora_expand":
218 return OpType.LORA_EXPAND
219 + # Adding support for fused moe lora, both in gate_up and down
220 + if s.lower() == "fused_moe_lora_gate_up_shrink": ## Gate/Up variant with shrink
221 + return OpType.FUSED_MOE_LORA_GATE_UP_SHRINK
222 + if s.lower() == "fused_moe_lora_gate_up_expand": ## Gate/Up variant with expand
223 + return OpType.FUSED_MOE_LORA_GATE_UP_EXPAND
224 + if s.lower() == "fused_moe_lora_down_shrink": ## Down variant with shrink
225 + return OpType.FUSED_MOE_LORA_DOWN_SHRINK
226 + if s.lower() == "fused_moe_lora_down_expand": ## Down variant with expand
227 + return OpType.FUSED_MOE_LORA_DOWN_EXPAND
228 raise ValueError(f"Unrecognized str {s} to convert to OpType")
229
230 def is_shrink_fn(self) -> bool:
@@ -206,19 +233,56 @@ class OpType(Enum):
233 def is_expand_fn(self) -> bool:
234 return self in [OpType.LORA_EXPAND]
235
236 + def is_fused_moe_lora_fn(self) -> bool: ## adding for fused MoE LoRA
237 + return self in [
238 + OpType.FUSED_MOE_LORA_GATE_UP_SHRINK,
239 + OpType.FUSED_MOE_LORA_DOWN_SHRINK,
240 + OpType.FUSED_MOE_LORA_GATE_UP_EXPAND,
241 + OpType.FUSED_MOE_LORA_DOWN_EXPAND,
242 + ]
243 +
244 + def is_fused_moe_lora_gate_up_fn(
245 + self,
246 + ) -> bool: ## adding for fused MoE LoRA Gate/Up
247 + return self in [
248 + OpType.FUSED_MOE_LORA_GATE_UP_SHRINK,
249 + OpType.FUSED_MOE_LORA_GATE_UP_EXPAND,
250 + ]
251 +
252 + def is_fused_moe_lora_down_fn(self) -> bool: ## adding for fused MoE LoRA Down
253 + return self in [
254 + OpType.FUSED_MOE_LORA_DOWN_SHRINK,
255 + OpType.FUSED_MOE_LORA_DOWN_EXPAND,
256 + ]
257 +
258 + def is_fused_moe_lora_shrink_fn(self) -> bool:
259 + return self in [
260 + OpType.FUSED_MOE_LORA_GATE_UP_SHRINK,
261 + OpType.FUSED_MOE_LORA_DOWN_SHRINK,
262 + ]
263 +
264 + def is_fused_moe_lora_expand_fn(self) -> bool:
265 + return self in [
266 + OpType.FUSED_MOE_LORA_GATE_UP_EXPAND,
267 + OpType.FUSED_MOE_LORA_DOWN_EXPAND,
268 + ]
269 +
270 def num_slices(self) -> list[int]:
271 + if self.is_fused_moe_lora_gate_up_fn():
272 + return [2]
273 + elif self.is_fused_moe_lora_down_fn():
274 + return [1]
275 return [1, 2, 3]
276
277 def mkn(
278 self, batch_size: int, seq_length: int, hidden_size: int, lora_rank: int
279 ) -> tuple[int, int, int]:
280 num_tokens = batch_size * seq_length
216 - if self.is_shrink_fn():
281 + if self.is_shrink_fn() or self.is_fused_moe_lora_fn():
282 m = num_tokens
283 k = hidden_size
284 n = lora_rank
220 - else:
221 - assert self.is_expand_fn()
285 + elif self.is_expand_fn():
286 m = num_tokens
287 k = lora_rank
288 n = hidden_size
@@ -232,9 +296,36 @@ class OpType(Enum):
296 """
297 if self.is_shrink_fn():
298 return op_dtype, op_dtype, torch.float32
235 - else:
236 - assert self.is_expand_fn()
299 + elif self.is_expand_fn():
300 return torch.float32, op_dtype, op_dtype
301 + else:
302 + assert self.is_fused_moe_lora_fn()
303 + return op_dtype, op_dtype, op_dtype
304 +
305 + def matmul_shapes_fused_moe_lora(
306 + self,
307 + m: int,
308 + n: int,
309 + k: int,
310 + num_loras: int,
311 + num_slices: int,
312 + top_k_num: int,
313 + num_experts: int,
314 + ) -> tuple[tuple[int], tuple[int], tuple[int], tuple[int]]:
315 + if self.is_fused_moe_lora_shrink_fn():
316 + input_shape = (
317 + (m * top_k_num, n)
318 + if self in [OpType.FUSED_MOE_LORA_DOWN_SHRINK]
319 + else (m, n)
320 + )
321 + output_shape = (num_slices, m, top_k_num, k)
322 + weight_shape = (num_loras, num_experts, k, n)
323 + else:
324 + assert self.is_fused_moe_lora_expand_fn()
325 + input_shape = (num_slices, m, top_k_num, k)
326 + output_shape = (m, top_k_num, n * num_slices)
327 + weight_shape = (num_loras, num_experts, n, k)
328 + return (input_shape, weight_shape, output_shape)
329
330 def matmul_shapes(
331 self,
@@ -244,6 +335,8 @@ class OpType(Enum):
335 lora_rank: int,
336 num_loras: int,
337 num_slices: int,
338 + top_k_num: int | None = None,
339 + num_experts: int | None = None,
340 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
341 """
342 Given num_slices, return the shapes of the A, B, and C matrices
@@ -258,6 +351,16 @@ class OpType(Enum):
351 if self in [OpType.LORA_EXPAND]:
352 # LoRA expand kernels support num_slices inherently in the kernel
353 return ((num_slices, m, k), b_shape, (m, n * num_slices))
354 + if self.is_fused_moe_lora_fn():
355 + return self.matmul_shapes_fused_moe_lora(
356 + m,
357 + k,
358 + n,
359 + num_loras,
360 + num_slices,
361 + top_k_num,
362 + num_experts,
363 + )
364 raise ValueError(f"Unrecognized op_type {self}")
365
366 def bench_fn(self) -> Callable:
@@ -265,6 +368,16 @@ class OpType(Enum):
368 return lora_shrink
369 if self == OpType.LORA_EXPAND:
370 return lora_expand
371 + if self in [
372 + OpType.FUSED_MOE_LORA_GATE_UP_SHRINK,
373 + OpType.FUSED_MOE_LORA_DOWN_SHRINK,
374 + ]:
375 + return fused_moe_lora_shrink
376 + if self in [
377 + OpType.FUSED_MOE_LORA_GATE_UP_EXPAND,
378 + OpType.FUSED_MOE_LORA_DOWN_EXPAND,
379 + ]:
380 + return fused_moe_lora_expand
381
382 raise ValueError(f"Unrecognized optype {self}")
383
@@ -318,6 +431,8 @@ class BenchmarkContext:
431 sort_by_lora_id: bool
432 dtype: torch.dtype
433 seq_length: int | None = None
434 + num_experts: int | None = None # num_experts for MoE based ops
435 + top_k_num: int | None = None # top_k for MoE based ops
436 num_slices: int | None = None # num_slices for slice based ops
437
438 def with_seq_length(self, seq_length: int) -> "BenchmarkContext":
@@ -373,6 +488,11 @@ class BenchmarkTensors:
488 f"{dtype_to_str(self.output.dtype)}"
489 )
490
491 + def get_num_tokens(self, size: int, top_k_num: int, op_type: OpType):
492 + return (
493 + size * top_k_num if op_type in [OpType.FUSED_MOE_LORA_DOWN_SHRINK] else size
494 + )
495 +
496 @staticmethod
497 def make(
498 ctx: BenchmarkContext, op_type: OpType, device: str = "cuda"
@@ -385,6 +505,8 @@ class BenchmarkTensors:
505 ctx.lora_rank,
506 ctx.num_loras,
507 ctx.num_slices,
508 + ctx.top_k_num,
509 + ctx.num_experts,
510 )
511 a_type, b_type, c_type = op_type.matmul_dtypes(ctx.dtype)
512 input_tensor, lora_weights, output_tensor = make_rand_tensors(
@@ -432,17 +554,27 @@ class BenchmarkTensors:
554 prompt_lora_indices_tensor,
555 )
556
435 - def sanity_check(self) -> None:
557 + def sanity_check(self, ctx: BenchmarkContext, op_type: OpType) -> None:
558 """
559 Fails asserts when non-conformality is detected.
560 """
439 - num_tokens = self.input.shape[-2]
561 + num_tokens = (
562 + self.input.shape[1]
563 + if op_type.is_fused_moe_lora_expand_fn()
564 + else self.input.shape[-2]
565 + )
566 # check metadata tensors
441 - assert torch.sum(self.seq_lens) == num_tokens
567 + ## In down shrink case, each token is repeated top_k_num times
568 + assert num_tokens == self.get_num_tokens(
569 + torch.sum(self.seq_lens), ctx.top_k_num, op_type
570 + ), f"Expected {num_tokens} tokens, but got {torch.sum(self.seq_lens)}"
571 num_seqs = self.seq_lens.shape[0]
572 # assert self.seq_start_loc.shape[0] == num_seqs
573 + ## In down shrink case, each prompt corresponds to top_k_num sequences
574 assert self.prompt_lora_mapping.shape[0] == num_seqs
445 - assert self.lora_kernel_meta.token_lora_mapping.shape[0] == num_tokens
575 + assert self.get_num_tokens(
576 + self.lora_kernel_meta.token_lora_mapping.shape[0], ctx.top_k_num, op_type
577 + )
578
579 def to_device(self, device: str):
580 """
@@ -471,21 +603,111 @@ class BenchmarkTensors:
603 to_device(field) if field_name != "no_lora_flag_cpu" else field,
604 )
605
474 - def metadata(self) -> tuple[int, int, int]:
606 + def metadata(self, ctx: BenchmarkContext, op_type: OpType) -> tuple[int, int, int]:
607 """
608 Return num_seqs, num_tokens and max_seq_len
609 """
610 num_seqs = self.seq_lens.shape[0]
479 - num_tokens = self.lora_kernel_meta.token_lora_mapping.shape[0]
611 + num_tokens = self.get_num_tokens(
612 + self.lora_kernel_meta.token_lora_mapping.shape[0], ctx.top_k_num, op_type
613 + )
614 max_seq_len = torch.max(self.seq_lens).item()
615 num_slices = len(self.lora_weights_lst)
616 return num_seqs, num_tokens, max_seq_len, num_slices
617
484 - def as_lora_shrink_kwargs(self) -> dict[str, Any]:
485 - self.sanity_check()
618 + def fused_moe_lora_data_prepare(
619 + self,
620 + block_size: int,
621 + token_lora_mapping: torch.Tensor,
622 + ctx: BenchmarkContext,
623 + ):
624 + def moe_lora_align_block_size(
625 + topk_ids: torch.Tensor,
626 + token_lora_mapping: torch.Tensor,
627 + block_size: int,
628 + num_experts: int,
629 + max_loras: int,
630 + expert_map: torch.Tensor | None = None,
631 + pad_sorted_ids: bool = False,
632 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
633 + """
634 + Aligns tokens and experts into block-sized chunks for LoRA-based
635 + mixture-of-experts (MoE) execution.
636 + """
637 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
638 + if pad_sorted_ids:
639 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
640 + sorted_ids = torch.empty(
641 + (max_loras * max_num_tokens_padded,),
642 + dtype=torch.int32,
643 + device=topk_ids.device,
644 + )
645 + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
646 + # Expert ids must be set default to -1 to prevent a blank block
647 + expert_ids = torch.empty(
648 + (max_loras * max_num_m_blocks,),
649 + dtype=torch.int32,
650 + device=topk_ids.device,
651 + )
652 + num_tokens_post_pad = torch.empty(
653 + (max_loras), dtype=torch.int32, device=topk_ids.device
654 + )
655 +
656 + ops.moe_lora_align_block_size(
657 + topk_ids,
658 + token_lora_mapping,
659 + num_experts,
660 + block_size,
661 + max_loras,
662 + max_num_tokens_padded,
663 + max_num_m_blocks,
664 + sorted_ids,
665 + expert_ids,
666 + num_tokens_post_pad,
667 + )
668 + if expert_map is not None:
669 + expert_ids = expert_map[expert_ids]
670 +
671 + return sorted_ids, expert_ids, num_tokens_post_pad
672 +
673 + num_tokens = ctx.batch_size
674 + curr_topk_ids = torch.randint(
675 + 0,
676 + ctx.num_experts,
677 + (num_tokens, ctx.top_k_num),
678 + device="cuda",
679 + dtype=torch.int32,
680 + )
681 + topk_weights = torch.randint(
682 + 0,
683 + ctx.num_experts,
684 + (num_tokens, ctx.top_k_num),
685 + device="cuda",
686 + dtype=torch.int32,
687 + )
688 +
689 + (sorted_token_ids_lora, expert_ids_lora, num_tokens_post_padded_lora) = (
690 + moe_lora_align_block_size(
691 + topk_ids=curr_topk_ids,
692 + token_lora_mapping=token_lora_mapping,
693 + block_size=block_size,
694 + num_experts=ctx.num_experts,
695 + max_loras=ctx.num_loras,
696 + )
697 + )
698 +
699 + sorted_token_ids = sorted_token_ids_lora.view(ctx.num_loras, -1)
700 + expert_ids = expert_ids_lora.view(ctx.num_loras, -1)
701 + num_tokens_post_padded = num_tokens_post_padded_lora
702 + return (topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded)
703 +
704 + def as_lora_shrink_kwargs(
705 + self, ctx: BenchmarkContext, op_type: OpType
706 + ) -> dict[str, Any]:
707 + self.sanity_check(ctx, op_type)
708 self.to_device(self.input.device)
709
488 - _, num_tokens, _, num_slices = self.metadata()
710 + _, num_tokens, _, num_slices = self.metadata(ctx, op_type)
711
712 # Sanity check matrix shapes.
713 i_shape, lw_shape, o_shape = (
@@ -520,11 +742,13 @@ class BenchmarkTensors:
742 "no_lora_flag_cpu": self.lora_kernel_meta.no_lora_flag_cpu,
743 }
744
523 - def as_lora_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:
524 - self.sanity_check()
745 + def as_lora_expand_kwargs(
746 + self, ctx: BenchmarkContext, op_type: OpType, add_inputs: bool
747 + ) -> dict[str, Any]:
748 + self.sanity_check(ctx, op_type)
749 self.to_device(self.input.device)
750
527 - _, num_tokens, _, num_slices = self.metadata()
751 + _, num_tokens, _, num_slices = self.metadata(ctx, op_type)
752
753 # Sanity check matrix shapes.
754 i_shape, lw_shape, o_shape = (
@@ -561,18 +785,173 @@ class BenchmarkTensors:
785 "no_lora_flag_cpu": self.lora_kernel_meta.no_lora_flag_cpu,
786 }
787
788 + def as_fused_moe_lora_shrink_kwargs(
789 + self, ctx: BenchmarkContext, op_type: OpType
790 + ) -> dict[str, Any]:
791 + self.sanity_check(ctx, op_type)
792 + self.to_device(self.input.device)
793 +
794 + _, num_tokens, _, num_slices = self.metadata(ctx, op_type)
795 +
796 + # Sanity check matrix shapes.
797 + i_shape, lw_shape, o_shape = (
798 + self.input.shape,
799 + self.lora_weights_lst[0].shape,
800 + self.output.shape,
801 + )
802 + # Expected input shape : [num_tokens, hidden_size] for gate_up
803 + # Expected input shape : [top_k_num * num_tokens, hidden_size] for down
804 + assert len(i_shape) == 2
805 + assert i_shape[0] == num_tokens
806 + hidden_size = i_shape[1]
807 + # Expected lora weight shape [max_lora, num_experts, lora_rank, hidden_size]
808 + assert len(lw_shape) == 4
809 + assert lw_shape[-1] == hidden_size
810 + lora_rank = lw_shape[-2]
811 + # Expected output shape : [num_slices, num_tokens, top_k_num, lora_rank]
812 + assert len(o_shape) == 4
813 + assert (
814 + o_shape
815 + == (num_slices, num_tokens // ctx.top_k_num, ctx.top_k_num, lora_rank)
816 + if op_type in [OpType.FUSED_MOE_LORA_DOWN_SHRINK]
817 + else o_shape == (num_slices, num_tokens, ctx.top_k_num, lora_rank)
818 + )
819 + kernel_config = get_lora_op_configs(
820 + op_type.name.lower(),
821 + max_loras=lw_shape[0],
822 + batch=num_tokens,
823 + hidden_size=hidden_size,
824 + rank=lora_rank,
825 + num_slices=num_slices,
826 + add_inputs=False,
827 + )
828 +
829 + (topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded) = (
830 + self.fused_moe_lora_data_prepare(
831 + block_size=kernel_config["BLOCK_SIZE_M"],
832 + token_lora_mapping=self.lora_kernel_meta.token_lora_mapping,
833 + ctx=ctx,
834 + )
835 + )
836 +
837 + return {
838 + "qcurr_hidden_states": self.input,
839 + "lora_a_stacked": self.lora_weights_lst,
840 + "a_intermediate_cache1": self.output,
841 + "topk_weights": topk_weights,
842 + "sorted_token_ids": sorted_token_ids,
843 + "expert_ids": expert_ids,
844 + "num_tokens_post_padded": num_tokens_post_padded,
845 + "top_k_num": ctx.top_k_num,
846 + "device": self.input.device,
847 + "N": lora_rank,
848 + "M": topk_weights.shape[0],
849 + "EM": sorted_token_ids.shape[1],
850 + "K": self.input.shape[1],
851 + "num_tokens": num_tokens,
852 + "num_experts": ctx.num_experts,
853 + "num_slices": num_slices,
854 + "shrink_block_size_m": kernel_config["BLOCK_SIZE_M"],
855 + "shrink_block_size_n": kernel_config["BLOCK_SIZE_N"],
856 + "shrink_block_size_k": kernel_config["BLOCK_SIZE_K"],
857 + "shrink_group_size_m": kernel_config["GROUP_SIZE_M"],
858 + "shrink_num_warps": kernel_config["NUM_WARPS"],
859 + "shrink_num_stages": kernel_config["NUM_STAGES"],
860 + "shrink_split_k": kernel_config.get("SPLIT_K", 1),
861 + "mul_routed_weight": op_type.is_fused_moe_lora_down_fn(),
862 + }
863 +
864 + def as_fused_moe_lora_expand_kwargs(
865 + self, ctx: BenchmarkContext, op_type: OpType
866 + ) -> dict[str, Any]:
867 + self.sanity_check(ctx, op_type)
868 + self.to_device(self.input.device)
869 +
870 + _, num_tokens, _, num_slices = self.metadata(ctx, op_type)
871 +
872 + # Sanity check matrix shapes.
873 + i_shape, lw_shape, o_shape = (
874 + self.input.shape,
875 + self.lora_weights_lst[0].shape,
876 + self.output.shape,
877 + )
878 +
879 + # Expected input shape : [num_slices, num_tokens, top_k_num, lora_rank]
880 + assert len(i_shape) == 4
881 + assert i_shape[0] == num_slices
882 + assert i_shape[1] == num_tokens
883 + lora_rank = i_shape[-1]
884 + # Expected lora weight shape : [num_loras, num_experts, hidden_size, lora_rank]
885 + assert len(lw_shape) == 4
886 + assert lw_shape[-1] == lora_rank
887 + hidden_size = lw_shape[-2]
888 + # Expected output shape : [num_tokens, top_k_num, hidden_size * num_slices]
889 + assert len(o_shape) == 3
890 + assert o_shape == (num_tokens, ctx.top_k_num, hidden_size * num_slices)
891 +
892 + kernel_config = get_lora_op_configs(
893 + op_type.name.lower(),
894 + max_loras=lw_shape[0],
895 + batch=num_tokens,
896 + hidden_size=hidden_size,
897 + rank=lora_rank,
898 + num_slices=num_slices,
899 + add_inputs=False,
900 + )
901 +
902 + (topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded) = (
903 + self.fused_moe_lora_data_prepare(
904 + block_size=kernel_config["BLOCK_SIZE_M"],
905 + token_lora_mapping=self.lora_kernel_meta.token_lora_mapping,
906 + ctx=ctx,
907 + )
908 + )
909 +
910 + return {
911 + "a_intermediate_cache1": self.input,
912 + "lora_b_stacked": self.lora_weights_lst,
913 + "output": self.output,
914 + "topk_weights": topk_weights,
915 + "sorted_token_ids": sorted_token_ids,
916 + "expert_ids": expert_ids,
917 + "num_tokens_post_padded": num_tokens_post_padded,
918 + "top_k_num": ctx.top_k_num,
919 + "device": self.input.device,
920 + "N": lora_rank,
921 + "M": topk_weights.shape[0],
922 + "EM": sorted_token_ids.shape[1],
923 + "K": self.input.shape[1],
924 + "num_tokens": num_tokens,
925 + "num_experts": ctx.num_experts,
926 + "num_slices": num_slices,
927 + "max_lora_rank": lora_rank,
928 + "w1_output_dim_size": lw_shape[2],
929 + "expand_block_size_m": kernel_config["BLOCK_SIZE_M"],
930 + "expand_block_size_n": kernel_config["BLOCK_SIZE_N"],
931 + "expand_block_size_k": kernel_config["BLOCK_SIZE_K"],
932 + "expand_group_size_m": kernel_config["GROUP_SIZE_M"],
933 + "expand_num_warps": kernel_config["NUM_WARPS"],
934 + "expand_num_stages": kernel_config["NUM_STAGES"],
935 + "expand_split_k": kernel_config.get("SPLIT_K", 1),
936 + "mul_routed_weight": op_type.is_fused_moe_lora_down_fn(),
937 + }
938 +
939 def bench_fn_kwargs(
565 - self, op_type: OpType, add_inputs: bool | None = None
940 + self, ctx: BenchmarkContext, op_type: OpType, add_inputs: bool | None = None
941 ) -> dict[str, Any]:
567 - if op_type.is_shrink_fn():
942 + if op_type.is_shrink_fn() or op_type.is_fused_moe_lora_fn():
943 assert add_inputs is None
944 else:
945 assert add_inputs is not None
946
947 if op_type == OpType.LORA_SHRINK:
573 - return self.as_lora_shrink_kwargs()
948 + return self.as_lora_shrink_kwargs(ctx, op_type)
949 if op_type == OpType.LORA_EXPAND:
575 - return self.as_lora_expand_kwargs(add_inputs)
950 + return self.as_lora_expand_kwargs(ctx, op_type, add_inputs)
951 + if op_type.is_fused_moe_lora_shrink_fn():
952 + return self.as_fused_moe_lora_shrink_kwargs(ctx, op_type)
953 + if op_type.is_fused_moe_lora_expand_fn():
954 + return self.as_fused_moe_lora_expand_kwargs(ctx, op_type)
955 raise ValueError(f"Unrecognized optype {self}")
956
957 def test_correctness(
@@ -617,7 +996,7 @@ def bench_optype(
996 test_correctness: bool = False,
997 ) -> TMeasurement:
998 assert arg_pool_size >= 1
620 - if op_type.is_shrink_fn():
999 + if op_type.is_shrink_fn() or op_type.is_fused_moe_lora_fn():
1000 assert expand_fn_add_inputs is None
1001 else:
1002 assert expand_fn_add_inputs is not None
@@ -627,23 +1006,30 @@ def bench_optype(
1006 BenchmarkTensors.make(ctx, op_type) for _ in range(arg_pool_size)
1007 ]
1008 for bt in bench_tensors:
630 - bt.sanity_check()
1009 + bt.sanity_check(ctx, op_type)
1010
1011 # Test correctness of our implementation.
1012 if test_correctness:
1013 + assert op_type in [OpType.LORA_SHRINK, OpType.LORA_EXPAND], (
1014 + f"Correctness testing is not supported for {op_type.name}."
1015 + )
1016 assert all(
635 - [bt.test_correctness(op_type, expand_fn_add_inputs) for bt in bench_tensors]
1017 + [
1018 + bt.test_correctness(ctx, op_type, expand_fn_add_inputs)
1019 + for bt in bench_tensors
1020 + ]
1021 )
1022
1023 # BenchmarkTensors -> dict (kwargs)
1024 kwargs_list = [
640 - bt.bench_fn_kwargs(op_type, add_inputs=expand_fn_add_inputs)
1025 + bt.bench_fn_kwargs(ctx, op_type, add_inputs=expand_fn_add_inputs)
1026 for bt in bench_tensors
1027 ]
1028
1029 # Clear LoRA optimization hash-maps.
1030 _LORA_A_PTR_DICT.clear()
1031 _LORA_B_PTR_DICT.clear()
1032 + _LORA_PTR_DICT.clear()
1033 # Run bench function so that _LORA_A_PTR_DICT and _LORA_B_PTR_DICT are set up
1034 for kwargs in kwargs_list:
1035 op_type.bench_fn()(**kwargs)
@@ -793,7 +1179,9 @@ def run(args: argparse.Namespace, bench_ctxs: list[BenchmarkContext]):
1179
1180 # Benchmark bench_op
1181 expand_fn_add_inputs = (
796 - [None] if bench_op.is_shrink_fn() else args.expand_fn_add_inputs
1182 + [None]
1183 + if bench_op.is_shrink_fn() or bench_op.is_fused_moe_lora_fn()
1184 + else args.expand_fn_add_inputs
1185 )
1186 for add_input_arg in expand_fn_add_inputs:
1187 seq_len_timers.append(
@@ -831,12 +1219,22 @@ def as_benchmark_contexts(
1219 hidden_sizes: list[int], lora_ranks: list[int], args: argparse.Namespace
1220 ) -> list[BenchmarkContext]:
1221 ctxs: list[BenchmarkContext] = []
834 - for batch_size, hidden_size, lora_rank, num_loras, sort_by_lora_id in product( # noqa
1222 + for (
1223 + batch_size,
1224 + hidden_size,
1225 + lora_rank,
1226 + num_loras,
1227 + sort_by_lora_id,
1228 + top_k_num,
1229 + num_experts,
1230 + ) in product( # noqa
1231 args.batch_sizes,
1232 list(hidden_sizes),
1233 lora_ranks,
1234 args.num_loras,
1235 args.sort_by_lora_id,
1236 + args.top_k_nums,
1237 + args.num_experts,
1238 ):
1239 ctxs.append(
1240 BenchmarkContext(
@@ -851,6 +1249,8 @@ def as_benchmark_contexts(
1249 seq_length=None,
1250 sort_by_lora_id=sort_by_lora_id,
1251 dtype=args.dtype,
1252 + top_k_num=top_k_num,
1253 + num_experts=num_experts,
1254 # To be filled based on the OpType to benchmark
1255 num_slices=None,
1256 )
@@ -1012,6 +1412,22 @@ if __name__ == "__main__":
1412 ),
1413 )
1414
1415 + p.add_argument(
1416 + "--top-k-nums",
1417 + nargs="+",
1418 + type=int,
1419 + default=DEFAULT_TOP_K_NUMS,
1420 + help="Top-K values for MoE LoRA operations",
1421 + )
1422 +
1423 + p.add_argument(
1424 + "--num-experts",
1425 + nargs="+",
1426 + type=int,
1427 + default=DEFAULT_NUM_EXPERTS,
1428 + help="Number of experts for MoE LoRA operations",
1429 + )
1430 +
1431 parser = FlexibleArgumentParser(
1432 description=f"""
1433 Benchmark LoRA kernels:
tests/lora/test_fused_moe_lora_kernel.py
+11
@@ -158,6 +158,8 @@ def use_fused_moe_lora_kernel(
158 "BLOCK_SIZE_N": 32,
159 "BLOCK_SIZE_K": 64,
160 "GROUP_SIZE_M": 1,
161 + "NUM_WARPS": 4,
162 + "NUM_STAGES": 3,
163 "SPLIT_K": 1,
164 }
165
@@ -182,6 +184,15 @@ def use_fused_moe_lora_kernel(
184 config["BLOCK_SIZE_N"],
185 config["BLOCK_SIZE_K"],
186 config["GROUP_SIZE_M"],
187 + config["NUM_WARPS"],
188 + config["NUM_STAGES"],
189 + config["SPLIT_K"],
190 + config["BLOCK_SIZE_M"],
191 + config["BLOCK_SIZE_N"],
192 + config["BLOCK_SIZE_K"],
193 + config["GROUP_SIZE_M"],
194 + config["NUM_WARPS"],
195 + config["NUM_STAGES"],
196 config["SPLIT_K"],
197 mul_routed_weight,
198 )
vllm/lora/layers/fused_moe.py
+83 -20
@@ -13,6 +13,7 @@ from vllm.distributed.parallel_state import (
13 get_tensor_model_parallel_world_size,
14 )
15 from vllm.lora.layers.base import BaseLayerWithLoRA
16 +from vllm.lora.ops.triton_ops.utils import get_lora_op_configs
17 from vllm.model_executor.layers.fused_moe import FusedMoE
18 from vllm.model_executor.layers.fused_moe.config import (
19 _get_config_dtype_str,
@@ -39,6 +40,64 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
40 self.device = base_layer.w2_weight.device
41 self._inject_lora_into_fused_moe()
42
43 + def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
44 + normalized_config = {}
45 + for key, value in config.items():
46 + if key.islower():
47 + if key.startswith("block_"):
48 + normalized_key = "BLOCK_SIZE_" + key.split("_")[-1].upper()
49 + else:
50 + normalized_key = key.upper()
51 + else:
52 + normalized_key = key
53 + normalized_config[normalized_key] = value
54 + return normalized_config
55 +
56 + def _get_lora_moe_configs(
57 + self,
58 + op_prefix: str,
59 + lora_a_stacked: torch.Tensor,
60 + lora_b_stacked: torch.Tensor,
61 + num_slices: int,
62 + M: int,
63 + layer: FusedMoE,
64 + top_k: int,
65 + config_dtype: str,
66 + ):
67 + if envs.VLLM_TUNED_CONFIG_FOLDER:
68 + shrink_config = get_lora_op_configs(
69 + op_type=f"fused_moe_lora_{op_prefix}_shrink",
70 + max_loras=lora_a_stacked.shape[0],
71 + batch=M,
72 + hidden_size=lora_a_stacked.shape[-1],
73 + rank=lora_a_stacked.shape[-2],
74 + num_slices=num_slices,
75 + moe_intermediate_size=lora_b_stacked.shape[-2],
76 + )
77 + expand_config = get_lora_op_configs(
78 + op_type=f"fused_moe_lora_{op_prefix}_expand",
79 + max_loras=lora_a_stacked.shape[0],
80 + batch=M,
81 + hidden_size=lora_a_stacked.shape[-1],
82 + rank=lora_a_stacked.shape[-2],
83 + num_slices=num_slices,
84 + moe_intermediate_size=lora_b_stacked.shape[-2],
85 + )
86 + else: # fall back to the default config
87 + get_config_func = functools.partial(
88 + try_get_optimal_moe_config,
89 + layer.w13_weight.size(),
90 + layer.w2_weight.size(),
91 + top_k,
92 + config_dtype,
93 + block_shape=layer.quant_method.moe_quant_config.block_shape,
94 + )
95 + shrink_config = get_config_func(M)
96 + expand_config = get_config_func(M)
97 + shrink_config = self._normalize_keys(shrink_config)
98 + expand_config = self._normalize_keys(expand_config)
99 + return shrink_config, expand_config
100 +
101 def _inject_lora_into_fused_moe(self):
102 moe_state_dict = {}
103 top_k = self.base_layer.top_k
@@ -90,17 +149,19 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
149 num_tokens = hidden_states.size(0)
150 M = min(num_tokens, CHUNK_SIZE)
151
93 - get_config_func = functools.partial(
94 - try_get_optimal_moe_config,
95 - layer.w13_weight.size(),
96 - layer.w2_weight.size(),
97 - top_k,
98 - config_dtype,
99 - block_shape=layer.quant_method.moe_quant_config.block_shape,
152 + shrink_config, expand_config = self._get_lora_moe_configs(
153 + op_prefix="w13",
154 + lora_a_stacked=self.w1_lora_a_stacked,
155 + lora_b_stacked=self.w1_lora_b_stacked,
156 + num_slices=2,
157 + M=M,
158 + layer=layer,
159 + top_k=top_k,
160 + config_dtype=config_dtype,
161 )
162
163 + # get the block size of m from customized config or default config
164 max_loras = self.w1_lora_a_stacked.shape[0]
103 - config = get_config_func(M)
165 (
166 sorted_token_ids_lora,
167 expert_ids_lora,
@@ -108,7 +169,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
169 ) = self.punica_wrapper.moe_lora_align_block_size(
170 curr_topk_ids,
171 num_tokens,
111 - config["BLOCK_SIZE_M"],
172 + shrink_config["BLOCK_SIZE_M"],
173 self.base_layer.local_num_experts,
174 max_loras,
175 self.adapter_enabled,
@@ -138,7 +199,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
199 num_tokens_post_padded_lora,
200 max_lora_rank,
201 top_k,
141 - config,
202 + shrink_config, ## pass the shrink config
203 + expand_config, ## pass the expand config
204 self.adapter_enabled,
205 )
206
@@ -164,17 +226,17 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
226 num_tokens = hidden_states.size(0)
227 M = min(num_tokens, CHUNK_SIZE)
228
167 - get_config_func = functools.partial(
168 - try_get_optimal_moe_config,
169 - layer.w13_weight.size(),
170 - layer.w2_weight.size(),
171 - top_k,
172 - config_dtype,
173 - block_shape=layer.quant_method.moe_quant_config.block_shape,
229 + shrink_config, expand_config = self._get_lora_moe_configs(
230 + op_prefix="w2",
231 + lora_a_stacked=self.w2_lora_a_stacked,
232 + lora_b_stacked=self.w2_lora_b_stacked,
233 + num_slices=1,
234 + M=M,
235 + layer=layer,
236 + top_k=top_k,
237 + config_dtype=config_dtype,
238 )
239
176 - config = get_config_func(M)
177 -
240 sorted_token_ids_lora = moe_state_dict["sorted_token_ids_lora"]
241 expert_ids_lora = moe_state_dict["expert_ids_lora"]
242 num_tokens_post_padded_lora = moe_state_dict[
@@ -197,7 +259,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
259 num_tokens_post_padded_lora,
260 max_lora_rank,
261 top_k,
200 - config,
262 + shrink_config, ## pass the shrink config
263 + expand_config, ## pass the expand config
264 self.adapter_enabled,
265 True,
266 )
vllm/lora/ops/triton_ops/README_TUNING.md
+10 -1
@@ -44,8 +44,17 @@ For `shrink`, the config file is named as `{gpu_name}_SHRINK.json`, e.g. `NVIDIA
44
45 For `expand`, the config fileis named as `{gpu_name}_EXPAND_{add_input}.json`, e.g. `NVIDIA_H200_EXPAND_TRUE.json`.
46
47 +For `fused_moe_lora_w13_shrink`, the config file is named as `{gpu_name}_FUSED_MOE_LORA_W13_SHRINK.json`, e.g. `NVIDIA_H200_FUSED_MOE_LORA_W13_SHRINK.json`.
48 +
49 +For `fused_moe_lora_w13_expand`, the config file is named as `{gpu_name}_FUSED_MOE_LORA_W13_EXPAND.json`, e.g. `NVIDIA_H200_FUSED_MOE_LORA_W13_EXPAND.json`.
50 +
51 +For `fused_moe_lora_w2_shrink`, the config file is named as `{gpu_name}_FUSED_MOE_LORA_W2_SHRINK.json`, e.g. `NVIDIA_H200_FUSED_MOE_LORA_W2_SHRINK.json`.
52 +
53 +For `fused_moe_lora_w2_expand`, the config file is named as `{gpu_name}_FUSED_MOE_LORA_W2_EXPAND.json`, e.g. `NVIDIA_H200_FUSED_MOE_LORA_W2_EXPAND.json`.
54 +
55 The `gpu_name` can be automatically detected by calling `torch.cuda.get_device_name()`
56
57 ### Json Structure
58
51 -Optimal kernel configuration files are saved as JSON files with the structure `config_data[max_loras][num_slices][m][k][n]`
59 +Optimal kernel configuration files are saved as JSON files with the structure `config_data[max_loras][num_slices][m][k][n][i]`
60 +where `i` is an optional dimension in the `fused_moe_lora` configuration, representing the intermediate size of the MoE layer.
vllm/lora/ops/triton_ops/__init__.py
+8 -1
@@ -1,7 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from vllm.lora.ops.triton_ops.fused_moe_lora_op import fused_moe_lora
4 +
5 +from vllm.lora.ops.triton_ops.fused_moe_lora_op import (
6 + fused_moe_lora,
7 + fused_moe_lora_expand,
8 + fused_moe_lora_shrink,
9 +)
10 from vllm.lora.ops.triton_ops.lora_expand_op import lora_expand
11 from vllm.lora.ops.triton_ops.lora_kernel_metadata import LoRAKernelMeta
12 from vllm.lora.ops.triton_ops.lora_shrink_op import lora_shrink
@@ -11,4 +16,6 @@ __all__ = [
16 "lora_shrink",
17 "LoRAKernelMeta",
18 "fused_moe_lora",
19 + "fused_moe_lora_shrink",
20 + "fused_moe_lora_expand",
21 ]
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+296 -60
@@ -176,88 +176,50 @@ def _fused_moe_lora_kernel(
176
177
178 @torch.inference_mode()
179 -def _fused_moe_lora(
180 - output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)
179 +def _fused_moe_lora_shrink(
180 + a_intermediate_cache1: torch.Tensor,
181 + # (num_slices, num_tokens, top_k_num, max_lora_rank)
182 qcurr_hidden_states: torch.Tensor, # (num_tokens, K,)
183 lora_a_stacked: list[
184 torch.Tensor
185 ], # [(max_loras, num_experts, max_lora_rank, K,),...]
185 - lora_b_stacked: list[
186 - torch.Tensor
187 - ], # [(max_loras, num_experts, N, max_lora_rank,),...]
186 topk_weights: torch.Tensor, # (num_tokens, top_k_num)
187 sorted_token_ids: torch.Tensor, # (max_loras, _)
188 expert_ids: torch.Tensor, # (max_loras, _ ,)
189 num_tokens_post_padded: torch.Tensor, # (max_loras, )
192 - max_lora_rank: int,
190 top_k_num: int,
191 lora_ids: torch.Tensor,
192 adapter_enabled: torch.Tensor,
193 + ## adding for kernel
194 + device: torch.device,
195 + N: int,
196 + M: int,
197 + EM: int,
198 + K: int,
199 + num_tokens: int,
200 + num_experts: int,
201 + num_slices: int,
202 block_size_m: int,
203 block_size_n: int,
204 block_size_k: int,
205 group_size_m: int,
206 + num_warps: int,
207 + num_stages: int,
208 split_k: int,
209 mul_routed_weight: bool = False,
210 ) -> None:
203 - assert len(lora_a_stacked) == len(lora_b_stacked) > 0
204 - assert (
205 - sorted_token_ids.dim()
206 - == expert_ids.dim()
207 - == topk_weights.dim()
208 - == qcurr_hidden_states.dim()
209 - == 2
210 - )
211 - assert (
212 - sorted_token_ids.shape[0]
213 - == expert_ids.shape[0]
214 - == num_tokens_post_padded.shape[0]
215 - )
216 - assert len(lora_b_stacked) * lora_b_stacked[0].shape[-2] == output.shape[-1]
217 - assert output.shape[0] == topk_weights.shape[0]
218 - assert top_k_num == topk_weights.shape[1]
219 -
220 - for lora_a, lora_b in zip(lora_a_stacked, lora_b_stacked):
221 - assert lora_a.dtype == lora_b.dtype == output.dtype == qcurr_hidden_states.dtype
222 - assert lora_a.dtype in [torch.float16, torch.bfloat16]
223 -
224 - device = qcurr_hidden_states.device
225 - num_slices = len(lora_a_stacked)
211 + w1_lora_a_stacked = lora_a_stacked[0]
212
227 - config = {
213 + shrink_config = {
214 "BLOCK_SIZE_M": block_size_m,
215 "BLOCK_SIZE_N": block_size_n,
216 "BLOCK_SIZE_K": block_size_k,
217 "GROUP_SIZE_M": group_size_m,
218 + "num_warps": num_warps,
219 + "num_stages": num_stages,
220 "SPLIT_K": split_k,
221 }
222
235 - w1_lora_a_stacked = lora_a_stacked[0]
236 - w1_lora_b_stacked = lora_b_stacked[0]
237 - num_experts = lora_a_stacked[0].shape[1]
238 -
239 - N = max_lora_rank
240 - M = topk_weights.shape[0]
241 - EM = sorted_token_ids.shape[1]
242 - K = qcurr_hidden_states.shape[1]
243 - num_tokens = M * top_k_num
244 - w1_output_dim_size = w1_lora_b_stacked.shape[2]
245 -
246 - lora_intermediate_cache1 = torch.zeros(
247 - (num_slices * M * top_k_num * (max_lora_rank + w1_output_dim_size)),
248 - dtype=output.dtype,
249 - device=device,
250 - )
251 -
252 - # slices
253 - a_intermediate_size = num_slices * M * top_k_num * max_lora_rank
254 - a_intermediate_cache1 = lora_intermediate_cache1[:a_intermediate_size].view(
255 - num_slices, M, top_k_num, max_lora_rank
256 - )
257 - b_intermediate_cache1 = lora_intermediate_cache1[a_intermediate_size:].view(
258 - num_slices, M, top_k_num, w1_output_dim_size
259 - )
260 -
223 b_ptr = _get_ptr(lora_a_stacked, device)
224
225 grid = lambda META: (
@@ -299,19 +261,70 @@ def _fused_moe_lora(
261 num_slice_c=num_slices,
262 top_k=1 if mul_routed_weight else top_k_num,
263 MUL_ROUTED_WEIGHT=False,
302 - **config,
264 + **shrink_config,
265 )
266
267 +
268 +@torch.inference_mode()
269 +def _fused_moe_lora_expand(
270 + output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)
271 + a_intermediate_cache1: torch.Tensor, # (num_slices, M, top_k_num, max_lora_rank)
272 + lora_b_stacked: list[
273 + torch.Tensor
274 + ], # [(max_loras, num_experts, max_lora_rank, K,),...]
275 + topk_weights: torch.Tensor, # (num_tokens, top_k_num)
276 + sorted_token_ids: torch.Tensor, # (max_loras, _)
277 + expert_ids: torch.Tensor, # (max_loras, _ ,)
278 + num_tokens_post_padded: torch.Tensor, # (max_loras, )
279 + top_k_num: int,
280 + lora_ids: torch.Tensor,
281 + adapter_enabled: torch.Tensor,
282 + ## adding for kernel
283 + device: torch.device,
284 + N: int,
285 + M: int,
286 + EM: int,
287 + K: int,
288 + num_tokens: int,
289 + num_experts: int,
290 + num_slices: int,
291 + max_lora_rank: int,
292 + w1_output_dim_size: int,
293 + block_size_m: int,
294 + block_size_n: int,
295 + block_size_k: int,
296 + group_size_m: int,
297 + num_warps: int,
298 + num_stages: int,
299 + split_k: int,
300 + mul_routed_weight: bool = False,
301 +) -> None:
302 b_ptr = _get_ptr(lora_b_stacked, device)
303 K = max_lora_rank
304 N = w1_output_dim_size
305
306 + w1_lora_b_stacked = lora_b_stacked[0]
307 +
308 a_intermediate_cache1 = a_intermediate_cache1.view(
309 -1, a_intermediate_cache1.shape[3]
310 )
311
313 - # Set split_k = 1 for expand calls
314 - config["SPLIT_K"] = 1
312 + b_intermediate_cache1 = torch.zeros(
313 + (num_slices, M, top_k_num, w1_output_dim_size),
314 + dtype=output.dtype,
315 + device=device,
316 + )
317 +
318 + expand_config = {
319 + "BLOCK_SIZE_M": block_size_m,
320 + "BLOCK_SIZE_N": block_size_n,
321 + "BLOCK_SIZE_K": block_size_k,
322 + "GROUP_SIZE_M": group_size_m,
323 + "num_warps": num_warps,
324 + "num_stages": num_stages,
325 + "SPLIT_K": split_k, # Set split_k = 1 for expand calls
326 + }
327 +
328 grid = lambda META: (
329 triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
330 len(lora_b_stacked),
@@ -348,12 +361,142 @@ def _fused_moe_lora(
361 num_slice_c=num_slices,
362 top_k=1,
363 MUL_ROUTED_WEIGHT=mul_routed_weight,
351 - **config,
364 + **expand_config,
365 )
366 for i in range(num_slices):
367 output[:, :, i * N : (i + 1) * N] += b_intermediate_cache1[i]
368
369
370 +@torch.inference_mode()
371 +def _fused_moe_lora(
372 + output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)
373 + qcurr_hidden_states: torch.Tensor, # (num_tokens, K,)
374 + lora_a_stacked: list[
375 + torch.Tensor
376 + ], # [(max_loras, num_experts, max_lora_rank, K,),...]
377 + lora_b_stacked: list[
378 + torch.Tensor
379 + ], # [(max_loras, num_experts, N, max_lora_rank,),...]
380 + topk_weights: torch.Tensor, # (num_tokens, top_k_num)
381 + sorted_token_ids: torch.Tensor, # (max_loras, _)
382 + expert_ids: torch.Tensor, # (max_loras, _ ,)
383 + num_tokens_post_padded: torch.Tensor, # (max_loras, )
384 + max_lora_rank: int,
385 + top_k_num: int,
386 + lora_ids: torch.Tensor,
387 + adapter_enabled: torch.Tensor,
388 + shrink_block_size_m: int,
389 + shrink_block_size_n: int,
390 + shrink_block_size_k: int,
391 + shrink_group_size_m: int,
392 + shrink_num_warps: int,
393 + shrink_num_stages: int,
394 + shrink_split_k: int,
395 + expand_block_size_m: int,
396 + expand_block_size_n: int,
397 + expand_block_size_k: int,
398 + expand_group_size_m: int,
399 + expand_num_warps: int,
400 + expand_num_stages: int,
401 + expand_split_k: int,
402 + mul_routed_weight: bool = False,
403 +) -> None:
404 + assert len(lora_a_stacked) == len(lora_b_stacked) > 0
405 + assert (
406 + sorted_token_ids.dim()
407 + == expert_ids.dim()
408 + == topk_weights.dim()
409 + == qcurr_hidden_states.dim()
410 + == 2
411 + )
412 + assert (
413 + sorted_token_ids.shape[0]
414 + == expert_ids.shape[0]
415 + == num_tokens_post_padded.shape[0]
416 + )
417 + assert len(lora_b_stacked) * lora_b_stacked[0].shape[-2] == output.shape[-1]
418 + assert output.shape[0] == topk_weights.shape[0]
419 + assert top_k_num == topk_weights.shape[1]
420 + device = qcurr_hidden_states.device
421 + num_slices = len(lora_a_stacked)
422 + w1_lora_b_stacked = lora_b_stacked[0]
423 + num_experts = lora_a_stacked[0].shape[1]
424 + N = max_lora_rank
425 + M = topk_weights.shape[0]
426 + EM = sorted_token_ids.shape[1]
427 + K = qcurr_hidden_states.shape[1]
428 + num_tokens = M * top_k_num
429 + w1_output_dim_size = w1_lora_b_stacked.shape[2]
430 +
431 + a_intermediate_cache1 = torch.zeros(
432 + (num_slices, M, top_k_num, max_lora_rank),
433 + dtype=output.dtype,
434 + device=device,
435 + )
436 +
437 + _fused_moe_lora_shrink(
438 + a_intermediate_cache1,
439 + qcurr_hidden_states,
440 + lora_a_stacked,
441 + topk_weights,
442 + sorted_token_ids,
443 + expert_ids,
444 + num_tokens_post_padded,
445 + top_k_num,
446 + lora_ids,
447 + adapter_enabled,
448 + ## adding for kernel
449 + device,
450 + N,
451 + M,
452 + EM,
453 + K,
454 + num_tokens,
455 + num_experts,
456 + num_slices,
457 + shrink_block_size_m,
458 + shrink_block_size_n,
459 + shrink_block_size_k,
460 + shrink_group_size_m,
461 + shrink_num_warps,
462 + shrink_num_stages,
463 + shrink_split_k,
464 + mul_routed_weight,
465 + )
466 +
467 + _fused_moe_lora_expand(
468 + output,
469 + a_intermediate_cache1,
470 + lora_b_stacked,
471 + topk_weights,
472 + sorted_token_ids,
473 + expert_ids,
474 + num_tokens_post_padded,
475 + top_k_num,
476 + lora_ids,
477 + adapter_enabled,
478 + ## adding for kernel
479 + device,
480 + N,
481 + M,
482 + EM,
483 + K,
484 + num_tokens,
485 + num_experts,
486 + num_slices,
487 + max_lora_rank,
488 + w1_output_dim_size,
489 + expand_block_size_m,
490 + expand_block_size_n,
491 + expand_block_size_k,
492 + expand_group_size_m,
493 + expand_num_warps,
494 + expand_num_stages,
495 + expand_split_k,
496 + mul_routed_weight,
497 + )
498 +
499 +
500 def _fused_moe_lora_fake(
501 output: torch.Tensor,
502 qcurr_hidden_states: torch.Tensor,
@@ -367,10 +510,84 @@ def _fused_moe_lora_fake(
510 top_k_num: int,
511 lora_ids: torch.Tensor,
512 adapter_enabled: torch.Tensor,
513 + shrink_block_size_m: int,
514 + shrink_block_size_n: int,
515 + shrink_block_size_k: int,
516 + shrink_group_size_m: int,
517 + shrink_num_warps: int,
518 + shrink_num_stages: int,
519 + shrink_split_k: int,
520 + expand_block_size_m: int,
521 + expand_block_size_n: int,
522 + expand_block_size_k: int,
523 + expand_group_size_m: int,
524 + expand_num_warps: int,
525 + expand_num_stages: int,
526 + expand_split_k: int,
527 + mul_routed_weight: bool = False,
528 +) -> None:
529 + return
530 +
531 +
532 +def _fused_moe_lora_shrink_fake(
533 + a_intermediate_cache1: torch.Tensor,
534 + qcurr_hidden_states: torch.Tensor,
535 + lora_a_stacked: list[torch.Tensor],
536 + topk_weights: torch.Tensor,
537 + sorted_token_ids: torch.Tensor,
538 + expert_ids: torch.Tensor,
539 + num_tokens_post_padded: torch.Tensor,
540 + top_k_num: int,
541 + lora_ids: torch.Tensor,
542 + adapter_enabled: torch.Tensor,
543 + device: torch.device,
544 + N: int,
545 + M: int,
546 + EM: int,
547 + K: int,
548 + num_tokens: int,
549 + num_experts: int,
550 + num_slices: int,
551 + block_size_m: int,
552 + block_size_n: int,
553 + block_size_k: int,
554 + group_size_m: int,
555 + num_warps: int,
556 + num_stages: int,
557 + split_k: int,
558 + mul_routed_weight: bool = False,
559 +) -> None:
560 + return
561 +
562 +
563 +def _fused_moe_lora_expand_fake(
564 + output: torch.Tensor,
565 + a_intermediate_cache1: torch.Tensor,
566 + lora_b_stacked: list[torch.Tensor],
567 + topk_weights: torch.Tensor,
568 + sorted_token_ids: torch.Tensor,
569 + expert_ids: torch.Tensor,
570 + num_tokens_post_padded: torch.Tensor,
571 + top_k_num: int,
572 + lora_ids: torch.Tensor,
573 + adapter_enabled: torch.Tensor,
574 + device: torch.device,
575 + N: int,
576 + M: int,
577 + EM: int,
578 + K: int,
579 + num_tokens: int,
580 + num_experts: int,
581 + num_slices: int,
582 + max_lora_rank: int,
583 + w1_output_dim_size: int,
584 block_size_m: int,
585 block_size_n: int,
586 block_size_k: int,
587 group_size_m: int,
588 + num_warps: int,
589 + num_stages: int,
590 + split_k: int,
591 mul_routed_weight: bool = False,
592 ) -> None:
593 return
@@ -383,7 +600,26 @@ try:
600 mutates_args=["output"],
601 fake_impl=_fused_moe_lora_fake,
602 )
603 +
604 + direct_register_custom_op(
605 + op_name="fused_moe_lora_shrink",
606 + op_func=_fused_moe_lora_shrink,
607 + mutates_args=["a_intermediate_cache1"],
608 + fake_impl=_fused_moe_lora_shrink_fake,
609 + )
610 +
611 + direct_register_custom_op(
612 + op_name="fused_moe_lora_expand",
613 + op_func=_fused_moe_lora_expand,
614 + mutates_args=["output"],
615 + fake_impl=_fused_moe_lora_expand_fake,
616 + )
617 +
618 fused_moe_lora = torch.ops.vllm.fused_moe_lora
619 + fused_moe_lora_shrink = torch.ops.vllm.fused_moe_lora_shrink
620 + fused_moe_lora_expand = torch.ops.vllm.fused_moe_lora_expand
621
622 except AttributeError:
623 fused_moe_lora = _fused_moe_lora
624 + fused_moe_lora_shrink = _fused_moe_lora_shrink
625 + fused_moe_lora_expand = _fused_moe_lora_expand
vllm/lora/ops/triton_ops/utils.py
+38 -5
@@ -154,13 +154,13 @@ def load_lora_op_config(op_type: str, add_inputs: bool | None) -> dict | None:
154 gpu_name = gpu_name.replace("-", "_")
155
156 config_fname = None
157 - if op_type == "shrink":
158 - config_fname = f"{gpu_name}_{op_type.upper()}.json"
159 - else:
160 - assert op_type == "expand"
157 + # only expand op needs to consider add_inputs
158 + if op_type == "expand":
159 config_fname = (
160 f"{gpu_name}_{op_type.upper()}_{str(add_inputs).upper()}.json"
161 )
162 + else:
163 + config_fname = f"{gpu_name}_{op_type.upper()}.json"
164
165 config_path = Path(f"{user_defined_config_folder}/{config_fname}")
166 if not config_path.exists():
@@ -186,8 +186,17 @@ def get_lora_op_configs(
186 rank: int,
187 num_slices: int,
188 add_inputs: bool | None = None,
189 + moe_intermediate_size: int | None = None,
190 ) -> dict[str, int | None]:
190 - assert op_type in ["shrink", "expand"]
191 + # Add support for fused_moe_lora ops
192 + assert op_type in [
193 + "shrink",
194 + "expand",
195 + "fused_moe_lora_w13_shrink",
196 + "fused_moe_lora_w13_expand",
197 + "fused_moe_lora_w2_shrink",
198 + "fused_moe_lora_w2_expand",
199 + ]
200
201 # default config
202 default = {}
@@ -203,6 +212,22 @@ def get_lora_op_configs(
212 "num_stages": 2,
213 "max_nreg": None,
214 }
215 + # The default config for fused_moe_lora ops
216 + elif op_type in [
217 + "fused_moe_lora_w13_shrink",
218 + "fused_moe_lora_w13_expand",
219 + "fused_moe_lora_w2_shrink",
220 + "fused_moe_lora_w2_expand",
221 + ]:
222 + default = {
223 + "block_m": 64,
224 + "block_n": 64,
225 + "block_k": 32,
226 + "num_warps": 4,
227 + "num_stages": 3,
228 + "group_size_m": 8,
229 + "split_k": 1,
230 + }
231 else:
232 default = {
233 "block_m": 64,
@@ -247,5 +272,13 @@ def get_lora_op_configs(
272 or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - n))]
273 )
274
275 + # slice by moe-intermediate-size if applicable
276 + if moe_intermediate_size is not None:
277 + i = moe_intermediate_size
278 + config_data = (
279 + config_data.get(str(i))
280 + or config_data[min(config_data.keys(), key=lambda x: abs(int(x) - i))]
281 + )
282 +
283 assert config_data is not None
284 return config_data
vllm/lora/punica_wrapper/punica_base.py
+2 -1
@@ -479,7 +479,8 @@ class PunicaWrapperBase(PunicaWrapperABC):
479 num_tokens_post_padded: torch.Tensor,
480 max_lora_rank: int,
481 top_k_num: int,
482 - config,
482 + shrink_config,
483 + expand_config,
484 adapter_enabled: torch.Tensor,
485 mul_routed_weight=False,
486 ):
vllm/lora/punica_wrapper/punica_gpu.py
+16 -6
@@ -367,7 +367,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
367 num_tokens_post_padded: torch.Tensor,
368 max_lora_rank: int,
369 top_k_num: int,
370 - config,
370 + shrink_config,
371 + expand_config,
372 adapter_enabled: torch.Tensor,
373 mul_routed_weight=False,
374 ):
@@ -388,10 +389,19 @@ class PunicaWrapperGPU(PunicaWrapperBase):
389 top_k_num,
390 lora_ids,
391 adapter_enabled,
391 - config["BLOCK_SIZE_M"],
392 - config["BLOCK_SIZE_N"],
393 - config["BLOCK_SIZE_K"],
394 - config["GROUP_SIZE_M"],
395 - config.get("SPLIT_K", 1),
392 + shrink_config.get("BLOCK_SIZE_M", 64),
393 + shrink_config.get("BLOCK_SIZE_N", 64),
394 + shrink_config.get("BLOCK_SIZE_K", 32),
395 + shrink_config.get("GROUP_SIZE_M", 8),
396 + shrink_config.get("NUM_WARPS", 4),
397 + shrink_config.get("NUM_STAGES", 3),
398 + shrink_config.get("SPLIT_K", 1),
399 + expand_config.get("BLOCK_SIZE_M", 64),
400 + expand_config.get("BLOCK_SIZE_N", 64),
401 + expand_config.get("BLOCK_SIZE_K", 32),
402 + expand_config.get("GROUP_SIZE_M", 8),
403 + expand_config.get("NUM_WARPS", 4),
404 + expand_config.get("NUM_STAGES", 3),
405 + expand_config.get("SPLIT_K", 1),
406 mul_routed_weight,
407 )