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

Fix CuTeDSL NVFP4 linear: correct scale assembly in custom op

Browse Source 
- pad_and_swizzle_single takes 1 arg (2D tensor), not 4
- Inline the scale assembly logic: pad x_sf → swizzle → unsqueeze for 1 group
- Remove unused CuTeDSLNvfp4Linear import from custom op impl
This commit is contained in:
biondizzle
2026-05-19 01:01:42 +00:00
parent e0eb436914
commit b81200f427
1 changed files with 81 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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