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 module
ps_metadata
Builds persistent-scheduling work-partition metadata for MLA-prefill on the host.
Faithful port of the reference's kn_generate_ps_metadata
(v1_2_host.cuh:61-241) + WorkInfo/QTile (ps.h). Copies the reference
exactly; does not invent.
Even-division-of-split-units load balancer (the prefill-ps batch-flatness
lever): query tiles (qlen_granularity tokens x all heads) are built per batch
with causal effective-KV, counted in kvlen_granularity units, and total units
are split evenly across available_tgs persistent workgroups, splitting a
tile's KV across TGs when one TG's quota is exceeded (-> partials + reduce).
This is a pure host/CPU module (no GPU): the persistent kernel (S2) builds
the work partition on the host, uploads work_indptr/work_info to device
buffers, and the persistent grid consumes them. Tested standalone by
test/gpu/structured_kernels/test_ps_metadata.mojo.
comptime values
WORKINFO_DW
comptime WORKINFO_DW = 8
Structs
-
PsMetadata: Holds the persistent-scheduling work partition built on the host. -
QTile: Per ps.h:69-86; qo_start/qo_end are global TOKEN offsets.
Functions
-
build_ps_metadata: Port ofget_ps_metadata_v1_2_host(v1_2_host.cuh:265-314): the host wrapper that GCD-clusters heads across TGs, then calls the per-clusterkn_generate_ps_metadataand concatenates. -
build_uniform: Build metadata for uniform-seqlen self-attention, FP8 causal MLA-prefill. -
ceil_div: Computes the ceiling ofxdivided byyfor signed 32-bit integers. -
kn_generate_ps_metadata: Faithful port of v1_2_host.cuh:61-241 (single-cluster: cluster_id=0, current_work_idx=0 -- the nkv=1 prefill case). -
pack_dword: Packs two 16-bit halves into one 32-bit word as(high<<16) | (low & 0xFFFF).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!