flashinfer.comm._ulysses_lowp.quant_qkv_pack

flashinfer.comm._ulysses_lowp.quant_qkv_pack(q: Tensor, k: Tensor, v: Tensor, k_mean_global: Tensor, q_amax_final: Tensor, k_amax_final: Tensor, v_scale_global: Tensor, *, rank: int, world_size: int, out: Tensor | None = None, enable_pdl: bool | None = None) → Tensor

Fused V2-G quantize-and-pack convenience entry point.

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.

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

  • q_amax_final – Contiguous FP32 [B, H, slots(L, 32)] Q maxima after the cross-rank boundary merge.

  • k_amax_final – Contiguous FP32 [B, H, slots(L, 64)] centered-K maxima after the cross-rank boundary merge.

  • v_scale_global – Contiguous FP32 [B, H, D] global per-channel V divisor, global_amax / 2.25; zero-amax channels keep scale 0 and quantize to zero.

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

  • out – Optional contiguous uint8 [world_size, chunk_bytes] output on the input device. Scale slots and alignment padding are initialized on every call.

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

Returns:

Contiguous uint8 CUDA payload of shape [world_size, chunk_bytes].