biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use torch.library.custom_op for CuTeDSL NVFP4 linear GEMM

Browse Source 
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.
This commit is contained in:
biondizzle
2026-05-19 00:50:43 +00:00
parent c043a11bcc
commit c609e9ba3c
1 changed files with 109 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
vllm/kernels/linear/nvfp4/cutedsl.py
View File

@@ -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