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 7, 15:15 UTC
use 'max_active_experts' for moe lora input size (#33197)
Signed-off-by: gnovack <gnovack@amazon.com>
gnovack committed
Feb 26, 2026 at 19:50 UTC
a532c838492aced7d4c9ac30f694368c57371050
2 files changed
+4
tests/lora/test_moe_lora_align_sum.py
+2
@@ -47,6 +47,8 @@ def test_moe_lora_align_block_size(
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
+ if topk_ids.numel() < num_experts:
51
+ max_num_tokens_padded = topk_ids.numel() * block_size
52
max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
53
54
# init output tensors
vllm/lora/punica_wrapper/punica_gpu.py
+2
@@ -351,6 +351,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
351
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
352
if pad_sorted_ids:
353
max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
354
+ if topk_ids.numel() < num_experts:
355
+ max_num_tokens_padded = topk_ids.numel() * block_size
356
sorted_ids = torch.empty(
357
(max_loras * max_num_tokens_padded,),
358
dtype=torch.int32,