Figure 4 · Token-sharded prefill

Before and after: shard the attention by token, not by head

Toy example: 4 GPUs, 8 query heads, 16 prompt tokens. K3.1 Pro at TP8: 8 GPUs, 96 heads (12 per GPU). Not to scale.

Method: sparse_mla_shard_q (on by default for prefill steps of at least 1,024 tokens). GPU r takes prompt pieces r and 2W − 1 − r. TTFT: 8× GB300, TP8, one request at a time.

heads 0–1 heads 2–3 heads 4–5 heads 6–7 repeated or padded work