[LoRA][1/N] Integrate flashinfer MoE LoRA for BF16 model (#48632)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jul 15, 2026 at 10:54 UTC
fdf2cf66d3aa34800aaf1f4382fb2028886f0bae
2 files changed
+425
vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
new
+387
@@ -0,0 +1,387 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+LoRA-aware FlashInfer TRT-LLM MoE experts (BF16).
5
+
6
+Reuses the routed API + ``gemm1_lora_delta`` path from FlashInfer PR #3153:
7
+
8
+ * The W13 (gate_up) LoRA delta is passed directly as ``gemm1_lora_delta`` to
9
+ the routed kernel, which fuses it into FC1 before SwiGLU (BiasType::Mn).
10
+ * The W2 (down) LoRA cannot be fused -- we take the FC1 activation output
11
+ returned by the kernel (``gemm1_activation_output``, permuted) together with
12
+ ``expanded_idx_to_permuted_idx``, unpermute it, compute the W2 delta out of
13
+ kernel via punica, and add it to the already-finalized output.
14
+
15
+Constraints (matching the PR support matrix; final gating lives in the oracle):
16
+ * SM100+ (Blackwell), gated SwiGLU, shuffled weights only;
17
+ * BF16 only;
18
+ * routing must be computed outside the MoE (the Modular path satisfies this).
19
+
20
+"""
21
+
22
+from abc import abstractmethod
23
+
24
+import torch
25
+
26
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
27
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
28
+from vllm.model_executor.layers.fused_moe.config import (
29
+ FusedMoEConfig,
30
+ FusedMoEParallelConfig,
31
+ FusedMoEQuantConfig,
32
+ RoutingMethodType,
33
+)
34
+from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
35
+ LoRAExpertsMixin,
36
+)
37
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
38
+ TopKWeightAndReduceNoOP,
39
+)
40
+from vllm.model_executor.layers.fused_moe.utils import (
41
+ trtllm_moe_pack_topk_ids_weights,
42
+)
43
+from vllm.platforms import current_platform
44
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
45
+
46
+
47
+class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
48
+ """LoRA-aware trtllm MoE experts"""
49
+
50
+ def __init__(
51
+ self,
52
+ moe_config: FusedMoEConfig,
53
+ quant_config: FusedMoEQuantConfig,
54
+ ):
55
+ super().__init__(moe_config, quant_config)
56
+ self.routing_method_type = moe_config.routing_method
57
+ self.topk = moe_config.experts_per_token
58
+ self.intermediate_size_per_partition = (
59
+ moe_config.intermediate_size_per_partition
60
+ )
61
+ self.hidden_dim = moe_config.hidden_dim
62
+ self.local_num_experts = moe_config.num_local_experts
63
+ self.ep_rank = moe_config.moe_parallel_config.ep_rank
64
+
65
+ @staticmethod
66
+ def activation_format() -> mk.FusedMoEActivationFormat:
67
+ return mk.FusedMoEActivationFormat.Standard
68
+
69
+ @staticmethod
70
+ def _supports_current_device() -> bool:
71
+ p = current_platform
72
+ return (
73
+ p.is_cuda()
74
+ and p.is_device_capability_family(100)
75
+ and has_flashinfer_trtllm_fused_moe()
76
+ )
77
+
78
+ def workspace_shapes(
79
+ self,
80
+ M,
81
+ N,
82
+ K,
83
+ topk,
84
+ global_num_experts,
85
+ local_num_experts,
86
+ expert_tokens_meta,
87
+ activation,
88
+ ):
89
+ # flashinfer manages its own workspace; only declare output (M, K)
90
+ return (0,), (0,), (M, K)
91
+
92
+ def moe_problem_size(self, a1, w1, w2, topk_ids):
93
+ """Override the base 3D-weight assumption.
94
+
95
+ FusedMoEKernel._fused_experts calls moe_problem_size before apply(),
96
+ but the base impl asserts ``len(w1.shape) == 3``. The flashinfer
97
+ trtllm path stores shuffled weights in 4D BlockMajorK layout, so we
98
+ derive the (E, M, N, K, topk) tuple from config + inputs instead.
99
+ The N/K here only feed workspace sizing, which we zero out in
100
+ workspace_shapes(); the real shapes are handled inside flashinfer.
101
+ """
102
+ E = self.local_num_experts
103
+ N = 2 * self.intermediate_size_per_partition
104
+ K = self.hidden_dim
105
+ M = a1.size(0) if a1.dim() == 2 else a1.size(1)
106
+ topk = topk_ids.size(1)
107
+ return E, M, N, K, topk
108
+
109
+ 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
111
+ return TopKWeightAndReduceNoOP()
112
+
113
+ @staticmethod
114
+ def _supports_parallel_config(
115
+ moe_parallel_config: FusedMoEParallelConfig,
116
+ ) -> bool:
117
+ return (
118
+ not moe_parallel_config.use_all2all_kernels
119
+ or moe_parallel_config.use_ag_rs_all2all_kernels
120
+ ) and not moe_parallel_config.enable_eplb
121
+
122
+ @staticmethod
123
+ def _supports_router_logits_dtype(router_logits_dtype, routing_method) -> bool:
124
+ return True
125
+
126
+ @staticmethod
127
+ def _supports_no_act_and_mul() -> bool:
128
+ return False # gated only
129
+
130
+ @property
131
+ def expects_unquantized_inputs(self) -> bool:
132
+ return True
133
+
134
+ @abstractmethod
135
+ def invoke_routed_moe(
136
+ self,
137
+ *,
138
+ hidden_states: torch.Tensor,
139
+ w1: torch.Tensor,
140
+ w2: torch.Tensor,
141
+ packed_topk_ids: torch.Tensor,
142
+ gemm1_lora_delta: torch.Tensor | None,
143
+ global_num_experts: int,
144
+ a1q_scale: torch.Tensor | None,
145
+ output: torch.Tensor,
146
+ ) -> list[torch.Tensor]:
147
+ """Call the dtype-specific trtllm_*_routed_moe and return list[Tensor].
148
+
149
+ Return contract (do_finalize=True):
150
+ gemm1_lora_delta is None -> [output]
151
+ otherwise -> [output, expanded_idx_to_permuted_idx,
152
+ gemm1_activation_output(permuted)]
153
+ """
154
+ raise NotImplementedError
155
+
156
+ def apply(
157
+ self,
158
+ output: torch.Tensor,
159
+ hidden_states: torch.Tensor,
160
+ w1: torch.Tensor,
161
+ w2: torch.Tensor,
162
+ topk_weights: torch.Tensor,
163
+ topk_ids: torch.Tensor,
164
+ activation: MoEActivation,
165
+ global_num_experts: int,
166
+ expert_map: torch.Tensor | None,
167
+ a1q_scale: torch.Tensor | None,
168
+ a2_scale: torch.Tensor | None,
169
+ workspace13: torch.Tensor,
170
+ workspace2: torch.Tensor,
171
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
172
+ apply_router_weight_on_input: bool,
173
+ ):
174
+ lora_context = self._lora_context
175
+ assert lora_context is not None, "LoRA context must be set"
176
+ num_tokens = hidden_states.size(0)
177
+ top_k = self.topk
178
+ intermediate_size = self.intermediate_size_per_partition
179
+ K = output.size(1)
180
+
181
+ # The LoRA tile-config heuristic (try_get_optimal_moe_config) unpacks
182
+ # w1/w2 as standard 3D MoE weights, but flashinfer stores shuffled
183
+ # 4D BlockMajorK weights. add_lora_w13/add_lora_w2 only read .shape
184
+ # from w1/w2 (the actual GEMM uses lora_a/b_stacked), so pass
185
+ # zero-storage meta tensors carrying the logical 3D shapes:
186
+ # w1: (E, 2I, H) w2: (E, H, I)
187
+ w1_cfg = torch.empty(
188
+ (self.local_num_experts, 2 * intermediate_size, K),
189
+ device="meta",
190
+ dtype=torch.bfloat16,
191
+ )
192
+ w2_cfg = torch.empty(
193
+ (self.local_num_experts, K, intermediate_size),
194
+ device="meta",
195
+ dtype=torch.bfloat16,
196
+ )
197
+
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
+
202
+ # ---- 1) W13 LoRA delta -> gemm1_lora_delta (bf16, [T, top_k, 2I]) ----
203
+ gemm1_lora_delta = None
204
+ w13_meta = (None, None, None, None)
205
+
206
+ gemm1_lora_delta = torch.zeros(
207
+ num_tokens,
208
+ top_k,
209
+ 2 * intermediate_size,
210
+ dtype=torch.bfloat16,
211
+ device=hidden_states.device,
212
+ )
213
+
214
+ lora_x = hidden_states
215
+ if not self.expects_unquantized_inputs:
216
+ orig = lora_context.original_hidden_states
217
+ assert orig is not None and orig.shape[0] == hidden_states.shape[0], (
218
+ "quantized trtllm LoRA path requires original_hidden_states"
219
+ )
220
+ lora_x = orig
221
+ # add_inputs=False: write the pure delta only (the base is fused in
222
+ # by the kernel) and do NOT multiply by the routing weight (it is a
223
+ # pre-SwiGLU bias).
224
+ w13_meta = self.apply_w13_lora(
225
+ lora_context,
226
+ y=gemm1_lora_delta,
227
+ x=lora_x,
228
+ topk_ids=topk_ids,
229
+ topk_weights=topk_weights,
230
+ expert_map=expert_map,
231
+ w1=w1_cfg,
232
+ w2=w2_cfg,
233
+ num_tokens=num_tokens,
234
+ top_k_num=top_k,
235
+ 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,
248
+ )
249
+
250
+ # ---- 2) Call the routed flashinfer kernel ----
251
+ ret = self.invoke_routed_moe(
252
+ hidden_states=hidden_states,
253
+ w1=w1,
254
+ w2=w2,
255
+ packed_topk_ids=packed_topk_ids,
256
+ gemm1_lora_delta=gemm1_lora_delta,
257
+ global_num_experts=global_num_experts,
258
+ a1q_scale=a1q_scale,
259
+ output=output,
260
+ )
261
+ # ---- 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
264
+ act = self._unpermute_activation(
265
+ gemm1_act_permuted,
266
+ expanded_idx_to_permuted_idx,
267
+ num_tokens,
268
+ top_k,
269
+ intermediate_size,
270
+ ) # (T*top_k, I) -- same layout as the triton path's intermediate_cache2
271
+
272
+ (
273
+ sorted_token_ids_lora,
274
+ expert_ids_lora,
275
+ num_tokens_post_padded_lora,
276
+ token_lora_mapping,
277
+ ) = w13_meta
278
+
279
+ w2_delta = torch.zeros(
280
+ num_tokens,
281
+ top_k,
282
+ K,
283
+ dtype=output.dtype,
284
+ device=output.device,
285
+ )
286
+ self.apply_w2_lora(
287
+ lora_context,
288
+ y=w2_delta,
289
+ x=act,
290
+ topk_weights=topk_weights,
291
+ sorted_token_ids_lora=sorted_token_ids_lora,
292
+ expert_ids_lora=expert_ids_lora,
293
+ num_tokens_post_padded_lora=num_tokens_post_padded_lora,
294
+ token_lora_mapping=token_lora_mapping,
295
+ num_tokens=num_tokens,
296
+ w1=w1_cfg,
297
+ w2=w2_cfg,
298
+ top_k_num=top_k,
299
+ add_inputs=False,
300
+ )
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))
306
+
307
+ @staticmethod
308
+ def _unpermute_activation(
309
+ act_permuted: torch.Tensor,
310
+ idx_map: torch.Tensor,
311
+ num_tokens: int,
312
+ top_k: int,
313
+ intermediate_size: int,
314
+ ) -> torch.Tensor:
315
+ """Permuted FC1 activation -> (num_tokens*top_k, I).
316
+
317
+ expanded_idx = token*top_k + k; idx_map[expanded_idx] = permuted_idx or -1.
318
+ TODO optimize these operations
319
+ """
320
+
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)
325
+
326
+
327
+# BF16 unquantized trtllm MoE + LoRA
328
+class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
329
+ """BF16 unquantized trtllm MoE + LoRA."""
330
+
331
+ @staticmethod
332
+ def _supports_quant_scheme(weight_key, activation_key) -> bool:
333
+ return weight_key is None and activation_key is None
334
+
335
+ @staticmethod
336
+ def _supports_activation(activation: MoEActivation) -> bool:
337
+ return activation in [MoEActivation.SILU]
338
+
339
+ @staticmethod
340
+ def _supports_routing_method(routing_method, weight_key, activation_key) -> bool:
341
+ return routing_method in [
342
+ RoutingMethodType.DeepSeekV3,
343
+ RoutingMethodType.Llama4,
344
+ RoutingMethodType.Renormalize,
345
+ RoutingMethodType.RenormalizeNaive,
346
+ ]
347
+
348
+ def invoke_routed_moe(
349
+ self,
350
+ *,
351
+ hidden_states: torch.Tensor,
352
+ w1: torch.Tensor,
353
+ w2: torch.Tensor,
354
+ packed_topk_ids: torch.Tensor,
355
+ gemm1_lora_delta: torch.Tensor | None,
356
+ global_num_experts: int,
357
+ a1q_scale: torch.Tensor | None,
358
+ output: torch.Tensor,
359
+ ) -> list[torch.Tensor]:
360
+ import flashinfer
361
+
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.
366
+ ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(
367
+ topk_ids=packed_topk_ids,
368
+ hidden_states=hidden_states,
369
+ gemm1_weights=w1,
370
+ gemm2_weights=w2,
371
+ gemm1_lora_delta=gemm1_lora_delta,
372
+ num_experts=global_num_experts,
373
+ top_k=self.topk,
374
+ n_group=None,
375
+ topk_group=None,
376
+ intermediate_size=self.intermediate_size_per_partition,
377
+ local_expert_offset=self.ep_rank * self.local_num_experts,
378
+ local_num_experts=self.local_num_experts,
379
+ routed_scaling_factor=None,
380
+ routing_method_type=self.routing_method_type,
381
+ do_finalize=True,
382
+ )
383
+ if isinstance(ret, (list, tuple)):
384
+ output.copy_(ret[0])
385
+ return list(ret)
386
+ output.copy_(ret)
387
+ return [output]
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+38
@@ -163,6 +163,33 @@ def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend
163
)
164
165
166
+def _trtllm_bf16_lora_supported(moe_config: FusedMoEConfig) -> bool:
167
+ """Gate for routing LoRA-enabled BF16 MoE to the FlashInfer TRT-LLM
168
+ gemm1_lora_delta path (PR #3153). Conservative: device + routing method;
169
+ the experts class's own _supports_* checks and the modular_kernel LoRA
170
+ gate provide the final filtering.
171
+ """
172
+ from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
173
+ TrtLlmBf16LoRAExperts,
174
+ )
175
+
176
+ if not TrtLlmBf16LoRAExperts._supports_current_device():
177
+ return False
178
+ if not TrtLlmBf16LoRAExperts._supports_routing_method(
179
+ moe_config.routing_method, None, None
180
+ ):
181
+ return False
182
+ if not TrtLlmBf16LoRAExperts._supports_parallel_config(
183
+ moe_config.moe_parallel_config
184
+ ):
185
+ return False
186
+ # The flashinfer trtllm fused-MoE kernel requires the per-partition
187
+ # intermediate size to be a multiple of 128. Plain TP shards the MoE
188
+ # intermediate dim (e.g. 768 -> 192 at tp=4), which would crash the kernel
189
+ # at runtime; fall back to Triton in that case.
190
+ return moe_config.intermediate_size_per_partition % 128 == 0
191
+
192
+
193
def select_unquantized_moe_backend(
194
moe_config: FusedMoEConfig,
195
) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
@@ -182,6 +209,17 @@ def select_unquantized_moe_backend(
209
return UnquantizedMoeBackend.OOT, None
210
211
if moe_config.is_lora_enabled:
212
+ if _trtllm_bf16_lora_supported(moe_config):
213
+ from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
214
+ TrtLlmBf16LoRAExperts,
215
+ )
216
+
217
+ logger.info_once(
218
+ "Using TrtLlmBf16LoRAExperts Unquantized MoE LoRA backend "
219
+ "(TrtLlmBf16LoRAExperts)."
220
+ )
221
+ return UnquantizedMoeBackend.FLASHINFER_TRTLLM, TrtLlmBf16LoRAExperts
222
+ logger.info_once("Using TRITON Unquantized MoE LoRA backend")
223
return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls(
224
UnquantizedMoeBackend.TRITON
225
)[0]