flashinfer.comm._ulysses_lowp.merge_boundary_amax

flashinfer.comm._ulysses_lowp.merge_boundary_amax(grouped_amax: Tensor, gathered_descriptors: Tensor, *, rank: int, local_sequence: int, group: int, world_size: int) → Tensor

Overwrite this rank’s boundary slots with the cross-rank max merge.

gathered_descriptors is [P, B, H, 2] in group-rank order. For each of this rank’s (at most two) boundary groups, the final amax is the max of the partial amax from every rank whose interval intersects that group. The merge uses only max (exact, order-independent), so the result is bit-identical on every participating rank. Returns grouped_amax (modified in place) for convenience.

Parameters:

  • grouped_amax – FP32 local maxima [B, H, slots(local_sequence, group)].

  • gathered_descriptors – FP32 [world_size, B, H, 2] descriptors in group-rank order.

  • rank – Sequence-shard rank in [0, world_size).

  • local_sequence – Number of tokens per source shard, including zero padding.

  • group – Positive token count per quantization group.

  • world_size – Ulysses group size, one of 2, 4, 6 or 8; 6 is supported structurally but has not been validated.

Returns:

The input grouped_amax tensor, with boundary slots updated in place.