H100 (Hopper)¶

The same suites as on the RTX 3090 pages, run on an NVIDIA H100, plus the measurements behind the device-specific defaults for Hopper. Times are medians in milliseconds; inputs are [batch, seq, heads, head_dim] and bf16 unless stated otherwise. OOM: out of device memory. fail†: XLA found no kernel configuration that fits during autotuning (NOT_FOUND), so the case did not run. unsup.: the implementation rejects the shape or backend. error: another failure (the JSON row has the message). —: not measured.

Setup¶

Device-specific defaults¶

flex_attention and dot_product_attention look up the GPU's compute capability in probjax.nn.layers.attention_device when every kernel argument is at its default. GPUs without an entry (RTX 3090 and other non-sm_90 devices) keep the generic rules. On sm_90:

implementation="triton" (per call), "cudnn" or "mosaic" (raise if they cannot compute the call exactly), PROBJAX_ATTENTION_CUDNN=0, PROBJAX_ATTENTION_MOSAIC=0 or PROBJAX_ATTENTION_DEVICE_DEFAULTS=0 (generic rules) opt out. The cuDNN path is not bitwise identical to the Triton kernels; outputs and gradients agree to bf16 precision (max |difference| ≤ 0.004 in tests/test_attention_device_defaults.py).

Why¶

The generic auto rule picks the fused backward whenever the two lengths differ by at most 4×; on H100 that is up to 1.92× slower than the split backward, which is never more than 1% slower than auto. cuDNN, which supports per-batch KV lengths, is faster than every Triton variant above the work thresholds. fwd+bwd, ms ("generic auto" is the dispatch before the sm_90 entry):

fwd+bwd per implementation and shape, bf16; the H100 default is the filled blue dot.

Few queries over long KV, measured with the H100 default:

Mosaic GPU¶

mha_flash runs JAX's Hopper FlashAttention-3 kernel (warpgroup MMA, TMA pipelining), vendored as pallas_kernels/kernels/attention_mgpu.py because upstream hard-codes float16 shared-memory buffers and rejects every bf16 call. The wrapper takes Lq ≠ Lkv and GQA and pads query length and head_dim exactly; with per-batch KV lengths it masks and skips padded keys, so any length works. Its default tiles are the best H100 configuration found (64 × 128 forward, backward dK/dV 128 × 64 and dQ 64 × 128, two compute warpgroups); deeper pipelines and 128 × 128 tiles with two warpgroups exceed shared memory at head_dim 128. Unmasked, fwd+bwd ms:

flex_attention(implementation="mosaic") uses it as a backend with KV lengths: KV tiles past a batch element's length are neither loaded nor computed in the forward, dK/dV programs past it store zeros without running their pipeline, rows without a valid key return zeros. With KV lengths, fwd+bwd ms:

For head_dim ≥ 64 Mosaic is up to 2.89× faster than the Triton kernels and close to cuDNN (Triton is faster at 4 × 16384 × 256, 8, 64 (KV lengths): few keys per query tile). cuDNN is faster where it applies, so the sm_90 default tries cuDNN first and lowers to Mosaic where cuDNN cannot compute the call (head_dim not a multiple of 8, or PROBJAX_ATTENTION_CUDNN=0). head_dim ≤ 32 is padded to 64 and therefore stays on Triton. In rows with a single valid key the exact dQ and dK are 0 and carry ~1e-1 rounding noise from the kernel's approximate exp2 (upstream approx_math); Triton's is ~1e-5. Masks other than KVLenMask and equal-length causal (local windows, segments, functional biases) are not in the Mosaic kernel yet.

Attention suite¶

bf16, batch 4, 8 heads, H100. Top row forward, bottom row forward + backward.

fp32, batch 4, 8 heads, H100. fp32 matmuls run in TF32 (JAX default precision).

fwd+bwd time of each implementation divided by the probjax default at the same shape; below the line is faster than probjax.

The probjax default (cuDNN-routed bf16) is 1.50×–1.75× faster than torch SDPA flash. fp32 has no flash path in torch; jax.nn XLA is the fastest at fp32 head_dim 128 (see known issues).

H100 vs RTX 3090¶

RTX 3090 time divided by H100 time, median over every row both GPUs ran (same shapes).

MoE dispatch¶

MoELayer with dropless routing, bf16, a large decoder MoE (d_model 512, 64 experts, top-8, expert width 128, shared expert 1024); train step = forward + gradients w.r.t. parameters and input, ms:

MoE train step by dispatch, H100.

"mosaic" runs jax's Hopper ragged-dot kernels (ragged_dot_mgpu, transposed_ragged_dot_mgpu for the weight gradient; wgmma + TMA) on the sorted dispatch; "pallas" resolves to it on sm_90 for 16-bit compute (the wgmma kernels need 16-bit operands; float32 runs the Triton grouped matmul). Per train step it is faster than fused_gsum (8% at 16,384, 14% at 131,072, 20% at 1,048,576 tokens); the forward alone is 56.3 vs 52.5 ms at 1,048,576 tokens, the gain is in the backward. Outputs and gradients agree with triton / sorted to bf16 precision (max relative difference ≤ 1.1%). The sorted dispatch keeps top-k copies of each token (N·K × D buffers), so it needs more activation memory than fused_gsum.

Against PyTorch¶

PyTorch 2.14.1 (CUDA 13), same inputs and cotangents as the probjax rows; every torch result is checked against the probjax reference (gradients agree to bf16 precision, max |difference| ≤ 0.06). torch.nn.attention.flex_attention is compiled (torch.compile, one compiled entry per shape) with a create_block_mask built from the KV lengths; SDPA is pinned to one backend with sdpa_kernel. SDPA flash takes no attention mask, so it cannot run the masked cases; the memory-efficient backend takes them only as a dense boolean mask. fwd+bwd, ms (gpu_attention_h100_device.json, labels h100_torch_*):

With head_dim ≥ 64 the probjax default is the fastest in every case (flex_attention 1.03×–1.62× slower). With head_dim 32 torch.nn.attention.flex_attention is 1.26×–1.61× faster than the probjax default (see known issues).

Known issues on H100¶

• Shared memory (fixed here). Large-SMEM GPUs kept the requested tiles unchecked; on H100 head_dim 256 (forward) and the fp32 head_dim 128 single-pass backward requested 327,680 and 262,144 bytes and failed to compile. Triton buffers streamed tiles once per pipeline stage on sm_90, which the sm_86 estimates missed; tiles are now shrunk only when they would not fit.
• Explicit backward_pass_impl="triton_fused" on strongly asymmetric lengths (fixed here). Matching the tile counts forced tiles below 16 rows, which fell back to Pallas interpret mode: 25 s instead of ~1 ms at 256 × 16384, and a wrong dQ at 32768 × 256. It now runs the split backward. "auto" never selected this configuration.
• Hopper MoE and MLA on unaligned shapes (fixed here). The Mosaic ragged matmul took float32 operands and row counts off its 8-row tiling and failed to lower; mla tiled a short non-power-of-two query or latent length (e.g. 48) as one block, which Triton cannot lower. Float32 now runs the Triton grouped matmul, rows are padded into the last group, and mla pads and masks unaligned lengths like mha.
• fp32 head_dim 128: jax.nn XLA (TF32 matmuls) is about 1.95× faster than the Triton kernel (11.6 vs 22.7 ms at T=4096); no fp32 rule yet.
• Sequential reference paths (implementation="scan" for Mamba and RLA) are 1.0–1.8× slower than on the RTX 3090 (latency-bound loops); the GPU defaults are 1.3–2.6× faster. The bf16 DiffusionTransformer train step at 256 tokens is slower too (4.8 vs 3.0 ms), at the size where launch overhead dominates.
• head_dim 32: torch.nn.attention.flex_attention is faster than every probjax path on the ISAB shapes above (inducing points over context rows and back); the cuDNN route and the Triton tiles are not tuned for head_dim 32 yet.
• Tests: pytest --gpu passes on H100. Tests that compare XLA float32 math against scipy/numpy or an equivalent formulation are marked float32_matmul and run in true float32 on GPU (tests/conftest.py; XLA's default is TF32); Pallas kernel-vs-reference tests, whose kernels use TF32 dots, compare at TF32-level tolerances (tests/_tolerance.py).