| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm._custom_ops import ( |
| 7 | cutlass_scaled_mm_supports_fp4, |
| 8 | ) |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.math_utils import round_up |
| 11 | |
| 12 | |
| 13 | def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor: |
| 14 | """Pad and block-interleave the FP4 block-scales so that they match the data |
| 15 | layout expected by the CUTLASS / FlashInfer kernels. |
| 16 | |
| 17 | Args: |
| 18 | scale: The FP4 block-scale tensor to swizzle. |
| 19 | |
| 20 | Returns: |
| 21 | The swizzled tensor with the same logical shape as `scale`. |
| 22 | |
| 23 | """ |
| 24 | assert scale.dtype == torch.float8_e4m3fn, ( |
| 25 | "swizzle_blockscale expects the input tensor to be in " |
| 26 | "torch.float8_e4m3fn format." |
| 27 | ) |
| 28 | |
| 29 | scale_ndim = scale.ndim |
| 30 | if scale_ndim == 2: |
| 31 | scale = scale.unsqueeze(0) # (1, M, K) |
| 32 | assert scale.ndim == 3, "Expected a 2-D or 3-D tensor for block scales." |
| 33 | |
| 34 | B, M, K = scale.shape |
| 35 | |
| 36 | M_padded = round_up(M, 128) |
| 37 | K_padded = round_up(K, 4) |
| 38 | |
| 39 | padded = torch.zeros( |
| 40 | (B, M_padded, K_padded), dtype=scale.dtype, device=scale.device |
| 41 | ) |
| 42 | padded[:B, :M, :K] = scale |
| 43 | |
| 44 | # Reshape / permute to the layout required by the kernel. |
| 45 | padded = padded.reshape(B, M_padded // 128, 4, 32, K_padded // 4, 4) |
| 46 | swizzled = padded.permute(0, 1, 4, 3, 2, 5).contiguous().cuda() |
| 47 | |
| 48 | if scale_ndim == 2: |
| 49 | return swizzled.reshape(M_padded, K_padded) |
| 50 | return swizzled.reshape(B, M_padded, K_padded) |
| 51 | |
| 52 | |
| 53 | def cutlass_fp4_supported() -> bool: |
| 54 | if not current_platform.is_cuda(): |
| 55 | return False |
| 56 | capability_tuple = current_platform.get_device_capability() |
| 57 | capability = -1 if capability_tuple is None else capability_tuple.to_int() |
| 58 | return cutlass_scaled_mm_supports_fp4(capability) |
| 59 | |
| 60 | |
| 61 | def pad_nvfp4_weight_for_cutlass( |
| 62 | weight: torch.Tensor, |
| 63 | alignment: int = 32, |
| 64 | ) -> tuple[torch.Tensor, int]: |
| 65 | """Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy |
| 66 | the alignment constraints required by CUTLASS / FlashInfer FP4 kernels. |
| 67 | |
| 68 | CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible |
| 69 | by 32 for aligned memory access and efficient tensor core operations. |
| 70 | """ |
| 71 | weight_current_rows = weight.shape[0] |
| 72 | |
| 73 | # Pad N dimension (rows) if not aligned |
| 74 | if weight_current_rows % alignment != 0: |
| 75 | total_rows = round_up(weight_current_rows, alignment) |
| 76 | pad_rows = total_rows - weight_current_rows |
| 77 | weight = torch.nn.functional.pad(weight, (0, 0, 0, pad_rows)).contiguous() |
| 78 | |
| 79 | # Check K dimension alignment |
| 80 | # 2 FP4 items are packed per byte in the input dimension |
| 81 | weight_current_col_bytes = weight.shape[1] |
| 82 | weight_current_col_elements = weight_current_col_bytes * 2 |
| 83 | |
| 84 | weights_padding_bytes = 0 |
| 85 | if weight_current_col_elements % alignment != 0: |
| 86 | total_cols = round_up(weight_current_col_elements, alignment) |
| 87 | pad_cols = total_cols - weight_current_col_elements |
| 88 | # Convert from FP4 element count to bytes (2 FP4 values per byte) |
| 89 | # pad_cols is always even since alignment=32 and current elements are even |
| 90 | pad_bytes = pad_cols // 2 |
| 91 | weight = torch.nn.functional.pad(weight, (0, pad_bytes, 0, 0)).contiguous() |
| 92 | weights_padding_bytes = pad_bytes |
| 93 | |
| 94 | return weight, weights_padding_bytes |
| 95 | |
| 96 | |
| 97 | def nvfp4_weight_padding_bytes(layer: torch.nn.Module) -> int: |
| 98 | """Compute K-dim padding (packed bytes) of an NVFP4 weight, |
| 99 | Aligned with `pad_nvfp4_weight_for_cutlass` above. |
| 100 | """ |
| 101 | return layer.weight.shape[1] - layer.input_size_per_partition // 2 |
| 102 | |
| 103 | |
| 104 | def pad_nvfp4_activation_for_cutlass( |
| 105 | x_fp4: torch.Tensor, |
| 106 | weights_padding_bytes: int, |
| 107 | ) -> torch.Tensor: |
| 108 | """Pad packed FP4 activations to match the K-dimension padding applied to weights. |
| 109 | The padding is in bytes (tensor dimension), not FP4 elements. |
| 110 | """ |
| 111 | if weights_padding_bytes > 0: |
| 112 | return torch.nn.functional.pad(x_fp4, (0, weights_padding_bytes)).contiguous() |
| 113 | return x_fp4 |
| 114 | |
| 115 | |
| 116 | def slice_nvfp4_output( |
| 117 | out: torch.Tensor, |
| 118 | output_size: int, |
| 119 | ) -> torch.Tensor: |
| 120 | """Slice the output tensor to remove padding in N dimension if weight was padded.""" |
| 121 | if out.shape[-1] != output_size: |
| 122 | return out[..., :output_size].contiguous() |
| 123 | return out |