From c926c4a597809b0aa3409892f0d86d49e205804c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 2 Jun 2026 17:57:33 +0000 Subject: [PATCH] =?UTF-8?q?P5:=20Fix=20mhc=5Frmsnorm=5Fquantize=5Fnvfp4=20?= =?UTF-8?q?=E2=80=94=20add=20proper=20function=20definition?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dsv4/ops/quantize.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/dsv4/ops/quantize.py b/dsv4/ops/quantize.py index 399ba922..eab0281a 100644 --- a/dsv4/ops/quantize.py +++ b/dsv4/ops/quantize.py @@ -395,6 +395,23 @@ def dequantize_nvfp4(x_fp4, x_sf, gsa, shape=None): return mod.dequant_nvfp4(x_fp4, x_sf, gsa) +def mhc_rmsnorm_quantize_nvfp4(X_l, A_l, norm_weight, eps=1e-6, divisor=6.0 * 448.0): + """Fused mHC pre_block + RMSNorm + NVFP4 quantize: 2 kernel launches total. + + Replaces: bmm (1 launch) + rmsnorm (4+ launches) + quantize (2 launches) + Total unfused: 7+ launches per site × 122 sites = 854+ launches/token + Fused: 2 launches per site × 122 sites = 244 launches → 610 launches saved/token. + + Args: + X_l: (M, n_hc, N) BF16 tensor. n_hc must be <= 4, N multiple of 16. + A_l: (M, n_hc) BF16 tensor. Softmax weights from mHC._dynamic_params. + norm_weight: (N,) FP32 RMSNorm weight. + eps: RMSNorm epsilon (default 1e-6). + divisor: gsa = amax / divisor. Default 6.0 * 448.0 = 2688.0. + + Returns: + QuantizedActivation with x_fp4, x_sf, gsa, inv_rms + """ from dsv4.kernels.cuda.loader import get_cuda_module mod = get_cuda_module("fused_mhc_rmsnorm_quantize", ["fused_mhc_rmsnorm_quantize.cu"]) x_fp4, x_sf, gsa, inv_rms = mod.mhc_rmsnorm_quantize_nvfp4(X_l, A_l, norm_weight, eps, divisor)