From bf765a0f2ac83ea9c05ac1168835d1ccb87cbedf Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 18 Jul 2026 07:58:36 +0800 Subject: [PATCH 1/2] optimize: avoid fla l2-norm recompilation by token count (#31558) Signed-off-by: Xiake Sun --- python/sglang/srt/layers/attention/fla/l2norm.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/fla/l2norm.py b/python/sglang/srt/layers/attention/fla/l2norm.py index 104a312a2053..8e981cb25c3e 100644 --- a/python/sglang/srt/layers/attention/fla/l2norm.py +++ b/python/sglang/srt/layers/attention/fla/l2norm.py @@ -151,13 +151,12 @@ def l2norm_fwd_kernel1( # ], # key=["D", "NB"], # ) -@triton.jit +@triton.jit(do_not_specialize=["T"]) def l2norm_fwd_kernel( x, y, eps, - NB: tl.constexpr, - T: tl.constexpr, + T, D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, @@ -191,7 +190,6 @@ def l2norm_fwd( raise RuntimeError("This layer doesn't support feature dim >= 64KB.") if D <= 512: - NB = triton.cdiv(T, 2048) def grid(meta): return (triton.cdiv(T, meta["BT"]),) @@ -200,7 +198,6 @@ def grid(meta): x, y, eps, - NB=NB, T=T, D=D, BD=BD, From c0f8a277e8d2c9593dc1e586d497ddd31b3f441c Mon Sep 17 00:00:00 2001 From: zhengyao Date: Fri, 31 Jul 2026 08:49:45 +0000 Subject: [PATCH 2/2] perf(gdn): fuse Q/K L2 normalization on AMD Reduce GDN prefill launches for benchmarked BF16/FP32 GPU shapes while preserving the original path on unsupported platforms, dtypes, layouts, and launch-bound wide shapes. Signed-off-by: Xiake Sun --- .../sglang/srt/layers/attention/fla/chunk.py | 19 +- .../sglang/srt/layers/attention/fla/l2norm.py | 132 +++++--- .../jit/benchmark/bench_gdn_l2norm.py | 56 ++++ test/registered/jit/test_fused_gdn_l2norm.py | 285 ++++++++++++++++++ 4 files changed, 443 insertions(+), 49 deletions(-) create mode 100644 test/registered/jit/benchmark/bench_gdn_l2norm.py create mode 100644 test/registered/jit/test_fused_gdn_l2norm.py diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index c243ba22fd60..2e5dfc92957b 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -14,13 +14,20 @@ from sglang.srt.layers.attention.fla.index import ( prepare_chunk_indices, ) -from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd +from sglang.srt.layers.attention.fla.l2norm import ( + can_fuse_l2norm_qk, + fused_l2norm_qk, + l2norm_fwd, +) from sglang.srt.layers.attention.fla.utils import ( SUPPRESS_LEVEL, autocast_custom_fwd, input_guard, is_intel, ) +from sglang.srt.utils import is_hip + +_is_hip = is_hip() if is_intel: from sglang.srt.hardware_backend.xpu.kernels.fla.chunk_delta_h import ( @@ -102,12 +109,12 @@ def forward( cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False, ): - q_orig = q - k_orig = k - if use_qk_l2norm_in_kernel: - q = l2norm_fwd(q) - k = l2norm_fwd(k) + if _is_hip and can_fuse_l2norm_qk(q, k): + q, k = fused_l2norm_qk(q, k) + else: + q = l2norm_fwd(q) + k = l2norm_fwd(k) chunk_indices = ( prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) diff --git a/python/sglang/srt/layers/attention/fla/l2norm.py b/python/sglang/srt/layers/attention/fla/l2norm.py index 8e981cb25c3e..ff0ed1ef3ebe 100644 --- a/python/sglang/srt/layers/attention/fla/l2norm.py +++ b/python/sglang/srt/layers/attention/fla/l2norm.py @@ -14,47 +14,89 @@ BT_LIST = [8, 16, 32, 64, 128] -@triton.jit +@triton.jit(do_not_specialize=["TQ", "TK"]) def fused_l2norm_qk_kernel( q, k, q_out, k_out, eps, - T: tl.constexpr, + TQ, + TK, D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, ): i_t = tl.program_id(0) + rows = i_t * BT + tl.arange(0, BT) cols = tl.arange(0, BD) - mask = (rows[:, None] < T) & (cols[None, :] < D) - offsets = rows[:, None] * D + cols[None, :] + col_mask = cols < D - q_block = tl.load(q + offsets, mask=mask, other=0.0).to(tl.float32) - q_block = q_block / tl.sqrt(tl.sum(q_block * q_block, axis=1) + eps)[:, None] - tl.store(q_out + offsets, q_block.to(q_out.dtype.element_ty), mask=mask) + q_mask = (rows[:, None] < TQ) & col_mask[None, :] + q_offs = rows[:, None] * D + cols[None, :] + b_q = tl.load(q + q_offs, mask=q_mask, other=0.0).to(tl.float32) + q_var = tl.sum(b_q * b_q, axis=1) + b_q_out = b_q / tl.sqrt(q_var + eps)[:, None] + tl.store(q_out + q_offs, b_q_out.to(q_out.dtype.element_ty), mask=q_mask) - k_block = tl.load(k + offsets, mask=mask, other=0.0).to(tl.float32) - k_block = k_block / tl.sqrt(tl.sum(k_block * k_block, axis=1) + eps)[:, None] - tl.store(k_out + offsets, k_block.to(k_out.dtype.element_ty), mask=mask) + k_mask = (rows[:, None] < TK) & col_mask[None, :] + k_offs = rows[:, None] * D + cols[None, :] + b_k = tl.load(k + k_offs, mask=k_mask, other=0.0).to(tl.float32) + k_var = tl.sum(b_k * b_k, axis=1) + b_k_out = b_k / tl.sqrt(k_var + eps)[:, None] + tl.store(k_out + k_offs, b_k_out.to(k_out.dtype.element_ty), mask=k_mask) -@triton.jit -def fused_l2norm_qk_kernel1(q, k, q_out, k_out, D, BD: tl.constexpr, eps): - row = tl.program_id(0) +@triton.jit(do_not_specialize=["TQ", "TK"]) +def fused_l2norm_qk_kernel1( + q, + k, + q_out, + k_out, + D, + BD: tl.constexpr, + eps, + TQ, + TK, +): + i_t = tl.program_id(0) cols = tl.arange(0, BD) mask = cols < D - offsets = row * D + cols - - q_block = tl.load(q + offsets, mask=mask, other=0.0).to(tl.float32) - q_block = q_block / tl.sqrt(tl.sum(q_block * q_block, axis=0) + eps) - tl.store(q_out + offsets, q_block.to(q_out.dtype.element_ty), mask=mask) - - k_block = tl.load(k + offsets, mask=mask, other=0.0).to(tl.float32) - k_block = k_block / tl.sqrt(tl.sum(k_block * k_block, axis=0) + eps) - tl.store(k_out + offsets, k_block.to(k_out.dtype.element_ty), mask=mask) + offs = i_t * D + cols + + if i_t < TQ: + b_q = tl.load(q + offs, mask=mask, other=0.0).to(tl.float32) + q_var = tl.sum(b_q * b_q, axis=0) + b_q_out = b_q / tl.sqrt(q_var + eps) + tl.store(q_out + offs, b_q_out.to(q_out.dtype.element_ty), mask=mask) + + if i_t < TK: + b_k = tl.load(k + offs, mask=mask, other=0.0).to(tl.float32) + k_var = tl.sum(b_k * b_k, axis=0) + b_k_out = b_k / tl.sqrt(k_var + eps) + tl.store(k_out + offs, b_k_out.to(k_out.dtype.element_ty), mask=mask) + + +def can_fuse_l2norm_qk(q: torch.Tensor, k: torch.Tensor) -> bool: + if q.device != k.device or q.dtype != k.dtype: + return False + if not q.is_cuda or q.dtype not in (torch.bfloat16, torch.float32): + return False + if q.shape != k.shape: + return False + q_rows = q.numel() // q.shape[-1] + # Wide-shape sweep is positive through D=512 except for the launch-bound + # D=512, rows<32 corner. Keep unbenchmarked larger head dims on the existing + # two-kernel path. + if q.shape[-1] > 512 or (q.shape[-1] == 512 and q_rows < 32): + return False + if q.stride(-1) != 1 or k.stride(-1) != 1: + return False + # `view(-1, D)` in the kernel wrapper must be valid without an implicit copy. + if not q.is_contiguous() or not k.is_contiguous(): + return False + return True def fused_l2norm_qk( @@ -63,55 +105,59 @@ def fused_l2norm_qk( eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None, ) -> tuple[torch.Tensor, torch.Tensor]: - q_shape, k_shape = q.shape, k.shape + if not can_fuse_l2norm_qk(q, k): + raise ValueError("Incompatible Q/K tensors for fused_l2norm_qk") + + q_shape_og = q.shape + k_shape_og = k.shape q_flat = q.view(-1, q.shape[-1]) k_flat = k.view(-1, k.shape[-1]) - if q_flat.shape != k_flat.shape: - raise ValueError( - f"Fused Q/K L2Norm requires matching shapes, got {q_shape} and {k_shape}" - ) + TQ, TK, D = q_flat.shape[0], k_flat.shape[0], q_flat.shape[-1] if output_dtype is None: - q_out, k_out = torch.empty_like(q_flat), torch.empty_like(k_flat) + q_out = torch.empty_like(q_flat) + k_out = torch.empty_like(k_flat) else: q_out = torch.empty_like(q_flat, dtype=output_dtype) k_out = torch.empty_like(k_flat, dtype=output_dtype) - tokens, dim = q_flat.shape - max_fused_size = 65536 // q.element_size() - block_dim = min(max_fused_size, triton.next_power_of_2(dim)) - if dim > block_dim: + MAX_FUSED_SIZE = 65536 // q.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: raise RuntimeError("This layer doesn't support feature dim >= 64KB.") - if dim <= 512: - block_tokens = 16 - fused_l2norm_qk_kernel[(triton.cdiv(tokens, block_tokens),)]( + if D <= 512: + BT = 16 + fused_l2norm_qk_kernel[(triton.cdiv(max(TQ, TK), BT),)]( q_flat, k_flat, q_out, k_out, eps, - T=tokens, - D=dim, - BT=block_tokens, - BD=block_dim, + TQ=TQ, + TK=TK, + D=D, + BT=BT, + BD=BD, num_warps=8, num_stages=3, ) else: - fused_l2norm_qk_kernel1[(tokens,)]( + fused_l2norm_qk_kernel1[(max(TQ, TK),)]( q_flat, k_flat, q_out, k_out, - D=dim, - BD=block_dim, eps=eps, + D=D, + BD=BD, + TQ=TQ, + TK=TK, num_warps=8, num_stages=3, ) - return q_out.view(q_shape), k_out.view(k_shape) + return q_out.view(q_shape_og), k_out.view(k_shape_og) # @triton.autotune( diff --git a/test/registered/jit/benchmark/bench_gdn_l2norm.py b/test/registered/jit/benchmark/bench_gdn_l2norm.py new file mode 100644 index 000000000000..b398fa0f6ccf --- /dev/null +++ b/test/registered/jit/benchmark/bench_gdn_l2norm.py @@ -0,0 +1,56 @@ +import torch + +from sglang.jit_kernel.benchmark import marker +from sglang.srt.layers.attention.fla.l2norm import fused_l2norm_qk, l2norm_fwd +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci( + est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) +register_amd_ci(est_time=8, stage="jit-kernel-benchmark", runner_config="amd") + + +def _run_fused(q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return fused_l2norm_qk(q, k) + + +def _run_separate( + q: torch.Tensor, k: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + return l2norm_fwd(q), l2norm_fwd(k) + + +FN_MAP = { + "fused": _run_fused, + "separate": _run_separate, +} + + +@marker.parametrize("tokens", [15, 16, 17, 257, 1024], [257]) +@marker.parametrize( + "local_heads,head_dim", + [ + (2, 128), + (4, 128), + (8, 128), + (16, 128), + (8, 256), + ], +) +@marker.benchmark("impl", ["fused", "separate"]) +def benchmark(tokens: int, local_heads: int, head_dim: int, impl: str): + q = torch.randn(tokens, local_heads, head_dim, dtype=torch.bfloat16, device="cuda") + k = torch.randn(tokens, local_heads, head_dim, dtype=torch.bfloat16, device="cuda") + return marker.do_bench( + FN_MAP[impl], + input_args=(q, k), + warmup_iters=80, + replay_iters=1200, + graph_clone_args=(0, 1), + memory_args=(q, k), + memory_output=None, + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/jit/test_fused_gdn_l2norm.py b/test/registered/jit/test_fused_gdn_l2norm.py new file mode 100644 index 000000000000..375342034a4b --- /dev/null +++ b/test/registered/jit/test_fused_gdn_l2norm.py @@ -0,0 +1,285 @@ +import sys + +import pytest +import torch + +from sglang.srt.layers.attention.fla import chunk as chunk_mod +from sglang.srt.layers.attention.fla.l2norm import ( + can_fuse_l2norm_qk, + fused_l2norm_qk, + l2norm_fwd, +) +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=8, stage="jit-kernel-unit", runner_config="amd") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("tokens", [1, 15, 16, 17, 257]) +@pytest.mark.parametrize("local_heads", [2, 4, 8, 16]) +def test_fused_l2norm_qk_matches_separate_qwen35_equal_heads( + dtype, tokens, local_heads +): + torch.manual_seed(2026) + head_dim = 128 + q = torch.randn(tokens, local_heads, head_dim, dtype=dtype, device="cuda") + k = torch.randn(tokens, local_heads, head_dim, dtype=dtype, device="cuda") + + q_ref = l2norm_fwd(q) + k_ref = l2norm_fwd(k) + q_fused, k_fused = fused_l2norm_qk(q, k) + + atol = 2e-2 if dtype == torch.bfloat16 else 1e-5 + rtol = 2e-2 if dtype == torch.bfloat16 else 1e-5 + torch.testing.assert_close(q_fused, q_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(k_fused, k_ref, atol=atol, rtol=rtol) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fused_l2norm_qk_generic_d256(dtype): + torch.manual_seed(99) + q = torch.randn(257, 8, 256, dtype=dtype, device="cuda") + k = torch.randn(257, 8, 256, dtype=dtype, device="cuda") + + q_ref = l2norm_fwd(q) + k_ref = l2norm_fwd(k) + q_fused, k_fused = fused_l2norm_qk(q, k) + + atol = 2e-2 if dtype == torch.bfloat16 else 1e-5 + rtol = 2e-2 if dtype == torch.bfloat16 else 1e-5 + torch.testing.assert_close(q_fused, q_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(k_fused, k_ref, atol=atol, rtol=rtol) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_can_fuse_l2norm_qk_rejects_asymmetric_rows(): + q = torch.randn(17, 16, 128, dtype=torch.bfloat16, device="cuda") + k = torch.randn(17, 8, 128, dtype=torch.bfloat16, device="cuda") + assert not can_fuse_l2norm_qk(q, k) + + +def test_can_fuse_l2norm_qk_rejects_cpu_tensors(): + q = torch.randn(4, 2, 128, dtype=torch.bfloat16) + k = torch.randn_like(q) + assert not can_fuse_l2norm_qk(q, k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_can_fuse_l2norm_qk_rejects_unvalidated_dtype(): + q = torch.randn(4, 2, 128, dtype=torch.float16, device="cuda") + k = torch.randn_like(q) + assert not can_fuse_l2norm_qk(q, k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_can_fuse_l2norm_qk_rejects_different_shapes_with_equal_rows(): + q = torch.randn(4, 2, 128, dtype=torch.bfloat16, device="cuda") + k = torch.randn(8, 1, 128, dtype=torch.bfloat16, device="cuda") + assert q.numel() == k.numel() + assert not can_fuse_l2norm_qk(q, k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +@pytest.mark.parametrize( + "rows,expected", + [ + (16, False), + (32, True), + ], +) +def test_can_fuse_l2norm_qk_d512_small_row_guard(rows, expected): + q = torch.randn(rows, 1, 512, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + assert can_fuse_l2norm_qk(q, k) is expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_can_fuse_l2norm_qk_rejects_unbenchmarked_large_head_dim(): + q = torch.randn(64, 1, 1024, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + assert not can_fuse_l2norm_qk(q, k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_chunk_uses_fused_qk_l2norm(monkeypatch): + seen = {"fused": 0, "separate": 0} + + def fake_can_fuse(q, k): + return True + + def fake_fused(q, k, eps=1e-6, output_dtype=None): + seen["fused"] += 1 + return q + 3, k + 5 + + def fake_l2norm(*args, **kwargs): + seen["separate"] += 1 + raise AssertionError("Fallback l2norm_fwd should not run in fused path") + + def fake_chunk_fwd(**kwargs): + q = kwargs["q"] + k = kwargs["k"] + assert torch.allclose(q, q_in + 3) + assert torch.allclose(k, k_in + 5) + o = torch.zeros_like(kwargs["v"]) + return ( + kwargs["g"], + o, + torch.empty(0, device=o.device), + None, + kwargs["initial_state"], + None, + ) + + monkeypatch.setattr(chunk_mod, "_is_hip", True, raising=False) + monkeypatch.setattr(chunk_mod, "can_fuse_l2norm_qk", fake_can_fuse) + monkeypatch.setattr(chunk_mod, "fused_l2norm_qk", fake_fused) + monkeypatch.setattr(chunk_mod, "l2norm_fwd", fake_l2norm) + monkeypatch.setattr(chunk_mod, "chunk_gated_delta_rule_fwd", fake_chunk_fwd) + + q_in = torch.randn(1, 4, 2, 128, dtype=torch.bfloat16, device="cuda") + k_in = torch.randn(1, 4, 2, 128, dtype=torch.bfloat16, device="cuda") + v = torch.randn(1, 4, 2, 64, dtype=torch.bfloat16, device="cuda") + g = torch.randn(1, 4, 2, dtype=torch.bfloat16, device="cuda") + beta = torch.sigmoid(torch.randn(1, 4, 2, dtype=torch.bfloat16, device="cuda")) + initial_state = torch.randn(1, 2, 64, 128, dtype=torch.float32, device="cuda") + initial_state_indices = torch.tensor([0], dtype=torch.int32, device="cuda") + + chunk_mod.chunk_gated_delta_rule( + q=q_in, + k=k_in, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + assert seen["fused"] == 1 + assert seen["separate"] == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_chunk_falls_back_to_two_l2norm_calls(monkeypatch): + seen = {"fused": 0, "separate": 0} + + def fake_can_fuse(q, k): + return False + + def fake_fused(*args, **kwargs): + seen["fused"] += 1 + raise AssertionError("fused_l2norm_qk should not run in fallback path") + + def fake_l2norm(x, eps=1e-6, output_dtype=None): + seen["separate"] += 1 + return x + 7 + + def fake_chunk_fwd(**kwargs): + q = kwargs["q"] + k = kwargs["k"] + assert torch.allclose(q, q_in + 7) + assert torch.allclose(k, k_in + 7) + o = torch.zeros_like(kwargs["v"]) + return ( + kwargs["g"], + o, + torch.empty(0, device=o.device), + None, + kwargs["initial_state"], + None, + ) + + monkeypatch.setattr(chunk_mod, "_is_hip", True, raising=False) + monkeypatch.setattr(chunk_mod, "can_fuse_l2norm_qk", fake_can_fuse) + monkeypatch.setattr(chunk_mod, "fused_l2norm_qk", fake_fused) + monkeypatch.setattr(chunk_mod, "l2norm_fwd", fake_l2norm) + monkeypatch.setattr(chunk_mod, "chunk_gated_delta_rule_fwd", fake_chunk_fwd) + + q_in = torch.randn(1, 4, 2, 128, dtype=torch.bfloat16, device="cuda") + k_in = torch.randn(1, 4, 2, 128, dtype=torch.bfloat16, device="cuda") + v = torch.randn(1, 4, 2, 64, dtype=torch.bfloat16, device="cuda") + g = torch.randn(1, 4, 2, dtype=torch.bfloat16, device="cuda") + beta = torch.sigmoid(torch.randn(1, 4, 2, dtype=torch.bfloat16, device="cuda")) + initial_state = torch.randn(1, 2, 64, 128, dtype=torch.float32, device="cuda") + initial_state_indices = torch.tensor([0], dtype=torch.int32, device="cuda") + + chunk_mod.chunk_gated_delta_rule( + q=q_in, + k=k_in, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + assert seen["fused"] == 0 + assert seen["separate"] == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_chunk_non_hip_uses_two_l2norm_even_if_fusible(monkeypatch): + seen = {"fused": 0, "separate": 0} + + def fake_can_fuse(q, k): + return True + + def fake_fused(*args, **kwargs): + seen["fused"] += 1 + return args[0], args[1] + + def fake_l2norm(x, eps=1e-6, output_dtype=None): + seen["separate"] += 1 + return x + 11 + + def fake_chunk_fwd(**kwargs): + assert torch.allclose(kwargs["q"], q_in + 11) + assert torch.allclose(kwargs["k"], k_in + 11) + o = torch.zeros_like(kwargs["v"]) + return ( + kwargs["g"], + o, + torch.empty(0, device=o.device), + None, + kwargs["initial_state"], + None, + ) + + monkeypatch.setattr(chunk_mod, "_is_hip", False, raising=False) + monkeypatch.setattr(chunk_mod, "can_fuse_l2norm_qk", fake_can_fuse) + monkeypatch.setattr(chunk_mod, "fused_l2norm_qk", fake_fused) + monkeypatch.setattr(chunk_mod, "l2norm_fwd", fake_l2norm) + monkeypatch.setattr(chunk_mod, "chunk_gated_delta_rule_fwd", fake_chunk_fwd) + + q_in = torch.randn(1, 4, 2, 128, dtype=torch.bfloat16, device="cuda") + k_in = torch.randn(1, 4, 2, 128, dtype=torch.bfloat16, device="cuda") + v = torch.randn(1, 4, 2, 64, dtype=torch.bfloat16, device="cuda") + g = torch.randn(1, 4, 2, dtype=torch.bfloat16, device="cuda") + beta = torch.sigmoid(torch.randn(1, 4, 2, dtype=torch.bfloat16, device="cuda")) + initial_state = torch.randn(1, 2, 64, 128, dtype=torch.float32, device="cuda") + initial_state_indices = torch.tensor([0], dtype=torch.int32, device="cuda") + + chunk_mod.chunk_gated_delta_rule( + q=q_in, + k=k_in, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + assert seen["fused"] == 0 + assert seen["separate"] == 2 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))