For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo package

sm100

NVIDIA SM100 (Blackwell) attention kernels.

Covers MHA (flash-attention v4) and MLA (multi-head latent attention) for both prefill and decode, including FP8 and block-scaled quantization variants.

Packages

• ​mha_depth512: Re-exports the depth-512 multi-head attention kernel package for NVIDIA SM100 GPUs.

Modules

• ​attention: FA4 (Flash Attention 4) configuration for SM100 (Blackwell) kernels.
• ​attention_utils: Shared SM100 attention primitives used by both MHA and MLA kernels.
• ​correction_warp: Correction warp group logic for FA4 (SM100 Flash Attention).
• ​dispatch: Provides the SM100 (Blackwell) flash-attention host-side dispatch entry point that selects a 1Q or 2Q FA4 kernel configuration, builds the Q/K/V/O TMA descriptors and tile scheduler, and enqueues the kernel onto the device.
• ​kernel: Implements the SM100 (Blackwell) warp-specialized FlashAttention-4 multi-head attention kernel with the two-query (2Q) variant.
• ​load_warp: TMA load warp logic for FA4 (SM100 Flash Attention).
• ​mha_1q: Implements single-query (decode) multi-head attention for NVIDIA SM100 (Blackwell) GPUs using warp-specialized UMMA and tensor-memory (TMEM) accumulators.
• ​mla_decode_combine: MLA Decode Split-K Combine Kernel for SM100 (B200).
• ​mla_decode_dispatch:
• ​mla_decode_kv_bf16: Implements the SM100 (Blackwell) MLA decode kernel operating on a BF16 key-value cache.
• ​mla_decode_kv_fp8: Implements the SM100 MLA decode kernel variant that loads KV cache in FP8 and converts to BF16 in shared memory before MMA.
• ​mla_decode_qkv_fp8: Native FP8 MLA decode kernel for SM100 (B200).
• ​mla_decode_qkv_fp8_layout_g: Native FP8 MLA decode kernel for SM100 (B200): Layout G fold path.
• ​mla_decode_qkv_fp8_per_token_scale_rope_aware: SnapMLA FP8+BF16 MLA decode kernel for SM100 (B200).
• ​mla_decode_sparse: Implements the sparse MLA decoding kernel for SM100 (Blackwell) GPUs.
• ​mla_decode_sparse_kv_bf16: SM100 (B200) sparse MLA decode kernel with BF16 KV cache.
• ​mla_decode_sparse_kv_fp8:
• ​mla_decode_sparse_qkv_fp8: Native FP8 sparse MLA decode kernel for SM100 (B200).
• ​mla_decode_utils: Provides shared utilities for SM100 MLA decode attention kernels.
• ​mla_prefill: SM100 (Blackwell) MLA prefill kernels.
• ​mla_prefill_blockscale: Implements the SM100 (Blackwell) MLA prefill attention kernel with blockwise FP8 scaling, converting block-scaled FP8 K_rope tiles to BF16 in shared memory before the MMA stage.
• ​mla_prefill_generic:
• ​mla_prefill_per_token_scale: Per-token-scale MLA prefill kernel.
• ​mla_prefill_sparse: SM100 (B200) sparse MLA prefill kernel with BF16 KV cache.
• ​mla_prefill_sparse_kv_fp8: SM100 (B200) sparse MLA prefill kernel with FP8 KV cache.
• ​mla_prefill_sparse_qkv_fp8: SM100 (B200) native-FP8 sparse MLA prefill kernel.
• ​mla_prefill_sparse_utils: SM100 (B200) sparse MLA prefill shared surface (BF16 + FP8 KV variants).
• ​mla_prefill_utils:
• ​mma_warp: MMA warp logic for FA4 (SM100 Flash Attention).
• ​smem: Shared memory layout for SM100 attention kernels.
• ​softmax_warp: Softmax warp group logic for FA4 (SM100 Flash Attention).

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit