flashinfer.comm._ulysses_lowp.q_grouped_amax

flashinfer.comm._ulysses_lowp.q_grouped_amax(q: Tensor, *, rank: int, world_size: int, enable_pdl: bool | None = None) → Tensor

Per-touched-global-group abs(Q) partial amax, [B, H, slots(L,32)].

Valid slots are [0, touched); the surplus allocation stays zero.

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.

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

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

Returns:

FP32 CUDA partial maxima of shape [B, H, slots(L, 32)].