Supported Ops and Hardware

Attention Variants

Op Hardware PyTorch-TPU API Includes Backward Pass Summary
splash_attention TPU ✅ ✅ Computes block-sparse Flash Attention (Splash Attention) for long-context Transformer training and inference.
batched_rpa TPU ✅ Computes paged attention (batched_ragged_paged_attention) over packed variable-length sequences and a paged KV cache for LLM inference serving.
causal_conv1d_gated_delta_rule TPU ✅ Fuses causal 1D convolution with Gated DeltaNet (GDN) or Kimi Delta Attention (KDA) linear attention state updates for inference in hybrid linear-attention models (e.g., Qwen 3.5, Kimi Linear).
DeepSeek Sparse Attention (DSA) deepseek_rope TPU ✅ Applies interleaved Rotary Position Embedding (RoPE) (optionally fused with RMSNorm or quantization) to queries and keys in DeepSeek-V4 attention layers.
proj_and_save_state TPU ✅ Projects hidden states into KV and score states and writes them to a paged state cache as the first stage of KV cache compression in DeepSeek-V4.
dsv4_compressor TPU ✅ Compresses windows of cached token states and writes the compressed representations into the KV cache for DeepSeek-V4.
lightning_indexer TPU ✅ Computes lightweight query-key relevance scores over a paged KV cache and selects the top-k KV positions per query token (invoking topk at the end) in GLM-5.2, DeepSeek-V3.2, and DeepSeek-V4.
topk TPU ✅ Selects the top-k indices per row (top_k) across variable-length score rows on TPU SparseCore; invoked as the selection stage of lightning_indexer and also available standalone.
csa_gather TPU ✅ Gathers and repacks compressed KV cache entries on TPU SparseCore for selected top-k token indices; invoked as the gather stage of sparse_mla and also available standalone.
sparse_mla TPU ✅ Computes Sparse Multi-Head Latent Attention (SparseMLA) over selected top-k compressed KV entries (invoking csa_gather first to gather them) for DeepSeek-V4 inference.
masked_dense_mla TPU ✅ Computes Multi-Head Latent Attention (MLA) restricted to a selected top-k subset of KV positions via a dense mask, used for shorter sequences in GLM-5.2, DeepSeek-V3.2, and DeepSeek-V4.
grouped_lowrank_proj TPU ✅ Fuses inverse RoPE rotation, quantization, and grouped low-rank output projection in the DeepSeek-V4 attention output path.
attention GPU, TPU ✅ Computes standard scaled dot-product attention (FlashAttention) for Transformer training and inference.
kda TPU ✅ Computes Kimi Delta Attention (KDA) (kimi_delta_attention), a linear attention mechanism with delta-rule state updates used in Kimi Linear.
flex_attention (experimental) GPU ✅ Computes scaled dot-product attention with user-defined score modification and masking functions (FlexAttention).
mla (experimental) TPU Computes Multi-Head Latent Attention (MLA) over compressed latent KV caches for DeepSeek-V2, DeepSeek-V3, and DeepSeek-R1 inference.

Mixture-of-Experts (MoE) Components

Op Hardware PyTorch-TPU API Includes Backward Pass Summary
ragged_dot GPU, TPU ✅ ✅ Performs grouped matrix multiplication (MegaBlocks; drop-in for jax.lax.ragged_dot, with gmm_v2 as one of its TPU backends) to compute expert projections over routed token groups.
ragged_gather TPU ✅ Gathers and permutes token activations into expert-grouped order during MoE token dispatch.
ragged_gather_reduce TPU ✅ Gathers expert outputs and performs a weighted reduction per token during the MoE combine (unpermute) step.
router_topk TPU ✅ Selects the top-k expert weights and indices per token from router scores for MoE routing.
fused_moe (experimental) TPU ✅ Executes a fused end-to-end expert-parallel MoE layer (routing, dispatch, expert projections, activation, and collectives) for LLM inference.
fused_moe_rs (experimental) TPU Fuses MoE expert computation with cross-device reduce-scatter into a single TPU kernel for inference.
routing_moe (experimental) TPU Implements a pipelined MoE layer that overlaps token gather/scatter with expert matrix multiplication.

Miscellaneous

Op Hardware PyTorch-TPU API Includes Backward Pass Summary
linear_softmax_xent_loss TPU ✅ ✅ Fuses the final linear projection with softmax cross-entropy loss (linear_softmax_cross_entropy_loss) to avoid materializing the full vocabulary logits tensor in memory.
mhc TPU ✅ Implements Manifold-Constrained Hyper-Connections (mHC) (mhc_pre, mhc_post, and mhc_fused_post_pre) for mixing multiple residual streams across sublayers in DeepSeek-V4.
gated_linear_unit GPU ✅ Computes a fused Gated Linear Unit (e.g., SwiGLU, GeGLU) for Transformer feed-forward networks.
normalization GPU ✅ Computes Layer Normalization (LayerNorm) or Root Mean Square Normalization (RMSNorm).
triangle_multiplication GPU ✅ Computes the triangle multiplicative update on 2D pair representations for structure prediction models such as AlphaFold 2 and AlphaFold 3.
fused_fp8_matmul (experimental) TPU ✅ Performs FP8 matrix multiplication with fused dynamic activation quantization and weight scaling.