flashinfer.comm.quantized_all_reduce

flashinfer.comm.quantized_all_reduce(inp: Tensor, group: ProcessGroup, *, scale_group: int = 256, block_size: int | None = None, num_warps: int | None = None, max_num_blocks: int | None = None, p2p_phase3: bool | None = None, output: Tensor | None = None) → Tensor

FP8 quantized two-shot AllReduce with per-block scaling.

Quantizes activations to FP8 with per-block scales before P2P transfer via symmetric memory, halving inter-GPU bandwidth during the collective.

Requires SM90+ (Hopper or later) for FP8 E4M3 support and NVSwitch for symmetric memory. World size must divide block_size (default 4096), so world_size must be a power of 2 up to 16.

Best suited for tensors >= 2MB where NVLink bandwidth dominates. For smaller tensors, torch.ops.symm_mem.multimem_all_reduce_ or flashinfer.comm.allreduce_fusion may be faster due to lower fixed overhead.

Note

The first call allocates a symmetric memory buffer (collective across all ranks). Call once with the largest expected tensor during warmup to avoid re-allocation during inference.

Parameters:

  • inp – Input tensor (BF16, contiguous, numel divisible by 8).

  • group – Process group for the collective.

  • scale_group – Elements per scale group (default 256). Smaller values give better accuracy at marginal bandwidth cost.

  • block_size – Compute tile size. None = auto (4096 for <= 64MB, 16384 for larger).

  • num_warps – Warps per thread block. None = auto (16).

  • max_num_blocks – Max SMs to use. None = auto-tuned by tensor size.

  • p2p_phase3 – Phase 3 read strategy. None = auto (True for >= 32MB).

  • output – Optional pre-allocated output tensor. If None, allocates a new tensor.

Returns:

AllReduced output tensor (same shape/dtype as input).