flashinfer.comm._ulysses_lowp.k_sum_v_amax

flashinfer.comm._ulysses_lowp.k_sum_v_amax(k: Tensor, v: Tensor, *, out: Tuple[Tensor, Tensor] | None = None, enable_pdl: bool | None = None) → Tuple[Tensor, Tensor]

Compute local FP32 K sum and V absolute maximum over sequence.

Both inputs use canonical NHD [B, L, H, D] storage. The returned [B, H, D] tensors are local statistics; callers must synchronize them across the Ulysses process group before quantizing K or V.

The reduction is two-stage sequence-parallel: stage 1 reduces fixed KSUM_CHUNK_TOKENS-token chunks into fp32 partials — the [B, H, ceil(L/256), D] fp32 workspaces are allocated here and passed to the kernel — and stage 2 combines the chunk partials in FIXED ascending chunk order, so results are deterministic (bit-identical run to run). NOTE: the fp32 k_sum association differs from a single-pass sequential sum by ULPs (the one deliberate bit change of the two-stage launch-occupancy fix); v_amax is max-reduced (order-independent) and stays byte-identical to the single-pass form.

Parameters:

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

  • out – Optional pair of contiguous FP32 [B, H, D] output tensors on the input device.

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

Returns:

Pair (k_sum, v_amax) of FP32 CUDA tensors of shape [B, H, D].