flashinfer.comm._ulysses_lowp.local_stats

flashinfer.comm._ulysses_lowp.local_stats(q: Tensor, k: Tensor, v: Tensor, *, rank: int, world_size: int, used_sequence: int | None = None, enable_pdl: bool | None = None) → Tuple[Tensor, StatsContext]

Everything the single boundary-statistics AllGather must carry.

Returns a contiguous 1-D fp32 tensor (identical numel on every rank) and the context for finalize_stats(). Every shard sends the K per-channel sum and V per-channel amax, plus the Q boundary descriptors and the per-channel raw-K min/max of this rank’s two K boundary slices (min/max over LIVE rows only when used_sequence is given: a zero padding row would contribute min=0/max=0, and abs(0 - mean) would re-create the tail-group scale pollution at the derive).

Parameters:

  • q – BF16/FP16 CUDA Q shard of shape [B, L, H, D]; D is 64 or 128. The last dimension must be dense and every row 16-byte aligned.

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

  • v – BF16/FP16 CUDA V 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.

  • enable_pdl – Whether to enable Programmatic Dependent Launch. None selects according to device support.

Returns:

Tuple of a contiguous one-dimensional FP32 stats send tensor and its StatsContext; this function does not issue a collective.