feat: add indexed attention training kernels

README.md

@@ -13,6 +13,7 @@ CUDA kernels for Prime Intellect training stacks, shipped as one wheel, `prime-k
     │   ├── __init__.py       # Python surface: op wrappers, fake tensors
     │   ├── mxfp8.py
     │   └── csrc/             # the C++/CUDA sources compiled into prime_kernels.flash_moe._C
+    ├── indexed_attention/    # Python-only TileLang indexed GQA forward + backward
     ├── mxfp8_moe/            # Python-only MXFP8 MoE runtime kernels
     └── rmsnorm/
         ├── __init__.py
@@ -55,6 +56,10 @@ neither built nor shipped in the wheel, and the registry does not list it.
 `flash_moe` is currently dormant in prime-rl. `mxfp8_moe` provides differentiable MXFP8
 grouped GEMM and MXFP8 expert-parallel transport. It is registered as Python-only because
 it orchestrates PyTorch and torchao kernels rather than compiling a `_C` extension here.
+`indexed_attention` provides differentiable grouped-query attention over an explicit token
+selection for each query. Its TileLang kernels compute selection scores and radix selection
+as well as attention, and accept different query and KV lengths so the caller can gather KV
+for context parallelism without gathering queries.
 
 ## Installing
 
@@ -99,7 +104,9 @@ cxx-std = 20
 
 For a Python-only kernel, set `python-only = true`, omit `ops` and `sources`, and expose
 the differentiable Python surface from `__init__.py`. Optional import requirements belong
-in the manifest's `requires` list so `is_available()` fails during setup.
+in the manifest's `requires` list so `is_available()` fails during setup. Python-only ops
+may use `torch.library.custom_op`; register fake and autograd implementations so they remain
+visible to `torch.compile` and training.
 
 Whatever the kernel requires of its inputs — block sizes, alignments, layouts — belongs
 here, not in the caller: `TORCH_CHECK` it in the binding, and export the constants

prime_kernels/__init__.py

@@ -21,6 +21,7 @@
 ]
 
 _SPECS: dict[str, KernelSpec] = _load_manifest(Path(__file__).parent)
+_LOADED_MODULES: dict[str, ModuleType] = {}
 
 KERNELS: tuple[str, ...] = tuple(_SPECS)
 
@@ -61,10 +62,16 @@ def is_available(name: str, device: int | None = None) -> bool:
 
 
 def load(name: str, device: int | None = None) -> ModuleType:
+    try:
+        return _LOADED_MODULES[name]
+    except KeyError:
+        pass
     reason = unavailable_reason(name, device)
     if reason is not None:
         raise RuntimeError(reason)
-    return importlib.import_module(spec(name).module)
+    module = importlib.import_module(spec(name).module)
+    _LOADED_MODULES[name] = module
+    return module
 
 
 def status(device: int | None = None) -> dict[str, str]:

prime_kernels/indexed_attention/__init__.py

@@ -0,0 +1,4 @@
+from prime_kernels.indexed_attention.forward import indexed_attention, unsupported_shape_reason
+from prime_kernels.indexed_attention.selection import select_indexed_blocks
+
+__all__ = ["indexed_attention", "select_indexed_blocks", "unsupported_shape_reason"]