From 74dd3afb61956d7e2aab38f0314d78de5aa7cb79 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 2 Jul 2026 11:31:28 +0300 Subject: [PATCH] nvfp4: decompose the GEMM probe - input quantization vs FP4 kernel First live diagnostics run showed the whole quantized path at 0.26x of bf16 despite a healthy cuda backend and zero emulated layers. The single combined timing cannot say WHICH half loses: the dynamic input quantization wrapped around every matmul (a fixed per-call cost times 1232 layers/step) or the FP4 GEMM kernel not engaging. The probe now times them separately at M=1024 and M=16384 (realistic video token count) and names the culprit in the verdict. Co-Authored-By: Claude Fable 5 --- nvfp4_patch.py | 37 +++++++++++++++++++++++++------------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/nvfp4_patch.py b/nvfp4_patch.py index b809ecdc..b5a98319 100644 --- a/nvfp4_patch.py +++ b/nvfp4_patch.py @@ -226,7 +226,6 @@ def log_nvfp4_diagnostics(model, device=None): import comfy.quant_ops as quant_ops d = torch.device(device) if device is not None else torch.device("cuda") w = torch.randn(4096, 4096, dtype=torch.bfloat16, device=d) - x = torch.randn(1024, 4096, dtype=torch.bfloat16, device=d) qw = quant_ops.QuantizedTensor.from_float(w, "TensorCoreNVFP4Layout", scale="recalculate") def _timed(fn, n=30): @@ -239,18 +238,32 @@ def log_nvfp4_diagnostics(model, device=None): torch.cuda.synchronize(d) return (time.perf_counter() - t0) / n - t_bf16 = _timed(lambda: torch.nn.functional.linear(x, w)) - - def _q(): + # Decomposed probe at two token counts: M=1024 (small scene chunk) and + # M=16384 (realistic video token count). Times the input quantization and + # the FP4 GEMM SEPARATELY — "no speedup" has two very different fixes + # depending on which one is at fault: + # * GEMM alone beats bf16, quantization dominates -> per-call overhead + # problem (x1232 layers/step); needs static input scales or fusion. + # * GEMM alone loses to bf16 -> the FP4 kernel isn't engaging at all. + for m in (1024, 16384): + x = torch.randn(m, 4096, dtype=torch.bfloat16, device=d) + t_bf16 = _timed(lambda: torch.nn.functional.linear(x, w)) + t_quant = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout")) qx = quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout") - torch.nn.functional.linear(qx, qw) - - t_q = _timed(_q) - verdict = ("FAST PATH ACTIVE" if t_q < t_bf16 * 0.8 - else "NO SPEEDUP — emulated or overhead-bound") - lines.append(f" GEMM probe 4096x4096 @ bs1024 (incl. dynamic input quant): " - f"bf16 {t_bf16 * 1e3:.3f} ms vs nvfp4 {t_q * 1e3:.3f} ms " - f"({t_bf16 / max(t_q, 1e-9):.2f}x) -> {verdict}") + t_gemm = _timed(lambda: torch.nn.functional.linear(qx, qw)) + total = t_quant + t_gemm + if t_gemm < t_bf16 * 0.8 and total < t_bf16 * 0.8: + verdict = "FAST PATH ACTIVE (net win incl. input quant)" + elif t_gemm < t_bf16 * 0.8: + verdict = "FP4 GEMM fast but INPUT QUANTIZATION dominates — overhead-bound" + else: + verdict = "FP4 GEMM NOT faster than bf16 — kernel not engaging" + lines.append( + f" probe 4096x4096 @ M={m}: bf16 {t_bf16 * 1e3:.3f} ms | " + f"input-quant {t_quant * 1e3:.3f} ms + fp4-gemm {t_gemm * 1e3:.3f} ms " + f"= {total * 1e3:.3f} ms ({t_bf16 / max(total, 1e-9):.2f}x net, " + f"{t_bf16 / max(t_gemm, 1e-9):.2f}x gemm-only) -> {verdict}") + del x, qx except Exception as e: # noqa: BLE001 lines.append(f" GEMM probe failed: {e}") else: