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
[Performance][LoRA] add context varying params to 'do_not_specialize' in fused moe lora (#27445)
Signed-off-by: gnovack <gnovack@amazon.com>
gnovack committed
Oct 26, 2025 at 23:31 UTC
a806c14cc7cf465bb30f34cd9782bd28ae3eb08e
1 file changed
+16
-7
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+16
-7
@@ -31,7 +31,16 @@ def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
31
return _LORA_PTR_DICT.get(key)
32
33
34
-@triton.jit
34
+@triton.jit(
35
+ do_not_specialize=[
36
+ "num_valid_tokens",
37
+ "EM",
38
+ "stride_tl",
39
+ "stride_el",
40
+ "slice_a_size",
41
+ "slice_c_size",
42
+ ]
43
+)
44
def _fused_moe_lora_kernel(
45
a_ptr,
46
b_ptr,
@@ -60,11 +69,11 @@ def _fused_moe_lora_kernel(
69
stride_cn,
70
stride_tl,
71
stride_el,
72
+ slice_a_size,
73
+ slice_c_size,
74
# Meta-parameters
75
num_slice_a: tl.constexpr,
76
num_slice_c: tl.constexpr,
66
- slice_a_size: tl.constexpr,
67
- slice_c_size: tl.constexpr,
77
top_k: tl.constexpr,
78
MUL_ROUTED_WEIGHT: tl.constexpr,
79
BLOCK_SIZE_M: tl.constexpr,
@@ -256,10 +265,10 @@ def _fused_moe_lora(
265
a_intermediate_cache1.stride(3),
266
sorted_token_ids.stride(0),
267
expert_ids.stride(0),
259
- num_slice_a=1,
260
- num_slice_c=num_slices,
268
slice_a_size=qcurr_hidden_states.numel(),
269
slice_c_size=a_intermediate_cache1.numel() // num_slices,
270
+ num_slice_a=1,
271
+ num_slice_c=num_slices,
272
top_k=1 if mul_routed_weight else top_k_num,
273
MUL_ROUTED_WEIGHT=False,
274
**config,
@@ -305,10 +314,10 @@ def _fused_moe_lora(
314
b_intermediate_cache1.stride(3),
315
sorted_token_ids.stride(0),
316
expert_ids.stride(0),
308
- num_slice_a=num_slices,
309
- num_slice_c=num_slices,
317
slice_a_size=a_intermediate_cache1.numel() // num_slices,
318
slice_c_size=b_intermediate_cache1.numel() // num_slices,
319
+ num_slice_a=num_slices,
320
+ num_slice_c=num_slices,
321
top_k=1,
322
MUL_ROUTED_WEIGHT=mul_routed_weight,
323
**config,