| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm._custom_ops import scaled_fp4_quant |
| 6 | from vllm.scalar_type import scalar_types |
| 7 | |
| 8 | FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max() |
| 9 | FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max |
| 10 | |
| 11 | kE2M1ToFloat = torch.tensor( |
| 12 | [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size): |
| 17 | m_tiles = (m + 128 - 1) // 128 |
| 18 | f = block_size * 4 |
| 19 | k_tiles = (k + f - 1) // f |
| 20 | tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4)) |
| 21 | tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) |
| 22 | out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size) |
| 23 | return out[0:m, 0:k] |
| 24 | |
| 25 | |
| 26 | def convert_swizzled_8x4_layout_to_linear( |
| 27 | a_sf_swizzled: torch.Tensor, m, k, block_size |
| 28 | ): |
| 29 | m_tiles = (m + 8 - 1) // 8 |
| 30 | f = block_size * 4 |
| 31 | k_tiles = (k + f - 1) // f |
| 32 | tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 8, 4)) |
| 33 | tmp = torch.permute(tmp, (0, 1, 3, 2, 4)) |
| 34 | out = tmp.reshape(m_tiles * 8, k_tiles * f // block_size) |
| 35 | return out[0:m, 0:k] |
| 36 | |
| 37 | |
| 38 | def dequantize_nvfp4_to_dtype( |
| 39 | tensor_fp4, |
| 40 | tensor_sf, |
| 41 | global_scale, |
| 42 | dtype, |
| 43 | device, |
| 44 | block_size=16, |
| 45 | is_sf_128x4_layout=True, |
| 46 | is_sf_linear_layout=False, |
| 47 | ): |
| 48 | """Dequantize the fp4 tensor back to high precision.""" |
| 49 | # Two fp4 values are packed into one uint8. |
| 50 | assert tensor_fp4.dtype == torch.uint8 |
| 51 | m, packed_k = tensor_fp4.shape |
| 52 | k = packed_k * 2 |
| 53 | tensor_f32 = break_fp4_bytes(tensor_fp4, dtype) |
| 54 | tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) |
| 55 | tensor_sf = tensor_sf.view(torch.float8_e4m3fn) |
| 56 | if is_sf_linear_layout: |
| 57 | tensor_sf = tensor_sf[:m, : k // block_size] |
| 58 | elif is_sf_128x4_layout: |
| 59 | tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) |
| 60 | else: |
| 61 | tensor_sf = convert_swizzled_8x4_layout_to_linear(tensor_sf, m, k, block_size) |
| 62 | |
| 63 | tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale |
| 64 | |
| 65 | # scale the tensor |
| 66 | out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) |
| 67 | return out.to(dtype=dtype) |
| 68 | |
| 69 | |
| 70 | def break_fp4_bytes(a, dtype): |
| 71 | assert a.dtype == torch.uint8 |
| 72 | m, n = a.shape |
| 73 | |
| 74 | # Vectorized nibble processing |
| 75 | a_flat = a.flatten() |
| 76 | high = (a_flat & 0xF0) >> 4 # Upper nibbles |
| 77 | low = a_flat & 0x0F # Lower nibbles |
| 78 | |
| 79 | # Combine nibbles for batch processing |
| 80 | combined = torch.stack((low, high), dim=1).flatten() |
| 81 | |
| 82 | # Vectorized sign and magnitude extraction |
| 83 | signs = (combined & 0x08).to(torch.bool) # Sign bits |
| 84 | abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices |
| 85 | |
| 86 | # Device-aware lookup and sign application |
| 87 | kE2M1 = kE2M1ToFloat.to(device=a.device) |
| 88 | values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) |
| 89 | |
| 90 | # Reshape to final form |
| 91 | return values.reshape(m, n * 2).to(dtype=dtype) |
| 92 | |
| 93 | |
| 94 | def dequant_nvfp4_kv_cache( |
| 95 | fp4_data: torch.Tensor, |
| 96 | block_scale: torch.Tensor, |
| 97 | global_scale: float, |
| 98 | head_size: int, |
| 99 | block_size: int, |
| 100 | swizzled_scales: bool = True, |
| 101 | ) -> torch.Tensor: |
| 102 | """Dequantize an NVFP4 KV cache. |
| 103 | |
| 104 | The input must be in HND layout so that the last two dims are |
| 105 | (block_size, last_dim). For NHD caches, permute to HND first. |
| 106 | |
| 107 | Args: |
| 108 | fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4. |
| 109 | block_scale: [..., num_heads, block_size, head_size//16] fp8 block |
| 110 | scales (as uint8 or float8_e4m3fn). |
| 111 | global_scale: checkpoint dequant scale (k_scale or v_scale). |
| 112 | head_size: head dimension. |
| 113 | block_size: page size. |
| 114 | swizzled_scales: whether ``block_scale`` is in the 4x4-swizzled |
| 115 | layout of the SM100 trtllm-gen kernel. ``reshape_and_cache_flash`` |
| 116 | writes K scales linearly and V scales swizzled, so pass ``False`` |
| 117 | for K and ``True`` for V. |
| 118 | |
| 119 | Returns: |
| 120 | [..., num_heads, block_size, head_size] float32. |
| 121 | |
| 122 | """ |
| 123 | data_dim = head_size // 2 |
| 124 | scale_dim = head_size // 16 |
| 125 | |
| 126 | fp4_packed = fp4_data |
| 127 | sf_bytes = block_scale.view(torch.uint8) |
| 128 | batch_shape = sf_bytes.shape[:-2] |
| 129 | T, S = block_size, scale_dim |
| 130 | |
| 131 | if swizzled_scales: |
| 132 | # Unswizzle 4x4 block scales on (block_size, scale_dim) plane. |
| 133 | # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S] |
| 134 | sg = S // 4 |
| 135 | sf_reshape = sf_bytes.reshape(*batch_shape, T // 4, 4, sg, 4) |
| 136 | ndim = sf_reshape.ndim |
| 137 | # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg) |
| 138 | perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2] |
| 139 | sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S) |
| 140 | else: |
| 141 | sf_linear = sf_bytes |
| 142 | sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32) |
| 143 | |
| 144 | # Unpack fp4 |
| 145 | shape = fp4_packed.shape # [..., T, data_dim] |
| 146 | fp4_flat = fp4_packed.reshape(-1, data_dim) |
| 147 | fp4_vals = break_fp4_bytes(fp4_flat, torch.float32) |
| 148 | fp4_vals = fp4_vals.reshape(*shape[:-1], head_size) |
| 149 | |
| 150 | # Dequant: fp4_val * block_scale * global_scale per 16-element group |
| 151 | return ( |
| 152 | fp4_vals.reshape(*shape[:-1], scale_dim, 16) |
| 153 | * (sf_f32 * global_scale).unsqueeze(-1) |
| 154 | ).reshape(*shape[:-1], head_size) |
| 155 | |
| 156 | |
| 157 | def get_nvfp4_global_scale(a: torch.Tensor): |
| 158 | return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32) |
| 159 | |
| 160 | |
| 161 | def quant_nvfp4_tensor( |
| 162 | a: torch.Tensor, |
| 163 | *, |
| 164 | is_sf_swizzled_layout: bool = True, |
| 165 | ): |
| 166 | a_global_scale = get_nvfp4_global_scale(a) |
| 167 | a_quant, a_block_scale = scaled_fp4_quant( |
| 168 | a, |
| 169 | a_global_scale, |
| 170 | is_sf_swizzled_layout=is_sf_swizzled_layout, |
| 171 | ) |
| 172 | return a_quant, a_block_scale, a_global_scale |