[LoRA] Optimize TrtLlmLoRAExperts (#48759)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jul 18, 2026 at 10:26 UTC
da64db78b989d820effd090d8ac4b65ade2cbe4f
2 files changed
+221
-41
vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py
+11
-1
@@ -46,6 +46,7 @@ class LoRAExpertsMixin:
46
num_tokens: int,
47
top_k_num: int,
48
add_inputs: bool = True,
49
+ swap_w13_slices: bool = False,
50
) -> tuple[
51
torch.Tensor | None,
52
torch.Tensor | None,
@@ -53,6 +54,7 @@ class LoRAExpertsMixin:
54
torch.Tensor | None,
55
]:
56
w13_lora_a_stacked = lora_context.w13_lora_a_stacked
57
+ w13_lora_b_stacked = lora_context.w13_lora_b_stacked
58
if lora_context.enable_moe_shared_loras:
59
# w13 lora_A is shared across experts (collapsed expert-dim 1);
60
# broadcast to local_num_experts via a stride-0 view. The kernel
@@ -61,11 +63,19 @@ class LoRAExpertsMixin:
63
a.expand(-1, lora_context.local_num_experts, -1, -1)
64
for a in w13_lora_a_stacked
65
)
66
+ if swap_w13_slices:
67
+ # The expand kernel writes slice j into the j-th half of y's last
68
+ # dim. Reversing the (gate, up) slice tuples makes it emit
69
+ # [up, gate] order directly -- used by the FlashInfer trtllm path,
70
+ # whose SwiGLU expects the up half first, to avoid an out-of-place
71
+ # concat swap afterwards.
72
+ w13_lora_a_stacked = w13_lora_a_stacked[::-1]
73
+ w13_lora_b_stacked = w13_lora_b_stacked[::-1]
74
return lora_context.punica_wrapper.add_lora_w13(
75
y,
76
x,
77
w13_lora_a_stacked,
68
- lora_context.w13_lora_b_stacked,
78
+ w13_lora_b_stacked,
79
topk_ids,
80
topk_weights,
81
expert_map,
vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
+210
-40
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.config import (
31
FusedMoEQuantConfig,
32
RoutingMethodType,
33
)
34
+from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext
35
from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
36
LoRAExpertsMixin,
37
)
@@ -41,9 +42,75 @@ from vllm.model_executor.layers.fused_moe.utils import (
42
trtllm_moe_pack_topk_ids_weights,
43
)
44
from vllm.platforms import current_platform
45
+from vllm.triton_utils import tl, triton
46
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
47
48
49
+@triton.jit
50
+def _unpermute_activation_kernel(
51
+ act_ptr, # act_permuted: (num_permuted, num_cols)
52
+ idx_ptr, # idx_map: (num_rows,), values in [0, num_permuted) or -1
53
+ out_ptr, # out: (num_rows, num_cols)
54
+ num_cols,
55
+ stride_ar,
56
+ stride_or,
57
+ BLOCK_I: tl.constexpr,
58
+):
59
+ row = tl.program_id(0)
60
+ col_offs = tl.program_id(1) * BLOCK_I + tl.arange(0, BLOCK_I)
61
+ col_mask = col_offs < num_cols
62
+
63
+ idx = tl.load(idx_ptr + row)
64
+ out_ptrs = out_ptr + row * stride_or + col_offs
65
+ if idx >= 0:
66
+ vals = tl.load(act_ptr + idx * stride_ar + col_offs, mask=col_mask, other=0.0)
67
+ tl.store(out_ptrs, vals, mask=col_mask)
68
+ else:
69
+ zeros = tl.zeros((BLOCK_I,), dtype=out_ptr.dtype.element_ty)
70
+ tl.store(out_ptrs, zeros, mask=col_mask)
71
+
72
+
73
+@triton.jit
74
+def _finalize_lora_kernel(
75
+ gemm2_ptr, # (num_permuted, K) base FC2 output, permuted, unweighted
76
+ weight_ptr, # (num_tokens * top_k,) routing weights (expanded order)
77
+ idx_ptr, # (num_tokens * top_k,) expanded_idx -> permuted_idx or -1
78
+ delta_ptr, # (num_tokens, top_k, K) W2 LoRA delta, already routing-weighted
79
+ out_ptr, # (num_tokens, K)
80
+ K,
81
+ stride_g0,
82
+ stride_d0,
83
+ stride_d1,
84
+ stride_o0,
85
+ scale,
86
+ TOP_K: tl.constexpr,
87
+ BLOCK_K: tl.constexpr,
88
+):
89
+ token = tl.program_id(0)
90
+ col = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K)
91
+ mask = col < K
92
+
93
+ acc_base = tl.zeros((BLOCK_K,), dtype=tl.float32)
94
+ acc_delta = tl.zeros((BLOCK_K,), dtype=tl.float32)
95
+ for k in tl.static_range(TOP_K):
96
+ eid = token * TOP_K + k
97
+ pidx = tl.load(idx_ptr + eid)
98
+ if pidx >= 0:
99
+ w = tl.load(weight_ptr + eid).to(tl.float32)
100
+ base = tl.load(gemm2_ptr + pidx * stride_g0 + col, mask=mask, other=0.0).to(
101
+ tl.float32
102
+ )
103
+ acc_base += w * base
104
+ acc_delta += tl.load(
105
+ delta_ptr + token * stride_d0 + k * stride_d1 + col, mask=mask, other=0.0
106
+ ).to(tl.float32)
107
+
108
+ out = acc_base * scale + acc_delta
109
+ tl.store(
110
+ out_ptr + token * stride_o0 + col, out.to(out_ptr.dtype.element_ty), mask=mask
111
+ )
112
+
113
+
114
class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
115
"""LoRA-aware trtllm MoE experts"""
116
@@ -107,7 +174,8 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
174
return E, M, N, K, topk
175
176
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
110
- # do_finalize=True: the kernel already does moe_sum, so this is a No-Op
177
+ # apply() writes the fully finalized result into `output` (fused base
178
+ # finalize + W2 LoRA reduction), so this is a No-Op.
179
return TopKWeightAndReduceNoOP()
180
181
@staticmethod
@@ -146,9 +214,13 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
214
) -> list[torch.Tensor]:
215
"""Call the dtype-specific trtllm_*_routed_moe and return list[Tensor].
216
149
- Return contract (do_finalize=True):
150
- gemm1_lora_delta is None -> [output]
151
- otherwise -> [output, expanded_idx_to_permuted_idx,
217
+ The LoRA path always sets gemm1_lora_delta and runs with
218
+ do_finalize=False so the base finalize can be fused with the W2 LoRA
219
+ reduction (see _finalize_with_w2_lora). Return contract:
220
+ gemm1_lora_delta is None -> [output] (do_finalize=True)
221
+ otherwise -> [gemm2_output(permuted, unweighted),
222
+ expert_weights,
223
+ expanded_idx_to_permuted_idx,
224
gemm1_activation_output(permuted)]
225
"""
226
raise NotImplementedError
@@ -178,6 +250,27 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
250
intermediate_size = self.intermediate_size_per_partition
251
K = output.size(1)
252
253
+ # Routing is computed outside the MoE; pack it into the
254
+ # (eid<<16)|w.bf16 format the routed API expects.
255
+ packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
256
+
257
+ # ---- Base-model fast path ----
258
+ # When no token in the batch selects a LoRA adapter, skip the LoRA machinery
259
+ # and run the plain base MoE with do_finalize=True, which writes the finalized
260
+ # result straight into `output`.
261
+ if self._batch_has_no_lora(lora_context):
262
+ self.invoke_routed_moe(
263
+ hidden_states=hidden_states,
264
+ w1=w1,
265
+ w2=w2,
266
+ packed_topk_ids=packed_topk_ids,
267
+ gemm1_lora_delta=None, # without LoRA, no delta
268
+ global_num_experts=global_num_experts,
269
+ a1q_scale=a1q_scale,
270
+ output=output,
271
+ )
272
+ return
273
+
274
# The LoRA tile-config heuristic (try_get_optimal_moe_config) unpacks
275
# w1/w2 as standard 3D MoE weights, but flashinfer stores shuffled
276
# 4D BlockMajorK weights. add_lora_w13/add_lora_w2 only read .shape
@@ -195,14 +288,13 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
288
dtype=torch.bfloat16,
289
)
290
198
- # Routing is computed outside the MoE; pack it into the
199
- # (eid<<16)|w.bf16 format the routed API expects.
200
- packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
201
-
291
# ---- 1) W13 LoRA delta -> gemm1_lora_delta (bf16, [T, top_k, 2I]) ----
292
gemm1_lora_delta = None
293
w13_meta = (None, None, None, None)
294
295
+ # zeros (not empty): under EP the punica expand kernel only writes
296
+ # slots whose expert is local to this rank; non-local (token, top_k)
297
+ # slots must stay 0 so they contribute no bias when fed to flashinfer.
298
gemm1_lora_delta = torch.zeros(
299
num_tokens,
300
top_k,
@@ -221,6 +313,10 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
313
# add_inputs=False: write the pure delta only (the base is fused in
314
# by the kernel) and do NOT multiply by the routing weight (it is a
315
# pre-SwiGLU bias).
316
+ # swap_w13_slices=True: apply_w13_lora writes in vLLM's w13 order
317
+ # (gate=w1 first, up=w3 second), but FlashInfer's gemm1_lora_delta
318
+ # expects [up, gate]; reversing the slices emits that order directly,
319
+ # avoiding an out-of-place concat swap.
320
w13_meta = self.apply_w13_lora(
321
lora_context,
322
y=gemm1_lora_delta,
@@ -233,18 +329,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
329
num_tokens=num_tokens,
330
top_k_num=top_k,
331
add_inputs=False,
236
- )
237
-
238
- # apply_w13_lora writes the delta in vLLM's w13 order (gate=w1 first,
239
- # up=w3 second), but FlashInfer's gemm1_lora_delta expects the halves
240
- # in [up, gate] order. Swap them so the delta lands on the matching
241
- # SwiGLU branch.
242
- gemm1_lora_delta = torch.cat(
243
- [
244
- gemm1_lora_delta[..., intermediate_size:],
245
- gemm1_lora_delta[..., :intermediate_size],
246
- ],
247
- dim=-1,
332
+ swap_w13_slices=True,
333
)
334
335
# ---- 2) Call the routed flashinfer kernel ----
@@ -259,8 +344,14 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
344
output=output,
345
)
346
# ---- 3) W2 LoRA (computed out of kernel) ----
262
- expanded_idx_to_permuted_idx = ret[1]
263
- gemm1_act_permuted = ret[2] # [max_padded, I], post-act
347
+ # do_finalize=False: flashinfer returns the *unfinalized* base output.
348
+ # ret = [gemm2_output(permuted, unweighted),
349
+ # expert_weights, expanded_idx_to_permuted_idx,
350
+ # gemm1_activation_output(permuted)]
351
+ gemm2_permuted = ret[0]
352
+ expert_weights = ret[1]
353
+ expanded_idx_to_permuted_idx = ret[2]
354
+ gemm1_act_permuted = ret[3] # [max_padded, I], post-act
355
act = self._unpermute_activation(
356
gemm1_act_permuted,
357
expanded_idx_to_permuted_idx,
@@ -298,11 +389,37 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
389
top_k_num=top_k,
390
add_inputs=False,
391
)
301
- # The base output is already finalized (routing-weighted + summed over
302
- # top_k); the W2 delta is likewise already routing-weighted, so sum it
303
- # over top_k and add.
304
- # TODO(verify): if routed_scaling_factor is not None, scale to match.
305
- output.add_(w2_delta.sum(dim=1))
392
+ # ---- 4) Fused finalize: reduce the base path over top_k (with routing
393
+ # weights) and add the already-weighted W2 delta, in a single kernel.
394
+ # This replaces flashinfer's internal finalize launch plus a separate
395
+ # w2_delta.sum(dim=1) + add_.
396
+ self._finalize_with_w2_lora(
397
+ output,
398
+ gemm2_permuted,
399
+ expert_weights,
400
+ expanded_idx_to_permuted_idx,
401
+ w2_delta,
402
+ num_tokens,
403
+ top_k,
404
+ scale=1.0,
405
+ )
406
+
407
+ @staticmethod
408
+ def _batch_has_no_lora(lora_context: MoELoRAContext) -> bool:
409
+ """True when no token in the batch selects a LoRA adapter.
410
+
411
+ Mirrors the no-lora fast path in
412
+ ``PunicaWrapperGPU.add_lora_fused_moe``: the punica kernel metadata
413
+ carries a CPU ``no_lora_flag`` computed once per forward from the
414
+ token->LoRA mapping. Reading it is a host-only check (no device sync),
415
+ and under CUDA graphs the branch is frozen at capture time against the
416
+ graph's ``has_lora`` dispatch key, so it stays correct on replay.
417
+ """
418
+ meta = getattr(lora_context.punica_wrapper, "token_mapping_meta", None)
419
+ if meta is None:
420
+ return False
421
+ flag = meta.no_lora_flag_cpu
422
+ return bool(flag.numel() == 1 and flag.item())
423
424
@staticmethod
425
def _unpermute_activation(
@@ -315,13 +432,63 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
432
"""Permuted FC1 activation -> (num_tokens*top_k, I).
433
434
expanded_idx = token*top_k + k; idx_map[expanded_idx] = permuted_idx or -1.
318
- TODO optimize these operations
435
+ Fused gather + drop-masking: each output row copies the matching
436
+ permuted row, or is zeroed when idx_map < 0.
437
"""
438
+ num_rows = num_tokens * top_k
439
+ out = torch.empty(
440
+ (num_rows, intermediate_size),
441
+ dtype=act_permuted.dtype,
442
+ device=act_permuted.device,
443
+ )
444
+ BLOCK_I = 1024
445
+ grid = (num_rows, triton.cdiv(intermediate_size, BLOCK_I))
446
+ _unpermute_activation_kernel[grid](
447
+ act_permuted,
448
+ idx_map,
449
+ out,
450
+ intermediate_size,
451
+ act_permuted.stride(0),
452
+ out.stride(0),
453
+ BLOCK_I=BLOCK_I,
454
+ )
455
+ return out
456
321
- valid = idx_map >= 0
322
- safe_idx = idx_map.clamp_min(0).long()
323
- gathered = act_permuted[safe_idx]
324
- return gathered * valid.unsqueeze(1).to(act_permuted.dtype)
457
+ @staticmethod
458
+ def _finalize_with_w2_lora(
459
+ output: torch.Tensor,
460
+ gemm2_permuted: torch.Tensor,
461
+ expert_weights: torch.Tensor,
462
+ idx_map: torch.Tensor,
463
+ w2_delta: torch.Tensor,
464
+ num_tokens: int,
465
+ top_k: int,
466
+ scale: float = 1.0,
467
+ ) -> None:
468
+ """Fused base finalize + W2 LoRA reduction, written into ``output``.
469
+
470
+ For each token: sum the routing-weighted permuted base rows over top_k
471
+ (``expert_weights`` in expanded order, ``idx_map < 0`` dropped), scale by
472
+ ``scale``, and add the already-weighted ``w2_delta`` reduced over top_k.
473
+ """
474
+ K = gemm2_permuted.size(1)
475
+ BLOCK_K = 512
476
+ grid = (num_tokens, triton.cdiv(K, BLOCK_K))
477
+ _finalize_lora_kernel[grid](
478
+ gemm2_permuted,
479
+ expert_weights.reshape(-1),
480
+ idx_map,
481
+ w2_delta,
482
+ output,
483
+ K,
484
+ gemm2_permuted.stride(0),
485
+ w2_delta.stride(0),
486
+ w2_delta.stride(1),
487
+ output.stride(0),
488
+ scale,
489
+ TOP_K=top_k,
490
+ BLOCK_K=BLOCK_K,
491
+ )
492
493
494
# BF16 unquantized trtllm MoE + LoRA
@@ -359,10 +526,12 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
526
) -> list[torch.Tensor]:
527
import flashinfer
528
362
- # Unlike the fp8/mxint4 routed APIs, trtllm_bf16_routed_moe has no
363
- # `output=` kwarg: it returns the finalized tensor (or a list whose
364
- # [0] is it when gemm1_lora_delta is set). Copy it into the caller's
365
- # buffer so the modular-kernel output plumbing sees the result.
529
+ # With gemm1_lora_delta set (the LoRA path) run do_finalize=False and
530
+ # return the unfinalized permuted base output so apply() can fuse the
531
+ # finalize with the W2 LoRA reduction (see _finalize_with_w2_lora).
532
+ # Without a delta (base path), run do_finalize=True and hand flashinfer
533
+ # the caller's buffer via output= so it finalizes in place -- no copy.
534
+ do_finalize = gemm1_lora_delta is None
535
ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(
536
topk_ids=packed_topk_ids,
537
hidden_states=hidden_states,
@@ -378,10 +547,11 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
547
local_num_experts=self.local_num_experts,
548
routed_scaling_factor=None,
549
routing_method_type=self.routing_method_type,
381
- do_finalize=True,
550
+ do_finalize=do_finalize,
551
+ output=output if do_finalize else None,
552
)
383
- if isinstance(ret, (list, tuple)):
384
- output.copy_(ret[0])
553
+ if not do_finalize:
554
+ # [gemm2_output, expert_weights, expanded_idx, gemm1_activation]
555
return list(ret)
386
- output.copy_(ret)
556
+ # do_finalize=True finalized directly into `output`.
557
return [output]