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 15, 01:00 UTC
[Kernel][Performance] Tweak MoE Batched silu_mul_fp8_quant_deep_gemm kernel (#21193)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Jul 19, 2025 at 11:39 UTC
dcc6cfb991cd76369aad96e04424f29c8fecdbd8
1 file changed
+4
-5
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+4
-5
@@ -55,6 +55,7 @@ def _silu_mul_fp8_quant_deep_gemm(
55
56
# Meta ---------------------------------------------------------------
57
BLOCK: tl.constexpr,
58
+ NUM_STAGES: tl.constexpr,
59
):
60
G = H // GROUP_SIZE
61
@@ -73,8 +74,7 @@ def _silu_mul_fp8_quant_deep_gemm(
74
cols = cols.to(tl.int64)
75
mask_h = cols < BLOCK
76
76
- t = tl.zeros([], tl.int64)
77
- while t < n_tokens:
77
+ for t in tl.range(0, n_tokens, num_stages=NUM_STAGES):
78
base_i_offset = (e * stride_i_e + t * stride_i_t +
79
g * GROUP_SIZE * stride_i_h)
80
base_yq_offset = (e * stride_yq_e + t * stride_yq_t +
@@ -102,8 +102,6 @@ def _silu_mul_fp8_quant_deep_gemm(
102
tl.store(y_q_ptr + base_yq_offset + cols * stride_yq_h, y_q, mask=mask)
103
tl.store(y_s_ptr + base_ys_offset, y_s)
104
105
- t += 1
106
-
105
106
def silu_mul_fp8_quant_deep_gemm(
107
y: torch.Tensor, # (E, T, 2*H) float32
@@ -180,7 +178,8 @@ def silu_mul_fp8_quant_deep_gemm(
178
fp8_max,
179
is_blackwell_deep_gemm_used(),
180
BLOCK=group_size,
183
- num_warps=4,
181
+ NUM_STAGES=8,
182
+ num_warps=1,
183
)
184
185
return y_q, y_s