add pad-aware swiglu limit kernel (#48287)
Signed-off-by: gnovack <novackgm@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
gnovack committed
Jul 13, 2026 at 16:48 UTC
21472f32ea6c1f08b6b3c12aec9d3df0a41b4575
4 files changed
+140
-4
vllm/model_executor/layers/fused_moe/activation.py
+19
-2
@@ -107,6 +107,22 @@ def activation_without_mul(activation: str) -> str:
107
return MoEActivation.from_str(activation).without_mul().value
108
109
110
+def silu_and_mul_with_clamp(
111
+ output: torch.Tensor,
112
+ input: torch.Tensor,
113
+ clamp_limit: float,
114
+ topk_ids: torch.Tensor | None = None,
115
+ expert_map: torch.Tensor | None = None,
116
+) -> None:
117
+ if topk_ids is not None and expert_map is not None:
118
+ from vllm.model_executor.layers.fused_moe.utils import swiglu_limit_func
119
+
120
+ swiglu_limit_func(output, input, clamp_limit, topk_ids, expert_map)
121
+ else:
122
+ # Fused silu(clamp(gate)) * clamp(up); equivalent to swiglu_limit_func.
123
+ torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, 1.0, 0.0)
124
+
125
+
126
def apply_moe_activation(
127
activation: MoEActivation,
128
output: torch.Tensor,
@@ -115,6 +131,8 @@ def apply_moe_activation(
131
clamp_limit: float | None = None,
132
alpha: float = 1.0,
133
beta: float = 0.0,
134
+ topk_ids: torch.Tensor | None = None,
135
+ expert_map: torch.Tensor | None = None,
136
) -> torch.Tensor:
137
"""Apply MoE activation function.
138
@@ -138,8 +156,7 @@ def apply_moe_activation(
156
# Activations with gated multiplication (gate × activation(up))
157
if activation == MoEActivation.SILU:
158
if clamp_limit is not None:
141
- # Fused silu(clamp(gate)) * clamp(up); equivalent to swiglu_limit_func.
142
- torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, 1.0, 0.0)
159
+ silu_and_mul_with_clamp(output, input, clamp_limit, topk_ids, expert_map)
160
else:
161
torch.ops._C.silu_and_mul(output, input)
162
elif activation == MoEActivation.GELU:
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+8
@@ -73,6 +73,7 @@ def _fused_marlin_moe(
73
num_tokens_post_padded: torch.Tensor,
74
activation: MoEActivation = MoEActivation.SILU,
75
activation_func: Callable[..., None] = apply_moe_activation,
76
+ topk_ids: torch.Tensor | None = None,
77
input_global_scale1: torch.Tensor | None = None,
78
input_global_scale2: torch.Tensor | None = None,
79
global_scale1: torch.Tensor | None = None,
@@ -168,6 +169,8 @@ def _fused_marlin_moe(
169
clamp_limit=clamp_limit,
170
alpha=gemm1_alpha,
171
beta=gemm1_beta,
172
+ topk_ids=topk_ids,
173
+ expert_map=expert_map,
174
)
175
176
if output is None:
@@ -346,6 +349,7 @@ def fused_marlin_moe(
349
w1_scale=w1_scale,
350
w2_scale=w2_scale,
351
topk_weights=topk_weights,
352
+ topk_ids=topk_ids,
353
num_topk=topk,
354
quant_type=quant_type,
355
apply_router_weight_on_input=apply_router_weight_on_input,
@@ -827,6 +831,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
831
clamp_limit: float | None = None,
832
alpha: float = 1.0,
833
beta: float = 0.0,
834
+ topk_ids: torch.Tensor | None = None,
835
+ expert_map: torch.Tensor | None = None,
836
) -> None:
837
# act_input = intermediate_cache1 (M*topk, 2N for gated)
838
# act_output = intermediate_cache2 (M*topk, N)
@@ -863,6 +869,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
869
clamp_limit=clamp_limit,
870
alpha=alpha,
871
beta=beta,
872
+ topk_ids=topk_ids,
873
+ expert_map=expert_map,
874
)
875
lora_state["cache2"] = act_output
876
vllm/model_executor/layers/fused_moe/modular_kernel.py
+10
-1
@@ -889,9 +889,18 @@ class FusedMoEExpertsModular(FusedMoEExperts):
889
clamp_limit: float | None = None,
890
alpha: float = 1.0,
891
beta: float = 0.0,
892
+ topk_ids: torch.Tensor | None = None,
893
+ expert_map: torch.Tensor | None = None,
894
) -> None:
895
apply_moe_activation(
894
- activation, output, input, clamp_limit=clamp_limit, alpha=alpha, beta=beta
896
+ activation,
897
+ output,
898
+ input,
899
+ clamp_limit=clamp_limit,
900
+ alpha=alpha,
901
+ beta=beta,
902
+ topk_ids=topk_ids,
903
+ expert_map=expert_map,
904
)
905
906
@abstractmethod
vllm/model_executor/layers/fused_moe/utils.py
+103
-1
@@ -454,7 +454,7 @@ def trtllm_moe_pack_topk_ids_weights(
454
455
456
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
457
-def swiglu_limit_func(
457
+def _swiglu_limit_torch(
458
output: torch.Tensor,
459
input: torch.Tensor, # first half is gate, second half is up
460
swiglu_limit: float = 0.0,
@@ -470,6 +470,108 @@ def swiglu_limit_func(
470
output.copy_(F.silu(gate) * up)
471
472
473
+@triton.jit
474
+def _swiglu_limit_pad_aware_kernel(
475
+ input_ptr, # [num_tokens, 2 * hidden_size]
476
+ output_ptr, # [num_tokens, hidden_size]
477
+ topk_ids_ptr, # [num_tokens, num_topk]
478
+ expert_map_ptr, # global -> local expert id, or -1 if non-local
479
+ hidden_size,
480
+ input_row_stride,
481
+ num_tokens,
482
+ swiglu_limit,
483
+ HAS_LIMIT: tl.constexpr,
484
+ HAS_EXPERT_MAP: tl.constexpr,
485
+ BLOCK_SIZE: tl.constexpr,
486
+):
487
+ # Persistent over rows: each CTA owns one column tile and processes a
488
+ # strided set of token assignments.
489
+ pid = tl.program_id(0)
490
+ row_stride = tl.num_programs(0)
491
+ column_tile = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
492
+ mask = column_tile < hidden_size
493
+
494
+ for row in tl.range(pid, num_tokens, row_stride):
495
+ expert_id = tl.load(topk_ids_ptr + row)
496
+ should_compute = expert_id != -1
497
+ if HAS_EXPERT_MAP:
498
+ local_expert_id = tl.load(
499
+ expert_map_ptr + expert_id,
500
+ mask=expert_id >= 0,
501
+ other=-1,
502
+ )
503
+ should_compute = should_compute & (local_expert_id != -1)
504
+
505
+ if should_compute:
506
+ gate_offsets = row.to(tl.int64) * input_row_stride + column_tile
507
+ up_offsets = gate_offsets + hidden_size
508
+
509
+ gate = tl.load(input_ptr + gate_offsets, mask=mask, other=0.0).to(
510
+ tl.float32
511
+ )
512
+
513
+ up = tl.load(input_ptr + up_offsets, mask=mask, other=0.0).to(tl.float32)
514
+
515
+ if HAS_LIMIT:
516
+ gate = tl.minimum(gate, swiglu_limit)
517
+ up = tl.maximum(up, -swiglu_limit)
518
+ up = tl.minimum(up, swiglu_limit)
519
+
520
+ silu_gate = gate / (1.0 + tl.exp(-gate))
521
+ result = silu_gate * up
522
+ tl.store(
523
+ output_ptr + row.to(tl.int64) * hidden_size + column_tile,
524
+ result.to(output_ptr.dtype.element_ty),
525
+ mask=mask,
526
+ )
527
+
528
+
529
+def _swiglu_limit_pad_aware(
530
+ output: torch.Tensor,
531
+ input: torch.Tensor,
532
+ topk_ids: torch.Tensor,
533
+ swiglu_limit: float,
534
+ expert_map: torch.Tensor | None = None,
535
+) -> None:
536
+ num_tokens, gate_up_size = input.shape
537
+ hidden_size = gate_up_size // 2
538
+ if num_tokens == 0:
539
+ return
540
+
541
+ BLOCK_SIZE = 1024
542
+ grid = (min(num_tokens, 256), triton.cdiv(hidden_size, BLOCK_SIZE))
543
+ _swiglu_limit_pad_aware_kernel[grid](
544
+ input,
545
+ output,
546
+ topk_ids,
547
+ expert_map,
548
+ hidden_size,
549
+ gate_up_size,
550
+ num_tokens,
551
+ swiglu_limit,
552
+ HAS_LIMIT=swiglu_limit > 0,
553
+ HAS_EXPERT_MAP=expert_map is not None,
554
+ BLOCK_SIZE=BLOCK_SIZE,
555
+ num_warps=4,
556
+ )
557
+
558
+
559
+def swiglu_limit_func(
560
+ output: torch.Tensor,
561
+ input: torch.Tensor, # first half is gate, second half is up
562
+ swiglu_limit: float = 0.0,
563
+ topk_ids: torch.Tensor | None = None,
564
+ expert_map: torch.Tensor | None = None,
565
+) -> None:
566
+ # The pad-aware Triton kernel skips unrouted token slots (topk_ids == -1)
567
+ # and, when expert_map is given, slots routed to non-local experts, so it
568
+ # requires topk_ids. Fall back to the torch implementation otherwise.
569
+ if topk_ids is not None:
570
+ _swiglu_limit_pad_aware(output, input, topk_ids, swiglu_limit, expert_map)
571
+ else:
572
+ _swiglu_limit_torch(output, input, swiglu_limit)
573
+
574
+
575
@functools.lru_cache
576
def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool:
577
return (