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
Bugfix - pass 'max_num_tokens_padded' into 'moe_lora_align_block_size' (#27311)
Signed-off-by: gnovack <gnovack@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
gnovack committed
Oct 22, 2025 at 05:23 UTC
8e4ca4d14e34305c0ab5640a9143871b25c91811
8 files changed
+17
-8
csrc/moe/moe_lora_align_sum_kernels.cu
+2
-6
@@ -124,18 +124,14 @@ __global__ void moe_lora_align_sum_kernel(
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,
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) {
132
const int topk_num = topk_ids.size(1);
133
133
- int max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1);
134
-
134
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);
135
136
int device_max_shared_mem;
137
auto dev = topk_ids.get_device();
csrc/moe/moe_ops.h
+2
-1
@@ -23,7 +23,8 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
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,
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);
csrc/moe/torch_bindings.cpp
+2
@@ -40,6 +40,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
40
" Tensor token_lora_mapping,"
41
" int num_experts,"
42
" int block_size, int max_loras, "
43
+ " int max_num_tokens_padded, "
44
+ " int max_num_m_blocks, "
45
" Tensor !sorted_token_ids,"
46
" Tensor !experts_ids,"
47
" Tensor !num_tokens_post_pad) -> () ");
tests/lora/test_fused_moe_lora_kernel.py
+2
@@ -142,6 +142,8 @@ def use_fused_moe_lora_kernel(
142
num_experts,
143
block_size,
144
max_loras,
145
+ max_num_tokens_padded,
146
+ max_num_m_blocks,
147
sorted_token_ids,
148
expert_ids,
149
num_tokens_post_padded,
tests/lora/test_gptoss.py
+1
-1
@@ -36,7 +36,7 @@ def test_gptoss20b_lora(gptoss20b_lora_files):
36
llm = vllm.LLM(
37
MODEL_PATH,
38
enable_lora=True,
39
- max_loras=1,
39
+ max_loras=4,
40
trust_remote_code=True,
41
)
42
tests/lora/test_moe_lora_align_sum.py
+2
@@ -68,6 +68,8 @@ def test_moe_lora_align_block_size(
68
num_experts,
69
block_size,
70
max_loras,
71
+ max_num_tokens_padded,
72
+ max_num_m_blocks,
73
sorted_token_ids,
74
expert_ids,
75
num_tokens_post_pad,
vllm/_custom_ops.py
+4
@@ -1801,6 +1801,8 @@ def moe_lora_align_block_size(
1801
num_experts: int,
1802
block_size: int,
1803
max_loras: int,
1804
+ max_num_tokens_padded: int,
1805
+ max_num_m_blocks: int,
1806
sorted_token_ids: torch.Tensor,
1807
experts_ids: torch.Tensor,
1808
num_tokens_post_pad: torch.Tensor,
@@ -1811,6 +1813,8 @@ def moe_lora_align_block_size(
1813
num_experts,
1814
block_size,
1815
max_loras,
1816
+ max_num_tokens_padded,
1817
+ max_num_m_blocks,
1818
sorted_token_ids,
1819
experts_ids,
1820
num_tokens_post_pad,
vllm/lora/punica_wrapper/punica_gpu.py
+2
@@ -341,6 +341,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
341
num_experts,
342
block_size,
343
max_loras,
344
+ max_num_tokens_padded,
345
+ max_num_m_blocks,
346
sorted_ids,
347
expert_ids,
348
num_tokens_post_pad,