From c609e9ba3c8927a7f627fb2d26a1016316c8acd6 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 00:50:43 +0000 Subject: [PATCH] Use torch.library.custom_op for CuTeDSL NVFP4 linear GEMM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dynamo in fullgraph mode traces through torch.autograd.Function, hitting CuTeDSL JIT internals (Path.cwd) and crashing. Registering as a custom op makes it opaque to Dynamo — tracing calls the fake impl, real impl only runs during inference. Custom op: cutedsl::nvfp4_gemm(x, mat_b, scale_b, global_scale_b, in_features, out_features, activation_global_scale) -> Tensor Store finalized weight tensors on the layer (from runner._mat_b etc.) instead of the runner object, since custom ops can only accept tensors. --- vllm/kernels/linear/nvfp4/cutedsl.py | 139 +++++++++++++++++++++------ 1 file changed, 109 insertions(+), 30 deletions(-) diff --git a/vllm/kernels/linear/nvfp4/cutedsl.py b/vllm/kernels/linear/nvfp4/cutedsl.py index 1dbf72b4..72f55a52 100644 --- a/vllm/kernels/linear/nvfp4/cutedsl.py +++ b/vllm/kernels/linear/nvfp4/cutedsl.py @@ -1,14 +1,14 @@ -# SPDX-License-Identifier: Apache-2.0 +# SPDX-License: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """CuTeDSL NVFP4 Linear Kernel for vLLM. -Registers as an NvFp4LinearKernel so that vLLM's kernel selection -mechanism (init_nvfp4_linear_kernel) picks it up on Blackwell GPUs. -Routes NVFP4 GEMM through the CuTeDSL framework, which uses MLIR-compiled -grouped GEMM kernels with Blackwell-specific TMA + wgmma instructions. +Registers as an NvFp4LinearKernel so that vLLM kernel selection +(init_nvfp4_linear_kernel) picks it up on Blackwell GPUs. +Routes NVFP4 GEMM through CuTeDSL's MLIR-compiled grouped GEMM. -CUDA-graph-compatible: all intermediate buffers are pre-allocated, -no CPU-GPU syncs, no dynamic shapes. +The GEMM is registered as a torch.library.custom_op so that +torch.compile/Dynamo treats it as opaque (CuTeDSL internals use +Path.cwd, JIT compilation, etc. which Dynamo cannot trace). """ import torch @@ -21,16 +21,86 @@ from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig logger = init_logger(__name__) +def _cutedsl_nvfp4_gemm_impl( + 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 + + num_tokens = x.shape[0] + padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128 + + # Quantize activation + x_fp4, x_sf = quantize_activation_nvfp4(x, activation_global_scale) + + # 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 + + # Assemble A-side scales in CuTeDSL layout + scale_a = pad_and_swizzle_single(x_sf, num_tokens, x_sf.shape[1], x.device) + + # Expert offsets for 1 group (all tokens in one group) + expert_offsets = torch.tensor([padded_rows], dtype=torch.int64, device=x.device) + + # Global scale for activation + global_scale_a = torch.tensor([activation_global_scale], dtype=torch.float32, 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, + ) + + return out[:num_tokens] + + +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) + + +# Register custom op (idempotent — safe to import multiple times) +if not hasattr(torch.ops, 'cutedsl') or not hasattr(torch.ops.cutedsl, 'nvfp4_gemm'): + _CUTEDSL_NVFP4_GEMM = torch.library.custom_op( + "cutedsl::nvfp4_gemm", + _cutedsl_nvfp4_gemm_impl, + mutates_args=(), + )(_cutedsl_nvfp4_gemm_impl) + _CUTEDSL_NVFP4_GEMM.register_fake(_cutedsl_nvfp4_gemm_fake) + + class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): """NVFP4 GEMM via the CuTeDSL framework (Blackwell SM100+). Uses CuTeDSL's ScaledGroupedGemmKernel with num_groups=1 for - single linear layers. Weight processing: - - uint8 packed FP4 → float4_e2m1fn_x2, permuted to (K, N) - - FP8 block scales permuted to (K_sf, N) - - Global scale stored as float32 - - Activation quantization is done internally (NVFP4 W4A4). + single linear layers. """ @classmethod @@ -49,10 +119,10 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """Convert NVFP4 weights into CuTeDSL kernel format. - Reads the layer's weight (uint8), weight_scale (fp8), and - weight_global_scale (float32) — all set up by - ModelOptNvFp4LinearMethod.process_weights_before our call. - Creates a CuTeDSLNvfp4Linear runner and stores it on the layer. + 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. """ from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear @@ -91,7 +161,7 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): sf_f32[:, split_point:] *= (gs1 / gs) sf = sf_f32.to(torch.float8_e4m3fn) - # Create CuTeDSL runner + # Create CuTeDSL runner to finalize weights (swizzle, TMA, etc.) runner = CuTeDSLNvfp4Linear( in_features=in_features, out_features=out_features, @@ -103,19 +173,22 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): runner.finalize_weights() # Compute activation global scale from input_global_scale_inv. - # ModelOptNvFp4LinearMethod sets: - # input_global_scale = input_scale.max() = amax/448 (small) - # input_global_scale_inv = 1/input_global_scale = 448/amax (large) - # Our quantize_activation_nvfp4(x, global_scale) normalizes: + # quantize_activation_nvfp4(x, global_scale) normalizes: # x_norm = x / global_scale - # So global_scale = amax/448 = input_global_scale = 1/inv. + # 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: - runner._activation_global_scale = 1.0 / inv + activation_global_scale = 1.0 / inv - # Store runner on the layer - layer._cutedsl_runner = runner + # Store the finalized weight tensors on the layer for the custom op. + 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) layer.weight = torch.nn.Parameter( @@ -124,9 +197,7 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): requires_grad=False, ) - # Clean up NVFP4 params that are now in the runner. - # Keep output_size_per_partition, logical_widths, input_size_per_partition - # which may be referenced by the layer's forward path. + # Clean up NVFP4 params that are now handled by our custom op. for attr in ("weight_scale", "weight_global_scale", "input_global_scale", "input_global_scale_inv", "alpha", "weights_padding_cols", "weight_scale_2", @@ -143,7 +214,15 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - result = layer._cutedsl_runner(x) + result = torch.ops.cutedsl.nvfp4_gemm( + 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: result = result + bias return result