| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | from vllm.scalar_type import scalar_types |
| 9 | from vllm.triton_utils import tl, triton |
| 10 | from vllm.utils.torch_utils import direct_register_custom_op |
| 11 | |
| 12 | __all__ = [ |
| 13 | "break_fp4_bytes", |
| 14 | "dequantize_to_dtype", |
| 15 | "nvfp4_gathered_bias", |
| 16 | "ref_nvfp4_quant", |
| 17 | ] |
| 18 | |
| 19 | FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max() |
| 20 | FLOAT4_E2M1_MAX_RECIPROCAL = 1 / FLOAT4_E2M1_MAX |
| 21 | |
| 22 | kE2M1ToFloat_handle = SimpleNamespace( |
| 23 | val=torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | @triton.jit |
| 28 | def _e2m1_inline(nibble): |
| 29 | """Decode an NVFP4 nibble (4 bits: 1 sign + 3 magnitude) to float32. |
| 30 | |
| 31 | Uses direct IEEE 754 bit construction. |
| 32 | For magnitudes 2-7 the FP32 bit pattern is 0x3F000000 + (mag << 22), |
| 33 | which is a single shift + add + bitcast. Magnitudes 0 (zero) and 1 |
| 34 | (E2M1 subnormal = 0.5) are patched with two tl.where ops. |
| 35 | """ |
| 36 | magnitude = nibble & 0x07 |
| 37 | sign = (nibble >> 3) & 1 |
| 38 | |
| 39 | fp32_bits = 0x3F000000 + (magnitude.to(tl.int32) << 22) |
| 40 | val = fp32_bits.to(tl.float32, bitcast=True) |
| 41 | |
| 42 | val = tl.where(magnitude == 0, 0.0, val) |
| 43 | val = tl.where(magnitude == 1, 0.5, val) |
| 44 | |
| 45 | return tl.where(sign == 1, -val, val) |
| 46 | |
| 47 | |
| 48 | @triton.jit |
| 49 | def _nvfp4_gathered_bias_kernel( |
| 50 | markov_ptr, |
| 51 | packed_weight_ptr, |
| 52 | weight_scale_ptr, |
| 53 | global_scale_ptr, |
| 54 | values_ptr, |
| 55 | index_ptr, |
| 56 | logits_ptr, |
| 57 | stride_markov_b, |
| 58 | stride_markov_k, |
| 59 | stride_weight_v, |
| 60 | stride_weight_k, |
| 61 | stride_scale_v, |
| 62 | stride_scale_k, |
| 63 | stride_values_b, |
| 64 | stride_values_j, |
| 65 | stride_index_b, |
| 66 | stride_index_j, |
| 67 | stride_logits_b, |
| 68 | stride_logits_v, |
| 69 | alpha, |
| 70 | TOPK: tl.constexpr, |
| 71 | K: tl.constexpr, |
| 72 | NUM_BLOCKS: tl.constexpr, |
| 73 | ): |
| 74 | """Apply selected NVFP4 Markov-projection rows to dense logits. |
| 75 | |
| 76 | BF16: ``weight[index] -> baddbmm_ -> scatter_``. |
| 77 | NVFP4: ``gather packed W2 row -> dequantize 16-element blocks -> dot with |
| 78 | Markov embedding -> add base value -> scatter corrected logit``. |
| 79 | |
| 80 | The fused NVFP4 path is needed because packed weights cannot be indexed as |
| 81 | BF16 rows, while dequantizing the full vocabulary would defeat top-k. |
| 82 | """ |
| 83 | candidate_id = tl.program_id(0) |
| 84 | request_id = candidate_id // TOPK |
| 85 | candidate = candidate_id - request_id * TOPK |
| 86 | |
| 87 | vocab_id = tl.load( |
| 88 | index_ptr + request_id * stride_index_b + candidate * stride_index_j |
| 89 | ).to(tl.int64) |
| 90 | |
| 91 | block_ids = tl.arange(0, NUM_BLOCKS) |
| 92 | byte_ids = tl.arange(0, 8) |
| 93 | block_mask = block_ids * 16 < K |
| 94 | |
| 95 | packed_offsets = block_ids[:, None] * 8 + byte_ids[None, :] |
| 96 | packed = tl.load( |
| 97 | packed_weight_ptr |
| 98 | + vocab_id * stride_weight_v |
| 99 | + packed_offsets * stride_weight_k, |
| 100 | mask=block_mask[:, None], |
| 101 | other=0, |
| 102 | ) |
| 103 | low = _e2m1_inline(packed & 0x0F) |
| 104 | high = _e2m1_inline((packed >> 4) & 0x0F) |
| 105 | decoded = tl.interleave(low, high) |
| 106 | |
| 107 | raw_weight_scale = tl.load( |
| 108 | weight_scale_ptr + vocab_id * stride_scale_v + block_ids * stride_scale_k, |
| 109 | mask=block_mask, |
| 110 | other=0, |
| 111 | ) |
| 112 | weight_scale = tl.cast( |
| 113 | raw_weight_scale, |
| 114 | tl.float8e4nv, |
| 115 | bitcast=True, |
| 116 | ).to(tl.float32) |
| 117 | global_scale = tl.load(global_scale_ptr).to(tl.float32) |
| 118 | |
| 119 | element_ids = tl.arange(0, 16) |
| 120 | k_offsets = block_ids[:, None] * 16 + element_ids[None, :] |
| 121 | k_mask = k_offsets < K |
| 122 | markov = tl.load( |
| 123 | markov_ptr + request_id * stride_markov_b + k_offsets * stride_markov_k, |
| 124 | mask=k_mask, |
| 125 | other=0.0, |
| 126 | ) |
| 127 | |
| 128 | dequant_weight = (decoded * weight_scale[:, None] * global_scale).to(tl.bfloat16) |
| 129 | products = dequant_weight.to(tl.float32) * markov.to(tl.float32) |
| 130 | block_sums = tl.sum(products, axis=1) |
| 131 | dot = tl.sum(block_sums, axis=0) |
| 132 | |
| 133 | base_value = tl.load( |
| 134 | values_ptr + request_id * stride_values_b + candidate * stride_values_j |
| 135 | ).to(tl.float32) |
| 136 | corrected = base_value + alpha * dot |
| 137 | tl.store( |
| 138 | logits_ptr + request_id * stride_logits_b + vocab_id * stride_logits_v, |
| 139 | corrected, |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | def _nvfp4_gathered_bias_impl( |
| 144 | markov_embed: torch.Tensor, |
| 145 | packed_weight: torch.Tensor, |
| 146 | weight_scale: torch.Tensor, |
| 147 | global_scale: torch.Tensor, |
| 148 | values: torch.Tensor, |
| 149 | index: torch.Tensor, |
| 150 | logits: torch.Tensor, |
| 151 | alpha: float, |
| 152 | ) -> None: |
| 153 | if index.numel() == 0: |
| 154 | return |
| 155 | |
| 156 | assert markov_embed.dtype == torch.bfloat16 |
| 157 | assert values.dtype == torch.bfloat16 |
| 158 | assert logits.dtype == torch.bfloat16 |
| 159 | assert packed_weight.dtype == torch.uint8 |
| 160 | assert weight_scale.dtype == torch.float8_e4m3fn |
| 161 | assert index.dtype == torch.int64 |
| 162 | assert markov_embed.is_cuda |
| 163 | assert all( |
| 164 | tensor.device == markov_embed.device |
| 165 | for tensor in ( |
| 166 | packed_weight, |
| 167 | weight_scale, |
| 168 | global_scale, |
| 169 | values, |
| 170 | index, |
| 171 | logits, |
| 172 | ) |
| 173 | ) |
| 174 | assert markov_embed.ndim == 2 |
| 175 | assert packed_weight.ndim == 2 |
| 176 | assert weight_scale.ndim == 2 |
| 177 | assert values.ndim == 2 |
| 178 | assert index.ndim == 2 |
| 179 | assert logits.ndim == 2 |
| 180 | assert global_scale.numel() == 1 |
| 181 | assert values.shape == index.shape |
| 182 | assert logits.shape[0] == index.shape[0] == markov_embed.shape[0] |
| 183 | |
| 184 | k = markov_embed.shape[-1] |
| 185 | if k % 16 != 0 or packed_weight.shape[-1] * 2 != k: |
| 186 | raise ValueError( |
| 187 | "nvfp4_gathered_bias requires k to be divisible by 16 and " |
| 188 | f"packed_weight width to equal k / 2; got k={k} and " |
| 189 | f"packed_weight width={packed_weight.shape[-1]}" |
| 190 | ) |
| 191 | if weight_scale.shape[-1] != k // 16: |
| 192 | raise ValueError( |
| 193 | "nvfp4_gathered_bias only supports an NVFP4 group size of 16; " |
| 194 | f"got {weight_scale.shape[-1]} scale blocks for k={k}" |
| 195 | ) |
| 196 | |
| 197 | topk = index.shape[1] |
| 198 | num_blocks = triton.next_power_of_2(k // 16) |
| 199 | scale_bytes = weight_scale.view(torch.uint8) |
| 200 | _nvfp4_gathered_bias_kernel[(index.numel(),)]( |
| 201 | markov_embed, |
| 202 | packed_weight, |
| 203 | scale_bytes, |
| 204 | global_scale, |
| 205 | values, |
| 206 | index, |
| 207 | logits, |
| 208 | markov_embed.stride(0), |
| 209 | markov_embed.stride(1), |
| 210 | packed_weight.stride(0), |
| 211 | packed_weight.stride(1), |
| 212 | scale_bytes.stride(0), |
| 213 | scale_bytes.stride(1), |
| 214 | values.stride(0), |
| 215 | values.stride(1), |
| 216 | index.stride(0), |
| 217 | index.stride(1), |
| 218 | logits.stride(0), |
| 219 | logits.stride(1), |
| 220 | alpha, |
| 221 | TOPK=topk, |
| 222 | K=k, |
| 223 | NUM_BLOCKS=num_blocks, |
| 224 | num_warps=4, |
| 225 | num_stages=2, |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | def _nvfp4_gathered_bias_fake( |
| 230 | markov_embed: torch.Tensor, |
| 231 | packed_weight: torch.Tensor, |
| 232 | weight_scale: torch.Tensor, |
| 233 | global_scale: torch.Tensor, |
| 234 | values: torch.Tensor, |
| 235 | index: torch.Tensor, |
| 236 | logits: torch.Tensor, |
| 237 | alpha: float, |
| 238 | ) -> None: |
| 239 | return None |
| 240 | |
| 241 | |
| 242 | direct_register_custom_op( |
| 243 | op_name="nvfp4_gathered_bias", |
| 244 | op_func=_nvfp4_gathered_bias_impl, |
| 245 | mutates_args=["logits"], |
| 246 | fake_impl=_nvfp4_gathered_bias_fake, |
| 247 | ) |
| 248 | |
| 249 | |
| 250 | def nvfp4_gathered_bias( |
| 251 | markov_embed: torch.Tensor, |
| 252 | packed_weight: torch.Tensor, |
| 253 | weight_scale: torch.Tensor, |
| 254 | global_scale: torch.Tensor, |
| 255 | values: torch.Tensor, |
| 256 | index: torch.Tensor, |
| 257 | logits: torch.Tensor, |
| 258 | alpha: float, |
| 259 | ) -> None: |
| 260 | torch.ops.vllm.nvfp4_gathered_bias( |
| 261 | markov_embed, |
| 262 | packed_weight, |
| 263 | weight_scale, |
| 264 | global_scale, |
| 265 | values, |
| 266 | index, |
| 267 | logits, |
| 268 | alpha, |
| 269 | ) |
| 270 | |
| 271 | |
| 272 | @triton.jit |
| 273 | def _dequantize_nvfp4_kernel( |
| 274 | fp4_ptr, |
| 275 | scale_ptr, |
| 276 | global_scale_ptr, |
| 277 | output_ptr, |
| 278 | rows_per_batch: tl.constexpr, |
| 279 | num_blocks: tl.constexpr, |
| 280 | BLOCK_SIZE: tl.constexpr, |
| 281 | has_batch_global_scale: tl.constexpr, |
| 282 | TILE_BLOCKS: tl.constexpr, |
| 283 | ): |
| 284 | """Triton kernel for NVFP4 dequantization (swizzle=False). |
| 285 | |
| 286 | Optimized with 2D tile processing + interleave for coalesced stores. |
| 287 | """ |
| 288 | BLOCK_PACKED: tl.constexpr = BLOCK_SIZE // 2 |
| 289 | |
| 290 | row_idx = tl.program_id(0).to(tl.int64) |
| 291 | tile_idx = tl.program_id(1) |
| 292 | |
| 293 | if has_batch_global_scale: |
| 294 | batch_idx = row_idx // rows_per_batch |
| 295 | global_scale = tl.load(global_scale_ptr + batch_idx).to(tl.float32) |
| 296 | else: |
| 297 | global_scale = tl.load(global_scale_ptr).to(tl.float32) |
| 298 | |
| 299 | fp4_row_offset = row_idx * num_blocks * BLOCK_PACKED |
| 300 | scale_row_offset = row_idx * num_blocks |
| 301 | output_row_offset = row_idx * num_blocks * BLOCK_SIZE |
| 302 | |
| 303 | start_block = tile_idx * TILE_BLOCKS |
| 304 | |
| 305 | # Load scales for this tile: [TILE_BLOCKS] |
| 306 | block_offsets = tl.arange(0, TILE_BLOCKS) |
| 307 | block_mask = (start_block + block_offsets) < num_blocks |
| 308 | |
| 309 | raw_scales = tl.load( |
| 310 | scale_ptr + scale_row_offset + start_block + block_offsets, |
| 311 | mask=block_mask, |
| 312 | other=0, |
| 313 | ) |
| 314 | scale_f32 = tl.cast(raw_scales, tl.float8e4nv, bitcast=True).to(tl.float32) |
| 315 | scale_values = (scale_f32 * global_scale)[:, None] |
| 316 | |
| 317 | # Load [TILE_BLOCKS, BLOCK_PACKED] packed bytes |
| 318 | packed_offsets = tl.arange(0, BLOCK_PACKED)[None, :] |
| 319 | byte_indices = ( |
| 320 | fp4_row_offset |
| 321 | + (start_block + block_offsets[:, None]) * BLOCK_PACKED |
| 322 | + packed_offsets |
| 323 | ) |
| 324 | elem_mask = block_mask[:, None] |
| 325 | raw_bytes = tl.load(fp4_ptr + byte_indices, mask=elem_mask, other=0) |
| 326 | |
| 327 | low_nibble = raw_bytes & 0x0F |
| 328 | high_nibble = (raw_bytes >> 4) & 0x0F |
| 329 | |
| 330 | low_result = _e2m1_inline(low_nibble) * scale_values |
| 331 | high_result = _e2m1_inline(high_nibble) * scale_values |
| 332 | |
| 333 | # Interleave for coalesced contiguous store |
| 334 | result = tl.interleave(low_result, high_result) |
| 335 | |
| 336 | elem_offsets = tl.arange(0, BLOCK_SIZE)[None, :] |
| 337 | out_indices = ( |
| 338 | output_row_offset |
| 339 | + (start_block + block_offsets[:, None]) * BLOCK_SIZE |
| 340 | + elem_offsets |
| 341 | ) |
| 342 | tl.store(output_ptr + out_indices, result, mask=block_mask[:, None]) |
| 343 | |
| 344 | |
| 345 | @triton.jit |
| 346 | def _round_to_fp4(x): |
| 347 | """Round float values to the nearest E2M1 representable value. |
| 348 | |
| 349 | Matches the thresholds in the Python ``cast_to_fp4`` exactly. |
| 350 | """ |
| 351 | sign = tl.where(x < 0.0, -1.0, 1.0) |
| 352 | abs_x = tl.abs(x) |
| 353 | result = tl.where(abs_x > 5.0, 6.0, 0.0) |
| 354 | result = tl.where((abs_x >= 3.5) & (abs_x <= 5.0), 4.0, result) |
| 355 | result = tl.where((abs_x > 2.5) & (abs_x < 3.5), 3.0, result) |
| 356 | result = tl.where((abs_x >= 1.75) & (abs_x <= 2.5), 2.0, result) |
| 357 | result = tl.where((abs_x > 1.25) & (abs_x < 1.75), 1.5, result) |
| 358 | result = tl.where((abs_x >= 0.75) & (abs_x <= 1.25), 1.0, result) |
| 359 | result = tl.where((abs_x > 0.25) & (abs_x < 0.75), 0.5, result) |
| 360 | return result * sign |
| 361 | |
| 362 | |
| 363 | @triton.jit |
| 364 | def _nvfp4_quant_dequant_kernel( |
| 365 | input_ptr, |
| 366 | output_ptr, |
| 367 | global_scale_ptr, |
| 368 | k: tl.constexpr, |
| 369 | num_blocks: tl.constexpr, |
| 370 | BLOCK_SIZE: tl.constexpr, |
| 371 | FP4_MAX_RECIPROCAL: tl.constexpr, |
| 372 | TILE_BLOCKS: tl.constexpr, |
| 373 | ): |
| 374 | """Fused NVFP4 quantize-dequantize kernel. |
| 375 | |
| 376 | Uses a 2D grid (rows x tiles) to parallelize across both rows |
| 377 | and quantization groups within a row. Each program handles |
| 378 | TILE_BLOCKS groups at once using vectorized 2D operations. |
| 379 | """ |
| 380 | row_idx = tl.program_id(0) |
| 381 | tile_idx = tl.program_id(1) |
| 382 | global_scale = tl.load(global_scale_ptr).to(tl.float32) |
| 383 | row_offset = row_idx * k |
| 384 | |
| 385 | start_block = tile_idx * TILE_BLOCKS |
| 386 | block_offsets = tl.arange(0, TILE_BLOCKS) |
| 387 | block_mask = (start_block + block_offsets) < num_blocks |
| 388 | |
| 389 | # Load [TILE_BLOCKS, BLOCK_SIZE] elements |
| 390 | indices = ( |
| 391 | row_offset |
| 392 | + (start_block + block_offsets[:, None]) * BLOCK_SIZE |
| 393 | + tl.arange(0, BLOCK_SIZE)[None, :] |
| 394 | ) |
| 395 | mask_2d = block_mask[:, None] |
| 396 | x = tl.load(input_ptr + indices, mask=mask_2d, other=0.0).to(tl.float32) |
| 397 | |
| 398 | # Per-group scale: [TILE_BLOCKS] |
| 399 | vec_max = tl.max(tl.abs(x), axis=1) |
| 400 | scale = global_scale * (vec_max * FP4_MAX_RECIPROCAL) |
| 401 | scale = tl.clamp(scale, -448.0, 448.0) |
| 402 | scale = scale.to(tl.float8e4nv).to(tl.float32) |
| 403 | |
| 404 | # Safe reciprocal, broadcast to [TILE_BLOCKS, 1] |
| 405 | output_scale = tl.where(scale == 0.0, 0.0, global_scale / scale)[:, None] |
| 406 | |
| 407 | # Quantize: scale, clamp, round to FP4 |
| 408 | scaled_x = tl.clamp(x * output_scale, -6.0, 6.0) |
| 409 | fp4_val = _round_to_fp4(scaled_x) |
| 410 | |
| 411 | # Dequantize: fp4_val * (scale / global_scale) |
| 412 | dequant_scale = (scale / global_scale)[:, None] |
| 413 | result = fp4_val * dequant_scale |
| 414 | |
| 415 | tl.store(output_ptr + indices, result, mask=mask_2d) |
| 416 | |
| 417 | |
| 418 | def _triton_nvfp4_quant_dequant( |
| 419 | x: torch.Tensor, |
| 420 | global_scale: torch.Tensor, |
| 421 | block_size: int, |
| 422 | ) -> torch.Tensor: |
| 423 | """Triton-accelerated NVFP4 quantize-dequantize.""" |
| 424 | x_m, x_k = x.shape |
| 425 | |
| 426 | if not torch.compiler.is_compiling(): |
| 427 | assert x_k % block_size == 0, ( |
| 428 | f"Weight shape K={x_k} is not divisible by block_size={block_size}" |
| 429 | ) |
| 430 | |
| 431 | output_dtype = x.dtype |
| 432 | num_blocks = x_k // block_size |
| 433 | |
| 434 | output = torch.empty(x_m, x_k, dtype=output_dtype, device=x.device) |
| 435 | |
| 436 | tile_blocks = min(64, triton.next_power_of_2(num_blocks)) |
| 437 | num_tiles = (num_blocks + tile_blocks - 1) // tile_blocks |
| 438 | grid = (x_m, num_tiles) |
| 439 | _nvfp4_quant_dequant_kernel[grid]( |
| 440 | x, |
| 441 | output, |
| 442 | global_scale, |
| 443 | x_k, |
| 444 | num_blocks, |
| 445 | block_size, |
| 446 | FLOAT4_E2M1_MAX_RECIPROCAL, |
| 447 | tile_blocks, |
| 448 | ) |
| 449 | |
| 450 | return output |
| 451 | |
| 452 | |
| 453 | def _triton_dequantize_nvfp4( |
| 454 | tensor_fp4: torch.Tensor, |
| 455 | tensor_sf: torch.Tensor, |
| 456 | global_scale: torch.Tensor, |
| 457 | dtype: torch.dtype, |
| 458 | block_size: int = 16, |
| 459 | ) -> torch.Tensor: |
| 460 | """Dequantize NVFP4 using Triton (swizzle=False only). |
| 461 | |
| 462 | Supports both 2D and 3D inputs: |
| 463 | - 2D: [m, packed_k] -> [m, k] |
| 464 | - 3D: [dim0, m, packed_k] -> [dim0, m, k] |
| 465 | """ |
| 466 | assert tensor_fp4.dtype == torch.uint8 |
| 467 | |
| 468 | is_3d = tensor_fp4.ndim == 3 |
| 469 | if is_3d: |
| 470 | dim0, m_per_batch, packed_k = tensor_fp4.shape |
| 471 | tensor_fp4_2d = tensor_fp4.reshape(-1, packed_k) |
| 472 | tensor_sf_2d = tensor_sf.reshape(-1, tensor_sf.shape[-1]) |
| 473 | total_rows_flat = dim0 * m_per_batch |
| 474 | else: |
| 475 | m_per_batch, packed_k = tensor_fp4.shape |
| 476 | tensor_fp4_2d = tensor_fp4 |
| 477 | tensor_sf_2d = tensor_sf |
| 478 | total_rows_flat = m_per_batch |
| 479 | |
| 480 | k = packed_k * 2 |
| 481 | num_blocks = k // block_size |
| 482 | |
| 483 | output = torch.empty(total_rows_flat, k, dtype=dtype, device=tensor_fp4.device) |
| 484 | |
| 485 | # View as uint8 so Triton can load raw bytes and bitcast to float8_e4m3fn |
| 486 | scale_raw = tensor_sf_2d.contiguous().view(torch.uint8) |
| 487 | |
| 488 | # Shape-adaptive tile sizing: for large row counts (3D), process |
| 489 | # entire row in one tile. For small row counts (2D), use smaller |
| 490 | # tiles to increase parallelism across CUs. |
| 491 | np2 = triton.next_power_of_2(num_blocks) |
| 492 | if total_rows_flat >= 4096: |
| 493 | # Many rows: maximize work per CTA, one tile per row |
| 494 | tile_blocks = np2 |
| 495 | nw = 1 |
| 496 | ns = 2 |
| 497 | elif total_rows_flat >= 2048: |
| 498 | # Medium-many rows: full row, 2 warps |
| 499 | tile_blocks = np2 |
| 500 | nw = 2 |
| 501 | ns = 2 |
| 502 | else: |
| 503 | # Few rows: use moderate tiles for CU utilization |
| 504 | tile_blocks = min(64, np2) |
| 505 | nw = 4 |
| 506 | ns = 2 |
| 507 | num_tiles = (num_blocks + tile_blocks - 1) // tile_blocks |
| 508 | grid = (total_rows_flat, num_tiles) |
| 509 | _dequantize_nvfp4_kernel[grid]( |
| 510 | tensor_fp4_2d, |
| 511 | scale_raw, |
| 512 | global_scale, |
| 513 | output, |
| 514 | m_per_batch, |
| 515 | num_blocks, |
| 516 | block_size, |
| 517 | is_3d, |
| 518 | tile_blocks, |
| 519 | num_warps=nw, |
| 520 | num_stages=ns, |
| 521 | ) |
| 522 | |
| 523 | if is_3d: |
| 524 | output = output.reshape(dim0, m_per_batch, k) |
| 525 | |
| 526 | return output |
| 527 | |
| 528 | |
| 529 | def break_fp4_bytes(a, dtype): |
| 530 | assert a.dtype == torch.uint8 |
| 531 | m, n = a.shape |
| 532 | # Vectorized nibble processing |
| 533 | a_flat = a.flatten() |
| 534 | high = (a_flat & 0xF0) >> 4 # Upper nibbles |
| 535 | low = a_flat & 0x0F # Lower nibbles |
| 536 | # Combine nibbles for batch processing |
| 537 | combined = torch.stack((low, high), dim=1).flatten() |
| 538 | # Vectorized sign and magnitude extraction |
| 539 | signs = (combined & 0x08).to(torch.bool) # Sign bits |
| 540 | abs_vals = (combined & 0x07).to(torch.long) |
| 541 | |
| 542 | kE2M1 = kE2M1ToFloat_handle.val.to(device=a.device) |
| 543 | # Device-aware lookup and sign application |
| 544 | values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) |
| 545 | # Reshape to final form |
| 546 | return values.reshape(m, n * 2).to(dtype=dtype) |
| 547 | |
| 548 | |
| 549 | def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size): |
| 550 | m_tiles = (m + 128 - 1) // 128 |
| 551 | f = block_size * 4 |
| 552 | k_tiles = (k + f - 1) // f |
| 553 | tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4)) |
| 554 | tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) |
| 555 | out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size) |
| 556 | return out[0:m, 0:k] |
| 557 | |
| 558 | |
| 559 | def dequantize_to_dtype( |
| 560 | tensor_fp4: torch.Tensor, |
| 561 | tensor_sf: torch.Tensor, |
| 562 | global_scale: torch.Tensor, |
| 563 | dtype: torch.dtype, |
| 564 | block_size: int = 16, |
| 565 | swizzle: bool | None = True, |
| 566 | ): |
| 567 | """Dequantize the fp4 tensor back to high precision. |
| 568 | |
| 569 | Supports both 2D and 3D inputs: |
| 570 | - 2D: [m, packed_k] -> [m, k] |
| 571 | - 3D: [dim0, m, packed_k] -> [dim0, m, k] |
| 572 | """ |
| 573 | # Two fp4 values are packed into one uint8. |
| 574 | assert tensor_fp4.dtype == torch.uint8 |
| 575 | |
| 576 | if not swizzle and tensor_fp4.is_cuda and current_platform.is_cuda_alike(): |
| 577 | return _triton_dequantize_nvfp4( |
| 578 | tensor_fp4, tensor_sf, global_scale, dtype, block_size |
| 579 | ) |
| 580 | |
| 581 | # We handle 3D tensors reshaping them to 2D. |
| 582 | is_3d = tensor_fp4.ndim == 3 |
| 583 | |
| 584 | if is_3d: |
| 585 | dim0, m, packed_k = tensor_fp4.shape |
| 586 | tensor_fp4 = tensor_fp4.reshape(-1, packed_k) |
| 587 | tensor_sf = tensor_sf.reshape(-1, tensor_sf.shape[-1]) |
| 588 | global_scale = global_scale[:, None, None] |
| 589 | else: |
| 590 | m, packed_k = tensor_fp4.shape |
| 591 | |
| 592 | k = packed_k * 2 |
| 593 | tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32) |
| 594 | tensor_f32 = tensor_f32.reshape(-1, k // block_size, block_size) |
| 595 | tensor_sf = tensor_sf.view(torch.float8_e4m3fn) |
| 596 | |
| 597 | if swizzle: |
| 598 | tensor_sf = convert_swizzled_to_linear( # noqa: E501 |
| 599 | tensor_sf, tensor_f32.size(0), k, block_size |
| 600 | ) |
| 601 | |
| 602 | if is_3d: |
| 603 | tensor_sf = tensor_sf.reshape(dim0, m, k // block_size) |
| 604 | tensor_sf_dtype = tensor_sf.to(torch.float32) * global_scale |
| 605 | |
| 606 | if is_3d: |
| 607 | tensor_f32 = tensor_f32.reshape(dim0, m, -1, block_size) |
| 608 | |
| 609 | # scale the tensor |
| 610 | out = tensor_f32 * tensor_sf_dtype.unsqueeze(-1) |
| 611 | out = out.reshape(*out.shape[:-2], -1) |
| 612 | |
| 613 | return out.to(dtype) |
| 614 | |
| 615 | |
| 616 | def get_reciprocal(x): |
| 617 | if isinstance(x, torch.Tensor): |
| 618 | # torch.where yields operation not permitted when stream is capturing. |
| 619 | return 1.0 / (x + (x == 0) * 1e8) |
| 620 | elif isinstance(x, (float, int)): |
| 621 | return 0.0 if x == 0 else 1.0 / x |
| 622 | else: |
| 623 | raise TypeError("Input must be a float, int, or a torch.Tensor.") |
| 624 | |
| 625 | |
| 626 | def cast_to_fp4(x): |
| 627 | sign = torch.sign(x) |
| 628 | x = torch.abs(x) |
| 629 | x[(x >= 0.0) & (x <= 0.25)] = 0.0 |
| 630 | x[(x > 0.25) & (x < 0.75)] = 0.5 |
| 631 | x[(x >= 0.75) & (x <= 1.25)] = 1.0 |
| 632 | x[(x > 1.25) & (x < 1.75)] = 1.5 |
| 633 | x[(x >= 1.75) & (x <= 2.5)] = 2.0 |
| 634 | x[(x > 2.5) & (x < 3.5)] = 3.0 |
| 635 | x[(x >= 3.5) & (x <= 5.0)] = 4.0 |
| 636 | x[x > 5.0] = 6.0 |
| 637 | return x * sign |
| 638 | |
| 639 | |
| 640 | def ref_nvfp4_quant(x, global_scale, block_size): |
| 641 | assert global_scale.dtype == torch.float32 |
| 642 | assert x.ndim == 2 |
| 643 | m, n = x.shape |
| 644 | x = torch.reshape(x, (m, n // block_size, block_size)) |
| 645 | vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) |
| 646 | scale = global_scale * (vec_max * FLOAT4_E2M1_MAX_RECIPROCAL) |
| 647 | scale = torch.clamp(scale, max=448, min=-448) |
| 648 | scale = scale.to(torch.float8_e4m3fn).to(torch.float32) |
| 649 | output_scale = get_reciprocal(scale * get_reciprocal(global_scale)) |
| 650 | |
| 651 | scaled_x = x.to(torch.float32) * output_scale |
| 652 | clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n) |
| 653 | # both outputs are float32 |
| 654 | return cast_to_fp4(clipped_x), scale.squeeze(-1) |
| 655 | |
| 656 | |
| 657 | def ref_nvfp4_quant_dequant( |
| 658 | x: torch.Tensor, global_scale: torch.Tensor, block_size: int |
| 659 | ) -> torch.Tensor: |
| 660 | """NVFP4 quantize-dequantize operation. |
| 661 | |
| 662 | `global_scale` is expected to have a single element. |
| 663 | """ |
| 664 | if current_platform.is_cuda_alike(): |
| 665 | return _triton_nvfp4_quant_dequant(x, global_scale, block_size) |
| 666 | |
| 667 | x_m, x_k = x.shape |
| 668 | output_dtype = x.dtype |
| 669 | |
| 670 | # quantize input to (FP4 and interleaved block scale) |
| 671 | x_fp4, x_blockscale = ref_nvfp4_quant(x, global_scale, block_size) |
| 672 | |
| 673 | # dequantize input |
| 674 | x_fp4 = x_fp4.reshape(x_m, x_k // block_size, block_size) |
| 675 | x_blockscale = x_blockscale.unsqueeze(-1) / global_scale |
| 676 | x_dq = (x_fp4 * x_blockscale).reshape(x_m, x_k).to(output_dtype) |
| 677 | |
| 678 | return x_dq |
| 679 | |
| 680 | |
| 681 | def run_nvfp4_emulations( |
| 682 | x: torch.Tensor, |
| 683 | input_global_scale: torch.Tensor, |
| 684 | weight: torch.Tensor, |
| 685 | weight_scale_swizzled: torch.Tensor, |
| 686 | weight_global_scale: torch.Tensor, |
| 687 | swizzle: bool | None = True, |
| 688 | ): |
| 689 | output_dtype = x.dtype |
| 690 | group_size = 16 |
| 691 | |
| 692 | x_dq = ref_nvfp4_quant_dequant(x, input_global_scale, block_size=group_size) |
| 693 | |
| 694 | # dequantize weight |
| 695 | w_fp4 = weight.data.view(torch.uint8) |
| 696 | w_dq = dequantize_to_dtype( |
| 697 | w_fp4, |
| 698 | weight_scale_swizzled.data, |
| 699 | weight_global_scale, |
| 700 | output_dtype, |
| 701 | group_size, |
| 702 | swizzle=swizzle, |
| 703 | ) |
| 704 | |
| 705 | # matmul |
| 706 | out = torch.matmul(x_dq, w_dq.t()) |
| 707 | return out |