test_fused_moe_lora_kernel.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,321 lines 37.4 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	import os
4	import random
5
6	import pytest
7	import torch
8
9	from tests.utils import ensure_current_vllm_config, multi_gpu_test
10	from vllm import _custom_ops as ops
11	from vllm.distributed import (
12	init_distributed_environment,
13	initialize_model_parallel,
14	tensor_model_parallel_all_gather,
15	tensor_model_parallel_all_reduce,
16	)
17	from vllm.distributed.parallel_state import (
18	get_tensor_model_parallel_world_size,
19	)
20	from vllm.lora.ops.triton_ops import fused_moe_lora
21	from vllm.platforms import current_platform
22	from vllm.utils.network_utils import get_open_port
23	from vllm.utils.torch_utils import set_random_seed
24
25	# The tensor-parallel cases build their distributed environment inside processes
26	# spawned by torch.multiprocessing.spawn, so this process never has one to tear down.
27	pytestmark = pytest.mark.skip_global_cleanup
28
29
30	@pytest.fixture(autouse=True)
31	def reset_device(reset_default_device):
32	pass
33
34
35	def round_up(x, base):
36	return ((x + base - 1) // base) * base
37
38
39	def CEILDIV(x, y):
40	return (x + y - 1) // y
41
42
43	def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int):
44	"""Split `num_tokens` into `num_sequences` sequences.
45	Each sequence randomly selects 1 LoRA index from [0, max_loras),
46	and all tokens in that sequence are assigned this LoRA index.
47
48	Args:
49	num_tokens (int): Total number of tokens.
50	num_sequences (int): Number of sequences to split the tokens into.
51	max_loras (int): Total number of available LoRA modules.
52
53	Returns:
54	torch.Tensor: 1D tensor of shape [num_tokens], where each value
55	is the LoRA index assigned to that token.
56
57	"""
58	assert num_sequences > 0 and max_loras > 0
59	assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences"
60
61	# Compute token distribution per sequence (distribute remainder evenly)
62	tokens_per_seq = num_tokens // num_sequences
63	remainder = num_tokens % num_sequences
64
65	token_lora_mapping = torch.empty(num_tokens, dtype=torch.int32)
66
67	start = 0
68	for seq_idx in range(num_sequences):
69	# Determine the token range for this sequence
70	end = start + tokens_per_seq + (1 if seq_idx < remainder else 0)
71
72	# Randomly select one LoRA ID for this sequence
73	lora_id = random.randint(0, max_loras - 1)
74
75	# Assign the same LoRA ID to all tokens in this sequence
76	token_lora_mapping[start:end] = lora_id
77
78	start = end
79
80	return token_lora_mapping
81
82
83	def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int):
84	"""For each token, randomly select `top_k_num` distinct experts out of
85	`num_experts`, and assign normalized random weights that sum to 1.
86
87	Args:
88	num_tokens (int): Total number of tokens.
89	num_experts (int): Total number of available experts.
90	top_k_num (int): Number of experts to select per token.
91
92	Returns:
93	expert_indices (torch.Tensor): shape [num_tokens, top_k_num],
94	expert index for each token.
95	expert_weights (torch.Tensor): shape [num_tokens, top_k_num],
96	normalized weights (sum = 1 per row).
97
98	"""
99	assert top_k_num <= num_experts, "top_k_num must be <= num_experts"
100
101	# Randomly select top_k_num distinct experts for each token
102	expert_indices = torch.empty((num_tokens, top_k_num), dtype=torch.int32)
103	for i in range(num_tokens):
104	# Randomly choose unique expert indices
105	selected = torch.randperm(num_experts)[:top_k_num]
106	expert_indices[i] = selected
107
108	# Generate random weights and normalize along dim=1
109	expert_weights = torch.rand((num_tokens, top_k_num), dtype=torch.float32)
110	expert_weights = expert_weights / expert_weights.sum(dim=1, keepdim=True)
111
112	return expert_indices, expert_weights
113
114
115	def sample_data(
116	num_tokens: int,
117	num_sequences: int,
118	max_loras: int,
119	num_experts: int,
120	top_k_num: int,
121	):
122	topk_ids, topk_weights = assign_experts_to_tokens(
123	num_tokens, num_experts, top_k_num
124	)
125	token_lora_mapping = assign_loras_to_tokens(num_tokens, num_sequences, max_loras)
126	active_lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
127	lora_ids = torch.unique(token_lora_mapping, sorted=True)
128	active_lora_ids[: lora_ids.size(0)].copy_(lora_ids, non_blocking=True)
129	return topk_ids, topk_weights, token_lora_mapping, active_lora_ids
130
131
132	def use_fused_moe_lora_kernel(
133	topk_ids,
134	topk_weights,
135	token_lora_mapping,
136	max_lora_rank,
137	top_k_num,
138	lora_ids,
139	lora_a_stacked,
140	lora_b_stacked,
141	hidden_states,
142	output,
143	max_loras,
144	num_experts,
145	block_size,
146	fully_sharded=False,
147	offset=0,
148	add_inputs=True,
149	):
150	max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
151	max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
152	max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
153
154	# init output tensors
155	sorted_token_ids = torch.empty(
156	(max_loras * max_num_tokens_padded,),
157	dtype=torch.int32,
158	)
159	expert_ids = torch.empty((max_loras * max_num_m_blocks,), dtype=torch.int32)
160	num_tokens_post_padded = torch.empty((max_loras,), dtype=torch.int32)
161	adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
162
163	# call kernel
164	ops.moe_lora_align_block_size(
165	topk_ids,
166	token_lora_mapping,
167	num_experts,
168	block_size,
169	max_loras,
170	max_num_tokens_padded,
171	max_num_m_blocks,
172	sorted_token_ids,
173	expert_ids,
174	num_tokens_post_padded,
175	adapter_enabled,
176	lora_ids,
177	)
178
179	config = {
180	"BLOCK_SIZE_M": block_size,
181	"BLOCK_SIZE_N": 32,
182	"BLOCK_SIZE_K": 64,
183	"GROUP_SIZE_M": 1,
184	"NUM_WARPS": 4,
185	"NUM_STAGES": 3,
186	"SPLIT_K": 1,
187	}
188
189	mul_routed_weight = False
190	expert_ids = expert_ids.view(max_loras, -1)
191	sorted_token_ids = sorted_token_ids.view(max_loras, -1)
192
193	# num_active_loras is the number of active LoRAs
194	# (max_loras + 1 to include no-lora case)
195	# Stored as CPU tensor to match the kernel API (torch.compile compatibility)
196	num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
197
198	fused_moe_lora(
199	output,
200	hidden_states,
201	lora_a_stacked,
202	lora_b_stacked,
203	topk_weights,
204	sorted_token_ids,
205	expert_ids,
206	num_tokens_post_padded,
207	token_lora_mapping,
208	max_lora_rank,
209	top_k_num,
210	lora_ids,
211	num_active_loras,
212	adapter_enabled,
213	config["BLOCK_SIZE_M"],
214	config["BLOCK_SIZE_N"],
215	config["BLOCK_SIZE_K"],
216	config["GROUP_SIZE_M"],
217	config["NUM_WARPS"],
218	config["NUM_STAGES"],
219	config["SPLIT_K"],
220	config["BLOCK_SIZE_M"],
221	config["BLOCK_SIZE_N"],
222	config["BLOCK_SIZE_K"],
223	config["GROUP_SIZE_M"],
224	config["NUM_WARPS"],
225	config["NUM_STAGES"],
226	config["SPLIT_K"],
227	mul_routed_weight,
228	fully_sharded=fully_sharded,
229	offset=offset,
230	add_inputs=add_inputs,
231	)
232
233
234	def use_torch(
235	hidden_states,
236	token_lora_mapping,
237	topk_ids,
238	lora_a_stacked,
239	lora_b_stacked,
240	top_k_num,
241	num_slices=1,
242	):
243	outputs = []
244	for i in range(hidden_states.shape[0]):
245	slice_tensors = []
246	for slice_id in range(num_slices):
247	lora_idx = token_lora_mapping[i]
248	expert_ids = topk_ids[i]
249	lora_a = lora_a_stacked[slice_id][lora_idx][expert_ids]
250	lora_b = lora_b_stacked[slice_id][lora_idx][expert_ids]
251	tensors = [
252	hidden_states[i] @ lora_a[x].T @ lora_b[x].T for x in range(top_k_num)
253	]
254	slice_tensors.append(torch.stack(tensors, dim=0))
255
256	outputs.append(torch.concat(slice_tensors, dim=-1))
257	return torch.stack(outputs, dim=0)
258
259
260	DEVICE_TYPE = current_platform.device_type
261	DTYPES = [torch.float16, torch.bfloat16]
262	DEVICES = [f"{DEVICE_TYPE}:{0}"]
263	SEED = [42]
264
265
266	@pytest.mark.parametrize("num_tokens", [100])
267	@pytest.mark.parametrize("top_k_num", [6, 12])
268	@pytest.mark.parametrize("num_experts", [64])
269	@pytest.mark.parametrize("max_loras", [4, 6, 16])
270	@pytest.mark.parametrize("N", [1408])
271	@pytest.mark.parametrize("K", [2048])
272	@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
273	@pytest.mark.parametrize("block_size", [16])
274	@pytest.mark.parametrize("num_slices", [1, 2])
275	@pytest.mark.parametrize("dtype", DTYPES)
276	@pytest.mark.parametrize("device", DEVICES)
277	@pytest.mark.parametrize("seed", SEED)
278	def test_fused_moe_lora_kernel(
279	num_tokens,
280	top_k_num,
281	num_experts,
282	max_loras,
283	N,
284	K,
285	max_lora_rank,
286	block_size,
287	num_slices,
288	dtype,
289	device,
290	seed,
291	):
292	torch.set_default_device(device)
293	set_random_seed(seed)
294	# the number of randomly generated sentences.
295	num_sequences = 10
296	# generate data
297	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
298	num_tokens, num_sequences, max_loras, num_experts, top_k_num
299	)
300
301	# init lora weights
302	lora_a_stacked = [
303	torch.rand(
304	(
305	max_loras,
306	num_experts,
307	max_lora_rank,
308	K,
309	),
310	dtype=dtype,
311	)
312	for _ in range(num_slices)
313	]
314	lora_b_stacked = [
315	torch.rand(
316	(
317	max_loras,
318	num_experts,
319	N // num_slices,
320	max_lora_rank,
321	),
322	dtype=dtype,
323	)
324	for _ in range(num_slices)
325	]
326	hidden_states = torch.rand(
327	(
328	num_tokens,
329	K,
330	),
331	dtype=dtype,
332	)
333
334	# fused_moe_lora_kernel output
335	output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
336	use_fused_moe_lora_kernel(
337	topk_ids,
338	topk_weights,
339	token_lora_mapping,
340	max_lora_rank,
341	top_k_num,
342	lora_ids,
343	lora_a_stacked,
344	lora_b_stacked,
345	hidden_states,
346	output,
347	max_loras,
348	num_experts,
349	block_size,
350	)
351	# pytorch output
352	output2 = use_torch(
353	hidden_states,
354	token_lora_mapping,
355	topk_ids,
356	lora_a_stacked,
357	lora_b_stacked,
358	top_k_num,
359	num_slices,
360	)
361
362	torch.testing.assert_close(output, output2, atol=1e-2, rtol=1e-2)
363
364
365	@pytest.mark.parametrize("max_lora_rank", [128, 256])
366	@pytest.mark.parametrize("add_inputs", [True, False])
367	@pytest.mark.parametrize("num_slices", [1, 2])
368	@pytest.mark.parametrize("device", DEVICES)
369	def test_fused_moe_lora_kernel_large_rank_fallback(
370	max_lora_rank,
371	add_inputs,
372	num_slices,
373	device,
374	):
375	"""Ranks above the one-shot limit fall back to shrink/expand kernels."""
376	torch.set_default_device(device)
377	set_random_seed(42)
378
379	num_tokens = 8
380	num_sequences = 4
381	top_k_num = 2
382	num_experts = 4
383	max_loras = 2
384	N = 256
385	K = 256
386	block_size = 16
387	dtype = torch.bfloat16
388
389	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
390	num_tokens, num_sequences, max_loras, num_experts, top_k_num
391	)
392	lora_a_stacked = [
393	torch.rand(
394	(max_loras, num_experts, max_lora_rank, K),
395	dtype=dtype,
396	)
397	for _ in range(num_slices)
398	]
399	lora_b_stacked = [
400	torch.rand(
401	(max_loras, num_experts, N // num_slices, max_lora_rank),
402	dtype=dtype,
403	)
404	for _ in range(num_slices)
405	]
406	hidden_states = torch.rand((num_tokens, K), dtype=dtype)
407
408	residual = torch.randn((num_tokens, top_k_num, N), dtype=dtype)
409	output = residual.clone()
410	use_fused_moe_lora_kernel(
411	topk_ids,
412	topk_weights,
413	token_lora_mapping,
414	max_lora_rank,
415	top_k_num,
416	lora_ids,
417	lora_a_stacked,
418	lora_b_stacked,
419	hidden_states,
420	output,
421	max_loras,
422	num_experts,
423	block_size,
424	add_inputs=add_inputs,
425	)
426
427	expected = use_torch(
428	hidden_states,
429	token_lora_mapping,
430	topk_ids,
431	lora_a_stacked,
432	lora_b_stacked,
433	top_k_num,
434	num_slices,
435	)
436	if add_inputs:
437	expected += residual
438
439	torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
440
441
442	def use_fused_moe_lora_kernel_naive(
443	topk_ids,
444	topk_weights,
445	token_lora_mapping,
446	max_lora_rank,
447	top_k_num,
448	lora_ids,
449	lora_a_stacked,
450	lora_b_stacked,
451	hidden_states,
452	output,
453	max_loras,
454	block_size,
455	fully_sharded=False,
456	offset=0,
457	add_inputs=True,
458	):
459	"""Test helper for naive_block_assignment path.
460	Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids.
461	"""
462	config = {
463	"BLOCK_SIZE_M": block_size,
464	"BLOCK_SIZE_N": 32,
465	"BLOCK_SIZE_K": 64,
466	"GROUP_SIZE_M": 1,
467	"NUM_WARPS": 4,
468	"NUM_STAGES": 3,
469	"SPLIT_K": 1,
470	}
471
472	mul_routed_weight = False
473
474	# In naive mode:
475	# - expert_ids = topk_ids.view(-1), shape: (num_tokens * top_k,)
476	# - sorted_token_ids = None
477	# - num_tokens_post_padded = None
478	expert_ids = topk_ids.reshape(-1)
479	sorted_token_ids = None
480	num_tokens_post_padded = None
481
482	adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
483
484	# num_active_loras is the number of active LoRAs
485	# (max_loras + 1 to include no-lora case)
486	# Stored as CPU tensor to match the kernel API (torch.compile compatibility)
487	num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
488
489	fused_moe_lora(
490	output,
491	hidden_states,
492	lora_a_stacked,
493	lora_b_stacked,
494	topk_weights,
495	sorted_token_ids,
496	expert_ids,
497	num_tokens_post_padded,
498	token_lora_mapping,
499	max_lora_rank,
500	top_k_num,
501	lora_ids,
502	num_active_loras,
503	adapter_enabled,
504	config["BLOCK_SIZE_M"],
505	config["BLOCK_SIZE_N"],
506	config["BLOCK_SIZE_K"],
507	config["GROUP_SIZE_M"],
508	config["NUM_WARPS"],
509	config["NUM_STAGES"],
510	config["SPLIT_K"],
511	config["BLOCK_SIZE_M"],
512	config["BLOCK_SIZE_N"],
513	config["BLOCK_SIZE_K"],
514	config["GROUP_SIZE_M"],
515	config["NUM_WARPS"],
516	config["NUM_STAGES"],
517	config["SPLIT_K"],
518	mul_routed_weight=mul_routed_weight,
519	fully_sharded=fully_sharded,
520	offset=offset,
521	add_inputs=add_inputs,
522	)
523
524
525	@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8])
526	@pytest.mark.parametrize("top_k_num", [1, 2])
527	@pytest.mark.parametrize("num_experts", [64, 128])
528	@pytest.mark.parametrize("max_loras", [4, 8])
529	@pytest.mark.parametrize("N", [1408])
530	@pytest.mark.parametrize("K", [2048])
531	@pytest.mark.parametrize("max_lora_rank", [16, 32])
532	@pytest.mark.parametrize("block_size", [16])
533	@pytest.mark.parametrize("num_slices", [1, 2])
534	@pytest.mark.parametrize("dtype", DTYPES)
535	@pytest.mark.parametrize("device", DEVICES)
536	@pytest.mark.parametrize("seed", SEED)
537	def test_fused_moe_lora_kernel_naive_block_assignment(
538	num_tokens,
539	top_k_num,
540	num_experts,
541	max_loras,
542	N,
543	K,
544	max_lora_rank,
545	block_size,
546	num_slices,
547	dtype,
548	device,
549	seed,
550	):
551	"""Test the naive_block_assignment path of the fused_moe_lora kernel.
552	This path is triggered when batch_size * top_k is much smaller than
553	num_experts * max_loras, and skips the moe_lora_align_block_size kernel.
554	"""
555	torch.set_default_device(device)
556	set_random_seed(seed)
557
558	# Verify this configuration would trigger naive_block_assignment
559	# (num_tokens * top_k * SPARSITY_FACTOR <= num_experts * max_loras)
560	SPARSITY_FACTOR = 8
561	assert num_tokens * top_k_num * SPARSITY_FACTOR <= num_experts * max_loras, (
562	f"Test configuration doesn't meet naive_block_assignment condition: "
563	f"{num_tokens} * {top_k_num} * {SPARSITY_FACTOR} > {num_experts} * {max_loras}"
564	)
565
566	# the number of randomly generated sentences.
567	num_sequences = min(num_tokens, 4)
568	# generate data
569	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
570	num_tokens, num_sequences, max_loras, num_experts, top_k_num
571	)
572
573	# init lora weights
574	lora_a_stacked = [
575	torch.rand(
576	(
577	max_loras,
578	num_experts,
579	max_lora_rank,
580	K,
581	),
582	dtype=dtype,
583	)
584	for _ in range(num_slices)
585	]
586	lora_b_stacked = [
587	torch.rand(
588	(
589	max_loras,
590	num_experts,
591	N // num_slices,
592	max_lora_rank,
593	),
594	dtype=dtype,
595	)
596	for _ in range(num_slices)
597	]
598	hidden_states = torch.rand(
599	(
600	num_tokens,
601	K,
602	),
603	dtype=dtype,
604	)
605
606	# fused_moe_lora_kernel output (naive path)
607	output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
608	use_fused_moe_lora_kernel_naive(
609	topk_ids,
610	topk_weights,
611	token_lora_mapping,
612	max_lora_rank,
613	top_k_num,
614	lora_ids,
615	lora_a_stacked,
616	lora_b_stacked,
617	hidden_states,
618	output,
619	max_loras,
620	block_size,
621	)
622
623	# pytorch reference output
624	output_ref = use_torch(
625	hidden_states,
626	token_lora_mapping,
627	topk_ids,
628	lora_a_stacked,
629	lora_b_stacked,
630	top_k_num,
631	num_slices,
632	)
633
634	torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
635
636
637	@multi_gpu_test(num_gpus=2)
638	@pytest.mark.parametrize("num_tokens", [100])
639	@pytest.mark.parametrize("top_k_num", [6])
640	@pytest.mark.parametrize("num_experts", [64])
641	@pytest.mark.parametrize("max_loras", [4])
642	@pytest.mark.parametrize("N", [1408])
643	@pytest.mark.parametrize("K", [2048])
644	@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
645	@pytest.mark.parametrize("block_size", [16])
646	@pytest.mark.parametrize("dtype", DTYPES)
647	@pytest.mark.parametrize("seed", SEED)
648	@pytest.mark.parametrize("column_parallel", [True, False])
649	def test_fused_moe_lora_kernel_fully_sharded(
650	num_tokens,
651	top_k_num,
652	num_experts,
653	max_loras,
654	N,
655	K,
656	max_lora_rank,
657	block_size,
658	dtype,
659	seed,
660	column_parallel,
661	):
662	set_random_seed(seed)
663	# the number of randomly generated sentences.
664	num_sequences = 10
665	# generate data
666	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
667	num_tokens, num_sequences, max_loras, num_experts, top_k_num
668	)
669
670	def run_torch_spawn(fn, nprocs):
671	torch.multiprocessing.spawn(
672	fn,
673	args=(
674	nprocs,
675	f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
676	dtype,
677	seed,
678	N,
679	K,
680	num_tokens,
681	topk_ids,
682	topk_weights,
683	token_lora_mapping,
684	max_lora_rank,
685	top_k_num,
686	lora_ids,
687	max_loras,
688	num_experts,
689	block_size,
690	column_parallel,
691	),
692	nprocs=nprocs,
693	)
694
695	run_torch_spawn(use_fused_moe_lora_kernel_tensor_parallel, nprocs=2)
696
697
698	def use_fused_moe_lora_kernel_tensor_parallel(
699	local_rank,
700	world_size,
701	init_method,
702	dtype,
703	seed,
704	N,
705	K,
706	num_tokens,
707	topk_ids,
708	topk_weights,
709	token_lora_mapping,
710	max_lora_rank,
711	top_k_num,
712	lora_ids,
713	max_loras,
714	num_experts,
715	block_size,
716	column_parallel,
717	):
718	def _get_shard_slice(shard_size):
719	return slice(local_rank * shard_size, (local_rank + 1) * shard_size)
720
721	set_random_seed(seed)
722
723	device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
724	torch.accelerator.set_device_index(device)
725	torch.set_default_device(device)
726	torch.set_default_dtype(dtype)
727
728	init_distributed_environment(
729	world_size=world_size,
730	rank=local_rank,
731	local_rank=local_rank,
732	distributed_init_method=init_method,
733	backend=current_platform.dist_backend,
734	)
735	with ensure_current_vllm_config():
736	initialize_model_parallel(world_size, 1)
737	tp_size = get_tensor_model_parallel_world_size()
738
739	input_dim = K if column_parallel else N
740	output_dim = N if column_parallel else K
741
742	# init lora weights
743	lora_a = torch.rand(
744	(
745	max_loras,
746	num_experts,
747	max_lora_rank,
748	input_dim,
749	),
750	dtype=dtype,
751	)
752	lora_b = torch.rand(
753	(
754	max_loras,
755	num_experts,
756	output_dim,
757	max_lora_rank,
758	),
759	dtype=dtype,
760	)
761
762	hidden_states = torch.rand(
763	(
764	num_tokens,
765	input_dim,
766	),
767	dtype=dtype,
768	)
769
770	output = torch.zeros((num_tokens, top_k_num, output_dim), dtype=dtype)
771	topk_ids = topk_ids.to(device)
772	topk_weights = topk_weights.to(device)
773	token_lora_mapping = token_lora_mapping.to(device)
774	lora_ids = lora_ids.to(device)
775
776	ref_output = use_torch(
777	hidden_states,
778	token_lora_mapping,
779	topk_ids,
780	[lora_a],
781	[lora_b],
782	top_k_num,
783	)
784
785	if column_parallel:
786	# Column parallel (e.g. gate_up_proj): LoRA A is sliced along the rank dim,
787	# and Lora B is sliced along the output dim
788	lora_a_shard_size = max_lora_rank // tp_size
789	lora_a = lora_a[:, :, _get_shard_slice(lora_a_shard_size), :]
790	max_lora_rank = lora_a_shard_size
791	offset = 0
792
793	lora_b_shard_size = output_dim // tp_size
794	lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]
795	output = output[:, :, _get_shard_slice(lora_b_shard_size)].contiguous()
796	else:
797	# Row parallel (e.g. down proj): LoRA A is sliced along the input dim,
798	# and LoRA B is sliced along the output dim
799	lora_a_shard_size = input_dim // tp_size
800	lora_a = lora_a[:, :, :, _get_shard_slice(lora_a_shard_size)]
801	hidden_states = hidden_states[:, _get_shard_slice(lora_a_shard_size)]
802
803	lora_b_shard_size = output_dim // tp_size
804	lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]
805	offset = lora_b_shard_size * local_rank
806
807	use_fused_moe_lora_kernel(
808	topk_ids,
809	topk_weights,
810	token_lora_mapping,
811	max_lora_rank,
812	top_k_num,
813	lora_ids,
814	[lora_a],
815	[lora_b],
816	hidden_states,
817	output,
818	max_loras,
819	num_experts,
820	block_size,
821	fully_sharded=True,
822	offset=offset,
823	)
824
825	if column_parallel:
826	output = tensor_model_parallel_all_gather(output)
827	else:
828	output = tensor_model_parallel_all_reduce(output)
829
830	torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
831
832
833	# -- one-shot fast-path coverage --------------------------------------------
834	# The fused shrink+expand one-shot kernel pads `BLOCK_R` to next_pow2(rank),
835	# with a floor of 16 (tensor-core minimum). Small ranks (4, 8) exercise the
836	# rank-dim masking and are not covered by the original tests, which start at
837	# rank=16. The legacy two-kernel path additionally fails on rank=4 in TMA
838	# mode because the rank-dim stride (rank * elem_size) is not 16-byte
839	# aligned; the one-shot fast path takes precedence whenever fully_sharded
840	# is False so this regression is hidden in normal use, but the test still
841	# ensures the one-shot logic is correct against the pytorch reference.
842
843
844	@pytest.mark.parametrize("num_tokens", [16, 100])
845	@pytest.mark.parametrize("top_k_num", [2])
846	@pytest.mark.parametrize("num_experts", [8, 64])
847	@pytest.mark.parametrize("max_loras", [4])
848	@pytest.mark.parametrize("N", [1408])
849	@pytest.mark.parametrize("K", [2048])
850	@pytest.mark.parametrize("max_lora_rank", [4, 8])
851	@pytest.mark.parametrize("block_size", [16, 64])
852	@pytest.mark.parametrize("num_slices", [1, 2])
853	@pytest.mark.parametrize("dtype", DTYPES)
854	@pytest.mark.parametrize("device", DEVICES)
855	@pytest.mark.parametrize("seed", SEED)
856	def test_fused_moe_lora_kernel_small_rank(
857	num_tokens,
858	top_k_num,
859	num_experts,
860	max_loras,
861	N,
862	K,
863	max_lora_rank,
864	block_size,
865	num_slices,
866	dtype,
867	device,
868	seed,
869	):
870	"""One-shot fast path covering rank<16 (padded to BLOCK_R=16 inside kernel)."""
871	torch.set_default_device(device)
872	set_random_seed(seed)
873	num_sequences = max(1, min(num_tokens, 8))
874	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
875	num_tokens, num_sequences, max_loras, num_experts, top_k_num
876	)
877
878	lora_a_stacked = [
879	torch.rand(
880	(max_loras, num_experts, max_lora_rank, K),
881	dtype=dtype,
882	)
883	for _ in range(num_slices)
884	]
885	lora_b_stacked = [
886	torch.rand(
887	(max_loras, num_experts, N // num_slices, max_lora_rank),
888	dtype=dtype,
889	)
890	for _ in range(num_slices)
891	]
892	hidden_states = torch.rand((num_tokens, K), dtype=dtype)
893
894	output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
895	use_fused_moe_lora_kernel(
896	topk_ids,
897	topk_weights,
898	token_lora_mapping,
899	max_lora_rank,
900	top_k_num,
901	lora_ids,
902	lora_a_stacked,
903	lora_b_stacked,
904	hidden_states,
905	output,
906	max_loras,
907	num_experts,
908	block_size,
909	)
910	output_ref = use_torch(
911	hidden_states,
912	token_lora_mapping,
913	topk_ids,
914	lora_a_stacked,
915	lora_b_stacked,
916	top_k_num,
917	num_slices,
918	)
919	torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
920
921
922	@pytest.mark.parametrize("num_tokens", [16, 64])
923	@pytest.mark.parametrize("top_k_num", [2])
924	@pytest.mark.parametrize("num_experts", [8])
925	@pytest.mark.parametrize("max_loras", [4])
926	@pytest.mark.parametrize("N", [2048])
927	@pytest.mark.parametrize("K", [4096])
928	@pytest.mark.parametrize("max_lora_rank", [8, 16, 32, 64])
929	@pytest.mark.parametrize("block_size", [64])
930	@pytest.mark.parametrize("num_slices", [2])
931	@pytest.mark.parametrize("dtype", [torch.bfloat16])
932	@pytest.mark.parametrize("device", DEVICES)
933	@pytest.mark.parametrize("seed", SEED)
934	def test_fused_moe_lora_kernel_npid_path(
935	num_tokens,
936	top_k_num,
937	num_experts,
938	max_loras,
939	N,
940	K,
941	max_lora_rank,
942	block_size,
943	num_slices,
944	dtype,
945	device,
946	seed,
947	):
948	"""Exercise the small-batch / NPID > 1 branch of the one-shot fast path.
949
950	With these sizes the one-shot wrapper computes NPID_FACTOR > 1 (base CTA count
951	< SM count), so each program covers only an outer chunk of N. The
952	cross-outer-block write mask is the correctness-critical bit.
953	"""
954	torch.set_default_device(device)
955	set_random_seed(seed)
956	num_sequences = max(1, min(num_tokens, 4))
957	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
958	num_tokens, num_sequences, max_loras, num_experts, top_k_num
959	)
960
961	lora_a_stacked = [
962	torch.rand(
963	(max_loras, num_experts, max_lora_rank, K),
964	dtype=dtype,
965	)
966	for _ in range(num_slices)
967	]
968	lora_b_stacked = [
969	torch.rand(
970	(max_loras, num_experts, N // num_slices, max_lora_rank),
971	dtype=dtype,
972	)
973	for _ in range(num_slices)
974	]
975	hidden_states = torch.rand((num_tokens, K), dtype=dtype)
976
977	output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
978	use_fused_moe_lora_kernel(
979	topk_ids,
980	topk_weights,
981	token_lora_mapping,
982	max_lora_rank,
983	top_k_num,
984	lora_ids,
985	lora_a_stacked,
986	lora_b_stacked,
987	hidden_states,
988	output,
989	max_loras,
990	num_experts,
991	block_size,
992	)
993	output_ref = use_torch(
994	hidden_states,
995	token_lora_mapping,
996	topk_ids,
997	lora_a_stacked,
998	lora_b_stacked,
999	top_k_num,
1000	num_slices,
1001	)
1002	torch.testing.assert_close(output, output_ref, atol=2e-2, rtol=2e-2)
1003
1004
1005	# -- one-shot corner-case coverage ------------------------------------------
1006	# Each of the following exercises a path where the kernel is launched but
1007	# every program early-exits, leaving the output unchanged. The contract is
1008	# additive (`output += contribution`), so an empty contribution must leave
1009	# the input residual untouched.
1010
1011
1012	def _build_one_shot_inputs(
1013	num_tokens,
1014	top_k_num,
1015	num_experts,
1016	max_loras,
1017	max_lora_rank,
1018	K,
1019	N,
1020	num_slices,
1021	block_size,
1022	dtype,
1023	):
1024	"""Common scaffolding for the corner-case tests below."""
1025	num_sequences = max(1, min(num_tokens, 4)) if num_tokens > 0 else 1
1026	if num_tokens > 0:
1027	topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
1028	num_tokens, num_sequences, max_loras, num_experts, top_k_num
1029	)
1030	else:
1031	# M=0 path: caller may still hand us empty tensors with the right shape.
1032	topk_ids = torch.empty((0, top_k_num), dtype=torch.int32)
1033	topk_weights = torch.empty((0, top_k_num), dtype=torch.float32)
1034	token_lora_mapping = torch.empty((0,), dtype=torch.int32)
1035	lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
1036
1037	lora_a_stacked = [
1038	torch.rand((max_loras, num_experts, max_lora_rank, K), dtype=dtype)
1039	for _ in range(num_slices)
1040	]
1041	lora_b_stacked = [
1042	torch.rand(
1043	(max_loras, num_experts, N // num_slices, max_lora_rank), dtype=dtype
1044	)
1045	for _ in range(num_slices)
1046	]
1047	hidden_states = torch.rand((max(num_tokens, 0), K), dtype=dtype)
1048	return (
1049	topk_ids,
1050	topk_weights.to(dtype),
1051	token_lora_mapping,
1052	lora_ids,
1053	lora_a_stacked,
1054	lora_b_stacked,
1055	hidden_states,
1056	)
1057
1058
1059	def _call_one_shot(
1060	output,
1061	hidden_states,
1062	lora_a_stacked,
1063	lora_b_stacked,
1064	topk_weights,
1065	sorted_token_ids,
1066	expert_ids,
1067	num_tokens_post_padded,
1068	token_lora_mapping,
1069	max_lora_rank,
1070	top_k_num,
1071	lora_ids,
1072	num_active_loras,
1073	adapter_enabled,
1074	block_size,
1075	add_inputs=True,
1076	):
1077	"""Direct call into fused_moe_lora with one-shot-routed defaults."""
1078	from vllm.lora.ops.triton_ops import fused_moe_lora as _op
1079
1080	_op(
1081	output,
1082	hidden_states,
1083	lora_a_stacked,
1084	lora_b_stacked,
1085	topk_weights,
1086	sorted_token_ids,
1087	expert_ids,
1088	num_tokens_post_padded,
1089	token_lora_mapping,
1090	max_lora_rank,
1091	top_k_num,
1092	lora_ids,
1093	num_active_loras,
1094	adapter_enabled,
1095	block_size,
1096	32,
1097	64,
1098	1,
1099	4,
1100	3,
1101	1,
1102	block_size,
1103	32,
1104	64,
1105	1,
1106	4,
1107	3,
1108	1,
1109	False,
1110	False,
1111	0,
1112	add_inputs,
1113	)
1114
1115
1116	@pytest.mark.parametrize(
1117	"trigger",
1118	["sorted_lora_ids_neg", "naive_mapping_neg", "naive_all_disabled"],
1119	)
1120	@pytest.mark.parametrize("device", DEVICES)
1121	def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device):
1122	"""one-shot must leave the residual byte-identical when every program
1123	must early-exit. Three trigger conditions are covered:
1124
1125	- "sorted_lora_ids_neg": sorted path, lora_ids all -1 (lora_id<0 check)
1126	- "naive_mapping_neg": naive path, token_lora_mapping all -1
1127	- "naive_all_disabled": naive path, adapter_enabled all 0
1128	"""
1129	torch.set_default_device(device)
1130	set_random_seed(0)
1131
1132	# Per-trigger shapes: naive_mapping_neg needs the naive dispatch gate
1133	# `num_tokenstop_k8 <= num_experts*max_loras` to hold, hence the
1134	# larger E/max_loras and smaller num_tokens.
1135	if trigger == "naive_mapping_neg":
1136	num_tokens, top_k, E, max_loras, R = 8, 2, 64, 8, 16
1137	elif trigger == "naive_all_disabled":
1138	num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 32
1139	else: # sorted_lora_ids_neg
1140	num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 16
1141	K, N = 1024, 1024
1142	block_size, num_slices, dtype = 16, 2, torch.bfloat16
1143
1144	(
1145	topk_ids,
1146	topk_weights,
1147	token_lora_mapping,
1148	lora_ids,
1149	lora_a_stacked,
1150	lora_b_stacked,
1151	hidden_states,
1152	) = _build_one_shot_inputs(
1153	num_tokens, top_k, E, max_loras, R, K, N, num_slices, block_size, dtype
1154	)
1155
1156	adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
1157	num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
1158
1159	if trigger == "sorted_lora_ids_neg":
1160	lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
1161	max_pad = topk_ids.numel() + E * (block_size - 1)
1162	max_pad = round_up(max_pad, block_size)
1163	max_blocks = CEILDIV(max_pad, block_size)
1164	sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)
1165	expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)
1166	num_post = torch.zeros((max_loras,), dtype=torch.int32)
1167	else:
1168	sorted_token_ids = None
1169	expert_ids = topk_ids.reshape(-1).contiguous()
1170	num_post = None
1171	if trigger == "naive_mapping_neg":
1172	token_lora_mapping = torch.full((num_tokens,), -1, dtype=torch.int32)
1173	lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
1174	else: # naive_all_disabled
1175	adapter_enabled = torch.zeros(max_loras + 1, dtype=torch.int32)
1176
1177	residual = torch.randn((num_tokens, top_k, N), dtype=dtype) * 0.1
1178	output = residual.clone()
1179
1180	_call_one_shot(
1181	output,
1182	hidden_states,
1183	lora_a_stacked,
1184	lora_b_stacked,
1185	topk_weights,
1186	sorted_token_ids,
1187	expert_ids,
1188	num_post,
1189	token_lora_mapping,
1190	R,
1191	top_k,
1192	lora_ids,
1193	num_active_loras,
1194	adapter_enabled,
1195	block_size,
1196	)
1197	torch.testing.assert_close(output, residual, atol=0, rtol=0)
1198
1199
1200	@pytest.mark.parametrize("device", DEVICES)
1201	def test_fused_moe_lora_kernel_zero_grid_no_crash(device):
1202	"""num_active_loras=0 (or num_slices=0) would otherwise launch a grid
1203	with a zero dimension. one-shot wrapper must short-circuit before launch."""
1204	torch.set_default_device(device)
1205	set_random_seed(0)
1206	num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024
1207	block_size, num_slices, dtype = 16, 2, torch.bfloat16
1208
1209	(
1210	topk_ids,
1211	topk_weights,
1212	token_lora_mapping,
1213	lora_ids,
1214	lora_a_stacked,
1215	lora_b_stacked,
1216	hidden_states,
1217	) = _build_one_shot_inputs(
1218	num_tokens,
1219	top_k,
1220	E,
1221	max_loras,
1222	R,
1223	K,
1224	N,
1225	num_slices,
1226	block_size,
1227	dtype,
1228	)
1229	adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
1230	num_active_loras = torch.tensor([0], dtype=torch.int32, device="cpu")
1231	residual = torch.randn((num_tokens, top_k, N), dtype=dtype) * 0.1
1232	output = residual.clone()
1233
1234	# sorted path is the one that uses num_active_loras for grid axis 2
1235	max_pad = topk_ids.numel() + E * (block_size - 1)
1236	max_pad = round_up(max_pad, block_size)
1237	max_blocks = CEILDIV(max_pad, block_size)
1238	sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)
1239	expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)
1240	num_post = torch.zeros((max_loras,), dtype=torch.int32)
1241	_call_one_shot(
1242	output,
1243	hidden_states,
1244	lora_a_stacked,
1245	lora_b_stacked,
1246	topk_weights,
1247	sorted_token_ids,
1248	expert_ids,
1249	num_post,
1250	token_lora_mapping,
1251	R,
1252	top_k,
1253	lora_ids,
1254	num_active_loras,
1255	adapter_enabled,
1256	block_size,
1257	)
1258	torch.testing.assert_close(output, residual, atol=0, rtol=0)
1259
1260
1261	@pytest.mark.parametrize("device", DEVICES)
1262	def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device):
1263	"""one-shot must surface a clear assertion when shrink_block_size_m is not
1264	a power of 2 / less than 16, instead of the cryptic Triton compile
1265	failure (`arange's range must be a power of 2`)."""
1266	torch.set_default_device(device)
1267	set_random_seed(0)
1268	num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024
1269	num_slices, dtype = 2, torch.bfloat16
1270	block_size = 24 # NOT a power of 2
1271
1272	(
1273	topk_ids,
1274	topk_weights,
1275	token_lora_mapping,
1276	lora_ids,
1277	lora_a_stacked,
1278	lora_b_stacked,
1279	hidden_states,
1280	) = _build_one_shot_inputs(
1281	num_tokens,
1282	top_k,
1283	E,
1284	max_loras,
1285	R,
1286	K,
1287	N,
1288	num_slices,
1289	16,
1290	dtype,
1291	)
1292	# Build sorted-mode metadata at block_size=16 so shapes are sane,
1293	# but pass block_size=24 to the op (the buggy combination).
1294	max_pad = topk_ids.numel() + E * (16 - 1)
1295	max_pad = round_up(max_pad, 16)
1296	max_blocks = CEILDIV(max_pad, 16)
1297	sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)
1298	expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)
1299	num_post = torch.zeros((max_loras,), dtype=torch.int32)
1300	adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
1301	num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
1302	output = torch.zeros((num_tokens, top_k, N), dtype=dtype)
1303
1304	with pytest.raises(AssertionError, match="shrink_block_size_m"):
1305	_call_one_shot(
1306	output,
1307	hidden_states,
1308	lora_a_stacked,
1309	lora_b_stacked,
1310	topk_weights,
1311	sorted_token_ids,
1312	expert_ids,
1313	num_post,
1314	token_lora_mapping,
1315	R,
1316	top_k,
1317	lora_ids,
1318	num_active_loras,
1319	adapter_enabled,
1320	block_size,
1321	)