[FEAT] Lightning Indexer by Micky774 · Pull Request #606 · ROCm/TransformerEngine

Micky774 · 2026-06-01T20:55:02Z

Description

This PR contains:

Lightning indexer fwd/bwd hybrid implementations, using partial fused Triton kernels
Fused implementation of indexer+topk
API for DSA/HCA built on stubs for temporary staging and design (pending externally sourced kernels)
Benchmark scripts for the lightning indexer and indexer+topk
Rough integration of Triton kernels into JAX for ROCm

Note that in the PR that will eventually target/land on main, only the API changes and triton kernel integration will persist, with kernels being relocated elsewhere.

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Co-authored-by: Leon Ling <leon.ling@amd.com>

Strip Deep Sparse Attention from the sparse_attention package so it exposes only the lightning indexer: - Delete transformer_engine/jax/sparse_attention/dsa.py and its test (tests/jax/test_sparse_attention.py). - Drop DSA exports from sparse_attention/__init__.py; update the docstring. - Remove the test_sparse_attention.py invocation from ci/jax.sh. - Reword the now-stale DSA reference in test_indexer.py's reference oracle. The full indexer+DSA state is preserved on branch zain/lightning-indexer-dsa. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

github-actions · 2026-07-07T20:40:47Z

+
+
+mlir.register_lowering(_score_reduce_p, _score_reduce_lowering, platform="rocm")
+mlir.register_lowering(_score_reduce_p, _score_reduce_lowering, platform="cuda")


All three primitives (_score_reduce_p, _score_dscores_chunk_p, _score_topk_p) register lowerings for both platform="rocm" and platform="cuda" (also lines 389 and 967). The autotune configs, however, carry AMD-only meta-args (matrix_instr_nonkdim, waves_per_eu) and the kernel comments call out gfx950 / MI355X specifics. triton_extensions/utils.py::_compile_triton_hip filters HIP-only fields via dataclasses.fields(cb_hip.HIPOptions), but the CUDA compile path does not — so these config kwargs reach the kernel signature and fail at Triton compile time on NVIDIA.

Since this file is untested on CUDA and the tuning is ROCm-specific, I'd recommend restricting registration to platform="rocm" here (and at the other two sites). Per the ROCm-fork guidance, ROCm-only kernels should not silently claim CUDA support.

github-actions · 2026-07-07T20:40:51Z

+    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)


With T_s=4096, S_PAD = _next_pow2(T_s) = 4096, and the lowering picks the single-sort kernel (_score_topk_single_kernel) since S_PAD <= _SINGLE_SORT_MAX = 4096. That means the streaming _score_topk_kernel — which uses the 1D-encoded-t_enc layout, the multi-outer streaming buffer, and the AMD-backend workaround called out in its docstring — is not exercised anywhere in the test suite. It is the more complex of the two and the more likely place for a regression.

Please add one case that pushes into the streaming path (e.g. T_s >= 8192, or drop _SINGLE_SORT_MAX via monkeypatch in a dedicated test).

github-actions · 2026-07-07T20:40:55Z

+    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


test_lightning_indexer_topk_mode only checks shape and dtype; a kernel returning uninitialized memory (or all-zeros from a buggy scatter path) would still pass. Please add a validity check on the indices:

assert bool((idx >= 0).all()) and bool((idx < T_s).all())

Ideally also a set-based comparison against jax.lax.top_k (score-normalized, as in test_topk_matches_reference) — otherwise the module-level topk path is essentially untested for correctness.

github-actions · 2026-07-07T20:40:58Z

@@ -0,0 +1,24 @@
+# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.


This file is new in this PR (first commit 30faf3c1, June 2026). The AMD copyright range 2024-2026 implies AMD authorship starting in 2024, which isn't the case. Per the copyright-check convention (Yfirst must be ≥ the year the AMD copyright was first added), new files should carry a single year:

Suggested change

# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.

github-actions · 2026-07-07T20:41:01Z

@@ -0,0 +1,164 @@
+# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.


New file in this PR — AMD copyright range 2024-2026 should be single year 2026 (the year AMD authorship began for this file).

Suggested change

# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.

github-actions · 2026-07-07T20:41:02Z

@@ -0,0 +1,1016 @@
+# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.


New file in this PR — AMD copyright range 2024-2026 should be single year 2026.

Suggested change

# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.

github-actions · 2026-07-07T20:41:05Z

 from . import flax
 from . import quantize

+# AMD lightning-indexer / sparse-attention staging module.


This file is now modified by ROCm (the sparse_attention import + __all__ entry added here) but still only carries the upstream NVIDIA copyright header (line 1). Per CLAUDE.md rule #2, an AMD copyright line should be added above the NVIDIA line:

# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Leave the NVIDIA range untouched.

github-actions · 2026-07-07T20:41:17Z

Claude review

Reviewed the lightning-indexer surface added by this PR — the sparse-attention API (transformer_engine/jax/sparse_attention/), the Triton kernels (transformer_engine/jax/triton_extensions/indexer.py), the HIP compile path added to triton_extensions/utils.py, and tests/jax/test_indexer.py. Skipped the large upstream-merge churn in the diff (dev rebases pulling in HipKittens / gfx1250 / router / permutation changes) since it isn't the PR's contribution.

Highlights

The three indexer primitives register lowerings for both platform="rocm" and platform="cuda", but the autotune configs carry AMD-only meta-args (matrix_instr_nonkdim, waves_per_eu) and the kernels are tuned/documented for gfx950/MI355X. The CUDA path in _compile_triton_hip's sibling won't filter these. Recommend rocm-only registration.
Streaming _score_topk_kernel (the more complex of the two top-k paths) is not exercised by the test suite — T_s=4096 always lands in the single-sort path. Please add a case with T_s ≥ 8192.
test_lightning_indexer_topk_mode only checks shape/dtype; a garbage-index return would pass. Add index-range and score-based validation.

Copyright headers: 4 files need updates — three new files use 2024-2026 where single-year 2026 is correct, and transformer_engine/jax/__init__.py was modified but is missing the AMD line above the NVIDIA header. See inline comments.

github-actions · 2026-07-07T20:42:23Z

Claude Walkthrough

Intent. Introduces a JAX lightning indexer for sparse attention on ROCm - a low-rank query/key projection followed by a fused Triton score+ReLU+weighted-H-reduction (and an optional fused top-k). The op is a staging ground for a DeepSeek-style DSA/HCA indexer; the PR description notes that when this eventually targets main only the public API and Triton integration will remain, with kernels relocated. To make it work, triton_extensions/utils.py is extended with an AMD/HIP Triton compile path.

Key changes.

New transformer_engine.jax.sparse_attention package with the user-facing LightningIndexer Flax module and functional indexer / indexer_topk ops - transformer_engine/jax/sparse_attention/indexer.py:1, wired through transformer_engine/jax/__init__.py:37.
New Triton kernel suite in transformer_engine/jax/triton_extensions/indexer.py:1 - _score_reduce_kernel (fwd), _score_dscores_chunk_kernel (chunked bwd), and two top-k variants (_score_topk_kernel streaming + _score_topk_single_kernel single-sort for small S_PAD), plus a custom_vjp and JAX primitive plumbing.
HIP backend for the Triton compile path in transformer_engine/jax/triton_extensions/utils.py:312 (new _compile_triton_hip), plus new float8_e4m3fnuz / float8_e5m2fnuz dtype mappings at transformer_engine/jax/triton_extensions/utils.py:248.
triton_call_lowering grows callable-grid support and optional num_warps / num_stages / num_ctas overrides (transformer_engine/jax/triton_extensions/utils.py:482) so autotuned kernels whose grid depends on BLOCK_T / BLOCK_S launch correctly per config.
New correctness suite tests/jax/test_indexer.py:1 and CI wiring in ci/jax.sh:66.

Walkthrough.

transformer_engine/jax/sparse_attention/indexer.py: defines the reference math (C_q = Q@W_dq, H_q = einsum(C_q, W_uq), H_k = K@W_k, W_o = Q@W_w, O = einsum(relu(H_q · H_k), W_o)) as four jnp.einsum calls that lower to hipBLASLt bf16 GEMMs, then hands (H_q, H_k, W_o) to a Triton kernel that fuses the score matmul, ReLU, and per-(t,h) weighted reduction so the (B, oH, T, H, S) pre-ReLU score tensor never lands in HBM. LightningIndexer (Flax module) owns the four projection weights and, via its topk field, dispatches to indexer_topk (fused top-k) or indexer (returns the full (..., T, S) score tensor).
transformer_engine/jax/triton_extensions/indexer.py: hosts the kernels and their JAX primitives. The forward _score_reduce_kernel streams over H with H_k loop-invariant in registers. The backward is FlashAttention-style - residuals are just (Hq, Hk, W_o); a lax.scan over H-chunks calls _score_dscores_chunk_kernel, which recomputes scores, applies the ReLU mask, and produces dscores_chunk + dW_o_chunk while bounding peak HBM to H_CHUNK/H of the full (B, oH, T, H, T_s) tensor. Top-k uses a packed-uint64 buffer with the query index encoded in the top 8 bits so a single 1-D tl.sort yields per-query groups - the file comment explains this sidesteps an AMD-backend bug in tl.gather + tl.sort on 2-D uint64 tensors (gfx950, Triton 3.4.0). Two variants coexist: streaming (2K buffer, N sorts) for large T_s, single-sort (S_PAD <= 4096) when everything fits in registers. Autotune configs pin matrix_instr_nonkdim=16 and num_stages=1 in the bwd kernel because other combos crash LLVM codegen on Triton 3.7.0 / gfx950; a lowering-side pre-prune is applied because the runtime autotuner picks the fastest config by wall time and would otherwise "win" on an invalid config that does zero work and returns garbage indices.
transformer_engine/jax/triton_extensions/utils.py: gains a _compile_triton_hip that (a) strips the :sramecc+:xnack- suffix from the arch string, (b) sets supported_fp8_dtypes per-arch to mirror what parse_options would do (bypassed when constructing HIPOptions directly), (c) routes AMD compile-hint constexprs (matrix_instr_nonkdim, waves_per_eu, kpack, cluster_dims) from autotune configs into HIPOptions instead of kernel constexprs, and (d) writes the HSACO blob to a temp file (jaxlib HIP TritonKernel ctor takes a path, not bytes) which is retained in a module-level list for the process lifetime. triton_call_lowering grows a grid that may be Callable[[merged_kwargs], tuple], re-evaluated per autotune config so BLOCK_T / BLOCK_S-dependent grids launch at the right cdiv.
tests/jax/test_indexer.py: JIT-compiles the pure-einsum reference (stable HLO so top-k tie-breaking is deterministic), checks fwd (rel-err < 5e-3), fwd top-k score-equality against jax.lax.top_k at k in {32..1024} on T_s = 4*k_max, jax.grad of the hybrid vs the reference (rel-err < 5e-2), and the LightningIndexer module round-trip. An autouse fixture sets NVTE_INDEXER_DISABLE_AUTOTUNE=1 to pin each kernel to its first prune-valid config, avoiding a multi-minute sweep per test.

Testing. tests/jax/test_indexer.py is new - five test functions covering hybrid fwd, top-k, hybrid bwd, module fwd, and module top-k - and is invoked from ci/jax.sh under the fused-attn-agnostic default config.

Notes for reviewers.

Per the PR description this is staging: the kernels are expected to move elsewhere before a main-targeted PR - the durable surface is the sparse_attention API and the HIP Triton compile path in utils.py.
The _HSACO_TEMP_FILES list intentionally leaks temp files for the process lifetime because jaxlib holds the path, not the bytes; consider whether long-lived processes need a cleanup story.
Autotune correctness on ROCm relies on the lowering-side early_config_prune in _score_topk_lowering. If someone later refactors that path, an invalid config could silently win and return zero-index tensors - the comment there is worth preserving.
The bwd tolerance is deliberately loose (5e-2) with a TODO to tighten once per-grad error is characterized on-device.
HIP FP8 dtype selection is hardcoded per arch (gfx942, gfx950, gfx12*); new archs will fall through to (fp8e5,) and silently drop e4m3 support.

Generated by Claude. To request a code review, comment /claude review.

ipanfilo · 2026-07-20T14:21:23Z

ipanfilo · 2026-07-20T14:25:54Z

+
+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


This comment probably more suitable for specific kernel

ipanfilo · 2026-07-20T14:27:07Z

+from .utils import triton_call_lowering
+
+
+def _autotune_disabled():


Consider adding caching

ipanfilo · 2026-07-20T14:30:33Z

    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)


nit: move up, before test_layer

Micky774 added 17 commits May 6, 2026 18:38

Messy iter commit

55eb955

Dirty commit

1e18252

Trimmed API

fdfc1a4

Updated benchmarks

9335ef0

Trimmed fp8 fragments

358d326

Added initial bwd kernels

7c55255

Updated bwd pass to chunked hybrid kernel for mem consideration

949aee6

Added T-tiling to fused score top-k op

37a8563

Added initial API and tests

22b168f

Trimmed and streamlined

7c01b2b

Added benchmarks for indexer

31b9a8d

Minimized diff

2a27163

Removed comment

40bd8cc

Minimize diff

e27f4d5

Refactored package structure for new sparse attention components

30faf3c

Updated benchmark scripts, added indexer class

03a7447

Corrected import

223ca5a

leonling-ll mentioned this pull request Jun 12, 2026

[DRAFT] Optimize on lightnling Indexer Triton kernels #624

Closed

Micky774 and others added 8 commits June 16, 2026 16:18

Merge branch 'dev' into zain/lightning-indexer

51a26e9

Incorporated changes from gh-624

f6d86b8

Co-authored-by: Leon Ling <leon.ling@amd.com>

API trim, benchmark removal

d9a4dd8

Merge branch 'dev' into zain/lightning-indexer

9e367a0

Added hip options bypass and manual config pruning

8b10946

Trimmed indexer implementation and added CI run for tests

f2f9fc2

Streamline tests, add autotune bypass for tests

bdc5128

Micky774 changed the title ~~WIP Lightning Indexer + DSA/HCA API~~ [FEAT] Lightning Indexer Jul 1, 2026

Micky774 marked this pull request as ready for review July 7, 2026 20:27

Micky774 requested review from ipanfilo and wenchenvincent as code owners July 7, 2026 20:27

Micky774 requested a review from wangye805 as a code owner July 7, 2026 20:27

wenchenvincent added claude-review claude-summary labels Jul 7, 2026

github-actions Bot reviewed Jul 7, 2026

View reviewed changes

github-actions Bot removed the claude-review label Jul 7, 2026

github-actions Bot removed the claude-summary label Jul 7, 2026

Formatting

0b6840f

ipanfilo reviewed Jul 20, 2026

View reviewed changes



		mlir.register_lowering(_score_reduce_p, _score_reduce_lowering, platform="rocm")
		mlir.register_lowering(_score_reduce_p, _score_reduce_lowering, platform="cuda")

		@@ -0,0 +1,24 @@
		# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

	# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
	# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.

		@@ -0,0 +1,164 @@
		# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

		@@ -0,0 +1,1016 @@
		# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

		from .utils import triton_call_lowering


		def _autotune_disabled():

Uh oh!

Conversation

Micky774 commented Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 7, 2026

Claude review

Uh oh!

github-actions Bot commented Jul 7, 2026

Claude Walkthrough

Uh oh!

ipanfilo Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

ipanfilo Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

ipanfilo Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

ipanfilo Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Micky774 commented Jun 1, 2026 •

edited

Loading