@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 294c805f1d

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

Early exit for MoE LoRA kernels (#27131)

Signed-off-by: gnovack <gnovack@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

gnovack committed Nov 3, 2025 at 04:22 UTC 294c805f1df9ddf62c2290989710da9d48ab4973

11 files changed +123 -34

csrc/moe/moe_lora_align_sum_kernels.cu

+16 -11

@@ -28,11 +28,16 @@ __global__ void moe_lora_align_sum_kernel(

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) {

31 + int topk_num, int32_t* total_tokens_post_pad, int32_t* adapter_enabled,

32 + int32_t* lora_ids) {

33 const size_t tokens_per_thread = div_ceil(numel, blockDim.x);

34 const size_t start_idx = threadIdx.x * tokens_per_thread;

35 - int lora_id = blockIdx.x;

36 + int lora_idx = blockIdx.x;

37 + int lora_id = lora_ids[lora_idx];

38 + if (lora_id == -1 || adapter_enabled[lora_id] == 0) {

39 + return;

40 + }

41 extern __shared__ int32_t shared_mem[];

42 int32_t* cumsum = shared_mem;

43 token_cnts_t* tokens_cnts = (token_cnts_t*)(shared_mem + num_experts + 1);

@@ -121,14 +126,13 @@ __global__ void moe_lora_align_sum_kernel(

126 }

127 }

128

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, int64_t max_num_tokens_padded,

128 - int64_t max_num_m_blocks,

129 - torch::Tensor sorted_token_ids,

130 - torch::Tensor expert_ids,

131 - torch::Tensor num_tokens_post_pad) {

129 +void moe_lora_align_block_size(

130 + torch::Tensor topk_ids, torch::Tensor token_lora_mapping,

131 + int64_t num_experts, int64_t block_size, int64_t max_loras,

132 + int64_t max_num_tokens_padded, int64_t max_num_m_blocks,

133 + torch::Tensor sorted_token_ids, torch::Tensor expert_ids,

134 + torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,

135 + torch::Tensor lora_ids) {

136 const int topk_num = topk_ids.size(1);

137

138 TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");

@@ -164,6 +168,7 @@ void moe_lora_align_block_size(torch::Tensor topk_ids,

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,

167 - num_tokens_post_pad.data_ptr<int32_t>());

171 + num_tokens_post_pad.data_ptr<int32_t>(),

172 + adapter_enabled.data_ptr<int32_t>(), lora_ids.data_ptr<int32_t>());

173 });

174 }

\ No newline at end of file

csrc/moe/moe_ops.h

+7 -8

@@ -20,14 +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);

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, int64_t max_num_tokens_padded,

27 - int64_t max_num_m_blocks,

28 - torch::Tensor sorted_token_ids,

29 - torch::Tensor expert_ids,

30 - torch::Tensor num_tokens_post_pad);

23 +void moe_lora_align_block_size(

24 + torch::Tensor topk_ids, torch::Tensor token_lora_mapping,

25 + int64_t num_experts, int64_t block_size, int64_t max_loras,

26 + int64_t max_num_tokens_padded, int64_t max_num_m_blocks,

27 + torch::Tensor sorted_token_ids, torch::Tensor expert_ids,

28 + torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,

29 + torch::Tensor lora_ids);

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

+3 -1

@@ -44,7 +44,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {

44 " int max_num_m_blocks, "

45 " Tensor !sorted_token_ids,"

46 " Tensor !experts_ids,"

47 - " Tensor !num_tokens_post_pad) -> () ");

47 + " Tensor !num_tokens_post_pad,"

48 + " Tensor !adapter_enabled,"

49 + " Tensor !lora_ids) -> () ");

50 m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size);

52 #ifndef USE_ROCM

tests/lora/test_fused_moe_lora_kernel.py

@@ -134,6 +134,8 @@ def use_fused_moe_lora_kernel(

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 + adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)

138 + lora_ids = torch.arange(max_loras + 2, dtype=torch.int32)

139

140 # call kernel

141 ops.moe_lora_align_block_size(

@@ -147,6 +149,8 @@ def use_fused_moe_lora_kernel(

149 sorted_token_ids,

150 expert_ids,

151 num_tokens_post_padded,

152 + adapter_enabled,

153 + lora_ids,

154 )

155

156 config = {

@@ -172,6 +176,8 @@ def use_fused_moe_lora_kernel(

176 num_tokens_post_padded,

177 max_lora_rank,

178 top_k_num,

179 + lora_ids,

180 + adapter_enabled,

181 config["BLOCK_SIZE_M"],

182 config["BLOCK_SIZE_N"],

183 config["BLOCK_SIZE_K"],

tests/lora/test_moe_lora_align_sum.py

@@ -60,6 +60,8 @@ def test_moe_lora_align_block_size(

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 + adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")

64 + lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device="cuda")

66 # call kernel

67 ops.moe_lora_align_block_size(

@@ -73,6 +75,8 @@ def test_moe_lora_align_block_size(

75 sorted_token_ids,

76 expert_ids,

77 num_tokens_post_pad,

78 + adapter_enabled,

79 + lora_ids,

80 )

82 # verify values

tests/lora/test_olmoe_tp.py

+43 -7

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +

5 import vllm

6 from vllm.lora.request import LoRARequest

@@ -28,8 +29,17 @@ EXPECTED_LORA_OUTPUT = [

29 "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501

30 ]

32 +EXPECTED_BASE_MODEL_OUTPUT = [

33 + "SELECT COUNT(Candidate_ID) FROM candidate",

34 + "SELECT COUNT(Candidate_ID) FROM candidate",

35 + "SELECT Candidate_ID, COUNT(*) as Total_Candidates\nFROM candidate\nINNER JOIN people ON candidate.People_ID = people.People_ID", # noqa: E501

36 + "SELECT Candidate_ID, Poll_Source FROM candidate WHERE People_ID IN (SELECT People_ID FROM people) ORDER BY COUNT(*) DESC LIMIT 1", # noqa: E501

37 +]

38 +

32 -def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:

40 +def generate_and_test(

41 + llm: vllm.LLM, lora_path: str, lora_id: list[int | None] | int | None

42 +) -> None:

43 prompts = [

44 PROMPT_TEMPLATE.format(context="How many candidates are there?"),

45 PROMPT_TEMPLATE.format(context="Count the number of candidates."),

@@ -40,12 +50,18 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:

50 context="Return the poll resource associated with the most candidates."

51 ),

52 ]

53 +

54 + lora_request = None

55 + if isinstance(lora_id, int):

56 + lora_request = LoRARequest(str(lora_id), lora_id, lora_path)

57 + elif isinstance(lora_id, list):

58 + lora_request = [

59 + LoRARequest(str(i), i, lora_path) if i is not None else None

60 + for i in lora_id

61 + ]

62 +

63 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 - )

64 + outputs = llm.generate(prompts, sampling_params, lora_request=lora_request)

65 # Print the outputs.

66 generated_texts: list[str] = []

67 for output in outputs:

@@ -55,7 +71,13 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:

71 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

73 for i in range(len(EXPECTED_LORA_OUTPUT)):

58 - assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])

74 + req_lora_id = lora_id[i] if isinstance(lora_id, list) else lora_id

75 + expected_output = (

76 + EXPECTED_LORA_OUTPUT[i]

77 + if req_lora_id is not None

78 + else EXPECTED_BASE_MODEL_OUTPUT[i]

79 + )

80 + assert generated_texts[i].startswith(expected_output)

83 def test_olmoe_lora(olmoe_lora_files):

@@ -75,6 +97,20 @@ def test_olmoe_lora(olmoe_lora_files):

97 generate_and_test(llm, olmoe_lora_files, lora_id=2)

100 +def test_olmoe_lora_mixed(olmoe_lora_files):

101 + llm = vllm.LLM(

102 + MODEL_PATH,

103 + max_model_len=1024,

104 + enable_lora=True,

105 + max_loras=4,

106 + enforce_eager=True,

107 + trust_remote_code=True,

108 + enable_chunked_prefill=True,

109 + )

110 +

111 + generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])

112 +

113 +

114 @multi_gpu_test(num_gpus=2)

115 def test_olmoe_lora_tp2(olmoe_lora_files):

116 llm = vllm.LLM(

vllm/_custom_ops.py

@@ -1823,6 +1823,8 @@ def moe_lora_align_block_size(

1823 sorted_token_ids: torch.Tensor,

1824 experts_ids: torch.Tensor,

1825 num_tokens_post_pad: torch.Tensor,

1826 + adapter_enabled: torch.Tensor,

1827 + lora_ids: torch.Tensor,

1828 ) -> None:

1829 torch.ops._moe_C.moe_lora_align_block_size(

1830 topk_ids,

@@ -1835,6 +1837,8 @@ def moe_lora_align_block_size(

1837 sorted_token_ids,

1838 experts_ids,

1839 num_tokens_post_pad,

1840 + adapter_enabled,

1841 + lora_ids,

1842 )

1843

1844

vllm/lora/layers/fused_moe.py

+10 -1

@@ -111,6 +111,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

111 config["BLOCK_SIZE_M"],

112 self.base_layer.local_num_experts,

113 max_loras,

114 + self.adapter_enabled,

115 expert_map,

116 )

117

@@ -138,6 +139,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

139 max_lora_rank,

140 top_k,

141 config,

142 + self.adapter_enabled,

143 )

144

145 result = func(*args, **kwargs)

@@ -196,6 +198,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

198 max_lora_rank,

199 top_k,

200 config,

201 + self.adapter_enabled,

202 True,

203 )

204

@@ -227,6 +230,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

230 ) -> None:

231 """Initializes lora matrices."""

232

233 + self.adapter_enabled = torch.tensor(

234 + [0] * (max_loras + 1), dtype=torch.int, device=self.device

235 + )

236 +

237 self.w1_lora_a_stacked = torch.zeros(

238 (

239 max_loras,

@@ -313,6 +320,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

320 self.w3_lora_b_stacked[index] = 0

321 self.w2_lora_a_stacked[index] = 0

322 self.w2_lora_b_stacked[index] = 0

323 + self.adapter_enabled[index] = 0

324

325 def set_lora(

326 self,

@@ -322,8 +330,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

330 embeddings_tensor: torch.Tensor | None,

331 bias: torch.Tensor | None = None,

332 ):

325 - self.reset_lora(index)

333 """Overwrites lora tensors at index."""

334 + self.reset_lora(index)

335 + self.adapter_enabled[index] = 1

336 for eid in range(len(lora_a) // 3):

337 w1_lora_a = lora_a[eid * 3]

338 w2_lora_a = lora_a[eid * 3 + 1]

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+20 -5

@@ -54,6 +54,8 @@ def _fused_moe_lora_kernel(

54 EM,

55 num_valid_tokens,

56 num_experts,

57 + lora_ids,

58 + adapter_enabled,

59 # The stride variables represent how much to increase the ptr by when

60 # moving by 1 element in a particular dimension. E.g. `stride_am` is

61 # how much to increase `a_ptr` by to get the element one row down

@@ -84,6 +86,11 @@ def _fused_moe_lora_kernel(

86 pid = tl.program_id(axis=0)

87 slice_id = tl.program_id(axis=1)

88 lora_idx = tl.program_id(axis=2)

89 + lora_id = tl.load(lora_ids + lora_idx)

90 + moe_enabled = tl.load(adapter_enabled + lora_id)

91 + if lora_id == -1 or moe_enabled == 0:

92 + # Early exit for the no-lora case.

93 + return

94 max_loras = tl.num_programs(axis=2)

95 grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)

@@ -100,12 +107,12 @@ def _fused_moe_lora_kernel(

107 pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)

108 pid_n = (pid_m_n % num_pid_in_group) // group_size_m

109

103 - num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_idx)

110 + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)

111 if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:

112 return

113

114 # get the expert_id to process curr shard

108 - ind = lora_idx * stride_el + pid_m

115 + ind = lora_id * stride_el + pid_m

116 expert_id = tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)

117 if expert_id == -1:

118 return

@@ -119,7 +126,7 @@ def _fused_moe_lora_kernel(

126 offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)

127

128 offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)

122 - token_ind = stride_tl * lora_idx + offs_token_id

129 + token_ind = stride_tl * lora_id + offs_token_id

130 offs_token = tl.load(

131 sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0

132 )

@@ -132,7 +139,7 @@ def _fused_moe_lora_kernel(

139

140 b_ptrs = (

141 cur_b_ptr

135 - + lora_idx * stride_bl

142 + + lora_id * stride_bl

143 + expert_id * stride_be

144 + offs_k[:, None] * stride_bk

145 + offs_bn[None, :] * stride_bn

@@ -184,6 +191,8 @@ def _fused_moe_lora(

191 num_tokens_post_padded: torch.Tensor, # (max_loras, )

192 max_lora_rank: int,

193 top_k_num: int,

194 + lora_ids: torch.Tensor,

195 + adapter_enabled: torch.Tensor,

196 block_size_m: int,

197 block_size_n: int,

198 block_size_k: int,

@@ -234,7 +243,7 @@ def _fused_moe_lora(

243 num_tokens = M * top_k_num

244 w1_output_dim_size = w1_lora_b_stacked.shape[2]

245

237 - lora_intermediate_cache1 = torch.empty(

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,

@@ -272,6 +281,8 @@ def _fused_moe_lora(

281 EM,

282 num_tokens,

283 num_experts,

284 + lora_ids,

285 + adapter_enabled,

286 qcurr_hidden_states.stride(0),

287 qcurr_hidden_states.stride(1),

288 w1_lora_a_stacked.stride(0),

@@ -319,6 +330,8 @@ def _fused_moe_lora(

330 EM,

331 num_tokens,

332 num_experts,

333 + lora_ids,

334 + adapter_enabled,

335 a_intermediate_cache1.stride(0),

336 a_intermediate_cache1.stride(1),

337 w1_lora_b_stacked.stride(0),

@@ -352,6 +365,8 @@ def _fused_moe_lora_fake(

365 num_tokens_post_padded: torch.Tensor,

366 max_lora_rank: int,

367 top_k_num: int,

368 + lora_ids: torch.Tensor,

369 + adapter_enabled: torch.Tensor,

370 block_size_m: int,

371 block_size_n: int,

372 block_size_k: int,

vllm/lora/punica_wrapper/punica_base.py

@@ -456,6 +456,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

456 block_size: int,

457 num_experts: int,

458 max_loras: int,

459 + adapter_enabled: torch.Tensor,

460 expert_map: torch.Tensor | None = None,

461 pad_sorted_ids: bool = False,

462 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

@@ -479,6 +480,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

480 max_lora_rank: int,

481 top_k_num: int,

482 config,

483 + adapter_enabled: torch.Tensor,

484 mul_routed_weight=False,

485 ):

486 """

vllm/lora/punica_wrapper/punica_gpu.py

+8 -1

@@ -305,6 +305,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

305 block_size: int,

306 num_experts: int,

307 max_loras: int,

308 + adapter_enabled: torch.Tensor,

309 expert_map: torch.Tensor | None = None,

310 pad_sorted_ids: bool = False,

311 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

@@ -331,7 +332,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

332 (max_loras), dtype=torch.int32, device=topk_ids.device

333 )

334

334 - (token_lora_mapping, _, _, _, _, _) = self.token_mapping_meta.meta_args(

335 + (token_lora_mapping, _, _, _, lora_ids, _) = self.token_mapping_meta.meta_args(

336 num_tokens

337 )

338

@@ -346,6 +347,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):

347 sorted_ids,

348 expert_ids,

349 num_tokens_post_pad,

350 + adapter_enabled,

351 + lora_ids,

352 )

353 if expert_map is not None:

354 expert_ids = expert_map[expert_ids]

@@ -365,11 +368,13 @@ class PunicaWrapperGPU(PunicaWrapperBase):

368 max_lora_rank: int,

369 top_k_num: int,

370 config,

371 + adapter_enabled: torch.Tensor,

372 mul_routed_weight=False,

373 ):

374 """

375 Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.

376 """

377 + (_, _, _, _, lora_ids, _) = self.token_mapping_meta.meta_args(x.size(0))

378 fused_moe_lora(

379 y,

380 x,

@@ -381,6 +386,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):

386 num_tokens_post_padded,

387 max_lora_rank,

388 top_k_num,

389 + lora_ids,

390 + adapter_enabled,

391 config["BLOCK_SIZE_M"],

392 config["BLOCK_SIZE_N"],

393 config["BLOCK_SIZE_K"],