@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 45232a454e
Fork 0

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 4, 23:45 UTC

[FEAT] [Perf] [Gemma4] Fused Gemma4 Routing Function Triton (#39083)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

TJian committed Apr 19, 2026 at 17:57 UTC 45232a454e4c57191633fefda5baa9fa1f5dc8fd
3 files changed +180 -16
pyproject.toml
+1
@@ -170,6 +170,7 @@ eles = "eles"
170 datas = "datas"
171 ser = "ser"
172 ure = "ure"
173 +VALU = "VALU"
174 # Walsh-Hadamard Transform
175 wht = "wht"
176 WHT = "WHT"
tests/kernels/moe/test_gemma4router.py new
+57
@@ -0,0 +1,57 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +
6 +from vllm.model_executor.models.gemma4 import (
7 + gemma4_fused_routing_kernel_triton,
8 + gemma4_routing_function_torch,
9 +)
10 +
11 +
12 +def sort_by_id(w, ids):
13 + order = ids.argsort(dim=-1)
14 + return w.gather(1, order), ids.gather(1, order)
15 +
16 +
17 +# Gemma4 MoE Model has context length of 250K
18 +# the minus 1 is to ensure that edge cases are tested
19 +@pytest.mark.parametrize("num_tokens", [1, 2, 2048, 250000])
20 +@pytest.mark.parametrize("num_experts", [128]) # gemma4 moe experts
21 +@pytest.mark.parametrize("topk", [8]) # gemma4 topk
22 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
23 +def test_gemma4_routing_kernel_triton(
24 + num_tokens: int,
25 + num_experts: int,
26 + topk: int,
27 + dtype: torch.dtype,
28 +):
29 + torch.manual_seed(0)
30 +
31 + gating = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda")
32 + scales = torch.rand(num_experts, dtype=torch.float32, device="cuda")
33 +
34 + ref_w, ref_ids = gemma4_routing_function_torch(gating, topk, scales)
35 + tri_w, tri_ids = gemma4_fused_routing_kernel_triton(gating, topk, scales)
36 +
37 + # Sort by expert id — to remove tie-breaking differences
38 + ref_ws, ref_is = sort_by_id(ref_w, ref_ids)
39 + tri_ws, tri_is = sort_by_id(tri_w, tri_ids)
40 +
41 + ids_match = (ref_is == tri_is).all().item()
42 + weights_match = torch.allclose(ref_ws, tri_ws, atol=1e-2, rtol=1e-2)
43 + all_match = ids_match and weights_match
44 + max_err = (ref_ws - tri_ws).abs().max().item()
45 + print(
46 + f"T={num_tokens:5d} E={num_experts:4d} K={topk} "
47 + f"{str(dtype).split('.')[-1]:7s} ids={ids_match} max_Δweight={max_err:.2e}"
48 + )
49 + if not all_match:
50 + bad = (ref_is != tri_is).any(dim=-1).nonzero(as_tuple=True)[0]
51 + if len(bad):
52 + r = bad[0].item()
53 + print(
54 + f" first bad row {r}: ref_ids={ref_ids[r].tolist()} "
55 + f"tri_ids={tri_ids[r].tolist()}"
56 + )
57 + assert all_match
vllm/model_executor/models/gemma4.py
+122 -16
@@ -57,7 +57,9 @@ from vllm.model_executor.model_loader.weight_utils import (
57 default_weight_loader,
58 maybe_remap_kv_scale_name,
59 )
60 +from vllm.platforms import current_platform
61 from vllm.sequence import IntermediateTensors
62 +from vllm.triton_utils import tl, triton
63 from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata
64
65 from .interfaces import (
@@ -79,6 +81,120 @@ from .utils import (
81 logger = init_logger(__name__)
82
83
84 +@triton.jit
85 +def _gemma4_routing_kernel(
86 + gating_ptr,
87 + per_expert_scale_ptr,
88 + topk_weights_ptr,
89 + topk_ids_ptr,
90 + E: tl.constexpr,
91 + K: tl.constexpr,
92 + BLOCK_E: tl.constexpr,
93 +):
94 + pid = tl.program_id(0)
95 + offs_e = tl.arange(0, BLOCK_E)
96 + valid = offs_e < E
97 +
98 + logits = tl.load(
99 + gating_ptr + pid * E + offs_e,
100 + mask=valid,
101 + other=-float("inf"),
102 + ).to(tl.float32)
103 +
104 + max_l = tl.max(logits, axis=0)
105 +
106 + # Float32 → ascending-sortable bijection
107 + MIN32 = -2147483648
108 + logit_bits = logits.to(tl.int32, bitcast=True)
109 + sign_b = logit_bits >> 31
110 + key = tl.where(sign_b == 0, logit_bits ^ -1, logit_bits ^ MIN32)
111 + key = tl.where(valid, key, 0x7FFFFFFF)
112 + sk64 = key.to(tl.int64) & 0x00000000FFFFFFFF
113 + packed = (sk64 << 32) | offs_e.to(tl.int64)
114 + sorted_p = tl.sort(packed, descending=False)
115 +
116 + # Vectorized extraction of ALL sorted elements — no K-loop, no cross-lane reductions
117 + all_keys = ((sorted_p >> 32) & 0x00000000FFFFFFFF).to(tl.int32)
118 + all_ids = (sorted_p & 0x00000000FFFFFFFF).to(tl.int32)
119 +
120 + # Inverse bijection: recover original logit bits
121 + sign_k = all_keys >> 31
122 + all_bits = tl.where(sign_k < 0, all_keys ^ -1, all_keys ^ MIN32)
123 + all_logits = all_bits.to(tl.float32, bitcast=True)
124 +
125 + # Compute raw_exp for ALL BLOCK_E elements — vectorized, ~2 VALU clocks
126 + all_raw_exp = tl.math.exp2((all_logits - max_l) * 1.4426950408889634)
127 +
128 + # Sum only top-K for renorm — ONE masked reduction
129 + top_mask = offs_e < K
130 + renorm_raw = tl.sum(tl.where(top_mask, all_raw_exp, 0.0), axis=0)
131 + renorm_raw = tl.where(renorm_raw > 0.0, renorm_raw, 1.0)
132 + inv_renorm = 1.0 / renorm_raw
133 +
134 + # Load scales for top-K only (masked gather; scale array is tiny → L1 cached)
135 + all_scales = tl.load(
136 + per_expert_scale_ptr + all_ids.to(tl.int64),
137 + mask=top_mask,
138 + other=1.0,
139 + ).to(tl.float32)
140 +
141 + # Final weights: vectorized multiply (only top-K will be stored)
142 + all_weights = (all_raw_exp * inv_renorm * all_scales).to(tl.float32)
143 +
144 + # Write results with TWO masked stores — replaces K × 2 serial scalar stores
145 + base_off = pid * K + offs_e
146 + tl.store(topk_ids_ptr + base_off, all_ids, mask=top_mask)
147 + tl.store(topk_weights_ptr + base_off, all_weights, mask=top_mask)
148 +
149 +
150 +def gemma4_fused_routing_kernel_triton(
151 + gating_output: torch.Tensor,
152 + topk: int,
153 + per_expert_scale: torch.Tensor,
154 + num_warps: int = 1,
155 +) -> tuple[torch.Tensor, torch.Tensor]:
156 + gating_output = gating_output.contiguous()
157 + per_expert_scale = per_expert_scale.contiguous()
158 + T, E = gating_output.shape
159 + weights = torch.empty(T, topk, dtype=torch.float32, device=gating_output.device)
160 + ids = torch.empty(T, topk, dtype=torch.int32, device=gating_output.device)
161 + BLOCK_E = triton.next_power_of_2(E)
162 + _gemma4_routing_kernel[(T,)](
163 + gating_output,
164 + per_expert_scale,
165 + weights,
166 + ids,
167 + E,
168 + topk,
169 + BLOCK_E,
170 + num_warps=num_warps,
171 + )
172 + return weights, ids
173 +
174 +
175 +def gemma4_routing_function_torch(
176 + gating_output: torch.Tensor,
177 + topk: int,
178 + per_expert_scale: torch.Tensor,
179 +) -> tuple[torch.Tensor, torch.Tensor]:
180 + _, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
181 + router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1)
182 + indicator = torch.nn.functional.one_hot(
183 + topk_ids, num_classes=gating_output.size(-1)
184 + ).sum(dim=-2)
185 + gate_weights = indicator * router_probabilities
186 + renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True)
187 + renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0)
188 + dispatch_weights = gate_weights / renorm_factor
189 +
190 + topk_weights = dispatch_weights.gather(1, topk_ids)
191 +
192 + # Fold per_expert_scale into routing weights
193 + expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype)
194 + topk_weights = topk_weights * expert_scales
195 + return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
196 +
197 +
198 def _get_text_config(config):
199 """Dereference text_config if config is a nested Gemma4Config.
200
@@ -216,22 +332,12 @@ class Gemma4MoE(nn.Module):
332 topk: int,
333 renormalize: bool,
334 ) -> tuple[torch.Tensor, torch.Tensor]:
219 - _, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
220 - router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1)
221 - indicator = torch.nn.functional.one_hot(
222 - topk_ids, num_classes=gating_output.size(-1)
223 - ).sum(dim=-2)
224 - gate_weights = indicator * router_probabilities
225 - renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True)
226 - renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0)
227 - dispatch_weights = gate_weights / renorm_factor
228 -
229 - topk_weights = dispatch_weights.gather(1, topk_ids)
230 -
231 - # Fold per_expert_scale into routing weights
232 - expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype)
233 - topk_weights = topk_weights * expert_scales
234 - return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
335 + if current_platform.is_cuda_alike() or current_platform.is_xpu():
336 + return gemma4_fused_routing_kernel_triton(
337 + gating_output, topk, per_expert_scale
338 + )
339 +
340 + return gemma4_routing_function_torch(gating_output, topk, per_expert_scale)
341
342 # FusedMoE experts with custom Gemma4 routing
343 self.experts = FusedMoE(