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 function

splitk_num_partitions

def splitk_num_partitions[config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_]](ws_num_partitions: UInt32) -> UInt32

The split-K partition count P this CTA must divide its KV range by.

Single source of truth for the two split-K mechanisms. All four FA4 warps (load, mma, softmax, correction) feed this to splitk_window and MUST derive the same window: a disagreement makes the producer over- or under-fill relative to its consumers, which HANGS rather than producing a wrong number.

Cluster/DSMEM split-K bakes P into the launch cluster, so it is the comptime config.splitk_partitions (or cluster_dim.x once the cluster dimension becomes dynamic).
Workspace (traditional/unfused) split-K keeps config.splitk_partitions == 1 -- no launch cluster -- and carries P at runtime by over-launching grid.x; ws_num_partitions is that count.

Parameters:

config (FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_]): The FA4 config, supplying the comptime cluster partition count.

Args:

ws_num_partitions (UInt32): Runtime partition count for the workspace scheme. Ignored when config.splitk_partitions > 1. Its 1 default makes this a no-op for a caller with no split-K at all.

Returns:

UInt32