@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5f6cbf60d6
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 14, 05:30 UTC

[Feature][Kernel]FusedMoE LoRA (#21229)

Signed-off-by: wuchen <cntryroa@gmail.com> Signed-off-by: banjuede <lmklhc@163.com> Signed-off-by: Chen Wu <cntryroa@gmail.com> Signed-off-by: Danielle Robinson <dmmaddix@amazon.com> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: bk-201 <joy25810@foxmail.com> Co-authored-by: wuchen <wuchen@zetyun.com> Co-authored-by: Nathan Van Gheem <vangheem@gmail.com> Co-authored-by: banjuede <lmklhc@163.com> Co-authored-by: Danielle Robinson <dmmaddix@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: bk-201 <joy25810@foxmail.com>

Chen Wu committed Oct 21, 2025 at 11:01 UTC 5f6cbf60d667408ef1bf3871c0b7563effa6094a
28 files changed +2084 -55
.buildkite/test-pipeline.yaml
+7 -1
@@ -384,7 +384,12 @@ steps:
384 --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT \
385 --ignore=lora/test_chatglm3_tp.py \
386 --ignore=lora/test_llama_tp.py \
387 - --ignore=lora/test_llm_with_multi_loras.py
387 + --ignore=lora/test_llm_with_multi_loras.py \
388 + --ignore=lora/test_olmoe_tp.py \
389 + --ignore=lora/test_deepseekv2_tp.py \
390 + --ignore=lora/test_gptoss.py \
391 + --ignore=lora/test_qwen3moe_tp.py
392 +
393 parallelism: 4
394
395 - label: PyTorch Compilation Unit Tests # 15min
@@ -1065,6 +1070,7 @@ steps:
1070 - pytest -v -s -x lora/test_chatglm3_tp.py
1071 - pytest -v -s -x lora/test_llama_tp.py
1072 - pytest -v -s -x lora/test_llm_with_multi_loras.py
1073 + - pytest -v -s -x lora/test_olmoe_tp.py
1074
1075
1076 - label: Weight Loading Multiple GPU Test # 33min
CMakeLists.txt
+1
@@ -883,6 +883,7 @@ target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
883 set(VLLM_MOE_EXT_SRC
884 "csrc/moe/torch_bindings.cpp"
885 "csrc/moe/moe_align_sum_kernels.cu"
886 + "csrc/moe/moe_lora_align_sum_kernels.cu"
887 "csrc/moe/topk_softmax_kernels.cu")
888
889 if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/moe/moe_lora_align_sum_kernels.cu new
+173
@@ -0,0 +1,173 @@
1 +#include <stdio.h>
2 +#include <stdlib.h>
3 +#include <time.h>
4 +#include <torch/all.h>
5 +#include <ATen/cuda/CUDAContext.h>
6 +#include <c10/cuda/CUDAGuard.h>
7 +
8 +#include <ATen/ATen.h>
9 +#include <ATen/cuda/Atomic.cuh>
10 +
11 +#include "../cuda_compat.h"
12 +#include "../dispatch_utils.h"
13 +#include "core/math.hpp"
14 +
15 +namespace {
16 +
17 +__device__ __forceinline__ int32_t index(int32_t total_col, int32_t row,
18 + int32_t col) {
19 + return row * total_col + col;
20 +}
21 +
22 +} // namespace
23 +
24 +// TODO: Refactor common parts with moe_align_sum_kernels
25 +template <typename scalar_t, typename token_cnts_t>
26 +__global__ void moe_lora_align_sum_kernel(
27 + scalar_t* __restrict__ topk_ids, int32_t* token_lora_mapping,
28 + int64_t block_size, int num_experts, int max_loras, size_t numel,
29 + int max_num_tokens_padded, int max_num_m_blocks,
30 + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
31 + int topk_num, int32_t* total_tokens_post_pad) {
32 + const size_t tokens_per_thread = div_ceil(numel, blockDim.x);
33 + const size_t start_idx = threadIdx.x * tokens_per_thread;
34 +
35 + int lora_id = blockIdx.x;
36 + extern __shared__ int32_t shared_mem[];
37 + int32_t* cumsum = shared_mem;
38 + token_cnts_t* tokens_cnts = (token_cnts_t*)(shared_mem + num_experts + 1);
39 +
40 + // Initialize sorted_token_ids with numel
41 + for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
42 + sorted_token_ids[lora_id * max_num_tokens_padded + it] = numel;
43 + }
44 +
45 + // Initialize expert_ids with -1
46 + for (size_t it = threadIdx.x; it < max_num_m_blocks; it += blockDim.x) {
47 + expert_ids[lora_id * max_num_m_blocks + it] = -1;
48 + }
49 +
50 + // Initialize total_tokens_post_pad with 0
51 + if (threadIdx.x == 0) {
52 + total_tokens_post_pad[lora_id] = 0;
53 + }
54 +
55 + for (int i = 0; i < num_experts; ++i) {
56 + tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
57 + }
58 +
59 + for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
60 + int mask = token_lora_mapping[i / topk_num] == lora_id;
61 + int idx = index(num_experts, threadIdx.x + 1, topk_ids[i]);
62 + tokens_cnts[idx] += mask;
63 + }
64 +
65 + __syncthreads();
66 +
67 + // For each expert we accumulate the token counts from the different threads.
68 + if (threadIdx.x < num_experts) {
69 + tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0;
70 + for (int i = 1; i <= blockDim.x; ++i) {
71 + tokens_cnts[index(num_experts, i, threadIdx.x)] +=
72 + tokens_cnts[index(num_experts, i - 1, threadIdx.x)];
73 + }
74 + }
75 +
76 + __syncthreads();
77 +
78 + // We accumulate the token counts of all experts in thread 0.
79 + if (threadIdx.x == 0) {
80 + cumsum[0] = 0;
81 + for (int i = 1; i <= num_experts; ++i) {
82 + cumsum[i] = cumsum[i - 1] +
83 + div_ceil(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
84 + block_size) *
85 + block_size;
86 + }
87 + total_tokens_post_pad[lora_id] = static_cast<int32_t>(cumsum[num_experts]);
88 + }
89 +
90 + __syncthreads();
91 +
92 + /**
93 + * For each expert, each thread processes the tokens of the corresponding
94 + * blocks and stores the corresponding expert_id for each block.
95 + */
96 + if (threadIdx.x < num_experts) {
97 + for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
98 + i += block_size) {
99 + expert_ids[index(max_num_m_blocks, lora_id, i / block_size)] =
100 + threadIdx.x;
101 + }
102 + }
103 +
104 + for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
105 + int32_t expert_id = topk_ids[i];
106 + /** The cumsum[expert_id] stores the starting index of the tokens that the
107 + * expert with expert_id needs to process, and
108 + * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
109 + * processed by the expert with expert_id within the current thread's token
110 + * shard.
111 + */
112 + int32_t rank_post_pad =
113 + tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
114 + cumsum[expert_id];
115 +
116 + int mask = (int)token_lora_mapping[i / topk_num] == lora_id;
117 + atomicAdd(
118 + &sorted_token_ids[index(max_num_tokens_padded, lora_id, rank_post_pad)],
119 + (i - numel) * mask);
120 + tokens_cnts[index(num_experts, threadIdx.x, expert_id)] += mask;
121 + }
122 +}
123 +
124 +void moe_lora_align_block_size(torch::Tensor topk_ids,
125 + torch::Tensor token_lora_mapping,
126 + int64_t num_experts, int64_t block_size,
127 + int64_t max_loras,
128 + torch::Tensor sorted_token_ids,
129 + torch::Tensor expert_ids,
130 + torch::Tensor num_tokens_post_pad) {
131 + const int topk_num = topk_ids.size(1);
132 +
133 + int max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1);
134 +
135 + TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
136 + max_num_tokens_padded = round_to_next_multiple_of(
137 + max_num_tokens_padded, static_cast<int>(block_size));
138 + int max_num_m_blocks = div_ceil(max_num_tokens_padded, block_size);
139 +
140 + int device_max_shared_mem;
141 + auto dev = topk_ids.get_device();
142 + cudaDeviceGetAttribute(&device_max_shared_mem,
143 + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
144 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
145 +
146 + const int32_t num_thread = max((int32_t)num_experts, 128); // WARP_SIZE,
147 + TORCH_CHECK(num_thread <= 1024,
148 + "num_thread must be less than 1024, "
149 + "and fallback is not implemented yet.");
150 + const int32_t shared_mem = (num_thread + 1) * num_experts * sizeof(int32_t) +
151 + (num_experts + 1) * sizeof(int32_t);
152 +
153 + if (shared_mem > device_max_shared_mem) {
154 + TORCH_CHECK(false,
155 + "Shared memory usage exceeds device limit, and global memory "
156 + "fallback is not implemented yet.");
157 + }
158 +
159 + VLLM_DISPATCH_INTEGRAL_TYPES(
160 + topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
161 + dim3 blockDim(num_thread);
162 + auto kernel = moe_lora_align_sum_kernel<scalar_t, int32_t>;
163 + AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
164 + (void*)kernel, shared_mem));
165 + kernel<<<max_loras, blockDim, shared_mem, stream>>>(
166 + topk_ids.data_ptr<scalar_t>(),
167 + token_lora_mapping.data_ptr<int32_t>(), block_size, num_experts,
168 + max_loras, topk_ids.numel(), max_num_tokens_padded,
169 + max_num_m_blocks, sorted_token_ids.data_ptr<int32_t>(),
170 + expert_ids.data_ptr<int32_t>(), topk_num,
171 + num_tokens_post_pad.data_ptr<int32_t>());
172 + });
173 +}
\ No newline at end of file
csrc/moe/moe_ops.h
+7
@@ -20,6 +20,13 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
20 torch::Tensor expert_ids,
21 torch::Tensor num_tokens_post_pad);
22
23 +void moe_lora_align_block_size(torch::Tensor topk_ids,
24 + torch::Tensor token_lora_mapping,
25 + int64_t num_experts, int64_t block_size,
26 + int64_t max_loras,
27 + torch::Tensor sorted_token_ids,
28 + torch::Tensor expert_ids,
29 + torch::Tensor num_tokens_post_pad);
30 #ifndef USE_ROCM
31 torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
32 torch::Tensor b_qweight, torch::Tensor b_scales,
csrc/moe/torch_bindings.cpp
+12
@@ -33,6 +33,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
33 m.impl("batched_moe_align_block_size", torch::kCUDA,
34 &batched_moe_align_block_size);
35
36 + // Aligning the number of tokens to be processed by each expert such
37 + // that it is divisible by the block size.
38 + m.def(
39 + "moe_lora_align_block_size(Tensor topk_ids,"
40 + " Tensor token_lora_mapping,"
41 + " int num_experts,"
42 + " int block_size, int max_loras, "
43 + " Tensor !sorted_token_ids,"
44 + " Tensor !experts_ids,"
45 + " Tensor !num_tokens_post_pad) -> () ");
46 + m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size);
47 +
48 #ifndef USE_ROCM
49 m.def(
50 "moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
tests/lora/conftest.py
+20
@@ -230,6 +230,26 @@ def tinyllama_lora_files():
230 return snapshot_download(repo_id="jashing/tinyllama-colorist-lora")
231
232
233 +@pytest.fixture(scope="session")
234 +def deepseekv2_lora_files():
235 + return snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA")
236 +
237 +
238 +@pytest.fixture(scope="session")
239 +def gptoss20b_lora_files():
240 + return snapshot_download(repo_id="LevinZheng/gpt-oss-20b-lora-adapter")
241 +
242 +
243 +@pytest.fixture(scope="session")
244 +def qwen3moe_lora_files():
245 + return snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider")
246 +
247 +
248 +@pytest.fixture(scope="session")
249 +def olmoe_lora_files():
250 + return snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider")
251 +
252 +
253 @pytest.fixture
254 def reset_default_device():
255 """
tests/lora/test_deepseekv2_tp.py new
+97
@@ -0,0 +1,97 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import vllm
5 +from vllm.lora.request import LoRARequest
6 +
7 +from ..utils import multi_gpu_test
8 +
9 +MODEL_PATH = "deepseek-ai/DeepSeek-V2-Lite-Chat"
10 +
11 +PROMPT_TEMPLATE = "<｜begin▁of▁sentence｜>You are a helpful assistant.\n\nUser: {context}\n\nAssistant:" # noqa: E501
12 +
13 +
14 +def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int):
15 + prompts = [
16 + PROMPT_TEMPLATE.format(context="Who are you?"),
17 + ]
18 + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
19 + outputs = llm.generate(
20 + prompts,
21 + sampling_params,
22 + lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
23 + )
24 + # Print the outputs.
25 + generated_texts: list[str] = []
26 + for output in outputs:
27 + prompt = output.prompt
28 + generated_text = output.outputs[0].text.strip()
29 + generated_texts.append(generated_text)
30 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
31 + # return generated_texts
32 + expected_lora_output = [
33 + "I am \u5f20\u5b50\u8c6a, an AI assistant developed by \u9648\u58eb\u680b.", # noqa: E501
34 + ]
35 + for i in range(len(expected_lora_output)):
36 + assert generated_texts[i].startswith(expected_lora_output[i])
37 +
38 +
39 +def test_deepseekv2_lora(deepseekv2_lora_files):
40 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
41 + # Otherwise, the lora-test will fail due to CUDA OOM.
42 + llm = vllm.LLM(
43 + MODEL_PATH,
44 + max_model_len=1024,
45 + enable_lora=True,
46 + max_loras=4,
47 + enforce_eager=True,
48 + trust_remote_code=True,
49 + enable_chunked_prefill=True,
50 + )
51 + generate_and_test(llm, deepseekv2_lora_files, 1)
52 +
53 +
54 +def test_deepseekv2(deepseekv2_lora_files):
55 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
56 + # Otherwise, the lora-test will fail due to CUDA OOM.
57 + llm = vllm.LLM(
58 + MODEL_PATH,
59 + max_model_len=1024,
60 + enable_lora=True,
61 + max_loras=4,
62 + enforce_eager=True,
63 + trust_remote_code=True,
64 + )
65 + generate_and_test(llm, deepseekv2_lora_files, 1)
66 +
67 +
68 +@multi_gpu_test(num_gpus=2)
69 +def test_deepseekv2_tp2(deepseekv2_lora_files):
70 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
71 + # Otherwise, the lora-test will fail due to CUDA OOM.
72 + llm = vllm.LLM(
73 + MODEL_PATH,
74 + max_model_len=1024,
75 + enable_lora=True,
76 + max_loras=4,
77 + enforce_eager=True,
78 + trust_remote_code=True,
79 + tensor_parallel_size=2,
80 + )
81 + generate_and_test(llm, deepseekv2_lora_files, 2)
82 +
83 +
84 +@multi_gpu_test(num_gpus=4)
85 +def test_deepseekv2_tp4(deepseekv2_lora_files):
86 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
87 + # Otherwise, the lora-test will fail due to CUDA OOM.
88 + llm = vllm.LLM(
89 + MODEL_PATH,
90 + max_model_len=1024,
91 + enable_lora=True,
92 + max_loras=4,
93 + enforce_eager=True,
94 + trust_remote_code=True,
95 + tensor_parallel_size=4,
96 + )
97 + generate_and_test(llm, deepseekv2_lora_files, 2)
tests/lora/test_fused_moe_lora_kernel.py new
+287
@@ -0,0 +1,287 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import random
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm import _custom_ops as ops
9 +from vllm.lora.ops.triton_ops import fused_moe_lora
10 +from vllm.platforms import current_platform
11 +
12 +
13 +@pytest.fixture(autouse=True)
14 +def reset_device(reset_default_device):
15 + pass
16 +
17 +
18 +def round_up(x, base):
19 + return ((x + base - 1) // base) * base
20 +
21 +
22 +def CEILDIV(x, y):
23 + return (x + y - 1) // y
24 +
25 +
26 +def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int):
27 + """
28 + Split `num_tokens` into `num_sequences` sequences.
29 + Each sequence randomly selects 1 LoRA index from [0, max_loras),
30 + and all tokens in that sequence are assigned this LoRA index.
31 +
32 + Args:
33 + num_tokens (int): Total number of tokens.
34 + num_sequences (int): Number of sequences to split the tokens into.
35 + max_loras (int): Total number of available LoRA modules.
36 +
37 + Returns:
38 + torch.Tensor: 1D tensor of shape [num_tokens], where each value
39 + is the LoRA index assigned to that token.
40 + """
41 + assert num_sequences > 0 and max_loras > 0
42 + assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences"
43 +
44 + # Compute token distribution per sequence (distribute remainder evenly)
45 + tokens_per_seq = num_tokens // num_sequences
46 + remainder = num_tokens % num_sequences
47 +
48 + token_lora_mapping = torch.empty(num_tokens, dtype=torch.int32)
49 +
50 + start = 0
51 + for seq_idx in range(num_sequences):
52 + # Determine the token range for this sequence
53 + end = start + tokens_per_seq + (1 if seq_idx < remainder else 0)
54 +
55 + # Randomly select one LoRA ID for this sequence
56 + lora_id = random.randint(0, max_loras - 1)
57 +
58 + # Assign the same LoRA ID to all tokens in this sequence
59 + token_lora_mapping[start:end] = lora_id
60 +
61 + start = end
62 +
63 + return token_lora_mapping
64 +
65 +
66 +def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int):
67 + """
68 + For each token, randomly select `top_k_num` distinct experts out of `num_experts`,
69 + and assign normalized random weights that sum to 1.
70 +
71 + Args:
72 + num_tokens (int): Total number of tokens.
73 + num_experts (int): Total number of available experts.
74 + top_k_num (int): Number of experts to select per token.
75 +
76 + Returns:
77 + expert_indices (torch.Tensor): shape [num_tokens, top_k_num],
78 + expert index for each token.
79 + expert_weights (torch.Tensor): shape [num_tokens, top_k_num],
80 + normalized weights (sum = 1 per row).
81 + """
82 + assert top_k_num <= num_experts, "top_k_num must be <= num_experts"
83 +
84 + # Randomly select top_k_num distinct experts for each token
85 + expert_indices = torch.empty((num_tokens, top_k_num), dtype=torch.int32)
86 + for i in range(num_tokens):
87 + # Randomly choose unique expert indices
88 + selected = torch.randperm(num_experts)[:top_k_num]
89 + expert_indices[i] = selected
90 +
91 + # Generate random weights and normalize along dim=1
92 + expert_weights = torch.rand((num_tokens, top_k_num), dtype=torch.float32)
93 + expert_weights = expert_weights / expert_weights.sum(dim=1, keepdim=True)
94 +
95 + return expert_indices, expert_weights
96 +
97 +
98 +def sample_data(
99 + num_tokens: int,
100 + num_sequences: int,
101 + max_loras: int,
102 + num_experts: int,
103 + top_k_num: int,
104 +):
105 + topk_ids, topk_weights = assign_experts_to_tokens(
106 + num_tokens, num_experts, top_k_num
107 + )
108 + token_lora_mapping = assign_loras_to_tokens(num_tokens, num_sequences, max_loras)
109 + return topk_ids, topk_weights, token_lora_mapping
110 +
111 +
112 +def use_fused_moe_lora_kernel(
113 + topk_ids,
114 + topk_weights,
115 + token_lora_mapping,
116 + max_lora_rank,
117 + top_k_num,
118 + lora_a_stacked,
119 + lora_b_stacked,
120 + hidden_states,
121 + output,
122 + max_loras,
123 + num_experts,
124 + block_size,
125 +):
126 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
127 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
128 + max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
129 +
130 + # init output tensors
131 + sorted_token_ids = torch.empty(
132 + (max_loras * max_num_tokens_padded,),
133 + dtype=torch.int32,
134 + )
135 + expert_ids = torch.empty((max_loras * max_num_m_blocks,), dtype=torch.int32)
136 + num_tokens_post_padded = torch.empty((max_loras,), dtype=torch.int32)
137 +
138 + # call kernel
139 + ops.moe_lora_align_block_size(
140 + topk_ids,
141 + token_lora_mapping,
142 + num_experts,
143 + block_size,
144 + max_loras,
145 + sorted_token_ids,
146 + expert_ids,
147 + num_tokens_post_padded,
148 + )
149 +
150 + config = {
151 + "BLOCK_SIZE_M": 16,
152 + "BLOCK_SIZE_N": 32,
153 + "BLOCK_SIZE_K": 64,
154 + "GROUP_SIZE_M": 1,
155 + }
156 +
157 + mul_routed_weight = False
158 + expert_ids = expert_ids.view(max_loras, -1)
159 + sorted_token_ids = sorted_token_ids.view(max_loras, -1)
160 +
161 + fused_moe_lora(
162 + output,
163 + hidden_states,
164 + lora_a_stacked,
165 + lora_b_stacked,
166 + topk_weights,
167 + sorted_token_ids,
168 + expert_ids,
169 + num_tokens_post_padded,
170 + max_lora_rank,
171 + top_k_num,
172 + config["BLOCK_SIZE_M"],
173 + config["BLOCK_SIZE_N"],
174 + config["BLOCK_SIZE_K"],
175 + config["GROUP_SIZE_M"],
176 + mul_routed_weight,
177 + )
178 +
179 + return output
180 +
181 +
182 +def use_torch(
183 + hidden_states,
184 + token_lora_mapping,
185 + topk_ids,
186 + lora_a_stacked,
187 + lora_b_stacked,
188 + top_k_num,
189 +):
190 + outputs = []
191 + for i in range(hidden_states.shape[0]):
192 + lora_idx = token_lora_mapping[i]
193 + expert_ids = topk_ids[i]
194 + lora_a = lora_a_stacked[0][lora_idx][expert_ids]
195 + lora_b = lora_b_stacked[0][lora_idx][expert_ids]
196 + tensors = [
197 + hidden_states[i] @ lora_a[x].T @ lora_b[x].T for x in range(top_k_num)
198 + ]
199 + outputs.append(torch.stack(tensors, dim=0))
200 + return torch.stack(outputs, dim=0)
201 +
202 +
203 +@pytest.mark.parametrize("num_tokens", [100])
204 +@pytest.mark.parametrize("top_k_num", [6, 12])
205 +@pytest.mark.parametrize("num_experts", [64])
206 +@pytest.mark.parametrize("max_loras", [4, 6, 16])
207 +@pytest.mark.parametrize("N", [1408])
208 +@pytest.mark.parametrize("K", [2048])
209 +@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
210 +@pytest.mark.parametrize("block_size", [16])
211 +def test_fused_moe_lora_kernel(
212 + num_tokens,
213 + top_k_num,
214 + num_experts,
215 + max_loras,
216 + N,
217 + K,
218 + max_lora_rank,
219 + block_size,
220 +):
221 + torch.set_default_device("cuda:0")
222 + current_platform.seed_everything(42)
223 + # the number of randomly generated sentences.
224 + num_sequences = 10
225 + # generate data
226 + topk_ids, topk_weights, token_lora_mapping = sample_data(
227 + num_tokens, num_sequences, max_loras, num_experts, top_k_num
228 + )
229 +
230 + # init lora weights
231 + lora_a_stacked = [
232 + torch.rand(
233 + (
234 + max_loras,
235 + num_experts,
236 + max_lora_rank,
237 + K,
238 + ),
239 + dtype=torch.bfloat16,
240 + )
241 + ]
242 + lora_b_stacked = [
243 + torch.rand(
244 + (
245 + max_loras,
246 + num_experts,
247 + N,
248 + max_lora_rank,
249 + ),
250 + dtype=torch.bfloat16,
251 + )
252 + ]
253 + hidden_states = torch.rand(
254 + (
255 + num_tokens,
256 + K,
257 + ),
258 + dtype=torch.bfloat16,
259 + )
260 +
261 + # fused_moe_lora_kernel output
262 + output = torch.zeros((num_tokens, top_k_num, N), dtype=torch.bfloat16)
263 + use_fused_moe_lora_kernel(
264 + topk_ids,
265 + topk_weights,
266 + token_lora_mapping,
267 + max_lora_rank,
268 + top_k_num,
269 + lora_a_stacked,
270 + lora_b_stacked,
271 + hidden_states,
272 + output,
273 + max_loras,
274 + num_experts,
275 + block_size,
276 + )
277 + # pytorch output
278 + output2 = use_torch(
279 + hidden_states,
280 + token_lora_mapping,
281 + topk_ids,
282 + lora_a_stacked,
283 + lora_b_stacked,
284 + top_k_num,
285 + )
286 +
287 + torch.testing.assert_close(output, output2, atol=1e-1, rtol=1e-1)
tests/lora/test_gptoss.py new
+52
@@ -0,0 +1,52 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import vllm
5 +from vllm.lora.request import LoRARequest
6 +
7 +MODEL_PATH = "openai/gpt-oss-20b"
8 +
9 +PROMPT_TEMPLATE = "<｜begin▁of▁sentence｜>You are a helpful assistant.\n\nUser: {context}\n\nAssistant:" # noqa: E501
10 +
11 +
12 +def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
13 + prompts = [
14 + PROMPT_TEMPLATE.format(context="Who are you?"),
15 + ]
16 + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
17 + outputs = llm.generate(
18 + prompts,
19 + sampling_params,
20 + lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
21 + )
22 + # Print the outputs.
23 + generated_texts: list[str] = []
24 + for output in outputs:
25 + prompt = output.prompt
26 + generated_text = output.outputs[0].text.strip()
27 + generated_texts.append(generated_text)
28 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
29 + return generated_texts
30 +
31 +
32 +# FIXME: Load gpt-oss adapter
33 +def test_gptoss20b_lora(gptoss20b_lora_files):
34 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
35 + # Otherwise, the lora-test will fail due to CUDA OOM.
36 + llm = vllm.LLM(
37 + MODEL_PATH,
38 + enable_lora=True,
39 + max_loras=1,
40 + trust_remote_code=True,
41 + )
42 +
43 + expected_lora_output = [
44 + "I am an AI language model developed by OpenAI. "
45 + "I am here to help you with any questions or "
46 + "tasks you may have."
47 + ]
48 +
49 + output1 = do_sample(llm, gptoss20b_lora_files, lora_id=1)
50 + print(output1)
51 + for i in range(len(expected_lora_output)):
52 + assert output1[i].startswith(expected_lora_output[i])
tests/lora/test_moe_lora_align_sum.py new
+90
@@ -0,0 +1,90 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import random
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm import _custom_ops as ops
9 +
10 +
11 +def round_up(x, base):
12 + return ((x + base - 1) // base) * base
13 +
14 +
15 +def CEILDIV(x, y):
16 + return (x + y - 1) // y
17 +
18 +
19 +def sample_data(num_experts, max_loras, num_tokens, topk_num):
20 + topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
21 + token_lora_mapping = torch.zeros((num_tokens,), dtype=torch.int32)
22 +
23 + for i in range(num_tokens):
24 + pool = list(range(num_experts))
25 + random.shuffle(pool)
26 + for j in range(topk_num):
27 + topk_ids[i, j] = pool[j]
28 + token_lora_mapping[i] = random.randint(0, max_loras - 1)
29 +
30 + return topk_ids.to("cuda"), token_lora_mapping.to("cuda")
31 +
32 +
33 +@pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096]) # 81920
34 +@pytest.mark.parametrize("topk_num", [6])
35 +@pytest.mark.parametrize("num_experts", [64, 128])
36 +@pytest.mark.parametrize("max_loras", [2, 32])
37 +@pytest.mark.parametrize("block_size", [16])
38 +def test_moe_lora_align_block_size(
39 + num_tokens, topk_num, num_experts, max_loras, block_size
40 +):
41 + # sample data
42 + random.seed(1)
43 + topk_ids, token_lora_mapping = sample_data(
44 + num_experts, max_loras, num_tokens, topk_num
45 + )
46 +
47 + # compute paddings
48 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
49 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
50 + max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
51 +
52 + # init output tensors
53 + sorted_token_ids = torch.full(
54 + (max_loras * max_num_tokens_padded,),
55 + topk_ids.numel(),
56 + dtype=torch.int32,
57 + device="cuda",
58 + )
59 + expert_ids = torch.full(
60 + (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
61 + )
62 + num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
63 +
64 + # call kernel
65 + ops.moe_lora_align_block_size(
66 + topk_ids,
67 + token_lora_mapping,
68 + num_experts,
69 + block_size,
70 + max_loras,
71 + sorted_token_ids,
72 + expert_ids,
73 + num_tokens_post_pad,
74 + )
75 +
76 + # verify values
77 + expert_ids = expert_ids.view(max_loras, -1)
78 + sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size)
79 +
80 + for lora_idx in range(max_loras):
81 + for token_idx in range(sorted_token_ids.size(1)):
82 + block = sorted_token_ids[lora_idx][token_idx]
83 + indices = block[block != topk_ids.numel()]
84 + if indices.numel() > 0:
85 + expert_id = expert_ids[lora_idx][token_idx]
86 + assert torch.all(topk_ids.view(-1)[indices] == expert_id)
87 +
88 +
89 +if __name__ == "__main__":
90 + pytest.main([__file__])
tests/lora/test_olmoe_tp.py new
+109
@@ -0,0 +1,109 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import vllm
5 +from vllm.lora.request import LoRARequest
6 +
7 +from ..utils import multi_gpu_test
8 +
9 +MODEL_PATH = "allenai/OLMoE-1B-7B-0125-Instruct"
10 +
11 +PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
12 +"
13 +##Instruction:
14 +candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
15 +Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
16 +The People_ID of candidate is the foreign key of People_ID of people.
17 +
18 +
19 +###Input:
20 +{context}
21 +
22 +###Response:""" # noqa: E501
23 +
24 +EXPECTED_LORA_OUTPUT = [
25 + "SELECT count(*) FROM candidate",
26 + "SELECT count(*) FROM candidate",
27 + "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
28 + "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
29 +]
30 +
31 +
32 +def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
33 + prompts = [
34 + PROMPT_TEMPLATE.format(context="How many candidates are there?"),
35 + PROMPT_TEMPLATE.format(context="Count the number of candidates."),
36 + PROMPT_TEMPLATE.format(
37 + context="Which poll resource provided the most number of candidate information?" # noqa: E501
38 + ),
39 + PROMPT_TEMPLATE.format(
40 + context="Return the poll resource associated with the most candidates."
41 + ),
42 + ]
43 + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
44 + outputs = llm.generate(
45 + prompts,
46 + sampling_params,
47 + lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
48 + )
49 + # Print the outputs.
50 + generated_texts: list[str] = []
51 + for output in outputs:
52 + prompt = output.prompt
53 + generated_text = output.outputs[0].text.strip()
54 + generated_texts.append(generated_text)
55 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
56 +
57 + for i in range(len(EXPECTED_LORA_OUTPUT)):
58 + assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
59 +
60 +
61 +def test_olmoe_lora(olmoe_lora_files):
62 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
63 + # Otherwise, the lora-test will fail due to CUDA OOM.
64 + llm = vllm.LLM(
65 + MODEL_PATH,
66 + max_model_len=1024,
67 + enable_lora=True,
68 + max_loras=4,
69 + enforce_eager=True,
70 + trust_remote_code=True,
71 + enable_chunked_prefill=True,
72 + )
73 +
74 + generate_and_test(llm, olmoe_lora_files, lora_id=1)
75 + generate_and_test(llm, olmoe_lora_files, lora_id=2)
76 +
77 +
78 +@multi_gpu_test(num_gpus=2)
79 +def test_olmoe_lora_tp2(olmoe_lora_files):
80 + llm = vllm.LLM(
81 + MODEL_PATH,
82 + max_model_len=1024,
83 + enable_lora=True,
84 + max_loras=4,
85 + enforce_eager=True,
86 + trust_remote_code=True,
87 + enable_chunked_prefill=True,
88 + tensor_parallel_size=2,
89 + )
90 +
91 + generate_and_test(llm, olmoe_lora_files, lora_id=1)
92 + generate_and_test(llm, olmoe_lora_files, lora_id=2)
93 +
94 +
95 +@multi_gpu_test(num_gpus=4)
96 +def test_olmoe_lora_tp4(olmoe_lora_files):
97 + llm = vllm.LLM(
98 + MODEL_PATH,
99 + max_model_len=1024,
100 + enable_lora=True,
101 + max_loras=4,
102 + enforce_eager=True,
103 + trust_remote_code=True,
104 + enable_chunked_prefill=True,
105 + tensor_parallel_size=4,
106 + )
107 +
108 + generate_and_test(llm, olmoe_lora_files, lora_id=1)
109 + generate_and_test(llm, olmoe_lora_files, lora_id=2)
tests/lora/test_qwen3moe_tp.py new
+111
@@ -0,0 +1,111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import vllm
5 +from vllm.lora.request import LoRARequest
6 +
7 +from ..utils import multi_gpu_test
8 +
9 +MODEL_PATH = "Qwen/Qwen3-30B-A3B"
10 +
11 +PROMPT_TEMPLATE = """<|im_start|>user
12 +I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
13 +"
14 +##Instruction:
15 +candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
16 +Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
17 +The People_ID of candidate is the foreign key of People_ID of people.
18 +
19 +
20 +###Input:
21 +{context}
22 +
23 +###Response:<|im_end|>
24 +<|im_start|>assistant""" # noqa: E501
25 +
26 +EXPECTED_LORA_OUTPUT = [
27 + "<think>\n\n</think>\n\nSELECT count(*) FROM candidate",
28 + "<think>\n\n</think>\n\nSELECT count(*) FROM candidate",
29 + "<think>\n\n</think>\n\nSELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
30 + "<think>\n\n</think>\n\nSELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
31 +]
32 +
33 +
34 +def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
35 + prompts = [
36 + PROMPT_TEMPLATE.format(context="How many candidates are there?"),
37 + PROMPT_TEMPLATE.format(context="Count the number of candidates."),
38 + PROMPT_TEMPLATE.format(
39 + context="Which poll resource provided the most number of candidate information?" # noqa: E501
40 + ),
41 + PROMPT_TEMPLATE.format(
42 + context="Return the poll resource associated with the most candidates."
43 + ),
44 + ]
45 + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
46 + outputs = llm.generate(
47 + prompts,
48 + sampling_params,
49 + lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
50 + )
51 + # Print the outputs.
52 + generated_texts: list[str] = []
53 + for output in outputs:
54 + prompt = output.prompt
55 + generated_text = output.outputs[0].text.strip()
56 + generated_texts.append(generated_text)
57 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
58 +
59 + for i in range(len(EXPECTED_LORA_OUTPUT)):
60 + assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
61 +
62 +
63 +def test_qwen3moe_lora(qwen3moe_lora_files):
64 + # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
65 + # Otherwise, the lora-test will fail due to CUDA OOM.
66 + llm = vllm.LLM(
67 + MODEL_PATH,
68 + max_model_len=1024,
69 + enable_lora=True,
70 + max_loras=4,
71 + enforce_eager=True,
72 + trust_remote_code=True,
73 + enable_chunked_prefill=True,
74 + )
75 +
76 + generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
77 + generate_and_test(llm, qwen3moe_lora_files, lora_id=2)
78 +
79 +
80 +@multi_gpu_test(num_gpus=2)
81 +def test_qwen3moe_lora_tp2(qwen3moe_lora_files):
82 + llm = vllm.LLM(
83 + MODEL_PATH,
84 + max_model_len=1024,
85 + enable_lora=True,
86 + max_loras=4,
87 + enforce_eager=True,
88 + trust_remote_code=True,
89 + enable_chunked_prefill=True,
90 + tensor_parallel_size=2,
91 + )
92 +
93 + generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
94 + generate_and_test(llm, qwen3moe_lora_files, lora_id=2)
95 +
96 +
97 +@multi_gpu_test(num_gpus=4)
98 +def test_qwen3moe_lora_tp4(qwen3moe_lora_files):
99 + llm = vllm.LLM(
100 + MODEL_PATH,
101 + max_model_len=1024,
102 + enable_lora=True,
103 + max_loras=4,
104 + enforce_eager=True,
105 + trust_remote_code=True,
106 + enable_chunked_prefill=True,
107 + tensor_parallel_size=4,
108 + )
109 +
110 + generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
111 + generate_and_test(llm, qwen3moe_lora_files, lora_id=2)
vllm/_custom_ops.py
+22
@@ -1795,6 +1795,28 @@ def batched_moe_align_block_size(
1795 )
1796
1797
1798 +def moe_lora_align_block_size(
1799 + topk_ids: torch.Tensor,
1800 + token_lora_mapping: torch.Tensor,
1801 + num_experts: int,
1802 + block_size: int,
1803 + max_loras: int,
1804 + sorted_token_ids: torch.Tensor,
1805 + experts_ids: torch.Tensor,
1806 + num_tokens_post_pad: torch.Tensor,
1807 +) -> None:
1808 + torch.ops._moe_C.moe_lora_align_block_size(
1809 + topk_ids,
1810 + token_lora_mapping,
1811 + num_experts,
1812 + block_size,
1813 + max_loras,
1814 + sorted_token_ids,
1815 + experts_ids,
1816 + num_tokens_post_pad,
1817 + )
1818 +
1819 +
1820 def moe_wna16_gemm(
1821 input: torch.Tensor,
1822 output: torch.Tensor,
vllm/lora/layers/__init__.py
+2
@@ -11,6 +11,7 @@ from vllm.lora.layers.column_parallel_linear import (
11 QKVParallelLinearWithLoRA,
12 QKVParallelLinearWithShardedLoRA,
13 )
14 +from vllm.lora.layers.fused_moe import FusedMoEWithLoRA
15 from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
16 from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA
17 from vllm.lora.layers.row_parallel_linear import (
@@ -36,4 +37,5 @@ __all__ = [
37 "RowParallelLinearWithShardedLoRA",
38 "ReplicatedLinearWithLoRA",
39 "LoRAMapping",
40 + "FusedMoEWithLoRA",
41 ]
vllm/lora/layers/fused_moe.py new
+410
@@ -0,0 +1,410 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 +
5 +import torch
6 +import torch.nn as nn
7 +from transformers import PretrainedConfig
8 +
9 +from vllm import envs
10 +from vllm.config.lora import LoRAConfig
11 +from vllm.distributed.parallel_state import (
12 + get_tensor_model_parallel_rank,
13 + get_tensor_model_parallel_world_size,
14 +)
15 +from vllm.lora.layers.base import BaseLayerWithLoRA
16 +from vllm.model_executor.layers.fused_moe import FusedMoE
17 +from vllm.model_executor.layers.fused_moe.config import (
18 + FUSED_MOE_UNQUANTIZED_CONFIG,
19 + _get_config_dtype_str,
20 + mxfp4_w4a16_moe_quant_config,
21 +)
22 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
23 + modular_marlin_fused_moe,
24 +)
25 +from vllm.model_executor.layers.fused_moe.fused_moe import (
26 + modular_triton_fused_moe,
27 + try_get_optimal_moe_config,
28 +)
29 +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config
30 +
31 +
32 +class FusedMoEWithLoRA(BaseLayerWithLoRA):
33 + def __init__(self, base_layer: FusedMoE) -> None:
34 + super().__init__()
35 + self.base_layer = base_layer
36 + self.tp_size = get_tensor_model_parallel_world_size()
37 + self.tp_rank = get_tensor_model_parallel_rank()
38 + self.device = base_layer.w2_weight.device
39 + self._inject_lora_into_fused_moe()
40 +
41 + def _inject_lora_into_fused_moe(self):
42 + moe_state_dict = {}
43 + top_k = self.base_layer.top_k
44 +
45 + if self.base_layer.quant_config is None:
46 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
47 + elif not isinstance(self.base_layer.quant_config, Mxfp4Config):
48 + quant_config = self.base_layer.quant_config
49 + else:
50 + quant_config = mxfp4_w4a16_moe_quant_config(
51 + w1_bias=self.base_layer.w13_bias,
52 + w2_bias=self.base_layer.w2_bias,
53 + w1_scale=self.base_layer.w13_weight_scale,
54 + w2_scale=self.base_layer.w2_weight_scale,
55 + )
56 +
57 + m_fused_moe_fn = (
58 + modular_triton_fused_moe(
59 + quant_config, shared_experts=self.base_layer.shared_experts
60 + )
61 + if not quant_config.use_mxfp4_w4a16
62 + else modular_marlin_fused_moe(
63 + quant_config, shared_experts=self.base_layer.shared_experts
64 + )
65 + )
66 +
67 + def fwd_decorator(layer, func):
68 + def wrapper(*args, **kwargs):
69 + moe_state_dict["hidden_states"] = kwargs["hidden_states"]
70 + moe_state_dict["topk_ids"] = kwargs["topk_ids"]
71 + moe_state_dict["topk_weights"] = kwargs["topk_weights"]
72 + moe_state_dict["global_num_experts"] = kwargs["global_num_experts"]
73 + moe_state_dict["expert_map"] = kwargs["expert_map"]
74 + moe_state_dict["apply_router_weight_on_input"] = kwargs[
75 + "apply_router_weight_on_input"
76 + ]
77 + moe_state_dict["max_loras"] = layer.w1_lora_a_stacked.shape[0]
78 + result = func(*args, **kwargs)
79 + return result
80 +
81 + return wrapper
82 +
83 + def act_decorator(layer, func):
84 + def wrapper(*args, **kwargs):
85 + _, output, input = args
86 +
87 + hidden_states = moe_state_dict["hidden_states"]
88 + topk_weights = moe_state_dict["topk_weights"]
89 + curr_topk_ids = moe_state_dict["topk_ids"]
90 + global_num_experts = moe_state_dict["global_num_experts"]
91 + expert_map = moe_state_dict["expert_map"]
92 + max_loras = moe_state_dict["max_loras"]
93 +
94 + config_dtype = _get_config_dtype_str(
95 + dtype=hidden_states.dtype,
96 + use_fp8_w8a8=False,
97 + use_int8_w8a16=False,
98 + use_int4_w4a16=False,
99 + )
100 + CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
101 + num_tokens = hidden_states.size(0)
102 + M = min(num_tokens, CHUNK_SIZE)
103 +
104 + get_config_func = functools.partial(
105 + try_get_optimal_moe_config,
106 + layer.w13_weight.size(),
107 + layer.w2_weight.size(),
108 + top_k,
109 + config_dtype,
110 + block_shape=layer.quant_method.moe_quant_config.block_shape,
111 + )
112 +
113 + config = get_config_func(M)
114 + (
115 + sorted_token_ids_lora,
116 + expert_ids_lora,
117 + num_tokens_post_padded_lora,
118 + ) = self.punica_wrapper.moe_lora_align_block_size(
119 + curr_topk_ids,
120 + num_tokens,
121 + config["BLOCK_SIZE_M"],
122 + global_num_experts,
123 + max_loras,
124 + expert_map,
125 + )
126 +
127 + moe_state_dict["sorted_token_ids_lora"] = sorted_token_ids_lora
128 + moe_state_dict["expert_ids_lora"] = expert_ids_lora
129 + moe_state_dict["num_tokens_post_padded_lora"] = (
130 + num_tokens_post_padded_lora
131 + )
132 +
133 + w13_lora_a_stacked = [self.w1_lora_a_stacked, self.w3_lora_a_stacked]
134 + w13_lora_b_stacked = [self.w1_lora_b_stacked, self.w3_lora_b_stacked]
135 + max_lora_rank = self.w1_lora_a_stacked.shape[-2]
136 + expert_ids_lora = expert_ids_lora.view(max_loras, -1)
137 + sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
138 +
139 + self.punica_wrapper.add_lora_fused_moe(
140 + input.view(-1, top_k, input.shape[-1]),
141 + hidden_states,
142 + w13_lora_a_stacked,
143 + w13_lora_b_stacked,
144 + topk_weights,
145 + sorted_token_ids_lora,
146 + expert_ids_lora,
147 + num_tokens_post_padded_lora,
148 + max_lora_rank,
149 + top_k,
150 + config,
151 + )
152 +
153 + result = func(*args, **kwargs)
154 +
155 + moe_state_dict["intermediate_cache2"] = output
156 + return result
157 +
158 + return wrapper
159 +
160 + def moe_sum_decorator(layer, func):
161 + def wrapper(*args, **kwargs):
162 + hidden_states = moe_state_dict["hidden_states"]
163 + topk_weights = moe_state_dict["topk_weights"]
164 + max_loras = moe_state_dict["max_loras"]
165 +
166 + config_dtype = _get_config_dtype_str(
167 + dtype=hidden_states.dtype,
168 + use_fp8_w8a8=False,
169 + use_int8_w8a16=False,
170 + use_int4_w4a16=False,
171 + )
172 + CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
173 + num_tokens = hidden_states.size(0)
174 + M = min(num_tokens, CHUNK_SIZE)
175 +
176 + get_config_func = functools.partial(
177 + try_get_optimal_moe_config,
178 + layer.w13_weight.size(),
179 + layer.w2_weight.size(),
180 + top_k,
181 + config_dtype,
182 + block_shape=layer.quant_method.moe_quant_config.block_shape,
183 + )
184 +
185 + config = get_config_func(M)
186 +
187 + sorted_token_ids_lora = moe_state_dict["sorted_token_ids_lora"]
188 + expert_ids_lora = moe_state_dict["expert_ids_lora"]
189 + num_tokens_post_padded_lora = moe_state_dict[
190 + "num_tokens_post_padded_lora"
191 + ]
192 +
193 + expert_ids_lora = expert_ids_lora.view(max_loras, -1)
194 + sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
195 + intermediate_cache2 = moe_state_dict["intermediate_cache2"]
196 + intermediate_cache3 = args[0]
197 + max_lora_rank = self.w1_lora_a_stacked.shape[-2]
198 + self.punica_wrapper.add_lora_fused_moe(
199 + intermediate_cache3,
200 + intermediate_cache2,
201 + [self.w2_lora_a_stacked],
202 + [self.w2_lora_b_stacked],
203 + topk_weights,
204 + sorted_token_ids_lora,
205 + expert_ids_lora,
206 + num_tokens_post_padded_lora,
207 + max_lora_rank,
208 + top_k,
209 + config,
210 + True,
211 + )
212 +
213 + result = func(*args, **kwargs)
214 + return result
215 +
216 + return wrapper
217 +
218 + fused_experts = m_fused_moe_fn.fused_experts
219 +
220 + m_fused_moe_fn.forward = fwd_decorator(self.base_layer, m_fused_moe_fn.forward)
221 + fused_experts.activation = act_decorator(
222 + self.base_layer, fused_experts.activation
223 + )
224 + fused_experts.moe_sum = moe_sum_decorator(
225 + self.base_layer, fused_experts.moe_sum
226 + )
227 +
228 + self.base_layer.quant_method.old_fused_experts = (
229 + self.base_layer.quant_method.fused_experts
230 + )
231 + self.base_layer.quant_method.fused_experts = m_fused_moe_fn
232 +
233 + def create_lora_weights(
234 + self,
235 + max_loras: int,
236 + lora_config: LoRAConfig,
237 + model_config: PretrainedConfig | None = None,
238 + ) -> None:
239 + """Initializes lora matrices."""
240 +
241 + assert not self.base_layer.use_ep, (
242 + "EP support for Fused MoE LoRA is not implemented yet."
243 + )
244 +
245 + self.w1_lora_a_stacked = torch.zeros(
246 + (
247 + max_loras,
248 + self.base_layer.global_num_experts,
249 + lora_config.max_lora_rank,
250 + self.base_layer.hidden_size,
251 + ),
252 + dtype=lora_config.lora_dtype,
253 + device=self.device,
254 + )
255 + self.w1_lora_b_stacked = torch.zeros(
256 + (
257 + max_loras,
258 + self.base_layer.global_num_experts,
259 + self.base_layer.intermediate_size_per_partition,
260 + lora_config.max_lora_rank,
261 + ),
262 + dtype=lora_config.lora_dtype,
263 + device=self.device,
264 + )
265 +
266 + self.w2_lora_a_stacked = torch.zeros(
267 + (
268 + max_loras,
269 + self.base_layer.global_num_experts,
270 + lora_config.max_lora_rank,
271 + self.base_layer.intermediate_size_per_partition,
272 + ),
273 + dtype=lora_config.lora_dtype,
274 + device=self.device,
275 + )
276 + self.w2_lora_b_stacked = torch.zeros(
277 + (
278 + max_loras,
279 + self.base_layer.global_num_experts,
280 + self.base_layer.hidden_size,
281 + lora_config.max_lora_rank,
282 + ),
283 + dtype=lora_config.lora_dtype,
284 + device=self.device,
285 + )
286 +
287 + self.w3_lora_a_stacked = torch.zeros(
288 + (
289 + max_loras,
290 + self.base_layer.global_num_experts,
291 + lora_config.max_lora_rank,
292 + self.base_layer.hidden_size,
293 + ),
294 + dtype=lora_config.lora_dtype,
295 + device=self.device,
296 + )
297 + self.w3_lora_b_stacked = torch.zeros(
298 + (
299 + max_loras,
300 + self.base_layer.global_num_experts,
301 + self.base_layer.intermediate_size_per_partition,
302 + lora_config.max_lora_rank,
303 + ),
304 + dtype=lora_config.lora_dtype,
305 + device=self.device,
306 + )
307 +
308 + self.base_layer.w1_lora_a_stacked = self.w1_lora_a_stacked
309 + self.base_layer.w1_lora_b_stacked = self.w1_lora_b_stacked
310 + self.base_layer.w2_lora_a_stacked = self.w2_lora_a_stacked
311 + self.base_layer.w2_lora_b_stacked = self.w2_lora_b_stacked
312 + self.base_layer.w3_lora_a_stacked = self.w3_lora_a_stacked
313 + self.base_layer.w3_lora_b_stacked = self.w3_lora_b_stacked
314 + # They will be used by 'LoRALayerWeights.create_dummy_lora_weights'
315 + # to create a dummy LoRA weights.
316 + self.lora_a_stacked = []
317 + self.lora_b_stacked = []
318 + for lora_id in range(max_loras):
319 + for experts_id in range(self.base_layer.global_num_experts):
320 + # gate_proj,down_proj,up_proj
321 + self.lora_a_stacked.append(self.w1_lora_a_stacked[lora_id][experts_id])
322 + self.lora_a_stacked.append(self.w2_lora_a_stacked[lora_id][experts_id])
323 + self.lora_a_stacked.append(self.w3_lora_a_stacked[lora_id][experts_id])
324 +
325 + self.lora_b_stacked.append(self.w1_lora_b_stacked[lora_id][experts_id])
326 + self.lora_b_stacked.append(self.w2_lora_b_stacked[lora_id][experts_id])
327 + self.lora_b_stacked.append(self.w3_lora_b_stacked[lora_id][experts_id])
328 +
329 + def reset_lora(self, index: int):
330 + """Resets the lora weights at index back to 0."""
331 + self.w1_lora_a_stacked[index] = 0
332 + self.w1_lora_b_stacked[index] = 0
333 + self.w3_lora_a_stacked[index] = 0
334 + self.w3_lora_b_stacked[index] = 0
335 + self.w2_lora_a_stacked[index] = 0
336 + self.w2_lora_b_stacked[index] = 0
337 +
338 + def set_lora(
339 + self,
340 + index: int,
341 + lora_a: torch.Tensor,
342 + lora_b: torch.Tensor,
343 + embeddings_tensor: torch.Tensor | None,
344 + bias: torch.Tensor | None = None,
345 + ):
346 + """Overwrites lora tensors at index."""
347 + for eid in range(len(lora_a) // 3):
348 + w1_lora_a = lora_a[eid * 3]
349 + w2_lora_a = lora_a[eid * 3 + 1]
350 + w3_lora_a = lora_a[eid * 3 + 2]
351 + w1_lora_b = lora_b[eid * 3]
352 + w2_lora_b = lora_b[eid * 3 + 1]
353 + w3_lora_b = lora_b[eid * 3 + 2]
354 +
355 + if self.tp_size > 1:
356 + shard_size = self.base_layer.intermediate_size_per_partition
357 + start_idx = self.tp_rank * shard_size
358 + end_idx = (self.tp_rank + 1) * shard_size
359 +
360 + w1_lora_b = w1_lora_b[start_idx:end_idx, :]
361 + w3_lora_b = w3_lora_b[start_idx:end_idx, :]
362 + w2_lora_a = w2_lora_a[:, start_idx:end_idx]
363 +
364 + self.w1_lora_a_stacked[
365 + index, eid, : w1_lora_a.shape[0], : w1_lora_a.shape[1]
366 + ].copy_(w1_lora_a, non_blocking=True)
367 +
368 + self.w3_lora_a_stacked[
369 + index, eid, : w3_lora_a.shape[0], : w3_lora_a.shape[1]
370 + ].copy_(w3_lora_a, non_blocking=True)
371 +
372 + self.w2_lora_b_stacked[
373 + index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
374 + ].copy_(w2_lora_b, non_blocking=True)
375 +
376 + self.w1_lora_b_stacked[
377 + index, eid, : w1_lora_b.shape[0], : w1_lora_b.shape[1]
378 + ].copy_(w1_lora_b, non_blocking=True)
379 + self.w3_lora_b_stacked[
380 + index, eid, : w3_lora_b.shape[0], : w3_lora_b.shape[1]
381 + ].copy_(w3_lora_b, non_blocking=True)
382 + self.w2_lora_a_stacked[
383 + index, eid, : w2_lora_a.shape[0], : w2_lora_a.shape[1]
384 + ].copy_(w2_lora_a, non_blocking=True)
385 +
386 + @classmethod
387 + def can_replace_layer(
388 + cls,
389 + source_layer: nn.Module,
390 + lora_config: LoRAConfig,
391 + packed_modules_list: list,
392 + model_config: PretrainedConfig | None,
393 + ) -> bool:
394 + """Returns True if the layer can be replaced by this LoRA layer."""
395 + # return type(source_layer) is FusedMoE
396 + return isinstance(source_layer, FusedMoE)
397 +
398 + def forward(self, *args, **kwargs):
399 + return self.base_layer.forward(*args, **kwargs)
400 +
401 + def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs):
402 + return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs)
403 +
404 + @property
405 + def _shared_experts(self):
406 + return self.base_layer._shared_experts
407 +
408 + @property
409 + def quant_method(self):
410 + return self.base_layer.quant_method
vllm/lora/models.py
+64 -20
@@ -13,7 +13,7 @@ from torch import nn
13
14 from vllm.config.lora import LoRAConfig
15 from vllm.logger import init_logger
16 -from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping
16 +from vllm.lora.layers import BaseLayerWithLoRA, FusedMoEWithLoRA, LoRAMapping
17 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
18 from vllm.lora.peft_helper import PEFTHelper
19 from vllm.lora.punica_wrapper import get_punica_wrapper
@@ -23,15 +23,14 @@ from vllm.lora.utils import (
23 get_supported_lora_modules,
24 is_regex_target_modules,
25 parse_fine_tuned_lora_name,
26 + process_packed_modules_mapping,
27 replace_submodule,
28 )
28 -from vllm.model_executor.layers.fused_moe import FusedMoE
29 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
30 from vllm.model_executor.models import SupportsLoRA, supports_multimodal
31 from vllm.model_executor.models.interfaces import is_pooling_model
32 from vllm.model_executor.models.module_mapping import MultiModelKeys
33 from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper
34 -from vllm.model_executor.utils import get_packed_modules_mapping
34 from vllm.utils import is_pin_memory_available
35 from vllm.utils.cache import LRUCache
36
@@ -60,18 +59,6 @@ def get_lora_id():
59 return _GLOBAL_LORA_ID
60
61
63 -def is_moe_model(model: nn.Module) -> bool:
64 - """Checks if the model contains FusedMoE layers and warns the user."""
65 - if any(isinstance(module, FusedMoE) for module in model.modules()):
66 - logger.warning_once(
67 - "For MoE models, vLLM currently does not support fused MoE LoRA "
68 - "inference. Please ensure that the loaded LoRA model does not "
69 - "contain expert weights."
70 - )
71 - return True
72 - return False
73 -
74 -
62 class LoRAModel:
63 """A LoRA fine-tuned model."""
64
@@ -229,9 +216,19 @@ class LoRAModel:
216 def check_unexpected_modules(modules: dict):
217 for lora_module in modules.keys(): # noqa
218 module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
232 - part_name = module_name.split(".")[-1]
233 - if part_name not in expected_lora_modules:
219 + # Handle FSDP file format where experts.base_layer is the
220 + # gate_up_proj and experts is the down_proj
221 + if "base_layer" in lora_module:
222 + continue
223 + # Case for expert lora weights
224 + if ".experts" in module_name:
225 + if not any(
226 + module_name.endswith(ele) for ele in expected_lora_modules
227 + ):
228 + unexpected_modules.append(module_name)
229 + elif module_name.split(".")[-1] not in expected_lora_modules:
230 unexpected_modules.append(module_name)
231 +
232 if unexpected_modules:
233 raise ValueError(
234 f"While loading {lora_dir}, expected"
@@ -371,7 +368,7 @@ class LoRAModelManager:
368 assert self.supported_lora_modules, "No supported LoRA modules found in"
369 f" {self.model.__class__.__name__}."
370
374 - self.packed_modules_mapping = get_packed_modules_mapping(self.model)
371 + self.packed_modules_mapping = process_packed_modules_mapping(self.model)
372 # Used to indicate whether the model is a multimodal model
373 self.supports_mm: bool = (
374 supports_multimodal(self.model)
@@ -380,7 +377,6 @@ class LoRAModelManager:
377 and hasattr(self.model, "get_mm_mapping")
378 )
379 self.is_pooling_model = is_pooling_model(self.model)
383 - self.is_moe_model = is_moe_model(self.model)
380 self.packed_modules: dict[str, list[str]] = {}
381 self.modules: dict[str, BaseLayerWithLoRA] = {}
382 # Dict instead of a set for compatibility with LRUCache.
@@ -431,6 +427,50 @@ class LoRAModelManager:
427 module_lora = self._get_lora_layer_weights(lora_model, module_name)
428 if module_lora:
429 module_lora.optimize()
430 + # Note (gnovack) - If MOE lora weights are not split into
431 + # num_experts chunks, we split them here
432 + if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(
433 + module_lora.lora_a
434 + ):
435 + # Handle FSDP file format where experts.base_layer is the
436 + # gate_up_proj and experts is the down_proj
437 + gate_up_proj_lora = self._get_lora_layer_weights(
438 + lora_model, module_name + ".base_layer"
439 + )
440 +
441 + assert gate_up_proj_lora is not None
442 + assert module_lora is not None
443 +
444 + down_proj_lora = module_lora
445 + num_experts = module_lora.lora_a.shape[0] // module_lora.rank
446 +
447 + gate_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
448 + up_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
449 +
450 + gate_proj_b = gate_up_proj_lora.lora_b[::2, ...].chunk(
451 + num_experts, dim=-1
452 + )
453 + up_proj_b = gate_up_proj_lora.lora_b[1::2, ...].chunk(
454 + num_experts, dim=-1
455 + )
456 +
457 + down_proj_a = down_proj_lora.lora_a.chunk(num_experts, dim=0)
458 + down_proj_b = down_proj_lora.lora_b.chunk(num_experts, dim=-1)
459 +
460 + lora_a = []
461 + lora_b = []
462 + for i in range(num_experts):
463 + lora_a.append(gate_proj_a[i])
464 + lora_a.append(down_proj_a[i])
465 + lora_a.append(up_proj_a[i])
466 +
467 + lora_b.append(gate_proj_b[i])
468 + lora_b.append(down_proj_b[i])
469 + lora_b.append(up_proj_b[i])
470 +
471 + module_lora.lora_a = lora_a
472 + module_lora.lora_b = lora_b
473 +
474 module.set_lora(
475 index,
476 module_lora.lora_a,
@@ -486,6 +526,7 @@ class LoRAModelManager:
526 for module_name, module in self.model.named_modules(remove_duplicate=False):
527 if isinstance(module, PPMissingLayer):
528 continue
529 +
530 if not self._match_target_modules(module_name):
531 continue
532 # A temporary approach for multimodal models to support LoRA
@@ -549,7 +590,10 @@ class LoRAModelManager:
590 new_module.set_mapping(self.punica_wrapper)
591
592 def register_module(self, module_name: str, module: "BaseLayerWithLoRA"):
552 - assert isinstance(module, BaseLayerWithLoRA)
593 + assert isinstance(module, BaseLayerWithLoRA), (
594 + f"Module {module_name} must be a BaseLayerWithLoRA instance,"
595 + )
596 + f" got {type(module)}"
597 self.modules[module_name] = module
598
599 def create_dummy_lora(
vllm/lora/ops/triton_ops/__init__.py
+2
@@ -1,6 +1,7 @@
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
5 from vllm.lora.ops.triton_ops.lora_expand_op import lora_expand
6 from vllm.lora.ops.triton_ops.lora_kernel_metadata import LoRAKernelMeta
7 from vllm.lora.ops.triton_ops.lora_shrink_op import lora_shrink
@@ -9,4 +10,5 @@ __all__ = [
10 "lora_expand",
11 "lora_shrink",
12 "LoRAKernelMeta",
13 + "fused_moe_lora",
14 ]
vllm/lora/ops/triton_ops/fused_moe_lora_op.py new
+350
@@ -0,0 +1,350 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +import triton
6 +import triton.language as tl
7 +
8 +from vllm.utils.torch_utils import direct_register_custom_op
9 +
10 +_LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
11 +
12 +
13 +def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
14 + """
15 + `_LORA_PTR_DICT` collects the required information during `profile_run`,
16 + After this, it remains constant and subsequent usage is through LUT.
17 + Refer to:
18 + https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
19 + """
20 + key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights)
21 +
22 + if (ptr_tensor := _LORA_PTR_DICT.get(key)) is not None:
23 + return ptr_tensor
24 +
25 + tensor_ptrs = []
26 + for lora_weight in lora_weights:
27 + tensor_ptrs.append(lora_weight.data_ptr())
28 + ptr_tensor = torch.tensor(tensor_ptrs, device=device)
29 +
30 + _LORA_PTR_DICT[key] = ptr_tensor
31 + return _LORA_PTR_DICT.get(key)
32 +
33 +
34 +@triton.jit
35 +def _fused_moe_lora_kernel(
36 + a_ptr,
37 + b_ptr,
38 + c_ptr,
39 + topk_weights_ptr,
40 + sorted_token_ids_ptr,
41 + expert_ids_ptr,
42 + num_tokens_post_padded_ptr,
43 + # Matrix dimensions
44 + N,
45 + K,
46 + EM,
47 + num_valid_tokens,
48 + num_experts,
49 + # The stride variables represent how much to increase the ptr by when
50 + # moving by 1 element in a particular dimension. E.g. `stride_am` is
51 + # how much to increase `a_ptr` by to get the element one row down
52 + # (A has M rows).
53 + stride_am,
54 + stride_ak,
55 + stride_bl,
56 + stride_be,
57 + stride_bk,
58 + stride_bn,
59 + stride_cm,
60 + stride_cn,
61 + stride_tl,
62 + stride_el,
63 + # Meta-parameters
64 + num_slice_a: tl.constexpr,
65 + num_slice_c: tl.constexpr,
66 + slice_a_size: tl.constexpr,
67 + slice_c_size: tl.constexpr,
68 + top_k: tl.constexpr,
69 + MUL_ROUTED_WEIGHT: tl.constexpr,
70 + BLOCK_SIZE_M: tl.constexpr,
71 + BLOCK_SIZE_N: tl.constexpr,
72 + BLOCK_SIZE_K: tl.constexpr,
73 + GROUP_SIZE_M: tl.constexpr,
74 +):
75 + pid = tl.program_id(axis=0)
76 + slice_id = tl.program_id(axis=1)
77 + lora_idx = tl.program_id(axis=2)
78 + max_loras = tl.num_programs(axis=2)
79 +
80 + # calculate pid_m,pid_n
81 + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
82 + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
83 + num_pid_in_group = GROUP_SIZE_M * num_pid_n
84 + group_id = pid // num_pid_in_group
85 + first_pid_m = group_id * GROUP_SIZE_M
86 + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
87 + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
88 + pid_n = (pid % num_pid_in_group) // group_size_m
89 +
90 + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_idx)
91 + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
92 + return
93 +
94 + # get the expert_id to process curr shard
95 + ind = lora_idx * stride_el + pid_m
96 + expert_id = tl.load(expert_ids_ptr + ind)
97 + if expert_id == -1:
98 + return
99 +
100 + # get a_ptr,b_ptr,c_ptr
101 + cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size
102 + cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(tl.bfloat16))
103 + cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size
104 +
105 + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
106 + offs_k = tl.arange(0, BLOCK_SIZE_K)
107 +
108 + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
109 + token_ind = stride_tl * lora_idx + offs_token_id
110 + offs_token = tl.load(
111 + sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0.0
112 + )
113 + token_mask = offs_token < num_valid_tokens
114 +
115 + # get a_ptrs,b_ptrs
116 + a_ptrs = cur_a_ptr + (
117 + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
118 + )
119 +
120 + b_ptrs = (
121 + cur_b_ptr
122 + + lora_idx * stride_bl
123 + + expert_id * stride_be
124 + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
125 + )
126 +
127 + # accumulator
128 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
129 + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
130 + a = tl.load(
131 + a_ptrs,
132 + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
133 + other=0.0,
134 + )
135 + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
136 + accumulator += tl.dot(a, b)
137 + # Advance the ptrs to the next K block.
138 + a_ptrs += BLOCK_SIZE_K * stride_ak
139 + b_ptrs += BLOCK_SIZE_K * stride_bk
140 +
141 + if MUL_ROUTED_WEIGHT:
142 + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
143 + accumulator = accumulator * moe_weight[:, None]
144 +
145 + accumulator = accumulator.to(tl.bfloat16)
146 + # Write back the block of the output
147 + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
148 + c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
149 + c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
150 + tl.store(c_ptrs, accumulator, mask=c_mask)
151 +
152 +
153 +@torch.inference_mode()
154 +def _fused_moe_lora(
155 + output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)
156 + qcurr_hidden_states: torch.Tensor, # (num_tokens, K,)
157 + lora_a_stacked: list[
158 + torch.Tensor
159 + ], # [(max_loras, num_experts, max_lora_rank, K,),...]
160 + lora_b_stacked: list[
161 + torch.Tensor
162 + ], # [(max_loras, num_experts, N, max_lora_rank,),...]
163 + topk_weights: torch.Tensor, # (num_tokens, top_k_num)
164 + sorted_token_ids: torch.Tensor, # (max_loras, _)
165 + expert_ids: torch.Tensor, # (max_loras, _ ,)
166 + num_tokens_post_padded: torch.Tensor, # (max_loras, )
167 + max_lora_rank: int,
168 + top_k_num: int,
169 + block_size_m: int,
170 + block_size_n: int,
171 + block_size_k: int,
172 + group_size_m: int,
173 + mul_routed_weight: bool = False,
174 +) -> None:
175 + assert len(lora_a_stacked) == len(lora_b_stacked) > 0
176 + assert (
177 + sorted_token_ids.dim()
178 + == expert_ids.dim()
179 + == topk_weights.dim()
180 + == qcurr_hidden_states.dim()
181 + == 2
182 + )
183 + assert (
184 + sorted_token_ids.shape[0]
185 + == expert_ids.shape[0]
186 + == num_tokens_post_padded.shape[0]
187 + )
188 + assert len(lora_b_stacked) * lora_b_stacked[0].shape[-2] == output.shape[-1]
189 + assert output.shape[0] == topk_weights.shape[0]
190 + assert top_k_num == topk_weights.shape[1]
191 +
192 + device = qcurr_hidden_states.device
193 + num_slices = len(lora_a_stacked)
194 +
195 + config = {
196 + "BLOCK_SIZE_M": block_size_m,
197 + "BLOCK_SIZE_N": block_size_n,
198 + "BLOCK_SIZE_K": block_size_k,
199 + "GROUP_SIZE_M": group_size_m,
200 + }
201 +
202 + w1_lora_a_stacked = lora_a_stacked[0]
203 + w1_lora_b_stacked = lora_b_stacked[0]
204 + num_experts = lora_a_stacked[0].shape[1]
205 +
206 + N = max_lora_rank
207 + M = topk_weights.shape[0]
208 + EM = sorted_token_ids.shape[1]
209 + K = qcurr_hidden_states.shape[1]
210 + num_tokens = M * top_k_num
211 + w1_output_dim_size = w1_lora_b_stacked.shape[2]
212 +
213 + lora_intermediate_cache1 = torch.zeros(
214 + (num_slices * M * top_k_num * (max_lora_rank + w1_output_dim_size)),
215 + dtype=torch.bfloat16,
216 + device=device,
217 + )
218 +
219 + # slices
220 + a_intermediate_size = num_slices * M * top_k_num * max_lora_rank
221 + a_intermediate_cache1 = lora_intermediate_cache1[:a_intermediate_size].view(
222 + num_slices, M, top_k_num, max_lora_rank
223 + )
224 + b_intermediate_cache1 = lora_intermediate_cache1[a_intermediate_size:].view(
225 + num_slices, M, top_k_num, w1_output_dim_size
226 + )
227 +
228 + b_ptr = _get_ptr(lora_a_stacked, device)
229 +
230 + grid = lambda META: (
231 + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
232 + len(lora_a_stacked),
233 + lora_a_stacked[0].shape[0],
234 + )
235 +
236 + _fused_moe_lora_kernel[grid](
237 + qcurr_hidden_states,
238 + b_ptr,
239 + a_intermediate_cache1,
240 + topk_weights,
241 + sorted_token_ids,
242 + expert_ids,
243 + num_tokens_post_padded,
244 + N,
245 + K,
246 + EM,
247 + num_tokens,
248 + num_experts,
249 + qcurr_hidden_states.stride(0),
250 + qcurr_hidden_states.stride(1),
251 + w1_lora_a_stacked.stride(0),
252 + w1_lora_a_stacked.stride(1),
253 + w1_lora_a_stacked.stride(3),
254 + w1_lora_a_stacked.stride(2),
255 + a_intermediate_cache1.stride(2),
256 + a_intermediate_cache1.stride(3),
257 + sorted_token_ids.stride(0),
258 + expert_ids.stride(0),
259 + num_slice_a=1,
260 + num_slice_c=num_slices,
261 + slice_a_size=qcurr_hidden_states.numel(),
262 + slice_c_size=a_intermediate_cache1.numel() // num_slices,
263 + top_k=1 if mul_routed_weight else top_k_num,
264 + MUL_ROUTED_WEIGHT=False,
265 + **config,
266 + )
267 +
268 + b_ptr = _get_ptr(lora_b_stacked, device)
269 + K = max_lora_rank
270 + N = w1_output_dim_size
271 +
272 + # a_intermediate_cache1 = a_intermediate_cache1.view(
273 + # M, -1, a_intermediate_cache1.shape[3]
274 + # )
275 +
276 + a_intermediate_cache1 = a_intermediate_cache1.view(
277 + -1, a_intermediate_cache1.shape[3]
278 + )
279 +
280 + grid = lambda META: (
281 + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
282 + len(lora_b_stacked),
283 + lora_b_stacked[0].shape[0],
284 + )
285 + _fused_moe_lora_kernel[grid](
286 + a_intermediate_cache1,
287 + b_ptr,
288 + b_intermediate_cache1,
289 + topk_weights,
290 + sorted_token_ids,
291 + expert_ids,
292 + num_tokens_post_padded,
293 + N,
294 + K,
295 + EM,
296 + num_tokens,
297 + num_experts,
298 + a_intermediate_cache1.stride(0),
299 + a_intermediate_cache1.stride(1),
300 + w1_lora_b_stacked.stride(0),
301 + w1_lora_b_stacked.stride(1),
302 + w1_lora_b_stacked.stride(3),
303 + w1_lora_b_stacked.stride(2),
304 + b_intermediate_cache1.stride(2),
305 + b_intermediate_cache1.stride(3),
306 + sorted_token_ids.stride(0),
307 + expert_ids.stride(0),
308 + num_slice_a=num_slices,
309 + num_slice_c=num_slices,
310 + slice_a_size=a_intermediate_cache1.numel() // num_slices,
311 + slice_c_size=b_intermediate_cache1.numel() // num_slices,
312 + top_k=1,
313 + MUL_ROUTED_WEIGHT=mul_routed_weight,
314 + **config,
315 + )
316 + for i in range(num_slices):
317 + output[:, :, i * N : (i + 1) * N] += b_intermediate_cache1[i]
318 +
319 +
320 +def _fused_moe_lora_fake(
321 + output: torch.Tensor,
322 + qcurr_hidden_states: torch.Tensor,
323 + lora_a_stacked: list[torch.Tensor],
324 + lora_b_stacked: list[torch.Tensor],
325 + topk_weights: torch.Tensor,
326 + sorted_token_ids: torch.Tensor,
327 + expert_ids: torch.Tensor,
328 + num_tokens_post_padded: torch.Tensor,
329 + max_lora_rank: int,
330 + top_k_num: int,
331 + block_size_m: int,
332 + block_size_n: int,
333 + block_size_k: int,
334 + group_size_m: int,
335 + mul_routed_weight: bool = False,
336 +) -> None:
337 + return
338 +
339 +
340 +try:
341 + direct_register_custom_op(
342 + op_name="fused_moe_lora",
343 + op_func=_fused_moe_lora,
344 + mutates_args=["output"],
345 + fake_impl=_fused_moe_lora_fake,
346 + )
347 + fused_moe_lora = torch.ops.vllm.fused_moe_lora
348 +
349 +except AttributeError:
350 + fused_moe_lora = _fused_moe_lora
vllm/lora/punica_wrapper/punica_base.py
+39
@@ -448,3 +448,42 @@ class PunicaWrapperBase(PunicaWrapperABC):
448 """
449 # TODO: implement it based on torch ops
450 raise NotImplementedError
451 +
452 + def moe_lora_align_block_size(
453 + self,
454 + topk_ids: torch.Tensor,
455 + num_tokens: int,
456 + block_size: int,
457 + num_experts: int,
458 + max_loras: int,
459 + expert_map: torch.Tensor | None = None,
460 + pad_sorted_ids: bool = False,
461 + ) -> tuple[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.
465 + """
466 + # TODO: implement it based on torch ops
467 + raise NotImplementedError
468 +
469 + def add_lora_fused_moe(
470 + self,
471 + y: torch.Tensor,
472 + x: torch.Tensor,
473 + lora_a_stacked: list[torch.Tensor],
474 + lora_b_stacked: list[torch.Tensor],
475 + topk_weights: torch.Tensor,
476 + sorted_token_ids: torch.Tensor,
477 + expert_ids: torch.Tensor,
478 + num_tokens_post_padded: torch.Tensor,
479 + max_lora_rank: int,
480 + top_k_num: int,
481 + config,
482 + mul_routed_weight=False,
483 + ):
484 + """
485 + Performs a fused forward computation for LoRA of
486 + Mixture-of-Experts (MoE) layer.
487 + """
488 + # TODO: implement it based on torch ops
489 + raise NotImplementedError
vllm/lora/punica_wrapper/punica_gpu.py
+98 -2
@@ -12,10 +12,18 @@ from typing import final
12 import torch
13
14 from vllm.lora.layers import LoRAMapping
15 -from vllm.triton_utils import HAS_TRITON
15 +from vllm.triton_utils import HAS_TRITON, triton
16 +from vllm.utils import round_up
17
18 if HAS_TRITON:
18 - from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_expand, lora_shrink
19 + from vllm.lora.ops.triton_ops import (
20 + LoRAKernelMeta,
21 + fused_moe_lora,
22 + lora_expand,
23 + lora_shrink,
24 + )
25 +
26 +from vllm import _custom_ops as ops
27
28 from .punica_base import PunicaWrapperBase
29
@@ -289,3 +297,91 @@ class PunicaWrapperGPU(PunicaWrapperBase):
297 add_inputs=True,
298 )
299 y = y.view_as(y_org)
300 +
301 + def moe_lora_align_block_size(
302 + self,
303 + topk_ids: torch.Tensor,
304 + num_tokens: int,
305 + block_size: int,
306 + num_experts: int,
307 + max_loras: int,
308 + expert_map: torch.Tensor | None = None,
309 + pad_sorted_ids: bool = False,
310 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
311 + """
312 + Aligns tokens and experts into block-sized chunks for LoRA-based
313 + mixture-of-experts (MoE) execution.
314 + """
315 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
316 + if pad_sorted_ids:
317 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
318 + sorted_ids = torch.empty(
319 + (max_loras * max_num_tokens_padded,),
320 + dtype=torch.int32,
321 + device=topk_ids.device,
322 + )
323 + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
324 + # Expert ids must be set default to -1 to prevent a blank block
325 + expert_ids = torch.empty(
326 + (max_loras * max_num_m_blocks,),
327 + dtype=torch.int32,
328 + device=topk_ids.device,
329 + )
330 + num_tokens_post_pad = torch.empty(
331 + (max_loras), dtype=torch.int32, device=topk_ids.device
332 + )
333 +
334 + (token_lora_mapping, _, _, _, _, _) = self.token_mapping_meta.meta_args(
335 + num_tokens
336 + )
337 +
338 + ops.moe_lora_align_block_size(
339 + topk_ids,
340 + token_lora_mapping,
341 + num_experts,
342 + block_size,
343 + max_loras,
344 + sorted_ids,
345 + expert_ids,
346 + num_tokens_post_pad,
347 + )
348 + if expert_map is not None:
349 + expert_ids = expert_map[expert_ids]
350 +
351 + return sorted_ids, expert_ids, num_tokens_post_pad
352 +
353 + def add_lora_fused_moe(
354 + self,
355 + y: torch.Tensor,
356 + x: torch.Tensor,
357 + lora_a_stacked: list[torch.Tensor],
358 + lora_b_stacked: list[torch.Tensor],
359 + topk_weights: torch.Tensor,
360 + sorted_token_ids: torch.Tensor,
361 + expert_ids: torch.Tensor,
362 + num_tokens_post_padded: torch.Tensor,
363 + max_lora_rank: int,
364 + top_k_num: int,
365 + config,
366 + mul_routed_weight=False,
367 + ):
368 + """
369 + Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
370 + """
371 + fused_moe_lora(
372 + y,
373 + x,
374 + lora_a_stacked,
375 + lora_b_stacked,
376 + topk_weights,
377 + sorted_token_ids,
378 + expert_ids,
379 + num_tokens_post_padded,
380 + max_lora_rank,
381 + top_k_num,
382 + config["BLOCK_SIZE_M"],
383 + config["BLOCK_SIZE_N"],
384 + config["BLOCK_SIZE_K"],
385 + config["GROUP_SIZE_M"],
386 + mul_routed_weight,
387 + )
vllm/lora/utils.py
+39
@@ -23,6 +23,7 @@ from vllm.lora.layers import (
23 BaseLayerWithLoRA,
24 ColumnParallelLinearWithLoRA,
25 ColumnParallelLinearWithShardedLoRA,
26 + FusedMoEWithLoRA,
27 LogitsProcessorWithLoRA,
28 MergedColumnParallelLinearWithLoRA,
29 MergedColumnParallelLinearWithShardedLoRA,
@@ -35,7 +36,9 @@ from vllm.lora.layers import (
36 RowParallelLinearWithShardedLoRA,
37 VocabParallelEmbeddingWithLoRA,
38 )
39 +from vllm.model_executor.layers.fused_moe import FusedMoE
40 from vllm.model_executor.layers.linear import LinearBase
41 +from vllm.model_executor.utils import get_moe_expert_mapping, get_packed_modules_mapping
42
43 if TYPE_CHECKING:
44 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -58,9 +61,18 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {
61 MergedColumnParallelLinearWithShardedLoRA,
62 MergedQKVParallelLinearWithShardedLoRA,
63 RowParallelLinearWithShardedLoRA,
64 + FusedMoEWithLoRA,
65 }
66
67
68 +def is_moe_model(model: nn.Module) -> bool:
69 + """Checks if the model contains FusedMoE layers and warns the user."""
70 + if any(isinstance(module, FusedMoE) for module in model.modules()):
71 + logger.info_once("MoE model detected. Using fused MoE LoRA implementation.")
72 + return True
73 + return False
74 +
75 +
76 def from_layer(
77 layer: nn.Module,
78 max_loras: int,
@@ -205,6 +217,9 @@ def get_supported_lora_modules(model: nn.Module) -> list[str]:
217 if isinstance(module, (LinearBase,)):
218 supported_lora_modules.add(name.split(".")[-1])
219
220 + if isinstance(module, (FusedMoE,)):
221 + supported_lora_modules.add(name.split(".")[-1])
222 +
223 return list(supported_lora_modules)
224
225
@@ -252,3 +267,27 @@ def get_adapter_absolute_path(lora_path: str) -> str:
267 return lora_path
268
269 return local_snapshot_path
270 +
271 +
272 +def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
273 + if is_moe_model(model):
274 + if moe_packed_mapping := get_moe_expert_mapping(model):
275 + # This method generates and returns a dictionary mapping packed module
276 + # names to lists of their corresponding submodule names. It includes
277 + # both static mappings and dynamic mappings for expert layers, where
278 + # the expert indices are expanded based on the configured number
279 + # of routed experts.
280 + packed_modules_mapping = get_packed_modules_mapping(model)
281 +
282 + packed_modules_mapping["experts"] = [
283 + weight_name.rstrip(".") for _, weight_name, _, _ in moe_packed_mapping
284 + ]
285 +
286 + return packed_modules_mapping
287 + else:
288 + raise AttributeError(
289 + "To support LoRA for MoE model, "
290 + "'get_expert_mapping' must be implemented"
291 + )
292 + else:
293 + return get_packed_modules_mapping(model)
vllm/lora/worker_manager.py
+2 -1
@@ -94,7 +94,8 @@ class WorkerLoRAManager:
94 expected_lora_modules.extend(packed_modules_mapping[module])
95 else:
96 expected_lora_modules.append(module)
97 -
97 + if module == "experts":
98 + expected_lora_modules.append(module)
99 expected_lora_modules = list(set(expected_lora_modules))
100 lora_path = get_adapter_absolute_path(lora_request.lora_path)
101
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+53 -18
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Fused MoE utilities for GPTQ."""
4
5 +from collections.abc import Callable
6 +
7 import torch
8
9 import vllm._custom_ops as ops
@@ -11,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
13 batched_moe_align_block_size,
14 moe_align_block_size,
15 )
16 +from vllm.model_executor.layers.fused_moe.prepare_finalize import (
17 + MoEPrepareAndFinalizeNoEP,
18 +)
19 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
20 TopKWeightAndReduceDelegate,
21 TopKWeightAndReduceNoOP,
@@ -24,6 +29,21 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
29 from vllm.scalar_type import ScalarType, scalar_types
30
31
32 +def default_activation_func(
33 + activation: str, output: torch.Tensor, input: torch.Tensor
34 +) -> None:
35 + if activation == "silu":
36 + torch.ops._C.silu_and_mul(output, input)
37 + elif activation == "swigluoai":
38 + # alpha = 1.702, limit = 7.0
39 + torch.ops._C.swigluoai_and_mul(output, input)
40 + else:
41 + raise ValueError(
42 + f"Unsupported activation: {activation}. "
43 + "Only silu and swigluoai activations are supported."
44 + )
45 +
46 +
47 def _fused_marlin_moe(
48 hidden_states: torch.Tensor,
49 w1: torch.Tensor,
@@ -36,12 +56,15 @@ def _fused_marlin_moe(
56 num_topk: int,
57 quant_type: ScalarType,
58 apply_router_weight_on_input: bool,
39 - activation: str,
59 expert_map: torch.Tensor | None,
60 block_size_m: int,
61 sorted_token_ids: torch.Tensor,
62 expert_ids: torch.Tensor,
63 num_tokens_post_padded: torch.Tensor,
64 + activation: str = "silu",
65 + activation_func: Callable[
66 + [str, torch.Tensor, torch.Tensor], None
67 + ] = default_activation_func,
68 global_scale1: torch.Tensor | None = None,
69 global_scale2: torch.Tensor | None = None,
70 g_idx1: torch.Tensor | None = None,
@@ -118,20 +141,9 @@ def _fused_marlin_moe(
141 is_zp_float=False,
142 )
143
121 - if activation == "silu":
122 - torch.ops._C.silu_and_mul(
123 - intermediate_cache2, intermediate_cache1.view(-1, 2 * N)
124 - )
125 - elif activation == "swigluoai":
126 - # alpha = 1.702, limit = 7.0
127 - torch.ops._C.swigluoai_and_mul(
128 - intermediate_cache2, intermediate_cache1.view(-1, 2 * N)
129 - )
130 - else:
131 - raise ValueError(
132 - f"Unsupported activation: {activation}. "
133 - "Only silu and swigluoai activations are supported."
134 - )
144 + activation_func(
145 + activation, intermediate_cache2, intermediate_cache1.view(-1, 2 * N)
146 + )
147
148 if output is None:
149 output = intermediate_cache3
@@ -185,7 +197,11 @@ def fused_marlin_moe(
197 quant_type_id: int,
198 apply_router_weight_on_input: bool = False,
199 global_num_experts: int = -1,
188 - activation: str | None = "silu",
200 + activation: str = "silu",
201 + activation_func: Callable[
202 + [str, torch.Tensor, torch.Tensor], None
203 + ] = default_activation_func,
204 + moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
205 expert_map: torch.Tensor | None = None,
206 global_scale1: torch.Tensor | None = None,
207 global_scale2: torch.Tensor | None = None,
@@ -290,12 +306,13 @@ def fused_marlin_moe(
306 num_topk=topk,
307 quant_type=quant_type,
308 apply_router_weight_on_input=apply_router_weight_on_input,
293 - activation=activation,
309 expert_map=expert_map,
310 block_size_m=block_size_m,
311 sorted_token_ids=sorted_token_ids,
312 expert_ids=expert_ids,
313 num_tokens_post_padded=num_tokens_post_padded,
314 + activation=activation,
315 + activation_func=activation_func,
316 global_scale1=global_scale1,
317 global_scale2=global_scale2,
318 g_idx1=g_idx1,
@@ -317,7 +334,10 @@ def fused_marlin_moe(
334 else:
335 output = torch.empty_like(hidden_states)
336
320 - return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output)
337 + if moe_sum is None:
338 + return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output)
339 + else:
340 + return moe_sum(moe_output, output)
341
342
343 def batched_fused_marlin_moe(
@@ -600,6 +620,8 @@ class MarlinExperts(MarlinExpertsBase):
620 apply_router_weight_on_input=apply_router_weight_on_input,
621 global_num_experts=global_num_experts,
622 activation=activation,
623 + activation_func=self.activation,
624 + moe_sum=self.moe_sum,
625 expert_map=expert_map,
626 output=output,
627 # Workspaces are swapped in workspace_shapes() to account for proper
@@ -608,6 +630,19 @@ class MarlinExperts(MarlinExpertsBase):
630 intermediate_cache2=workspace13,
631 )
632
633 + def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
634 + ops.moe_sum(input, output)
635 +
636 +
637 +def modular_marlin_fused_moe(
638 + quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None
639 +) -> mk.FusedMoEModularKernel:
640 + return mk.FusedMoEModularKernel(
641 + MoEPrepareAndFinalizeNoEP(),
642 + MarlinExperts(quant_config),
643 + shared_experts,
644 + )
645 +
646
647 class BatchedMarlinExperts(MarlinExpertsBase):
648 def __init__(
vllm/model_executor/layers/fused_moe/fused_moe.py
+7 -2
@@ -2135,13 +2135,18 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2135 B_bias=self.w2_bias,
2136 )
2137
2138 - ops.moe_sum(intermediate_cache3, output)
2138 + # separate function is required for MoE + LoRA
2139 + self.moe_sum(intermediate_cache3, output)
2140 +
2141 + def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
2142 + ops.moe_sum(input, output)
2143
2144
2145 def modular_triton_fused_moe(
2142 - quant_config: FusedMoEQuantConfig,
2146 + quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None
2147 ) -> mk.FusedMoEModularKernel:
2148 return mk.FusedMoEModularKernel(
2149 MoEPrepareAndFinalizeNoEP(),
2150 TritonExperts(quant_config),
2151 + shared_experts,
2152 )
vllm/model_executor/layers/fused_moe/modular_kernel.py
+3
@@ -557,6 +557,9 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
557 torch.ops._C.silu_and_mul(output, input)
558 elif activation == "gelu":
559 torch.ops._C.gelu_and_mul(output, input)
560 + elif activation == "swigluoai":
561 + # alpha = 1.702, limit = 7.0
562 + torch.ops._C.swigluoai_and_mul(output, input)
563 else:
564 raise ValueError(f"Unsupported FusedMoe activation: {activation}")
565
vllm/model_executor/models/deepseek_v2.py
+11
@@ -1313,6 +1313,17 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR
1313 logits = self.logits_processor(self.lm_head, hidden_states)
1314 return logits
1315
1316 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1317 + # Params for weights, fp8 weight scales, fp8 activation scales
1318 + # (param_name, weight_name, expert_id, shard_id)
1319 + return SharedFusedMoE.make_expert_params_mapping(
1320 + ckpt_gate_proj_name="gate_proj",
1321 + ckpt_down_proj_name="down_proj",
1322 + ckpt_up_proj_name="up_proj",
1323 + num_experts=self.config.n_routed_experts,
1324 + num_redundant_experts=0,
1325 + )
1326 +
1327 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1328 stacked_params_mapping = [
1329 # (param_name, shard_name, shard_id)
vllm/model_executor/models/gpt_oss.py
+13 -2
@@ -32,7 +32,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
32 from vllm.sequence import IntermediateTensors
33 from vllm.utils import cdiv
34
35 -from .interfaces import SupportsEagle3, SupportsPP
35 +from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
36 from .utils import (
37 AutoWeightsLoader,
38 WeightsMapper,
@@ -627,7 +627,7 @@ class GptOssModel(nn.Module):
627 )
628
629
630 -class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
630 +class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
631 packed_modules_mapping = {"qkv": ["q_proj", "k_proj", "v_proj"]}
632
633 hf_to_vllm_mapper = WeightsMapper(
@@ -696,6 +696,17 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
696 logits = self.logits_processor(self.lm_head, hidden_states)
697 return logits
698
699 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
700 + # Params for weights, weight scales, activation scales
701 + # (param_name, weight_name, expert_id, shard_id)
702 + return FusedMoE.make_expert_params_mapping(
703 + ckpt_gate_proj_name="gate_proj",
704 + ckpt_down_proj_name="down_proj",
705 + ckpt_up_proj_name="up_proj",
706 + num_experts=self.config.num_local_experts,
707 + num_redundant_experts=0,
708 + )
709 +
710 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
711 loader = AutoWeightsLoader(
712 self,
vllm/model_executor/models/olmoe.py
+3 -9
@@ -49,7 +49,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
49 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
50 from vllm.sequence import IntermediateTensors
51
52 -from .interfaces import SupportsPP
52 +from .interfaces import SupportsLoRA, SupportsPP
53 from .utils import (
54 AutoWeightsLoader,
55 is_pp_missing_parameter,
@@ -349,8 +349,6 @@ class OlmoeModel(nn.Module):
349 ("qkv_proj", "q_proj", "q"),
350 ("qkv_proj", "k_proj", "k"),
351 ("qkv_proj", "v_proj", "v"),
352 - ("gate_up_proj", "gate_proj", 0),
353 - ("gate_up_proj", "up_proj", 1),
352 ]
353
354 params_dict = dict(self.named_parameters())
@@ -433,17 +431,13 @@ class OlmoeModel(nn.Module):
431 return loaded_params
432
433
436 -class OlmoeForCausalLM(nn.Module, SupportsPP):
434 +class OlmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
435 packed_modules_mapping = {
436 "qkv_proj": [
437 "q_proj",
438 "k_proj",
439 "v_proj",
442 - ],
443 - "gate_up_proj": [
444 - "gate_proj",
445 - "up_proj",
446 - ],
440 + ]
441 }
442
443 def __init__(