Function shard_shapes 

pub(super) fn shard_shapes(
    shapes: &[(usize, usize)],
    worker_index: usize,
    worker_count: usize,
) -> Vec<(usize, usize)>

Partitions shapes into a per-worker shard using a stride assignment.

Worker worker_index receives shapes at positions worker_index, worker_index + worker_count, worker_index + 2*worker_count, … in the input slice order.

Why stride and not contiguous blocks?
The default warmup grid is ordered batch-major: {1,2,4,8,16,32} × {128,512,2048,8192}. Each consecutive group of four shapes belongs to one batch size, and within a group the sequence length grows monotonically. Stride assignment therefore spreads the work so each GPU receives one shape from each batch group at a different sequence length. The expensive _×8192 shapes land on different workers than each other (e.g. with 4 workers, worker 3 gets all 8192-seq shapes, which compile in parallel with the cheaper shapes on workers 0–2). Total wall-clock time is approximately the serial compile time for worker 3’s four shapes, compared to the serial time for all 24 - a rough 4× speedup on 4 GPUs.

Returns all shapes unchanged when worker_count ≤ 1.