diff --git a/benchmark/kernels/bench_fused_gate_sigmoid_mul_add.py b/benchmark/kernels/bench_fused_gate_sigmoid_mul_add.py new file mode 100644 index 000000000..14ffe9465 --- /dev/null +++ b/benchmark/kernels/bench_fused_gate_sigmoid_mul_add.py @@ -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) diff --git a/benchmark/kernels/bench_fused_sigmoid_mul.py b/benchmark/kernels/bench_fused_sigmoid_mul.py new file mode 100644 index 000000000..1f3cd6e67 --- /dev/null +++ b/benchmark/kernels/bench_fused_sigmoid_mul.py @@ -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) diff --git a/python/sglang/srt/layers/elementwise.py b/python/sglang/srt/layers/elementwise.py index 11095f7bf..4065e9dea 100644 --- a/python/sglang/srt/layers/elementwise.py +++ b/python/sglang/srt/layers/elementwise.py @@ -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, + ) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 736c2e424..d0186f870 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -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( diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 19ab50c2a..de3af00e1 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -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 diff --git a/test/manual/layers/test_fused_gate_sigmoid_mul_add.py b/test/manual/layers/test_fused_gate_sigmoid_mul_add.py new file mode 100644 index 000000000..bf46806d3 --- /dev/null +++ b/test/manual/layers/test_fused_gate_sigmoid_mul_add.py @@ -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"])) diff --git a/test/manual/layers/test_fused_sigmoid_mul.py b/test/manual/layers/test_fused_sigmoid_mul.py new file mode 100644 index 000000000..a585254a2 --- /dev/null +++ b/test/manual/layers/test_fused_sigmoid_mul.py @@ -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"]))