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 13, 14:30 UTC
[Kernel] Fix fused_gdn_gating (#28343)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Nov 10, 2025 at 05:26 UTC
c4768dcf47ae919257e31b49a03c00d383ba3c55
1 file changed
+5
-3
vllm/model_executor/models/qwen3_next.py
+5
-3
@@ -1367,8 +1367,10 @@ def fused_gdn_gating_kernel(
1367
blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
1368
tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
1369
# compute beta_output = sigmoid(b)
1370
- blk_beta = 1.0 / (1.0 + tl.exp(-blk_b.to(tl.float32)))
1371
- tl.store(beta_output + off, blk_beta.to(beta_output.dtype.element_ty), mask=mask)
1370
+ blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
1371
+ tl.store(
1372
+ beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
1373
+ )
1374
1375
1376
def fused_gdn_gating(
@@ -1389,7 +1391,7 @@ def fused_gdn_gating(
1391
seq_len = 1
1392
grid = (batch, seq_len, triton.cdiv(num_heads, 8))
1393
g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
1392
- beta_output = torch.empty(1, batch, num_heads, dtype=torch.float32, device=b.device)
1394
+ beta_output = torch.empty(1, batch, num_heads, dtype=b.dtype, device=b.device)
1395
fused_gdn_gating_kernel[grid](
1396
g,
1397
beta_output,