flashinfer.comm._ulysses_lowp.derive_k_boundary_amax

flashinfer.comm._ulysses_lowp.derive_k_boundary_amax(grouped_amax: Tensor, gathered_minmax: Tensor, k_mean_global: Tensor, *, rank: int, local_sequence: int, world_size: int) → Tensor

Overwrite this rank’s K boundary slots with the derived cross-rank amax.

gathered_minmax is [P, B, H, 2, 2, D] fp32 in group-rank order (from AllGather #1 under stats protocol 2). For each of this rank’s (at most two) boundary groups g and every rank r touching g:

contrib(r, g) = max_d max( rn(maxK[r,g,d] - m[d]),

rn(m[d] - minK[r,g,d]) ) floored at 1e-7

which is bit-equal to r’s abs(K - m) partial amax over its slice of g: fp32 round-to-nearest subtraction is monotone in its tensor argument, so max_t rn(k_t - m) = rn(max_t k_t - m) and likewise for the min side, and max_t |rn(k_t - m)| = max(rn(maxK - m), rn(m - minK)) exactly. The final slot value is max_r contrib(r, g) – the same values the retired AllGather #2 merge produced, computed from identical gathered inputs with identical exact ops on every rank, hence bit-identical group-wide without a second collective. A padding-only slice’s sentinels (min=+inf, max=-inf) yield -inf before the floor, i.e. contribute 1e-7 exactly as the two-collective path did. Non-crossing boundary groups (single toucher) reduce to the kernel’s own partial, so the unconditional overwrite is bit-neutral there.

Parameters:

  • grouped_amax – FP32 [B, H, slots(local_sequence, 64)] local centered-K maxima to update in place.

  • gathered_minmax – FP32 [world_size, B, H, 2, 2, D] raw-K boundary min/max in group-rank order.

  • k_mean_global – Contiguous [B, H, D] global K mean, rounded to the input dtype.

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

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

  • 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.