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

Mojo module

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.

comptime values

FA4_FORCE_2Q

comptime FA4_FORCE_2Q = (get_defined_int[StringSpan("FA4_FORCE_2Q"), Int(0)]() != Int(0))

FA4_WS_POISON

comptime FA4_WS_POISON = (get_defined_int[StringSpan("FA4_WS_POISON"), Int(0)]() != Int(0))

logger

comptime logger = Logger(stdout, prefix=String(""), source_location=False)

SK_WS_KEYS_PER_PARTITION

comptime SK_WS_KEYS_PER_PARTITION = UInt32(384)

WS_RAW_GRID_CLAMP

comptime WS_RAW_GRID_CLAMP = UInt32(4)

WS_SWEEP_MAX_RAW_GRID

comptime WS_SWEEP_MAX_RAW_GRID = UInt32(24)

Functions

• ​mha_sm100_dispatch: Dispatches the SM100 FA4 flash-attention kernel for a prefill or decode workload.
• ​ws_p_ceiling: Capture-invariant partition-count ceiling for the split-K crossover: follows one GPC wave's SM-fill, holds the plateau through WS_SWEEP_MAX_RAW_GRID, then decays toward one partition.