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, 21:30 UTC
[Bugfix] Fix lora_ids in FusedMoE LoRA test (#35135)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Feb 23, 2026 at 18:49 UTC
a7c9f7b7ec6546cf36eb217a551c4d5b91e6a064
1 file changed
+19
-10
tests/lora/test_fused_moe_lora_kernel.py
+19
-10
@@ -118,7 +118,10 @@ def sample_data(
118
num_tokens, num_experts, top_k_num
119
)
120
token_lora_mapping = assign_loras_to_tokens(num_tokens, num_sequences, max_loras)
121
- return topk_ids, topk_weights, token_lora_mapping
121
+ active_lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
122
+ lora_ids = torch.unique(token_lora_mapping, sorted=True)
123
+ active_lora_ids[: lora_ids.size(0)].copy_(lora_ids, non_blocking=True)
124
+ return topk_ids, topk_weights, token_lora_mapping, active_lora_ids
125
126
127
def use_fused_moe_lora_kernel(
@@ -127,6 +130,7 @@ def use_fused_moe_lora_kernel(
130
token_lora_mapping,
131
max_lora_rank,
132
top_k_num,
133
+ lora_ids,
134
lora_a_stacked,
135
lora_b_stacked,
136
hidden_states,
@@ -149,7 +153,6 @@ def use_fused_moe_lora_kernel(
153
expert_ids = torch.empty((max_loras * max_num_m_blocks,), dtype=torch.int32)
154
num_tokens_post_padded = torch.empty((max_loras,), dtype=torch.int32)
155
adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
152
- lora_ids = torch.arange(max_loras + 2, dtype=torch.int32)
156
157
# call kernel
158
ops.moe_lora_align_block_size(
@@ -168,7 +171,7 @@ def use_fused_moe_lora_kernel(
171
)
172
173
config = {
171
- "BLOCK_SIZE_M": 16,
174
+ "BLOCK_SIZE_M": block_size,
175
"BLOCK_SIZE_N": 32,
176
"BLOCK_SIZE_K": 64,
177
"GROUP_SIZE_M": 1,
@@ -275,7 +278,7 @@ def test_fused_moe_lora_kernel(
278
# the number of randomly generated sentences.
279
num_sequences = 10
280
# generate data
278
- topk_ids, topk_weights, token_lora_mapping = sample_data(
281
+ topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
282
num_tokens, num_sequences, max_loras, num_experts, top_k_num
283
)
284
@@ -318,6 +321,7 @@ def test_fused_moe_lora_kernel(
321
token_lora_mapping,
322
max_lora_rank,
323
top_k_num,
324
+ lora_ids,
325
lora_a_stacked,
326
lora_b_stacked,
327
hidden_states,
@@ -336,7 +340,7 @@ def test_fused_moe_lora_kernel(
340
top_k_num,
341
)
342
339
- torch.testing.assert_close(output, output2, atol=1e-1, rtol=1e-1)
343
+ torch.testing.assert_close(output, output2, atol=1e-2, rtol=1e-2)
344
345
346
def use_fused_moe_lora_kernel_naive(
@@ -345,6 +349,7 @@ def use_fused_moe_lora_kernel_naive(
349
token_lora_mapping,
350
max_lora_rank,
351
top_k_num,
352
+ lora_ids,
353
lora_a_stacked,
354
lora_b_stacked,
355
hidden_states,
@@ -379,7 +384,6 @@ def use_fused_moe_lora_kernel_naive(
384
num_tokens_post_padded = None
385
386
adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
382
- lora_ids = torch.arange(max_loras + 2, dtype=torch.int32)
387
388
# num_active_loras is the number of active LoRAs
389
# (max_loras + 1 to include no-lora case)
@@ -463,7 +467,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
467
# the number of randomly generated sentences.
468
num_sequences = min(num_tokens, 4)
469
# generate data
466
- topk_ids, topk_weights, token_lora_mapping = sample_data(
470
+ topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
471
num_tokens, num_sequences, max_loras, num_experts, top_k_num
472
)
473
@@ -506,6 +510,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
510
token_lora_mapping,
511
max_lora_rank,
512
top_k_num,
513
+ lora_ids,
514
lora_a_stacked,
515
lora_b_stacked,
516
hidden_states,
@@ -524,7 +529,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
529
top_k_num,
530
)
531
527
- torch.testing.assert_close(output, output_ref, atol=1e-1, rtol=1e-1)
532
+ torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
533
534
535
@multi_gpu_test(num_gpus=2)
@@ -556,7 +561,7 @@ def test_fused_moe_lora_kernel_fully_sharded(
561
# the number of randomly generated sentences.
562
num_sequences = 10
563
# generate data
559
- topk_ids, topk_weights, token_lora_mapping = sample_data(
564
+ topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
565
num_tokens, num_sequences, max_loras, num_experts, top_k_num
566
)
567
@@ -576,6 +581,7 @@ def test_fused_moe_lora_kernel_fully_sharded(
581
token_lora_mapping,
582
max_lora_rank,
583
top_k_num,
584
+ lora_ids,
585
max_loras,
586
num_experts,
587
block_size,
@@ -601,6 +607,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
607
token_lora_mapping,
608
max_lora_rank,
609
top_k_num,
610
+ lora_ids,
611
max_loras,
612
num_experts,
613
block_size,
@@ -660,6 +667,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
667
topk_ids = topk_ids.to(device)
668
topk_weights = topk_weights.to(device)
669
token_lora_mapping = token_lora_mapping.to(device)
670
+ lora_ids = lora_ids.to(device)
671
672
ref_output = use_torch(
673
hidden_states,
@@ -698,6 +706,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
706
token_lora_mapping,
707
max_lora_rank,
708
top_k_num,
709
+ lora_ids,
710
[lora_a],
711
[lora_b],
712
hidden_states,
@@ -714,4 +723,4 @@ def use_fused_moe_lora_kernel_tensor_parallel(
723
else:
724
output = tensor_model_parallel_all_reduce(output)
725
717
- torch.testing.assert_close(output, ref_output, atol=1e-1, rtol=1e-1)
726
+ torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)