flashinfer.msa_ops.msa_icp_comm_layout

flashinfer.msa_ops.msa_icp_comm_layout(world_size: int, heads_per_rank: int, capacity: int) → dict

Layout of the caller’s decode receive storage.

Returns a dict: storage_words (int32 words of one rank’s symmetric allocation), legacy_offset_bytes (0) and legacy_shape (int32 [3, W, hpr, capacity, 33]), tagged_offset_bytes and tagged_shape (int64 [3, W, hpr, capacity, 32]), and the int32 counter shapes generations_shape [3, W * hpr, capacity] and expected_shape [3, hpr, capacity].

The caller zeroes the storage and both counters once, on every rank, and never resets them. Every publish writes all T x H rows. Tagged rows need at least two sparse layers between uses of one slot set (including across steps); legacy rows are safe at any distance.