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 11, 06:15 UTC
fused_moe_lora PDL improvements (#30716)
Signed-off-by: gnovack <gnovack@amazon.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
gnovack committed
Dec 17, 2025 at 19:55 UTC
5a3adf581e372c60d6135a535561f4d491c4d046
1 file changed
+18
-12
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+18
-12
@@ -156,16 +156,22 @@ def _fused_moe_lora_kernel(
156
+ offs_bn[None, :] * stride_bn
157
)
158
159
+ if USE_GDC and IS_PRIMARY:
160
+ # GDC launch dependents hints the runtime system to launch dependent kernels.
161
+ tl.extra.cuda.gdc_launch_dependents()
162
+
163
# accumulator
164
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
165
+
166
+ # GDC wait waits for ALL programs in the prior kernel to complete
167
+ # before continuing.
168
+ if USE_GDC and not IS_PRIMARY:
169
+ tl.extra.cuda.gdc_wait()
170
+
171
for k in range(0, grid_k):
172
k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
173
# pre-fetch lora weight
174
b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
165
- # GDC wait waits for ALL programs in the prior kernel to complete
166
- # before continuing.
167
- if USE_GDC and not IS_PRIMARY:
168
- tl.extra.cuda.gdc_wait()
175
a = tl.load(
176
a_ptrs,
177
mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
@@ -179,9 +185,6 @@ def _fused_moe_lora_kernel(
185
if MUL_ROUTED_WEIGHT:
186
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
187
accumulator = accumulator * moe_weight[:, None]
182
- if USE_GDC and IS_PRIMARY:
183
- # GDC launch dependents hints the runtime system to launch dependent kernels.
184
- tl.extra.cuda.gdc_launch_dependents()
188
accumulator = accumulator.to(c_ptr.dtype.element_ty)
189
# Write back the block of the output
190
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
@@ -290,6 +293,7 @@ def _fused_moe_lora_shrink(
293
def _fused_moe_lora_expand(
294
output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)
295
a_intermediate_cache1: torch.Tensor, # (num_slices, M, top_k_num, max_lora_rank)
296
+ b_intermediate_cache1: torch.Tensor, # (num_slices, M, top_k_num, output_dim_size)
297
lora_b_stacked: list[
298
torch.Tensor
299
], # [(max_loras, num_experts, max_lora_rank, K,),...]
@@ -331,11 +335,6 @@ def _fused_moe_lora_expand(
335
-1, a_intermediate_cache1.shape[3]
336
)
337
334
- b_intermediate_cache1 = torch.zeros(
335
- (num_slices, M, top_k_num, w1_output_dim_size),
336
- dtype=output.dtype,
337
- device=device,
338
- )
338
use_gdc = supports_pdl(a_intermediate_cache1.device)
339
expand_config = {
340
"BLOCK_SIZE_M": block_size_m,
@@ -460,6 +459,12 @@ def _fused_moe_lora(
459
device=device,
460
)
461
462
+ b_intermediate_cache1 = torch.zeros(
463
+ (num_slices, M, top_k_num, w1_output_dim_size),
464
+ dtype=output.dtype,
465
+ device=device,
466
+ )
467
+
468
_fused_moe_lora_shrink(
469
a_intermediate_cache1,
470
qcurr_hidden_states,
@@ -506,6 +511,7 @@ def _fused_moe_lora(
511
_fused_moe_lora_expand(
512
output,
513
a_intermediate_cache1,
514
+ b_intermediate_cache1,
515
lora_b_stacked,
516
topk_weights,
517
sorted_token_ids,