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 9, 11:15 UTC
[lora/moe] Avoid extra intermediate buffer & Python slicing in expand phase when split_k == 1 (#32774)
Signed-off-by: 陈建华 <1647430658@qq.com>
cwazai committed
Jan 29, 2026 at 00:22 UTC
f210f0b7b1e4f179cf194c09c8805898134b70d8
1 file changed
+17
-14
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+17
-14
@@ -84,6 +84,7 @@ def _fused_moe_lora_kernel(
84
num_slice_c: tl.constexpr,
85
top_k: tl.constexpr,
86
MUL_ROUTED_WEIGHT: tl.constexpr,
87
+ ADD_INPUTS: tl.constexpr,
88
USE_B_L2_CACHE: tl.constexpr, # new, enable .ca load for B
89
BLOCK_SIZE_M: tl.constexpr,
90
BLOCK_SIZE_N: tl.constexpr,
@@ -211,7 +212,11 @@ def _fused_moe_lora_kernel(
212
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
213
214
if SPLIT_K == 1:
214
- tl.store(c_ptrs, accumulator, mask=c_mask)
215
+ if ADD_INPUTS:
216
+ prev = tl.load(c_ptrs, mask=c_mask, other=0.0)
217
+ tl.store(c_ptrs, prev + accumulator, mask=c_mask)
218
+ else:
219
+ tl.store(c_ptrs, accumulator, mask=c_mask)
220
else:
221
tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
222
@@ -305,6 +310,7 @@ def _fused_moe_lora_shrink(
310
num_slice_c=num_slices,
311
top_k=1 if mul_routed_weight else top_k_num,
312
MUL_ROUTED_WEIGHT=False,
313
+ ADD_INPUTS=False,
314
USE_B_L2_CACHE=True, # new
315
IS_PRIMARY=True,
316
**shrink_config,
@@ -315,7 +321,6 @@ def _fused_moe_lora_shrink(
321
def _fused_moe_lora_expand(
322
output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)
323
a_intermediate_cache1: torch.Tensor, # (num_slices, M, top_k_num, max_lora_rank)
318
- b_intermediate_cache1: torch.Tensor, # (num_slices, M, top_k_num, output_dim_size)
324
lora_b_stacked: list[
325
torch.Tensor
326
], # [(max_loras, num_experts, max_lora_rank, K,),...]
@@ -376,10 +381,15 @@ def _fused_moe_lora_expand(
381
## max_loras + 1 to handle the no-lora case (lora_id == -1)
382
lora_b_stacked[0].shape[0] + 1,
383
)
384
+
385
+ # Fast path: directly accumulate into the corresponding slice interval of output.
386
+ out_view = output[:, :, offset : offset + num_slices * N]
387
+ slice_c_size = N * out_view.stride(2)
388
+
389
_fused_moe_lora_kernel[grid](
390
a_intermediate_cache1,
391
b_ptr,
382
- b_intermediate_cache1,
392
+ out_view,
393
topk_weights,
394
sorted_token_ids,
395
expert_ids,
@@ -398,22 +408,21 @@ def _fused_moe_lora_expand(
408
w1_lora_b_stacked.stride(1),
409
w1_lora_b_stacked.stride(3),
410
w1_lora_b_stacked.stride(2),
401
- b_intermediate_cache1.stride(2),
402
- b_intermediate_cache1.stride(3),
411
+ out_view.stride(1),
412
+ out_view.stride(2),
413
sorted_token_ids.stride(0),
414
expert_ids.stride(0),
415
slice_a_size=a_intermediate_cache1.numel() // num_slices,
406
- slice_c_size=b_intermediate_cache1.numel() // num_slices,
416
+ slice_c_size=slice_c_size,
417
num_slice_a=num_slices,
418
num_slice_c=num_slices,
419
top_k=1,
420
MUL_ROUTED_WEIGHT=mul_routed_weight,
421
+ ADD_INPUTS=True,
422
USE_B_L2_CACHE=True, # new
423
IS_PRIMARY=False,
424
**expand_config,
425
)
415
- for i in range(num_slices):
416
- output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i]
426
427
428
@torch.inference_mode()
@@ -484,11 +493,6 @@ def _fused_moe_lora(
493
device=device,
494
)
495
487
- b_intermediate_cache1 = torch.zeros(
488
- (num_slices, M, top_k_num, w1_output_dim_size),
489
- dtype=output.dtype,
490
- device=device,
491
- )
496
use_gdc = supports_pdl(device) and not fully_sharded
497
_fused_moe_lora_shrink(
498
a_intermediate_cache1,
@@ -537,7 +541,6 @@ def _fused_moe_lora(
541
_fused_moe_lora_expand(
542
output,
543
a_intermediate_cache1,
540
- b_intermediate_cache1,
544
lora_b_stacked,
545
topk_weights,
546
sorted_token_ids,