diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 000000000..4315d3ed3 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,33 @@ +# Third-Party Notices + +This repository includes or adapts code from the following MIT-licensed +projects. The top-level `LICENSE` contains the MIT license text used for this +repository distribution; the original copyright notices below are retained in +the relevant source files. + +## Qwen FlashQLA + +Files under `tileops/kernels/gated_deltanet/gdn_prefill/` implement a TileOps +version of the CP-split Gated DeltaNet prefill schedule. The schedule-level +reference for the h-state / corrected-segment-start part of this implementation +comes from Qwen FlashQLA: + +- `__init__.py` +- `cp_fwd.py` +- `fused_fwd.py` +- `prepare_h.py` +- `tilelang_compat.py` +- `utils.py` + +Original notice: + +```text +Copyright (c) 2026 The Qwen team, Alibaba Group. +Licensed under the MIT License. +``` + +The TileOps versions are not direct wrappers around the FlashQLA kernels. They +adapt the CP-split scheduling idea into the TileOps operator API, BTHD dispatch, +TileLang compatibility layer, benchmarking, tests, and local replay/output +implementation. The utility file keeps the upstream copyright notice present in +the FlashQLA utility source. diff --git a/benchmarks/ops/bench_gated_deltanet_prefill.py b/benchmarks/ops/bench_gated_deltanet_prefill.py new file mode 100644 index 000000000..73263b836 --- /dev/null +++ b/benchmarks/ops/bench_gated_deltanet_prefill.py @@ -0,0 +1,151 @@ +"""Benchmark: TileOPs Gated DeltaNet inference prefill. + +When FLA is installed, record it as the independent baseline. Otherwise fall +back to a pure-torch reference so every benchmark row has a non-TileOps entry. + +The benchmark measures the serving-oriented BTHD layout because that is the +production fast path used by FLA/Qwen-style inference prefill. +""" + +import inspect +from typing import Any + +import pytest +import torch + +from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark +from benchmarks.ops.attention.manifest_params import manifest_params +from benchmarks.ops.bench_gated_deltanet import ( + _to_fla_layout, + compute_w_u_torch, + kernel2_gated_deltanet_torch, + prepare_wy_repr_gated_torch, +) +from tileops.manifest import load_workloads +from tileops.ops import GatedDeltaNetPrefillFwdOp +from workloads.gated_deltanet import GatedDeltaNetPrefillFwdTest + +_OP_NAME = "GatedDeltaNetPrefillFwdOp" + + +try: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule +except ImportError: + chunk_gated_delta_rule = None + + +class _GatedDeltaNetPrefillFwdTestBaseline(GatedDeltaNetPrefillFwdTest): + """Adds a pure-torch fallback baseline for benchmark profiling.""" + + def ref_program( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch, heads, seq_len, dim_k = k.shape + dim_v = v.shape[-1] + chunk_size = self.chunk_size + num_chunks = seq_len // chunk_size + g_cum = ( + g.float() + .reshape(batch, heads, num_chunks, chunk_size) + .cumsum(-1) + .reshape(batch, heads, seq_len) + .to(g.dtype) + ) + aw, au = prepare_wy_repr_gated_torch(k, g_cum, beta, chunk_size) + w, u = compute_w_u_torch(aw, au, k, v, beta, chunk_size) + initial_state = torch.zeros( + batch, heads, dim_k, dim_v, dtype=torch.float32, device=q.device + ) + final_state, o = kernel2_gated_deltanet_torch( + q, k, g_cum, w, u, initial_state, chunk_size + ) + return o.to(self.dtype), final_state.to(self.dtype) + + +def _fla_prefill_fwd(): + """Return the FLA prefill baseline callable, or None if unavailable.""" + if chunk_gated_delta_rule is None: + return None + + signature = inspect.signature(chunk_gated_delta_rule) + supports_output_final_state = "output_final_state" in signature.parameters + + def baseline_fn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ): + kwargs: dict[str, Any] = {"scale": 1.0} + if supports_output_final_state: + kwargs["output_final_state"] = True + return chunk_gated_delta_rule(q, k, v, g, beta, **kwargs) + + return baseline_fn + + +def _gdn_prefill_args(workload: dict[str, Any]) -> tuple[int, int, int, int, int, int]: + batch, seq_len, heads, dim_k = workload["q_shape"] + _, v_seq_len, v_heads, dim_v = workload["v_shape"] + if v_seq_len != seq_len: + raise ValueError("GDN prefill q_shape and v_shape must share seq_len") + if v_heads != heads: + raise ValueError("GDN prefill q_shape and v_shape must share heads") + return batch, heads, seq_len, dim_k, dim_v, workload.get("chunk_size", 64) + + +_BENCH_PARAMS = manifest_params(load_workloads(_OP_NAME), _gdn_prefill_args, tune=False) + + +@pytest.mark.parametrize( + "batch, heads, seq_len, dim_k, dim_v, chunk_size, dtype, tune", + _BENCH_PARAMS, +) +def test_gated_deltanet_prefill_fwd_bench( + batch: int, + heads: int, + seq_len: int, + dim_k: int, + dim_v: int, + chunk_size: int, + dtype: torch.dtype, + tune: bool, +) -> None: + test = _GatedDeltaNetPrefillFwdTestBaseline( + batch, heads, seq_len, dim_k, dim_v, chunk_size, dtype + ) + inputs = test.gen_inputs() + + op = GatedDeltaNetPrefillFwdOp( + batch, + heads, + seq_len, + dim_k, + dim_v, + chunk_size, + dtype, + layout="bthd", + tune=tune, + ) + bm = ManifestBenchmark(_OP_NAME, op, test) + bthd_inputs = _to_fla_layout(*inputs) + result = bm.profile(op, *bthd_inputs) + BenchmarkReport.record(op, locals(), result, tag="tileops") + + fla_fn = _fla_prefill_fwd() + if fla_fn is not None: + result_fla = bm.profile(fla_fn, *bthd_inputs) + BenchmarkReport.record(op, locals(), result_fla, tag="fla") + else: + result_ref = bm.profile(test.ref_program, *inputs) + BenchmarkReport.record(op, locals(), result_ref, tag="torch-ref") + + +if __name__ == "__main__": + pytest.main([__file__, "-vvs"]) diff --git a/tests/ops/test_gated_deltanet_prefill.py b/tests/ops/test_gated_deltanet_prefill.py new file mode 100644 index 000000000..24bf728e3 --- /dev/null +++ b/tests/ops/test_gated_deltanet_prefill.py @@ -0,0 +1,440 @@ +import pytest +import torch + +from tests.ops.test_gated_deltanet_fwd import ( + compute_w_u_torch, + kernel2_gated_deltanet_torch, + prepare_wy_repr_gated_torch, +) +from tests.test_base import FixtureBase, TestBase +from tileops.ops import GatedDeltaNetPrefillFwdOp +from workloads.gated_deltanet import ( + GatedDeltaNetPrefillFwdTest as _GatedDeltaNetPrefillFwdTestWorkload, +) + + +def dense_chunk_transitions_torch( + k: torch.Tensor, + g: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + chunk_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build dense per-chunk affine transitions for scan-path validation.""" + B, H, S, DK = k.shape + DV = u.shape[-1] + BC = chunk_size + num_chunks = S // BC + + k = k.float().reshape(B, H, num_chunks, BC, DK) + g = g.float().reshape(B, H, num_chunks, BC) + w = w.float().reshape(B, H, num_chunks, BC, DK) + u = u.float().reshape(B, H, num_chunks, BC, DV) + + A = torch.empty(B, H, num_chunks, DK, DK, dtype=torch.float32, device=k.device) + bias = torch.empty(B, H, num_chunks, DK, DV, dtype=torch.float32, device=k.device) + eye = torch.eye(DK, dtype=torch.float32, device=k.device).view(1, 1, DK, DK) + + for c in range(num_chunks): + k_c = k[:, :, c] + g_c = g[:, :, c] + w_c = w[:, :, c] + u_c = u[:, :, c] + g_last = g_c[:, :, -1] + + k_scaled = k_c * torch.exp(g_last.unsqueeze(-1) - g_c).unsqueeze(-1) + w_scaled = w_c * torch.exp(g_c + g_last.unsqueeze(-1)).unsqueeze(-1) + A[:, :, c] = torch.exp(g_last).view(B, H, 1, 1) * eye - torch.einsum( + "bhnk,bhnj->bhkj", k_scaled, w_scaled + ) + bias[:, :, c] = torch.einsum("bhnk,bhnv->bhkv", k_scaled, u_c) + + return A, bias + + +def compose_dense_transitions_torch( + prev_A: torch.Tensor, + prev_bias: torch.Tensor, + next_A: torch.Tensor, + next_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compose ``next`` after ``prev`` for H' = A @ H + B.""" + composed_A = torch.einsum("bhij,bhjk->bhik", next_A, prev_A) + composed_bias = torch.einsum("bhij,bhjv->bhiv", next_A, prev_bias) + next_bias + return composed_A, composed_bias + + +def serial_prefix_dense_transitions_torch( + A: torch.Tensor, + bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference prefix scan over dense chunk transitions.""" + B, H, num_chunks, DK, _ = A.shape + DV = bias.shape[-1] + running_A = torch.eye(DK, dtype=torch.float32, device=A.device).expand(B, H, DK, DK) + running_bias = torch.zeros(B, H, DK, DV, dtype=torch.float32, device=A.device) + prefix_A = torch.empty_like(A) + prefix_bias = torch.empty_like(bias) + + for c in range(num_chunks): + running_A, running_bias = compose_dense_transitions_torch( + running_A, running_bias, A[:, :, c], bias[:, :, c] + ) + prefix_A[:, :, c] = running_A + prefix_bias[:, :, c] = running_bias + + return prefix_A, prefix_bias + + +def butterfly_prefix_dense_transitions_torch( + A: torch.Tensor, + bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Hillis-Steele style inclusive scan over dense chunk transitions.""" + num_chunks = A.shape[2] + prefix_A = A.clone() + prefix_bias = bias.clone() + + offset = 1 + while offset < num_chunks: + prev_A = prefix_A.clone() + prev_bias = prefix_bias.clone() + for c in range(offset, num_chunks): + prefix_A[:, :, c], prefix_bias[:, :, c] = compose_dense_transitions_torch( + prev_A[:, :, c - offset], + prev_bias[:, :, c - offset], + prev_A[:, :, c], + prev_bias[:, :, c], + ) + offset *= 2 + + return prefix_A, prefix_bias + + +def structured_chunk_transitions_torch( + k: torch.Tensor, + g: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + chunk_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build structured per-chunk transitions A = alpha I + U V^T.""" + B, H, S, DK = k.shape + DV = u.shape[-1] + BC = chunk_size + num_chunks = S // BC + + k = k.float().reshape(B, H, num_chunks, BC, DK) + g = g.float().reshape(B, H, num_chunks, BC) + w = w.float().reshape(B, H, num_chunks, BC, DK) + u = u.float().reshape(B, H, num_chunks, BC, DV) + + alpha = torch.empty(B, H, num_chunks, dtype=torch.float32, device=k.device) + U = torch.empty(B, H, num_chunks, DK, BC, dtype=torch.float32, device=k.device) + V = torch.empty(B, H, num_chunks, DK, BC, dtype=torch.float32, device=k.device) + bias = torch.empty(B, H, num_chunks, DK, DV, dtype=torch.float32, device=k.device) + + for c in range(num_chunks): + k_c = k[:, :, c] + g_c = g[:, :, c] + w_c = w[:, :, c] + u_c = u[:, :, c] + g_last = g_c[:, :, -1] + + k_scaled = k_c * torch.exp(g_last.unsqueeze(-1) - g_c).unsqueeze(-1) + w_scaled = w_c * torch.exp(g_c + g_last.unsqueeze(-1)).unsqueeze(-1) + alpha[:, :, c] = torch.exp(g_last) + U[:, :, c] = -k_scaled.transpose(-1, -2) + V[:, :, c] = w_scaled.transpose(-1, -2) + bias[:, :, c] = torch.einsum("bhnk,bhnv->bhkv", k_scaled, u_c) + + return alpha, U, V, bias + + +def materialize_structured_A_torch( + alpha: torch.Tensor, + U: torch.Tensor, + V: torch.Tensor, +) -> torch.Tensor: + """Materialize a structured transition for reference checks.""" + DK = U.shape[-2] + eye = torch.eye(DK, dtype=torch.float32, device=U.device).view(1, 1, DK, DK) + return alpha[..., None, None] * eye + torch.einsum("...kr,...jr->...kj", U, V) + + +def compose_structured_transition_torch( + prev: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + next_: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compose ``next`` after ``prev`` while keeping A = alpha I + U V^T.""" + prev_alpha, prev_U, prev_V, prev_bias = prev + next_alpha, next_U, next_V, next_bias = next_ + + alpha = next_alpha * prev_alpha + cross = torch.einsum("bhkr,bhks->bhrs", prev_U, next_V) + prev_V_part = next_alpha.view(*next_alpha.shape, 1, 1) * prev_V + next_V_part = ( + prev_alpha.view(*prev_alpha.shape, 1, 1) * next_V + + torch.einsum("bhkr,bhrs->bhks", prev_V, cross) + ) + U = torch.cat([prev_U, next_U], dim=-1) + V = torch.cat([prev_V_part, next_V_part], dim=-1) + + next_inner = torch.einsum("bhkr,bhkv->bhrv", next_V, prev_bias) + bias = ( + next_alpha.view(*next_alpha.shape, 1, 1) * prev_bias + + torch.einsum("bhkr,bhrv->bhkv", next_U, next_inner) + + next_bias + ) + return alpha, U, V, bias + + +def butterfly_prefix_structured_transitions_torch( + alpha: torch.Tensor, + U: torch.Tensor, + V: torch.Tensor, + bias: torch.Tensor, +) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + """Hillis-Steele inclusive scan over structured chunk transitions.""" + num_chunks = alpha.shape[2] + prefix = [ + (alpha[:, :, c], U[:, :, c], V[:, :, c], bias[:, :, c]) + for c in range(num_chunks) + ] + + offset = 1 + while offset < num_chunks: + prev_prefix = prefix.copy() + for c in range(offset, num_chunks): + prefix[c] = compose_structured_transition_torch( + prev_prefix[c - offset], prev_prefix[c] + ) + offset *= 2 + + return prefix + + +class GatedDeltaNetPrefillFwdTest(_GatedDeltaNetPrefillFwdTestWorkload, TestBase): + def ref_program( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + B, H, S, DK = k.shape + _, _, _, DV = v.shape + BC = self.chunk_size + g_cum = g.float().reshape(B, H, S // BC, BC).cumsum(-1).reshape(B, H, S).to(g.dtype) + Aw, Au = prepare_wy_repr_gated_torch(k, g_cum, beta, BC) + w, u = compute_w_u_torch(Aw, Au, k, v, beta, BC) + S_0 = torch.zeros(B, H, DK, DV, dtype=torch.float32, device=q.device) + final_state, o = kernel2_gated_deltanet_torch(q, k, g_cum, w, u, S_0, BC) + return o.to(self.dtype), final_state.to(self.dtype) + + +def _get_tolerances(dtype: torch.dtype) -> dict: + if dtype == torch.float32: + return {"atol": 1e-2, "rtol": 1e-2} + if dtype == torch.float16: + return {"atol": 5e-2, "rtol": 5e-2} + return {"atol": 1e-1, "rtol": 1e-1} + + +class GatedDeltaNetPrefillFwdFixture(FixtureBase): + PARAMS = [ + ("batch, seq_len, heads, dim_k, dim_v, chunk_size, dtype, tune", [ + pytest.param(1, 64, 2, 64, 64, 32, torch.float32, False, marks=pytest.mark.smoke), + pytest.param(1, 64, 2, 64, 64, 32, torch.float16, False, marks=pytest.mark.smoke), + pytest.param(1, 64, 2, 64, 64, 32, torch.bfloat16, False, marks=pytest.mark.smoke), + pytest.param(1, 128, 4, 64, 64, 32, torch.float16, False, marks=pytest.mark.full), + pytest.param(1, 128, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full), + ]), + ] + + +@GatedDeltaNetPrefillFwdFixture +def test_gated_deltanet_prefill_fwd( + batch: int, + seq_len: int, + heads: int, + dim_k: int, + dim_v: int, + chunk_size: int, + dtype: torch.dtype, + tune: bool, +) -> None: + torch.manual_seed(42) + test = GatedDeltaNetPrefillFwdTest(batch, heads, seq_len, dim_k, dim_v, chunk_size, dtype) + op = GatedDeltaNetPrefillFwdOp( + batch, + heads, + seq_len, + dim_k, + dim_v, + chunk_size, + dtype, + layout="bhtd", + tune=tune, + ) + test.check(op, *test.gen_inputs(), **_get_tolerances(dtype)) + + +@pytest.mark.smoke +def test_gated_deltanet_prefill_bthd_layout_matches_bhtd() -> None: + torch.manual_seed(42) + B, H, S, DK, DV, BC = 1, 4, 128, 64, 64, 64 + dtype = torch.float16 + + test = GatedDeltaNetPrefillFwdTest(B, H, S, DK, DV, BC, dtype) + q, k, v, g, beta = test.gen_inputs() + + bhtd_op = GatedDeltaNetPrefillFwdOp( + B, + H, + S, + DK, + DV, + BC, + dtype, + layout="bhtd", + tune=False, + ) + bthd_op = GatedDeltaNetPrefillFwdOp( + B, + H, + S, + DK, + DV, + BC, + dtype, + tune=False, + ) + + o_bhtd, state_bhtd = bhtd_op(q, k, v, g, beta) + q_bthd = q.permute(0, 2, 1, 3) + k_bthd = k.permute(0, 2, 1, 3) + v_bthd = v.permute(0, 2, 1, 3) + g_bthd = g.permute(0, 2, 1) + beta_bthd = beta.permute(0, 2, 1) + assert not q_bthd.is_contiguous() + assert not g_bthd.is_contiguous() + o_bthd, state_bthd = bthd_op( + q_bthd, + k_bthd, + v_bthd, + g_bthd, + beta_bthd, + ) + + torch.testing.assert_close( + o_bthd.permute(0, 2, 1, 3).contiguous(), + o_bhtd, + **_get_tolerances(dtype), + ) + torch.testing.assert_close(state_bthd, state_bhtd, **_get_tolerances(dtype)) + + +@pytest.mark.smoke +def test_gated_deltanet_prefill_bthd_blocksolve_path_matches_bhtd() -> None: + torch.manual_seed(42) + B, H, S, DK, DV, BC = 1, 16, 128, 128, 128, 64 + dtype = torch.float16 + + test = GatedDeltaNetPrefillFwdTest(B, H, S, DK, DV, BC, dtype) + q, k, v, g, beta = test.gen_inputs() + + bhtd_op = GatedDeltaNetPrefillFwdOp( + B, + H, + S, + DK, + DV, + BC, + dtype, + layout="bhtd", + tune=False, + ) + bthd_op = GatedDeltaNetPrefillFwdOp( + B, + H, + S, + DK, + DV, + BC, + dtype, + tune=False, + layout="bthd", + ) + + o_bhtd, state_bhtd = bhtd_op(q, k, v, g, beta) + o_bthd, state_bthd = bthd_op( + q.permute(0, 2, 1, 3).contiguous(), + k.permute(0, 2, 1, 3).contiguous(), + v.permute(0, 2, 1, 3).contiguous(), + g.permute(0, 2, 1).contiguous(), + beta.permute(0, 2, 1).contiguous(), + ) + + torch.testing.assert_close( + o_bthd.permute(0, 2, 1, 3).contiguous(), + o_bhtd, + **_get_tolerances(dtype), + ) + torch.testing.assert_close(state_bthd, state_bhtd, **_get_tolerances(dtype)) + + +@pytest.mark.smoke +def test_gated_deltanet_prefill_dense_transition_reference() -> None: + torch.manual_seed(42) + B, H, S, DK, DV, BC = 1, 2, 64, 16, 16, 16 + dtype = torch.float32 + device = torch.device("cuda") + + test = GatedDeltaNetPrefillFwdTest(B, H, S, DK, DV, BC, dtype) + q, k, v, g, beta = test.gen_inputs() + q, k, v, g, beta = ( + q.to(device), + k.to(device), + v.to(device), + g.to(device), + beta.to(device), + ) + + g_cum = g.float().reshape(B, H, S // BC, BC).cumsum(-1).reshape(B, H, S) + Aw, Au = prepare_wy_repr_gated_torch(k, g_cum, beta, BC) + w, u = compute_w_u_torch(Aw, Au, k, v, beta, BC) + + S_0 = torch.randn(B, H, DK, DV, dtype=torch.float32, device=device) * 0.01 + serial_final_state, _ = kernel2_gated_deltanet_torch(q, k, g_cum, w, u, S_0, BC) + + A, bias = dense_chunk_transitions_torch(k, g_cum, w, u, BC) + prefix_A, prefix_bias = serial_prefix_dense_transitions_torch(A, bias) + butterfly_A, butterfly_bias = butterfly_prefix_dense_transitions_torch(A, bias) + structured_alpha, structured_U, structured_V, structured_bias = ( + structured_chunk_transitions_torch(k, g_cum, w, u, BC) + ) + structured_A = materialize_structured_A_torch( + structured_alpha, structured_U, structured_V + ) + structured_prefix = butterfly_prefix_structured_transitions_torch( + structured_alpha, structured_U, structured_V, structured_bias + ) + final_alpha, final_U, final_V, final_bias = structured_prefix[-1] + structured_final_A = materialize_structured_A_torch(final_alpha, final_U, final_V) + scan_final_state = torch.einsum( + "bhij,bhjv->bhiv", prefix_A[:, :, -1], S_0 + ) + prefix_bias[:, :, -1] + + torch.testing.assert_close(structured_A, A, atol=5e-4, rtol=5e-4) + torch.testing.assert_close(structured_bias, bias, atol=5e-4, rtol=5e-4) + torch.testing.assert_close(butterfly_A, prefix_A, atol=5e-4, rtol=5e-4) + torch.testing.assert_close(butterfly_bias, prefix_bias, atol=5e-4, rtol=5e-4) + torch.testing.assert_close(structured_final_A, prefix_A[:, :, -1], atol=5e-4, rtol=5e-4) + torch.testing.assert_close(final_bias, prefix_bias[:, :, -1], atol=5e-4, rtol=5e-4) + torch.testing.assert_close(scan_final_state, serial_final_state, atol=5e-4, rtol=5e-4) + + +if __name__ == "__main__": + pytest.main([__file__, "-vvs"]) diff --git a/tests/perf/test_gated_deltanet_prefill_roofline.py b/tests/perf/test_gated_deltanet_prefill_roofline.py new file mode 100644 index 000000000..dec5454d2 --- /dev/null +++ b/tests/perf/test_gated_deltanet_prefill_roofline.py @@ -0,0 +1,70 @@ +import pytest + +from tileops.perf.formulas import gated_deltanet_prefill_fwd_roofline + +pytestmark = pytest.mark.smoke + + +def _expected_roofline( + batch: int, + heads: int, + seq_len: int, + dim_k: int, + dim_v: int, + chunk_size: int, + elem_bytes: int, +) -> tuple[int, int]: + num_chunks = seq_len // chunk_size + state_flops = 4 * batch * heads * num_chunks * chunk_size * dim_k * dim_v + intra_flops = 4 * batch * heads * num_chunks * chunk_size * chunk_size * ( + dim_k + dim_v + ) + input_elems = ( + 3 * batch * heads * seq_len * dim_k + + batch * heads * seq_len * dim_v + + 2 * batch * heads * seq_len + ) + output_elems = batch * heads * seq_len * dim_v + batch * heads * dim_k * dim_v + return state_flops + intra_flops, (input_elems + output_elems) * elem_bytes + + +def test_gated_deltanet_prefill_roofline_manifest_bthd_layout() -> None: + flops, nbytes = gated_deltanet_prefill_fwd_roofline( + q_shape=[1, 512, 16, 128], + v_shape=[1, 512, 16, 128], + chunk_size=64, + layout="bthd", + dtype="float16", + ) + + assert (flops, nbytes) == _expected_roofline( + batch=1, + heads=16, + seq_len=512, + dim_k=128, + dim_v=128, + chunk_size=64, + elem_bytes=2, + ) + assert flops > 0 + assert flops > 0 + + +def test_gated_deltanet_prefill_roofline_head_major_layout() -> None: + flops, nbytes = gated_deltanet_prefill_fwd_roofline( + q_shape=[1, 16, 512, 128], + v_shape=[1, 16, 512, 128], + chunk_size=64, + layout="bhtd", + dtype="float16", + ) + + assert (flops, nbytes) == _expected_roofline( + batch=1, + heads=16, + seq_len=512, + dim_k=128, + dim_v=128, + chunk_size=64, + elem_bytes=2, + ) diff --git a/tileops/kernels/__init__.py b/tileops/kernels/__init__.py index 3ba6780fc..d931cf55f 100644 --- a/tileops/kernels/__init__.py +++ b/tileops/kernels/__init__.py @@ -50,7 +50,11 @@ from .fft import FFTC2CKernel from .fp8_lighting_indexer import FP8LightingIndexerKernel from .fp8_quant import FP8QuantKernel -from .gated_deltanet import GatedDeltaNetBwdKernel, GatedDeltaNetFwdKernel +from .gated_deltanet import ( + GatedDeltaNetBwdKernel, + GatedDeltaNetFwdKernel, + GatedDeltaNetPrefillFwdKernel, +) from .gated_deltanet_recurrence import ( GatedDeltaNetDecodeFP32Kernel, GatedDeltaNetDecodeKernel, @@ -137,6 +141,7 @@ "GatedDeltaNetDecodeKernel", "GatedDeltaNetDecodeRawCudaFlaStyleKernel", "GatedDeltaNetFwdKernel", + "GatedDeltaNetPrefillFwdKernel", "GemmKernel", "GemvKernel", "GroupConv1dKernel", diff --git a/tileops/kernels/gated_deltanet/__init__.py b/tileops/kernels/gated_deltanet/__init__.py index 0345504e4..ffb12705f 100644 --- a/tileops/kernels/gated_deltanet/__init__.py +++ b/tileops/kernels/gated_deltanet/__init__.py @@ -1,7 +1,9 @@ from .gated_deltanet_bwd import GatedDeltaNetBwdKernel from .gated_deltanet_fwd import GatedDeltaNetFwdKernel +from .gated_deltanet_prefill import GatedDeltaNetPrefillFwdKernel __all__ = [ "GatedDeltaNetBwdKernel", "GatedDeltaNetFwdKernel", + "GatedDeltaNetPrefillFwdKernel", ] diff --git a/tileops/kernels/gated_deltanet/gated_deltanet_fwd.py b/tileops/kernels/gated_deltanet/gated_deltanet_fwd.py index 656b7d031..cdea50948 100644 --- a/tileops/kernels/gated_deltanet/gated_deltanet_fwd.py +++ b/tileops/kernels/gated_deltanet/gated_deltanet_fwd.py @@ -59,7 +59,7 @@ def _h_recurrence_tl( @tilelang.jit( out_idx=[-2, -1], pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, }, compile_flags=["-O3", "-DENABLE_BF16"], ) @@ -81,7 +81,6 @@ def h_recurrence_kernel( u_c = T.alloc_shared([block_C, BV], dtype) h_c = T.alloc_shared([dim_k, BV], dtype) v_new_c = T.alloc_shared([block_C, BV], dtype) - k_scaled_s = T.alloc_shared([block_C, dim_k], dtype) ws_frag = T.alloc_fragment([block_C, BV], accum_dtype) h_next_frag = T.alloc_fragment([dim_k, BV], accum_dtype) @@ -118,13 +117,20 @@ def h_recurrence_kernel( v_offset : v_offset + BV], disable_tma=True) - # h_tile_next = h_tile * exp(g_last) + k_scaled^T @ v_new_tile - for n, kk in T.Parallel(block_C, dim_k): - k_scaled_s[n, kk] = k_c[n, kk] * T.exp2( + # h_tile_next = h_tile * exp(g_last) + k^T @ scaled_v_new_tile + for n, j in T.Parallel(block_C, BV): + v_new_c[n, j] = v_new_c[n, j] * T.exp2( (g_c[block_C - 1] - g_c[n]) * _LOG2E) for i, j in T.Parallel(dim_k, BV): - h_next_frag[i, j] = h_c[i, j] * T.exp2(g_c[block_C - 1] * _LOG2E) - T.gemm(k_scaled_s, v_new_c, h_next_frag, transpose_A=True) + h_next_frag[i, j] = h_c[i, j] * T.exp2( + g_c[block_C - 1] * _LOG2E) + T.gemm( + k_c, + v_new_c, + h_next_frag, + transpose_A=True, + policy=T.GemmWarpPolicy.FullRow, + ) T.copy(h_next_frag, h_c) for i, j in T.Parallel(dim_k, BV): S[bid, hid, t + 1, i, v_offset + j] = h_c[i, j] @@ -160,7 +166,7 @@ def _output_o_tl( @tilelang.jit( out_idx=[-1], pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, }, compile_flags=["-O3", "-DENABLE_BF16"], ) diff --git a/tileops/kernels/gated_deltanet/gated_deltanet_prefill.py b/tileops/kernels/gated_deltanet/gated_deltanet_prefill.py new file mode 100644 index 000000000..fbca8d189 --- /dev/null +++ b/tileops/kernels/gated_deltanet/gated_deltanet_prefill.py @@ -0,0 +1,1770 @@ +"""Gated DeltaNet inference prefill kernel. + +The prefill path shares the chunkwise state/output kernels with training +forward, but uses an inference-only prepare kernel that produces only ``w`` and +``u``. Training-only ``Aw``/``Au`` matrices are kept in shared memory and are +not written to global memory. +""" + +import functools +import math +import os +from typing import Optional, Tuple + +import tilelang +import tilelang.language as T +import torch + +from tileops.kernels.kernel_base import Kernel + +from .gated_deltanet_fwd import ( + _LOG2E, + _chunk_local_cumsum, + _h_recurrence_tl, + _output_o_tl, +) + +__all__ = ["GatedDeltaNetPrefillFwdKernel"] + + +def _normalize_prefill_layout(layout: str) -> str: + layout = layout.lower() + if layout == "bhsd": + return "bhtd" + if layout in ("bhtd", "bthd"): + return layout + raise ValueError(f"Unsupported layout: {layout}") + + +@functools.lru_cache(maxsize=32) +def _prefill_single_batch_cu_seqlens( + num_tokens: int, + device_idx: int, +) -> torch.Tensor: + return torch.tensor([0, num_tokens], dtype=torch.int32, device=f"cuda:{device_idx}") + + +def _prefill_auto_cp_local_chunks(num_chunks: int, num_heads: int) -> int: + env_max_local_chunks = os.environ.get( + "TILEOPS_GDN_PREFILL_MAX_LOCAL_CHUNKS", + os.environ.get("TILEOPS_GDN_PREFILL_CP_MAX_LOCAL_CHUNKS"), + ) + if env_max_local_chunks: + max_local_chunks = int(env_max_local_chunks) + else: + sm_count = torch.cuda.get_device_properties().multi_processor_count + max_local_chunks = 2 ** round( + math.log2(math.sqrt(num_heads * num_chunks / sm_count) * 3) + ) + if num_heads >= 64 and num_chunks >= 512: + max_local_chunks = max(max_local_chunks, 256) + return max(max_local_chunks, 4) + + +def _prefill_partitioned_initial_state_bthd( + k: torch.Tensor, + v: torch.Tensor, + A: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int, +) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + from tileops.kernels.gated_deltanet.gdn_prefill import ( + correct_initial_states, + fused_gdr_h, + get_warmup_chunks, + ) + + batch, num_tokens, num_heads, _ = k.shape + assert batch == 1 + raw_cu_seqlens = _prefill_single_batch_cu_seqlens(num_tokens, k.device.index) + num_chunks = tilelang.cdiv(num_tokens, chunk_size) + max_local_chunks = _prefill_auto_cp_local_chunks(num_chunks, num_heads) + + has_partition = num_chunks > max_local_chunks + force_partition = ( + os.environ.get( + "TILEOPS_GDN_PREFILL_FORCE_PARTITION", + os.environ.get("TILEOPS_GDN_PREFILL_CP_FORCE", "0"), + ) + == "1" + ) + use_partition = has_partition and ( + force_partition or num_heads <= 40 or (num_heads <= 64 and num_chunks >= 128) + ) + if not use_partition: + return None, raw_cu_seqlens, None, None + + cp_cu_seqlens = [] + ht_mask = [] + seq_map_c2r = [] + seq_map_r2c = [0] + split_tokens = max_local_chunks * chunk_size + start = 0 + while start < num_tokens: + cp_cu_seqlens.append(start) + ht_mask.append(False) + seq_map_c2r.append(0) + start += split_tokens + ht_mask[-1] = True + seq_map_r2c.append(len(cp_cu_seqlens)) + cp_cu_seqlens.append(num_tokens) + + cp_cu_seqlens_t = torch.tensor( + cp_cu_seqlens, dtype=torch.int32, device=k.device, requires_grad=False + ) + seq_map_c2r_t = torch.tensor(seq_map_c2r, dtype=torch.int32, device=k.device) + seq_map_r2c_t = torch.tensor( + seq_map_r2c, dtype=torch.int32, device=k.device, requires_grad=False + ) + ht_mask_t = torch.tensor( + ht_mask, dtype=torch.bool, device=k.device, requires_grad=False + ) + + num_warmup_chunks, fallback_mask = get_warmup_chunks( + g=g, + cu_seqlens=cp_cu_seqlens_t, + ht_mask=ht_mask_t, + chunk_size=chunk_size, + threshold=-10.0, + ) + _, ht, mt = fused_gdr_h( + k=k, + v=v, + a=A, + g=g, + b=beta, + initial_state=None, + output_final_state=True, + output_h=False, + cu_seqlens=cp_cu_seqlens_t, + num_warmup_chunks=num_warmup_chunks, + ) + cp_h0 = correct_initial_states( + raw_h0=None, + ht_buffer=ht, + mt_buffer=mt, + fallback_mask=fallback_mask, + seq_map_r2c=seq_map_r2c_t, + ) + return cp_h0, cp_cu_seqlens_t, seq_map_c2r_t, raw_cu_seqlens + + +@functools.lru_cache(maxsize=32) +def _prefill_chunk_local_cumsum_bhtd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dtype: str, +): + num_chunks = seq_len // chunk_size + + @tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(): + @T.prim_func + def prefill_chunk_cumsum_bhtd( + g: T.Tensor([batch, head, seq_len], dtype), + out: T.Tensor([batch, head, seq_len], dtype), + ): + with T.Kernel(batch * head * num_chunks, threads=128) as bx: + bid = bx // (head * num_chunks) + hid = (bx // num_chunks) % head + cid = bx % num_chunks + base = cid * chunk_size + + g_s = T.alloc_shared([chunk_size], dtype) + out_s = T.alloc_shared([chunk_size], "float32") + + T.copy(g[bid, hid, base : base + chunk_size], g_s, disable_tma=True) + + out_s[0] = T.cast(g_s[0], "float32") + for i in T.Serial(1, chunk_size): + out_s[i] = out_s[i - 1] + T.cast(g_s[i], "float32") + for i in T.Parallel(chunk_size): + out[bid, hid, base + i] = T.cast(out_s[i], dtype) + + return prefill_chunk_cumsum_bhtd + + return _func() + + +@functools.lru_cache(maxsize=32) +def _prefill_prepare_w_u_bhtd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str = "float32", +): + """Compute Gated DeltaNet prefill ``w`` and ``u`` without global Aw/Au.""" + accum_dtype = "float32" + block_C = chunk_size + num_rounds = int(math.ceil(math.log2(chunk_size))) if chunk_size > 1 else 0 + + @tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _fused_func(num_stages, threads=128): + @T.prim_func + def prefill_prepare_w_u_bhtd( + k: T.Tensor([batch, head, seq_len, dim_k], dtype), + v: T.Tensor([batch, head, seq_len, dim_v], dtype), + g: T.Tensor([batch, head, seq_len], dtype), + beta: T.Tensor([batch, head, seq_len], dtype), + w: T.Tensor([batch, head, seq_len, dim_k], dtype), + u: T.Tensor([batch, head, seq_len, dim_v], dtype), + ): + with T.Kernel(batch, head, seq_len // block_C, threads=threads) as (bid, hid, by): + k_shared = T.alloc_shared([block_C, dim_k], dtype) + g_shared = T.alloc_shared([block_C], dtype) + beta_shared = T.alloc_shared([block_C], dtype) + k_beta_shared = T.alloc_shared([block_C, dim_k], dtype) + v_beta_shared = T.alloc_shared([block_C, dim_v], dtype) + S_shared = T.alloc_shared([block_C, block_C], dtype) + P_shared = T.alloc_shared([block_C, block_C], dtype) + g_pos_shared = T.alloc_shared([block_C], accum_dtype) + g_neg_shared = T.alloc_shared([block_C], accum_dtype) + + gram_frag = T.alloc_fragment([block_C, block_C], accum_dtype) + temp_frag = T.alloc_fragment([block_C, block_C], accum_dtype) + w_frag = T.alloc_fragment([block_C, dim_k], accum_dtype) + u_frag = T.alloc_fragment([block_C, dim_v], accum_dtype) + + T.copy( + k[bid, hid, by * block_C : (by + 1) * block_C, :], + k_shared, + disable_tma=True, + ) + T.copy( + g[bid, hid, by * block_C : (by + 1) * block_C], + g_shared, + disable_tma=True, + ) + T.copy( + beta[bid, hid, by * block_C : (by + 1) * block_C], + beta_shared, + disable_tma=True, + ) + + for i_s, i_k in T.Parallel(block_C, dim_k): + k_beta_shared[i_s, i_k] = k_shared[i_s, i_k] * beta_shared[i_s] + for i in T.Parallel(block_C): + g_pos_shared[i] = T.exp2(g_shared[i] * _LOG2E) + g_neg_shared[i] = T.exp2(-g_shared[i] * _LOG2E) + T.clear(gram_frag) + T.gemm(k_beta_shared, k_shared, gram_frag, transpose_B=True) + + for i, j in T.Parallel(block_C, block_C): + P_shared[i, j] = T.if_then_else( + i > j, + -gram_frag[i, j] + * g_pos_shared[i] + * g_neg_shared[j], + T.float32(0.0), + ) + T.clear(S_shared) + for i in T.Parallel(block_C): + S_shared[i, i] = T.float32(1.0) + + for _r in T.Serial(num_rounds): + T.clear(temp_frag) + T.gemm(P_shared, S_shared, temp_frag) + for i, j in T.Parallel(block_C, block_C): + S_shared[i, j] = S_shared[i, j] + temp_frag[i, j] + T.clear(temp_frag) + T.gemm(P_shared, P_shared, temp_frag) + T.copy(temp_frag, P_shared) + + T.clear(w_frag) + T.gemm(S_shared, k_beta_shared, w_frag) + T.copy( + w_frag, + w[bid, hid, by * block_C : (by + 1) * block_C, :], + disable_tma=True, + ) + + for i, j in T.Parallel(block_C, dim_v): + v_beta_shared[i, j] = ( + v[bid, hid, by * block_C + i, j] * beta_shared[i] + ) + T.clear(u_frag) + T.gemm(S_shared, v_beta_shared, u_frag) + T.copy( + u_frag, + u[bid, hid, by * block_C : (by + 1) * block_C, :], + disable_tma=True, + ) + + return prefill_prepare_w_u_bhtd + + return _fused_func + + +@functools.lru_cache(maxsize=32) +def _prefill_chunk_local_cumsum_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dtype: str, +): + num_chunks = seq_len // chunk_size + + @tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(): + @T.prim_func + def chunk_cumsum_bthd_kernel( + g: T.Tensor([batch, seq_len, head], dtype), + out: T.Tensor([batch, seq_len, head], dtype), + ): + with T.Kernel(batch, num_chunks, threads=128) as (bid, cid): + base = cid * chunk_size + acc_s = T.alloc_shared([head], "float32") + + for hid in T.Parallel(head): + acc_s[hid] = T.float32(0.0) + for i in T.Serial(chunk_size): + for hid in T.Parallel(head): + acc_s[hid] = acc_s[hid] + T.cast( + g[bid, base + i, hid], "float32" + ) + out[bid, base + i, hid] = T.cast(acc_s[hid], dtype) + + return chunk_cumsum_bthd_kernel + + return _func() + + +@functools.lru_cache(maxsize=32) +def _prefill_blocksolve_A_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dtype: str, +): + if chunk_size != 64 or dim_k != 128: + raise ValueError("TileLang blocksolve-A currently expects chunk64 and K=128") + + block_t = 64 + block_c = 16 + block_k = 64 + accum_dtype = "float32" + solve_dtype = dtype + + @tilelang.jit( + out_idx=[], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(threads=32): + @T.prim_func + def prefill_blocksolve_A_bthd_tl( + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + g: T.Tensor([batch, seq_len, head], dtype), + beta: T.Tensor([batch, seq_len, head], dtype), + A: T.Tensor([batch, seq_len, head, chunk_size], dtype), + ): + with T.Kernel(batch, head, seq_len // block_t, threads=threads) as ( + bid, + hid, + cid, + ): + base = cid * block_t + k0 = T.alloc_shared([block_c, block_k], dtype) + k1 = T.alloc_shared([block_c, block_k], dtype) + k2 = T.alloc_shared([block_c, block_k], dtype) + k3 = T.alloc_shared([block_c, block_k], dtype) + g_s = T.alloc_shared([block_t], dtype) + beta_s = T.alloc_shared([block_t], dtype) + gate0_s = T.alloc_shared([block_t], dtype) + gate1_s = T.alloc_shared([block_t], dtype) + gate2_s = T.alloc_shared([block_t], dtype) + gate3_s = T.alloc_shared([block_t], dtype) + beta_f_s = T.alloc_shared([block_t], dtype) + a_s = T.alloc_shared([10, block_c, block_c], solve_dtype) + i_s = T.alloc_shared([4, block_c, block_c], solve_dtype) + work_s = T.alloc_shared([1, block_c, block_c], solve_dtype) + + G00 = T.alloc_fragment([block_c, block_c], accum_dtype) + G10 = T.alloc_fragment([block_c, block_c], accum_dtype) + G11 = T.alloc_fragment([block_c, block_c], accum_dtype) + G20 = T.alloc_fragment([block_c, block_c], accum_dtype) + G21 = T.alloc_fragment([block_c, block_c], accum_dtype) + G22 = T.alloc_fragment([block_c, block_c], accum_dtype) + G30 = T.alloc_fragment([block_c, block_c], accum_dtype) + G31 = T.alloc_fragment([block_c, block_c], accum_dtype) + G32 = T.alloc_fragment([block_c, block_c], accum_dtype) + G33 = T.alloc_fragment([block_c, block_c], accum_dtype) + tmp = T.alloc_fragment([block_c, block_c], accum_dtype) + + T.annotate_layout({ + k0: tilelang.layout.make_swizzled_layout(k0), + k1: tilelang.layout.make_swizzled_layout(k1), + k2: tilelang.layout.make_swizzled_layout(k2), + k3: tilelang.layout.make_swizzled_layout(k3), + a_s: tilelang.layout.make_swizzled_layout(a_s), + i_s: tilelang.layout.make_swizzled_layout(i_s), + work_s: tilelang.layout.make_swizzled_layout(work_s), + }) + + T.copy(g[bid, base : base + block_t, hid], g_s, disable_tma=True) + T.copy( + beta[bid, base : base + block_t, hid], + beta_s, + disable_tma=True, + ) + + T.clear(G00) + T.clear(G10) + T.clear(G11) + T.clear(G20) + T.clear(G21) + T.clear(G22) + T.clear(G30) + T.clear(G31) + T.clear(G32) + T.clear(G33) + + for kt in T.Serial(dim_k // block_k): + koff = kt * block_k + T.async_copy( + k[bid, base : base + block_c, hid, koff : koff + block_k], + k0, + ) + T.async_copy( + k[ + bid, + base + block_c : base + 2 * block_c, + hid, + koff : koff + block_k, + ], + k1, + ) + T.async_copy( + k[ + bid, + base + 2 * block_c : base + 3 * block_c, + hid, + koff : koff + block_k, + ], + k2, + ) + T.async_copy( + k[ + bid, + base + 3 * block_c : base + 4 * block_c, + hid, + koff : koff + block_k, + ], + k3, + ) + T.ptx_wait_group(0) + T.sync_threads() + + T.gemm(k0, k0, G00, transpose_B=True) + T.gemm(k1, k0, G10, transpose_B=True) + T.gemm(k1, k1, G11, transpose_B=True) + T.gemm(k2, k0, G20, transpose_B=True) + T.gemm(k2, k1, G21, transpose_B=True) + T.gemm(k2, k2, G22, transpose_B=True) + T.gemm(k3, k0, G30, transpose_B=True) + T.gemm(k3, k1, G31, transpose_B=True) + T.gemm(k3, k2, G32, transpose_B=True) + T.gemm(k3, k3, G33, transpose_B=True) + + for t in T.Parallel(block_t): + g_val = T.cast(g_s[t], accum_dtype) + gate0_s[t] = T.exp2( + (g_val - T.cast(g_s[0], accum_dtype)) * _LOG2E + ) + gate1_s[t] = T.exp2( + (g_val - T.cast(g_s[block_c], accum_dtype)) * _LOG2E + ) + gate2_s[t] = T.exp2( + (g_val - T.cast(g_s[2 * block_c], accum_dtype)) * _LOG2E + ) + gate3_s[t] = T.exp2( + (g_val - T.cast(g_s[3 * block_c], accum_dtype)) * _LOG2E + ) + beta_f_s[t] = beta_s[t] + T.sync_threads() + + for i, j in T.Parallel(block_c, block_c): + s00 = ( + T.cast(beta_f_s[i], accum_dtype) + * T.cast(gate0_s[i], accum_dtype) + / T.cast(gate0_s[j], accum_dtype) + ) + s11 = ( + T.cast(beta_f_s[block_c + i], accum_dtype) + * T.cast(gate1_s[block_c + i], accum_dtype) + / T.cast(gate1_s[block_c + j], accum_dtype) + ) + s22 = ( + T.cast(beta_f_s[2 * block_c + i], accum_dtype) + * T.cast(gate2_s[2 * block_c + i], accum_dtype) + / T.cast(gate2_s[2 * block_c + j], accum_dtype) + ) + s33 = ( + T.cast(beta_f_s[3 * block_c + i], accum_dtype) + * T.cast(gate3_s[3 * block_c + i], accum_dtype) + / T.cast(gate3_s[3 * block_c + j], accum_dtype) + ) + s10 = ( + T.cast(beta_f_s[block_c + i], accum_dtype) + * T.cast(gate0_s[block_c + i], accum_dtype) + / T.cast(gate0_s[j], accum_dtype) + ) + s20 = ( + T.cast(beta_f_s[2 * block_c + i], accum_dtype) + * T.cast(gate0_s[2 * block_c + i], accum_dtype) + / T.cast(gate0_s[j], accum_dtype) + ) + s21 = ( + T.cast(beta_f_s[2 * block_c + i], accum_dtype) + * T.cast(gate1_s[2 * block_c + i], accum_dtype) + / T.cast(gate1_s[block_c + j], accum_dtype) + ) + s30 = ( + T.cast(beta_f_s[3 * block_c + i], accum_dtype) + * T.cast(gate0_s[3 * block_c + i], accum_dtype) + / T.cast(gate0_s[j], accum_dtype) + ) + s31 = ( + T.cast(beta_f_s[3 * block_c + i], accum_dtype) + * T.cast(gate1_s[3 * block_c + i], accum_dtype) + / T.cast(gate1_s[block_c + j], accum_dtype) + ) + s32 = ( + T.cast(beta_f_s[3 * block_c + i], accum_dtype) + * T.cast(gate2_s[3 * block_c + i], accum_dtype) + / T.cast(gate2_s[2 * block_c + j], accum_dtype) + ) + a_s[0, i, j] = T.if_then_else( + i > j, + -G00[i, j] * s00, + T.float32(0.0), + ) + a_s[2, i, j] = T.if_then_else( + i > j, + -G11[i, j] * s11, + T.float32(0.0), + ) + a_s[5, i, j] = T.if_then_else( + i > j, + -G22[i, j] * s22, + T.float32(0.0), + ) + a_s[9, i, j] = T.if_then_else( + i > j, + -G33[i, j] * s33, + T.float32(0.0), + ) + a_s[1, i, j] = G10[i, j] * s10 + a_s[3, i, j] = G20[i, j] * s20 + a_s[4, i, j] = G21[i, j] * s21 + a_s[6, i, j] = G30[i, j] * s30 + a_s[7, i, j] = G31[i, j] * s31 + a_s[8, i, j] = G32[i, j] * s32 + i_s[0, i, j] = T.if_then_else(i == j, T.float32(1.0), T.float32(0.0)) + i_s[1, i, j] = T.if_then_else(i == j, T.float32(1.0), T.float32(0.0)) + i_s[2, i, j] = T.if_then_else(i == j, T.float32(1.0), T.float32(0.0)) + i_s[3, i, j] = T.if_then_else(i == j, T.float32(1.0), T.float32(0.0)) + T.sync_threads() + + for _r in T.Serial(1): + T.clear(tmp) + T.gemm(a_s[0, :, :], i_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + i_s[0, i, j] = i_s[0, i, j] + tmp[i, j] + T.clear(tmp) + T.gemm(a_s[0, :, :], a_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[0, i, j] = tmp[i, j] + + T.clear(tmp) + T.gemm(a_s[2, :, :], i_s[1, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + i_s[1, i, j] = i_s[1, i, j] + tmp[i, j] + T.clear(tmp) + T.gemm(a_s[2, :, :], a_s[2, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[2, i, j] = tmp[i, j] + + T.clear(tmp) + T.gemm(a_s[5, :, :], i_s[2, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + i_s[2, i, j] = i_s[2, i, j] + tmp[i, j] + T.clear(tmp) + T.gemm(a_s[5, :, :], a_s[5, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[5, i, j] = tmp[i, j] + + T.clear(tmp) + T.gemm(a_s[9, :, :], i_s[3, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + i_s[3, i, j] = i_s[3, i, j] + tmp[i, j] + T.clear(tmp) + T.gemm(a_s[9, :, :], a_s[9, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[9, i, j] = tmp[i, j] + T.sync_threads() + + T.clear(tmp) + T.gemm(i_s[1, :, :], a_s[1, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = tmp[i, j] + T.sync_threads() + T.clear(tmp) + T.gemm(work_s[0, :, :], i_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[1, i, j] = -tmp[i, j] + + T.clear(tmp) + T.gemm(i_s[2, :, :], a_s[4, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = tmp[i, j] + T.sync_threads() + T.clear(tmp) + T.gemm(work_s[0, :, :], i_s[1, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[4, i, j] = -tmp[i, j] + + T.clear(tmp) + T.gemm(a_s[3, :, :], i_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = tmp[i, j] + T.clear(tmp) + T.gemm(a_s[4, :, :], a_s[1, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = work_s[0, i, j] + tmp[i, j] + T.sync_threads() + T.clear(tmp) + T.gemm(i_s[2, :, :], work_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[3, i, j] = -tmp[i, j] + + T.clear(tmp) + T.gemm(a_s[6, :, :], i_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = tmp[i, j] + T.clear(tmp) + T.gemm(a_s[7, :, :], a_s[1, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = work_s[0, i, j] + tmp[i, j] + T.clear(tmp) + T.gemm(a_s[8, :, :], a_s[3, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = work_s[0, i, j] + tmp[i, j] + T.sync_threads() + T.clear(tmp) + T.gemm(i_s[3, :, :], work_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[6, i, j] = -tmp[i, j] + + T.clear(tmp) + T.gemm(a_s[7, :, :], i_s[1, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = tmp[i, j] + T.clear(tmp) + T.gemm(a_s[8, :, :], a_s[4, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = work_s[0, i, j] + tmp[i, j] + T.sync_threads() + T.clear(tmp) + T.gemm(i_s[3, :, :], work_s[0, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[7, i, j] = -tmp[i, j] + + T.clear(tmp) + T.gemm(i_s[3, :, :], a_s[8, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + work_s[0, i, j] = tmp[i, j] + T.sync_threads() + T.clear(tmp) + T.gemm(work_s[0, :, :], i_s[2, :, :], tmp) + for i, j in T.Parallel(block_c, block_c): + a_s[8, i, j] = -tmp[i, j] + T.sync_threads() + + for i, j in T.Parallel(block_c, block_c): + A[bid, base + i, hid, j] = T.cast(i_s[0, i, j], dtype) + A[bid, base + block_c + i, hid, j] = T.cast(a_s[1, i, j], dtype) + A[bid, base + block_c + i, hid, block_c + j] = T.cast(i_s[1, i, j], dtype) + A[bid, base + 2 * block_c + i, hid, j] = T.cast(a_s[3, i, j], dtype) + A[bid, base + 2 * block_c + i, hid, block_c + j] = T.cast(a_s[4, i, j], dtype) + A[bid, base + 2 * block_c + i, hid, 2 * block_c + j] = T.cast(i_s[2, i, j], dtype) + A[bid, base + 3 * block_c + i, hid, j] = T.cast(a_s[6, i, j], dtype) + A[bid, base + 3 * block_c + i, hid, block_c + j] = T.cast(a_s[7, i, j], dtype) + A[bid, base + 3 * block_c + i, hid, 2 * block_c + j] = T.cast(a_s[8, i, j], dtype) + A[bid, base + 3 * block_c + i, hid, 3 * block_c + j] = T.cast(i_s[3, i, j], dtype) + + return prefill_blocksolve_A_bthd_tl + + return _func(32) + + +def _prefill_blocksolve_A_bthd( + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int, +) -> torch.Tensor: + batch, seq_len, head, dim_k = k.shape + A = torch.zeros(batch, seq_len, head, chunk_size, dtype=k.dtype, device=k.device) + kernel = _prefill_blocksolve_A_bthd_tl( + batch, + head, + seq_len, + chunk_size, + dim_k, + str(k.dtype).split(".")[-1], + ) + kernel(k, g, beta, A) + return A + +@functools.lru_cache(maxsize=32) +def _prefill_recompute_w_u_from_A_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str, +): + block_c = chunk_size + block_d = 64 + num_k_tiles = dim_k // block_d + num_v_tiles = dim_v // block_d + accum_dtype = "float32" + + @tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(threads=128): + @T.prim_func + def recompute_w_u_from_A_bthd( + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + v: T.Tensor([batch, seq_len, head, dim_v], dtype), + beta: T.Tensor([batch, seq_len, head], dtype), + A: T.Tensor([batch, seq_len, head, chunk_size], dtype), + w: T.Tensor([batch, seq_len, head, dim_k], dtype), + u: T.Tensor([batch, seq_len, head, dim_v], dtype), + ): + with T.Kernel(batch, head, seq_len // block_c, threads=threads) as ( + bid, + hid, + by, + ): + base = by * block_c + A_s = T.alloc_shared([block_c, block_c], dtype) + beta_s = T.alloc_shared([block_c], dtype) + x_s = T.alloc_shared([block_c, block_d], dtype) + out_s = T.alloc_shared([block_c, block_d], dtype) + out_frag = T.alloc_fragment([block_c, block_d], accum_dtype) + + T.annotate_layout({ + A_s: tilelang.layout.make_swizzled_layout(A_s), + x_s: tilelang.layout.make_swizzled_layout(x_s), + out_s: tilelang.layout.make_swizzled_layout(out_s), + }) + + T.async_copy(A[bid, base : base + block_c, hid, 0:block_c], A_s) + T.ptx_wait_group(0) + T.copy( + beta[bid, base : base + block_c, hid], + beta_s, + disable_tma=True, + ) + T.sync_threads() + + for vt in T.Serial(num_v_tiles): + v_offset = vt * block_d + T.async_copy( + v[ + bid, + base : base + block_c, + hid, + v_offset : v_offset + block_d, + ], + x_s, + ) + T.ptx_wait_group(0) + T.sync_threads() + for i, d in T.Parallel(block_c, block_d): + x_s[i, d] = x_s[i, d] * beta_s[i] + T.clear(out_frag) + T.gemm(A_s, x_s, out_frag) + T.copy(out_frag, out_s) + T.copy( + out_s, + u[ + bid, + base : base + block_c, + hid, + v_offset : v_offset + block_d, + ], + disable_tma=True, + ) + + for kt in T.Serial(num_k_tiles): + k_offset = kt * block_d + T.async_copy( + k[ + bid, + base : base + block_c, + hid, + k_offset : k_offset + block_d, + ], + x_s, + ) + T.ptx_wait_group(0) + T.sync_threads() + for i, d in T.Parallel(block_c, block_d): + x_s[i, d] = x_s[i, d] * beta_s[i] + T.clear(out_frag) + T.gemm(A_s, x_s, out_frag) + T.copy(out_frag, out_s) + T.copy( + out_s, + w[ + bid, + base : base + block_c, + hid, + k_offset : k_offset + block_d, + ], + disable_tma=True, + ) + + return recompute_w_u_from_A_bthd + + return _func + + +@functools.lru_cache(maxsize=32) +def _prefill_prepare_w_u_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str = "float32", +): + accum_dtype = "float32" + block_C = chunk_size + num_rounds = int(math.ceil(math.log2(chunk_size))) if chunk_size > 1 else 0 + + @tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _fused_func(num_stages, threads=128): + @T.prim_func + def prefill_prepare_w_u_bthd( + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + v: T.Tensor([batch, seq_len, head, dim_v], dtype), + g: T.Tensor([batch, seq_len, head], dtype), + beta: T.Tensor([batch, seq_len, head], dtype), + w: T.Tensor([batch, seq_len, head, dim_k], dtype), + u: T.Tensor([batch, seq_len, head, dim_v], dtype), + ): + with T.Kernel(batch, head, seq_len // block_C, threads=threads) as ( + bid, + hid, + by, + ): + base = by * block_C + k_shared = T.alloc_shared([block_C, dim_k], dtype) + g_shared = T.alloc_shared([block_C], dtype) + beta_shared = T.alloc_shared([block_C], dtype) + k_beta_shared = T.alloc_shared([block_C, dim_k], dtype) + v_beta_shared = T.alloc_shared([block_C, dim_v], dtype) + S_shared = T.alloc_shared([block_C, block_C], dtype) + P_shared = T.alloc_shared([block_C, block_C], dtype) + g_pos_shared = T.alloc_shared([block_C], accum_dtype) + g_neg_shared = T.alloc_shared([block_C], accum_dtype) + + gram_frag = T.alloc_fragment([block_C, block_C], accum_dtype) + temp_frag = T.alloc_fragment([block_C, block_C], accum_dtype) + w_frag = T.alloc_fragment([block_C, dim_k], accum_dtype) + u_frag = T.alloc_fragment([block_C, dim_v], accum_dtype) + + for i, d in T.Parallel(block_C, dim_k): + k_shared[i, d] = k[bid, base + i, hid, d] + for i in T.Parallel(block_C): + g_shared[i] = g[bid, base + i, hid] + beta_shared[i] = beta[bid, base + i, hid] + + for i, d in T.Parallel(block_C, dim_k): + k_beta_shared[i, d] = k_shared[i, d] * beta_shared[i] + for i in T.Parallel(block_C): + g_pos_shared[i] = T.exp2(g_shared[i] * _LOG2E) + g_neg_shared[i] = T.exp2(-g_shared[i] * _LOG2E) + T.clear(gram_frag) + T.gemm(k_beta_shared, k_shared, gram_frag, transpose_B=True) + + for i, j in T.Parallel(block_C, block_C): + P_shared[i, j] = T.if_then_else( + i > j, + -gram_frag[i, j] + * g_pos_shared[i] + * g_neg_shared[j], + T.float32(0.0), + ) + T.clear(S_shared) + for i in T.Parallel(block_C): + S_shared[i, i] = T.float32(1.0) + + for _r in T.Serial(num_rounds): + T.clear(temp_frag) + T.gemm(P_shared, S_shared, temp_frag) + for i, j in T.Parallel(block_C, block_C): + S_shared[i, j] = S_shared[i, j] + temp_frag[i, j] + T.clear(temp_frag) + T.gemm(P_shared, P_shared, temp_frag) + T.copy(temp_frag, P_shared) + + T.clear(w_frag) + T.gemm(S_shared, k_beta_shared, w_frag) + for i, d in T.Parallel(block_C, dim_k): + w[bid, base + i, hid, d] = w_frag[i, d] + + for i, d in T.Parallel(block_C, dim_v): + v_beta_shared[i, d] = ( + v[bid, base + i, hid, d] * beta_shared[i] + ) + T.clear(u_frag) + T.gemm(S_shared, v_beta_shared, u_frag) + for i, d in T.Parallel(block_C, dim_v): + u[bid, base + i, hid, d] = u_frag[i, d] + + return prefill_prepare_w_u_bthd + + return _fused_func + + +@functools.lru_cache(maxsize=32) +def _prefill_h_recurrence_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str = "float32", + block_v: int = 0, +): + accum_dtype = "float32" + block_C = chunk_size + num_chunks = seq_len // block_C + BV = dim_v if block_v <= 0 else block_v + num_v_tiles = dim_v // BV + + @tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(num_stages, threads=128): + @T.prim_func + def h_recurrence_bthd( + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + g: T.Tensor([batch, seq_len, head], dtype), + w: T.Tensor([batch, seq_len, head, dim_k], dtype), + u: T.Tensor([batch, seq_len, head, dim_v], dtype), + S_0: T.Tensor([batch, head, dim_k, dim_v], dtype), + S: T.Tensor([batch, head, num_chunks + 1, dim_k, dim_v], dtype), + v_new: T.Tensor([batch, seq_len, head, dim_v], dtype), + ): + with T.Kernel(num_v_tiles, batch, head, threads=threads) as ( + vid, + bid, + hid, + ): + k_c = T.alloc_shared([block_C, dim_k], dtype) + g_c = T.alloc_shared([block_C], dtype) + w_c = T.alloc_shared([block_C, dim_k], dtype) + u_c = T.alloc_shared([block_C, BV], dtype) + h_c = T.alloc_shared([dim_k, BV], dtype) + v_new_c = T.alloc_shared([block_C, BV], dtype) + + ws_frag = T.alloc_fragment([block_C, BV], accum_dtype) + h_next_frag = T.alloc_fragment([dim_k, BV], accum_dtype) + + v_offset = vid * BV + + T.copy(S_0[bid, hid, :, v_offset : v_offset + BV], h_c, disable_tma=True) + for i, j in T.Parallel(dim_k, BV): + S[bid, hid, 0, i, v_offset + j] = h_c[i, j] + + for t in T.Pipelined(num_chunks, num_stages=num_stages): + base = t * block_C + for i, d in T.Parallel(block_C, dim_k): + k_c[i, d] = k[bid, base + i, hid, d] + w_c[i, d] = w[bid, base + i, hid, d] + for i in T.Parallel(block_C): + g_c[i] = g[bid, base + i, hid] + for i, d in T.Parallel(block_C, BV): + u_c[i, d] = u[bid, base + i, hid, v_offset + d] + + T.clear(ws_frag) + T.gemm(w_c, h_c, ws_frag) + for i, j in T.Parallel(block_C, BV): + v_new_c[i, j] = u_c[i, j] - ws_frag[i, j] * T.exp2( + (g_c[i] + g_c[block_C - 1]) * _LOG2E + ) + + for i, j in T.Parallel(block_C, BV): + v_new[bid, base + i, hid, v_offset + j] = v_new_c[i, j] + + for n, j in T.Parallel(block_C, BV): + v_new_c[n, j] = v_new_c[n, j] * T.exp2( + (g_c[block_C - 1] - g_c[n]) * _LOG2E + ) + for i, j in T.Parallel(dim_k, BV): + h_next_frag[i, j] = h_c[i, j] * T.exp2( + g_c[block_C - 1] * _LOG2E + ) + T.gemm( + k_c, + v_new_c, + h_next_frag, + transpose_A=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(h_next_frag, h_c) + for i, j in T.Parallel(dim_k, BV): + S[bid, hid, t + 1, i, v_offset + j] = h_c[i, j] + + return h_recurrence_bthd + + return _func + + +@functools.lru_cache(maxsize=32) +def _prefill_output_o_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str = "float32", +): + accum_dtype = "float32" + block_C = chunk_size + num_chunks = seq_len // block_C + + @tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(threads=128): + @T.prim_func + def output_o_bthd( + q: T.Tensor([batch, seq_len, head, dim_k], dtype), + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + g: T.Tensor([batch, seq_len, head], dtype), + S: T.Tensor([batch, head, num_chunks + 1, dim_k, dim_v], dtype), + v_new: T.Tensor([batch, seq_len, head, dim_v], dtype), + o: T.Tensor([batch, seq_len, head, dim_v], dtype), + ): + with T.Kernel(num_chunks, batch, head, threads=threads) as (tid, bid, hid): + base = tid * block_C + q_c = T.alloc_shared([block_C, dim_k], dtype) + k_c = T.alloc_shared([block_C, dim_k], dtype) + g_c = T.alloc_shared([block_C], dtype) + h_c = T.alloc_shared([dim_k, dim_v], dtype) + v_new_c = T.alloc_shared([block_C, dim_v], dtype) + attn = T.alloc_shared([block_C, block_C], dtype) + + o_frag = T.alloc_fragment([block_C, dim_v], accum_dtype) + attn_frag = T.alloc_fragment([block_C, block_C], accum_dtype) + + for i, d in T.Parallel(block_C, dim_k): + q_c[i, d] = q[bid, base + i, hid, d] + k_c[i, d] = k[bid, base + i, hid, d] + for i in T.Parallel(block_C): + g_c[i] = g[bid, base + i, hid] + T.copy(S[bid, hid, tid, :, :], h_c, disable_tma=True) + for i, d in T.Parallel(block_C, dim_v): + v_new_c[i, d] = v_new[bid, base + i, hid, d] + + T.clear(o_frag) + T.gemm(q_c, h_c, o_frag) + for i, j in T.Parallel(block_C, dim_v): + o_frag[i, j] = o_frag[i, j] * T.exp2(g_c[i] * _LOG2E) + + T.clear(attn_frag) + T.gemm(q_c, k_c, attn_frag, transpose_B=True) + for i, j in T.Parallel(block_C, block_C): + attn[i, j] = T.if_then_else( + i >= j, + attn_frag[i, j] * T.exp2((g_c[i] - g_c[j]) * _LOG2E), + T.float32(0.0), + ) + + T.gemm(attn, v_new_c, o_frag) + for i, d in T.Parallel(block_C, dim_v): + o[bid, base + i, hid, d] = o_frag[i, d] + + return output_o_bthd + + return _func + + +@functools.lru_cache(maxsize=32) +def _prefill_group_transition_summary_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + group_chunks: int, + dim_k: int, + dim_v: int, + dtype: str = "float16", + block_v: int = 64, +): + accum_dtype = "float32" + block_C = chunk_size + num_chunks = seq_len // block_C + num_groups = num_chunks // group_chunks + dim_aug = dim_v + dim_k + BV = block_v + num_v_tiles = dim_aug // BV + + @tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(num_stages, threads=256): + @T.prim_func + def group_transition_summary_bthd( + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + g: T.Tensor([batch, seq_len, head], dtype), + w: T.Tensor([batch, seq_len, head, dim_k], dtype), + u: T.Tensor([batch, seq_len, head, dim_v], dtype), + summary: T.Tensor([batch, head, num_groups, dim_k, dim_aug], dtype), + ): + with T.Kernel(num_v_tiles * num_groups, batch, head, threads=threads) as ( + vgid, + bid, + hid, + ): + k_c = T.alloc_shared([block_C, dim_k], dtype) + g_c = T.alloc_shared([block_C], dtype) + w_c = T.alloc_shared([block_C, dim_k], dtype) + u_c = T.alloc_shared([block_C, BV], dtype) + h_c = T.alloc_shared([dim_k, BV], dtype) + v_new_c = T.alloc_shared([block_C, BV], dtype) + + ws_frag = T.alloc_fragment([block_C, BV], accum_dtype) + h_next_frag = T.alloc_fragment([dim_k, BV], accum_dtype) + + vid = vgid % num_v_tiles + gid = vgid // num_v_tiles + v_offset = vid * BV + start_chunk = gid * group_chunks + + for i, j in T.Parallel(dim_k, BV): + aug_col = v_offset + j + if aug_col < dim_v: + h_c[i, j] = 0.0 + else: + a_col = aug_col - dim_v + if i == a_col: + h_c[i, j] = 1.0 + else: + h_c[i, j] = 0.0 + + for local_t in T.Pipelined(group_chunks, num_stages=num_stages): + t = start_chunk + local_t + base = t * block_C + for i, d in T.Parallel(block_C, dim_k): + k_c[i, d] = k[bid, base + i, hid, d] + w_c[i, d] = w[bid, base + i, hid, d] + for i in T.Parallel(block_C): + g_c[i] = g[bid, base + i, hid] + for i, j in T.Parallel(block_C, BV): + aug_col = v_offset + j + if aug_col < dim_v: + u_c[i, j] = u[bid, base + i, hid, aug_col] + else: + u_c[i, j] = 0.0 + + T.clear(ws_frag) + T.gemm(w_c, h_c, ws_frag) + for i, j in T.Parallel(block_C, BV): + v_new_c[i, j] = u_c[i, j] - ws_frag[i, j] * T.exp2( + (g_c[i] + g_c[block_C - 1]) * _LOG2E + ) + + for n, j in T.Parallel(block_C, BV): + v_new_c[n, j] = v_new_c[n, j] * T.exp2( + (g_c[block_C - 1] - g_c[n]) * _LOG2E + ) + for i, j in T.Parallel(dim_k, BV): + h_next_frag[i, j] = h_c[i, j] * T.exp2( + g_c[block_C - 1] * _LOG2E + ) + T.gemm( + k_c, + v_new_c, + h_next_frag, + transpose_A=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(h_next_frag, h_c) + + for i, j in T.Parallel(dim_k, BV): + summary[bid, hid, gid, i, v_offset + j] = h_c[i, j] + + return group_transition_summary_bthd + + return _func + + +@functools.lru_cache(maxsize=32) +def _prefill_dense_group_start_scan_bthd_tl( + batch: int, + head: int, + num_groups: int, + dim_k: int, + dim_v: int, + dtype: str = "float16", + block_v: int = 64, +): + accum_dtype = "float32" + dim_aug = dim_v + dim_k + BV = block_v + num_v_tiles = dim_v // BV + + @tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(num_stages, threads=256): + @T.prim_func + def dense_group_start_scan_bthd( + summary: T.Tensor([batch, head, num_groups, dim_k, dim_aug], dtype), + group_start: T.Tensor([batch, head, num_groups, dim_k, dim_v], dtype), + ): + with T.Kernel(num_v_tiles, batch, head, threads=threads) as ( + vid, + bid, + hid, + ): + a_c = T.alloc_shared([dim_k, dim_k], dtype) + b_c = T.alloc_shared([dim_k, BV], dtype) + h_c = T.alloc_shared([dim_k, BV], dtype) + h_next_frag = T.alloc_fragment([dim_k, BV], accum_dtype) + + v_offset = vid * BV + + for i, j in T.Parallel(dim_k, BV): + h_c[i, j] = 0.0 + + for gid in T.Pipelined(num_groups, num_stages=num_stages): + for i, j in T.Parallel(dim_k, BV): + group_start[bid, hid, gid, i, v_offset + j] = h_c[i, j] + b_c[i, j] = summary[bid, hid, gid, i, v_offset + j] + for i, j in T.Parallel(dim_k, dim_k): + a_c[i, j] = summary[bid, hid, gid, i, dim_v + j] + + T.clear(h_next_frag) + T.gemm(a_c, h_c, h_next_frag) + for i, j in T.Parallel(dim_k, BV): + h_c[i, j] = h_next_frag[i, j] + b_c[i, j] + + return dense_group_start_scan_bthd + + return _func + + +@functools.lru_cache(maxsize=32) +def _prefill_grouped_replay_bthd_tl( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + group_chunks: int, + dim_k: int, + dim_v: int, + dtype: str = "float16", + block_v: int = 64, +): + accum_dtype = "float32" + block_C = chunk_size + num_chunks = seq_len // block_C + num_groups = num_chunks // group_chunks + BV = block_v + num_v_tiles = dim_v // BV + + @tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16"], + ) + def _func(num_stages, threads=256): + @T.prim_func + def h_grouped_replay_bthd( + k: T.Tensor([batch, seq_len, head, dim_k], dtype), + g: T.Tensor([batch, seq_len, head], dtype), + w: T.Tensor([batch, seq_len, head, dim_k], dtype), + u: T.Tensor([batch, seq_len, head, dim_v], dtype), + group_start: T.Tensor([batch, head, num_groups, dim_k, dim_v], dtype), + S: T.Tensor([batch, head, num_chunks + 1, dim_k, dim_v], dtype), + v_new: T.Tensor([batch, seq_len, head, dim_v], dtype), + ): + with T.Kernel(num_v_tiles * num_groups, batch, head, threads=threads) as ( + vgid, + bid, + hid, + ): + k_c = T.alloc_shared([block_C, dim_k], dtype) + g_c = T.alloc_shared([block_C], dtype) + w_c = T.alloc_shared([block_C, dim_k], dtype) + u_c = T.alloc_shared([block_C, BV], dtype) + h_c = T.alloc_shared([dim_k, BV], dtype) + v_new_c = T.alloc_shared([block_C, BV], dtype) + + ws_frag = T.alloc_fragment([block_C, BV], accum_dtype) + h_next_frag = T.alloc_fragment([dim_k, BV], accum_dtype) + + vid = vgid % num_v_tiles + gid = vgid // num_v_tiles + v_offset = vid * BV + start_chunk = gid * group_chunks + + T.copy( + group_start[bid, hid, gid, :, v_offset : v_offset + BV], + h_c, + disable_tma=True, + ) + if gid == 0: + for i, j in T.Parallel(dim_k, BV): + S[bid, hid, 0, i, v_offset + j] = h_c[i, j] + + for local_t in T.Pipelined(group_chunks, num_stages=num_stages): + t = start_chunk + local_t + base = t * block_C + for i, d in T.Parallel(block_C, dim_k): + k_c[i, d] = k[bid, base + i, hid, d] + w_c[i, d] = w[bid, base + i, hid, d] + for i in T.Parallel(block_C): + g_c[i] = g[bid, base + i, hid] + for i, d in T.Parallel(block_C, BV): + u_c[i, d] = u[bid, base + i, hid, v_offset + d] + + T.clear(ws_frag) + T.gemm(w_c, h_c, ws_frag) + for i, j in T.Parallel(block_C, BV): + v_new_c[i, j] = u_c[i, j] - ws_frag[i, j] * T.exp2( + (g_c[i] + g_c[block_C - 1]) * _LOG2E + ) + + for i, j in T.Parallel(block_C, BV): + v_new[bid, base + i, hid, v_offset + j] = v_new_c[i, j] + + for n, j in T.Parallel(block_C, BV): + v_new_c[n, j] = v_new_c[n, j] * T.exp2( + (g_c[block_C - 1] - g_c[n]) * _LOG2E + ) + for i, j in T.Parallel(dim_k, BV): + h_next_frag[i, j] = h_c[i, j] * T.exp2( + g_c[block_C - 1] * _LOG2E + ) + T.gemm( + k_c, + v_new_c, + h_next_frag, + transpose_A=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(h_next_frag, h_c) + for i, j in T.Parallel(dim_k, BV): + S[bid, hid, t + 1, i, v_offset + j] = h_c[i, j] + + return h_grouped_replay_bthd + + return _func + + +@torch.library.custom_op("tileops::gated_deltanet_prefill_fwd_kernel", mutates_args=()) +def _gated_deltanet_prefill_wrapped_kernel( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str, + fused_num_stages: int, + fused_threads: int, + h_num_stages: int, + h_threads: int, + h_block_v: int, + o_threads: int, + layout: str, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + layout = _normalize_prefill_layout(layout) + if layout == "bthd": + g_cum = _prefill_chunk_local_cumsum_bthd_tl( + batch, head, seq_len, chunk_size, dtype + )(g) + use_blocksolve_prepare = ( + chunk_size == 64 + and dim_k == 128 + and dim_v == 128 + and dtype == "float16" + ) + use_partitioned_prefill = ( + os.environ.get( + "TILEOPS_GDN_PREFILL_PARTITIONED", + os.environ.get("TILEOPS_GDN_PREFILL_CP_SPLIT", "1"), + ) + != "0" + and batch == 1 + and use_blocksolve_prepare + ) + if use_partitioned_prefill: + from tileops.kernels.gated_deltanet.gdn_prefill import fused_gdr_fwd + + g_zero = torch.zeros_like(g_cum) + A = _prefill_blocksolve_A_bthd(k, g_zero, beta, chunk_size) + initial_state, cu_seqlens, cp_seq_map, raw_cu_seqlens = ( + _prefill_partitioned_initial_state_bthd( + k=k, + v=v, + A=A, + g=g_cum, + beta=beta, + chunk_size=chunk_size, + ) + ) + o, _h, final_state = fused_gdr_fwd( + q=q, + k=k, + v=v, + a=A, + g=g_cum, + b=beta, + scale=1.0, + initial_state=initial_state, + output_final_state=True, + output_h=False, + output_o=True, + cu_seqlens=cu_seqlens, + cp_seq_map=cp_seq_map, + raw_cu_seqlens=raw_cu_seqlens, + ) + return o, final_state.to(q.dtype) + + o_fn = _prefill_output_o_bthd_tl( + batch, head, seq_len, chunk_size, dim_k, dim_v, dtype + )(o_threads) + if use_blocksolve_prepare: + A = _prefill_blocksolve_A_bthd(k, g_cum, beta, chunk_size) + recompute_fn = _prefill_recompute_w_u_from_A_bthd_tl( + batch, head, seq_len, chunk_size, dim_k, dim_v, dtype + )(128) + w, u = recompute_fn(k, v, beta, A) + else: + prepare_fn = _prefill_prepare_w_u_bthd_tl( + batch, head, seq_len, chunk_size, dim_k, dim_v, dtype + )(fused_num_stages, fused_threads) + w, u = prepare_fn(k, v, g_cum, beta) + group_chunks = 64 + use_scan_h = ( + batch == 1 + and head == 16 + and seq_len >= 65536 + and chunk_size == 64 + and dim_k == 128 + and dim_v == 128 + and dtype == "float16" + and seq_len % (chunk_size * group_chunks) == 0 + ) + if use_scan_h: + scan_block_v = 64 + scan_stages = 2 + scan_threads = 256 + summary_fn = _prefill_group_transition_summary_bthd_tl( + batch, + head, + seq_len, + chunk_size, + group_chunks, + dim_k, + dim_v, + dtype, + block_v=scan_block_v, + )(scan_stages, scan_threads) + scan_fn = _prefill_dense_group_start_scan_bthd_tl( + batch, + head, + seq_len // chunk_size // group_chunks, + dim_k, + dim_v, + dtype, + block_v=scan_block_v, + )(scan_stages, scan_threads) + grouped_h_fn = _prefill_grouped_replay_bthd_tl( + batch, + head, + seq_len, + chunk_size, + group_chunks, + dim_k, + dim_v, + dtype, + block_v=scan_block_v, + )(scan_stages, scan_threads) + summary = summary_fn(k, g_cum, w, u) + group_start = scan_fn(summary) + states, v_new = grouped_h_fn(k, g_cum, w, u, group_start) + else: + h_fn = _prefill_h_recurrence_bthd_tl( + batch, + head, + seq_len, + chunk_size, + dim_k, + dim_v, + dtype, + block_v=h_block_v, + )(h_num_stages, h_threads) + S_0 = torch.zeros(batch, head, dim_k, dim_v, dtype=q.dtype, device=q.device) + states, v_new = h_fn(k, g_cum, w, u, S_0) + o = o_fn(q, k, g_cum, states, v_new) + final_state = states[:, :, -1, :, :].contiguous() + return o, final_state + + if chunk_size == 64 and batch * head > 64: + g_cum = _prefill_chunk_local_cumsum_bhtd_tl( + batch, head, seq_len, chunk_size, dtype + )(g) + else: + g_cum = _chunk_local_cumsum(g.float(), chunk_size).to(g.dtype) + prepare_fn = _prefill_prepare_w_u_bhtd_tl( + batch, head, seq_len, chunk_size, dim_k, dim_v, dtype + )(fused_num_stages, fused_threads) + h_fn = _h_recurrence_tl( + batch, + head, + seq_len, + chunk_size, + dim_k, + dim_v, + dtype, + block_v=h_block_v, + )(h_num_stages, h_threads) + o_fn = _output_o_tl(batch, head, seq_len, chunk_size, dim_k, dim_v, dtype)( + o_threads + ) + S_0 = torch.zeros(batch, head, dim_k, dim_v, dtype=q.dtype, device=q.device) + w, u = prepare_fn(k, v, g_cum, beta) + states, v_new = h_fn(k, g_cum, w, u, S_0) + o = o_fn(q, k, g_cum, states, v_new) + final_state = states[:, :, -1, :, :].contiguous() + return o, final_state + + +@_gated_deltanet_prefill_wrapped_kernel.register_fake +def _gated_deltanet_prefill_wrapped_kernel_fake( + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str, + fused_num_stages: int, + fused_threads: int, + h_num_stages: int, + h_threads: int, + h_block_v: int, + o_threads: int, + layout: str, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + del dtype, fused_num_stages, fused_threads, h_num_stages, h_threads, h_block_v, o_threads + del k, v, g, beta + layout = _normalize_prefill_layout(layout) + if layout == "bthd": + o = torch.empty(batch, seq_len, head, dim_v, dtype=q.dtype, device=q.device) + else: + o = torch.empty(batch, head, seq_len, dim_v, dtype=q.dtype, device=q.device) + final_state = torch.empty(batch, head, dim_k, dim_v, dtype=q.dtype, device=q.device) + return o, final_state + + +class GatedDeltaNetPrefillFwdKernel(Kernel): + """Gated DeltaNet zero-state prefill.""" + + supported_archs: list[int] = [80, 89, 90] + + def __init__( + self, + batch: int, + head: int, + seq_len: int, + chunk_size: int, + dim_k: int, + dim_v: int, + dtype: str = "float32", + config: Optional[dict] = None, + layout: str = "bhtd", + tune: bool = False, + ): + super().__init__() + layout = _normalize_prefill_layout(layout) + self.batch = batch + self.head = head + self.seq_len = seq_len + self.chunk_size = chunk_size + self.dim_k = dim_k + self.dim_v = dim_v + self.dtype = dtype + self.layout = layout + self.init_config(config, tune) + + @property + def default_config(self) -> dict: + streams = self.batch * self.head + if self.layout == "bthd": + if self.chunk_size == 64 and streams >= 64: + h_block_v = 32 + h_num_stages = 2 + h_threads = 256 + else: + h_block_v = 16 + h_num_stages = 2 + h_threads = 128 + return { + "fused_num_stages": 2, + "fused_threads": 128, + "h_num_stages": h_num_stages, + "h_threads": h_threads, + "h_block_v": h_block_v, + "o_threads": 256, + } + + if self.chunk_size >= 128 and 1 < streams <= 8: + h_block_v = 8 + h_num_stages = 2 + h_threads = 64 + elif self.chunk_size >= 64 and streams <= 4: + h_block_v = 8 + h_num_stages = 2 + h_threads = 128 + elif self.chunk_size >= 64 and streams <= 48: + h_block_v = 16 + h_num_stages = 2 + h_threads = 128 + elif self.chunk_size == 64 and streams <= 66: + h_block_v = 64 + h_num_stages = 3 + h_threads = 256 + elif self.chunk_size >= 64: + h_block_v = 16 + h_num_stages = 3 + h_threads = 128 + else: + h_block_v = 0 + h_num_stages = 2 + h_threads = 256 + fused_threads = 128 if self.chunk_size >= 64 else 256 + o_threads = 128 if self.chunk_size == 64 and streams > 64 else 256 + return { + "fused_num_stages": 2, + "fused_threads": fused_threads, + "h_num_stages": h_num_stages, + "h_threads": h_threads, + "h_block_v": h_block_v, + "o_threads": o_threads, + } + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return _gated_deltanet_prefill_wrapped_kernel( + self.batch, + self.head, + self.seq_len, + self.chunk_size, + self.dim_k, + self.dim_v, + self.dtype_str, + self.config["fused_num_stages"], + self.config["fused_threads"], + self.config["h_num_stages"], + self.config["h_threads"], + self.config.get("h_block_v", 0), + self.config["o_threads"], + self.layout, + q, + k, + v, + g, + beta, + ) diff --git a/tileops/kernels/gated_deltanet/gdn_prefill/__init__.py b/tileops/kernels/gated_deltanet/gdn_prefill/__init__.py new file mode 100644 index 000000000..b64851d4a --- /dev/null +++ b/tileops/kernels/gated_deltanet/gdn_prefill/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under the MIT License; see THIRD_PARTY_NOTICES.md for details. +# Adapted and modified for TileOps GatedDeltaNet prefill integration. + +from .tilelang_compat import install_gemm_v1_compat + +install_gemm_v1_compat() + +from .cp_fwd import correct_initial_states, get_warmup_chunks # noqa: E402 +from .fused_fwd import fused_gdr_fwd # noqa: E402 +from .prepare_h import fused_gdr_h # noqa: E402 + +__all__ = [ + "correct_initial_states", + "fused_gdr_fwd", + "fused_gdr_h", + "get_warmup_chunks", +] diff --git a/tileops/kernels/gated_deltanet/gdn_prefill/cp_fwd.py b/tileops/kernels/gated_deltanet/gdn_prefill/cp_fwd.py new file mode 100644 index 000000000..825281048 --- /dev/null +++ b/tileops/kernels/gated_deltanet/gdn_prefill/cp_fwd.py @@ -0,0 +1,350 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under the MIT License; see THIRD_PARTY_NOTICES.md for details. +# Adapted and modified for TileOps GatedDeltaNet prefill integration. + +import tilelang +import tilelang.language as T +import torch + + +@tilelang.jit() +def tilelang_get_warmup_chunks( + num_heads, + chunk_size, + threshold, + accum_dtype, + g_dtype, + mask_dtype, + seqlen_dtype, +): + batch_size = T.dynamic("batch_size") + num_tokens = T.dynamic("num_tokens") + num_threads = tilelang.cdiv(num_heads, 32) * 32 + + @T.prim_func + def tilelang_get_warmup_chunks_kernel( + g: T.Tensor([1, num_tokens, num_heads], dtype=g_dtype), + ht_mask: T.Tensor([batch_size], dtype=mask_dtype), + cu_seqlens: T.Tensor([batch_size + 1], dtype=seqlen_dtype), + num_warmup_chunks: T.Tensor([batch_size, num_heads], dtype=seqlen_dtype), + fallback_mask: T.Tensor([batch_size, num_heads], dtype=mask_dtype), + ): + with T.Kernel(batch_size, threads=num_threads) as (bb,): + if ht_mask[bb]: + for i_h in T.Parallel(num_heads): + num_warmup_chunks[bb, i_h] = 0 + else: + seq_start_idx = T.alloc_var("int32") + seq_end_idx = T.alloc_var("int32") + num_iters = T.alloc_var("int32") + seq_start_idx = cu_seqlens[bb] + seq_end_idx = cu_seqlens[bb + 1] + num_iters = (seq_end_idx - seq_start_idx) // chunk_size + + g_fragment = T.alloc_fragment((num_heads), dtype=accum_dtype) + g_cumsum = T.alloc_fragment((num_heads), dtype=accum_dtype) + n_fragment = T.alloc_fragment((num_heads), dtype=seqlen_dtype) + f_fragment = T.alloc_fragment((num_heads), dtype=mask_dtype) + T.clear(g_cumsum) + T.fill(n_fragment, num_iters) + T.fill(f_fragment, True) + + for i_s in T.serial(num_iters): + for i_h in T.Parallel(num_heads): + g_fragment[i_h] = g[0, seq_end_idx - i_s * chunk_size - 1, i_h] + for i_h in T.Parallel(num_heads): + g_cumsum[i_h] += g_fragment[i_h] + for i_h in T.Parallel(num_heads): + if g_cumsum[i_h] < threshold and n_fragment[i_h] == num_iters: + n_fragment[i_h] = i_s + 1 + f_fragment[i_h] = False + + for i_h in T.Parallel(num_heads): + num_warmup_chunks[bb, i_h] = n_fragment[i_h] + for i_h in T.Parallel(num_heads): + fallback_mask[bb, i_h] = f_fragment[i_h] + + return tilelang_get_warmup_chunks_kernel + + +def get_warmup_chunks( + g: torch.Tensor, # [1, num_total_tokens, num_v_heads] + cu_seqlens: torch.Tensor, # [cp_real_batch_size + 1] + ht_mask: torch.Tensor, # [cp_real_batch_size] + chunk_size: int = 64, + threshold: float = -10.0, +): + batch_size, num_tokens, num_heads = g.shape + real_batch_size = ht_mask.shape[0] + assert cu_seqlens.shape[0] == real_batch_size + 1 + assert batch_size == 1 + assert chunk_size == 64 + + tilelang_get_warmup_chunks_kernel = tilelang_get_warmup_chunks( + num_heads=num_heads, + chunk_size=chunk_size, + threshold=threshold, + accum_dtype="float32", + g_dtype=g.dtype, + mask_dtype=ht_mask.dtype, + seqlen_dtype=cu_seqlens.dtype, + ) + num_warmup_chunks = torch.empty( + [real_batch_size, num_heads], dtype=cu_seqlens.dtype, device=cu_seqlens.device + ) + fallback_mask = torch.empty( + [real_batch_size, num_heads], dtype=ht_mask.dtype, device=cu_seqlens.device + ) + tilelang_get_warmup_chunks_kernel( + g, ht_mask, cu_seqlens, num_warmup_chunks, fallback_mask + ) + + return num_warmup_chunks, fallback_mask + + +@tilelang.jit() +def tilelang_correct_h0( + H, + DK, + DV, + res_dtype, + accum_dtype, + buffer_dtype, + seqlen_dtype, + mask_dtype, + use_raw_h0, + block_DV: int = 32, +): + cp_batch_size = T.dynamic("cp_batch_size") + raw_batch_size = T.dynamic("raw_batch_size") + + @T.macro + def kernel_body( + bb, + bh, + bv, + seq_start_idx, + seq_end_idx, + num_iters, + ht_buffer, + mt_buffer, + fallback_mask, + seq_map_r2c, + cp_h0, + h_fragment, + ): + h_shared = T.alloc_shared((DK, block_DV), dtype=buffer_dtype) + hd_shared = T.alloc_shared((DK, block_DV), dtype=buffer_dtype) + m_shared = T.alloc_shared((DK, DK), dtype=buffer_dtype) + + T.copy( + h_fragment, + cp_h0[seq_start_idx, bh, 0:DK, bv * block_DV : (bv + 1) * block_DV], + ) + + for i_s in T.Pipelined(num_iters - 1, num_stages=2): + if fallback_mask[seq_start_idx + i_s, bh]: + T.copy(h_fragment, hd_shared) + T.copy( + ht_buffer[ + seq_start_idx + i_s, bh, 0:DK, bv * block_DV : (bv + 1) * block_DV + ], + h_shared, + ) + T.copy(h_shared, h_fragment) + if fallback_mask[seq_start_idx + i_s, bh]: + T.copy(mt_buffer[seq_start_idx + i_s, bh, 0:DK, 0:DK], m_shared) + T.gemm(m_shared, hd_shared, h_fragment, clear_accum=False) + T.copy( + h_fragment, + cp_h0[ + seq_start_idx + i_s + 1, + bh, + 0:DK, + bv * block_DV : (bv + 1) * block_DV, + ], + ) + + if use_raw_h0: + + @T.prim_func + def tilelang_correct_h0_kernel( + raw_h0: T.Tensor([raw_batch_size, H, DK, DV], dtype=res_dtype), + ht_buffer: T.Tensor([cp_batch_size, H, DK, DV], dtype=buffer_dtype), + mt_buffer: T.Tensor([cp_batch_size, H, DK, DK], dtype=buffer_dtype), + fallback_mask: T.Tensor([cp_batch_size, H], dtype=mask_dtype), + seq_map_r2c: T.Tensor([raw_batch_size + 1], dtype=seqlen_dtype), + cp_h0: T.Tensor([cp_batch_size, H, DK, DV], dtype=res_dtype), + ): + with T.Kernel( + T.ceildiv(DV, block_DV) * H * raw_batch_size, threads=128 + ) as (bbhv,): + bbh, bv = ( + bbhv // T.ceildiv(DV, block_DV), + bbhv % T.ceildiv(DV, block_DV), + ) + bb, bh = bbh // H, bbh % H + + seq_start_idx = seq_map_r2c[bb] + seq_end_idx = seq_map_r2c[bb + 1] + num_iters = seq_end_idx - seq_start_idx + + h_fragment = T.alloc_fragment((DK, block_DV), dtype=accum_dtype) + T.copy( + raw_h0[bb, bh, 0:DK, bv * block_DV : (bv + 1) * block_DV], + h_fragment, + ) + + kernel_body( + bb, + bh, + bv, + seq_start_idx, + seq_end_idx, + num_iters, + ht_buffer, + mt_buffer, + fallback_mask, + seq_map_r2c, + cp_h0, + h_fragment, + ) + + else: + + @T.prim_func + def tilelang_correct_h0_kernel( + ht_buffer: T.Tensor([cp_batch_size, H, DK, DV], dtype=buffer_dtype), + mt_buffer: T.Tensor([cp_batch_size, H, DK, DK], dtype=buffer_dtype), + fallback_mask: T.Tensor([cp_batch_size, H], dtype=mask_dtype), + seq_map_r2c: T.Tensor([raw_batch_size + 1], dtype=seqlen_dtype), + cp_h0: T.Tensor([cp_batch_size, H, DK, DV], dtype=res_dtype), + ): + with T.Kernel( + T.ceildiv(DV, block_DV) * H * raw_batch_size, threads=128 + ) as (bbhv,): + bbh, bv = ( + bbhv // T.ceildiv(DV, block_DV), + bbhv % T.ceildiv(DV, block_DV), + ) + bb, bh = bbh // H, bbh % H + + seq_start_idx = seq_map_r2c[bb] + seq_end_idx = seq_map_r2c[bb + 1] + num_iters = seq_end_idx - seq_start_idx + + h_fragment = T.alloc_fragment((DK, block_DV), dtype=accum_dtype) + T.clear(h_fragment) + + kernel_body( + bb, + bh, + bv, + seq_start_idx, + seq_end_idx, + num_iters, + ht_buffer, + mt_buffer, + fallback_mask, + seq_map_r2c, + cp_h0, + h_fragment, + ) + + return tilelang_correct_h0_kernel + + +@tilelang.jit() +def tilelang_zero_initial_cp_h0( + H, + DK, + DV, + res_dtype, + seqlen_dtype, + block_DV: int = 32, +): + cp_batch_size = T.dynamic("cp_batch_size") + raw_batch_size = T.dynamic("raw_batch_size") + + @T.prim_func + def tilelang_zero_initial_cp_h0_kernel( + seq_map_r2c: T.Tensor([raw_batch_size + 1], dtype=seqlen_dtype), + cp_h0: T.Tensor([cp_batch_size, H, DK, DV], dtype=res_dtype), + ): + with T.Kernel( + T.ceildiv(DV, block_DV) * H * raw_batch_size, threads=128 + ) as (bbhv,): + bbh, bv = ( + bbhv // T.ceildiv(DV, block_DV), + bbhv % T.ceildiv(DV, block_DV), + ) + bb, bh = bbh // H, bbh % H + seq_start_idx = seq_map_r2c[bb] + for j_k, j_v in T.Parallel(DK, block_DV): + cp_h0[seq_start_idx, bh, j_k, bv * block_DV + j_v] = 0.0 + + return tilelang_zero_initial_cp_h0_kernel + + +def correct_initial_states( + raw_h0: torch.Tensor + | None, # [raw_batch_size, num_v_heads, k_head_dim, v_head_dim] + ht_buffer: torch.Tensor, # [cp_batch_size, num_v_heads, k_head_dim, v_head_dim] + mt_buffer: torch.Tensor, # [cp_batch_size, num_v_heads, k_head_dim, k_head_dim] + fallback_mask: torch.Tensor, # [cp_batch_size, num_v_heads] + seq_map_r2c: torch.Tensor, # [raw_batch_size + 1] +): + cp_batch_size = fallback_mask.shape[0] + _, num_heads, k_head_dim, v_head_dim = ht_buffer.shape + assert k_head_dim == v_head_dim == 128 + + if raw_h0 is None: + res_dtype = torch.float32 + use_raw_h0 = False + else: + res_dtype = raw_h0.dtype + use_raw_h0 = True + + tilelang_correct_h0_kernel = tilelang_correct_h0( + H=num_heads, + DK=k_head_dim, + DV=v_head_dim, + res_dtype=res_dtype, + accum_dtype="float32", + buffer_dtype=ht_buffer.dtype, + seqlen_dtype=seq_map_r2c.dtype, + mask_dtype=fallback_mask.dtype, + use_raw_h0=use_raw_h0, + ) + cp_h0 = torch.empty( + (cp_batch_size, num_heads, k_head_dim, v_head_dim), + dtype=res_dtype, + device=ht_buffer.device, + ) + tilelang_zero_initial_cp_h0_kernel = tilelang_zero_initial_cp_h0( + H=num_heads, + DK=k_head_dim, + DV=v_head_dim, + res_dtype=res_dtype, + seqlen_dtype=seq_map_r2c.dtype, + ) + tilelang_zero_initial_cp_h0_kernel(seq_map_r2c, cp_h0) + if use_raw_h0: + tilelang_correct_h0_kernel( + raw_h0, + ht_buffer, + mt_buffer, + fallback_mask, + seq_map_r2c, + cp_h0, + ) + else: + tilelang_correct_h0_kernel( + ht_buffer, + mt_buffer, + fallback_mask, + seq_map_r2c, + cp_h0, + ) + + return cp_h0 diff --git a/tileops/kernels/gated_deltanet/gdn_prefill/fused_fwd.py b/tileops/kernels/gated_deltanet/gdn_prefill/fused_fwd.py new file mode 100644 index 000000000..cd3d98939 --- /dev/null +++ b/tileops/kernels/gated_deltanet/gdn_prefill/fused_fwd.py @@ -0,0 +1,678 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under the MIT License; see THIRD_PARTY_NOTICES.md for details. +# Adapted and modified for TileOps GatedDeltaNet prefill integration. + +import os + +import tilelang +import tilelang.language as T +import torch + +from .utils import prepare_chunk_offsets + +MULTI_PROCESSOR_COUNT = torch.cuda.get_device_properties().multi_processor_count +TARGET_NUM_CTAS = int(MULTI_PROCESSOR_COUNT * 0.7) + + +@tilelang.jit( + # out_idx=[-3, -2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + # tilelang.PassConfigKey.TL_DISABLE_THREAD_STORAGE_SYNC: True, + }, + compile_flags=["-O3", "-DENABLE_BF16", "-include", "tl_templates/cuda/gemm.h"], +) +def tilelang_fused_chunk_gdr_fwd( + H, + Hg, + DK, + DV, + chunk_size, + scale, + accum_dtype, + qkva_dtype, + g_dtype, + b_dtype, + h0_dtype, + ht_dtype, + h_dtype, + o_dtype, + seqlen_dtype, + use_initial_state, + store_final_state, + store_h, + store_o, + is_varlen, + is_cp, + block_DV=128, +): + batch_size = T.dynamic("batch_size") + num_tokens = T.dynamic("num_tokens") + num_chunks = T.dynamic("num_chunks") + raw_batch_size = T.dynamic("raw_batch_size") + block_S = chunk_size + + if is_varlen: + q_shape = (1, num_tokens, Hg, DK) + k_shape = (1, num_tokens, Hg, DK) + v_shape = (1, num_tokens, H, DV) + o_shape = (1, num_tokens, H, DV) + a_shape = (1, num_tokens, H, chunk_size) + g_shape = (1, num_tokens, H) + b_shape = (1, num_tokens, H) + h_shape = (1, num_chunks, H, DK, DV) + else: + q_shape = (batch_size, num_tokens, Hg, DK) + k_shape = (batch_size, num_tokens, Hg, DK) + v_shape = (batch_size, num_tokens, H, DV) + o_shape = (batch_size, num_tokens, H, DV) + a_shape = (batch_size, num_tokens, H, chunk_size) + g_shape = (batch_size, num_tokens, H) + b_shape = (batch_size, num_tokens, H) + h_shape = (batch_size, num_chunks, H, DK, DV) + h0_shape = (batch_size, H, DK, DV) + ht_shape = (raw_batch_size, H, DK, DV) + + @T.prim_func + def tilelang_fused_chunk_gdr_fwd_kernel( + q: T.Tensor(q_shape, dtype=qkva_dtype), + k: T.Tensor(k_shape, dtype=qkva_dtype), + v: T.Tensor(v_shape, dtype=qkva_dtype), + a: T.Tensor(a_shape, dtype=qkva_dtype), + g: T.Tensor(g_shape, dtype=g_dtype), + b: T.Tensor(b_shape, dtype=b_dtype), + h0: T.Tensor(h0_shape, dtype=h0_dtype), + cu_seqlens: T.Tensor([batch_size + 1], dtype=seqlen_dtype), + chunk_offsets: T.Tensor([batch_size + 1], dtype=seqlen_dtype), + cp_seq_map: T.Tensor([batch_size], dtype=seqlen_dtype), + raw_cu_seqlens: T.Tensor([raw_batch_size + 1], dtype=seqlen_dtype), + o: T.Tensor(o_shape, dtype=o_dtype), + h: T.Tensor(h_shape, dtype=h_dtype), + ht: T.Tensor(ht_shape, dtype=ht_dtype), + ): + with T.Kernel(T.ceildiv(DV, block_DV) * batch_size * H, threads=512) as (bbhv,): + bbh, bv = bbhv // T.ceildiv(DV, block_DV), bbhv % T.ceildiv(DV, block_DV) + bb, bh = bbh // H, bbh % H + bhg = bh // (H // Hg) + + batch_idx = T.alloc_var("int32") + seq_start_idx = T.alloc_var("int32") + seq_end_idx = T.alloc_var("int32") + seq_split_idx = T.alloc_var("int32") + chunk_start_idx = T.alloc_var("int32") + chunk_split_idx = T.alloc_var("int32") + + batch_idx = 0 if is_varlen else bb + seq_start_idx = cu_seqlens[bb] if is_varlen else 0 + seq_end_idx = cu_seqlens[bb + 1] if is_varlen else num_tokens + chunk_start_idx = chunk_offsets[bb] if is_varlen else 0 + + raw_batch_idx = T.alloc_var("int32") + raw_seq_end_idx = T.alloc_var("int32") + need_store_final_state = T.alloc_var("bool") + raw_batch_idx = cp_seq_map[bb] if is_cp else bb + raw_seq_end_idx = ( + raw_cu_seqlens[raw_batch_idx + 1] if is_cp else seq_end_idx + ) + need_store_final_state = store_final_state & ( + raw_seq_end_idx == seq_end_idx + ) + + num_iters = T.alloc_var("int32") + num_unmasked_iters = T.alloc_var("int32") + num_iters = T.ceildiv(seq_end_idx - seq_start_idx, block_S) + num_unmasked_iters = (seq_end_idx - seq_start_idx) // block_S + + q_shared = T.alloc_shared((2, block_S, DK), dtype=qkva_dtype) + k_shared = T.alloc_shared((2, block_S, DK), dtype=qkva_dtype) + v_shared = T.alloc_shared((2, block_S, block_DV), dtype=qkva_dtype) + a_shared = T.alloc_shared((2, block_S, block_S), dtype=qkva_dtype) + g_shared = T.alloc_shared((2, block_S), dtype=accum_dtype, scope="shared") + b_shared = T.alloc_shared((2, block_S), dtype=accum_dtype, scope="shared") + + o_shared = T.alloc_shared((block_S, block_DV), dtype=o_dtype) + h_shared = T.alloc_shared((DK, block_DV), dtype=qkva_dtype) + vd_shared = T.alloc_shared((block_S, block_DV), dtype=qkva_dtype) + vn_shared = T.alloc_shared((block_S, block_DV), dtype=qkva_dtype) + p_shared = T.alloc_shared((block_S, block_S), dtype=qkva_dtype) + g_exp_shared = T.alloc_shared((block_S), dtype=accum_dtype, scope="shared") + g_rev_exp_shared = T.alloc_shared( + (block_S), dtype=accum_dtype, scope="shared" + ) + + h_fragment = T.alloc_fragment((DK, block_DV), dtype=accum_dtype) + o_fragment = T.alloc_fragment((block_S, block_DV), dtype=accum_dtype) + v_fragment = T.alloc_fragment((block_S, block_DV), dtype=accum_dtype) + u_fragment = T.alloc_fragment((block_S, block_DV), dtype=accum_dtype) + p_fragment = T.alloc_fragment((block_S, block_S), dtype=accum_dtype) + a_fragment = T.alloc_fragment((block_S, block_S), dtype=accum_dtype) + g_fragment = T.alloc_fragment((block_S, block_S), dtype=accum_dtype) + g_last_local = T.alloc_local((1), dtype=accum_dtype) + + data_is_ready = T.alloc_barrier(arrive_count=[96] * 2) + data_is_free = T.alloc_barrier(arrive_count=[384] * 2) + + bar_o = T.alloc_barrier(arrive_count=128) + bar_0 = T.alloc_barrier(arrive_count=416) + bar_1 = T.alloc_barrier(arrive_count=256) + _bar_2 = T.alloc_barrier(arrive_count=128) + bar_3 = T.alloc_barrier(arrive_count=128) + bar_4 = T.alloc_barrier(arrive_count=128) + bar_5 = T.alloc_barrier(arrive_count=416) + + T.use_swizzle(10) + + tx = T.get_thread_binding() + + PRODUCER_NREG = 32 + CONSUMER_V_NREG = 128 + CONSUMER_S_NREG = 160 + CONSUMER_O_NREG = 128 + + if tx < 128: + T.set_max_nreg(CONSUMER_S_NREG, 1) + + # Initialize S + if use_initial_state: + T.copy( + h0[bb, bh, 0:DK, bv * block_DV : (bv + 1) * block_DV], + h_fragment, + ) + else: + T.clear(h_fragment) + + # Main Loop + for i_s in T.serial(num_iters): + # [STAGE 0] + T.barrier_wait(data_is_ready[i_s % 2], (i_s // 2 + 0) % 2) + T.barrier_arrive(bar_0) + + # [STAGE 0] 0 + T.barrier_wait(bar_0, i_s % 2) + # S4[S] S + T.copy(h_fragment, h_shared) + T.barrier_arrive(bar_1) + + # [STAGE 0] 2, 3, 4 + T.barrier_wait(bar_1, i_s % 2) + # S = g_last * S + g_last_local[0] = g_exp_shared[block_S - 1] + for j_k, j_v in T.Parallel(DK, block_DV): + h_fragment[j_k, j_v] *= g_last_local[0] + T.barrier_arrive(bar_5) + + # [STAGE 0] 5 + T.barrier_wait(bar_5, i_s % 2) + # S += K^T @ V' + T.gemm_v1( + k_shared[i_s % 2, :, :], + vn_shared, + h_fragment, + transpose_A=True, + clear_accum=False, + ) + + T.barrier_arrive(data_is_free[i_s % 2]) + + # Store final S + if need_store_final_state: + T.copy( + h_fragment, + ht[ + raw_batch_idx, bh, 0:DK, bv * block_DV : (bv + 1) * block_DV + ], + ) + + elif tx < 256: + T.set_max_nreg(CONSUMER_V_NREG, 1) + + # Main Loop + for i_s in T.serial(num_iters): + # [STAGE 0] + T.barrier_wait(data_is_ready[i_s % 2], (i_s // 2 + 0) % 2) + T.barrier_arrive(bar_0) + + # [STAGE 0] 0 + T.barrier_wait(bar_0, i_s % 2) + # Precompute g, g_last/g + for j_s in T.Parallel(block_S): + g_exp_shared[j_s] = T.exp2(g_shared[i_s % 2, j_s] * 1.442695) + for j_s in T.Parallel(block_S): + g_rev_exp_shared[j_s] = T.if_then_else( + seq_start_idx + i_s * block_S + j_s < seq_end_idx, + T.exp2( + ( + g_shared[i_s % 2, block_S - 1] + - g_shared[i_s % 2, j_s] + ) + * 1.442695 + ), + 0.0, + ) + T.barrier_arrive(bar_1) + + # [STAGE 0] 1 + T.barrier_wait(bar_1, i_s % 2) + # U = K @ S + T.gemm_v1( + k_shared[i_s % 2, :, :], h_shared, u_fragment, clear_accum=True + ) + + # [STAGE 0] 2 + # W = V - g * U + for j_s, j_v in T.Parallel(block_S, block_DV): + u_fragment[j_s, j_v] *= -g_exp_shared[j_s] + for j_s, j_v in T.Parallel(block_S, block_DV): + u_fragment[j_s, j_v] += v_shared[i_s % 2, j_s, j_v] + # S2[V] W + for j_s, j_v in T.Parallel(block_S, block_DV): + v_shared[i_s % 2, j_s, j_v] = u_fragment[j_s, j_v] + + # [STAGE 0] 3 + T.barrier_wait(bar_3, i_s % 2) + # Vd = Ag @ W + T.gemm_v1( + a_shared[i_s % 2, :, :], + v_shared[i_s % 2, :, :], + v_fragment, + clear_accum=True, + ) + # S2[2] Vd + T.copy(v_fragment, vd_shared) + T.barrier_arrive(bar_4) + + # [STAGE 0] 4 + # V' = g_last/g Vd + for j_s, j_v in T.Parallel(block_S, block_DV): + v_fragment[j_s, j_v] *= g_rev_exp_shared[j_s] + # S2[1] V' + T.copy(v_fragment, vn_shared) + T.barrier_arrive(bar_5) + + T.barrier_wait(bar_5, i_s % 2) + + T.barrier_arrive(data_is_free[i_s % 2]) + + elif tx < 384: + T.set_max_nreg(CONSUMER_O_NREG, 1) + + # Main Loop + for i_s in T.serial(num_iters): + # [STAGE 0] + T.barrier_wait(data_is_ready[i_s % 2], (i_s // 2 + 0) % 2) + T.barrier_arrive(bar_0) + + # [STAGE 0] 0 + T.barrier_wait(bar_0, i_s % 2) + # P = Q K^T + T.gemm_v1( + q_shared[i_s % 2, :, :], + k_shared[i_s % 2, :, :], + p_fragment, + transpose_B=True, + clear_accum=True, + ) + + # [STAGE 0] 1 + # G = Lower(diag(g) @ I @ diag(1/g)) + for j_s, j_t in T.Parallel(block_S, block_S): + g_fragment[j_s, j_t] = ( + g_shared[i_s % 2, j_s] - g_shared[i_s % 2, j_t] + ) + for j_s, j_t in T.Parallel(block_S, block_S): + if j_s >= j_t: + g_fragment[j_s, j_t] = T.exp2( + g_fragment[j_s, j_t] * 1.442695 + ) + else: + g_fragment[j_s, j_t] = 0 + # Ag = G * Ar * b + for j_s, j_t in T.Parallel(block_S, block_S): + a_fragment[j_s, j_t] = a_shared[i_s % 2, j_s, j_t] + for j_s, j_t in T.Parallel(block_S, block_S): + a_fragment[j_s, j_t] *= g_fragment[j_s, j_t] + for j_s, j_t in T.Parallel(block_S, block_S): + a_fragment[j_s, j_t] *= b_shared[i_s % 2, j_t] + for j_s, j_t in T.Parallel(block_S, block_S): + a_shared[i_s % 2, j_s, j_t] = a_fragment[j_s, j_t] + + # [STAGE 0] 2 + T.barrier_wait(bar_1, i_s % 2) + # O = Q @ S + T.gemm_v1( + q_shared[i_s % 2, :, :], h_shared, o_fragment, clear_accum=True + ) + + # [STAGE 0] 3 + # Pg = s * G * P + for j_s, j_t in T.Parallel(block_S, block_S): + p_fragment[j_s, j_t] *= scale * g_fragment[j_s, j_t] + # S1[1] Pg + T.copy(p_fragment, p_shared) + T.barrier_arrive(bar_3) + # O = s * g * O + for j_s, j_k in T.Parallel(block_S, DK): + o_fragment[j_s, j_k] *= scale * g_exp_shared[j_s] + + # [STAGE 0] 4 + T.barrier_wait(bar_4, i_s % 2) + # O += Pg @ Vd + T.gemm_v1(p_shared, vd_shared, o_fragment, clear_accum=False) + T.barrier_arrive(bar_5) + + # [STAGE 0] 5 + T.barrier_wait(bar_5, i_s % 2) + # S2[S] O + T.copy(o_fragment, o_shared) + + T.barrier_arrive(data_is_free[i_s % 2]) + + T.barrier_arrive(bar_o) + + else: + T.set_max_nreg(PRODUCER_NREG, 0) + + if tx < 384 + 32: + for i_s in T.serial(num_iters): + T.barrier_wait(data_is_free[i_s % 2], (i_s // 2 + 1) % 2) + left = seq_start_idx + i_s * block_S + right = left + block_S + + # Load Q + T.tma_copy( + q[batch_idx, left:right, bhg, 0:DK], + q_shared[i_s % 2, :, :], + barrier=data_is_ready[i_s % 2], + ) + # Load K + T.tma_copy( + k[batch_idx, left:right, bhg, 0:DK], + k_shared[i_s % 2, :, :], + barrier=data_is_ready[i_s % 2], + ) + + T.barrier_arrive(data_is_ready[i_s % 2]) + + elif tx < 384 + 64: + for i_s in T.serial(num_iters): + T.barrier_wait(data_is_free[i_s % 2], (i_s // 2 + 1) % 2) + left = seq_start_idx + i_s * block_S + right = left + block_S + + # Load V + T.tma_copy( + v[ + batch_idx, + left:right, + bh, + bv * block_DV : (bv + 1) * block_DV, + ], + v_shared[i_s % 2, :, :], + barrier=data_is_ready[i_s % 2], + ) + # Load beta + if right <= seq_end_idx: + for j_s in T.Parallel(block_S): + b_shared[i_s % 2, j_s] = b[batch_idx, left + j_s, bh] + else: + for j_s in T.Parallel(block_S): + if left + j_s < seq_end_idx: + b_shared[i_s % 2, j_s] = b[ + batch_idx, left + j_s, bh + ] + else: + b_shared[i_s % 2, j_s] = 0 + + T.barrier_arrive(data_is_ready[i_s % 2]) + + elif tx < 384 + 96: + for i_s in T.serial(num_iters): + T.barrier_wait(data_is_free[i_s % 2], (i_s // 2 + 1) % 2) + left = seq_start_idx + i_s * block_S + right = left + block_S + + # Load A + T.tma_copy( + a[batch_idx, left:right, bh, 0:block_S], + a_shared[i_s % 2, :, :], + barrier=data_is_ready[i_s % 2], + ) + # Load gamma + if right <= seq_end_idx: + for j_s in T.Parallel(block_S): + g_shared[i_s % 2, j_s] = g[batch_idx, left + j_s, bh] + else: + for j_s in T.Parallel(block_S): + if left + j_s < seq_end_idx: + g_shared[i_s % 2, j_s] = g[ + batch_idx, left + j_s, bh + ] + else: + g_shared[i_s % 2, j_s] = g[ + batch_idx, seq_end_idx - 1, bh + ] + + T.barrier_arrive(data_is_ready[i_s % 2]) + + else: + for i_s in T.serial(num_unmasked_iters): + right = seq_start_idx + i_s * block_S + left = right - block_S + + T.barrier_arrive(bar_0) + + T.barrier_wait(bar_0, i_s % 2) + # Store O + if i_s > 0 and store_o: + T.copy( + o_shared, + o[ + batch_idx, + left:right, + bh, + bv * block_DV : (bv + 1) * block_DV, + ], + ) + T.barrier_arrive(bar_5) + + T.barrier_wait(bar_1, i_s % 2) + # Store S + if store_h: + T.copy( + h_shared, + h[ + batch_idx, + chunk_start_idx + i_s, + bh, + 0:DK, + bv * block_DV : (bv + 1) * block_DV, + ], + ) + + if num_unmasked_iters < num_iters: + seq_split_idx = seq_start_idx + num_unmasked_iters * block_S + chunk_split_idx = chunk_start_idx + num_unmasked_iters + + T.barrier_arrive(bar_0) + + T.barrier_wait(bar_0, num_unmasked_iters % 2) + # Store O + if num_unmasked_iters > 0 and store_o: + T.copy( + o_shared, + o[ + batch_idx, + seq_split_idx - block_S : seq_split_idx, + bh, + bv * block_DV : (bv + 1) * block_DV, + ], + ) + T.barrier_arrive(bar_5) + + T.barrier_wait(bar_1, num_unmasked_iters % 2) + # Store S + if store_h: + T.copy( + h_shared, + h[ + batch_idx, + chunk_split_idx, + bh, + 0:DK, + bv * block_DV : (bv + 1) * block_DV, + ], + ) + + seq_split_idx = seq_start_idx + (num_iters - 1) * block_S + + # Store O + T.barrier_wait(bar_o, 0) + if store_o: + for j_s, j_v in T.Parallel(block_S, block_DV): + with T.If(seq_split_idx + j_s < seq_end_idx): # noqa: SIM117 + with T.Then(): + o[ + batch_idx, + seq_split_idx + j_s, + bh, + bv * block_DV + j_v, + ] = o_shared[j_s, j_v] + + return tilelang_fused_chunk_gdr_fwd_kernel + + +def fused_gdr_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + g: torch.Tensor, + b: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = True, + output_h: bool = False, + output_o: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cp_seq_map: torch.LongTensor | None = None, + raw_cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, +): + batch_size, num_tokens, Hg, K = k.shape + _, _, H, V = v.shape + scale = scale or K ** (-0.5) + assert K == V == 128 + assert chunk_size == 64 + + if cu_seqlens is None: + real_batch_size = batch_size + num_chunks = tilelang.cdiv(num_tokens, chunk_size) if output_h else 0 + cu_seqlens = torch.empty((batch_size + 1), dtype=torch.int32, device=k.device) + chunk_offsets = torch.empty( + (batch_size + 1), dtype=torch.int32, device=k.device + ) + seqlen_dtype = torch.int32 + is_varlen = False + else: + real_batch_size = len(cu_seqlens) - 1 + chunk_offsets, num_chunks = prepare_chunk_offsets(cu_seqlens, chunk_size) + chunk_offsets = chunk_offsets.to(cu_seqlens.dtype) + num_chunks = num_chunks if output_h else 0 + seqlen_dtype = cu_seqlens.dtype + is_varlen = True + + if cp_seq_map is None: + cp_seq_map = torch.empty( + (real_batch_size,), dtype=seqlen_dtype, device=k.device + ) + is_cp = False + else: + is_cp = True + + use_initial_state = initial_state is not None + if initial_state is None: + initial_state = torch.empty( + (real_batch_size, H, K, V), dtype=torch.float32, device=k.device + ) + h = torch.empty((batch_size, num_chunks, H, K, V), dtype=k.dtype, device=k.device) + if raw_cu_seqlens is None: + raw_cu_seqlens = torch.empty( + (real_batch_size + 1,), dtype=seqlen_dtype, device=k.device + ) + final_state = torch.empty( + (real_batch_size, H, K, V), dtype=torch.float32, device=k.device + ) + else: + final_state = torch.empty( + (raw_cu_seqlens.shape[0] - 1, H, K, V), dtype=torch.float32, device=k.device + ) + o = torch.empty_like(v) + + grid_size = real_batch_size * H + block_dv_override = os.environ.get( + "TILEOPS_GDN_PREFILL_BLOCK_DV", + os.environ.get("TILEOPS_GDN_PREFILL_CP_BLOCK_DV"), + ) + if block_dv_override: + block_DV = int(block_dv_override) + if block_DV not in (32, 64, 128) or V % block_DV != 0: + raise ValueError( + "TILEOPS_GDN_PREFILL_BLOCK_DV must be one of 32, 64, 128 " + "and divide DV" + ) + elif grid_size >= TARGET_NUM_CTAS: + block_DV = 128 + elif grid_size * 2 >= TARGET_NUM_CTAS: + block_DV = 64 + else: + block_DV = 32 + + tilelang_fused_chunk_gdr_fwd_kernel = tilelang_fused_chunk_gdr_fwd( + H, + Hg, + K, + V, + chunk_size, + scale, + qkva_dtype=q.dtype, + g_dtype=g.dtype, + b_dtype=b.dtype, + h0_dtype=initial_state.dtype, + ht_dtype=final_state.dtype, + h_dtype=h.dtype, + o_dtype=o.dtype, + seqlen_dtype=seqlen_dtype, + accum_dtype="float32", + use_initial_state=use_initial_state, + store_final_state=output_final_state, + store_h=output_h, + store_o=output_o, + is_varlen=is_varlen, + is_cp=is_cp, + block_DV=block_DV, + ) + tilelang_fused_chunk_gdr_fwd_kernel( + q, + k, + v, + a, + g, + b, + initial_state, + cu_seqlens, + chunk_offsets, + cp_seq_map, + raw_cu_seqlens, + o, + h, + final_state, + ) + + if not output_final_state: + final_state = None + if not output_h: + h = None + if not output_o: + o = None + + return o, h, final_state diff --git a/tileops/kernels/gated_deltanet/gdn_prefill/prepare_h.py b/tileops/kernels/gated_deltanet/gdn_prefill/prepare_h.py new file mode 100644 index 000000000..7c7439684 --- /dev/null +++ b/tileops/kernels/gated_deltanet/gdn_prefill/prepare_h.py @@ -0,0 +1,563 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under the MIT License; see THIRD_PARTY_NOTICES.md for details. +# Adapted and modified for TileOps GatedDeltaNet prefill integration. + +import tilelang +import tilelang.language as T +import torch + +from .utils import prepare_chunk_offsets + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + compile_flags=["-O3", "-DENABLE_BF16", "-include", "tl_templates/cuda/gemm.h"], +) +def tilelang_prepare_h( + H, + Hg, + DK, + DV, + chunk_size, + accum_dtype, + qkva_dtype, + g_dtype, + b_dtype, + h0_dtype, + ht_dtype, + h_dtype, + seqlen_dtype, + use_initial_state, + store_final_state, + store_h, + is_varlen, + is_cp, + num_stages=2, +): + batch_size = T.dynamic("batch_size") + num_tokens = T.dynamic("num_tokens") + num_chunks = T.dynamic("num_chunks") + block_S = chunk_size + + if is_varlen: + k_shape = (1, num_tokens, Hg, DK) + v_shape = (1, num_tokens, H, DV) + a_shape = (1, num_tokens, H, chunk_size) + g_shape = (1, num_tokens, H) + b_shape = (1, num_tokens, H) + h_shape = (1, num_chunks, H, DK, DV) + else: + k_shape = (batch_size, num_tokens, Hg, DK) + v_shape = (batch_size, num_tokens, H, DV) + a_shape = (batch_size, num_tokens, H, chunk_size) + g_shape = (batch_size, num_tokens, H) + b_shape = (batch_size, num_tokens, H) + h_shape = (batch_size, num_chunks, H, DK, DV) + h0_shape = (batch_size, H, DK, DV) + ht_shape = (batch_size, H, DK, DV) + m_shape = (batch_size, H, DK, DK) + + @T.prim_func + def tilelang_prepare_h_kernel( + k: T.Tensor(k_shape, dtype=qkva_dtype), + v: T.Tensor(v_shape, dtype=qkva_dtype), + a: T.Tensor(a_shape, dtype=qkva_dtype), + g: T.Tensor(g_shape, dtype=g_dtype), + b: T.Tensor(b_shape, dtype=b_dtype), + h0: T.Tensor(h0_shape, dtype=h0_dtype), + cu_seqlens: T.Tensor([batch_size + 1], dtype=seqlen_dtype), + chunk_offsets: T.Tensor([batch_size + 1], dtype=seqlen_dtype), + num_warmup_chunks: T.Tensor([batch_size, H], dtype=seqlen_dtype), + h: T.Tensor(h_shape, dtype=h_dtype), + ht: T.Tensor(ht_shape, dtype=ht_dtype), + mt: T.Tensor(m_shape, dtype=ht_dtype), + ): + with T.Kernel(batch_size * H, threads=512) as (bbh,): + bb, bh = bbh // H, bbh % H + bhg = bh // (H // Hg) + + batch_idx = T.alloc_var("int32") + seq_start_idx = T.alloc_var("int32") + seq_end_idx = T.alloc_var("int32") + _seq_split_idx = T.alloc_var("int32") + chunk_start_idx = T.alloc_var("int32") + _chunk_split_idx = T.alloc_var("int32") + + batch_idx = 0 if is_varlen else bb + seq_start_idx = cu_seqlens[bb] if is_varlen else 0 + seq_end_idx = cu_seqlens[bb + 1] if is_varlen else num_tokens + chunk_start_idx = chunk_offsets[bb] if is_varlen else 0 + + num_iters = T.alloc_var("int32") + num_iters = ( + num_warmup_chunks[bb, bh] + if is_cp + else T.ceildiv(seq_end_idx - seq_start_idx, block_S) + ) + + calc_mt = T.alloc_var("bool") + calc_mt = is_cp and num_iters >= T.ceildiv( + seq_end_idx - seq_start_idx, block_S + ) + seq_start_idx = ( + seq_end_idx - num_iters * block_S if is_cp else seq_start_idx + ) + + k_shared = T.alloc_shared((num_stages, block_S, DK), dtype=qkva_dtype) + v_shared = T.alloc_shared((num_stages, block_S, DV), dtype=qkva_dtype) + a_shared = T.alloc_shared((num_stages, block_S, block_S), dtype=qkva_dtype) + g_shared = T.alloc_shared( + (num_stages, block_S), dtype=accum_dtype, scope="shared" + ) + b_shared = T.alloc_shared( + (num_stages, block_S), dtype=accum_dtype, scope="shared" + ) + h_shared = T.alloc_shared((DK, DV), dtype=qkva_dtype) + x_shared = T.alloc_shared((block_S, DK), dtype=qkva_dtype) + y_shared = T.alloc_shared((block_S, DV), dtype=qkva_dtype) + m_shared_L = T.alloc_shared((DK, DK // 2), dtype=qkva_dtype) + m_shared_R = T.alloc_shared((DK, DK // 2), dtype=qkva_dtype) + z_shared_L = T.alloc_shared((block_S, DK // 2), dtype=qkva_dtype) + z_shared_R = T.alloc_shared((block_S, DK // 2), dtype=qkva_dtype) + g_rev_exp_shared = T.alloc_shared( + (block_S), dtype=accum_dtype, scope="shared" + ) + + h_fragment = T.alloc_fragment((DK, DV), dtype=accum_dtype) + x_fragment = T.alloc_fragment((block_S, DK), dtype=accum_dtype) + y_fragment = T.alloc_fragment((block_S, DV), dtype=accum_dtype) + m_fragment_L = T.alloc_fragment((DK, DK // 2), dtype=accum_dtype) + m_fragment_R = T.alloc_fragment((DK, DK // 2), dtype=accum_dtype) + z_fragment_L = T.alloc_fragment((block_S, DK // 2), dtype=accum_dtype) + z_fragment_R = T.alloc_fragment((block_S, DK // 2), dtype=accum_dtype) + g_last_local_S = T.alloc_local((1), dtype=accum_dtype) + g_last_local_X = T.alloc_local((1), dtype=accum_dtype) + g_last_local_Y = T.alloc_local((1), dtype=accum_dtype) + g_prod_X = T.alloc_fragment((1), dtype=accum_dtype) + g_prod_Y = T.alloc_fragment((1), dtype=accum_dtype) + + data_is_ready = T.alloc_barrier(arrive_count=[96] * num_stages) + data_is_free = T.alloc_barrier(arrive_count=[384] * num_stages) + + bar_0 = T.alloc_barrier(arrive_count=416) + bar_1 = T.alloc_barrier(arrive_count=256) + bar_2 = T.alloc_barrier(arrive_count=384) + bar_3 = T.alloc_barrier(arrive_count=128) + + T.use_swizzle(10) + + tx = T.get_thread_binding() + + PRODUCER_NREG = 24 + CONSUMER_S_NREG = 168 + CONSUMER_X_NREG = 160 + CONSUMER_Y_NREG = 160 + + if tx < 128: + T.set_max_nreg(CONSUMER_S_NREG, 1) + + # Initialize S + if use_initial_state: + T.copy(h0[bb, bh, 0:DK, 0:DV], h_fragment) + else: + T.clear(h_fragment) + + # Main Loop + for i_s in T.serial(num_iters): + # [STAGE = i_s % num_stages] + T.barrier_wait( + data_is_ready[i_s % num_stages], (i_s // num_stages + 0) % 2 + ) + T.barrier_arrive(bar_0) + + # [STAGE = i_s % num_stages] 0 + T.barrier_wait(bar_0, i_s % 2) + # S4[1] S + T.copy(h_fragment, h_shared) + T.barrier_arrive(bar_1) + + # [STAGE = i_s % num_stages] 1 + T.barrier_wait(bar_1, i_s % 2) + # S = g_last * S + g_last_local_S[0] = T.exp2( + g_shared[i_s % num_stages, block_S - 1] * 1.442695 + ) + for j_k, j_v in T.Parallel(DK, DV): + h_fragment[j_k, j_v] *= g_last_local_S[0] + T.barrier_arrive(bar_2) + + # [STAGE = i_s % num_stages] 2 + T.barrier_wait(bar_2, i_s % 2) + # S += X^T @ Y + T.gemm_v1( + x_shared, + y_shared, + h_fragment, + transpose_A=True, + clear_accum=False, + ) + T.barrier_arrive(bar_3) + + T.barrier_arrive(data_is_free[i_s % num_stages]) + + # Store final S + if store_final_state: + T.copy(h_fragment, ht[bb, bh, 0:DK, 0:DV]) + + elif tx < 256: + T.set_max_nreg(CONSUMER_X_NREG, 1) + + if calc_mt: + for j_k, j_v in T.Parallel(DK, DK // 2): + if j_k == j_v + DK // 2: + m_fragment_R[j_k, j_v] = 1 + else: + m_fragment_R[j_k, j_v] = 0 + g_prod_X[0] = 0 + + # Main Loop + for i_s in T.serial(num_iters): + # [STAGE = i_s % num_stages] + T.barrier_wait( + data_is_ready[i_s % num_stages], (i_s // num_stages + 0) % 2 + ) + T.barrier_arrive(bar_0) + + # [STAGE = i_s % num_stages] 0 + T.barrier_wait(bar_0, i_s % 2) + # X = A^T @ K + T.gemm_v1( + a_shared[i_s % num_stages, :, :], + k_shared[i_s % num_stages, :, :], + x_fragment, + transpose_A=True, + clear_accum=True, + ) + + # [STAGE = i_s % num_stages] 1 + # X = - b * X + for j_s, j_k in T.Parallel(block_S, DK): + x_fragment[j_s, j_k] *= -b_shared[i_s % num_stages, j_s] + # S2[1] X + T.copy(x_fragment, x_shared) + T.barrier_arrive(bar_2) + + if calc_mt: + # [STAGE = i_s % num_stages] 2 + g_prod_X[0] += g_shared[i_s % num_stages, block_S - 1] + # S4[2] M + T.copy(m_fragment_R, m_shared_R) + + # [STAGE = i_s % num_stages] 3 + T.barrier_wait(bar_3, i_s % 2) + # Z = K @ M + T.gemm_v1( + k_shared[i_s % num_stages, :, :], + m_shared_R, + z_fragment_R, + clear_accum=True, + ) + # S4[2] Z + T.copy(z_fragment_R, z_shared_R) + # M += X^T @ Z + T.gemm_v1( + x_shared, + z_shared_R, + m_fragment_R, + transpose_A=True, + clear_accum=False, + ) + + T.barrier_arrive(data_is_free[i_s % num_stages]) + + if calc_mt: + g_last_local_X[0] = T.exp2(g_prod_X[0] * 1.442695) + for j_k, j_v in T.Parallel(DK, DK // 2): + m_fragment_R[j_k, j_v] *= g_last_local_X[0] + T.copy(m_fragment_R, mt[bb, bh, 0:DK, DK // 2 :]) + + elif tx < 384: + T.set_max_nreg(CONSUMER_Y_NREG, 1) + + if calc_mt: + for j_k, j_v in T.Parallel(DK, DK // 2): + if j_k == j_v: + m_fragment_L[j_k, j_v] = 1 + else: + m_fragment_L[j_k, j_v] = 0 + g_prod_Y[0] = 0 + + # Main Loop + for i_s in T.serial(num_iters): + # [STAGE = i_s % num_stages] + T.barrier_wait( + data_is_ready[i_s % num_stages], (i_s // num_stages + 0) % 2 + ) + T.barrier_arrive(bar_0) + + # [STAGE = i_s % num_stages] 0 + T.barrier_wait(bar_0, i_s % 2) + # Precompute g_last/g + g_last_local_Y[0] = g_shared[i_s % num_stages, block_S - 1] + for j_s in T.Parallel(block_S): + g_rev_exp_shared[j_s] = T.exp2( + (g_last_local_Y[0] - g_shared[i_s % num_stages, j_s]) + * 1.442695 + ) + g_last_local_Y[0] = T.exp2(g_last_local_Y[0] * 1.442695) + T.barrier_arrive(bar_1) + + # [STAGE = i_s % num_stages] 1 + T.barrier_wait(bar_1, i_s % 2) + # U = K @ S + T.gemm_v1( + k_shared[i_s % num_stages, :, :], + h_shared, + y_fragment, + clear_accum=True, + ) + # Y = g_last * U - g_last/g * V + for j_s, j_v in T.Parallel(block_S, DV): + y_fragment[j_s, j_v] *= g_last_local_Y[0] + for j_s, j_v in T.Parallel(block_S, DV): + y_fragment[j_s, j_v] -= ( + v_shared[i_s % num_stages, j_s, j_v] * g_rev_exp_shared[j_s] + ) + # S2[2] Y + T.copy(y_fragment, y_shared) + T.barrier_arrive(bar_2) + + if calc_mt: + # [STAGE = i_s % num_stages] 2 + g_prod_Y[0] += g_shared[i_s % num_stages, block_S - 1] + # S4[2] M + T.copy(m_fragment_L, m_shared_L) + + # [STAGE = i_s % num_stages] 3 + T.barrier_wait(bar_3, i_s % 2) + # Z = K @ M + T.gemm_v1( + k_shared[i_s % num_stages, :, :], + m_shared_L, + z_fragment_L, + clear_accum=True, + ) + # S4[2] Z + T.copy(z_fragment_L, z_shared_L) + # M += X^T @ Z + T.gemm_v1( + x_shared, + z_shared_L, + m_fragment_L, + transpose_A=True, + clear_accum=False, + ) + + T.barrier_arrive(data_is_free[i_s % num_stages]) + + if calc_mt: + g_last_local_Y[0] = T.exp2(g_prod_Y[0] * 1.442695) + for j_k, j_v in T.Parallel(DK, DK // 2): + m_fragment_L[j_k, j_v] *= g_last_local_Y[0] + T.copy(m_fragment_L, mt[bb, bh, 0:DK, : DK // 2]) + + else: + T.set_max_nreg(PRODUCER_NREG, 0) + + if tx < 384 + 32: + for i_s in T.serial(num_iters): + T.barrier_wait( + data_is_free[i_s % num_stages], (i_s // num_stages + 1) % 2 + ) + left = seq_start_idx + i_s * block_S + right = left + block_S + + # Load K + T.tma_copy( + k[batch_idx, left:right, bhg, 0:DK], + k_shared[i_s % num_stages, :, :], + barrier=data_is_ready[i_s % num_stages], + ) + + T.barrier_arrive(data_is_ready[i_s % num_stages]) + + elif tx < 384 + 64: + for i_s in T.serial(num_iters): + T.barrier_wait( + data_is_free[i_s % num_stages], (i_s // num_stages + 1) % 2 + ) + left = seq_start_idx + i_s * block_S + right = left + block_S + + # Load V + T.tma_copy( + v[batch_idx, left:right, bh, 0:DV], + v_shared[i_s % num_stages, :, :], + barrier=data_is_ready[i_s % num_stages], + ) + # Load A TODO: Mask A for the last chunk + T.tma_copy( + a[batch_idx, left:right, bh, 0:block_S], + a_shared[i_s % num_stages, :, :], + barrier=data_is_ready[i_s % num_stages], + ) + + T.barrier_arrive(data_is_ready[i_s % num_stages]) + + elif tx < 384 + 96: + for i_s in T.serial(num_iters): + T.barrier_wait( + data_is_free[i_s % num_stages], (i_s // num_stages + 1) % 2 + ) + left = seq_start_idx + i_s * block_S + right = left + block_S + + # Load gamma + if right <= seq_end_idx: + for j_s in T.Parallel(block_S): + g_shared[i_s % num_stages, j_s] = g[ + batch_idx, left + j_s, bh + ] + else: + for j_s in T.Parallel(block_S): + if left + j_s < seq_end_idx: + g_shared[i_s % num_stages, j_s] = g[ + batch_idx, left + j_s, bh + ] + else: + g_shared[i_s % num_stages, j_s] = g[ + batch_idx, seq_end_idx - 1, bh + ] + # Load beta + if right <= seq_end_idx: + for j_s in T.Parallel(block_S): + b_shared[i_s % num_stages, j_s] = b[ + batch_idx, left + j_s, bh + ] + else: + for j_s in T.Parallel(block_S): + if left + j_s < seq_end_idx: + b_shared[i_s % num_stages, j_s] = b[ + batch_idx, left + j_s, bh + ] + else: + b_shared[i_s % num_stages, j_s] = 0 + + T.barrier_arrive(data_is_ready[i_s % num_stages]) + + else: + for i_s in T.serial(num_iters): + T.barrier_arrive(bar_0) + + T.barrier_wait(bar_0, i_s % 2) + T.barrier_wait(bar_1, i_s % 2) + # Store S + if store_h: + T.copy( + h_shared, + h[batch_idx, chunk_start_idx + i_s, bh, 0:DK, 0:DV], + ) + + return tilelang_prepare_h_kernel + + +def fused_gdr_h( + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + g: torch.Tensor, + b: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = True, + output_h: bool = True, + chunk_size: int = 64, + cu_seqlens: torch.LongTensor | None = None, + num_warmup_chunks: torch.LongTensor | None = None, +): + batch_size, num_tokens, Hg, K = k.shape + _, _, H, V = v.shape + assert K == V == 128 + assert chunk_size == 64 + + if cu_seqlens is None: + assert num_warmup_chunks is None + real_batch_size = batch_size + num_chunks = tilelang.cdiv(num_tokens, chunk_size) if output_h else 0 + cu_seqlens = torch.empty((batch_size + 1), dtype=torch.int32, device=k.device) + chunk_offsets = torch.empty( + (batch_size + 1), dtype=torch.int32, device=k.device + ) + is_varlen = False + is_cp = False + else: + real_batch_size = len(cu_seqlens) - 1 + chunk_offsets, num_chunks = prepare_chunk_offsets(cu_seqlens, chunk_size) + chunk_offsets = chunk_offsets.to(cu_seqlens.dtype) + num_chunks = num_chunks if output_h else 0 + is_varlen = True + if num_warmup_chunks is None: + num_warmup_chunks = torch.empty( + (real_batch_size, H), dtype=cu_seqlens.dtype, device=k.device + ) + is_cp = False + else: + is_cp = True + + use_initial_state = initial_state is not None + if initial_state is None: + initial_state = torch.empty( + (real_batch_size, H, K, V), dtype=torch.float32, device=k.device + ) + h = torch.empty((batch_size, num_chunks, H, K, V), dtype=k.dtype, device=k.device) + ht_dtype = k.dtype if is_cp else torch.float32 + final_state = torch.empty( + (real_batch_size, H, K, V), dtype=ht_dtype, device=k.device + ) + final_correction = torch.empty( + (real_batch_size, H, K, K), dtype=ht_dtype, device=k.device + ) + + tilelang_prepare_h_kernel = tilelang_prepare_h( + H, + Hg, + K, + V, + chunk_size, + qkva_dtype=k.dtype, + g_dtype=g.dtype, + b_dtype=b.dtype, + h0_dtype=initial_state.dtype, + ht_dtype=final_state.dtype, + h_dtype=h.dtype, + seqlen_dtype=cu_seqlens.dtype, + accum_dtype="float32", + use_initial_state=use_initial_state, + store_final_state=output_final_state, + store_h=output_h, + is_varlen=is_varlen, + is_cp=is_cp, + ) + tilelang_prepare_h_kernel( + k, + v, + a, + g, + b, + initial_state, + cu_seqlens, + chunk_offsets, + num_warmup_chunks, + h, + final_state, + final_correction, + ) + + if not output_final_state: + final_state = None + final_correction = None + if not output_h: + h = None + + return h, final_state, final_correction diff --git a/tileops/kernels/gated_deltanet/gdn_prefill/tilelang_compat.py b/tileops/kernels/gated_deltanet/gdn_prefill/tilelang_compat.py new file mode 100644 index 000000000..c0f3ccf13 --- /dev/null +++ b/tileops/kernels/gated_deltanet/gdn_prefill/tilelang_compat.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under the MIT License; see THIRD_PARTY_NOTICES.md for details. +# Adapted and modified for TileOps GatedDeltaNet prefill integration. + +from __future__ import annotations + +import os +from typing import Any + +import tilelang.language as T + + +def _shape_dim(buf: Any, idx: int) -> int: + region = getattr(buf, "region", None) + if region is not None: + extents = [int(r.extent) for r in region] + while len(extents) > 2 and extents[0] == 1: + extents.pop(0) + return extents[idx] + regions = getattr(buf, "regions", None) + if regions is not None: + extents = [int(r.extent) for r in regions] + while len(extents) > 2 and extents[0] == 1: + extents.pop(0) + return extents[idx] + return int(buf.shape[idx]) + + +def _dtype_of(buf: Any) -> str: + dtype = getattr(buf, "dtype", None) + if dtype is not None: + return str(dtype) + inner = getattr(buf, "buffer", None) + if inner is not None: + inner_dtype = getattr(inner, "dtype", None) + if inner_dtype is not None: + return str(inner_dtype) + return "" + + +def _read_ptr(buf: Any): + if hasattr(buf, "access_ptr"): + return buf.access_ptr("r") + inner = getattr(buf, "buffer", None) + region = getattr(buf, "region", None) + if inner is not None and region is not None: + return T.address_of(inner[tuple(r.min for r in region)]) + return T.address_of(buf[0, 0]) + + +@T.macro +def _gemm_ss_compat( + a, + b, + c, + m: int, + n: int, + k: int, + *, + transpose_a: bool = False, + transpose_b: bool = False, + clear_accum: bool = False, + lda: int = 0, + ldb: int = 0, +): + if lda == 0: + lda = m if transpose_a else k + if ldb == 0: + ldb = k if transpose_b else n + name = ( + f"tl::gemm_ss<{m}, {n}, {k}, 4, 1, " + f"{int(transpose_a)}, {int(transpose_b)}, {int(clear_accum)}, " + f"{lda}, {ldb}, 0, 0, true>" + ) + T.sync_threads() + T.fence_proxy_async() + T.call_extern("handle", name, _read_ptr(a), _read_ptr(b), c.data) + + +@T.macro +def _wgmma_gemm_sync_compat( + a, + b, + c, + *, + transpose_a: bool = False, + transpose_b: bool = False, + policy: Any = None, + clear_accum: bool = False, + num_regs: int = 1, +): + if policy is None: + policy = T.GemmWarpPolicy.Square + T.wgmma_gemm( + a, + b, + c, + transpose_A=transpose_a, + transpose_B=transpose_b, + policy=policy, + clear_accum=clear_accum, + ) + T.wait_wgmma(0) + T.warpgroup_fence_operand(c, num_regs=num_regs) + + +def install_gemm_v1_compat() -> None: + original_gemm = T.gemm + + def gemm_v1_compat( + a, + b, + c, + transpose_A: bool = False, + transpose_B: bool = False, + policy: Any = None, + clear_accum: bool = False, + k_pack: int = 1, + mbar: Any = None, + ): + del k_pack, mbar + m = _shape_dim(c, 0) + n = _shape_dim(c, 1) + k = _shape_dim(a, 0) if transpose_A else _shape_dim(a, 1) + a_dtype = _dtype_of(a) + b_dtype = _dtype_of(b) + mode = os.environ.get( + "TILEOPS_GDN_PREFILL_GEMM_V1_MODE", + os.environ.get("FLASHQLA_TL019_GEMM_V1_MODE", "default"), + ) + if mode == "wgmma" and a_dtype in ("float16", "bfloat16") and a_dtype == b_dtype: + return _wgmma_gemm_sync_compat( + a, + b, + c, + transpose_a=transpose_A, + transpose_b=transpose_B, + policy=policy, + clear_accum=clear_accum, + num_regs=max(1, (m * n) // 128), + ) + if ( + mode == "legacy" + and a_dtype == "float16" + and b_dtype == "float16" + and m % 64 == 0 + and n in (32, 64, 128) + and k >= 16 + and k % 16 == 0 + ): + return _gemm_ss_compat( + a, + b, + c, + m, + n, + k, + transpose_a=transpose_A, + transpose_b=transpose_B, + clear_accum=clear_accum, + ) + return original_gemm( + a, + b, + c, + transpose_A=transpose_A, + transpose_B=transpose_B, + clear_accum=clear_accum, + ) + + T.gemm_v1 = gemm_v1_compat # type: ignore[attr-defined] diff --git a/tileops/kernels/gated_deltanet/gdn_prefill/utils.py b/tileops/kernels/gated_deltanet/gdn_prefill/utils.py new file mode 100644 index 000000000..49138e11a --- /dev/null +++ b/tileops/kernels/gated_deltanet/gdn_prefill/utils.py @@ -0,0 +1,140 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Licensed under the MIT License; see THIRD_PARTY_NOTICES.md for details. +# Adapted and modified for TileOps GatedDeltaNet prefill integration. + +import functools +from collections import OrderedDict +from collections.abc import Callable +from typing import Any + +import tilelang +import tilelang.language as T +import torch + + +def tensor_cache( + fn: Callable[..., torch.Tensor], +) -> Callable[..., torch.Tensor]: + """ + A decorator that caches the most recent results of a function with tensor inputs. + + This decorator will store the output of the decorated function for the most recent set of input tensors. + The cache is limited to a fixed size (default is 256). When the cache is full, the oldest entry will be removed. + + Args: + fn (Callable[..., torch.Tensor]): + The function to be decorated. It should take tensor inputs and return tensor outputs. + + Returns: + Callable[..., torch.Tensor]: + A wrapped version of the input function with single-entry caching. + """ + + cache: "OrderedDict[tuple[tuple[int, ...], tuple[tuple[str, int], ...]], tuple[tuple[Any, ...], dict[str, Any], Any]]" = OrderedDict() + cache_size = 256 + + def get_id(x: Any): + if (type(x) is int) or (type(x) is float) or (type(x) is str): + return x + else: + return id(x) + + def make_identity_key( + args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> tuple[tuple[int, ...], tuple[tuple[str, int], ...]]: + args_key = tuple(get_id(a) for a in args) + kwargs_key = tuple(sorted((k, get_id(v)) for k, v in kwargs.items())) + return args_key, kwargs_key + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal cache, cache_size + key = make_identity_key(args, kwargs) + if key in cache: + cache.move_to_end(key, last=True) + _, _, cached_result = cache[key] + return cached_result + + result = fn(*args, **kwargs) + cache[key] = (args, kwargs, result) + cache.move_to_end(key, last=True) + if len(cache) > cache_size: + cache.popitem(last=False) + return result + + return wrapper + + +@tensor_cache +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return torch.diff(cu_seqlens) + + +@tensor_cache +def prepare_chunk_indices( + cu_seqlens: torch.LongTensor, + chunk_size: int, +) -> torch.LongTensor: + # TODO: tilelang kernel + indices = torch.cat( + [ + torch.arange(n) + for n in tilelang.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() + ] + ) + return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) + + +@tilelang.jit() +def tilelang_prepare_chunk_offsets( + chunk_size, + block_size, + dtype, +): + batch_size_plus_1 = T.dynamic("batch_size_plus_1") + num_threads = min(max(block_size, 32), 128) + + @T.prim_func + def tilelang_prepare_chunk_offsets_kernel( + cu_seqlens: T.Tensor([batch_size_plus_1], dtype=dtype), + chunk_offsets: T.Tensor([batch_size_plus_1], dtype=dtype), + ): + with T.Kernel(1, threads=num_threads) as (bb,): + _batch_size = T.alloc_var("int32") + _batch_size = batch_size_plus_1 - 1 + + seqlen_start_fragment = T.alloc_fragment((block_size), dtype=dtype) + seqlen_end_fragment = T.alloc_fragment((block_size), dtype=dtype) + chunk_offset_fragment = T.alloc_fragment((block_size), dtype=dtype) + + T.copy(cu_seqlens[: batch_size_plus_1 - 1], seqlen_start_fragment) + T.copy(cu_seqlens[1:], seqlen_end_fragment) + + for i in T.Parallel(block_size): + chunk_offset_fragment[i] = ( + seqlen_end_fragment[i] - seqlen_start_fragment[i] + ) + chunk_offset_fragment[i] = ( + chunk_offset_fragment[i] + chunk_size - 1 + ) // chunk_size + T.cumsum(src=chunk_offset_fragment, dim=0) + + chunk_offsets[0] = 0 + T.copy(chunk_offset_fragment, chunk_offsets[1:]) + + return tilelang_prepare_chunk_offsets_kernel + + +@tensor_cache +def prepare_chunk_offsets( + cu_seqlens: torch.LongTensor, + chunk_size: int, +) -> torch.LongTensor: + chunk_offsets = torch.empty_like(cu_seqlens) + tilelang_prepare_chunk_offsets_kernel = tilelang_prepare_chunk_offsets( + chunk_size=chunk_size, + block_size=tilelang.next_power_of_2(cu_seqlens.shape[0] - 1), + dtype=cu_seqlens.dtype, + ) + tilelang_prepare_chunk_offsets_kernel(cu_seqlens, chunk_offsets) + return chunk_offsets, chunk_offsets[-1].item() diff --git a/tileops/manifest/linear_attention.yaml b/tileops/manifest/linear_attention.yaml index d00233ef1..87cef890c 100644 --- a/tileops/manifest/linear_attention.yaml +++ b/tileops/manifest/linear_attention.yaml @@ -10,21 +10,21 @@ GatedDeltaNetPrefillFwdOp: # consumed by decode. ref_api: "none" family: linear_attention - status: spec-only + status: implemented signature: inputs: - q: {dtype: "float16 | bfloat16 | float32", shape: "[B, H, S, DK]"} - k: {dtype: "same_as(q)", shape: "[B, H, S, DK]"} - v: {dtype: "same_as(q)", shape: "[B, H, S, DV]"} - g: {dtype: "same_as(q)", shape: "[B, H, S]"} - beta: {dtype: "same_as(q)", shape: "[B, H, S]"} + q: {dtype: "float16 | bfloat16 | float32", shape: "[B, S, H, DK]"} + k: {dtype: "same_as(q)", shape: "[B, S, H, DK]"} + v: {dtype: "same_as(q)", shape: "[B, S, H, DV]"} + g: {dtype: "same_as(q)", shape: "[B, S, H]"} + beta: {dtype: "same_as(q)", shape: "[B, S, H]"} outputs: - o: {dtype: "same_as(q)", shape: "[B, H, S, DV]"} + o: {dtype: "same_as(q)", shape: "[B, S, H, DV]"} final_state: {dtype: "same_as(q)", shape: "[B, H, DK, DV]"} params: chunk_size: {type: "int | None", default: null} - layout: {type: str, default: "bhtd"} + layout: {type: str, default: "bthd"} shape_rules: - "layout in ('bhtd', 'bhsd', 'bthd')" - "layout == 'bthd' or q.shape == (B, H, S, DK)" @@ -43,12 +43,11 @@ GatedDeltaNetPrefillFwdOp: - "chunk_size is None or S % chunk_size == 0" workloads: - # Inference-oriented BHSD rows. Unit tests should use smaller - # branch-covering shapes in tests/ops, not these benchmark workloads. - - {q_shape: [1, 16, 512, 128], k_shape: [1, 16, 512, 128], v_shape: [1, 16, 512, 128], g_shape: [1, 16, 512], beta_shape: [1, 16, 512], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b1-s512-h16-d128"} - - {q_shape: [4, 16, 2048, 64], k_shape: [4, 16, 2048, 64], v_shape: [4, 16, 2048, 128], g_shape: [4, 16, 2048], beta_shape: [4, 16, 2048], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b4-s2k-h16-dk64-dv128"} - - {q_shape: [1, 16, 2048, 128], k_shape: [1, 16, 2048, 128], v_shape: [1, 16, 2048, 128], g_shape: [1, 16, 2048], beta_shape: [1, 16, 2048], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b1-s2k-h16-d128"} - - {q_shape: [1, 16, 8192, 128], k_shape: [1, 16, 8192, 128], v_shape: [1, 16, 8192, 128], g_shape: [1, 16, 8192], beta_shape: [1, 16, 8192], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b1-s8k-h16-d128"} + # Serving-oriented BTHD rows, matching FLA/Qwen inference prefill. + - {q_shape: [1, 512, 16, 128], k_shape: [1, 512, 16, 128], v_shape: [1, 512, 16, 128], g_shape: [1, 512, 16], beta_shape: [1, 512, 16], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b1-s512-h16-d128"} + - {q_shape: [4, 2048, 16, 64], k_shape: [4, 2048, 16, 64], v_shape: [4, 2048, 16, 128], g_shape: [4, 2048, 16], beta_shape: [4, 2048, 16], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b4-s2k-h16-dk64-dv128"} + - {q_shape: [1, 2048, 16, 128], k_shape: [1, 2048, 16, 128], v_shape: [1, 2048, 16, 128], g_shape: [1, 2048, 16], beta_shape: [1, 2048, 16], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b1-s2k-h16-d128"} + - {q_shape: [1, 8192, 16, 128], k_shape: [1, 8192, 16, 128], v_shape: [1, 8192, 16, 128], g_shape: [1, 8192, 16], beta_shape: [1, 8192, 16], chunk_size: 64, dtypes: [float16, bfloat16], label: "gdn-prefill-b1-s8k-h16-d128"} roofline: func: "tileops.perf.formulas.gated_deltanet_prefill_fwd_roofline" diff --git a/tileops/ops/__init__.py b/tileops/ops/__init__.py index b56d88872..6060039e7 100644 --- a/tileops/ops/__init__.py +++ b/tileops/ops/__init__.py @@ -43,6 +43,7 @@ GatedDeltaNetDecodeOp, GatedDeltaNetFwdOp, GatedDeltaNetOp, + GatedDeltaNetPrefillFwdOp, ) from .gated_linear_attn import GLADecodeOp from .gemm import GemmOp @@ -138,6 +139,7 @@ "GatedDeltaNetDecodeOp", "GatedDeltaNetFwdOp", "GatedDeltaNetOp", + "GatedDeltaNetPrefillFwdOp", "GLABwdOp", "GLADecodeOp", "GLAFwdOp", diff --git a/tileops/ops/gated_deltanet.py b/tileops/ops/gated_deltanet.py index 309b09f91..9e4ae61a4 100644 --- a/tileops/ops/gated_deltanet.py +++ b/tileops/ops/gated_deltanet.py @@ -5,6 +5,7 @@ from tileops.kernels.gated_deltanet import ( GatedDeltaNetBwdKernel, GatedDeltaNetFwdKernel, + GatedDeltaNetPrefillFwdKernel, ) from tileops.kernels.gated_deltanet_recurrence import ( GatedDeltaNetDecodeFP32Kernel, @@ -16,7 +17,13 @@ from .op_base import Op -__all__ = ["GatedDeltaNetBwdOp", "GatedDeltaNetDecodeOp", "GatedDeltaNetFwdOp", "GatedDeltaNetOp"] +__all__ = [ + "GatedDeltaNetBwdOp", + "GatedDeltaNetDecodeOp", + "GatedDeltaNetFwdOp", + "GatedDeltaNetOp", + "GatedDeltaNetPrefillFwdOp", +] class GatedDeltaNetFwdOp(Op): @@ -121,6 +128,181 @@ def forward( return o, S, Aw, Au +class GatedDeltaNetPrefillFwdOp(Op): + """Gated DeltaNet inference prefill operator. + + This is the serving-oriented zero-state prefill interface: + ``(q, k, v, g, beta) -> (o, final_state)``. It intentionally does not + expose backward-only training artifacts such as ``Aw`` and ``Au``. + ``layout="bthd"`` follows the official FLA/Qwen convention + (``q/k/v/o [B, T, H, D]``, ``g/beta [B, T, H]``). ``layout="bhtd"`` + selects the TileOps head-major convention (``q/k/v/o [B, H, T, D]``, + ``g/beta [B, H, T]``). ``layout="bhsd"`` is accepted as a backward-compatible + alias for ``"bhtd"``. + When ``chunk_size`` is not specified, the op uses a small-stream serving + default: 128 for ``batch * heads <= 8`` when the sequence length allows it, + otherwise 64. + """ + + def __init__( + self, + batch: int, + heads: int, + seq_len: int, + dim_k: int, + dim_v: int, + chunk_size: Optional[int] = None, + dtype: torch.dtype = torch.float32, + kernel_map: Optional[Dict[str, Kernel]] = None, + tune: bool = False, + layout: str = "bthd", + ) -> None: + layout = self._normalize_layout(layout) + if chunk_size is None: + streams = batch * heads + chunk_size = 128 if streams <= 8 and seq_len % 128 == 0 else 64 + + self.batch = batch + self.heads = heads + self.seq_len = seq_len + self.dim_k = dim_k + self.dim_v = dim_v + self.chunk_size = chunk_size + self.dtype = dtype + self.layout = layout + + if seq_len % chunk_size != 0: + raise ValueError( + f"seq_len ({seq_len}) must be divisible by chunk_size ({chunk_size})" + ) + + self.dispatch_kernel(kernel_map) + + kernel_cls = self.kernel_map["GatedDeltaNetPrefillFwdKernel"] + kernel_dtype = Kernel.dtype_to_str(dtype) + self.kernel = kernel_cls( + batch, + heads, + seq_len, + chunk_size, + dim_k, + dim_v, + dtype=kernel_dtype, + layout=layout, + tune=tune, + ) + + @property + def default_kernel_map(self) -> Dict[str, Kernel]: + return { + "GatedDeltaNetPrefillFwdKernel": GatedDeltaNetPrefillFwdKernel, + } + + @staticmethod + def _normalize_layout(layout: str) -> str: + layout = layout.lower() + if layout == "bhsd": + return "bhtd" + if layout in ("bhtd", "bthd"): + return layout + raise ValueError(f"Unsupported layout: {layout}") + + def _infer_output_shapes( + self, + q_shape: tuple[int, ...], + k_shape: tuple[int, ...], + v_shape: tuple[int, ...], + g_shape: tuple[int, ...], + beta_shape: tuple[int, ...], + ) -> dict[str, tuple[int, ...]]: + del k_shape, g_shape, beta_shape + layout = self._normalize_layout(getattr(self, "layout", "bthd")) + if layout == "bthd": + return { + "o": (q_shape[0], q_shape[1], q_shape[2], v_shape[-1]), + "final_state": ( + q_shape[0], + q_shape[2], + q_shape[-1], + v_shape[-1], + ), + } + return { + "o": tuple(q_shape[:-1]) + (v_shape[-1],), + "final_state": ( + q_shape[0], + q_shape[1], + q_shape[-1], + v_shape[-1], + ), + } + + def _validate_dtypes( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> None: + if self.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise ValueError(f"Unsupported dtype: {self.dtype}") + for name, tensor in (("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)): + if tensor.dtype != self.dtype: + raise ValueError(f"{name}.dtype must be {self.dtype}, got {tensor.dtype}") + + def _validate_shapes( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> None: + if self.layout == "bthd": + q_shape = (self.batch, self.seq_len, self.heads, self.dim_k) + v_shape = (self.batch, self.seq_len, self.heads, self.dim_v) + gate_shape = (self.batch, self.seq_len, self.heads) + else: + q_shape = (self.batch, self.heads, self.seq_len, self.dim_k) + v_shape = (self.batch, self.heads, self.seq_len, self.dim_v) + gate_shape = (self.batch, self.heads, self.seq_len) + if tuple(q.shape) != q_shape: + raise ValueError(f"q must have shape {q_shape}, got {tuple(q.shape)}") + if tuple(k.shape) != q_shape: + raise ValueError(f"k must have shape {q_shape}, got {tuple(k.shape)}") + if tuple(v.shape) != v_shape: + raise ValueError(f"v must have shape {v_shape}, got {tuple(v.shape)}") + if tuple(g.shape) != gate_shape: + raise ValueError(f"g must have shape {gate_shape}, got {tuple(g.shape)}") + if tuple(beta.shape) != gate_shape: + raise ValueError(f"beta must have shape {gate_shape}, got {tuple(beta.shape)}") + if not all(tensor.is_cuda for tensor in (q, k, v, g, beta)): + raise ValueError("q, k, v, g, and beta must be CUDA tensors") + + def eval_roofline(self) -> tuple[int, int]: + from tileops.perf.formulas import gated_deltanet_prefill_fwd_roofline + + return gated_deltanet_prefill_fwd_roofline(self) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + g = g.contiguous() + beta = beta.contiguous() + self._validate_dtypes(q, k, v, g, beta) + self._validate_shapes(q, k, v, g, beta) + return self.kernel(q, k, v, g, beta) + + class GatedDeltaNetBwdOp(Op): """Gated DeltaNet backward operator. diff --git a/tileops/perf/formulas.py b/tileops/perf/formulas.py index 6c0f3afbd..81475cfc8 100644 --- a/tileops/perf/formulas.py +++ b/tileops/perf/formulas.py @@ -32,6 +32,7 @@ "eq_fwd_roofline", "floor_divide_fwd_roofline", "fused_moe_fwd_bytes", + "gated_deltanet_prefill_fwd_roofline", "ge_fwd_roofline", "gqa_bwd_roofline", "gqa_decode_paged_roofline", @@ -212,6 +213,59 @@ def gqa_prefill_fwd_roofline(op: Any | None = None, **kwargs: Any) -> tuple[int, return int(flops), int(nbytes) +def gated_deltanet_prefill_fwd_roofline(op: Any | None = None, **kwargs: Any) -> tuple[int, int]: + """Approximate roofline for Gated DeltaNet zero-state prefill. + + This models the dominant chunkwise matmul work in the current + implementation. It is intentionally conservative; the helper exists so the + manifest and benchmark share one explicit cost-model hook. + """ + data = _shape_or_attrs(op, kwargs) + if "q_shape" in data: + layout = str(data.get("layout", "bthd")).lower() + q_shape = data["q_shape"] + v_shape = data["v_shape"] + if layout == "bthd": + batch, seq_len, heads, dim_k = q_shape + _, v_seq_len, v_heads, dim_v = v_shape + elif layout in ("bhtd", "bhsd"): + batch, heads, seq_len, dim_k = q_shape + _, v_heads, v_seq_len, dim_v = v_shape + else: + raise ValueError(f"Unsupported GDN prefill layout: {layout}") + if v_seq_len != seq_len or v_heads != heads: + raise ValueError( + "GDN prefill q_shape and v_shape must share seq_len and heads" + ) + chunk_size = data.get("chunk_size", 64) or 64 + else: + batch, heads, seq_len, dim_k, dim_v, chunk_size = ( + data["batch"], + data["heads"], + data["seq_len"], + data["dim_k"], + data["dim_v"], + data["chunk_size"] or 64, + ) + elem_bytes = _dtype_itemsize(data.get("dtype", data.get("dtypes", "float16"))) + + num_chunks = seq_len // chunk_size + state_flops = 4 * batch * heads * num_chunks * chunk_size * dim_k * dim_v + intra_flops = 4 * batch * heads * num_chunks * chunk_size * chunk_size * ( + dim_k + dim_v + ) + flops = state_flops + intra_flops + + input_elems = ( + 3 * batch * heads * seq_len * dim_k + + batch * heads * seq_len * dim_v + + 2 * batch * heads * seq_len + ) + output_elems = batch * heads * seq_len * dim_v + batch * heads * dim_k * dim_v + nbytes = (input_elems + output_elems) * elem_bytes + return int(flops), int(nbytes) + + def gqa_prefill_fp8_tensor_core_roofline( op: Any | None = None, **kwargs: Any, diff --git a/workloads/gated_deltanet.py b/workloads/gated_deltanet.py index 34ee48808..1200802bc 100644 --- a/workloads/gated_deltanet.py +++ b/workloads/gated_deltanet.py @@ -33,6 +33,23 @@ def gen_inputs(self) -> tuple[torch.Tensor, ...]: return q, k, v, g, beta +class GatedDeltaNetPrefillFwdTest(GatedDeltaNetFwdTest): + """Inference prefill workload for Gated DeltaNet.""" + + def __init__( + self, + batch: int, + heads: int, + seq_len: int, + dim_k: int, + dim_v: int, + chunk_size: int, + dtype: torch.dtype, + ) -> None: + super().__init__(batch, heads, seq_len, dim_k, dim_v, chunk_size, dtype) + self.shape = (batch, heads, seq_len, dim_k) + + class GatedDeltaNetDecodeTest(WorkloadBase): def __init__(