flashinfer.comm._ulysses_lowp.unpack_for_sage

flashinfer.comm._ulysses_lowp.unpack_for_sage(recv_u8: Tensor, *, batch_size: int, local_sequence: int, local_heads: int, head_dim: int, world_size: int, aligned: bool | None = None, scale_sequence: int | None = None, out: Tuple[Tensor, Tensor, Tensor, Tensor, Tensor] | None = None, enable_pdl: bool | None = None) → Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]

Unpack source chunks into pre-quant compute inputs on the global grid.

Returns contiguous logical Q/K [B,S,h,D], globally packed V [B,D,h,ceil(S/64)*64] with a zero global tail, and global-grid Q/K scale tensors at the requested consumer width. Extra Q scale slots outside the live prefix contain initialized finite values.

The Sage kernel derives its per-head scale stride from the rows it is given, so a caller attending over a live prefix used < S must hand it scales of exactly scale_widths(used) slots. Pass scale_sequence=used and the kernel emits the scale tensors at that width directly (byte-identical to the full-width tensors’ leading slots), instead of the caller narrowing them with a copy per layer. None keeps the full logical width. Q/K/V are always emitted for the full logical sequence.

aligned=True asserts L % 128 == 0; None and False accept any legal shard. All values use the general receiver. Extra Q scale slots outside the live prefix are initialized but need not be zero.

Parameters:

  • recv_u8 – Contiguous CUDA uint8 source-major buffer [world_size, chunk_bytes] received by all-to-all.

  • batch_size – Positive batch size B.

  • local_sequence – Number of tokens per source shard, including zero padding.

  • local_heads – Heads owned by this destination after the all-to-all, H // world_size.

  • head_dim – Head dimension, 64 or 128; must agree on every rank.

  • world_size – Ulysses group size, one of 2, 4, 6 or 8; 6 is supported structurally but has not been validated.

  • aligned – True additionally requires local_sequence % 128 == 0; False and None accept arbitrary legal shards.

  • scale_sequence – Live prefix length determining Q/K scale widths; None uses the full global sequence.

  • out – Optional tuple (q8, k8, v8, q_scale, k_scale) of contiguous tensors on the receive device, with shapes and dtypes matching the return value.

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

Returns:

INT8 Q/K [B, S, h, D], FP8 E4M3 V [B, D, h, ceil(S/64)*64], and FP32 scales [B, h, width]. V scale remains with the caller and must be passed separately to attention.

Return type:

Tuple (q8, k8, v8, q_scale, k_scale)