diff --git a/ci/jax.sh b/ci/jax.sh index dbca249c9b..ae1330b8b5 100755 --- a/ci/jax.sh +++ b/ci/jax.sh @@ -67,6 +67,7 @@ run_test_config() { run_default_fa 1 test_layer.py # it effectively always uses unfused attention run_default_fa 1 test_sanity_import.py run_default_fa 1 test_softmax.py + run_default_fa 1 test_indexer.py # lightning indexer ops (fused-attn agnostic) } run_test_config_mgpu() { diff --git a/tests/jax/test_indexer.py b/tests/jax/test_indexer.py new file mode 100644 index 0000000000..64b469f4a0 --- /dev/null +++ b/tests/jax/test_indexer.py @@ -0,0 +1,162 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +"""Correctness tests for the lightning-indexer JAX ops. +""" + +import functools + +import jax +import jax.numpy as jnp +import pytest + +from transformer_engine.jax.sparse_attention.indexer import ( + LightningIndexer, + indexer, + indexer_topk, +) + + +@pytest.fixture(autouse=True) +def _disable_indexer_autotune(monkeypatch): + """Pin each indexer Triton kernel to a single (prune-valid) config so the + suite skips the multi-minute autotune sweep.""" + monkeypatch.setenv("NVTE_INDEXER_DISABLE_AUTOTUNE", "1") + + +@functools.partial(jax.jit, static_argnames=("out_dtype",)) +def _indexer_reference(Q, K, W_uq, W_dq, W_k, W_w, out_dtype=None): + """Pure-einsum lightning-indexer reference (test oracle). + + Materializes the (..., T, H, S) pre-relu score tensor, unlike the hybrid + Triton op under test. Shapes: Q [..., T, d], K [..., S, d], W_dq [d, d_c], + W_uq [H, d_c, d_i], W_k [d, d_i], W_w [d, H]. Returns O [..., T, S]. + + JIT-compiled so its HLO (reduction order / bf16 rounding) is stable — the + top-k test feeds these reference scores to ``jax.lax.top_k``, whose + tie-breaking is sensitive to sub-ULP score perturbations. + """ + C_q = jnp.einsum("...td,dc->...tc", Q, W_dq) + H_q = jnp.einsum("...tc,hci->...thi", C_q, W_uq) + H_k = jnp.einsum("...sd,di->...si", K, W_k) + W_o = jnp.einsum("...td,dh->...th", Q, W_w) + H = jax.nn.relu(jnp.einsum("...thi,...si->...ths", H_q, H_k)) # (..., T, H, S) + O = jnp.einsum("...ths,...th->...ts", H, W_o) # (..., T, S) + if out_dtype is not None: + O = O.astype(out_dtype) + return O + + +def _indexer_inputs(B, oH, T_t, T_s, d, d_c, H, d_i, seed): + keys = jax.random.split(jax.random.PRNGKey(seed), 6) + Q = jax.random.normal(keys[0], (B, oH, T_t, d), dtype=jnp.bfloat16) + K = jax.random.normal(keys[1], (B, oH, T_s, d), dtype=jnp.bfloat16) + W_uq = jax.random.normal(keys[2], (H, d_c, d_i), dtype=jnp.bfloat16) + W_dq = jax.random.normal(keys[3], (d, d_c), dtype=jnp.bfloat16) + W_k = jax.random.normal(keys[4], (d, d_i), dtype=jnp.bfloat16) + W_w = jax.random.normal(keys[5], (d, H), dtype=jnp.bfloat16) + return Q, K, W_uq, W_dq, W_k, W_w + + +def _rel_err(actual, ref): + actual = actual.astype(jnp.float32) + ref = ref.astype(jnp.float32) + return float(jnp.linalg.norm(actual - ref) / (jnp.linalg.norm(ref) + 1e-30)) + + +@pytest.mark.parametrize("B,oH", [(2, 3), (1, 1), (1, 4)]) +def test_hybrid_matches_reference(B, oH): + """Hybrid Triton score-reduce matches the pure-einsum reference forward.""" + args = _indexer_inputs(B, oH, T_t=64, T_s=64, d=32, d_c=32, H=8, d_i=32, seed=100) + o_ref = _indexer_reference(*args) + o_hyb = indexer(*args) + assert o_hyb.shape == o_ref.shape + assert _rel_err(o_hyb, o_ref) < 5e-3 + + +@pytest.mark.parametrize("k", [32, 64, 128, 256, 512, 1024]) +def test_topk_matches_reference(k): + """Fused top-k selects the same scores as reference + ``jax.lax.top_k``. + + Index set-equality is too strict (backends break ties differently), so the + check is on the *scores* at the fused-selected indices, compared rank-by-rank + against the reference top-k. ``T_s`` is ``4 * max(k)`` so the largest ``k`` + (1024) sits at the top quartile — deep enough to exercise the streaming + top-k path (2K candidate buffer) that small ``k`` / ``T_s`` never reaches. + + The gap is normalized by the overall score *scale* (max reference score), not + per element: as ``k`` grows into the near-zero ReLU tail, per-element relative + error is dominated by ties the fp32/bf16 paths break differently (denominators + ~0 blow it up), while the absolute gap stays ~0.1% of the max score. + + Leading dims (B, oH, T_t) are kept small so the reference — which materializes + the (B, oH, T_t, H, T_s) pre-relu score tensor — stays a few MB even at + T_s=4096; a larger footprint tips shared-GPU GEMMs into resource errors. + """ + B, oH, T_t = 1, 2, 32 + args = _indexer_inputs(B, oH, T_t, T_s=4096, d=32, d_c=32, H=16, d_i=32, seed=200) + o_ref = _indexer_reference(*args).astype(jnp.float32) + topk_idx = indexer_topk(*args, k=k) + assert topk_idx.shape == (B, oH, T_t, k) + + ref_vals = jax.lax.top_k(o_ref, k=k)[0] + scale = float(ref_vals.max()) + assert scale > 0, "degenerate test: all top-k scores are zero" + picked = jnp.take_along_axis(o_ref, topk_idx, axis=-1) + picked_sorted = jnp.sort(picked, axis=-1)[..., ::-1] + max_gap = float(jnp.abs(ref_vals - picked_sorted).max()) / scale + assert max_gap < 1e-2, f"fused top-k scores diverge: max_gap={max_gap:.3e} (k={k})" + + +@pytest.mark.parametrize("B,oH", [(2, 3), (1, 2)]) +def test_hybrid_backward_matches_reference_grad(B, oH): + """``jax.grad`` through the hybrid backend matches grad through reference. + + Tolerance is 5e-2 (bf16 projections + Triton score recompute) — looser than + the 5e-3 forward tolerance; tighten once per-grad error is characterized + on-device. + """ + args = _indexer_inputs(B, oH, T_t=32, T_s=32, d=32, d_c=32, H=8, d_i=32, seed=300) + + def _loss(fn): + def inner(*a): + return jnp.sum(fn(*a).astype(jnp.float32)) + return inner + + argnums = (0, 1, 2, 3, 4, 5) + grads_ref = jax.grad(_loss(_indexer_reference), argnums=argnums)(*args) + grads_hyb = jax.grad(_loss(indexer), argnums=argnums)(*args) + for gr, gh in zip(grads_ref, grads_hyb): + assert _rel_err(gh, gr) < 5e-2 + + +def test_lightning_indexer_module_matches_functional(): + """``LightningIndexer`` (Flax module) reproduces the functional ``indexer`` + when fed the module's own initialized weights.""" + B, oH, T_t, T_s, d, d_c, H, d_i = 2, 3, 64, 64, 32, 32, 8, 32 + keys = jax.random.split(jax.random.PRNGKey(7), 3) + Q = jax.random.normal(keys[0], (B, oH, T_t, d), dtype=jnp.bfloat16) + K = jax.random.normal(keys[1], (B, oH, T_s, d), dtype=jnp.bfloat16) + + mod = LightningIndexer(num_heads=H, d_c=d_c, d_i=d_i) + variables = mod.init(keys[2], Q, K) + o_mod = mod.apply(variables, Q, K) + assert o_mod.shape == (B, oH, T_t, T_s) + + p = variables["params"] + o_fn = indexer(Q, K, p["W_uq"], p["W_dq"], p["W_k"], p["W_w"]) + assert _rel_err(o_mod, o_fn) < 1e-5 + + +def test_lightning_indexer_topk_mode(): + """``LightningIndexer(topk=k)`` returns fused top-k indices of shape (..., T, k).""" + B, oH, T_t, T_s, d, d_c, H, d_i, k = 2, 3, 64, 128, 32, 32, 16, 32, 32 + keys = jax.random.split(jax.random.PRNGKey(9), 2) + Q = jax.random.normal(keys[0], (B, oH, T_t, d), dtype=jnp.bfloat16) + K = jax.random.normal(keys[1], (B, oH, T_s, d), dtype=jnp.bfloat16) + + mod = LightningIndexer(num_heads=H, d_c=d_c, d_i=d_i, topk=k) + variables = mod.init(jax.random.PRNGKey(0), Q, K) + idx = mod.apply(variables, Q, K) + assert idx.shape == (B, oH, T_t, k) + assert idx.dtype == jnp.int32 diff --git a/transformer_engine/jax/__init__.py b/transformer_engine/jax/__init__.py index d0afc1ff25..020373d596 100644 --- a/transformer_engine/jax/__init__.py +++ b/transformer_engine/jax/__init__.py @@ -34,6 +34,8 @@ from . import flax from . import quantize +from . import sparse_attention + from .quantize import autocast, fp8_autocast, update_collections from .quantize import NVTE_FP8_COLLECTION_NAME @@ -51,4 +53,5 @@ "MeshResource", "flax", "quantize", + "sparse_attention", ] diff --git a/transformer_engine/jax/sparse_attention/__init__.py b/transformer_engine/jax/sparse_attention/__init__.py new file mode 100644 index 0000000000..92dc876724 --- /dev/null +++ b/transformer_engine/jax/sparse_attention/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +"""Sparse attention primitives. + +Currently exposes the lightning indexer: + + * :mod:`~transformer_engine.jax.sparse_attention.indexer` — the lightning + indexer op (``indexer`` / ``indexer_topk``) and the :class:`LightningIndexer` + Flax module. + +The Triton kernel backends live in +:mod:`transformer_engine.jax.triton_extensions` alongside the other Triton +kernels. +""" + +from . import indexer + +from .indexer import LightningIndexer + +__all__ = [ + "indexer", + "LightningIndexer", +] diff --git a/transformer_engine/jax/sparse_attention/indexer.py b/transformer_engine/jax/sparse_attention/indexer.py new file mode 100644 index 0000000000..15fdc74592 --- /dev/null +++ b/transformer_engine/jax/sparse_attention/indexer.py @@ -0,0 +1,164 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +"""Indexer op (forward only), bf16 inputs. + +The op runs a hybrid backend: einsum projections (C_q, H_q, H_k, W_o) — +which lower to hipBLASLt bf16 GEMMs — followed by a fused Triton kernel that +does score+relu+H-reduction in registers. This avoids materializing the +(B, oH, T, H, S) pre-relu score tensor in HBM. + +Functional entry point: ``indexer(Q, K, W_uq, W_dq, W_k, W_w)``. +User-facing Flax module: :class:`LightningIndexer`, which owns the projection +weights and delegates to ``indexer`` / ``indexer_topk``. + +Math (low-rank form: Q is hidden state; query heads are produced by a +down-projection (d -> d_c) followed by an up-projection (d_c -> H * d_i); +output weights are produced from Q via a learnable d -> H projection): + + C_q = Q @ W_dq # (..., T, d_c) + H_q = einsum("...tc,hci->...thi", C_q, W_uq) # (..., T, H, d_i) + H_k = K @ W_k # (..., S, d_i) + W_o = Q @ W_w # (..., T, H) + H = relu(einsum("...thi,...si->...ths", H_q, H_k)) # (..., T, H, S) + O = einsum("...ths,...th->...ts", H, W_o) # (..., T, S) +""" + +import functools +from typing import Optional + +import jax +import jax.numpy as jnp +from flax import linen as nn + + +def _indexer_projections(Q, K, W_uq, W_dq, W_k, W_w): + """Low-rank indexer projections shared by the score and top-k paths. + + Returns (H_q, H_k, W_o) with shapes + (..., T, H, d_i), (..., S, d_i), (..., T, H). + """ + C_q = jnp.einsum("...td,dc->...tc", Q, W_dq) + H_q = jnp.einsum("...tc,hci->...thi", C_q, W_uq) + H_k = jnp.einsum("...sd,di->...si", K, W_k) + W_o = jnp.einsum("...td,dh->...th", Q, W_w) + return H_q, H_k, W_o + + +def _indexer_impl_hybrid(Q, K, W_uq, W_dq, W_k, W_w, out_dtype=None): + """Einsum projections + Triton score-relu-reduce. + + Runs the four projections (which lower to hipBLASLt bf16 GEMMs), then + hands Hq / Hk / W_o to a fused Triton kernel that does + score+relu+H-reduction in registers — eliminating the + (B, oH, T, H, S) pre-relu-score HBM round-trip a pure-einsum path pays. + """ + from transformer_engine.jax.triton_extensions.indexer import score_reduce_triton + + H_q, H_k, W_o = _indexer_projections(Q, K, W_uq, W_dq, W_k, W_w) + return score_reduce_triton(H_q, H_k, W_o, + out_dtype=out_dtype if out_dtype else Q.dtype) + + +@functools.partial(jax.jit, static_argnames=("k",)) +def indexer_topk(Q, K, W_uq, W_dq, W_k, weights, *, k): + """Lightning-indexer + top-k (fused). + + Same projections as ``indexer()`` (reference math), then a single Triton + kernel that computes the score row, ReLU, weighted H-reduction, and + streaming top-k all in one pass — the (B, oH, T_t, T_s) score matrix is + never materialized. + + Args: + Q, K, W_uq, W_dq, W_k, weights: same as ``indexer()``. + k: number of top scores to return per (B, oH, T_t) row. + Must be a power of 2 and <= S. + + Returns: + Topk_idx: (..., T_t, k) int32 — top-k indices into the S axis, + in descending score order. + """ + from transformer_engine.jax.triton_extensions.indexer import score_topk_triton + H_q, H_k, W_o = _indexer_projections(Q, K, W_uq, W_dq, W_k, weights) + return score_topk_triton(H_q, H_k, W_o, k=k) + + +@functools.partial(jax.jit, static_argnames=("out_dtype",)) +def indexer(Q, K, W_uq, W_dq, W_k, weights, *, out_dtype=None): + """Low-rank lightning-indexer (bf16), hybrid Triton backend. + + Args: + Q: (..., T, d) hidden state (per token) + K: (..., S, d) key hidden state + W_uq: (H, d_c, d_i) up-projection: d_c -> d_i (per head) + W_dq: (d, d_c) down-projection: d -> d_c + W_k: (d, d_i) key projection + weights: (d, H) learnable output-weight projection + (W_o = Q @ weights inside the impl) + out_dtype: output dtype override (defaults to Q.dtype). + + Returns: + O of shape (..., T, S). + """ + return _indexer_impl_hybrid(Q, K, W_uq, W_dq, W_k, weights, out_dtype=out_dtype) + + +class LightningIndexer(nn.Module): # pylint: disable=too-few-public-methods + """Lightning-indexer Flax module — the user-facing indexer API. + + Owns the low-rank indexer projection weights (``W_dq``, ``W_uq``, ``W_k``, + ``W_w``) and delegates to the functional :func:`indexer` / :func:`indexer_topk` + ops. Weight shapes mirror :func:`indexer`'s ``Args`` and are inferred from the + trailing hidden dimension ``d`` of ``Q`` at call time. + + Parameters + ---------- + num_heads : int + Number of indexer-internal heads (``H``). + d_c : int + Down-projection rank (``d -> d_c``). + d_i : int + Inner head dimension (``d_i``). + topk : Optional[int], default ``None`` + If set, :meth:`__call__` returns the fused top-``k`` indices + (``(..., T, k)`` int32) via :func:`indexer_topk`, and ``out_dtype`` is + ignored (top-k always uses the fused Triton kernel). + If ``None``, :meth:`__call__` returns the full score tensor + ``(..., T, S)`` (hybrid Triton backend). + out_dtype : Optional[jnp.dtype] + Output dtype override; defaults to ``Q.dtype``. Unused when ``topk`` is set. + dtype : Optional[jnp.dtype] + Parameter dtype. Defaults to the input dtype. + """ + + num_heads: int + d_c: int + d_i: int + topk: Optional[int] = None + out_dtype: Optional[jnp.dtype] = None + dtype: Optional[jnp.dtype] = None + + @nn.compact + def __call__(self, Q: jax.Array, K: jax.Array) -> jax.Array: + """Run the indexer on ``Q`` / ``K``. + + Args: + Q: ``(..., T, d)`` query-side hidden state. + K: ``(..., S, d)`` key-side hidden state. + + Returns: + ``(..., T, S)`` scores if ``topk is None``, else ``(..., T, k)`` + int32 top-k indices (in descending score order). + """ + d = Q.shape[-1] + param_dtype = self.dtype if self.dtype is not None else Q.dtype + init = nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal") + + W_dq = self.param("W_dq", init, (d, self.d_c), param_dtype) + W_uq = self.param("W_uq", init, (self.num_heads, self.d_c, self.d_i), param_dtype) + W_k = self.param("W_k", init, (d, self.d_i), param_dtype) + W_w = self.param("W_w", init, (d, self.num_heads), param_dtype) + + if self.topk is not None: + return indexer_topk(Q, K, W_uq, W_dq, W_k, W_w, k=self.topk) + return indexer(Q, K, W_uq, W_dq, W_k, W_w, out_dtype=self.out_dtype) diff --git a/transformer_engine/jax/triton_extensions/__init__.py b/transformer_engine/jax/triton_extensions/__init__.py index 150a5fbf12..68b869a029 100644 --- a/transformer_engine/jax/triton_extensions/__init__.py +++ b/transformer_engine/jax/triton_extensions/__init__.py @@ -61,3 +61,4 @@ def lowering(ctx, x, **kwargs): from .utils import * from .permutation import * +from .indexer import score_reduce_triton, score_topk_triton diff --git a/transformer_engine/jax/triton_extensions/indexer.py b/transformer_engine/jax/triton_extensions/indexer.py new file mode 100644 index 0000000000..f845bfc637 --- /dev/null +++ b/transformer_engine/jax/triton_extensions/indexer.py @@ -0,0 +1,1014 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +"""Triton score-relu-reduce kernel for the lightning-indexer hybrid backend. + +The hybrid backend computes the four projections (C_q, H_q, H_k, W_o) via +``jnp.einsum`` (which lowers to hipBLASLt bf16 GEMMs) and then hands the +results to this kernel for the score matmul + ReLU + per-(t, h) weighted +H-reduction: + + scores = relu(einsum("...thi,...si->...ths", H_q, H_k)) # never written + O = einsum("...ths,...th->...ts", scores, W_o) + +The kernel keeps each per-head score tile in registers, avoiding the +(B, oH, T, H, S) HBM round-trip that an einsum-only implementation pays +on the pre-relu score tensor. +""" + +import functools +import os + +import jax +import jax.numpy as jnp +import triton +import triton.language as tl + +from jax import core +from jax.extend import core as extend_core +from jax.interpreters import mlir, xla + +from .utils import triton_call_lowering + + +def _autotune_disabled(): + """True when ``NVTE_INDEXER_DISABLE_AUTOTUNE=1``. + + When set, each kernel's lowering collapses its autotune sweep to the first + (still prune-valid) config, so no time is spent compiling and benchmarking + every candidate. Intended for the test suite — a full sweep at large k/T_s + costs many minutes and only picks the fastest config, not a more correct + one. Read at lowering time so a test fixture can toggle it per process.""" + return os.environ.get("NVTE_INDEXER_DISABLE_AUTOTUNE", "0") == "1" + + +def _score_reduce_autotune_configs(): + # The kernel is dominated by Hq reads (one (BLOCK_T, d_i) load per H + # iteration). Bigger BLOCK_T ⇒ fewer T tiles ⇒ less total Hq traffic. + # Bigger BLOCK_S ⇒ more Hk reuse but bigger per-CTA footprint. + # + # BLOCK_T=512 was tried and consistently failed to launch on MI355X + # (resource exhaustion — VGPR/LDS budget for 64-iter H-loop with that + # large an accumulator). Capped at 256. + cfgs = [ + triton.Config({"BLOCK_T": 32, "BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_T": 32, "BLOCK_S": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_T": 256, "BLOCK_S": 32}, num_warps=8, num_stages=2), + ] + cfgs += [ + triton.Config({"BLOCK_T": bt, "BLOCK_S": bs, "matrix_instr_nonkdim": nk_dim, "waves_per_eu": wpe}, num_warps=nw, num_stages=ns) + for bt in (32, 64) + for bs in (256, 512) + for nk_dim in (16,) + for wpe in (0, 2) # 0 means let backend compiler decide + for nw in (4,) + for ns in (2,) + ] + return cfgs + + +@triton.autotune(configs=_score_reduce_autotune_configs(), key=["H", "d_i"]) +@triton.jit +def _score_reduce_kernel( + Hq_ptr, # (B, oH, T_t, H, d_i) — produced by einsum("...tc,hci->...thi") + Hk_ptr, # (B, oH, T_s, d_i) + W_o_ptr, # (B, oH, T_t, H) + O_ptr, # (B, oH, T_t, T_s) + B: tl.constexpr, + oH: tl.constexpr, + T_t: tl.constexpr, + T_s: tl.constexpr, + H: tl.constexpr, + d_i: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Compute one (BLOCK_T, BLOCK_S) tile of O for one (b, h_outer) slice. + + Grid order: (cdiv(T_s, BLOCK_S), cdiv(T_t, BLOCK_T), B * oH). + + S is the fastest-dispatching axis so consecutive CTAs share (B*oH, T) + and vary only in S — they all read the same per-head Hq slab, hitting + L2 instead of HBM. Hq layout is the natural einsum output + (..., T, H, d_i); per-head loads are strided in T (stride H*d_i). + """ + pid_s = tl.program_id(0) + pid_t = tl.program_id(1) + pid_bh = tl.program_id(2) + + # int64 indexing — Hq alone has B*oH*T*H*d_i = 4.3 B elements at T=S=4096, + # exceeds int32 range. + b = (pid_bh // oH).to(tl.int64) + h_outer = (pid_bh % oH).to(tl.int64) + + rt = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) + rs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + rdi = tl.arange(0, d_i) + + rt_mask = rt < T_t + rs_mask = rs < T_s + + hq_base = b * (oH * T_t * H * d_i) + h_outer * (T_t * H * d_i) + hk_base = b * (oH * T_s * d_i) + h_outer * (T_s * d_i) + wo_base = b * (oH * T_t * H) + h_outer * (T_t * H) + o_base = b * (oH * T_t * T_s) + h_outer * (T_t * T_s) + + # Load the (BLOCK_S, d_i) Hk slab once — it is loop-invariant over H. + hk_ptrs = Hk_ptr + hk_base + rs[:, None] * d_i + rdi[None, :] + Hk_tile = tl.load(hk_ptrs, mask=rs_mask[:, None], other=0.0) + Hk_T = tl.trans(Hk_tile) # (d_i, BLOCK_S) + + acc = tl.zeros((BLOCK_T, BLOCK_S), dtype=tl.float32) + + for h in range(H): + hq_ptrs = (Hq_ptr + hq_base + + rt[:, None] * (H * d_i) + h * d_i + rdi[None, :]) + Hq_h = tl.load(hq_ptrs, mask=rt_mask[:, None], other=0.0) + + wo_ptrs = W_o_ptr + wo_base + rt * H + h + w_h = tl.load(wo_ptrs, mask=rt_mask, other=0.0) + + score = tl.dot(Hq_h, Hk_T) + score = tl.maximum(score, 0.0) + acc += score * w_h[:, None].to(tl.float32) + + o_ptrs = O_ptr + o_base + rt[:, None] * T_s + rs[None, :] + tl.store(o_ptrs, acc.to(O_ptr.dtype.element_ty), + mask=rt_mask[:, None] & rs_mask[None, :]) + + +_score_reduce_p = extend_core.Primitive("te_indexer_score_reduce_triton") +_score_reduce_p.multiple_results = True + + +@_score_reduce_p.def_abstract_eval +def _score_reduce_abstract(Hq, Hk, W_o, *, out_dtype): + del W_o + # Hq layout: (B, oH, T_t, H, d_i) + B, oH, T_t, _H, _d_i = Hq.shape + T_s = Hk.shape[2] + return [core.ShapedArray((B, oH, T_t, T_s), out_dtype)] + + +_score_reduce_p.def_impl(functools.partial(xla.apply_primitive, _score_reduce_p)) + + +def _score_reduce_lowering(ctx, Hq, Hk, W_o, *, out_dtype): + del out_dtype + Hq_aval = ctx.avals_in[0] + Hk_aval = ctx.avals_in[1] + B, oH, T_t, H, d_i = Hq_aval.shape + T_s = Hk_aval.shape[2] + + def grid_fn(merged_kwargs): + bt = merged_kwargs.get("BLOCK_T", 64) + bs = merged_kwargs.get("BLOCK_S", 64) + # S as grid_x (fastest-dispatching) so per-(B*oH, T-tile) S workgroups + # cluster in time and hit L2 on the shared Hq slab. + return (triton.cdiv(T_s, bs), triton.cdiv(T_t, bt), B * oH) + + saved_configs = _score_reduce_kernel.configs + if _autotune_disabled(): + _score_reduce_kernel.configs = saved_configs[:1] + try: + return triton_call_lowering( + ctx, + _score_reduce_kernel, + Hq, Hk, W_o, + grid=grid_fn, + constexprs={ + "B": B, + "oH": oH, + "T_t": T_t, + "T_s": T_s, + "H": H, + "d_i": d_i, + }, + ) + finally: + _score_reduce_kernel.configs = saved_configs + + +mlir.register_lowering(_score_reduce_p, _score_reduce_lowering, platform="rocm") +mlir.register_lowering(_score_reduce_p, _score_reduce_lowering, platform="cuda") + + +# --- Chunked score-tile kernel for hybrid bwd -------------------------------- +# +# Produces dscores_chunk[B, oH, T, H_CHUNK, T_s] and dW_o_chunk[B, oH, T, H_CHUNK] +# for ONE h-chunk. Caller loops over H/H_CHUNK chunks and feeds dscores_chunk +# to hipBLASLt einsums for dHq/dHk reductions. Bounds peak materialization to +# H/H_CHUNK fraction of the full (B, oH, T, H, T_s) score tensor. +# +# Fuses score recompute + relu + mask + dO*W_o broadcast in registers -- +# nothing of size (B, oH, T, H, T_s) ever lands in HBM at full size. dW_o is +# reduced inline (sum_s of h_relu * dO) so h_relu also never materializes. + + +_HBWD_BLOCK_T = 64 + + +def _score_dscores_chunk_autotune_configs(): + # matrix_instr_nonkdim is pinned to 16. On Triton 3.7.0 / gfx950 the + # H_CHUNK-unrolled score matmul crashes the compiler (uncatchable + # std::bad_alloc / heap corruption the autotuner cannot skip) for both + # nonkdim=32 and the backend-default nonkdim (0); only 16x16 MFMA tiles + # compile. num_stages is pinned to 1: pipelining the s-loop crashes LLVM + # codegen on the same toolchain. + cfgs = [] + cfgs += [ + triton.Config( + {"BLOCK_T": bt, "BLOCK_S": bs, + "matrix_instr_nonkdim": 16, "waves_per_eu": wpe}, + num_warps=nw, num_stages=1) + for bt in (32, 64, 128) + for bs in (128, 256) + for wpe in (0, 2) + for nw in (4, 8) + ] + # larger BLOCK_S for long T_s + cfgs += [ + triton.Config({"BLOCK_T": bt, "BLOCK_S": 512, "matrix_instr_nonkdim": 16}, + num_warps=4, num_stages=1) + for bt in (32, 64) + ] + return cfgs + + +@triton.autotune(configs=_score_dscores_chunk_autotune_configs(), + key=["T", "T_s", "H_CHUNK", "d_i"]) +@triton.jit +def _score_dscores_chunk_kernel( + Hq_chunk_ptr, # input (B, oH, T, H_CHUNK, d_i) bf16 + Hk_ptr, # input (B, oH, T_s, d_i) bf16 + W_o_chunk_ptr, # input (B, oH, T, H_CHUNK) bf16 + dO_ptr, # input (B, oH, T, T_s) fp32 + dscores_chunk_ptr, # output (B, oH, T, H_CHUNK, T_s) bf16 + dWo_chunk_ptr, # output (B, oH, T, H_CHUNK) bf16 + B: tl.constexpr, + oH: tl.constexpr, + T: tl.constexpr, + T_s: tl.constexpr, + H_CHUNK: tl.constexpr, + d_i: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """One CTA handles (T_tile, all H_CHUNK heads) for one (b, h_outer). + + Grid: (cdiv(T, BLOCK_T), B * oH). For each s-chunk we load dO_chunk and + Hk_chunk ONCE and reuse them across every head in the chunk -- the key + saving vs the original (which spun a separate CTA per head, each re-reading + dO/Hk). dW_o is reduced in registers (sum over s) per head, so h_relu never + lands in HBM. + """ + pid_t = tl.program_id(0) + pid_bh = tl.program_id(1) + b = (pid_bh // oH).to(tl.int64) + h_outer = (pid_bh % oH).to(tl.int64) + + rt = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) + rdi = tl.arange(0, d_i) + rhc = tl.arange(0, H_CHUNK) + rt_mask = rt < T + + hq_base = b * (oH * T * H_CHUNK * d_i) + h_outer * (T * H_CHUNK * d_i) + hk_base = b * (oH * T_s * d_i) + h_outer * (T_s * d_i) + wo_base = b * (oH * T * H_CHUNK) + h_outer * (T * H_CHUNK) + do_base = b * (oH * T * T_s) + h_outer * (T * T_s) + ds_base = b * (oH * T * H_CHUNK * T_s) + h_outer * (T * H_CHUNK * T_s) + + # Per-head dW_o accumulators packed as (BLOCK_T, H_CHUNK), reduced over s. + dWo_acc = tl.zeros((BLOCK_T, H_CHUNK), dtype=tl.float32) + + for s_start in range(0, T_s, BLOCK_S): + rs = s_start + tl.arange(0, BLOCK_S) + rs_mask = rs < T_s + + # Load Hk[..., s_chunk, :] and dO[..., t_tile, s_chunk] ONCE per s-chunk + # -- shared across all H_CHUNK heads below. + hk_ptrs = Hk_ptr + hk_base + rs[:, None] * d_i + rdi[None, :] + Hk_chunk = tl.load(hk_ptrs, mask=rs_mask[:, None], other=0.0) + Hk_T = tl.trans(Hk_chunk) # (d_i, BLOCK_S) + + do_ptrs = dO_ptr + do_base + rt[:, None] * T_s + rs[None, :] + dO_chunk = tl.load( + do_ptrs, mask=rt_mask[:, None] & rs_mask[None, :], other=0.0, + ) + + for h in tl.static_range(H_CHUNK): + # Hq/w for head h (small, L2-resident across s-chunks). + Hq_h = tl.load( + Hq_chunk_ptr + hq_base + rt[:, None] * (H_CHUNK * d_i) + + h * d_i + rdi[None, :], + mask=rt_mask[:, None], other=0.0, + ) + w_h = tl.load( + W_o_chunk_ptr + wo_base + rt * H_CHUNK + h, + mask=rt_mask, other=0.0, + ).to(tl.float32) + + scores = tl.dot(Hq_h, Hk_T) # (BLOCK_T, BLOCK_S) + relu_mask = scores > 0 + h_relu = tl.where(relu_mask, scores, 0.0) + + # dW_o[..., h] += sum_s (h_relu * dO); accumulate into column h. + dwo_h = tl.sum(h_relu * dO_chunk, axis=1) # (BLOCK_T,) + dWo_acc += tl.where(rhc[None, :] == h, dwo_h[:, None], 0.0) + + # dscores[..., h, s] = relu_mask * (dO * W_o) + dscores = tl.where(relu_mask, dO_chunk * w_h[:, None], 0.0) + ds_ptrs = (dscores_chunk_ptr + ds_base + + rt[:, None] * (H_CHUNK * T_s) + h * T_s + rs[None, :]) + tl.store( + ds_ptrs, dscores.to(dscores_chunk_ptr.dtype.element_ty), + mask=rt_mask[:, None] & rs_mask[None, :], + ) + + # Store dW_o[..., t_tile, :] for all heads. + dwo_out_ptrs = dWo_chunk_ptr + wo_base + rt[:, None] * H_CHUNK + rhc[None, :] + tl.store( + dwo_out_ptrs, dWo_acc.to(dWo_chunk_ptr.dtype.element_ty), + mask=rt_mask[:, None], + ) + + +_score_dscores_chunk_p = extend_core.Primitive("te_indexer_score_dscores_chunk") +_score_dscores_chunk_p.multiple_results = True + + +@_score_dscores_chunk_p.def_abstract_eval +def _score_dscores_chunk_abstract(Hq_chunk, Hk, W_o_chunk, dO): + del Hk, W_o_chunk + B, oH, T, H_CHUNK, _ = Hq_chunk.shape + T_s = dO.shape[-1] + return [ + core.ShapedArray((B, oH, T, H_CHUNK, T_s), Hq_chunk.dtype), # dscores + core.ShapedArray((B, oH, T, H_CHUNK), Hq_chunk.dtype), # dW_o + ] + + +_score_dscores_chunk_p.def_impl( + functools.partial(xla.apply_primitive, _score_dscores_chunk_p) +) + + +def _score_dscores_chunk_lowering(ctx, Hq_chunk, Hk, W_o_chunk, dO): + Hq_aval = ctx.avals_in[0] + dO_aval = ctx.avals_in[3] + B, oH, T, H_CHUNK, d_i = Hq_aval.shape + T_s = dO_aval.shape[-1] + + # Grid: (T-tiles, B*oH) -- one CTA per (T_tile, b, h_outer) covers all + # H_CHUNK heads (dO/Hk shared across heads). Depends on autotuned BLOCK_T. + def grid_fn(merged_kwargs): + bt = merged_kwargs.get("BLOCK_T", _HBWD_BLOCK_T) + return ((T + bt - 1) // bt, B * oH) + + saved_configs = _score_dscores_chunk_kernel.configs + if _autotune_disabled(): + _score_dscores_chunk_kernel.configs = saved_configs[:1] + try: + return triton_call_lowering( + ctx, + _score_dscores_chunk_kernel, + Hq_chunk, Hk, W_o_chunk, dO, + grid=grid_fn, + constexprs={ + "B": B, "oH": oH, "T": T, "T_s": T_s, + "H_CHUNK": H_CHUNK, "d_i": d_i, + }, + ) + finally: + _score_dscores_chunk_kernel.configs = saved_configs + + +mlir.register_lowering(_score_dscores_chunk_p, _score_dscores_chunk_lowering, platform="rocm") +mlir.register_lowering(_score_dscores_chunk_p, _score_dscores_chunk_lowering, platform="cuda") + + +# --- Public score_reduce_triton with custom_vjp ------------------------------ + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(3,)) +def _score_reduce_with_vjp(Hq, Hk, W_o, out_dtype): + return _score_reduce_p.bind(Hq, Hk, W_o, out_dtype=out_dtype)[0] + + +def _score_reduce_fwd(Hq, Hk, W_o, out_dtype): + out = _score_reduce_p.bind(Hq, Hk, W_o, out_dtype=out_dtype)[0] + return out, (Hq, Hk, W_o) + + +_BWD_H_CHUNK = 8 # peak (B, oH, T, H_CHUNK, T_s) tile -- bounds materialization + + +def _score_reduce_bwd(out_dtype, residuals, dO): + del out_dtype + Hq, Hk, W_o = residuals + B, oH, T, H, d_i = Hq.shape + + # Hybrid scheme with bounded materialization: + # For each h-chunk of size H_CHUNK (driven by lax.scan, NOT Python + # unroll, so intermediates are freed between iterations): + # 1. Triton kernel fuses (score recompute + relu + mask + dO*W_o + # broadcast) and writes dscores_chunk[B,oH,T,H_CHUNK,T_s] to HBM. + # h_relu is consumed in-register to also produce dWo_chunk + # without ever materializing the (B,oH,T,H,T_s) h_relu tensor. + # 2. hipBLASLt einsums on dscores_chunk give dHq_chunk and a partial + # dHk contribution. + # Peak HBM intermediate stays at H_CHUNK/H fraction of the full score. + if H % _BWD_H_CHUNK == 0: + H_CHUNK = _BWD_H_CHUNK + else: + H_CHUNK = 1 + for c in (4, 2): + if H % c == 0: + H_CHUNK = c + break + n_chunks = H // H_CHUNK + + Hq_r = Hq.reshape(B, oH, T, n_chunks, H_CHUNK, d_i) + Wo_r = W_o.reshape(B, oH, T, n_chunks, H_CHUNK) + # Move chunk axis to leading for scan over axis 0. + Hq_s = jnp.moveaxis(Hq_r, -3, 0) # (n_chunks, B, oH, T, H_CHUNK, d_i) + Wo_s = jnp.moveaxis(Wo_r, -2, 0) # (n_chunks, B, oH, T, H_CHUNK) + + def step(dHk_acc, chunk): + Hq_c, Wo_c = chunk + # Triton: dscores_chunk + dWo_chunk; no full (B,oH,T,H,T_s) tensor + # ever exists in HBM. + dscores_c, dWo_c = _score_dscores_chunk_p.bind(Hq_c, Hk, Wo_c, dO) + dHq_c = jnp.einsum("...ths,...si->...thi", dscores_c, Hk) + dHk_c = jnp.einsum("...ths,...thi->...si", dscores_c, Hq_c) + new_dHk_acc = dHk_acc + dHk_c.astype(jnp.float32) + return new_dHk_acc, (dHq_c, dWo_c) + + init = jnp.zeros(Hk.shape, dtype=jnp.float32) + dHk_acc, (dHq_chunks, dWo_chunks) = jax.lax.scan( + step, init, (Hq_s, Wo_s), + ) + # dHq_chunks: (n_chunks, B, oH, T, H_CHUNK, d_i) + # dWo_chunks: (n_chunks, B, oH, T, H_CHUNK) + dHq = jnp.moveaxis(dHq_chunks, 0, -3).reshape(B, oH, T, H, d_i) + dWo = jnp.moveaxis(dWo_chunks, 0, -2).reshape(B, oH, T, H) + dHk = dHk_acc.astype(Hk.dtype) + + return dHq.astype(Hq.dtype), dHk, dWo.astype(W_o.dtype) + + +_score_reduce_with_vjp.defvjp(_score_reduce_fwd, _score_reduce_bwd) + + +def score_reduce_triton(Hq, Hk, W_o, *, out_dtype=None): + """Triton fused score-matmul + relu + per-(t, h) weighted H-reduction. + + Replaces the pattern: + + scores = relu(jnp.einsum("...thi,...si->...ths", Hq, Hk)) # never write + O = jnp.einsum("...ths,...th->...ts", scores, W_o) + + with a single kernel that holds the per-head score tile in registers, + avoiding the (B, oH, T, H, S) HBM round-trip an einsum+XLA chain pays. + + Differentiable via two backward kernels (FlashAttention-style: residuals + are just (Hq, Hk, W_o); the (T, H, S) score tensor is recomputed inside + backward, never materialized). + + Args: + Hq: (B, oH, T_t, H, d_i) + Hk: (B, oH, T_s, d_i) + W_o: (B, oH, T_t, H) + out_dtype: defaults to Hq.dtype. + + Returns: + O: (B, oH, T_t, T_s) + """ + if Hq.ndim != 5: + raise ValueError( + f"Hq must be rank-5 (B, oH, T_t, H, d_i); got shape {Hq.shape}" + ) + if Hk.ndim != 4: + raise ValueError( + f"Hk must be rank-4 (B, oH, T_s, d_i); got shape {Hk.shape}" + ) + if W_o.ndim != 4: + raise ValueError( + f"W_o must be rank-4 (B, oH, T_t, H); got shape {W_o.shape}" + ) + + B, oH, T_t, H, d_i = Hq.shape + Bk, oHk, T_s, d_i_k = Hk.shape + Bw, oHw, T_t_w, H_w = W_o.shape + if (Bk, oHk) != (B, oH): + raise ValueError( + f"(B, oH) mismatch: Hq has {(B, oH)}, Hk has {(Bk, oHk)}" + ) + if d_i != d_i_k: + raise ValueError(f"d_i mismatch: Hq has {d_i}, Hk has {d_i_k}") + if (Bw, oHw, T_t_w, H_w) != (B, oH, T_t, H): + raise ValueError( + f"W_o shape {W_o.shape} does not match expected " + f"(B={B}, oH={oH}, T_t={T_t}, H={H})" + ) + + if out_dtype is None: + out_dtype = Hq.dtype + + return _score_reduce_with_vjp(Hq, Hk, W_o, jnp.dtype(out_dtype)) + + +# --- Streaming top-k variant ---------------------------------------------------- +# +# Same einsum-projected (Hq, Hk, W_o) inputs, but fuses top-k indices into the +# kernel: one CTA per (B, oH, T_t) query token, score row never materialized. +# +# Algorithm (mirrors TileLang dsa_sparse_finetune/indexer_topk_reducesum): +# - Maintain a 2K-sized buffer of (score_bits, index) packed uint64 +# - Stream over T_s in BLOCK_S chunks; each chunk computes BLOCK_S new scores +# - Place chunk into buffer[K:K+BLOCK_S], zero buffer[K+BLOCK_S:2K] +# - tl.sort descending; top half is the running top-K +# - After all chunks: buffer[:K] is the answer +# +# tl.sort returns values only, so we pack (score_bits << 32) | index into uint64. +# Post-ReLU scores are >= 0, so fp32 bit pattern is monotone in value. + + +# Autotune sweep for _score_topk_kernel. +# +# BLOCK_T: number of query tokens per CTA. BLOCK_T>1 amortizes the Hk_chunk +# load across BLOCK_T queries — the single biggest lever at large T_s. At +# BLOCK_T=1 (original), each CTA reloads all of Hk for its (b, oH) slab, +# causing L2 thrash. BLOCK_T=2 halves Hk HBM traffic; BLOCK_T=4 quarters it, +# but grows per-CTA register pressure (Hq_token, top_packed, logits all +# scale with BLOCK_T). +# +# BLOCK_S knobs the inner-chunk size; bigger BLOCK_S = better matmul +# arithmetic intensity, but bigger per-CTA transient footprint +# (logits[BLOCK_S, BLOCK_T*H] fp32 + Hk_chunk[BLOCK_S, d_i] bf16). +# +# Constraint: BLOCK_S must divide K (so INNER = K // BLOCK_S is an integer +# >= 1). Configs whose BLOCK_S exceeds K or doesn't divide K must be pruned +# (see _prune_topk_configs) — otherwise the autotuner would time them +# as zero-work (fast) and pick a bogus winner that returns all-zero indices. +_SCORE_TOPK_CONFIGS = [ + triton.Config({"BLOCK_S": bs, "BLOCK_T": bt, "waves_per_eu": wpe}, num_warps=nw, num_stages=ns) + for bt in (1, 2) + for bs in (32, 64, 128, 256) + for wpe in (0, 2, 4) + for nw in (4, 8) + for ns in (1, 2) +] + [ + # BLOCK_T=4 only at smaller BLOCK_S — at BLOCK_S=256 the logits + # intermediate [256, 4*H=256] fp32 = 256 KB overflows reliably. + triton.Config({"BLOCK_S": bs, "BLOCK_T": 4}, num_warps=nw, num_stages=ns) + for bs in (32, 64, 128) + for nw in (4, 8) + for ns in (1, 2) +] + + +def _prune_topk_configs(configs, named_args, **kwargs): + """early_config_prune for _score_topk_kernel. Keep only configs where + BLOCK_S divides K (INNER = K//BLOCK_S >= 1) and BLOCK_T divides T_t. The + runtime values arrive in named_args or kwargs depending on call style.""" + vals = {**named_args, **kwargs} + k = vals["K"] + T_t = vals["T_t"] + return [ + c for c in configs + if c.kwargs["BLOCK_S"] <= k + and k % c.kwargs["BLOCK_S"] == 0 + and T_t % c.kwargs["BLOCK_T"] == 0 + ] + + +@triton.autotune( + configs=_SCORE_TOPK_CONFIGS, + key=["H", "d_i", "T_s", "K"], + prune_configs_by={"early_config_prune": _prune_topk_configs}, +) +@triton.jit +def _score_topk_kernel( + Hq_ptr, # (B, oH, T_t, H, d_i) bf16 + Hk_ptr, # (B, oH, T_s, d_i) bf16 + W_o_ptr, # (B, oH, T_t, H) bf16 + Topk_idx_ptr, # (B, oH, T_t, K) int32 OUTPUT + B: tl.constexpr, + oH: tl.constexpr, + T_t: tl.constexpr, + T_s: tl.constexpr, + H: tl.constexpr, + d_i: tl.constexpr, + K: tl.constexpr, + S_PAD: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Per-CTA: BLOCK_T consecutive query tokens, all sharing Hk loads. + + Grid: (cdiv(T_t, BLOCK_T), B * oH). Each CTA does: + - Pre-load Hq[..., rt, :, :] for BLOCK_T contiguous query tokens + - For each S chunk: load Hk_chunk ONCE, do one [BLOCK_S, d_i] @ + [d_i, BLOCK_T*H] matmul, weighted-H-reduce per T + - Maintain a single 1D top buffer of size BLOCK_T*2K, with T encoded + in the top 8 bits of each packed entry. After global sort desc, + per-T entries stay grouped together so per-T top-K can be sliced + from fixed offsets. + + Note on layout (1D vs 2D top buffer): + A 2D [BLOCK_T, 2K] top buffer with per-row sort is the natural + design, but `tl.gather + tl.sort(dim=1)` on uint64 2D tensors trips + `TritonGPUOptimizeThreadLocality` on the AMD backend (gfx950, Triton + 3.4.0). The 1D-with-encoded-T workaround sidesteps this — it pays a + ~1.5x sort-cost penalty (one sort of BLOCK_T*2K vs BLOCK_T sorts of + 2K) for BLOCK_T=2, but unblocks Hk-load amortization across queries. + """ + pid_t = tl.program_id(0) + pid_bh = tl.program_id(1) + # int64 indexing — Hq has B*oH*T*H*d_i = 4.3 B elements at T=S=4096. + b = (pid_bh // oH).to(tl.int64) + h_outer = (pid_bh % oH).to(tl.int64) + + rh = tl.arange(0, H) + rdi = tl.arange(0, d_i) + rs_chunk = tl.arange(0, BLOCK_S) + rk = tl.arange(0, K) + rt_local = tl.arange(0, BLOCK_T) + + rt = pid_t * BLOCK_T + rt_local + rt_64 = rt.to(tl.int64) + rt_mask = rt < T_t + + # Load Hq[b, h_outer, rt, :, :] -> [BLOCK_T, H, d_i]. + hq_base = b * (oH * T_t * H * d_i) + h_outer * (T_t * H * d_i) + Hq_token = tl.load( + Hq_ptr + hq_base + + rt_64[:, None, None] * (H * d_i) + + rh[None, :, None] * d_i + + rdi[None, None, :], + mask=rt_mask[:, None, None], + other=0.0, + ) + + # Load w_o[b, h_outer, rt, :] -> [BLOCK_T, H] + wo_base = b * (oH * T_t * H) + h_outer * (T_t * H) + w_o = tl.load( + W_o_ptr + wo_base + rt_64[:, None] * H + rh[None, :], + mask=rt_mask[:, None], + other=0.0, + ).to(tl.float32) + + # Flatten Hq for one big matmul per Hk_chunk: [BLOCK_T * H, d_i] -> trans + Hq_flat = tl.reshape(Hq_token, (BLOCK_T * H, d_i)) + Hq_T = tl.trans(Hq_flat) # [d_i, BLOCK_T * H] + w_o_flat = tl.reshape(w_o, (BLOCK_T * H,)) + + hk_base = b * (oH * T_s * d_i) + h_outer * (T_s * d_i) + + TOP_BUF: tl.constexpr = 2 * K + INNER: tl.constexpr = K // BLOCK_S # chunks per sort + N_OUTER: tl.constexpr = S_PAD // K # number of sorts per CTA + BIG_BUF: tl.constexpr = BLOCK_T * TOP_BUF + + # Initialize 1D top buffer with t-encoding pre-applied so per-T regions + # stay grouped after global sort. Each slot at position rb gets: + # t_pos = rb // TOP_BUF -> which T this slot belongs to + # t_enc = BLOCK_T - t_pos -> 1..BLOCK_T (never 0 → never collides with + # reserved init pattern) + # packed = (t_enc << 56) | 0 -> score=0 (sortable=0), index=0 + # Real candidates also get tagged with their t_enc; after global sort + # desc, all entries with t_enc=BLOCK_T (i.e. t=0) come first, then + # t_enc=BLOCK_T-1, etc. Within each t group, ordered by score then index. + rb = tl.arange(0, BIG_BUF) + rb_t = rb // TOP_BUF # [BIG_BUF] in [0, BLOCK_T) + rb_pos = rb % TOP_BUF # [BIG_BUF] in [0, TOP_BUF) + t_enc_per_slot = (BLOCK_T - rb_t).to(tl.uint64) + top_packed = t_enc_per_slot << 56 + + # Pre-compute the per-slot (t, pos)-to-flat-chunk-index map used in + # scatter: for each rb, identify the (t, j) in chunk_packed_flat to pull + # from. j depends on `chunk_offset` (varies per inner iter), so the + # gather index is recomputed each iter. + + for o in tl.static_range(N_OUTER): + for i in tl.static_range(INNER): + c = o * INNER + i + s_start = c * BLOCK_S + rs = s_start + rs_chunk # [BLOCK_S] + rs_mask = rs < T_s + + # Load Hk_chunk[BLOCK_S, d_i] ONCE — shared across BLOCK_T queries. + hk_ptrs = Hk_ptr + hk_base + rs[:, None] * d_i + rdi[None, :] + Hk_chunk = tl.load(hk_ptrs, mask=rs_mask[:, None], other=0.0) + + # One big matmul: [BLOCK_S, d_i] @ [d_i, BLOCK_T*H] -> [BLOCK_S, BLOCK_T*H] + logits = tl.dot(Hk_chunk, Hq_T) + logits = tl.maximum(logits, 0.0) + + # Weighted reduce over H per (s, t): + # chunk_scores[s, t] = sum_h logits[s, t*H + h] * w_o[t, h] + weighted = logits * w_o_flat[None, :] + weighted_3d = tl.reshape(weighted, (BLOCK_S, BLOCK_T, H)) + chunk_scores = tl.sum(weighted_3d, axis=2) # [BLOCK_S, BLOCK_T] + chunk_scores_T = tl.trans(chunk_scores) # [BLOCK_T, BLOCK_S] + + # Radix-flip: fp32 bit pattern -> sortable uint32 across full sign + # range (positives: flip sign bit; negatives: flip all bits). + # See https://stereopsis.com/radix.html + bits = chunk_scores_T.to(tl.uint32, bitcast=True) + sign = bits >> 31 + flip_mask = (0 - sign.to(tl.int32)).to(tl.uint32) | 0x80000000 + sortable = bits ^ flip_mask + sortable = tl.where(rs_mask[None, :], sortable, 0) + + # Pack: (t_enc<<56) | (sortable<<24) | (index in low 24 bits). + # 24-bit index supports T_s up to 16M, far above our regime. + t_enc_chunk = (BLOCK_T - rt_local).to(tl.uint64) # [BLOCK_T] + rs_2d = tl.broadcast_to(rs[None, :], (BLOCK_T, BLOCK_S)) + chunk_packed_2d = ( + (t_enc_chunk[:, None] << 56) + | (sortable.to(tl.uint64) << 24) + | rs_2d.to(tl.uint64) + ) # [BLOCK_T, BLOCK_S] + # Flatten to 1D for the scatter (1D gather + 1D sort sidesteps + # the AMD-backend bug with 2D gather+sort combos). + chunk_packed_flat = tl.reshape(chunk_packed_2d, (BLOCK_T * BLOCK_S,)) + + # Scatter into top_packed[t*TOP_BUF + K+i*BLOCK_S : ...] for each t. + # For each rb in [0, BIG_BUF): + # t = rb // TOP_BUF + # pos = rb % TOP_BUF + # in_slot = (pos >= K + i*BLOCK_S) & (pos < K + (i+1)*BLOCK_S) + # flat_idx = t * BLOCK_S + (pos - (K + i*BLOCK_S)) + chunk_offset = K + i * BLOCK_S + in_slot = (rb_pos >= chunk_offset) & (rb_pos < chunk_offset + BLOCK_S) + j = rb_pos - chunk_offset + flat_idx = tl.where(in_slot, rb_t * BLOCK_S + j, 0).to(tl.int32) + gathered = tl.gather(chunk_packed_flat, flat_idx, axis=0) + top_packed = tl.where(in_slot, gathered, top_packed) + + # 1D sort of the entire buffer. Per-T regions stay grouped via t_enc. + top_packed = tl.sort(top_packed, descending=True) + + # Extract per-T top K. After sort desc, t=0's top K is at positions + # [0, K), t=1's at [TOP_BUF, TOP_BUF+K), etc. — i.e. base = t*TOP_BUF. + out_idx = rt_local[:, None] * TOP_BUF + rk[None, :] # [BLOCK_T, K] + out_idx_flat = tl.reshape(out_idx, (BLOCK_T * K,)).to(tl.int32) + top_k_packed_flat = tl.gather(top_packed, out_idx_flat, axis=0) + top_k_packed = tl.reshape(top_k_packed_flat, (BLOCK_T, K)) + # Strip the t_enc and sortable bits, keep low 24 bits (index). + top_k_idx = (top_k_packed & 0xFFFFFF).to(tl.int32) + + out_base = b * (oH * T_t * K) + h_outer * (T_t * K) + out_ptrs = Topk_idx_ptr + out_base + rt_64[:, None] * K + rk[None, :] + tl.store(out_ptrs, top_k_idx, mask=rt_mask[:, None]) + + +# --- Single-sort top-k (for S_PAD that fits in registers) ------------- +# +# The streaming kernel above sorts a 2K buffer N_OUTER = S_PAD/K times. When k +# is a large fraction of T_s (e.g. k = T_s/2), that's several sorts of the 2K +# buffer. If all S_PAD candidates fit in registers, scattering them into one +# BLOCK_T*S_PAD buffer and doing a SINGLE descending sort is ~2x less sort work. +_SINGLE_SORT_MAX = 4096 + + +# Configs for the single-sort kernel. BLOCK_S must divide S_PAD; BLOCK_T must +# divide T_t (pruned below). Tuned winner on gfx950 is BLOCK_T=1, BLOCK_S=128. +_SINGLE_TOPK_CONFIGS = [ + triton.Config({"BLOCK_S": bs, "BLOCK_T": bt, "waves_per_eu": wpe}, num_warps=nw, num_stages=1) + for bs in (64, 128, 256) + for bt in (1, 2) + for wpe in (0, 2, 3, 4) + for nw in (4, 8) +] + + +def _prune_single_topk_configs(configs, named_args, **kwargs): + """early_config_prune for _score_topk_single_kernel. Keep only configs where + BLOCK_S divides S_PAD (the static chunk loop tiles it exactly) and BLOCK_T + divides T_t.""" + vals = {**named_args, **kwargs} + S_PAD = vals["S_PAD"] + T_t = vals["T_t"] + return [ + c for c in configs + if S_PAD % c.kwargs["BLOCK_S"] == 0 + and T_t % c.kwargs["BLOCK_T"] == 0 + ] + + +@triton.autotune( + configs=_SINGLE_TOPK_CONFIGS, + key=["H", "d_i", "T_s", "K"], + prune_configs_by={"early_config_prune": _prune_single_topk_configs}, +) +@triton.jit +def _score_topk_single_kernel( + Hq_ptr, Hk_ptr, W_o_ptr, Topk_idx_ptr, + B: tl.constexpr, oH: tl.constexpr, T_t: tl.constexpr, T_s: tl.constexpr, + H: tl.constexpr, d_i: tl.constexpr, K: tl.constexpr, S_PAD: tl.constexpr, + BLOCK_S: tl.constexpr, BLOCK_T: tl.constexpr, +): + """Like ``_score_topk_kernel`` but holds all S_PAD candidates and sorts once. + + Grid: (cdiv(T_t, BLOCK_T), B * oH). Buffer is BLOCK_T*S_PAD packed uint64 + with T encoded in the high bits (same 1D-sort-groups-per-T trick). Requires + BLOCK_S | S_PAD (so the static chunk loop tiles S_PAD exactly); no BLOCK_S|K + constraint is needed since there is no 2K streaming buffer. + """ + pid_t = tl.program_id(0) + pid_bh = tl.program_id(1) + b = (pid_bh // oH).to(tl.int64) + h_outer = (pid_bh % oH).to(tl.int64) + rh = tl.arange(0, H) + rdi = tl.arange(0, d_i) + rs_chunk = tl.arange(0, BLOCK_S) + rk = tl.arange(0, K) + rt_local = tl.arange(0, BLOCK_T) + rt = pid_t * BLOCK_T + rt_local + rt_64 = rt.to(tl.int64) + rt_mask = rt < T_t + + hq_base = b * (oH * T_t * H * d_i) + h_outer * (T_t * H * d_i) + Hq_token = tl.load( + Hq_ptr + hq_base + rt_64[:, None, None] * (H * d_i) + + rh[None, :, None] * d_i + rdi[None, None, :], + mask=rt_mask[:, None, None], other=0.0) + wo_base = b * (oH * T_t * H) + h_outer * (T_t * H) + w_o = tl.load(W_o_ptr + wo_base + rt_64[:, None] * H + rh[None, :], + mask=rt_mask[:, None], other=0.0).to(tl.float32) + Hq_flat = tl.reshape(Hq_token, (BLOCK_T * H, d_i)) + Hq_T = tl.trans(Hq_flat) + w_o_flat = tl.reshape(w_o, (BLOCK_T * H,)) + hk_base = b * (oH * T_s * d_i) + h_outer * (T_s * d_i) + + N_CHUNK: tl.constexpr = S_PAD // BLOCK_S + BIG: tl.constexpr = BLOCK_T * S_PAD + rb = tl.arange(0, BIG) + rb_t = rb // S_PAD + rb_pos = rb % S_PAD + t_enc_per_slot = (BLOCK_T - rb_t).to(tl.uint64) + top_packed = (t_enc_per_slot << 56) | rb_pos.to(tl.uint64) + + for c in tl.static_range(N_CHUNK): + rs = c * BLOCK_S + rs_chunk + rs_mask = rs < T_s + hk_ptrs = Hk_ptr + hk_base + rs[:, None] * d_i + rdi[None, :] + Hk_chunk = tl.load(hk_ptrs, mask=rs_mask[:, None], other=0.0) + logits = tl.dot(Hk_chunk, Hq_T) + logits = tl.maximum(logits, 0.0) + weighted = logits * w_o_flat[None, :] + weighted_3d = tl.reshape(weighted, (BLOCK_S, BLOCK_T, H)) + chunk_scores = tl.sum(weighted_3d, axis=2) + chunk_scores_T = tl.trans(chunk_scores) + bits = chunk_scores_T.to(tl.uint32, bitcast=True) + sign = bits >> 31 + flip_mask = (0 - sign.to(tl.int32)).to(tl.uint32) | 0x80000000 + sortable = bits ^ flip_mask + sortable = tl.where(rs_mask[None, :], sortable, 0) + t_enc_chunk = (BLOCK_T - rt_local).to(tl.uint64) + rs_2d = tl.broadcast_to(rs[None, :], (BLOCK_T, BLOCK_S)) + chunk_packed_2d = ((t_enc_chunk[:, None] << 56) + | (sortable.to(tl.uint64) << 24) | rs_2d.to(tl.uint64)) + chunk_packed_flat = tl.reshape(chunk_packed_2d, (BLOCK_T * BLOCK_S,)) + chunk_offset = c * BLOCK_S + in_slot = (rb_pos >= chunk_offset) & (rb_pos < chunk_offset + BLOCK_S) + j = rb_pos - chunk_offset + flat_idx = tl.where(in_slot, rb_t * BLOCK_S + j, 0).to(tl.int32) + gathered = tl.gather(chunk_packed_flat, flat_idx, axis=0) + top_packed = tl.where(in_slot, gathered, top_packed) + + top_packed = tl.sort(top_packed, descending=True) # SINGLE sort + out_idx = rt_local[:, None] * S_PAD + rk[None, :] + out_idx_flat = tl.reshape(out_idx, (BLOCK_T * K,)).to(tl.int32) + top_k_packed_flat = tl.gather(top_packed, out_idx_flat, axis=0) + top_k_packed = tl.reshape(top_k_packed_flat, (BLOCK_T, K)) + top_k_idx = (top_k_packed & 0xFFFFFF).to(tl.int32) + out_base = b * (oH * T_t * K) + h_outer * (T_t * K) + out_ptrs = Topk_idx_ptr + out_base + rt_64[:, None] * K + rk[None, :] + tl.store(out_ptrs, top_k_idx, mask=rt_mask[:, None]) + + +_score_topk_p = extend_core.Primitive("te_indexer_score_topk_triton") +_score_topk_p.multiple_results = True + + +def _next_pow2(n): + p = 1 + while p < n: + p *= 2 + return p + + +@_score_topk_p.def_abstract_eval +def _score_topk_abstract(Hq, Hk, W_o, *, k): + del Hk, W_o + B, oH, T_t, _H, _d_i = Hq.shape + return [core.ShapedArray((B, oH, T_t, k), jnp.int32)] + + +_score_topk_p.def_impl(functools.partial(xla.apply_primitive, _score_topk_p)) + + +def _score_topk_lowering(ctx, Hq, Hk, W_o, *, k): + Hq_aval = ctx.avals_in[0] + Hk_aval = ctx.avals_in[1] + B, oH, T_t, H, d_i = Hq_aval.shape + T_s = Hk_aval.shape[2] + S_PAD = _next_pow2(T_s) + + # Both kernels are self-autotuned (@triton.autotune at definition); invalid + # configs are dropped by their early_config_prune hooks, so the lowering just + # picks the right kernel and launches. Single-sort path when all S_PAD + # candidates fit in registers (~1.5x faster: one sort instead of S_PAD/K + # streaming sorts of a 2K buffer); streaming kernel for very large T_s. + autotuned_kernel = (_score_topk_single_kernel if S_PAD <= _SINGLE_SORT_MAX + else _score_topk_kernel) + + def grid_fn(merged_kwargs): + bt = merged_kwargs.get("BLOCK_T", 1) + return (triton.cdiv(T_t, bt), B * oH) + + constexprs = { + "B": B, "oH": oH, "T_t": T_t, "T_s": T_s, + "H": H, "d_i": d_i, + "K": k, "S_PAD": S_PAD, + } + + # Apply the kernel's early_config_prune ourselves. triton_call_lowering hands + # every config to jaxlib's runtime autotuner, which picks the fastest by + # timing -- not correctness -- and does not run the prune hook. An invalid + # config (e.g. a BLOCK_S that doesn't divide S_PAD, leaving the static chunk + # loop with zero iterations) does no work, "wins" on speed, and returns the + # kernel's uninitialized buffer. Prune here, before lowering, so only configs + # valid for this S_PAD/K/T_t reach the autotuner. + valid_configs = autotuned_kernel.early_config_prune(autotuned_kernel.configs, constexprs) + if _autotune_disabled(): + valid_configs = valid_configs[:1] + saved_configs = autotuned_kernel.configs + autotuned_kernel.configs = valid_configs + try: + return triton_call_lowering( + ctx, + autotuned_kernel, + Hq, Hk, W_o, + grid=grid_fn, + constexprs=constexprs, + ) + finally: + autotuned_kernel.configs = saved_configs + + +mlir.register_lowering(_score_topk_p, _score_topk_lowering, platform="rocm") +mlir.register_lowering(_score_topk_p, _score_topk_lowering, platform="cuda") + + +def score_topk_triton(Hq, Hk, W_o, *, k): + """Fused score-relu-reduce + streaming top-k. + + Computes the same scores as ``score_reduce_triton`` but never materializes the + (B, oH, T_t, T_s) score matrix — instead, returns the top-k indices into the + T_s axis directly. + + Args: + Hq: (B, oH, T_t, H, d_i) + Hk: (B, oH, T_s, d_i) + W_o: (B, oH, T_t, H) + k: number of top scores to return per (b, oH, T_t) row. Must be a + power of 2 and <= T_s. + + Returns: + Topk_idx: (B, oH, T_t, k) int32 — top-k indices into T_s axis, in + descending score order. + + Notes: + Streaming: maintains a 2K candidate buffer and bitonic-sorts on each + chunk. For k >> S/8 (e.g., k=S/2), this is algorithmically slower than a + single full-row sort but matches the TileLang reference structure and + generalizes to large S without per-CTA registers scaling with S. + """ + if Hq.ndim != 5: + raise ValueError(f"Hq must be rank-5; got shape {Hq.shape}") + if Hk.ndim != 4: + raise ValueError(f"Hk must be rank-4; got shape {Hk.shape}") + if W_o.ndim != 4: + raise ValueError(f"W_o must be rank-4; got shape {W_o.shape}") + + B, oH, T_t, H, d_i = Hq.shape + Bk, oHk, T_s, d_i_k = Hk.shape + Bw, oHw, T_t_w, H_w = W_o.shape + if (Bk, oHk) != (B, oH): + raise ValueError(f"(B, oH) mismatch: Hq has {(B, oH)}, Hk has {(Bk, oHk)}") + if d_i != d_i_k: + raise ValueError(f"d_i mismatch: Hq has {d_i}, Hk has {d_i_k}") + if (Bw, oHw, T_t_w, H_w) != (B, oH, T_t, H): + raise ValueError(f"W_o shape {W_o.shape} != expected (B, oH, T_t, H)") + + if k <= 0 or (k & (k - 1)) != 0: + raise ValueError(f"k must be a positive power of 2; got {k}") + if k > T_s: + raise ValueError(f"k={k} must be <= T_s={T_s}") + + return _score_topk_p.bind(Hq, Hk, W_o, k=k)[0] diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index f94d353940..52312ceb2e 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. @@ -35,8 +36,10 @@ Default is "0" (silent compatibility fallback). """ +import dataclasses import hashlib import os +import tempfile import warnings from typing import Any, Callable, Mapping import zlib @@ -47,6 +50,7 @@ from jaxlib.mlir import ir import jax import jax.numpy as jnp +from transformer_engine.jax.util import is_hip_extension from ..version_utils import ( TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION, @@ -199,6 +203,12 @@ def _check_triton_compatibility(): ) from e +# AMD/HIP backend imports are additive: the NVIDIA path above is left untouched. +if is_hip_extension(): + from triton.backends.amd import compiler as cb_hip # noqa: E402 + from triton.backends.compiler import GPUTarget as _TritonGPUTarget # noqa: E402 + + __all__ = ["triton_call_lowering", "get_triton_info"] # Triton kernel cache (module-level, shared across all kernels) @@ -256,6 +266,9 @@ def get_triton_dtype(aval): jnp.dtype("float16"): "fp16", jnp.dtype("float8_e4m3fn"): "fp8e4nv", jnp.dtype("float8_e5m2"): "fp8e5", + # AMD gfx942 "FNUZ" variants — Triton calls these fp8e4b8/fp8e5b16. + jnp.dtype("float8_e4m3fnuz"): "fp8e4b8", + jnp.dtype("float8_e5m2fnuz"): "fp8e5b16", jnp.dtype("int64"): "i64", jnp.dtype("int32"): "i32", jnp.dtype("int16"): "i16", @@ -317,6 +330,22 @@ def compile_triton( if cache_key in _TRITON_KERNEL_CACHE: return _TRITON_KERNEL_CACHE[cache_key] + # AMD/HIP uses a separate compilation path; the NVIDIA path below is the + # unchanged upstream implementation. + if is_hip_extension(): + kernel = _compile_triton_hip( + kernel_fn, + signature, + constants, + num_warps, + num_stages, + num_ctas, + compute_capability, + enable_fp_fusion, + ) + _TRITON_KERNEL_CACHE[cache_key] = kernel + return kernel + # Compile kernel cuda_option_kwargs = {} if version.parse(_TRITON_VERSION) < version.parse("3.6.0"): @@ -377,6 +406,96 @@ def compile_triton( return kernel +# Track HSACO temp files for the lifetime of the process so the kernel paths +# we hand to jaxlib don't get garbage-collected. +_HSACO_TEMP_FILES: list[str] = [] + + +def _compile_triton_hip( + kernel_fn, + signature, + constants, + num_warps, + num_stages, + num_ctas, + compute_capability, + enable_fp_fusion, +): + # AMD/HIP returns an arch string like "gfx950:sramecc+:xnack-"; strip the + # target-feature suffix -> "gfx950". + arch = gpu_triton.get_arch_details(0).split(":", 1)[0] + # Mirror what triton's parse_options would do per-arch: the default + # HIPOptions.supported_fp8_dtypes is just ("fp8e5",), and constructing + # HIPOptions directly bypasses the per-arch augmentation. Set it + # explicitly so FP8 e4m3 kernels compile on gfx942/gfx950. + if arch == "gfx942": + fp8_dtypes = ("fp8e4b8", "fp8e4nv", "fp8e5", "fp8e5b16") + elif arch == "gfx950" or arch.startswith("gfx12"): + fp8_dtypes = ("fp8e4nv", "fp8e5") + else: + fp8_dtypes = ("fp8e5",) + hip_option_kwargs = dict( + num_warps=num_warps, + num_stages=num_stages, + num_ctas=num_ctas, + debug=False, + enable_fp_fusion=enable_fp_fusion, + arch=arch, + supported_fp8_dtypes=fp8_dtypes, + ) + + # Autotune configs may carry AMD compile hints (matrix_instr_nonkdim, + # waves_per_eu, kpack, ...) mixed in with real kernel constexprs. Route any + # constant that names a HIPOptions field into the compile options; the rest + # stay as kernel constexprs. (cluster_dims is gated the same way — it was + # dropped from HIPOptions in Triton 3.7.1.) + hip_option_fields = {f.name for f in dataclasses.fields(cb_hip.HIPOptions)} + kernel_constants = {} + for name, value in constants.items(): + if name in hip_option_fields: + hip_option_kwargs[name] = value + else: + kernel_constants[name] = value + if "cluster_dims" in hip_option_fields: + hip_option_kwargs.setdefault("cluster_dims", (1, 1, 1)) + options = cb_hip.HIPOptions(**hip_option_kwargs) + + # Mark constants as constexpr in signature (mirrors the NVIDIA path). + signature_with_constexpr = dict(signature) + for const_name in kernel_constants: + if const_name in signature_with_constexpr: + signature_with_constexpr[const_name] = "constexpr" + + src = tc.ASTSource( + fn=kernel_fn, + constexprs=kernel_constants, + signature=signature_with_constexpr, + ) + compiled = tc.compile( + src, + target=_TritonGPUTarget("hip", arch, warp_size=64), + options=options.__dict__, + ) + + # jaxlib's HIP TritonKernel ctor takes a path to an HSACO blob, not bytes. + fd, hsaco_path = tempfile.mkstemp(suffix=".hsaco", prefix=f"te_{compiled.name}_") + with os.fdopen(fd, "wb") as f: + f.write(compiled.asm["hsaco"]) + _HSACO_TEMP_FILES.append(hsaco_path) + + return gpu_triton.TritonKernel( + compiled.name, + num_warps, + compiled.metadata.shared, + hsaco_path, + str(compiled.asm.get("ttir", "")), + compute_capability, + 1, + 1, + 1, + ) + + def triton_call_lowering( ctx, kernel_fn: Callable,