flashinfer.comm._ulysses_lowp.k_boundary_minmax

flashinfer.comm._ulysses_lowp.k_boundary_minmax(k: Tensor, *, rank: int, world_size: int, used_sequence: int | None = None) → Tensor

Per-channel raw-K min/max of this rank’s two K boundary slices.

Returns [B, H, 2, 2, D] fp32: dim2 is the boundary slot (0 = first touched global 64-group’s slice, 1 = last touched group’s; when the rank touches a single group both slots describe the same slice, and the downstream max-combine is idempotent). dim3 is (min, max).

Stats-protocol 2 gathers this INSTEAD of the mean-dependent K boundary amax: raw-K min/max needs no global mean, so it rides in AllGather #1, and after k_mean_global is derived every rank reconstructs each boundary group’s exact abs(K - mean) amax locally (see derive_k_boundary_amax()), eliminating the second collective.

fp32 transport is exact: every BF16/FP16 value converts to fp32 without rounding, and min/max are selections. used_sequence applies the live-rows rule one stage earlier: rows [used, S) are zero-filled padding whose inclusion would contribute min=0/max=0 and hence |0 - mean| = |mean| after the derive – exactly the tail-group pollution the k_grouped_amax fix removes. A slice that is entirely padding gets the sentinels min=+inf, max=-inf; the derive turns those into -inf before its 1e-7 floor, i.e. the same “this rank contributes nothing” value the two-collective path used.

Padding may span multiple groups or whole ranks. Empty boundary slices contribute only the amax floor after derivation.

Parameters:

  • k – BF16/FP16 CUDA K shard of shape [B, L, H, D], with the same layout and alignment requirements as Q.

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

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

  • used_sequence – Live global token count. None uses all L * world_size rows. Any remaining input rows must be zero-filled.

Returns:

FP32 CUDA boundary min/max of shape [B, H, 2, 2, D].