[FEAT] Lightning Indexer

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() {

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

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",
 ]

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",
+]

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)

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