diff --git a/vllm/kernels/linear/nvfp4/cutedsl.py b/vllm/kernels/linear/nvfp4/cutedsl.py index 53fb83c3..4e639b0e 100644 --- a/vllm/kernels/linear/nvfp4/cutedsl.py +++ b/vllm/kernels/linear/nvfp4/cutedsl.py @@ -1,4 +1,4 @@ -# SPDX-License: Apache-2.0 +# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """CuTeDSL NVFP4 Linear Kernel for vLLM. @@ -21,73 +21,87 @@ from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig logger = init_logger(__name__) -# Register custom op (idempotent — safe to import multiple times) -if not hasattr(torch.ops, 'cutedsl') or not hasattr(torch.ops.cutedsl, 'nvfp4_gemm'): +@torch.library.custom_op("cutedsl::nvfp4_linear", mutates_args=()) +def _cutedsl_nvfp4_linear( + x: torch.Tensor, + mat_b: torch.Tensor, + scale_b: torch.Tensor, + global_scale_b: torch.Tensor, + activation_global_scale: float, +) -> torch.Tensor: + """Run a single-group NVFP4 GEMM via CuTeDSL. - @torch.library.custom_op("cutedsl::nvfp4_gemm", mutates_args=()) - def _cutedsl_nvfp4_gemm( - x: torch.Tensor, - mat_b: torch.Tensor, - scale_b: torch.Tensor, - global_scale_b: torch.Tensor, - in_features: int, - out_features: int, - activation_global_scale: float, - ) -> torch.Tensor: - """Run a single-group NVFP4 GEMM via CuTeDSL.""" - from cutedsl.bridge import (pad_and_swizzle_single, - quantize_activation_nvfp4, - run_nvfp4_grouped_gemm) - from cutedsl.nvfp4_linear import cutedsl_ceil_div + This is registered as a torch.library.custom_op so Dynamo treats + it as opaque. The real implementation calls into CuTeDSL's + pre-assembled weight tensors and grouped GEMM kernel. - num_tokens = x.shape[0] - padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128 + Args: + x: BF16 input (M, K) + mat_b: Pre-assembled FP4 weight (1, K_padded, N_packed) float4_e2m1fn_x2 + scale_b: Pre-assembled block scales (1, K_sf_padded, N_sf_packed) fp8 + global_scale_b: Weight global scale (1,) float32 + activation_global_scale: Activation global scale (float) + """ + from cutedsl.bridge import (quantize_activation_nvfp4, + run_nvfp4_grouped_gemm) + from cutedsl.kernel.moe.torch_scaled_grouped_mm import pad_and_swizzle_single + from cutedsl.nvfp4_linear import cutedsl_ceil_div - # Quantize activation - x_fp4, x_sf = quantize_activation_nvfp4(x, activation_global_scale) + num_tokens = x.shape[0] + in_features = x.shape[1] + out_features = mat_b.shape[2] # packed N in float4 elements - # Pad activation to 128-row alignment for TMA - if num_tokens < padded_rows: - x_fp4_padded = torch.zeros(padded_rows, x_fp4.shape[1], - dtype=x_fp4.dtype, device=x.device) - x_fp4_padded[:num_tokens] = x_fp4 - else: - x_fp4_padded = x_fp4 + # Quantize activation + x_fp4, x_sf = quantize_activation_nvfp4(x, activation_global_scale) - # Assemble A-side scales in CuTeDSL layout - scale_a = pad_and_swizzle_single(x_sf, num_tokens, x_sf.shape[1], x.device) + # Pad activation to 128-row alignment for TMA + padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128 + if num_tokens < padded_rows: + x_fp4_padded = torch.zeros(padded_rows, x_fp4.shape[1], + dtype=x_fp4.dtype, device=x.device) + x_fp4_padded[:num_tokens] = x_fp4 + else: + x_fp4_padded = x_fp4 - # Expert offsets for 1 group (all tokens in one group) - expert_offsets = torch.tensor([padded_rows], dtype=torch.int64, device=x.device) + # Assemble A-side scales: pad + swizzle for single group + num_rows, num_cols = x_sf.shape + padded_rows_sf = cutedsl_ceil_div(num_rows, 128) * 128 + padded_cols_sf = cutedsl_ceil_div(num_cols, 4) * 4 + buf = torch.zeros(padded_rows_sf, padded_cols_sf, dtype=torch.float8_e4m3fn, device=x_sf.device) + buf[:num_rows, :num_cols] = x_sf + scale_a = pad_and_swizzle_single(buf).unsqueeze(0) # (1, padded_rows, padded_cols) - # Global scale for activation - global_scale_a = torch.tensor([activation_global_scale], dtype=torch.float32, device=x.device) + # Expert offsets for 1 group + expert_offsets = torch.tensor([padded_rows], dtype=torch.int64, device=x.device) - # Run the CuTeDSL grouped GEMM (1 group) - out = run_nvfp4_grouped_gemm( - mat_a=x_fp4_padded, - mat_b=mat_b, - scale_a=scale_a, - scale_b=scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, - global_scale_b=global_scale_b, - ) + # Global scale for activation + global_scale_a = torch.tensor([activation_global_scale], dtype=torch.float32, device=x.device) - return out[:num_tokens] + # Run the CuTeDSL grouped GEMM (1 group) + out = run_nvfp4_grouped_gemm( + mat_a=x_fp4_padded, + mat_b=mat_b, + scale_a=scale_a, + scale_b=scale_b, + expert_offsets=expert_offsets, + global_scale_a=global_scale_a, + global_scale_b=global_scale_b, + ) - @_cutedsl_nvfp4_gemm.register_fake - def _cutedsl_nvfp4_gemm_fake( - x: torch.Tensor, - mat_b: torch.Tensor, - scale_b: torch.Tensor, - global_scale_b: torch.Tensor, - in_features: int, - out_features: int, - activation_global_scale: float, - ) -> torch.Tensor: - return torch.empty((*x.shape[:-1], out_features), dtype=torch.bfloat16, - device=x.device) + return out[:num_tokens] + + +@_cutedsl_nvfp4_linear.register_fake +def _cutedsl_nvfp4_linear_fake( + x: torch.Tensor, + mat_b: torch.Tensor, + scale_b: torch.Tensor, + global_scale_b: torch.Tensor, + activation_global_scale: float, +) -> torch.Tensor: + out_features = mat_b.shape[2] + return torch.empty((*x.shape[:-1], out_features), dtype=torch.bfloat16, + device=x.device) class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): @@ -107,17 +121,11 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): return False, "CuTeDSL NVFP4 requires SM100+ (Blackwell)" @classmethod - def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[ bool, str | None]: return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """Convert NVFP4 weights into CuTeDSL kernel format. - - After ModelOptNvFp4LinearMethod.process_weights_after_loading - sets up input_global_scale, weight_global_scale, alpha, etc., - this method converts the weights into CuTeDSL's swizzled TMA - format and stores the finalized tensors on the layer. - """ + """Convert NVFP4 weights into CuTeDSL kernel format.""" from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear w_uint8 = layer.weight.data # (out, in//2) uint8 packed E2M1 @@ -134,12 +142,10 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): sf = sf.to(torch.float8_e4m3fn) sf = sf.permute(1, 0).contiguous() - # Global scale (set by ModelOptNvFp4LinearMethod.process_weights_after_loading) + # Global scale gs = layer.weight_global_scale.data.item() - # Handle fused projections (MergedColumnParallelLinear with dual gs). - # When weight_global_scale has 2 elements (e.g. fused_wqa_wkv), - # normalize to max(gs1, gs2) and fold ratio into block scales. + # Handle fused projections with dual global scales if layer.weight_global_scale.numel() == 2: gs0 = layer.weight_global_scale[0].item() gs1 = layer.weight_global_scale[1].item() @@ -167,21 +173,19 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): runner.finalize_weights() # Compute activation global scale from input_global_scale_inv. - # quantize_activation_nvfp4(x, global_scale) normalizes: - # x_norm = x / global_scale - # global_scale = amax/448 = input_global_scale = 1/inv. activation_global_scale = 1.0 / 2688.0 # default fallback if hasattr(layer, 'input_global_scale_inv') and layer.input_global_scale_inv is not None: inv = layer.input_global_scale_inv.data.item() if inv != 0: activation_global_scale = 1.0 / inv - # Store the finalized weight tensors on the layer for the custom op. + # Store pre-assembled weight tensors on the layer for the custom op. + # mat_b shape: (1, K_padded, N_packed) float4_e2m1fn_x2 + # scale_b shape: (1, K_sf_padded, N_sf_packed) fp8 + # gsb shape: (1,) float32 layer._cutedsl_mat_b = runner._mat_b layer._cutedsl_scale_b = runner._scale_b layer._cutedsl_global_scale_b = runner._gsb - layer._cutedsl_in_features = in_features - layer._cutedsl_out_features = out_features layer._cutedsl_activation_global_scale = activation_global_scale # Replace weight with dummy BF16 (vLLM module introspection may need it) @@ -208,13 +212,11 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - result = torch.ops.cutedsl.nvfp4_gemm( + result = torch.ops.cutedsl.nvfp4_linear( x, layer._cutedsl_mat_b, layer._cutedsl_scale_b, layer._cutedsl_global_scale_b, - layer._cutedsl_in_features, - layer._cutedsl_out_features, layer._cutedsl_activation_global_scale, ) if bias is not None: