IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

ws_p_ceiling

def ws_p_ceiling[sm_count: Int](raw_grid: UInt32) -> UInt32

Capture-invariant partition-count ceiling for the split-K crossover (M4): one GPC wave's SM-fill sm_count // raw_grid, flattened once raw_grid passes WS_RAW_GRID_CLAMP.

raw_grid is capture-invariant (see _bucket_ws), so this is a single fixed value per CUDA-graph capture.

Small raw_grid follows one wave exactly; large raw_grid -- where one wave alone is too small to be worth splitting to -- flattens at sm_count // WS_RAW_GRID_CLAMP. Written as one floor-divide rather than clamp(target, one_wave, sm_count) against a separate absolute target constant, because those two spellings are algebraically identical for any target of the form sm_count // k (max(sm_count // raw_grid, sm_count // k) == sm_count // min(raw_grid, k)) and the clamp form invites retuning target to a value that is NOT of that form -- which silently breaks the small-raw_grid passthrough below for every raw_grid < k, with no test to catch it.

Why flatten instead of scaling with one wave, and why 4: checked against a measured B200 ragged partition sweep (batch sweep at cache=131072). A flat MULTIPLE of one_wave reproduces the moderate-batch optimum (3*one_wave at batch=22, one_wave==6 -> 20, close to the measured workspace optimum) but ALSO scales up the already-good small-raw_grid case into a measured regression (that same 3x turns batch=4's one_wave==37 into 111, ~45% slower than 37). Clamping raw_grid instead leaves batch=4's 37 untouched -- it is exactly sm_count // 4 -- and lifts a small one-wave (batch=22's 6) to 37. The forced workspace P-sweep at batch=22 peaked at P=47, so 37 undershoots that peak by ~10-25% rather than matching it; 4 is the largest clamp that does not push batch=4 past its own measured optimum (P=49 costs batch=4 ~7-10% versus its P=37 peak). All four route arms in that sweep peaked in roughly the same absolute-P band, so this is shared across 1Q/WS-G/WS-E rather than per-route.

Returns:

UInt32