| |
| |
|
|
| """Transformer Engine NVFP4 quantization helper. |
| |
| This file is intended as a customer-facing example for preparing KV tensors |
| for the KVFP4 attention kernel: |
| - BF16/FP16 K/V input |
| - packed E2M1 FP4 data from Transformer Engine |
| - E4M3 block scales in cuBLAS/cuDNN 128x4 tiled layout |
| - one FP32 tensor/global scale per tensor |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Tuple |
|
|
| import torch |
|
|
|
|
| NVFP4_BLOCK_SIZE = 16 |
| NVFP4_FP4_MAX = 6.0 |
| NVFP4_FP8_E4M3_MAX = 448.0 |
|
|
|
|
| @dataclass(frozen=True) |
| class Nvfp4QuantizedTensor: |
| """Packed NVFP4 tensor plus dequantization metadata. |
| |
| Attributes |
| ---------- |
| data : torch.Tensor |
| Packed E2M1 FP4 data from Transformer Engine. The last dimension is |
| half of the original logical last dimension because each byte stores |
| two FP4 values. |
| scale_128x4 : torch.Tensor |
| E4M3 block scales in cuBLAS/cuDNN 128x4 tiled rowwise storage. |
| global_scale : torch.Tensor |
| FP32 tensor/global dequant scale. |
| logical_scale_shape : tuple[int, int] |
| Logical 2D scale shape ``(rows, cols)`` before 128x4 swizzling. |
| original_shape : tuple[int, ...] |
| Original BF16/FP16 tensor shape before quantization. |
| """ |
|
|
| data: torch.Tensor |
| scale_128x4: torch.Tensor |
| global_scale: torch.Tensor |
| logical_scale_shape: Tuple[int, int] |
| original_shape: Tuple[int, ...] |
|
|
|
|
| def _round_up(x: int, multiple: int) -> int: |
| return ((int(x) + multiple - 1) // multiple) * multiple |
|
|
|
|
| def nvfp4_scale_128x4_offset( |
| row: torch.Tensor, |
| col: torch.Tensor, |
| scale_cols: int, |
| ) -> torch.Tensor: |
| """Return flat offsets for cuBLAS/cuDNN 128x4 rowwise scale storage. |
| |
| Parameters |
| ---------- |
| row : torch.Tensor |
| Logical row indices. |
| col : torch.Tensor |
| Logical scale-column indices. |
| scale_cols : int |
| Logical number of scale columns before padding to a multiple of 4. |
| |
| Returns |
| ------- |
| torch.Tensor |
| Flat offsets into the padded 128x4 tiled storage. |
| """ |
|
|
| tiles_n = _round_up(scale_cols, 4) // 4 |
| tile_m = row // 128 |
| tile_n = col // 4 |
| outer = row % 128 |
| inner = col % 4 |
| return ( |
| (tile_m * tiles_n + tile_n) * 512 |
| + (outer % 32) * 16 |
| + (outer // 32) * 4 |
| + inner |
| ) |
|
|
|
|
| def swizzle_nvfp4_scale_to_128x4( |
| scale: torch.Tensor, |
| *, |
| rows: int, |
| cols: int, |
| ) -> torch.Tensor: |
| """Convert TE logical rowwise scales to cuBLAS/cuDNN 128x4 tiled layout. |
| |
| Parameters |
| ---------- |
| scale : torch.Tensor |
| Logical rowwise scale tensor with at least shape ``[rows, cols]``. |
| rows : int |
| Number of logical rows to convert. |
| cols : int |
| Number of logical scale columns to convert. |
| |
| Returns |
| ------- |
| torch.Tensor |
| Scale tensor padded to ``round_up(rows, 128)`` by ``round_up(cols, 4)`` |
| and swizzled into 128x4 tiled storage. |
| """ |
|
|
| if scale.ndim != 2: |
| raise ValueError(f"scale must be 2D, got shape {tuple(scale.shape)}") |
|
|
| rows = int(rows) |
| cols = int(cols) |
| padded_rows = _round_up(rows, 128) |
| padded_cols = _round_up(cols, 4) |
| if scale.shape[0] < rows or scale.shape[1] < cols: |
| raise ValueError( |
| "scale is smaller than the requested logical shape: " |
| f"got {tuple(scale.shape)}, need at least {(rows, cols)}" |
| ) |
|
|
| logical = scale[:rows, :cols].contiguous() |
| if logical.shape != (padded_rows, padded_cols): |
| logical = torch.nn.functional.pad( |
| logical.to(torch.float32), |
| (0, padded_cols - cols, 0, padded_rows - rows), |
| ).to(scale.dtype) |
| swizzled = torch.empty_like(logical) |
|
|
| row = torch.arange(padded_rows, device=scale.device, dtype=torch.int64)[:, None] |
| col = torch.arange(padded_cols, device=scale.device, dtype=torch.int64)[None, :] |
| offset = nvfp4_scale_128x4_offset(row, col, padded_cols).reshape(-1) |
| swizzled.reshape(-1)[offset] = logical.reshape(-1) |
| return swizzled |
|
|
|
|
| def nvfp4_global_scale_from_amax(amax: torch.Tensor) -> torch.Tensor: |
| """Compute TE NVFP4 tensor/global dequant scale from rowwise amax. |
| |
| Parameters |
| ---------- |
| amax : torch.Tensor |
| Rowwise absolute maxima returned by Transformer Engine. |
| |
| Returns |
| ------- |
| torch.Tensor |
| FP32 global scale equal to ``amax / (448 * 6)``. |
| """ |
|
|
| return amax.to(torch.float32) / (NVFP4_FP8_E4M3_MAX * NVFP4_FP4_MAX) |
|
|
|
|
| def _import_te_nvfp4_quantizer(): |
| try: |
| from transformer_engine.pytorch.tensor import NVFP4Quantizer |
| except Exception as exc: |
| raise RuntimeError( |
| "Transformer Engine NVFP4 quantization is unavailable. Install a " |
| "Transformer Engine build with its PyTorch dependencies, including " |
| "FlashAttention v3 when required by that TE build." |
| ) from exc |
| return NVFP4Quantizer |
|
|
|
|
| def quantize_bf16_to_nvfp4_128x4(x: torch.Tensor) -> Nvfp4QuantizedTensor: |
| """Quantize a BF16/FP16 tensor to NVFP4 using Transformer Engine. |
| |
| TE returns rowwise scales in logical padded layout. This helper returns |
| the scales in physical 128x4 tiled storage, so the attention kernel can |
| load them with ``nvfp4_scale_128x4_offset``. |
| |
| Parameters |
| ---------- |
| x : torch.Tensor |
| CUDA BF16 or FP16 tensor. The last dimension must be divisible by 16, |
| and the flattened row dimension ``prod(x.shape[:-1])`` must also be |
| divisible by 16. |
| |
| Returns |
| ------- |
| Nvfp4QuantizedTensor |
| Packed FP4 data, 128x4-swizzled block scales, global scale, and shape |
| metadata needed by the KVFP4 attention kernel or by reference |
| dequantization. |
| """ |
|
|
| if not x.is_cuda: |
| raise ValueError("NVFP4 quantization requires a CUDA tensor") |
| if x.dtype not in (torch.bfloat16, torch.float16): |
| raise TypeError(f"x must be bf16 or fp16, got {x.dtype}") |
| if x.ndim < 2: |
| raise ValueError(f"x must have at least 2 dimensions, got {x.ndim}") |
| if x.shape[-1] % NVFP4_BLOCK_SIZE != 0: |
| raise ValueError( |
| f"last dimension must be divisible by {NVFP4_BLOCK_SIZE}, got {x.shape[-1]}" |
| ) |
|
|
| rows = 1 |
| for dim in x.shape[:-1]: |
| rows *= int(dim) |
| if rows % NVFP4_BLOCK_SIZE != 0: |
| raise ValueError( |
| "flattened row dimension must be divisible by " |
| f"{NVFP4_BLOCK_SIZE}, got {rows}" |
| ) |
|
|
| NVFP4Quantizer = _import_te_nvfp4_quantizer() |
| quantizer = NVFP4Quantizer(rowwise=True, columnwise=False) |
| qx = quantizer.quantize(x.contiguous()) |
| meta = qx.get_metadata() |
|
|
| data = meta["rowwise_data"] |
| if data.dtype is not torch.uint8: |
| data = data.view(torch.uint8) |
| logical_scale = meta["rowwise_scale_inv"] |
| amax = meta["amax_rowwise"] |
| scale_cols = int(x.shape[-1]) // NVFP4_BLOCK_SIZE |
| scale_128x4 = swizzle_nvfp4_scale_to_128x4( |
| logical_scale, |
| rows=rows, |
| cols=scale_cols, |
| ) |
| global_scale = nvfp4_global_scale_from_amax(amax).contiguous() |
|
|
| return Nvfp4QuantizedTensor( |
| data=data, |
| scale_128x4=scale_128x4, |
| global_scale=global_scale, |
| logical_scale_shape=(rows, scale_cols), |
| original_shape=tuple(int(v) for v in x.shape), |
| ) |
|
|
|
|
| def quantize_kv_bf16_to_nvfp4_128x4( |
| k: torch.Tensor, |
| v: torch.Tensor, |
| ) -> tuple[Nvfp4QuantizedTensor, Nvfp4QuantizedTensor]: |
| """Quantize BF16/FP16 K and V tensors independently for KVFP4 attention. |
| |
| Parameters |
| ---------- |
| k : torch.Tensor |
| CUDA BF16 or FP16 K tensor. |
| v : torch.Tensor |
| CUDA BF16 or FP16 V tensor. |
| |
| Returns |
| ------- |
| tuple[Nvfp4QuantizedTensor, Nvfp4QuantizedTensor] |
| Quantized K and V tensors with independent scales. |
| """ |
|
|
| return quantize_bf16_to_nvfp4_128x4(k), quantize_bf16_to_nvfp4_128x4(v) |
|
|
|
|
| def dequantize_nvfp4_128x4_to_bf16( |
| qx: Nvfp4QuantizedTensor, |
| *, |
| include_global_scale: bool = True, |
| ) -> torch.Tensor: |
| """Reference dequantization for validation. |
| |
| This mirrors the kernel contract: |
| x = e2m1 * E4M3_block_scale_1x16 * FP32_global_scale |
| |
| Parameters |
| ---------- |
| qx : Nvfp4QuantizedTensor |
| Quantized tensor returned by ``quantize_bf16_to_nvfp4_128x4``. |
| include_global_scale : bool, optional |
| If True, multiply by ``qx.global_scale`` after applying per-block |
| scales. |
| |
| Returns |
| ------- |
| torch.Tensor |
| BF16 tensor with shape ``qx.original_shape``. |
| """ |
|
|
| data = qx.data if qx.data.dtype is torch.uint8 else qx.data.view(torch.uint8) |
| if data.shape[-1] * 2 != qx.original_shape[-1]: |
| raise ValueError( |
| "packed data last dimension does not match original shape: " |
| f"{data.shape[-1]} packed vs {qx.original_shape[-1]} logical" |
| ) |
|
|
| rows, scale_cols = qx.logical_scale_shape |
| logical_dim = int(qx.original_shape[-1]) |
| if scale_cols * NVFP4_BLOCK_SIZE != logical_dim: |
| raise ValueError( |
| "logical scale columns do not match original last dimension: " |
| f"{scale_cols} scale cols vs dim {logical_dim}" |
| ) |
|
|
| fp4_lut = torch.tensor( |
| [ |
| 0.0, |
| 0.5, |
| 1.0, |
| 1.5, |
| 2.0, |
| 3.0, |
| 4.0, |
| 6.0, |
| -0.0, |
| -0.5, |
| -1.0, |
| -1.5, |
| -2.0, |
| -3.0, |
| -4.0, |
| -6.0, |
| ], |
| dtype=torch.float32, |
| device=data.device, |
| ) |
| packed = data.reshape(rows, logical_dim // 2) |
| lo = packed & 0x0F |
| hi = packed >> 4 |
| values = torch.empty((rows, logical_dim), dtype=torch.float32, device=data.device) |
| values[:, 0::2] = fp4_lut[lo.long()] |
| values[:, 1::2] = fp4_lut[hi.long()] |
|
|
| row = torch.arange(rows, device=data.device, dtype=torch.int64)[:, None] |
| col = torch.arange(scale_cols, device=data.device, dtype=torch.int64)[None, :] |
| offset = nvfp4_scale_128x4_offset(row, col, scale_cols) |
| scale_u8 = qx.scale_128x4.reshape(-1)[offset.reshape(-1)].reshape(rows, scale_cols) |
| scale = scale_u8.view(torch.float8_e4m3fn).to(torch.float32) |
| scale = scale.repeat_interleave(NVFP4_BLOCK_SIZE, dim=1) |
| out = values * scale |
| if include_global_scale: |
| global_scale = qx.global_scale.reshape(-1)[0].to(torch.float32) |
| out = out * global_scale |
| return out.reshape(qx.original_shape).to(torch.bfloat16) |
|
|
|
|
| def _example() -> None: |
| device = torch.device("cuda") |
| k = torch.randn(128, 2, 128, device=device, dtype=torch.bfloat16) |
| v = torch.randn_like(k) |
| k_q, v_q = quantize_kv_bf16_to_nvfp4_128x4(k, v) |
| print("K FP4 data:", tuple(k_q.data.shape), k_q.data.dtype) |
| print("K scale 128x4:", tuple(k_q.scale_128x4.shape), k_q.scale_128x4.dtype) |
| print("K global scale:", tuple(k_q.global_scale.shape), k_q.global_scale.dtype) |
| print("V FP4 data:", tuple(v_q.data.shape), v_q.data.dtype) |
| print("V scale 128x4:", tuple(v_q.scale_128x4.shape), v_q.scale_128x4.dtype) |
| print("V global scale:", tuple(v_q.global_scale.shape), v_q.global_scale.dtype) |
|
|
|
|
| if __name__ == "__main__": |
| if not torch.cuda.is_available(): |
| raise RuntimeError("quantize.py requires CUDA") |
| _example() |
|
|
