flashinfer.comm._ulysses_lowp.k_grouped_amax

flashinfer.comm._ulysses_lowp.k_grouped_amax(k: Tensor, k_mean_global: Tensor, *, rank: int, world_size: int, used_sequence: int | None = None, enable_pdl: bool | None = None) → Tensor

Per-touched-global-group abs(K - global mean) partial amax, [B, H, slots(L,64)].

used_sequence is the live global row count when the caller’s packed sequence carries zero-filled tail padding in rows [used, S). Unlike Q and V, whose amax is taken on the raw values (a zero row contributes 0 and cannot raise any group’s max), the K amax is taken on the MEAN-SUBTRACTED Kc = K - k_mean_global: a zero padding row contributes |0 - k_mean| = |k_mean|, which under the smooth-k regime (large channel means, small deviations – the very regime smooth_k exists for) dominates the live rows’ amax and silently inflates the tail group’s scale. The CUDA kernel has no notion of used, so when padding is present this recomputes the partial live group’s slot over the live rows only, with the exact pinned math (fp32 convert, subtract, abs, order-independent max, 1e-7 floor), so the pack kernel’s per-group scales see the corrected value. With used_sequence is None or used == S nothing is touched, keeping the no-padding path bit-identical to the original kernel output.

Tail padding may span many trailing K groups. Only ONE group can mix live and padding rows – the partial live group (used-1)//64 – and only there does the zero padding pollute a scale that live rows consume. Fully-padded groups beyond it are never read by the compute (rows are sliced to used); their kernel-produced values are deterministic don’t-cares.

Parameters:

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

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

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

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