vllm.v1.attention.ops.triton_per_tensor_fp8_quant

Fused dynamic per-tensor FP8 quantization of up to three attention inputs.

Each [tokens, heads, dim] input (possibly strided, e.g. V as a view of the kv_b_proj output) gets its own amax and FP32 descale. A partial-amax pass and a reduce-and-quantize pass cover all inputs in two launches; the kernel boundary is the only global barrier, so no atomics or zero-initialization are needed and the ops are CUDA-graph safe. Compile-time specialization does not depend on the token counts (only the grid size does), so varying batch shapes do not recompile.

One grid covers all the inputs in each launch, so a program can land past the end of a smaller input. Such programs do no work and compute no address, and every load and store is masked regardless, so no out-of-range pointer is formed or dereferenced.

Adapted from the fused QKV quantizer in ROCm/ATOM (atom/model_ops/ triton_fused_qkv_quant.py, MIT license).

Functions:

• fused_per_tensor_fp8_quant –

  Quantize each input to FP8 with its own dynamic per-tensor scale.

fused_per_tensor_fp8_quant(*tensors, fp8_dtype=torch.float8_e4m3fn)

Quantize each input to FP8 with its own dynamic per-tensor scale.

Parameters:

• tensors

  (Tensor, default: () ) –

  One to three BF16/FP16 [tokens, heads, dim] tensors on the same device. Arbitrary strides are supported; token counts may differ and may be zero.

• fp8_dtype

  (dtype, default: float8_e4m3fn ) –

  The FP8 output dtype.

Returns:

• tuple[Tensor, ...] –

  (outputs, descales): contiguous FP8 tensors with the input shapes

• tuple[Tensor, ...] –

  and FP32 shape-[1] descales (amax / fp8_max, or 1e-6 for

• tuple[tuple[Tensor, ...], tuple[Tensor, ...]] –

  all-zero or empty inputs).

Source code in vllm/v1/attention/ops/triton_per_tensor_fp8_quant.py

def fused_per_tensor_fp8_quant(
    *tensors: torch.Tensor,
    fp8_dtype: torch.dtype = torch.float8_e4m3fn,
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]:
    """Quantize each input to FP8 with its own dynamic per-tensor scale.

    Args:
        tensors: One to three BF16/FP16 ``[tokens, heads, dim]`` tensors on the
            same device. Arbitrary strides are supported; token counts may
            differ and may be zero.
        fp8_dtype: The FP8 output dtype.

    Returns:
        ``(outputs, descales)``: contiguous FP8 tensors with the input shapes
        and FP32 shape-``[1]`` descales (``amax / fp8_max``, or 1e-6 for
        all-zero or empty inputs).

    """
    if not 1 <= len(tensors) <= 3:
        raise ValueError(f"Expected 1 to 3 tensors, got {len(tensors)}")
    first = tensors[0]
    for x in tensors:
        if x.ndim != 3 or x.shape[1] == 0 or x.shape[2] == 0:
            raise ValueError(
                f"Expected [tokens, nonzero heads, nonzero dim], got {x.shape}"
            )
        if x.device != first.device or x.dtype not in (
            torch.bfloat16,
            torch.float16,
        ):
            raise ValueError("Inputs must be BF16/FP16 tensors on the same device")

    outputs = tuple(
        torch.empty(x.shape, dtype=fp8_dtype, device=x.device) for x in tensors
    )
    scales = torch.empty(len(tensors), dtype=torch.float32, device=first.device)
    numels = [x.numel() for x in tensors]
    partial = torch.empty(
        (len(tensors), _PARTS), dtype=torch.float32, device=first.device
    )

    # Unused slots alias the first tensor; their programs never run.
    pad = 3 - len(tensors)
    xs = (*tensors, *(first,) * pad)
    ys = (*outputs, *(outputs[0],) * pad)
    ns = (*numels, *(0,) * pad)
    shapes = tuple(tuple(x.shape[1:]) for x in xs)
    strides = tuple(x.stride() for x in xs)

    _fused_amax_kernel[(_PARTS, len(tensors))](
        *xs, partial, *ns, shapes, strides, _PARTS, _BLOCK, num_warps=4
    )
    _fused_quant_kernel[(max(1, triton.cdiv(max(numels), _BLOCK)), len(tensors))](
        *xs,
        *ys,
        partial,
        scales,
        *ns,
        shapes,
        strides,
        _PARTS,
        _BLOCK,
        torch.finfo(fp8_dtype).max,
        _MIN_DESCALE,
        num_warps=4,
    )
    return outputs, tuple(scales[i : i + 1] for i in range(len(tensors)))