minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[4/N] Qwen3.5Opt: Overlap mamba verify update with draft extend (#26924)

Browse Source
This commit is contained in:
Yuhao Yang
2026-06-13 20:29:20 +08:00
committed by GitHub
parent f7041c9dee
commit aea0e30853
7 changed files with 534 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/bench_fused_gate_sigmoid_mul_add.py
+60
View File
@@ -0,0 +1,60 @@
"""Benchmark fused_gate_sigmoid_mul_add: Triton kernel vs PyTorch eager.
Compares the fused Triton kernel against a plain PyTorch implementation
over the Qwen3.5 MoE target hidden size.
"""
import torch
import triton
from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add
HIDDEN_DIMS = [4096]
def _pytorch_reference(hidden_states, gate_weight, shared_output, final_hidden_states):
gate = hidden_states @ gate_weight
final_hidden_states += torch.sigmoid(gate).unsqueeze(1) * shared_output
def make_bench(hidden_dim):
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens"],
x_vals=[1, 2, 4, 8, 16, 1024, 2048, 4096, 8192],
line_arg="impl",
line_vals=["triton", "pytorch"],
line_names=["Triton fused", "PyTorch eager"],
styles=[("blue", "-"), ("orange", "--")],
ylabel="us",
plot_name=f"fused_gate_sigmoid_mul_add-hidden{hidden_dim}",
args={"hidden_dim": hidden_dim},
)
)
def bench(num_tokens, impl, hidden_dim, dtype=torch.bfloat16):
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate_weight = torch.randn(hidden_dim, dtype=dtype, device="cuda")
shared_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
final_hidden_states = torch.randn(
num_tokens, hidden_dim, dtype=dtype, device="cuda"
)
if impl == "triton":
fn = lambda: fused_gate_sigmoid_mul_add(
hidden_states, gate_weight, shared_output, final_hidden_states
)
else:
fn = lambda: _pytorch_reference(
hidden_states, gate_weight, shared_output, final_hidden_states
)
ms = triton.testing.do_bench(fn, warmup=100, rep=200)
return ms * 1000 # convert to us
return bench
if __name__ == "__main__":
for d in HIDDEN_DIMS:
print(f"\n===== hidden_dim={d} =====")
make_bench(d).run(print_data=True)
benchmark/kernels/bench_fused_sigmoid_mul.py
+59
View File
@@ -0,0 +1,59 @@
"""Benchmark fused_sigmoid_mul: auto-dispatch vs PyTorch eager.
The auto-dispatch path uses a strided Triton kernel for Qwen3.5 MoE attention
output gates.
Both paths start from a strided 3D gate (from torch.chunk) to ensure
a fair comparison — the reshape/contiguous cost is included.
"""
import torch
import triton
from sglang.srt.layers.elementwise import fused_sigmoid_mul
NUM_HEADS = 32
HEAD_DIM = 256
HIDDEN_DIM = NUM_HEADS * HEAD_DIM # 8192
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens"],
x_vals=[1, 2, 4, 8, 16, 1024, 2048, 4096, 8192],
line_arg="impl",
line_vals=["auto", "auto_inplace", "pytorch_from_strided"],
line_names=[
"fused_sigmoid_mul (auto)",
"fused_sigmoid_mul (auto, inplace)",
"PyTorch eager (incl. reshape)",
],
styles=[("blue", "-"), ("green", "-"), ("orange", "--")],
ylabel="us",
plot_name="fused_sigmoid_mul_qwen3_5_moe_target",
args={},
)
)
def bench(num_tokens, impl, dtype=torch.bfloat16):
q_gate = torch.randn(
num_tokens, NUM_HEADS, 2 * HEAD_DIM, dtype=dtype, device="cuda"
)
_, gate_strided = torch.chunk(q_gate, 2, dim=-1)
attn_output = torch.randn(num_tokens, HIDDEN_DIM, dtype=dtype, device="cuda")
if impl == "auto":
fn = lambda: fused_sigmoid_mul(attn_output, gate_strided, inplace=False)
elif impl == "auto_inplace":
fn = lambda: fused_sigmoid_mul(attn_output, gate_strided, inplace=True)
else:
# Fair comparison: include reshape cost in every iteration
def fn():
g = gate_strided.contiguous().view(num_tokens, HIDDEN_DIM)
return attn_output * torch.sigmoid(g)
ms = triton.testing.do_bench(fn, warmup=100, rep=200)
return ms * 1000
if __name__ == "__main__":
bench.run(print_data=True)
python/sglang/srt/layers/elementwise.py
+172
View File
@@ -4,6 +4,7 @@ import torch
import triton
import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.layers.triton_ops.softcap import softcap_out as fused_softcap
from sglang.srt.utils import is_hip
from sglang.srt.utils.custom_op import register_custom_op
@@ -528,3 +529,174 @@ def silu_and_mul_triton(
return out_hidden_states, out_scales
else:
return out_hidden_states, None
@triton.jit
def _fused_sigmoid_mul_kernel(
output_ptr,
attn_output_ptr,
gate_ptr,
gate_stride_row,
gate_stride_head,
hidden_dim: tl.constexpr,
HEAD_DIM: tl.constexpr,
BLOCK_H: tl.constexpr,
):
"""Fuse sigmoid(gate) * attn_output into a single kernel."""
pid_row = tl.program_id(0).to(tl.int64)
pid_block = tl.program_id(1)
offsets = pid_block * BLOCK_H + tl.arange(0, BLOCK_H)
mask = offsets < hidden_dim
head = offsets // HEAD_DIM
d = offsets - head * HEAD_DIM
attn_off = pid_row * hidden_dim + offsets
attn = tl.load(attn_output_ptr + attn_off, mask=mask, other=0.0).to(tl.float32)
gate_off = pid_row * gate_stride_row + head * gate_stride_head + d
g = tl.load(gate_ptr + gate_off, mask=mask, other=0.0).to(tl.float32)
result = attn * tl.sigmoid(g)
tl.store(output_ptr + attn_off, result, mask=mask)
def fused_sigmoid_mul(
attn_output: torch.Tensor,
gate: torch.Tensor,
inplace: bool = False,
) -> torch.Tensor:
"""
Fused sigmoid-mul for attention output gating.
Equivalent to: attn_output * sigmoid(gate)
The production Qwen3.5 path passes a 3D strided gate. A single hidden-block
Triton kernel handles both that path and flat contiguous inputs.
When inplace=True, writes result back to attn_output and returns it.
Supports strided gate: if gate is 3D (num_tokens, num_heads, head_dim)
and attn_output is 2D (num_tokens, hidden_dim), the kernel reads gate
via explicit strides without requiring a contiguous copy.
"""
if gate.ndim == 3 and attn_output.ndim == 2:
# Strided gate path: gate is 3D (num_tokens, num_heads, head_dim)
num_tokens, num_heads, head_dim = gate.shape
hidden_dim = num_heads * head_dim
assert attn_output.shape == (num_tokens, hidden_dim)
gate_stride_row = gate.stride(0)
gate_stride_head = gate.stride(1)
else:
# Flat path: both tensors have the same shape
assert (
attn_output.shape == gate.shape
), "attn_output and gate must have the same shape"
hidden_dim = attn_output.shape[-1]
num_tokens = attn_output.numel() // hidden_dim
head_dim = hidden_dim
gate_stride_row = hidden_dim
gate_stride_head = hidden_dim
out = attn_output if inplace else torch.empty_like(attn_output)
block_h = 1024 if num_tokens < 1024 else 2048
grid = (num_tokens, triton.cdiv(hidden_dim, block_h))
_fused_sigmoid_mul_kernel[grid](
out,
attn_output,
gate,
gate_stride_row,
gate_stride_head,
hidden_dim,
HEAD_DIM=head_dim,
BLOCK_H=block_h,
num_warps=4,
)
return out
@triton.jit
def _fused_gate_sigmoid_mul_add_kernel(
hidden_states_ptr, # [num_tokens, hidden_dim]
gate_weight_ptr, # [hidden_dim]
shared_output_ptr, # [num_tokens, hidden_dim]
final_hidden_states_ptr, # [num_tokens, hidden_dim]
hidden_dim: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
USE_PDL: tl.constexpr = False,
):
pid = tl.program_id(axis=0).to(tl.int64)
row_offset = pid * hidden_dim
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < hidden_dim
w = tl.load(gate_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
if USE_PDL:
tl.extra.cuda.gdc_wait()
h = tl.load(hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0).to(
tl.float32
)
s = tl.load(shared_output_ptr + row_offset + offsets, mask=mask, other=0.0).to(
tl.float32
)
f = tl.load(
final_hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0
).to(tl.float32)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
gate_val = tl.sigmoid(tl.sum(h * w, axis=0))
result = f + gate_val * s
tl.store(final_hidden_states_ptr + row_offset + offsets, result, mask=mask)
def fused_gate_sigmoid_mul_add(
hidden_states: torch.Tensor,
gate_weight: torch.Tensor,
shared_output: torch.Tensor,
final_hidden_states: torch.Tensor,
) -> None:
"""
Fused gate-sigmoid-mul-add for MoE shared expert gating.
Equivalent to:
gate = hidden_states @ gate_weight
final_hidden_states += sigmoid(gate).unsqueeze(1) * shared_output
"""
assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
assert gate_weight.is_contiguous(), "gate_weight must be contiguous"
assert shared_output.is_contiguous(), "shared_output must be contiguous"
assert final_hidden_states.is_contiguous(), "final_hidden_states must be contiguous"
num_tokens, hidden_dim = hidden_states.shape
assert gate_weight.shape == (hidden_dim,)
assert shared_output.shape == (num_tokens, hidden_dim)
assert final_hidden_states.shape == (num_tokens, hidden_dim)
max_warps = 16 if _is_hip else 32
config = {
"BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
"num_warps": max(
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4
),
}
if num_tokens >= 1024:
config["num_warps"] = min(config["num_warps"], 8)
pdl_kwargs = {"USE_PDL": True, "launch_pdl": True} if is_arch_support_pdl() else {}
_fused_gate_sigmoid_mul_add_kernel[(num_tokens,)](
hidden_states,
gate_weight,
shared_output,
final_hidden_states,
hidden_dim=hidden_dim,
**config,
**pdl_kwargs,
)
python/sglang/srt/models/qwen2_moe.py
+33 -6
View File
@@ -53,6 +53,7 @@ from sglang.srt.layers.dp_attention import (
get_attention_tp_size,
is_dp_attention_enabled,
)
from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
MergedColumnParallelLinear,
@@ -265,6 +266,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
layer_id=layer_id,
)
# Disable inplace MoE when fused gate will need hidden_states after experts
_needs_hidden_after_experts = (
config.shared_expert_intermediate_size > 0
and not self.enable_shared_expert_fusion
)
self.experts = get_moe_impl_class(quant_config)(
layer_id=self.layer_id,
top_k=(
@@ -285,6 +291,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
prefix=add_prefix("experts", prefix),
routing_method_type=RoutingMethodType.RenormalizeNaive,
num_fused_shared_experts=self.num_fused_shared_experts,
inplace=not _needs_hidden_after_experts,
)
self.gate = ReplicatedLinear(
@@ -400,11 +407,13 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
router_logits=topk_output.router_logits,
)
def _forward_shared_experts(self, hidden_states: torch.Tensor):
def _forward_shared_experts(
self, hidden_states: torch.Tensor, apply_gate: bool = True
):
shared_output = None
if self.shared_expert is not None:
shared_output = self.shared_expert(hidden_states)
if self.shared_expert_gate is not None:
if self.shared_expert_gate is not None and apply_gate:
if use_intel_amx_backend(self.shared_expert_gate):
shared_output = torch.ops.sgl_kernel.fused_linear_sigmoid_mul(
hidden_states,
@@ -476,11 +485,14 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
use_fused_gate: bool = False,
) -> torch.Tensor:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
shared_output = (
self._forward_shared_experts(hidden_states.clone())
self._forward_shared_experts(
hidden_states.clone(), apply_gate=not use_fused_gate
)
if self.shared_expert is not None
else None
)
@@ -525,6 +537,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
if get_moe_a2a_backend().is_deepep():
return self._forward_deepep(hidden_states, forward_batch)
use_fused_gate = (
self.shared_expert_gate is not None
and not use_intel_amx_backend(self.shared_expert_gate)
)
if hidden_states.shape[0] == 0:
# M=0 guard for idle DP ranks: skip shared_experts and gate
# (which crash on empty tensors in FP4 GEMM), but still call
@@ -534,14 +551,24 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
final_hidden_states = self.experts(hidden_states, topk_output)
elif self.alt_stream is not None and get_is_capture_mode():
final_hidden_states, shared_output = self.forward_normal_dual_stream(
hidden_states
hidden_states, use_fused_gate=use_fused_gate
)
else:
shared_output = self._forward_shared_experts(hidden_states)
shared_output = self._forward_shared_experts(
hidden_states, apply_gate=not use_fused_gate
)
final_hidden_states = self._forward_router_experts(hidden_states)
if shared_output is not None:
final_hidden_states += shared_output
if use_fused_gate:
fused_gate_sigmoid_mul_add(
hidden_states,
self.shared_expert_gate.weight.squeeze(),
shared_output,
final_hidden_states,
)
else:
final_hidden_states += shared_output
if (
self.tp_size > 1
and not should_skip_post_experts_all_reduce(
python/sglang/srt/models/qwen3_5.py
+3 -10
View File
@@ -47,6 +47,7 @@ from sglang.srt.layers.dp_attention import (
get_attention_tp_size,
is_dp_attention_enabled,
)
from sglang.srt.layers.elementwise import fused_sigmoid_mul
# Layers - Others
from sglang.srt.layers.layernorm import GemmaRMSNorm
@@ -887,7 +888,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
q_gate = q_gate.view(*orig_shape, self.num_heads, -1)
q, gate = torch.chunk(q_gate, 2, dim=-1)
q = q.reshape(*orig_shape, -1)
gate = gate.reshape(*orig_shape, -1)
# gate stays as 3D strided view; fused_sigmoid_mul handles it directly
else:
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
gate = None
@@ -974,15 +975,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
attn_output = self.attn(q, k, v, forward_batch)
if self.attn_output_gate:
if _is_hip:
from sglang.jit_kernel.triton.sigmoid_gate_mul import (
sigmoid_gate_mul,
)
attn_output = sigmoid_gate_mul(attn_output, gate)
else:
gate = torch.sigmoid(gate)
attn_output = attn_output * gate
attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True)
output, _ = self.o_proj(attn_output)
return output
test/manual/layers/test_fused_gate_sigmoid_mul_add.py
+73
View File
@@ -0,0 +1,73 @@
import itertools
import pytest
import torch
from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add
DTYPES = [torch.float16, torch.bfloat16]
TOKEN_COUNTS = [1, 2, 4, 8, 16, 64, 512, 1024, 2048, 4096, 8192]
HIDDEN_DIMS = [2048, 3072, 4096, 6144]
def _reference(hidden_states, gate_weight, shared_output, final_hidden_states):
gate = hidden_states @ gate_weight
final_hidden_states += torch.sigmoid(gate).unsqueeze(1) * shared_output
@pytest.fixture(autouse=True)
def seed():
torch.manual_seed(42)
@pytest.mark.parametrize(
"num_tokens, hidden_dim, dtype",
list(itertools.product(TOKEN_COUNTS, HIDDEN_DIMS, DTYPES)),
)
def test_correctness(num_tokens, hidden_dim, dtype):
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate_weight = torch.randn(hidden_dim, dtype=dtype, device="cuda")
shared_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
final_ref = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
final_test = final_ref.clone()
_reference(hidden_states, gate_weight, shared_output, final_ref)
fused_gate_sigmoid_mul_add(hidden_states, gate_weight, shared_output, final_test)
torch.testing.assert_close(final_test, final_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", DTYPES)
def test_gate_near_zero(dtype):
num_tokens, hidden_dim = 16, 2048
hs = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gw = torch.zeros(hidden_dim, dtype=dtype, device="cuda")
so = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
f_ref = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
f_test = f_ref.clone()
_reference(hs, gw, so, f_ref)
fused_gate_sigmoid_mul_add(hs, gw, so, f_test)
torch.testing.assert_close(f_test, f_ref, rtol=1e-2, atol=1e-2)
def test_inplace_semantics():
num_tokens, hidden_dim = 32, 2048
hs = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
gw = torch.randn(hidden_dim, dtype=torch.float16, device="cuda")
so = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
fhs = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
original_ptr = fhs.data_ptr()
fused_gate_sigmoid_mul_add(hs, gw, so, fhs)
assert fhs.data_ptr() == original_ptr
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v"]))
test/manual/layers/test_fused_sigmoid_mul.py
+134
View File
@@ -0,0 +1,134 @@
import itertools
import pytest
import torch
from sglang.srt.layers.elementwise import fused_sigmoid_mul
DTYPES = [torch.float16, torch.bfloat16]
TOKEN_COUNTS = [1, 2, 4, 8, 16, 64, 512, 1024, 2048, 4096, 8192]
HIDDEN_DIMS = [2048, 3072, 4096, 6144]
NUM_HEADS = [1, 28]
def _reference(attn_output, gate):
return attn_output * torch.sigmoid(gate)
@pytest.fixture(autouse=True)
def seed():
torch.manual_seed(42)
@pytest.mark.parametrize(
"num_tokens, hidden_dim, dtype",
list(itertools.product(TOKEN_COUNTS, HIDDEN_DIMS, DTYPES)),
)
def test_correctness(num_tokens, hidden_dim, dtype):
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
ref = _reference(attn_output, gate)
out = fused_sigmoid_mul(attn_output, gate)
torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize(
"num_tokens, num_heads, dtype",
list(itertools.product(TOKEN_COUNTS, NUM_HEADS, DTYPES)),
)
def test_3d_shape(num_tokens, num_heads, dtype):
"""Test with 3D contiguous tensors (num_tokens, num_heads, head_dim)."""
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
head_dim = 128
attn_output = torch.randn(
num_tokens, num_heads, head_dim, dtype=dtype, device="cuda"
)
gate = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype, device="cuda")
ref = _reference(attn_output, gate)
out = fused_sigmoid_mul(attn_output, gate)
torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize(
"num_tokens, num_heads, dtype",
list(itertools.product(TOKEN_COUNTS, NUM_HEADS, DTYPES)),
)
def test_strided_gate(num_tokens, num_heads, dtype):
"""Test strided gate path: attn_output is 2D, gate is 3D non-contiguous from chunk."""
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
head_dim = 128
hidden_dim = num_heads * head_dim
# Simulate the real pattern: chunk produces non-contiguous views
q_gate = torch.randn(
num_tokens, num_heads, 2 * head_dim, dtype=dtype, device="cuda"
)
_, gate = torch.chunk(q_gate, 2, dim=-1)
# gate is non-contiguous when num_tokens > 1 or num_heads > 1
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate_flat = gate.reshape(num_tokens, hidden_dim)
ref = _reference(attn_output, gate_flat)
out = fused_sigmoid_mul(attn_output, gate, inplace=False)
torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize(
"num_tokens, dtype", list(itertools.product(TOKEN_COUNTS, DTYPES))
)
def test_qwen3_5_moe_target_strided_gate(num_tokens, dtype):
"""Qwen3.5 MoE target config: 32 attention heads, head_dim 256."""
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
num_heads, head_dim = 32, 256
hidden_dim = num_heads * head_dim
q_gate = torch.randn(
num_tokens, num_heads, 2 * head_dim, dtype=dtype, device="cuda"
)
_, gate = torch.chunk(q_gate, 2, dim=-1)
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
ref = _reference(attn_output, gate.reshape(num_tokens, hidden_dim))
out = fused_sigmoid_mul(attn_output, gate, inplace=False)
torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", DTYPES)
def test_gate_near_zero(dtype):
num_tokens, hidden_dim = 16, 2048
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate = torch.zeros(num_tokens, hidden_dim, dtype=dtype, device="cuda")
ref = _reference(attn_output, gate)
out = fused_sigmoid_mul(attn_output, gate)
torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2)
def test_returns_new_tensor():
num_tokens, hidden_dim = 32, 2048
attn_output = torch.randn(
num_tokens, hidden_dim, dtype=torch.float16, device="cuda"
)
gate = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
out = fused_sigmoid_mul(attn_output, gate)
assert out.data_ptr() != attn_output.data_ptr()
assert out.data_ptr() != gate.data_ptr()
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v"]))