From dd6d2c8880891ff9274656e6ad34f8644012c25a Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 21 Aug 2026 05:21:18 -0700 Subject: [PATCH 1/4] gemm: add a dense GELU MLP torch op --- docs/operations/GeluMLP.md | 50 ++++ llms.txt | 1 + python/cudnn/README.md | 6 +- python/cudnn/gemm/__init__.py | 5 +- python/cudnn/gemm/ops/__init__.py | 1 + python/cudnn/gemm/ops/_gelu_mlp.py | 406 +++++++++++++++++++++++++++++ test/python/gemm/test_gelu_mlp.py | 387 +++++++++++++++++++++++++++ 7 files changed, 852 insertions(+), 4 deletions(-) create mode 100644 docs/operations/GeluMLP.md create mode 100644 python/cudnn/gemm/ops/_gelu_mlp.py create mode 100644 test/python/gemm/test_gelu_mlp.py diff --git a/docs/operations/GeluMLP.md b/docs/operations/GeluMLP.md new file mode 100644 index 000000000..d8b03b874 --- /dev/null +++ b/docs/operations/GeluMLP.md @@ -0,0 +1,50 @@ +# GELU MLP + +`cudnn.gemm.ops.gelu_mlp` is a PyTorch-facing cuDNN implementation of the +dense feed-forward block used by DiT and ViT models: + +```python +from cudnn.gemm.ops import gelu_mlp + +y = gelu_mlp(x, w1, b1, w2, b2) +``` + +Its observable computation matches two ordinary `torch.nn.Linear` layers +separated by tanh-approximate GELU: + +```python +h = torch.nn.functional.gelu( + torch.nn.functional.linear(x, w1, b1), approximate="tanh" +) +y = torch.nn.functional.linear(h, w2, b2) +``` + +`x` has shape `[..., H]`; weights use the standard `nn.Linear` layouts +`w1[I, H]` and `w2[O, I]`; biases have shapes `b1[I]` and `b2[O]`. The output +has shape `[..., O]`. + +## Fusion and numerical boundaries + +Forward executes two cuDNN graphs. The first fuses the first matrix multiply, +column bias, and tanh-GELU. The second fuses the output matrix multiply and +bias. The post-bias first-layer value is rounded to BF16 before GELU, matching +the visible boundary of eager BF16 `Linear`; GEMMs accumulate in FP32. + +First-order autograd is supported for all five inputs. Backward fuses +`dout @ w2` with GELU backward, uses cuDNN GEMMs for activation and weight +gradients, and currently uses PyTorch reductions for bias gradients. Higher +order gradients are not supported and fail explicitly. + +## Current support + +- SM100 GPUs +- contiguous BF16 CUDA tensors on one device +- rank-two-or-higher `x` +- tanh-approximate GELU, with biases and no dropout inside the operation +- eager PyTorch execution and first-order autograd + +The first call for a new shape/device/stream builds and autotunes plans, so warm +the operation before timing. This release does not claim `torch.compile` or +CUDA Graph capture compatibility. Unsupported shapes, layouts, dtypes, +devices, or architectures fail rather than copying, converting, or falling +back silently. diff --git a/llms.txt b/llms.txt index b9866ae70..7653b4bf1 100644 --- a/llms.txt +++ b/llms.txt @@ -18,6 +18,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [Convolutions](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Convolutions.md) - [Normalizations (LayerNorm, RMSNorm, BatchNorm, InstanceNorm)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Normalizations.md) - [MoE Grouped Matmul](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/MoeGroupedMatmul.md) +- [Dense GELU MLP PyTorch operation](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/GeluMLP.md) - [Pointwise](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Pointwise.md) - [Block Scaling (MXFP8/NVFP4 quantization)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/BlockScaling.md) - [RoPE](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/RoPE.md) diff --git a/python/cudnn/README.md b/python/cudnn/README.md index dd72a170f..d94f0be8b 100644 --- a/python/cudnn/README.md +++ b/python/cudnn/README.md @@ -47,8 +47,10 @@ python/cudnn/gemm/ └── reference/ # pure-PyTorch correctness engine ``` -Every public GEMM symbol is re-exported at the top level (`cudnn.`), which -is the supported entry point — the directory layout is an implementation detail. +Backend-independent PyTorch operation contracts are available from +`cudnn.gemm.ops` and re-exported from `cudnn.gemm`. For example, +`cudnn.gemm.ops.gelu_mlp` implements the biased tanh-GELU feed-forward block; +see [GELU MLP](../../docs/operations/GeluMLP.md). ## diff --git a/python/cudnn/gemm/__init__.py b/python/cudnn/gemm/__init__.py index dc777717d..8375fd46e 100644 --- a/python/cudnn/gemm/__init__.py +++ b/python/cudnn/gemm/__init__.py @@ -11,13 +11,14 @@ cudnn.gemm.ops backend-independent op contracts cudnn.gemm.reference pure-PyTorch correctness engine -Every public symbol is also re-exported at the top level (``cudnn.``), -which is the supported entry point for users. +Operation symbols are re-exported from ``cudnn.gemm``; their defining +``cudnn.gemm.ops`` namespace remains supported as well. """ from typing import Any _LAZY_EXPORTS = { + "gelu_mlp": ("cudnn.gemm.ops", "gelu_mlp"), "moe_grouped_matmul": ("cudnn.gemm.ops", "moe_grouped_matmul"), "swiglu_mlp": ("cudnn.gemm.ops", "swiglu_mlp"), } diff --git a/python/cudnn/gemm/ops/__init__.py b/python/cudnn/gemm/ops/__init__.py index 2e189eb4f..8ef8d4756 100644 --- a/python/cudnn/gemm/ops/__init__.py +++ b/python/cudnn/gemm/ops/__init__.py @@ -9,6 +9,7 @@ # cudnn.gemm.ops`` stays frontend-only and torch is imported only when a # specific operation is first accessed. Mirrors cudnn/gemm/__init__.py. _LAZY_EXPORTS = { + "gelu_mlp": ("cudnn.gemm.ops._gelu_mlp", "gelu_mlp"), "moe_grouped_matmul": ("cudnn.gemm.ops.moe_grouped_matmul", "moe_grouped_matmul"), "swiglu_mlp": ("cudnn.gemm.ops.swiglu_mlp", "swiglu_mlp"), } diff --git a/python/cudnn/gemm/ops/_gelu_mlp.py b/python/cudnn/gemm/ops/_gelu_mlp.py new file mode 100644 index 000000000..143502bfc --- /dev/null +++ b/python/cudnn/gemm/ops/_gelu_mlp.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Private implementation of the dense bf16 GELU MLP cuDNN op. + +The public contract matches two ordinary ``torch.nn.Linear`` layers separated +by ``GELU(approximate="tanh")``:: + + h = gelu(x @ w1.T + b1, approximate="tanh") + y = h @ w2.T + b2 + +The input must have rank two or greater. ``w1`` and ``w2`` use the natural ``nn.Linear`` layouts +``[intermediate, in_features]`` and ``[out_features, intermediate]``. The +first graph fuses matmul, per-column bias and tanh-GELU; the second graph fuses +matmul and bias. The post-bias first-layer value is explicitly bf16-rounded +before GELU, preserving the observable boundary of eager bf16 ``Linear``. + +Backward uses cuDNN GEMMs for the input and weight gradients. Its first stage +is one ``dout @ w2`` + tanh-GELU-backward graph, avoiding a materialized ``dh``; +bias gradients are reductions in PyTorch for now. + +The implementation is intentionally narrow: dense contiguous bf16 CUDA tensors +on SM100. Unsupported layouts, devices, architectures, or dtypes raise instead +of being silently copied or reinterpreted. +""" + +from __future__ import annotations + +import torch + +import cudnn + +_BF16 = cudnn.data_type.BFLOAT16 +_FP32 = cudnn.data_type.FLOAT +_AUTOTUNE_ITERS = 20 + +# A cuDNN handle, execution plans, and scratch workspaces are private to one +# (device, stream). A stream serializes reuse of its workspace, while distinct +# streams never race through one handle or allocation. +_HANDLES = {} +_LINEAR_CACHE = {} +_MM_CACHE = {} +_DGELU_CACHE = {} + +_GRAD_X = 1 << 0 +_GRAD_W1 = 1 << 1 +_GRAD_B1 = 1 << 2 +_GRAD_W2 = 1 << 3 +_GRAD_B2 = 1 << 4 +_GRAD_FC1 = _GRAD_X | _GRAD_W1 | _GRAD_B1 + + +def _handle(device: torch.device): + stream = torch.cuda.current_stream(device).cuda_stream + key = (device.index, stream) + handle = _HANDLES.get(key) + if handle is None: + with torch.cuda.device(device): + handle = cudnn.create_handle() + cudnn.set_stream(handle=handle, stream=stream) + _HANDLES[key] = handle + return handle, stream + + +def _autotune(graph, handle, variant_pack): + """Build and time every viable graph plan, returning plan and workspace.""" + graph.check_support() + graph.build_plans(cudnn.build_plan_policy.ALL) + count = graph.get_execution_plan_count() + if count == 0: + raise RuntimeError("cudnn.gemm.gelu_mlp: no execution plan was generated for this graph") + + device = next(iter(variant_pack.values())).device + elapsed = [float("inf")] * count + errors = {} + with torch.cuda.device(device): + workspace = torch.empty( + max(graph.get_workspace_size_plan_at_index(i) for i in range(count)), + device=device, + dtype=torch.uint8, + ) + start = torch.cuda.Event(enable_timing=True) + stop = torch.cuda.Event(enable_timing=True) + for index in range(count): + try: + graph.execute_plan_at_index(variant_pack, workspace, index=index, handle=handle) + torch.cuda.synchronize(device) + start.record() + for _ in range(_AUTOTUNE_ITERS): + graph.execute_plan_at_index(variant_pack, workspace, index=index, handle=handle) + stop.record() + stop.synchronize() + elapsed[index] = start.elapsed_time(stop) / _AUTOTUNE_ITERS + except Exception as exc: # noqa: BLE001 -- one invalid plan must not suppress viable plans + errors[index] = repr(exc) + + best = min(range(count), key=elapsed.__getitem__) + if elapsed[best] == float("inf"): + raise RuntimeError(f"cudnn.gemm.gelu_mlp: all {count} autotune plans failed to execute; errors: {errors}") + return best, workspace + + +def _mm(a2: torch.Tensor, b2: torch.Tensor) -> torch.Tensor: + """Execute ``[M,K] @ [K,N]`` as an autotuned cuDNN graph.""" + with torch.cuda.device(a2.device): + handle, stream = _handle(a2.device) + av = a2.unsqueeze(0) + bv = b2.unsqueeze(0) + key = ( + tuple(av.shape), + tuple(av.stride()), + tuple(bv.shape), + tuple(bv.stride()), + a2.dtype, + a2.device.index, + stream, + ) + entry = _MM_CACHE.get(key) + if entry is None: + graph = cudnn.pygraph(handle=handle, compute_data_type=_FP32) + A = graph.tensor(dim=list(av.shape), stride=list(av.stride()), data_type=_BF16) + B = graph.tensor(dim=list(bv.shape), stride=list(bv.stride()), data_type=_BF16) + C = graph.matmul(name="mm", A=A, B=B, compute_data_type=_FP32) + C.set_output(True).set_data_type(_BF16) + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + output = torch.empty((1, av.shape[1], bv.shape[2]), device=a2.device, dtype=a2.dtype) + best, workspace = _autotune(graph, handle, {A: av, B: bv, C: output}) + entry = (graph, A, B, C, best, workspace) + _MM_CACHE[key] = entry + + graph, A, B, C, best, workspace = entry + output = torch.empty((1, av.shape[1], bv.shape[2]), device=a2.device, dtype=a2.dtype) + graph.execute_plan_at_index({A: av, B: bv, C: output}, workspace, index=best, handle=handle) + return output.squeeze(0) + + +def _linear_bias( + x2: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + *, + gelu: bool, + save_pre_activation: bool = False, +): + """Run one fused linear epilogue and return ``(output, pre_activation)``. + + With ``gelu=True``, ``pre_activation`` is the bf16-rounded post-bias value + and is emitted only when backward needs it. The h-only and h+z graphs use + distinct cache entries because output taps are part of a graph's contract. + """ + if save_pre_activation and not gelu: + raise ValueError("cudnn.gemm.gelu_mlp: a pre-activation tap is only valid for the GELU layer") + + with torch.cuda.device(x2.device): + handle, stream = _handle(x2.device) + xv = x2.unsqueeze(0) + wv = weight.t().unsqueeze(0) + bv = bias.view(1, 1, -1) + key = ( + bool(gelu), + bool(save_pre_activation), + tuple(xv.shape), + tuple(xv.stride()), + tuple(wv.shape), + tuple(wv.stride()), + tuple(bv.stride()), + x2.dtype, + x2.device.index, + stream, + ) + entry = _LINEAR_CACHE.get(key) + if entry is None: + graph = cudnn.pygraph(handle=handle, compute_data_type=_FP32) + X = graph.tensor(dim=list(xv.shape), stride=list(xv.stride()), data_type=_BF16) + W = graph.tensor(dim=list(wv.shape), stride=list(wv.stride()), data_type=_BF16) + BIAS = graph.tensor(dim=list(bv.shape), stride=list(bv.stride()), data_type=_BF16) + mm = graph.matmul(name="linear", A=X, B=W, compute_data_type=_FP32) + pre_activation = graph.bias(input=mm, bias=BIAS, name="bias") + # Preserve eager bf16 Linear -> GELU semantics even while the two + # operations stay in one graph and need not round-trip through HBM. + pre_activation.set_data_type(_BF16) + if gelu: + output_tensor = graph.gelu_approx_tanh(input=pre_activation, name="gelu_tanh") + output_tensor.set_output(True).set_data_type(_BF16) + if save_pre_activation: + pre_activation.set_output(True) + else: + output_tensor = pre_activation + output_tensor.set_output(True) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + output = torch.empty((1, xv.shape[1], wv.shape[2]), device=x2.device, dtype=x2.dtype) + variant_pack = {X: xv, W: wv, BIAS: bv, output_tensor: output} + if gelu and save_pre_activation: + saved_pre_activation = torch.empty_like(output) + variant_pack[pre_activation] = saved_pre_activation + best, workspace = _autotune(graph, handle, variant_pack) + entry = ( + graph, + X, + W, + BIAS, + output_tensor, + pre_activation, + best, + workspace, + ) + _LINEAR_CACHE[key] = entry + + graph, X, W, BIAS, output_tensor, pre_activation, best, workspace = entry + output = torch.empty((1, xv.shape[1], wv.shape[2]), device=x2.device, dtype=x2.dtype) + variant_pack = {X: xv, W: wv, BIAS: bv, output_tensor: output} + if gelu and save_pre_activation: + saved_pre_activation = torch.empty_like(output) + variant_pack[pre_activation] = saved_pre_activation + else: + saved_pre_activation = None + graph.execute_plan_at_index(variant_pack, workspace, index=best, handle=handle) + return output.squeeze(0), None if saved_pre_activation is None else saved_pre_activation.squeeze(0) + + +def _linear_dgelu(dout2: torch.Tensor, w2: torch.Tensor, pre_activation: torch.Tensor): + """Fuse ``dh = dout @ w2`` with tanh-GELU backward into one cuDNN graph.""" + with torch.cuda.device(dout2.device): + handle, stream = _handle(dout2.device) + dyv = dout2.unsqueeze(0) + wv = w2.unsqueeze(0) + zv = pre_activation.unsqueeze(0) + key = ( + tuple(dyv.shape), + tuple(dyv.stride()), + tuple(wv.shape), + tuple(wv.stride()), + tuple(zv.stride()), + dout2.dtype, + dout2.device.index, + stream, + ) + entry = _DGELU_CACHE.get(key) + if entry is None: + graph = cudnn.pygraph(handle=handle, compute_data_type=_FP32) + DY = graph.tensor(dim=list(dyv.shape), stride=list(dyv.stride()), data_type=_BF16) + W2 = graph.tensor(dim=list(wv.shape), stride=list(wv.stride()), data_type=_BF16) + Z = graph.tensor(dim=list(zv.shape), stride=list(zv.stride()), data_type=_BF16) + dh = graph.matmul(name="dhidden", A=DY, B=W2, compute_data_type=_FP32) + # Match the bf16 gradient crossing the eager second Linear boundary. + dh.set_data_type(_BF16) + DZ = graph.gelu_approx_tanh_backward(loss=dh, input=Z, name="dgelu_tanh") + DZ.set_output(True).set_data_type(_BF16) + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + dz = torch.empty_like(zv) + best, workspace = _autotune(graph, handle, {DY: dyv, W2: wv, Z: zv, DZ: dz}) + entry = (graph, DY, W2, Z, DZ, best, workspace) + _DGELU_CACHE[key] = entry + + graph, DY, W2, Z, DZ, best, workspace = entry + dz = torch.empty_like(zv) + graph.execute_plan_at_index( + {DY: dyv, W2: wv, Z: zv, DZ: dz}, + workspace, + index=best, + handle=handle, + ) + return dz.squeeze(0) + + +class _GeluMLP(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w1, b1, w2, b2, grad_mask): + input_shape = x.shape + x2 = x.reshape(-1, input_shape[-1]) + need_fc1_grad = bool(grad_mask & _GRAD_FC1) + hidden, pre_activation = _linear_bias( + x2, + w1, + b1, + gelu=True, + save_pre_activation=need_fc1_grad, + ) + output, _ = _linear_bias(hidden, w2, b2, gelu=False) + + saved_names = [] + saved_tensors = [] + + def save(name, tensor): + saved_names.append(name) + saved_tensors.append(tensor) + + if grad_mask & _GRAD_W1: + save("x2", x2) + if grad_mask & _GRAD_X: + save("w1", w1) + if need_fc1_grad: + save("w2", w2) + save("pre_activation", pre_activation) + if grad_mask & _GRAD_W2: + save("hidden", hidden) + + ctx.save_for_backward(*saved_tensors) + ctx.saved_names = tuple(saved_names) + ctx.grad_mask = grad_mask + ctx.input_shape = input_shape + ctx.out_features = w2.shape[0] + return output.reshape(*input_shape[:-1], w2.shape[0]) + + @staticmethod + def backward(ctx, dout): + if torch.is_grad_enabled(): + raise NotImplementedError("cudnn.gemm.gelu_mlp: double backward is not supported") + saved = dict(zip(ctx.saved_names, ctx.saved_tensors)) + grad_mask = ctx.grad_mask + dout2 = dout.reshape(-1, ctx.out_features).contiguous() + + dw2 = _mm(dout2.t(), saved["hidden"]) if grad_mask & _GRAD_W2 else None + db2 = dout2.sum(dim=0) if grad_mask & _GRAD_B2 else None + + if grad_mask & _GRAD_FC1: + dz = _linear_dgelu(dout2, saved["w2"], saved["pre_activation"]) + else: + dz = None + + dx = _mm(dz, saved["w1"]).reshape(ctx.input_shape) if grad_mask & _GRAD_X else None + dw1 = _mm(dz.t(), saved["x2"]) if grad_mask & _GRAD_W1 else None + db1 = dz.sum(dim=0) if grad_mask & _GRAD_B1 else None + return dx, dw1, db1, dw2, db2, None + + +def _validate(x, w1, b1, w2, b2): + operands = (("x", x), ("w1", w1), ("b1", b1), ("w2", w2), ("b2", b2)) + for name, tensor in operands: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"cudnn.gemm.gelu_mlp: {name} must be a torch.Tensor") + if tensor.dtype != torch.bfloat16: + raise TypeError(f"cudnn.gemm.gelu_mlp: {name} must be bfloat16, got {tensor.dtype}") + if tensor.device.type != "cuda": + raise ValueError(f"cudnn.gemm.gelu_mlp: {name} must be a CUDA tensor, got device {tensor.device}") + + if any(tensor.device != x.device for _, tensor in operands[1:]): + devices = ", ".join(f"{name}={tensor.device}" for name, tensor in operands) + raise ValueError(f"cudnn.gemm.gelu_mlp: all operands must be on the same CUDA device; got {devices}") + + if x.dim() < 2 or w1.dim() != 2 or b1.dim() != 1 or w2.dim() != 2 or b2.dim() != 1: + raise ValueError( + "cudnn.gemm.gelu_mlp: expected x[...,H], w1[I,H], b1[I], " + f"w2[O,I], b2[O]; got x{tuple(x.shape)}, w1{tuple(w1.shape)}, " + f"b1{tuple(b1.shape)}, w2{tuple(w2.shape)}, b2{tuple(b2.shape)}" + ) + + in_features = x.shape[-1] + intermediate = w1.shape[0] + out_features = w2.shape[0] + if w1.shape[1] != in_features or b1.shape[0] != intermediate or w2.shape[1] != intermediate or b2.shape[0] != out_features: + raise ValueError( + f"cudnn.gemm.gelu_mlp: shape mismatch for x[...,{in_features}], " + f"w1{tuple(w1.shape)}, b1{tuple(b1.shape)}, w2{tuple(w2.shape)}, b2{tuple(b2.shape)}; " + f"expected w1=[I,{in_features}], b1=[I], w2=[O,I], b2=[O]" + ) + if x.numel() == 0 or in_features == 0 or intermediate == 0 or out_features == 0: + raise ValueError("cudnn.gemm.gelu_mlp: zero-sized dimensions are not supported") + + noncontiguous = [name for name, tensor in operands if not tensor.is_contiguous()] + if noncontiguous: + raise ValueError("cudnn.gemm.gelu_mlp: operands must use dense contiguous nn.Linear layouts; " f"noncontiguous: {', '.join(noncontiguous)}") + + capability = torch.cuda.get_device_capability(x.device) + if capability != (10, 0): + raise NotImplementedError("cudnn.gemm.gelu_mlp: this implementation requires SM100; " f"got sm_{capability[0]}{capability[1]} on {x.device}") + + +def gelu_mlp(x, w1, b1, w2, b2): + """Run a dense bf16 tanh-GELU MLP on cuDNN. + + Args: + x: Dense rank-two-or-higher input activations ``[..., H]``. + w1: First ``nn.Linear`` weight ``[I, H]``. + b1: First ``nn.Linear`` bias ``[I]``. + w2: Second ``nn.Linear`` weight ``[O, I]``. + b2: Second ``nn.Linear`` bias ``[O]``. + + Returns: + Dense bf16 output ``[..., O]``. First-order gradients are supported + with respect to all five tensor inputs; double backward is not. + + All inputs must be contiguous bf16 tensors on the same SM100 CUDA device. + GELU always uses ``approximate="tanh"``. + """ + _validate(x, w1, b1, w2, b2) + + grad_mask = 0 + if torch.is_grad_enabled(): + for bit, tensor in ( + (_GRAD_X, x), + (_GRAD_W1, w1), + (_GRAD_B1, b1), + (_GRAD_W2, w2), + (_GRAD_B2, b2), + ): + if tensor.requires_grad: + grad_mask |= bit + return _GeluMLP.apply(x, w1, b1, w2, b2, grad_mask) diff --git a/test/python/gemm/test_gelu_mlp.py b/test/python/gemm/test_gelu_mlp.py new file mode 100644 index 000000000..60b970f4e --- /dev/null +++ b/test/python/gemm/test_gelu_mlp.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Correctness and contract tests for ``cudnn.gemm.ops.gelu_mlp``.""" + +import importlib +import subprocess +import sys + +import pytest +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from cudnn.gemm.ops import gelu_mlp + + +def _cc(): + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor + + +_SM100 = pytest.mark.skipif( + not (torch.cuda.is_available() and _cc() == 100), + reason="cuDNN GELU-MLP requires SM100", +) +_TOL = 2e-2 + + +def _ref(x, w1, b1, w2, b2): + hidden = F.gelu(F.linear(x, w1, b1), approximate="tanh") + return F.linear(hidden, w2, b2) + + +def _rel_l2(actual, expected): + return (actual.float() - expected.float()).norm().item() / max(expected.float().norm().item(), 1e-9) + + +def _inputs(*, requires=(False, False, False, False, False)): + torch.manual_seed(0) + M, H, intermediate, O = 128, 256, 512, 192 + base = ( + torch.randn(2, M, H, device="cuda", dtype=torch.bfloat16), + torch.randn(intermediate, H, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(intermediate, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(O, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02, + torch.randn(O, device="cuda", dtype=torch.bfloat16) * 0.02, + ) + return tuple(t.detach().requires_grad_(need) for t, need in zip(base, requires)) + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_full_autograd_parity(): + required = (True,) * 5 + args = _inputs(requires=required) + refs = tuple(t.detach().clone().requires_grad_(True) for t in args) + dout = torch.randn(2, 128, 192, device="cuda", dtype=torch.bfloat16) + + output = gelu_mlp(*args) + reference = _ref(*refs) + output.backward(dout) + reference.backward(dout) + + assert _rel_l2(output, reference) < _TOL + for name, actual, expected in zip(("dx", "dw1", "db1", "dw2", "db2"), args, refs): + assert _rel_l2(actual.grad, expected.grad) < _TOL, name + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_sum_loss_backward_parity(): + """Exercise the expanded zero-stride gradient produced by a sum loss.""" + args = _inputs(requires=(True,) * 5) + refs = tuple(t.detach().clone().requires_grad_(True) for t in args) + + gelu_mlp(*args).sum().backward() + _ref(*refs).sum().backward() + + for name, actual, expected in zip(("dx", "dw1", "db1", "dw2", "db2"), args, refs): + assert _rel_l2(actual.grad, expected.grad) < _TOL, name + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_double_backward_fails_closed(): + args = _inputs(requires=(True,) * 5) + output = gelu_mlp(*args) + + with pytest.raises(NotImplementedError, match="double backward is not supported"): + torch.autograd.grad(output.sum(), args, create_graph=True) + + +@pytest.mark.L0 +@_SM100 +@pytest.mark.parametrize("tokens", [512, 4096], ids=["text-512", "image-4096"]) +def test_gelu_mlp_qwen_image_sequence_length_parity(tokens): + """Cover both token counts used by the Qwen-Image image/text FFNs.""" + torch.manual_seed(11) + H, intermediate = 128, 256 + x = torch.randn(1, tokens, H, device="cuda", dtype=torch.bfloat16) + w1 = torch.randn(intermediate, H, device="cuda", dtype=torch.bfloat16) * 0.02 + b1 = torch.randn(intermediate, device="cuda", dtype=torch.bfloat16) * 0.02 + w2 = torch.randn(H, intermediate, device="cuda", dtype=torch.bfloat16) * 0.02 + b2 = torch.randn(H, device="cuda", dtype=torch.bfloat16) * 0.02 + + with torch.no_grad(): + output = gelu_mlp(x, w1, b1, w2, b2) + reference = _ref(x, w1, b1, w2, b2) + + assert _rel_l2(output, reference) < _TOL + + +@pytest.mark.L0 +@_SM100 +@pytest.mark.parametrize( + "required,expected_saved", + [ + ((True, False, False, False, False), ("w1", "w2", "pre_activation")), + ((False, True, False, False, False), ("x2", "w2", "pre_activation")), + ((False, False, True, False, False), ("w2", "pre_activation")), + ((False, False, False, True, False), ("hidden",)), + ((False, False, False, False, True), ()), + ((False, True, True, True, True), ("x2", "w2", "pre_activation", "hidden")), + ], + ids=["x-only", "w1-only", "b1-only", "w2-only", "b2-only", "weights-and-biases"], +) +def test_gelu_mlp_partial_grad(required, expected_saved): + args = _inputs(requires=required) + refs = tuple(t.detach().clone().requires_grad_(need) for t, need in zip(args, required)) + dout = torch.randn(2, 128, 192, device="cuda", dtype=torch.bfloat16) + + output = gelu_mlp(*args) + reference = _ref(*refs) + assert tuple(output.grad_fn.saved_names) == expected_saved + output.backward(dout) + reference.backward(dout) + + assert _rel_l2(output, reference) < _TOL + for name, actual, expected, need in zip(("dx", "dw1", "db1", "dw2", "db2"), args, refs, required): + if need: + assert _rel_l2(actual.grad, expected.grad) < _TOL, name + else: + assert actual.grad is None, name + + +@pytest.mark.L0 +@_SM100 +@pytest.mark.parametrize( + "grad_mode", + [torch.no_grad, torch.inference_mode], + ids=["no-grad", "inference-mode"], +) +def test_gelu_mlp_inference_omits_pre_activation(monkeypatch, grad_mode): + module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + original = module._linear_bias + observed = [] + + def record(*args, **kwargs): + observed.append((kwargs["gelu"], kwargs.get("save_pre_activation", False))) + return original(*args, **kwargs) + + monkeypatch.setattr(module, "_linear_bias", record) + args = _inputs(requires=(True,) * 5) + with grad_mode(): + output = module.gelu_mlp(*args) + reference = _ref(*args) + + assert not output.requires_grad + assert observed == [(True, False), (False, False)] + assert _rel_l2(output, reference) < _TOL + + +@pytest.mark.L0 +@_SM100 +@pytest.mark.parametrize("use_reentrant", [True, False], ids=["reentrant", "non-reentrant"]) +def test_gelu_mlp_checkpoint_parity(use_reentrant): + args = _inputs(requires=(True,) * 5) + refs = tuple(t.detach().clone().requires_grad_(True) for t in args) + dout = torch.randn(2, 128, 192, device="cuda", dtype=torch.bfloat16) + + output = torch.utils.checkpoint.checkpoint(gelu_mlp, *args, use_reentrant=use_reentrant) + reference = _ref(*refs) + output.backward(dout) + reference.backward(dout) + + assert _rel_l2(output, reference) < _TOL + for actual, expected in zip(args, refs): + assert _rel_l2(actual.grad, expected.grad) < _TOL + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_cache_is_stream_local(): + module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + module._HANDLES.clear() + module._LINEAR_CACHE.clear() + module._MM_CACHE.clear() + module._DGELU_CACHE.clear() + args = _inputs() + + gelu_mlp(*args) + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + gelu_mlp(*args) + torch.cuda.current_stream().wait_stream(side_stream) + + streams = {key[-1] for key in module._LINEAR_CACHE} + assert len(module._HANDLES) == 2 + assert len(streams) == 2 + assert len(module._LINEAR_CACHE) == 4 # two stages on each stream + + +@pytest.mark.L0 +@_SM100 +@pytest.mark.parametrize("save_pre_activation", [False, True], ids=["inference", "training"]) +def test_gelu_mlp_fc1_is_one_fused_kernel(save_pre_activation): + """The first Linear, bias, bf16 boundary, and tanh-GELU stay one launch.""" + module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + x, w1, b1, _, _ = _inputs() + x2 = x.reshape(-1, x.shape[-1]) + module._linear_bias( + x2, + w1, + b1, + gelu=True, + save_pre_activation=save_pre_activation, + ) + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: + module._linear_bias( + x2, + w1, + b1, + gelu=True, + save_pre_activation=save_pre_activation, + ) + torch.cuda.synchronize() + + launches = sum(event.count for event in profiler.key_averages() if event.self_device_time_total > 0) + stream = torch.cuda.current_stream().cuda_stream + matching = [ + entry + for key, entry in module._LINEAR_CACHE.items() + if key[-1] == stream and key[0] is True and key[1] is save_pre_activation and key[2] == (1, x2.shape[0], x2.shape[1]) + ] + assert len(matching) == 1 + graph, *_, best, _workspace = matching[0] + plan_name = graph.get_plan_name_at_index(best) + assert launches == 1, f"expected one FC1 fusion launch from {plan_name}, saw {launches}" + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_inference_forward_is_exactly_two_kernels(): + """The complete forward is fused FC1+bias+GELU plus FC2+bias.""" + args = _inputs() + with torch.no_grad(): + gelu_mlp(*args) + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: + with torch.no_grad(): + gelu_mlp(*args) + torch.cuda.synchronize() + + launches = sum(event.count for event in profiler.key_averages() if event.self_device_time_total > 0) + assert launches == 2, f"expected exactly two GELU-MLP forward launches, saw {launches}" + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_dhidden_dgelu_is_one_fused_kernel(): + """Backward must not materialize ``dhidden`` between GEMM and dGELU.""" + module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + x, w1, b1, w2, _ = _inputs() + x2 = x.reshape(-1, x.shape[-1]) + _, pre_activation = module._linear_bias( + x2, + w1, + b1, + gelu=True, + save_pre_activation=True, + ) + dout2 = torch.randn(x2.shape[0], w2.shape[0], device="cuda", dtype=torch.bfloat16) + module._linear_dgelu(dout2, w2, pre_activation) + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: + module._linear_dgelu(dout2, w2, pre_activation) + torch.cuda.synchronize() + + launches = sum(event.count for event in profiler.key_averages() if event.self_device_time_total > 0) + stream = torch.cuda.current_stream().cuda_stream + matching = [entry for key, entry in module._DGELU_CACHE.items() if key[-1] == stream and key[0] == (1, x2.shape[0], w2.shape[0])] + assert len(matching) == 1 + graph, *_, best, _workspace = matching[0] + plan_name = graph.get_plan_name_at_index(best) + assert launches == 1, f"expected one dLinear+dGELU launch from {plan_name}, saw {launches}" + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_backward_runtime_failure_is_not_hidden(monkeypatch): + module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + + def fail(*_args, **_kwargs): + raise RuntimeError("synthetic fused backward launch failure") + + monkeypatch.setattr(module, "_linear_dgelu", fail) + args = _inputs(requires=(True, False, False, False, False)) + output = module.gelu_mlp(*args) + with pytest.raises(RuntimeError, match="synthetic fused backward launch failure"): + output.sum().backward() + + +@pytest.mark.L0 +def test_gelu_mlp_is_lazy_public_export(): + assert callable(gelu_mlp) + import cudnn.gemm + + assert cudnn.gemm.gelu_mlp is gelu_mlp + + +@pytest.mark.L0 +def test_gelu_mlp_public_export_survives_internal_module_first_import(): + code = """ +import importlib +import cudnn.gemm.ops as ops +importlib.import_module('cudnn.gemm.ops._gelu_mlp') +assert callable(ops.gelu_mlp) +import cudnn.gemm +assert cudnn.gemm.gelu_mlp is ops.gelu_mlp +""" + subprocess.run([sys.executable, "-c", code], check=True) + + +@pytest.mark.parametrize( + "mutate,error,match", + [ + (lambda xs: xs.__setitem__(0, xs[0].float()), TypeError, "x must be bfloat16"), + ], + ids=["dtype"], +) +@pytest.mark.L0 +def test_gelu_mlp_cpu_validation(mutate, error, match): + tensors = [ + torch.empty(2, 3, 4, dtype=torch.bfloat16), + torch.empty(8, 4, dtype=torch.bfloat16), + torch.empty(8, dtype=torch.bfloat16), + torch.empty(6, 8, dtype=torch.bfloat16), + torch.empty(6, dtype=torch.bfloat16), + ] + mutate(tensors) + with pytest.raises(error, match=match): + gelu_mlp(*tensors) + + +@pytest.mark.L0 +def test_gelu_mlp_rejects_cpu_operands(): + tensors = [ + torch.empty(2, 3, 4, dtype=torch.bfloat16), + torch.empty(8, 4, dtype=torch.bfloat16), + torch.empty(8, dtype=torch.bfloat16), + torch.empty(6, 8, dtype=torch.bfloat16), + torch.empty(6, dtype=torch.bfloat16), + ] + with pytest.raises(ValueError, match="x must be a CUDA tensor"): + gelu_mlp(*tensors) + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_rejects_wrong_rank(): + x, w1, b1, w2, b2 = _inputs() + with pytest.raises(ValueError, match=r"expected x\[\.\.\.,H\]"): + gelu_mlp(x, w1, b1.reshape(1, -1), w2, b2) + + +@pytest.mark.L0 +@_SM100 +def test_gelu_mlp_rejects_noncontiguous_nn_linear_weight(): + x, w1, b1, w2, b2 = _inputs() + square = torch.empty(w1.shape[0], w1.shape[0], device="cuda", dtype=torch.bfloat16) + w1_bad = square.t() + x_bad = torch.empty(*x.shape[:-1], w1_bad.shape[1], device="cuda", dtype=torch.bfloat16) + w2_bad_shape = torch.empty(w2.shape[0], w1_bad.shape[0], device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="noncontiguous: w1"): + gelu_mlp(x_bad, w1_bad, b1, w2_bad_shape, b2) From 2d65aef64241d237abacd3223df793ecd3abb6ca Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 21 Aug 2026 05:21:27 -0700 Subject: [PATCH 2/4] benchmark: add a Qwen-Image MLP and SDPA matrix --- benchmark/e2e/Qwen-Image/run_bf16.py | 288 +++++++++++++++----- benchmark/e2e/Qwen-Image/run_model.py | 173 +++++++++++- benchmark/e2e/tests/test_qwen_image_spec.py | 98 +++++++ 3 files changed, 487 insertions(+), 72 deletions(-) diff --git a/benchmark/e2e/Qwen-Image/run_bf16.py b/benchmark/e2e/Qwen-Image/run_bf16.py index 625f896ed..2f78639d5 100644 --- a/benchmark/e2e/Qwen-Image/run_bf16.py +++ b/benchmark/e2e/Qwen-Image/run_bf16.py @@ -2,12 +2,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Balanced BF16 Qwen-Image joint-attention A/B benchmark. +"""Balanced BF16 Qwen-Image attention x GELU-MLP factorial benchmark. This is one numerical-recipe leaf: conservative BF16. It compares explicitly -forced PyTorch FlashAttention with the FE public cuDNN backend graph while keeping Diffusers' -Q/K/V projections, QK norm, RoPE, output projections, AdaLN, GELU FFNs, and all -other transformer work identical. +forced PyTorch FlashAttention with the FE public cuDNN backend graph and the +pinned Diffusers GELU FFN with ``cudnn.gemm.ops.gelu_mlp``. All other +transformer work remains identical. """ from __future__ import annotations @@ -33,13 +33,19 @@ FACTORIAL_PATH = E2E_DIR / "_factorial.py" sys.path.insert(0, str(E2E_DIR)) -from _factorial import config_fingerprint, paired_stats, percentile, williams_orders # noqa: E402 +from _factorial import config_fingerprint, factorial_main_effects, paired_stats, percentile, shapley_savings, williams_orders # noqa: E402 PROTOCOL_DEFAULTS = { "smoke": {"warmup": 1, "rounds": 8, "repeats": 1}, "formal": {"warmup": 3, "rounds": 40, "repeats": 3}, } -VARIANTS = (("0", "torch_flash"), ("1", "cudnn")) +VARIANTS = ( + ("00", "torch", "torch_flash"), + ("01", "torch", "cudnn"), + ("10", "cudnn", "torch_flash"), + ("11", "cudnn", "cudnn"), +) +AXIS_MASKS = {"mlp": 2, "attn": 1} def _utc_now(): @@ -107,8 +113,8 @@ def _resolve_protocol(args): protocol[name] = value if any(not isinstance(value, int) or value <= 0 for value in protocol.values()): raise ValueError(f"protocol values must be positive integers, got {protocol}") - if protocol["rounds"] % 2: - raise ValueError("rounds must be a multiple of 2 for the complete two-treatment Williams design") + if protocol["rounds"] % 4: + raise ValueError("rounds must be a multiple of 4 for the complete four-treatment Williams design") return protocol @@ -117,9 +123,11 @@ def _pick_device(torch, mode): for index in range(torch.cuda.device_count()): properties = torch.cuda.get_device_properties(index) candidates.append(f"cuda:{index}={properties.name}/{properties.multi_processor_count}SM") - if properties.major == 10 and (mode != "formal" or (properties.name == "NVIDIA B200" and properties.multi_processor_count == 148)): + if (properties.major, properties.minor) == (10, 0) and ( + mode != "formal" or (properties.name == "NVIDIA B200" and properties.multi_processor_count == 148) + ): return torch.device(f"cuda:{index}"), properties - requirement = "a full 148-SM NVIDIA B200" if mode == "formal" else "an SM100-family GPU" + requirement = "a full 148-SM NVIDIA B200" if mode == "formal" else "an SM100 GPU" raise RuntimeError(f"{mode} mode requires {requirement}; visible: {', '.join(candidates)}") @@ -168,26 +176,33 @@ def _compare(current, previous): if current_fp != previous_fp: raise ValueError(f"comparison fingerprint mismatch: current={current_fp}, previous={previous_fp}") arms = {} - for bits in ("0", "1"): + for bits, _, _ in VARIANTS: new = float(current["summary"][bits]["p50_ms"]) old = float(previous["summary"][bits]["p50_ms"]) if not all(math.isfinite(value) and value > 0 for value in (new, old)): raise ValueError("comparison p50 values must be finite and positive") - arms[bits] = {"previous_p50_ms": old, "current_p50_ms": new, "change_percent": (new / old - 1.0) * 100.0} + arms[bits] = { + "previous_p50_ms": old, + "current_p50_ms": new, + "change_ms": new - old, + "change_percent": (new / old - 1.0) * 100.0, + } return { "paired_across_runs": False, "arms": arms, - "previous_within_run_ratio": previous["comparison_within_run"]["paired_ratio_p50"], - "current_within_run_ratio": current["comparison_within_run"]["paired_ratio_p50"], + "previous_within_run_ratio": previous["comparisons"]["all_vs_baseline"]["paired_ratio_p50"], + "current_within_run_ratio": current["comparisons"]["all_vs_baseline"]["paired_ratio_p50"], } def _render_markdown(metadata, raw_name, raw_hash): config = metadata["config"] - paired = metadata["comparison_within_run"] + paired = metadata["comparisons"]["all_vs_baseline"] + effects = metadata["main_effects"] + shapley = metadata["shapley"] smoke = config["mode"] == "smoke" lines = [ - "# Qwen-Image BF16 joint-attention benchmark", + "# Qwen-Image BF16 attention x GELU-MLP benchmark", "", f"Generated: `{metadata['completed_utc']}` ", f"Mode: `{config['mode']}` ", @@ -208,19 +223,79 @@ def _render_markdown(metadata, raw_name, raw_hash): lines += [ "## Result", "", - f"Direct FE/cuDNN is `{ratio:.5f}x` the paired forced-PyTorch-Flash elapsed time " + f"The both-cuDNN-treatment `11` arm is `{ratio:.5f}x` the paired `00` baseline elapsed time " f"({(1 - ratio) * 100:.2f}% lower, `{1 / ratio:.3f}x` speedup; {paired['wins']}/{paired['batches']} wins).", "", ] lines += [ - "| transformer forward (SDPA treatment) | p10 | p50 | p90 | paired ratio vs PyTorch Flash |", - "|---|---:|---:|---:|---:|", + "| bits (M/A) | GELU MLP | joint attention | p10 | p50 | p90 | paired ratio vs 00 |", + "|---|---|---|---:|---:|---:|---:|", ] - for bits, label in VARIANTS: + for bits, mlp_backend, attention_backend in VARIANTS: value = metadata["summary"][bits] - ratio = 1.0 if bits == "0" else value["paired_ratio_p50"] - lines.append(f"| {label} | {value['p10_ms']:.3f} ms | {value['p50_ms']:.3f} ms | {value['p90_ms']:.3f} ms | {ratio:.5f} |") + lines.append( + f"| `{bits}` | {'cuDNN' if mlp_backend == 'cudnn' else 'Torch'} | " + f"{'cuDNN backend' if attention_backend == 'cudnn' else 'forced PyTorch Flash'} | " + f"{value['p10_ms']:.3f} ms | {value['p50_ms']:.3f} ms | {value['p90_ms']:.3f} ms | " + f"{value['paired_ratio_p50']:.5f} |" + ) + if not smoke: + contrast_labels = ( + ("attention_with_torch_mlp", "attention: 01 / 00 (Torch MLP)"), + ("attention_with_cudnn_mlp", "attention: 11 / 10 (cuDNN MLP)"), + ("mlp_with_flash_attention", "MLP: 10 / 00 (Flash attention)"), + ("mlp_with_cudnn_attention", "MLP: 11 / 01 (cuDNN attention)"), + ) + lines += [ + "", + "## Paired contrasts", + "", + "| contrast | paired ratio (p50) | speedup | paired delta (p50) | wins |", + "|---|---:|---:|---:|---:|", + ] + for key, label in contrast_labels: + contrast = metadata["comparisons"][key] + contrast_ratio = contrast["paired_ratio_p50"] + lines.append( + f"| {label} | {contrast_ratio:.5f} | {1 / contrast_ratio:.3f}x | " + f"{contrast['paired_delta_p50_ms']:+.3f} ms | {contrast['wins']}/{contrast['batches']} |" + ) + lines += [ + "", + "## Factorial attribution", + "", + "Conditional ratios average an axis over both paired contexts in log space. Shapley savings allocate " + "the paired `00` minus `11` elapsed-time saving, including their interaction; they are not module-time shares.", + "", + "| axis | conditional ratio (p50) | conditional speedup | conditional delta (p50) | Shapley saving (p50) |", + "|---|---:|---:|---:|---:|", + ] + for axis in AXIS_MASKS: + axis_ratio = effects[axis]["conditional_geomean_ratio"]["p50"] + axis_delta = effects[axis]["conditional_delta_ms"]["p50"] + axis_shapley = shapley["saving_ms"][axis]["p50"] + lines.append(f"| {axis} | {axis_ratio:.5f} | {1 / axis_ratio:.3f}x | {axis_delta:+.3f} ms | {axis_shapley:.3f} ms |") + comparison = metadata.get("comparison_across_runs") + if comparison is not None: + lines += [ + "", + "## Cross-run comparison", + "", + "**Not paired across runs.** Each arm compares independent p50 estimates; the headline values are each paired only within their own run.", + "", + f"The within-run paired `11/00` ratio changed from `{comparison['previous_within_run_ratio']:.5f}` to " + f"`{comparison['current_within_run_ratio']:.5f}`.", + "", + "| bits | previous p50 | current p50 | non-paired change | change |", + "|---|---:|---:|---:|---:|", + ] + for bits, arm in sorted(comparison["arms"].items()): + lines.append( + f"| `{bits}` | {arm['previous_p50_ms']:.3f} ms | {arm['current_p50_ms']:.3f} ms | " + f"{arm['change_ms']:+.3f} ms | {arm['change_percent']:+.2f}% |" + ) shape = config["shape"] + correctness = metadata["correctness"]["model_output_rel_l2"] lines += [ "", "## Scope and gates", @@ -229,10 +304,14 @@ def _render_markdown(metadata, raw_name, raw_hash): f"H={shape['hidden']}, heads={shape['heads']}x{shape['head_dim']}, FFN={shape['ffn']}, repeated layers={shape['layers']}/60.", f"- Recipe: `{config['numerical_recipe']['id']}` ({config['numerical_recipe']['parameter_dtype']} parameters and activations).", f"- Workload: `{config['workload']}`. One conditional transformer forward; no text encoder, VAE, scheduler, checkpoint weights, or full denoising loop.", + "- Treatment scope: `00` turns cuDNN off only for the two measured axes (stock Torch GELU FFN plus forced PyTorch Flash SDPA); " + "it does not claim that unrelated Torch operators never use cuDNN.", f"- PyTorch route: natural `{metadata['route']['torch_probe']['natural_choice_name']}`, timed treatment forced " - f"`{metadata['route']['torch_probe']['forced_choice_name']}`; calls Flash/cuDNN: " - f"`{metadata['route']['calls']['torch_flash']}/{metadata['route']['calls']['cudnn']}`.", - f"- Full-model output relative L2: `{metadata['correctness']['model_output_rel_l2']:.6g}`; " + f"`{metadata['route']['torch_probe']['forced_choice_name']}`; attention calls Flash/cuDNN: " + f"`{metadata['route']['attention_calls']['torch_flash']}/{metadata['route']['attention_calls']['cudnn']}`.", + f"- GELU-MLP route calls Torch/cuDNN: `{metadata['route']['mlp_calls']['torch']}/{metadata['route']['mlp_calls']['cudnn']}`; " + "the proxy has one image-stream and one text-stream FFN per repeated block.", + f"- Full-model output relative L2 by arm vs `00`: `{json.dumps(correctness, sort_keys=True)}`; " f"right-padded B=2 mask adapter relative L2: `{metadata['correctness']['padding_adapter']['rel_l2']:.6g}`.", "- Joint mask semantics: every query sees valid text and every image token; only padded text key columns are rejected.", "", @@ -240,6 +319,15 @@ def _render_markdown(metadata, raw_name, raw_hash): "", f"- Model: [`{config['model_anchor']['id']}@{config['model_anchor']['revision']}`]({config['model_anchor']['config_url']})", f"- Implementation: [`diffusers@{config['diffusers_anchor']['commit']}`]({config['diffusers_anchor']['url']})", + "", + "## Provenance", + "", + "| source | path | sha256 |", + "|---|---|---|", + ] + for name, source in sorted(metadata["provenance"]["sources"].items()): + lines.append(f"| {name} | `{source['path']}` | `{source['sha256']}` |") + lines += [ "", "This is a random-weight, depth-reduced transformer-shape proxy, not image-quality evidence or complete Qwen-Image pipeline throughput.", "", @@ -256,6 +344,9 @@ def main(): if os.environ.get("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "0").lower() in ("1", "true", "yes", "on"): raise RuntimeError("disable global FROST engines: the FE arm is defined as the cuDNN backend graph") torch, _, cudnn, sdpamod, diffusers, qwen_module = model_api.load_runtime() + gelu_mlp_module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + diffusers_attention_module = importlib.import_module("diffusers.models.attention") + diffusers_activations_module = importlib.import_module("diffusers.models.activations") loaded_diffusers_source = _source_record(qwen_module.__file__) expected_diffusers_sha = model_api.DIFFUSERS_ANCHOR["source_sha256"] if loaded_diffusers_source["sha256"] != expected_diffusers_sha: @@ -263,6 +354,14 @@ def main(): "loaded Diffusers Qwen-Image source does not match the pinned implementation: " f"got {loaded_diffusers_source['sha256']}, expected {expected_diffusers_sha}" ) + loaded_supporting_sources = { + "attention": _source_record(diffusers_attention_module.__file__), + "activations": _source_record(diffusers_activations_module.__file__), + } + for name, source in loaded_supporting_sources.items(): + expected = model_api.DIFFUSERS_ANCHOR["supporting_sources"][name]["source_sha256"] + if source["sha256"] != expected: + raise RuntimeError(f"loaded Diffusers {name} source does not match the pinned implementation: " f"got {source['sha256']}, expected {expected}") if not torch.cuda.is_available(): raise RuntimeError("CUDA is required") device, properties = _pick_device(torch, args.mode) @@ -272,48 +371,65 @@ def main(): padding_check = _focused_padding_check(torch, qwen_module, model_api, device) model = model_api.build_model(torch, qwen_module, device, layers=shape["layers"]) inputs = model_api.make_inputs(torch, device, shape) - calls = {"torch_flash": 0, "cudnn": 0} + attention_calls = {"torch_flash": 0, "cudnn": 0} + mlp_calls = {"torch": 0, "cudnn": 0} torch_probe = {} - select, restore, calls, torch_probe = model_api.install_joint_attention_dispatch( - qwen_module, text_tokens=shape["text_tokens"], counters=calls, torch_probe=torch_probe + select_attention, restore_attention, attention_calls, torch_probe = model_api.install_joint_attention_dispatch( + qwen_module, text_tokens=shape["text_tokens"], counters=attention_calls, torch_probe=torch_probe ) + select_mlp, restore_mlp, mlp_calls = model_api.install_gelu_mlp_dispatch(model, counters=mlp_calls) - def step(backend): - select(backend) + def step(mlp_backend, attention_backend): + select_mlp(mlp_backend) + select_attention(attention_backend) with torch.inference_mode(): return model_api.forward(model, inputs) try: - for bits, backend in VARIANTS: - before = dict(calls) + for bits, mlp_backend, attention_backend in VARIANTS: + attention_before = dict(attention_calls) + mlp_before = dict(mlp_calls) for _ in range(protocol["warmup"]): - step(backend) + step(mlp_backend, attention_backend) torch.cuda.synchronize(device) - delta = calls[backend] - before[backend] - expected = protocol["warmup"] * shape["layers"] - other = "cudnn" if backend == "torch_flash" else "torch_flash" - if delta != expected or calls[other] != before[other]: - raise RuntimeError(f"{backend} warm route mismatch: delta={delta}, expected={expected}, calls={calls}, before={before}") - - expected_output = step("torch_flash") - actual_output = step("cudnn") - torch.cuda.synchronize(device) - if not bool(torch.isfinite(expected_output).all()) or not bool(torch.isfinite(actual_output).all()): + attention_delta = {name: attention_calls[name] - attention_before[name] for name in attention_calls} + mlp_delta = {name: mlp_calls[name] - mlp_before[name] for name in mlp_calls} + expected_attention = protocol["warmup"] * shape["layers"] + expected_mlp = 2 * protocol["warmup"] * shape["layers"] + if attention_delta != { + "torch_reference": 0, + "torch_flash": expected_attention if attention_backend == "torch_flash" else 0, + "cudnn": expected_attention if attention_backend == "cudnn" else 0, + }: + raise RuntimeError(f"{bits} attention warm route mismatch: delta={attention_delta}") + if mlp_delta != {"torch": expected_mlp if mlp_backend == "torch" else 0, "cudnn": expected_mlp if mlp_backend == "cudnn" else 0}: + raise RuntimeError(f"{bits} MLP warm route mismatch: delta={mlp_delta}") + + expected_output = step("torch", "torch_flash") + if not bool(torch.isfinite(expected_output).all()): raise RuntimeError("non-finite model output") - model_rel_l2 = _rel_l2(actual_output, expected_output) - if not math.isfinite(model_rel_l2) or model_rel_l2 > 0.02: - raise RuntimeError(f"model output mismatch: rel_l2={model_rel_l2}") - - orders = williams_orders(2) - raw = {bits: [] for bits, _ in VARIANTS} - batches = {bits: [] for bits, _ in VARIANTS} + model_rel_l2 = {"00": 0.0} + for bits, mlp_backend, attention_backend in VARIANTS[1:]: + actual_output = step(mlp_backend, attention_backend) + torch.cuda.synchronize(device) + if not bool(torch.isfinite(actual_output).all()): + raise RuntimeError(f"non-finite model output in arm {bits}") + rel = _rel_l2(actual_output, expected_output) + if not math.isfinite(rel) or rel > 0.02: + raise RuntimeError(f"model output mismatch in arm {bits}: rel_l2={rel}") + model_rel_l2[bits] = rel + + orders = williams_orders(4) + raw = {bits: [] for bits, _, _ in VARIANTS} + batches = {bits: [] for bits, _, _ in VARIANTS} timing_started = time.time() for batch in range(protocol["rounds"]): for variant_index in orders[batch % len(orders)]: - bits, backend = VARIANTS[variant_index] + bits, mlp_backend, attention_backend = VARIANTS[variant_index] samples = [] for _ in range(protocol["repeats"]): - select(backend) + select_mlp(mlp_backend) + select_attention(attention_backend) start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) start.record() with torch.inference_mode(): @@ -328,9 +444,9 @@ def step(backend): batches[bits].append(statistics.median(samples)) print(f"BATCH {batch + 1}/{protocol['rounds']} elapsed_s={time.time() - timing_started:.1f}", flush=True) - paired = paired_stats(batches["1"], batches["0"]) + paired = paired_stats(batches["11"], batches["00"]) summary = {} - for bits, _ in VARIANTS: + for bits, _, _ in VARIANTS: values = batches[bits] result = { "p10_ms": percentile(values, 0.1), @@ -339,28 +455,48 @@ def step(backend): "mean_ms": statistics.mean(values), "batches": len(values), } - arm_paired = paired_stats(values, batches["0"]) + arm_paired = paired_stats(values, batches["00"]) result.update(arm_paired) summary[bits] = result - calls_per_backend = protocol["warmup"] + 1 + protocol["rounds"] * protocol["repeats"] - expected_calls = calls_per_backend * shape["layers"] - if calls != {"torch_reference": 0, "torch_flash": expected_calls, "cudnn": expected_calls}: - raise RuntimeError(f"final attention route mismatch: calls={calls}, expected_each={expected_calls}") + integer_batches = {int(bits, 2): values for bits, values in batches.items()} + main_effects = factorial_main_effects(integer_batches, axis_masks=AXIS_MASKS) + shapley = shapley_savings(integer_batches, axis_masks=AXIS_MASKS) + comparisons = { + "all_vs_baseline": paired, + "attention_with_torch_mlp": paired_stats(batches["01"], batches["00"]), + "attention_with_cudnn_mlp": paired_stats(batches["11"], batches["10"]), + "mlp_with_flash_attention": paired_stats(batches["10"], batches["00"]), + "mlp_with_cudnn_attention": paired_stats(batches["11"], batches["01"]), + } + + common_calls = protocol["warmup"] + 1 + protocol["rounds"] * protocol["repeats"] + expected_attention_calls = 2 * common_calls * shape["layers"] + expected_mlp_calls = 4 * common_calls * shape["layers"] + expected_attention_routes = {"torch_reference": 0, "torch_flash": expected_attention_calls, "cudnn": expected_attention_calls} + expected_mlp_routes = {"torch": expected_mlp_calls, "cudnn": expected_mlp_calls} + if attention_calls != expected_attention_routes: + raise RuntimeError(f"final attention route mismatch: calls={attention_calls}, expected={expected_attention_routes}") + if mlp_calls != expected_mlp_routes: + raise RuntimeError(f"final MLP route mismatch: calls={mlp_calls}, expected={expected_mlp_routes}") if torch_probe.get("forced_choice_name") != "FLASH_ATTENTION": raise RuntimeError(f"forced PyTorch FlashAttention treatment route changed: {torch_probe}") finally: - restore() + restore_mlp() + restore_attention() sources = { "runner": _source_record(Path(__file__)), "model_adapter": _source_record(MODEL_PATH), "statistics": _source_record(FACTORIAL_PATH), "diffusers_qwen_image": loaded_diffusers_source, + "diffusers_attention": loaded_supporting_sources["attention"], + "diffusers_activations": loaded_supporting_sources["activations"], "cudnn_sdpa": _source_record(sdpamod.__file__), + "cudnn_gelu_mlp": _source_record(gelu_mlp_module.__file__), } config = { - "schema_version": 1, + "schema_version": 2, "mode": args.mode, "timing_role": "validation_only" if args.mode == "smoke" else "formal_performance", "performance_claim_eligible": args.mode == "formal", @@ -383,7 +519,8 @@ def step(backend): "numerical_recipe": dict(model_api.NUMERICAL_RECIPE), "model_anchor": dict(model_api.OFFICIAL_MODEL), "diffusers_anchor": dict(model_api.DIFFUSERS_ANCHOR), - "variants": [{"bits": bits, "backend": backend} for bits, backend in VARIANTS], + "axis_masks": dict(AXIS_MASKS), + "variants": [{"bits": bits, "mlp": mlp_backend, "attention": attention_backend} for bits, mlp_backend, attention_backend in VARIANTS], "williams_orders": orders, } comparable = { @@ -409,6 +546,7 @@ def step(backend): "numerical_recipe", "model_anchor", "diffusers_anchor", + "axis_masks", "variants", "williams_orders", ) @@ -417,17 +555,25 @@ def step(backend): config["comparability_fingerprint"] = {"inputs": comparable, "sha256": config_fingerprint(comparable)} config["build_fingerprint"] = {"inputs": build, "sha256": config_fingerprint(build)} metadata = { - "schema_version": 1, + "schema_version": 2, "started_utc": started_utc, "completed_utc": _utc_now(), "arguments": {name: str(value) if isinstance(value, Path) else value for name, value in vars(args).items()}, "config": config, "correctness": {"model_output_rel_l2": model_rel_l2, "padding_adapter": padding_check}, "summary": summary, - "comparison_within_run": paired, + "comparisons": comparisons, + "main_effects": main_effects, + "shapley": shapley, "batch_medians_ms": batches, "raw_ms": raw, - "route": {"calls": calls, "expected_calls_each": expected_calls, "torch_probe": torch_probe}, + "route": { + "attention_calls": attention_calls, + "expected_attention_calls_each": expected_attention_calls, + "mlp_calls": mlp_calls, + "expected_mlp_calls_each": expected_mlp_calls, + "torch_probe": torch_probe, + }, "provenance": {"git": build["git"], "sources": sources}, } if args.compare is not None: @@ -446,7 +592,17 @@ def step(backend): raw_path.write_text(json.dumps(metadata, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8") raw_hash = _sha256(raw_path) report_path.write_text(_render_markdown(metadata, raw_path.name, raw_hash), encoding="utf-8") - print("RESULT " + json.dumps({"torch_flash_p50_ms": summary["0"]["p50_ms"], "cudnn_p50_ms": summary["1"]["p50_ms"], **paired}, sort_keys=True)) + print( + "RESULT " + + json.dumps( + { + "baseline_00_p50_ms": summary["00"]["p50_ms"], + "cudnn_treatments_11_p50_ms": summary["11"]["p50_ms"], + **paired, + }, + sort_keys=True, + ) + ) print(f"RAW_JSON {raw_path} sha256={raw_hash}") print(f"MARKDOWN {report_path} sha256={_sha256(report_path)}") diff --git a/benchmark/e2e/Qwen-Image/run_model.py b/benchmark/e2e/Qwen-Image/run_model.py index 751395db5..dbf31af01 100644 --- a/benchmark/e2e/Qwen-Image/run_model.py +++ b/benchmark/e2e/Qwen-Image/run_model.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Qwen-Image transformer-shape proxy and switchable BF16 joint attention. +"""Qwen-Image transformer-shape proxy with switchable BF16 attention and MLP. The proxy instantiates Diffusers' real ``QwenImageTransformer2DModel`` class with the published hidden/head/FFN dimensions. It reduces only the repeated @@ -44,6 +44,16 @@ "2f7e0154a9db246e95c9ede43edba7db5b130805/" "src/diffusers/models/transformers/transformer_qwenimage.py" ), + "supporting_sources": { + "attention": { + "path": "src/diffusers/models/attention.py", + "source_sha256": "3c61df6cc4832149eb654c1e82220f4a6b91daca13741c957c4e0faff7810adf", + }, + "activations": { + "path": "src/diffusers/models/activations.py", + "source_sha256": "ab1767e8e44e7d4bf1cb18299ba33654329e2711fa983b1423fc4fe12de2c3ab", + }, + }, } NUMERICAL_RECIPE = { "id": "qwen-image-conservative-bf16-v1", @@ -270,6 +280,151 @@ def restore(): return select, restore, counters, torch_probe +def install_gelu_mlp_dispatch(model, *, counters=None): + """Install strict Torch/cuDNN dispatchers for Qwen-Image's two GELU FFNs. + + The treatment is deliberately benchmark-local: it accepts only the pinned + Diffusers ``FeedForward`` structure used by this proxy, + ``Linear+bias -> GELU(approximate=\"tanh\") -> Dropout(0) -> Linear+bias``. + Any hook, parametrization, training state, or structural drift fails closed + instead of silently changing model semantics. Returns + ``(select, restore, counters)`` where ``select`` accepts ``\"torch\"`` or + ``\"cudnn\"``. + """ + import types + + import torch + + from cudnn.gemm.ops import gelu_mlp + + if counters is None: + counters = {} + for name in ("torch", "cudnn"): + counters.setdefault(name, 0) + + if model.__class__.__module__ != "diffusers.models.transformers.transformer_qwenimage" or model.__class__.__name__ != "QwenImageTransformer2DModel": + raise TypeError(f"expected pinned Diffusers QwenImageTransformer2DModel, got {type(model)!r}") + if model.training or any(parameter.requires_grad for parameter in model.parameters()): + raise NotImplementedError("the Qwen-Image GELU-MLP treatment is inference-only") + if getattr(model, "_compiled_call_impl", None) is not None: + raise NotImplementedError("torch.compile is outside the Qwen-Image GELU-MLP treatment") + if getattr(model, "peft_config", None): + raise NotImplementedError("LoRA/PEFT adapters are outside the Qwen-Image GELU-MLP treatment") + blocks = getattr(model, "transformer_blocks", None) + if not isinstance(blocks, torch.nn.ModuleList) or not blocks: + raise TypeError("expected a non-empty Qwen-Image transformer_blocks ModuleList") + + hook_fields = ("_forward_hooks", "_forward_pre_hooks", "_backward_hooks", "_backward_pre_hooks") + for module_name, child in model.named_modules(): + qualified = module_name or "model" + if getattr(child, "_compiled_call_impl", None) is not None: + raise NotImplementedError(f"torch.compile on {qualified} is outside the Qwen-Image GELU-MLP treatment") + if any(getattr(child, field, None) for field in hook_fields): + raise NotImplementedError(f"hooks on {qualified} are outside the Qwen-Image GELU-MLP treatment") + entries = [] + for block_index, block in enumerate(blocks): + if block.__class__.__module__ != "diffusers.models.transformers.transformer_qwenimage" or block.__class__.__name__ != "QwenImageTransformerBlock": + raise TypeError(f"transformer_blocks[{block_index}] is not the pinned QwenImageTransformerBlock: {type(block)!r}") + for stream in ("img", "txt"): + name = f"transformer_blocks[{block_index}].{stream}_mlp" + module = getattr(block, f"{stream}_mlp", None) + if module is None or module.__class__.__module__ != "diffusers.models.attention" or module.__class__.__name__ != "FeedForward": + raise TypeError(f"{name} is not the pinned Diffusers FeedForward: {type(module)!r}") + if module.training: + raise NotImplementedError(f"{name} must be in eval mode") + if "forward" in module.__dict__: + raise NotImplementedError(f"{name} already has an instance-level forward override") + net = getattr(module, "net", None) + if not isinstance(net, torch.nn.ModuleList) or len(net) != 3: + raise TypeError(f"{name}.net must contain exactly activation, dropout, and output projection") + activation, dropout, output = net + first = getattr(activation, "proj", None) + if ( + activation.__class__.__module__ != "diffusers.models.activations" + or activation.__class__.__name__ != "GELU" + or getattr(activation, "approximate", None) != "tanh" + or type(first) is not torch.nn.Linear + or type(dropout) is not torch.nn.Dropout + or type(output) is not torch.nn.Linear + ): + raise TypeError(f"{name} is not Linear -> GELU(tanh) -> Dropout -> Linear") + if dropout.p != 0.0 or dropout.inplace: + raise NotImplementedError(f"{name} requires non-inplace Dropout(0), got p={dropout.p}, inplace={dropout.inplace}") + if first.bias is None or output.bias is None: + raise NotImplementedError(f"{name} requires bias on both linear projections") + expected_first = (PUBLISHED_SHAPE["ffn"], PUBLISHED_SHAPE["hidden"]) + expected_second = (PUBLISHED_SHAPE["hidden"], PUBLISHED_SHAPE["ffn"]) + expected_first_bias = (PUBLISHED_SHAPE["ffn"],) + expected_second_bias = (PUBLISHED_SHAPE["hidden"],) + if ( + tuple(first.weight.shape) != expected_first + or tuple(first.bias.shape) != expected_first_bias + or tuple(output.weight.shape) != expected_second + or tuple(output.bias.shape) != expected_second_bias + ): + raise ValueError( + f"{name} projection shapes changed: got weights {tuple(first.weight.shape)}/{tuple(output.weight.shape)} " + f"and biases {tuple(first.bias.shape)}/{tuple(output.bias.shape)}; expected " + f"{expected_first}/{expected_second} and {expected_first_bias}/{expected_second_bias}" + ) + tensors = (first.weight, first.bias, output.weight, output.bias) + if any(tensor.dtype != torch.bfloat16 or tensor.device.type != "cuda" or not tensor.is_contiguous() for tensor in tensors): + raise NotImplementedError(f"{name} requires contiguous bf16 CUDA weights and biases") + if len({tensor.device for tensor in tensors}) != 1: + raise ValueError(f"{name} weights and biases must share one CUDA device") + for child_name, child in module.named_modules(): + qualified = name if not child_name else f"{name}.{child_name}" + if child.training: + raise NotImplementedError(f"{qualified} must be in eval mode") + if any(getattr(child, field, None) for field in hook_fields): + raise NotImplementedError(f"hooks on {qualified} are outside the benchmark treatment") + if torch.nn.utils.parametrize.is_parametrized(child): + raise NotImplementedError(f"parametrizations on {qualified} are outside the benchmark treatment") + entries.append((name, module, first, output, module.forward)) + + installed = {} + for name, module, first, output, original in entries: + + def torch_forward(self, hidden_states, *args, _name=name, _original=original, **kwargs): + if args or kwargs: + raise NotImplementedError(f"{_name} benchmark dispatcher accepts only the hidden_states argument") + counters["torch"] += 1 + return _original(hidden_states) + + def cudnn_forward(self, hidden_states, *args, _name=name, _first=first, _output=output, **kwargs): + if args or kwargs: + raise NotImplementedError(f"{_name} benchmark dispatcher accepts only the hidden_states argument") + if hidden_states.ndim != 3 or hidden_states.shape[-1] != PUBLISHED_SHAPE["hidden"]: + raise ValueError(f"{_name} expected [B,S,{PUBLISHED_SHAPE['hidden']}], got {tuple(hidden_states.shape)}") + counters["cudnn"] += 1 + return gelu_mlp(hidden_states, _first.weight, _first.bias, _output.weight, _output.bias) + + installed[module] = { + "original": original, + "torch": types.MethodType(torch_forward, module), + "cudnn": types.MethodType(cudnn_forward, module), + } + + def select(name): + if name not in ("torch", "cudnn"): + raise ValueError(f"unknown GELU-MLP backend {name!r}") + for module, forwards in installed.items(): + if module.forward not in forwards.values(): + raise RuntimeError("a Qwen-Image FeedForward was modified after installing the benchmark dispatcher") + for module, forwards in installed.items(): + module.forward = forwards[name] + + def restore(): + for module, forwards in installed.items(): + if module.forward in forwards.values(): + # ``forward`` originally came from the class. Removing our + # instance override restores that exact lookup state and lets + # a later benchmark install validate the same model again. + module.__dict__.pop("forward", None) + + return select, restore, counters + + def build_model(torch, qwen_module, device, *, layers): torch.manual_seed(0) with torch.cuda.device(device): @@ -318,6 +473,7 @@ def _main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--mode", choices=MODE_DEFAULTS, default="smoke") parser.add_argument("--backend", choices=("torch_reference", "torch_flash", "cudnn"), default="cudnn") + parser.add_argument("--mlp-backend", choices=("torch", "cudnn"), default="torch") parser.add_argument("--layers", type=int) parser.add_argument("--image-tokens", type=int) parser.add_argument("--text-tokens", type=int) @@ -332,15 +488,17 @@ def _main(): raise RuntimeError("CUDA is required") device = torch.device("cuda") properties = torch.cuda.get_device_properties(device) - if properties.major != 10: + if (properties.major, properties.minor) != (10, 0): raise RuntimeError(f"Qwen-Image proxy requires SM100, got {properties.name}") if cudnn.backend_version() < 92100: raise RuntimeError(f"joint d128 SDPA requires a current cuDNN backend, got {cudnn.backend_version()}") model = build_model(torch, qwen_module, device, layers=shape["layers"]) inputs = make_inputs(torch, device, shape) - select, restore, counters, probe = install_joint_attention_dispatch(qwen_module, text_tokens=shape["text_tokens"]) + select_attention, restore_attention, attention_counters, probe = install_joint_attention_dispatch(qwen_module, text_tokens=shape["text_tokens"]) + select_mlp, restore_mlp, mlp_counters = install_gelu_mlp_dispatch(model) try: - select(args.backend) + select_attention(args.backend) + select_mlp(args.mlp_backend) samples = [] with torch.inference_mode(): for _ in range(2): @@ -359,16 +517,19 @@ def _main(): json.dumps( { "backend": args.backend, + "mlp_backend": args.mlp_backend, "p50_ms": statistics.median(samples), "shape": shape, - "calls": counters, + "attention_calls": attention_counters, + "mlp_calls": mlp_counters, "torch_probe": probe, }, sort_keys=True, ) ) finally: - restore() + restore_mlp() + restore_attention() if __name__ == "__main__": diff --git a/benchmark/e2e/tests/test_qwen_image_spec.py b/benchmark/e2e/tests/test_qwen_image_spec.py index bfdcc69cd..cca712cef 100644 --- a/benchmark/e2e/tests/test_qwen_image_spec.py +++ b/benchmark/e2e/tests/test_qwen_image_spec.py @@ -4,6 +4,7 @@ import importlib.util from pathlib import Path import sys +from types import SimpleNamespace import unittest E2E_DIR = Path(__file__).resolve().parents[1] @@ -15,6 +16,14 @@ sys.modules[SPEC.name] = MODEL SPEC.loader.exec_module(MODEL) +RUNNER_PATH = E2E_DIR / "Qwen-Image" / "run_bf16.py" +RUNNER_SPEC = importlib.util.spec_from_file_location("qwen_image_bf16_cpu_test", RUNNER_PATH) +if RUNNER_SPEC is None or RUNNER_SPEC.loader is None: + raise RuntimeError(f"cannot load {RUNNER_PATH}") +RUNNER = importlib.util.module_from_spec(RUNNER_SPEC) +sys.modules[RUNNER_SPEC.name] = RUNNER +RUNNER_SPEC.loader.exec_module(RUNNER) + class QwenImageSpecTest(unittest.TestCase): def test_formal_shape_preserves_published_kernel_dimensions(self): @@ -39,9 +48,98 @@ def test_bf16_leaf_has_immutable_upstream_anchors(self): self.assertEqual(len(MODEL.OFFICIAL_MODEL["revision"]), 40) self.assertEqual(len(MODEL.DIFFUSERS_ANCHOR["commit"]), 40) self.assertEqual(len(MODEL.DIFFUSERS_ANCHOR["source_sha256"]), 64) + self.assertEqual(set(MODEL.DIFFUSERS_ANCHOR["supporting_sources"]), {"attention", "activations"}) + for source in MODEL.DIFFUSERS_ANCHOR["supporting_sources"].values(): + self.assertEqual(len(source["source_sha256"]), 64) self.assertEqual(MODEL.NUMERICAL_RECIPE["id"], "qwen-image-conservative-bf16-v1") self.assertEqual(MODEL.NUMERICAL_RECIPE["scope"], "inference_transformer_forward") + def test_bf16_leaf_is_complete_mlp_attention_factorial(self): + self.assertEqual( + RUNNER.VARIANTS, + ( + ("00", "torch", "torch_flash"), + ("01", "torch", "cudnn"), + ("10", "cudnn", "torch_flash"), + ("11", "cudnn", "cudnn"), + ), + ) + self.assertEqual(RUNNER.AXIS_MASKS, {"mlp": 2, "attn": 1}) + self.assertEqual(len(RUNNER.williams_orders(4)), 4) + + def test_bf16_protocol_requires_complete_four_arm_cycles(self): + args = SimpleNamespace(mode="formal", warmup=None, rounds=None, repeats=None) + self.assertEqual(RUNNER._resolve_protocol(args), {"warmup": 3, "rounds": 40, "repeats": 3}) + args.rounds = 6 + with self.assertRaisesRegex(ValueError, "multiple of 4"): + RUNNER._resolve_protocol(args) + + def test_bf16_report_names_off_on_scope_and_both_axes(self): + batches = { + "00": [10.0, 10.2], + "01": [8.0, 8.1], + "10": [9.8, 10.0], + "11": [7.8, 7.9], + } + summary = {} + for bits, values in batches.items(): + paired = RUNNER.paired_stats(values, batches["00"]) + summary[bits] = { + "p10_ms": min(values), + "p50_ms": sum(values) / len(values), + "p90_ms": max(values), + **paired, + } + integer_batches = {int(bits, 2): values for bits, values in batches.items()} + comparisons = { + "all_vs_baseline": RUNNER.paired_stats(batches["11"], batches["00"]), + "attention_with_torch_mlp": RUNNER.paired_stats(batches["01"], batches["00"]), + "attention_with_cudnn_mlp": RUNNER.paired_stats(batches["11"], batches["10"]), + "mlp_with_flash_attention": RUNNER.paired_stats(batches["10"], batches["00"]), + "mlp_with_cudnn_attention": RUNNER.paired_stats(batches["11"], batches["01"]), + } + metadata = { + "completed_utc": "2026-08-21T00:00:00Z", + "config": { + "mode": "formal", + "comparability_fingerprint": {"sha256": "comparable"}, + "build_fingerprint": {"sha256": "build"}, + "shape": MODEL.resolve_shape("formal"), + "numerical_recipe": dict(MODEL.NUMERICAL_RECIPE), + "workload": "single_conditional_transformer_forward_no_checkpoint", + "model_anchor": dict(MODEL.OFFICIAL_MODEL), + "diffusers_anchor": dict(MODEL.DIFFUSERS_ANCHOR), + }, + "summary": summary, + "comparisons": comparisons, + "main_effects": RUNNER.factorial_main_effects(integer_batches, axis_masks=RUNNER.AXIS_MASKS), + "shapley": RUNNER.shapley_savings(integer_batches, axis_masks=RUNNER.AXIS_MASKS), + "correctness": { + "model_output_rel_l2": {bits: 0.0 for bits in batches}, + "padding_adapter": {"rel_l2": 0.0}, + }, + "route": { + "torch_probe": {"natural_choice_name": "CUDNN_ATTENTION", "forced_choice_name": "FLASH_ATTENTION"}, + "attention_calls": {"torch_flash": 8, "cudnn": 8}, + "mlp_calls": {"torch": 16, "cudnn": 16}, + }, + "provenance": { + "sources": { + "cudnn_gelu_mlp": {"path": "_gelu_mlp.py", "sha256": "face"}, + "diffusers_attention": {"path": "attention.py", "sha256": "feed"}, + "diffusers_activations": {"path": "activations.py", "sha256": "beef"}, + } + }, + } + report = RUNNER._render_markdown(metadata, "result.json", "cafe") + self.assertIn("bits (M/A)", report) + self.assertIn("both-cuDNN-treatment `11`", report) + self.assertIn("turns cuDNN off only for the two measured axes", report) + self.assertIn("natural `CUDNN_ATTENTION`", report) + self.assertIn("cudnn_gelu_mlp", report) + self.assertIn("diffusers_attention", report) + self.assertIn("diffusers_activations", report) + if __name__ == "__main__": unittest.main() From 15665c1bd877c9a098e7444a9f15c05403043b88 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 21 Aug 2026 05:21:39 -0700 Subject: [PATCH 3/4] benchmark: add a ModelOpt-anchored Qwen-Image NVFP4 path --- .gitignore | 1 + THIRD_PARTY_LICENSES.txt | 17 + benchmark/e2e/Qwen-Image/modelopt_nvfp4.py | 2017 +++++++++++++++++ benchmark/e2e/Qwen-Image/run_nvfp4.py | 681 ++++++ benchmark/e2e/README.md | 169 +- .../e2e/tests/test_qwen_image_nvfp4_spec.py | 497 ++++ pyproject.toml | 4 +- python/cudnn/gemm/ops/_nvfp4_quantize.py | 225 ++ .../gemm/ops/csrc/nvfp4_quantize_sm100.cu | 73 + .../ops/csrc/nvfp4_smooth_quantize_sm100.cuh | 605 +++++ test/python/gemm/test_nvfp4_quantize.py | 155 ++ 11 files changed, 4417 insertions(+), 27 deletions(-) create mode 100644 benchmark/e2e/Qwen-Image/modelopt_nvfp4.py create mode 100644 benchmark/e2e/Qwen-Image/run_nvfp4.py create mode 100644 benchmark/e2e/tests/test_qwen_image_nvfp4_spec.py create mode 100644 python/cudnn/gemm/ops/_nvfp4_quantize.py create mode 100644 python/cudnn/gemm/ops/csrc/nvfp4_quantize_sm100.cu create mode 100644 python/cudnn/gemm/ops/csrc/nvfp4_smooth_quantize_sm100.cuh create mode 100644 test/python/gemm/test_nvfp4_quantize.py diff --git a/.gitignore b/.gitignore index caafff936..739d0e071 100644 --- a/.gitignore +++ b/.gitignore @@ -54,6 +54,7 @@ cover/ # Local end-to-end benchmark artifacts qwen3.8-factorial-results/ qwen-image-bf16-results/ +qwen-image-nvfp4-results/ # Translations *.mo diff --git a/THIRD_PARTY_LICENSES.txt b/THIRD_PARTY_LICENSES.txt index f9aa22a8e..02d5a3238 100644 --- a/THIRD_PARTY_LICENSES.txt +++ b/THIRD_PARTY_LICENSES.txt @@ -264,3 +264,20 @@ FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +------------------------------------------------------------------------------- +10. FlashInfer + https://github.com/flashinfer-ai/flashinfer + +The benchmark-private SM100 NVFP4 activation quantizer under +python/cudnn/gemm/ops/csrc/ is derived from FlashInfer's +include/flashinfer/gemm/nvfp4_smooth_quantize_sm100.cuh at commit +f212ec8230486e3615502b8af75fe7022c60b2f3. The upstream source in turn +credits NVIDIA TensorRT-LLM's nvfp4SmoothQuantize kernel and quantization +helpers. The derived source retains its NVIDIA copyright and provenance. +Local changes are limited to provenance and formatting annotations. + +FlashInfer +Copyright 2025-2026 NVIDIA +Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) + +License: Apache License 2.0 (same text as this repository's LICENSE.txt) diff --git a/benchmark/e2e/Qwen-Image/modelopt_nvfp4.py b/benchmark/e2e/Qwen-Image/modelopt_nvfp4.py new file mode 100644 index 000000000..572f8c69e --- /dev/null +++ b/benchmark/e2e/Qwen-Image/modelopt_nvfp4.py @@ -0,0 +1,2017 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark-local ModelOpt 0.46.0 NVFP4 adapter for Qwen-Image. + +This module is deliberately not a general quantization framework. It recognizes +the exact fourteen ``nn.Linear`` roles in the pinned Diffusers Qwen-Image block, +collects one synthetic full-precision max-calibration pass, freezes ModelOpt's +two-level NVFP4 scales, and installs three strict benchmark treatments: + +``A`` + Original bf16 linears/GELU FFNs. +``B`` + Original bf16 non-MLP linears plus ``cudnn.gemm.ops.gelu_mlp``. +``C`` + Native NVFP4 for all fourteen logical linears. The two FFNs use a private + FROST FC1+bias+GELU+requant graph so the FC2 input quantization is fused. + +Weights are quantized once during setup. Activations use fixed max-calibrated +global scales and are quantized at runtime. There is no fallback: structure, +shape, route, cache-sharing, or kernel-contract drift raises. +""" + +from __future__ import annotations + +from collections import Counter +from contextlib import contextmanager +from dataclasses import dataclass +import copy +import types + +MODELOPT_RECIPE = { + "id": "qwen-image-modelopt-0.46.0-nvfp4-max-interior-v1", + "project": "NVIDIA/Model-Optimizer", + "repo_url": "https://github.com/NVIDIA/Model-Optimizer", + "release": "0.46.0", + "commit": "43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a", + "upstream_anchor_args": "--model qwen-image --format fp4 --quant-algo max", + "proxy_overrides": { + "model_dtype": "BFloat16 (ModelOpt CLI default is Half)", + "trt_high_precision_dtype": "BFloat16 (ModelOpt CLI default is Half)", + "calibration": "one deterministic synthetic pass instead of the upstream calibration workload", + "depth": "four representative blocks projected onto full-model blocks [2, 20, 39, 57]", + "weights": "offline one-time weight prepacking/compression during benchmark setup", + }, + "alignment_scope": "Linear placement plus NVFP4 format/block scaling and max-policy; not exact dtype, calibration state, or workload", + "quantize_mha": False, + "linear_format": "NVFP4 E2M1, block_size=16, E4M3 block scales", + "calibration": "one deterministic synthetic bf16 max-observer pass; frozen before timing", + "attention": "bf16 core; ModelOpt quantize_mha is false", + "scope": "representative quantized middle transformer blocks", + "full_model_quantized_blocks": [2, 57], + "full_model_excluded_blocks": [0, 1, 58, 59], + "numerical_claim_eligible": False, + "sources": { + "selection": "examples/diffusers/quantization/quantize.py", + "qwen_defaults": "examples/diffusers/quantization/models_utils.py", + "preset": "modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml", + "numerics": "modelopt_recipes/configs/numerics/nvfp4.yaml", + "mha_policy": "examples/diffusers/quantization/utils.py", + "real_backend": "modelopt/torch/quantization/backends/nvfp4_gemm.py", + }, + "source_permalinks": { + "selection": "https://github.com/NVIDIA/Model-Optimizer/blob/43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a/examples/diffusers/quantization/quantize.py", + "qwen_defaults": "https://github.com/NVIDIA/Model-Optimizer/blob/43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a/examples/diffusers/quantization/models_utils.py", + "preset": "https://github.com/NVIDIA/Model-Optimizer/blob/43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml", + "numerics": "https://github.com/NVIDIA/Model-Optimizer/blob/43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a/modelopt_recipes/configs/numerics/nvfp4.yaml", + "mha_policy": "https://github.com/NVIDIA/Model-Optimizer/blob/43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a/examples/diffusers/quantization/utils.py", + "real_backend": "https://github.com/NVIDIA/Model-Optimizer/blob/43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a/modelopt/torch/quantization/backends/nvfp4_gemm.py", + }, +} + +ARM_CONFIGS = { + "A": {"generic_linear": "bf16", "mlp": "torch"}, + "B": {"generic_linear": "bf16", "mlp": "cudnn_bf16"}, + "C": {"generic_linear": "nvfp4", "mlp": "nvfp4"}, +} + +ROLE_ORDER = ( + "img_mod.1", + "txt_mod.1", + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.add_q_proj", + "attn.add_k_proj", + "attn.add_v_proj", + "attn.to_out.0", + "attn.to_add_out", + "img_mlp.net.0.proj", + "img_mlp.net.2", + "txt_mlp.net.0.proj", + "txt_mlp.net.2", +) + +MLP_ROLES = frozenset( + { + "img_mlp.net.0.proj", + "img_mlp.net.2", + "txt_mlp.net.0.proj", + "txt_mlp.net.2", + } +) + +_MODELOPT_FP8_MAX = 448.0 +_NVFP4_E2M1_MAX = 6.0 +_BLOCK_SIZE = 16 +_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) +_LINEAR_REFERENCE_RTOL = 2.0e-2 +_LINEAR_REFERENCE_ATOL = 2.0e-1 +_FUSED_HIDDEN_REFERENCE_REL_L2 = 2.5e-1 +_FUSED_HIDDEN_ORACLE_REL_L2 = 1.2e-1 + + +def _ceil_to(value, multiple): + return ((int(value) + multiple - 1) // multiple) * multiple + + +def representative_middle_blocks(layers): + """Evenly map a depth-reduced proxy onto Qwen-Image's full blocks 2..57.""" + if not isinstance(layers, int) or not 1 <= layers <= 56: + raise ValueError(f"representative middle-block count must be in [1, 56], got {layers!r}") + if layers == 1: + return [30] + lo, span = 2, 55 + # Round halves up, not Python's round-to-even. Four layers map exactly to + # the declared review anchors [2, 20, 39, 57]. + result = [lo + (2 * i * span + layers - 1) // (2 * (layers - 1)) for i in range(layers)] + if len(set(result)) != layers or result[0] != 2 or result[-1] != 57: + raise AssertionError(f"invalid representative middle-block mapping: {result}") + return result + + +def three_arm_orders(): + """Position- and first-order-carryover-balanced design for three arms.""" + orders = ( + (0, 1, 2), + (1, 2, 0), + (2, 0, 1), + (2, 1, 0), + (0, 2, 1), + (1, 0, 2), + ) + positions = Counter((position, treatment) for order in orders for position, treatment in enumerate(order)) + carryover = Counter(pair for order in orders for pair in zip(order, order[1:])) + if set(positions.values()) != {2} or len(positions) != 9: + raise AssertionError(f"unbalanced three-arm positions: {positions}") + if set(carryover.values()) != {2} or len(carryover) != 6: + raise AssertionError(f"unbalanced three-arm carryover: {carryover}") + return orders + + +def expected_input_shapes(shape): + """Exact activation shape at each Linear boundary for one proxy block.""" + batch = shape["bs"] + hidden = shape["hidden"] + ffn = shape["ffn"] + image = shape["image_tokens"] + text = shape["text_tokens"] + return { + "img_mod.1": (batch, hidden), + "txt_mod.1": (batch, hidden), + "attn.to_q": (batch, image, hidden), + "attn.to_k": (batch, image, hidden), + "attn.to_v": (batch, image, hidden), + "attn.add_q_proj": (batch, text, hidden), + "attn.add_k_proj": (batch, text, hidden), + "attn.add_v_proj": (batch, text, hidden), + "attn.to_out.0": (batch, image, hidden), + "attn.to_add_out": (batch, text, hidden), + "img_mlp.net.0.proj": (batch, image, hidden), + "img_mlp.net.2": (batch, image, ffn), + "txt_mlp.net.0.proj": (batch, text, hidden), + "txt_mlp.net.2": (batch, text, ffn), + } + + +def expected_weight_shapes(shape): + hidden, ffn = shape["hidden"], shape["ffn"] + return { + "img_mod.1": (6 * hidden, hidden), + "txt_mod.1": (6 * hidden, hidden), + "attn.to_q": (hidden, hidden), + "attn.to_k": (hidden, hidden), + "attn.to_v": (hidden, hidden), + "attn.add_q_proj": (hidden, hidden), + "attn.add_k_proj": (hidden, hidden), + "attn.add_v_proj": (hidden, hidden), + "attn.to_out.0": (hidden, hidden), + "attn.to_add_out": (hidden, hidden), + "img_mlp.net.0.proj": (ffn, hidden), + "img_mlp.net.2": (hidden, ffn), + "txt_mlp.net.0.proj": (ffn, hidden), + "txt_mlp.net.2": (hidden, ffn), + } + + +def expected_plan_contracts(shape): + """The seven distinct low-precision plan contracts in the pinned proxy.""" + hidden, ffn = shape["hidden"], shape["ffn"] + image, text = shape["image_tokens"], shape["text_tokens"] + contracts = ( + (1, 6 * hidden, hidden, "linear_bias"), + (text, hidden, hidden, "linear_bias"), + (image, hidden, hidden, "linear_bias"), + (text, hidden, ffn, "linear_bias"), + (image, hidden, ffn, "linear_bias"), + (text, ffn, hidden, "linear_bias_gelu_nvfp4"), + (image, ffn, hidden, "linear_bias_gelu_nvfp4"), + ) + if len(set(contracts)) != 7: + raise ValueError(f"proxy dimensions collapse distinct NVFP4 plan contracts: {contracts}") + return contracts + + +def expected_route_delta(arm, layers): + if arm not in ARM_CONFIGS: + raise ValueError(f"unknown arm {arm!r}") + zero_roles = {f"transformer_blocks.{block}.{role}": 0 for block in range(layers) for role in ROLE_ORDER} + common = { + "weight_pack_calls": 0, + "plan_build_calls": 0, + "fallback_calls": 0, + "forward_scopes": 1, + "nvfp4_linear_by_role": zero_roles, + } + if arm == "A": + return { + **common, + "bf16_linear_calls": 14 * layers, + "nvfp4_linear_calls": 0, + "activation_quant_logical": 0, + "activation_quant_physical": 0, + "activation_quant_standalone": 0, + "activation_quant_fused": 0, + "activation_cache_hits": 0, + "mlp_calls": {"torch": 2 * layers, "cudnn_bf16": 0, "nvfp4": 0}, + } + if arm == "B": + return { + **common, + "bf16_linear_calls": 14 * layers, + "nvfp4_linear_calls": 0, + "activation_quant_logical": 0, + "activation_quant_physical": 0, + "activation_quant_standalone": 0, + "activation_quant_fused": 0, + "activation_cache_hits": 0, + "mlp_calls": {"torch": 0, "cudnn_bf16": 2 * layers, "nvfp4": 0}, + } + common["nvfp4_linear_by_role"] = {f"transformer_blocks.{block}.{role}": 1 for block in range(layers) for role in ROLE_ORDER} + physical = 1 + 8 * layers + return { + **common, + "bf16_linear_calls": 0, + "nvfp4_linear_calls": 14 * layers, + "activation_quant_logical": 14 * layers, + "activation_quant_physical": physical, + "activation_quant_standalone": 1 + 6 * layers, + "activation_quant_fused": 2 * layers, + "activation_cache_hits": 14 * layers - physical, + "mlp_calls": {"torch": 0, "cudnn_bf16": 0, "nvfp4": 2 * layers}, + } + + +def counter_delta(after, before): + """Subtract two counter snapshots with identical nested structure.""" + if set(after) != set(before): + raise ValueError("counter snapshots have different keys") + result = {} + for key in after: + if isinstance(after[key], dict): + if not isinstance(before[key], dict): + raise ValueError(f"counter type changed at {key}") + result[key] = counter_delta(after[key], before[key]) + else: + result[key] = int(after[key]) - int(before[key]) + return result + + +def _resolve_path(module, path): + value = module + for component in path.split("."): + value = value[int(component)] if component.isdigit() else getattr(value, component) + return value + + +def _activation_group(block_index, role): + if role in ("img_mod.1", "txt_mod.1"): + return "all_blocks.modulation" + if role in ("attn.to_q", "attn.to_k", "attn.to_v"): + return f"block{block_index}.image_qkv" + if role in ("attn.add_q_proj", "attn.add_k_proj", "attn.add_v_proj"): + return f"block{block_index}.text_qkv" + return f"block{block_index}.{role}" + + +@dataclass +class _LinearEntry: + qualified_name: str + block_index: int + role: str + module: object + original_forward: object + input_shape: tuple + weight_shape: tuple + activation_group: str + activation_amax: object = None + activation_global_scale: object = None + weight_amax: object = None + weight_global_scale: object = None + alpha: object = None + packed_weight: object = None + weight_scale_factors: object = None + + @property + def m(self): + result = 1 + for value in self.input_shape[:-1]: + result *= value + return result + + @property + def k(self): + return self.input_shape[-1] + + @property + def n(self): + return self.weight_shape[0] + + +def _collect_entries(model, shape, torch): + if model.__class__.__module__ != "diffusers.models.transformers.transformer_qwenimage" or model.__class__.__name__ != "QwenImageTransformer2DModel": + raise TypeError(f"expected pinned Diffusers QwenImageTransformer2DModel, got {type(model)!r}") + if model.training or any(parameter.requires_grad for parameter in model.parameters()): + raise NotImplementedError("ModelOpt NVFP4 benchmark treatment is inference-only") + if getattr(model, "_compiled_call_impl", None) is not None or getattr(model, "peft_config", None): + raise NotImplementedError("torch.compile and PEFT/LoRA are outside the ModelOpt NVFP4 benchmark treatment") + blocks = getattr(model, "transformer_blocks", None) + if not isinstance(blocks, torch.nn.ModuleList) or len(blocks) != shape["layers"]: + raise TypeError(f"expected {shape['layers']} Qwen-Image transformer blocks, got {type(blocks)!r}/{len(blocks) if blocks is not None else None}") + + input_shapes = expected_input_shapes(shape) + weight_shapes = expected_weight_shapes(shape) + hook_fields = ( + "_forward_hooks", + "_forward_pre_hooks", + "_backward_hooks", + "_backward_pre_hooks", + ) + entries = [] + for block_index, block in enumerate(blocks): + if block.__class__.__module__ != "diffusers.models.transformers.transformer_qwenimage" or block.__class__.__name__ != "QwenImageTransformerBlock": + raise TypeError(f"transformer_blocks[{block_index}] is not the pinned QwenImageTransformerBlock: {type(block)!r}") + if getattr(block, "zero_cond_t", None) is not False: + raise NotImplementedError("shared modulation quantization requires pinned zero_cond_t=False") + actual_linears = {name for name, child in block.named_modules() if type(child) is torch.nn.Linear} + if actual_linears != set(ROLE_ORDER): + raise TypeError(f"transformer_blocks[{block_index}] Linear roles changed: got {sorted(actual_linears)}, expected {sorted(ROLE_ORDER)}") + for role in ROLE_ORDER: + module = _resolve_path(block, role) + qualified = f"transformer_blocks.{block_index}.{role}" + if "forward" in module.__dict__: + raise NotImplementedError(f"{qualified} already has an instance-level forward override") + if module.bias is None or tuple(module.weight.shape) != weight_shapes[role] or tuple(module.bias.shape) != (weight_shapes[role][0],): + raise ValueError( + f"{qualified} shape changed: weight={tuple(module.weight.shape)}, bias={None if module.bias is None else tuple(module.bias.shape)}, " + f"expected {weight_shapes[role]}/{(weight_shapes[role][0],)}" + ) + for tensor_name, tensor in ( + ("weight", module.weight), + ("bias", module.bias), + ): + if tensor.dtype != torch.bfloat16 or tensor.device.type != "cuda" or not tensor.is_contiguous(): + raise NotImplementedError(f"{qualified}.{tensor_name} must be contiguous bf16 CUDA") + if module.training or any(getattr(module, field, None) for field in hook_fields) or torch.nn.utils.parametrize.is_parametrized(module): + raise NotImplementedError(f"training, hooks, or parametrizations on {qualified} are outside the benchmark treatment") + entries.append( + _LinearEntry( + qualified_name=qualified, + block_index=block_index, + role=role, + module=module, + original_forward=module.forward, + input_shape=input_shapes[role], + weight_shape=weight_shapes[role], + activation_group=_activation_group(block_index, role), + ) + ) + return entries + + +def collect_max_calibration(model, shape, forward_call): + """Collect one untimed BF16 max-observer pass for all exact Linear inputs.""" + import torch + + entries = _collect_entries(model, shape, torch) + by_module = {entry.module: entry for entry in entries} + maxima = {} + calls = Counter() + handles = [] + + def observe(module, args): + entry = by_module[module] + if len(args) != 1 or tuple(args[0].shape) != entry.input_shape: + raise ValueError(f"{entry.qualified_name} calibration input changed: got {[tuple(x.shape) for x in args]}, expected {entry.input_shape}") + value = args[0] + if value.dtype != torch.bfloat16 or value.device.type != "cuda" or not value.is_contiguous(): + raise NotImplementedError(f"{entry.qualified_name} calibration input must be contiguous bf16 CUDA") + observed = value.detach().abs().amax().float() + maxima[entry.qualified_name] = observed if entry.qualified_name not in maxima else torch.maximum(maxima[entry.qualified_name], observed) + calls[entry.qualified_name] += 1 + + try: + for entry in entries: + handles.append(entry.module.register_forward_pre_hook(observe)) + with torch.inference_mode(): + output = forward_call() + finally: + for handle in handles: + handle.remove() + + expected_names = {entry.qualified_name for entry in entries} + if set(maxima) != expected_names or set(calls.values()) != {1}: + raise RuntimeError(f"calibration did not visit every Linear exactly once: calls={dict(calls)}") + if not bool(torch.isfinite(output).all()): + raise RuntimeError("synthetic calibration produced non-finite model output") + for name, value in maxima.items(): + if not bool(torch.isfinite(value)) or not bool(value > 0): + raise RuntimeError(f"invalid calibration amax for {name}: {value}") + + # These groups receive mathematically identical BF16 values in the pinned + # call graph. Exact equality lets one packed activation serve all consumers. + grouped = {} + for entry in entries: + grouped.setdefault(entry.activation_group, []).append(entry.qualified_name) + for group, names in grouped.items(): + values = [float(maxima[name].item()) for name in names] + if len(names) > 1 and any(value != values[0] for value in values[1:]): + raise RuntimeError(f"shared activation group {group} calibrated different maxima: {dict(zip(names, values))}") + return { + "amax": maxima, + "calls": dict(calls), + "metadata": { + "method": "synthetic_bf16_max", + "passes": 1, + "frozen_before_timing": True, + "amax": {name: float(value.item()) for name, value in maxima.items()}, + }, + } + + +def _tensor_signature(tensor): + """Identity plus metadata for a buffer baked into a resolved binding.""" + return ( + tensor.data_ptr(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + ) + + +def _run_resolved_with_temporary_output(compiled, resolved, output_id, output, *, stream): + """Bind one per-call output without retaining it in the prepared mapping.""" + if resolved.get(output_id) is not None: + raise RuntimeError("prepared NVFP4 output slot is already occupied") + resolved[output_id] = output + try: + return compiled.run_resolved(resolved, stream=stream) + finally: + # Do not let the per-entry cache pin every full-model activation. + resolved.pop(output_id, None) + + +@dataclass +class _PreparedNvfp4Binding: + entry: object + resolved: dict + resolved_refs: dict + resolved_signatures: dict + activation_packed: object + activation_scale_factors: object + activation_global_scale: object + packed_weight: object + weight_scale_factors: object + alpha: object + bias: object + activation_packed_signature: tuple + activation_scale_signature: tuple + activation_global_scale_signature: tuple + packed_weight_signature: tuple + weight_scale_signature: tuple + alpha_signature: tuple + bias_signature: tuple + output_id: int | None = None + hidden_global_scale: object = None + hidden_global_scale_source: object = None + hidden_global_scale_signature: tuple | None = None + hidden_global_scale_source_signature: tuple | None = None + output_packed: object = None + output_scale_factors: object = None + output_packed_signature: tuple | None = None + output_scale_signature: tuple | None = None + + +def _validate_resolved_cache(prepared, label): + if set(prepared.resolved) != set(prepared.resolved_refs): + raise RuntimeError(f"{label} resolved binding keys changed after preparation") + for tensor_id, expected in prepared.resolved_refs.items(): + current = prepared.resolved.get(tensor_id) + if current is not expected or _tensor_signature(current) != prepared.resolved_signatures[tensor_id]: + raise RuntimeError(f"{label} resolved binding changed after preparation") + + +def _validate_cached_tensor(label, current, cached, signature): + if current is not cached or _tensor_signature(current) != signature: + raise RuntimeError(f"{label} changed after NVFP4 preparation") + + +class _Nvfp4LinearPlan: + def __init__(self, torch, cudnn, build_gemm_plan, *, m, n, k, device): + self.torch = torch + self.m, self.n, self.k = m, n, k + sf_k = k // _BLOCK_SIZE + fp4, fp8 = cudnn.data_type.FP4_E2M1, cudnn.data_type.FP8_E4M3 + reorder = dict(reordering_type=cudnn.tensor_reordering.F8_128x4) + with torch.cuda.device(device): + graph = cudnn.pygraph( + io_data_type=cudnn.data_type.BFLOAT16, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + self.A = graph.tensor(name="A", dim=[1, m, k], stride=[m * k, k, 1], data_type=fp4) + # Graph descriptors remain logical. The variant pack carries the + # larger F8_128x4 physical blob produced by nvfp4_quantize. + self.SFA = graph.tensor( + name="SFA", + dim=[1, m, sf_k], + stride=[m * sf_k, sf_k, 1], + data_type=fp8, + **reorder, + ) + self.B = graph.tensor(name="B", dim=[1, k, n], stride=[k * n, 1, k], data_type=fp4) + self.SFB = graph.tensor( + name="SFB", + dim=[1, sf_k, n], + stride=[sf_k * n, 1, sf_k], + data_type=fp8, + **reorder, + ) + self.ALPHA = graph.tensor( + name="ALPHA", + dim=[1, 1, 1], + stride=[1, 1, 1], + data_type=cudnn.data_type.FLOAT, + ) + self.BIAS = graph.tensor( + name="BIAS", + dim=[1, 1, n], + stride=[n, n, 1], + data_type=cudnn.data_type.BFLOAT16, + ) + ad = graph.block_scale_dequantize(input=self.A, descale=self.SFA, block_size=[1, _BLOCK_SIZE]) + bd = graph.block_scale_dequantize(input=self.B, descale=self.SFB, block_size=[_BLOCK_SIZE, 1]) + acc = graph.matmul(A=ad, B=bd, name="linear") + corrected = graph.mul(a=acc, b=self.ALPHA, name="modelopt_alpha") + # ModelOpt's real backend requests a BF16 GEMM result, then performs + # the official Linear bias add. Preserve both observable boundaries. + corrected.set_data_type(cudnn.data_type.BFLOAT16) + self.Y = graph.bias(input=corrected, bias=self.BIAS, name="bias") + self.Y.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + self.compiled = build_gemm_plan(graph) + self.device = device + self._prepared = {} + + def prepare(self, resolve_variant_pack, activation_buffers, entry): + """Resolve the stable per-entry bindings once; Y remains per-call.""" + torch = self.torch + packed, scale_factors = activation_buffers + if id(entry) in self._prepared: + raise RuntimeError(f"duplicate prepared NVFP4 binding for {entry.qualified_name}") + variant_pack = { + self.A: packed.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFA: scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.B: entry.packed_weight.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFB: entry.weight_scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.ALPHA: entry.alpha, + self.BIAS: entry.module.bias.view(1, 1, self.n), + self.Y: None, + } + resolved = resolve_variant_pack(variant_pack, self.compiled.binding) + expected = {id(tensor) for tensor in self.compiled.bound} + if set(resolved) != expected or resolved.get(id(self.Y)) is not None: + raise RuntimeError(f"incomplete prepared NVFP4 binding for {entry.qualified_name}") + resolved.pop(id(self.Y)) + resolved_refs = dict(resolved) + self._prepared[id(entry)] = _PreparedNvfp4Binding( + entry=entry, + resolved=resolved, + resolved_refs=resolved_refs, + resolved_signatures={tensor_id: _tensor_signature(tensor) for tensor_id, tensor in resolved_refs.items()}, + activation_packed=packed, + activation_scale_factors=scale_factors, + activation_global_scale=entry.activation_global_scale, + packed_weight=entry.packed_weight, + weight_scale_factors=entry.weight_scale_factors, + alpha=entry.alpha, + bias=entry.module.bias, + activation_packed_signature=_tensor_signature(packed), + activation_scale_signature=_tensor_signature(scale_factors), + activation_global_scale_signature=_tensor_signature(entry.activation_global_scale), + packed_weight_signature=_tensor_signature(entry.packed_weight), + weight_scale_signature=_tensor_signature(entry.weight_scale_factors), + alpha_signature=_tensor_signature(entry.alpha), + bias_signature=_tensor_signature(entry.module.bias), + output_id=id(self.Y), + ) + + def validate_prepared(self, activation_buffers, entry): + prepared = self._prepared.get(id(entry)) + if prepared is None or prepared.entry is not entry: + raise RuntimeError(f"missing prepared NVFP4 binding for {entry.qualified_name}") + packed, scale_factors = activation_buffers + _validate_cached_tensor( + f"{entry.qualified_name} A", + packed, + prepared.activation_packed, + prepared.activation_packed_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} SFA", + scale_factors, + prepared.activation_scale_factors, + prepared.activation_scale_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} activation global scale", + entry.activation_global_scale, + prepared.activation_global_scale, + prepared.activation_global_scale_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} B", + entry.packed_weight, + prepared.packed_weight, + prepared.packed_weight_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} SFB", + entry.weight_scale_factors, + prepared.weight_scale_factors, + prepared.weight_scale_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} alpha", + entry.alpha, + prepared.alpha, + prepared.alpha_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} bias", + entry.module.bias, + prepared.bias, + prepared.bias_signature, + ) + if prepared.output_id in prepared.resolved: + raise RuntimeError(f"{entry.qualified_name} prepared output slot retained a tensor") + _validate_resolved_cache(prepared, entry.qualified_name) + + def _binding(self, activation, entry, alpha, bias): + prepared = self._prepared.get(id(entry)) + if prepared is None or prepared.entry is not entry: + raise RuntimeError(f"missing prepared NVFP4 binding for {entry.qualified_name}") + if ( + activation.packed is not prepared.activation_packed + or activation.scale_factors is not prepared.activation_scale_factors + or activation.global_scale is not prepared.activation_global_scale + ): + raise RuntimeError(f"{entry.qualified_name} activation buffers changed after NVFP4 preparation") + if ( + entry.packed_weight is not prepared.packed_weight + or entry.weight_scale_factors is not prepared.weight_scale_factors + or alpha is not prepared.alpha + or bias is not prepared.bias + ): + raise RuntimeError(f"{entry.qualified_name} weight-scale or bias binding changed after NVFP4 preparation") + return prepared + + def __call__(self, activation, weight, alpha, bias): + torch = self.torch + prepared = self._binding(activation, weight, alpha, bias) + output = torch.empty((1, self.m, self.n), dtype=torch.bfloat16, device=self.device) + stream = torch.cuda.current_stream(self.device).cuda_stream + _run_resolved_with_temporary_output( + self.compiled, + prepared.resolved, + prepared.output_id, + output, + stream=stream, + ) + return output.squeeze(0) + + def run_unprepared(self, activation, weight, alpha, bias): + """Dynamic-buffer path used only by the untimed numerical contract gate.""" + torch = self.torch + output = torch.empty((1, self.m, self.n), dtype=torch.bfloat16, device=self.device) + stream = torch.cuda.current_stream(self.device).cuda_stream + self.compiled( + { + self.A: activation.packed.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFA: activation.scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.B: weight.packed_weight.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFB: weight.weight_scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.ALPHA: alpha, + self.BIAS: bias.view(1, 1, self.n), + self.Y: output, + }, + stream=stream, + ) + return output.squeeze(0) + + +class _Nvfp4FusedFc1Plan: + def __init__(self, torch, cudnn, build_gemm_plan, *, m, n, k, device): + self.torch = torch + self.m, self.n, self.k = m, n, k + pm = _ceil_to(m, 128) + sf_k = k // _BLOCK_SIZE + sn = _ceil_to(n // _BLOCK_SIZE, 4) + fp4, fp8 = cudnn.data_type.FP4_E2M1, cudnn.data_type.FP8_E4M3 + reorder = dict(reordering_type=cudnn.tensor_reordering.F8_128x4) + with torch.cuda.device(device): + graph = cudnn.pygraph( + io_data_type=cudnn.data_type.BFLOAT16, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + self.A = graph.tensor(name="A", dim=[1, m, k], stride=[m * k, k, 1], data_type=fp4) + self.SFA = graph.tensor( + name="SFA", + dim=[1, m, sf_k], + stride=[m * sf_k, sf_k, 1], + data_type=fp8, + **reorder, + ) + self.B = graph.tensor(name="B", dim=[1, k, n], stride=[k * n, 1, k], data_type=fp4) + self.SFB = graph.tensor( + name="SFB", + dim=[1, sf_k, n], + stride=[sf_k * n, 1, sf_k], + data_type=fp8, + **reorder, + ) + self.ALPHA = graph.tensor( + name="ALPHA", + dim=[1, 1, 1], + stride=[1, 1, 1], + data_type=cudnn.data_type.FLOAT, + ) + self.BIAS = graph.tensor( + name="BIAS", + dim=[1, 1, n], + stride=[n, n, 1], + data_type=cudnn.data_type.BFLOAT16, + ) + self.HIDDEN_GLOBAL_SCALE = graph.tensor( + name="HIDDEN_GLOBAL_SCALE", + dim=[1, 1, 1], + stride=[1, 1, 1], + data_type=cudnn.data_type.FLOAT, + ) + ad = graph.block_scale_dequantize(input=self.A, descale=self.SFA, block_size=[1, _BLOCK_SIZE]) + bd = graph.block_scale_dequantize(input=self.B, descale=self.SFB, block_size=[_BLOCK_SIZE, 1]) + acc = graph.matmul(A=ad, B=bd, name="fc1") + corrected = graph.mul(a=acc, b=self.ALPHA, name="modelopt_alpha") + corrected.set_data_type(cudnn.data_type.BFLOAT16) + pre = graph.bias(input=corrected, bias=self.BIAS, name="bias") + pre.set_data_type(cudnn.data_type.BFLOAT16) + hidden = graph.gelu_approx_tanh(input=pre, name="gelu_tanh") + # Diffusers' eager GELU returns BF16; FC2 observes and quantizes that + # value, not an unrounded FP32 epilogue intermediate. + hidden.set_data_type(cudnn.data_type.BFLOAT16) + hidden_scaled = graph.mul( + a=hidden, + b=self.HIDDEN_GLOBAL_SCALE, + name="hidden_modelopt_global_scale", + ) + self.QH, self.SH = graph.block_scale_quantize( + input=hidden_scaled, + block_size=_BLOCK_SIZE, + axis=-1, + name="hidden_nvfp4", + ) + self.QH.set_output(True).set_data_type(fp4) + self.SH.set_dim([1, pm, sn]).set_stride([pm * sn, sn, 1]) + self.SH.set_output(True).set_data_type(fp8).set_reordering_type(cudnn.tensor_reordering.F8_128x4) + self.compiled = build_gemm_plan(graph) + self.qh = torch.empty((m, n // 2), dtype=torch.uint8, device=device) + self.sh = torch.empty((pm, sn), dtype=torch.uint8, device=device) + self.device = device + self._prepared = {} + + def prepare(self, resolve_variant_pack, activation_buffers, entry, hidden_global_scale): + """Resolve stable input, weight, and fixed output taps for one FC1.""" + torch = self.torch + packed, scale_factors = activation_buffers + if id(entry) in self._prepared: + raise RuntimeError(f"duplicate prepared NVFP4 binding for {entry.qualified_name}") + hidden_global_scale_source = hidden_global_scale + hidden_global_scale = hidden_global_scale_source.reshape(1, 1, 1) + variant_pack = { + self.A: packed.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFA: scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.B: entry.packed_weight.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFB: entry.weight_scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.ALPHA: entry.alpha, + self.BIAS: entry.module.bias.view(1, 1, self.n), + self.HIDDEN_GLOBAL_SCALE: hidden_global_scale, + self.QH: self.qh.view(torch.int8).unsqueeze(0), + self.SH: self.sh.view(torch.float8_e4m3fn).unsqueeze(0), + } + resolved = resolve_variant_pack(variant_pack, self.compiled.binding) + expected = {id(tensor) for tensor in self.compiled.bound} + if set(resolved) != expected: + raise RuntimeError(f"incomplete prepared NVFP4 binding for {entry.qualified_name}") + resolved_refs = dict(resolved) + self._prepared[id(entry)] = _PreparedNvfp4Binding( + entry=entry, + resolved=resolved, + resolved_refs=resolved_refs, + resolved_signatures={tensor_id: _tensor_signature(tensor) for tensor_id, tensor in resolved_refs.items()}, + activation_packed=packed, + activation_scale_factors=scale_factors, + activation_global_scale=entry.activation_global_scale, + packed_weight=entry.packed_weight, + weight_scale_factors=entry.weight_scale_factors, + alpha=entry.alpha, + bias=entry.module.bias, + activation_packed_signature=_tensor_signature(packed), + activation_scale_signature=_tensor_signature(scale_factors), + activation_global_scale_signature=_tensor_signature(entry.activation_global_scale), + packed_weight_signature=_tensor_signature(entry.packed_weight), + weight_scale_signature=_tensor_signature(entry.weight_scale_factors), + alpha_signature=_tensor_signature(entry.alpha), + bias_signature=_tensor_signature(entry.module.bias), + hidden_global_scale=hidden_global_scale, + hidden_global_scale_source=hidden_global_scale_source, + hidden_global_scale_signature=_tensor_signature(hidden_global_scale), + hidden_global_scale_source_signature=_tensor_signature(hidden_global_scale_source), + output_packed=self.qh, + output_scale_factors=self.sh, + output_packed_signature=_tensor_signature(self.qh), + output_scale_signature=_tensor_signature(self.sh), + ) + + def validate_prepared(self, activation_buffers, entry, hidden_global_scale_source): + prepared = self._prepared.get(id(entry)) + if prepared is None or prepared.entry is not entry: + raise RuntimeError(f"missing prepared NVFP4 binding for {entry.qualified_name}") + packed, scale_factors = activation_buffers + _validate_cached_tensor( + f"{entry.qualified_name} A", + packed, + prepared.activation_packed, + prepared.activation_packed_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} SFA", + scale_factors, + prepared.activation_scale_factors, + prepared.activation_scale_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} activation global scale", + entry.activation_global_scale, + prepared.activation_global_scale, + prepared.activation_global_scale_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} B", + entry.packed_weight, + prepared.packed_weight, + prepared.packed_weight_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} SFB", + entry.weight_scale_factors, + prepared.weight_scale_factors, + prepared.weight_scale_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} alpha", + entry.alpha, + prepared.alpha, + prepared.alpha_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} bias", + entry.module.bias, + prepared.bias, + prepared.bias_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} HGS source", + hidden_global_scale_source, + prepared.hidden_global_scale_source, + prepared.hidden_global_scale_source_signature, + ) + if _tensor_signature(prepared.hidden_global_scale) != prepared.hidden_global_scale_signature: + raise RuntimeError(f"{entry.qualified_name} HGS view changed after NVFP4 preparation") + _validate_cached_tensor( + f"{entry.qualified_name} QH", + self.qh, + prepared.output_packed, + prepared.output_packed_signature, + ) + _validate_cached_tensor( + f"{entry.qualified_name} SH", + self.sh, + prepared.output_scale_factors, + prepared.output_scale_signature, + ) + _validate_resolved_cache(prepared, entry.qualified_name) + + def _binding(self, activation, entry, alpha, bias): + prepared = self._prepared.get(id(entry)) + if prepared is None or prepared.entry is not entry: + raise RuntimeError(f"missing prepared NVFP4 binding for {entry.qualified_name}") + if ( + activation.packed is not prepared.activation_packed + or activation.scale_factors is not prepared.activation_scale_factors + or activation.global_scale is not prepared.activation_global_scale + ): + raise RuntimeError(f"{entry.qualified_name} activation buffers changed after NVFP4 preparation") + if ( + entry.packed_weight is not prepared.packed_weight + or entry.weight_scale_factors is not prepared.weight_scale_factors + or alpha is not prepared.alpha + or bias is not prepared.bias + or self.qh is not prepared.output_packed + or self.sh is not prepared.output_scale_factors + ): + raise RuntimeError(f"{entry.qualified_name} weight-scale or bias binding changed after NVFP4 preparation") + return prepared + + def __call__(self, activation, weight, alpha, bias): + prepared = self._binding(activation, weight, alpha, bias) + stream = self.torch.cuda.current_stream(self.device).cuda_stream + self.compiled.run_resolved(prepared.resolved, stream=stream) + return _QuantizedActivation( + self.qh, + self.sh, + prepared.hidden_global_scale_source, + source=None, + group="fused_hidden", + ) + + def run_unprepared(self, activation, weight, alpha, bias, hidden_global_scale): + """Dynamic-buffer path used only by the untimed numerical contract gate.""" + torch = self.torch + stream = torch.cuda.current_stream(self.device).cuda_stream + self.compiled( + { + self.A: activation.packed.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFA: activation.scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.B: weight.packed_weight.view(torch.float4_e2m1fn_x2).unsqueeze(0), + self.SFB: weight.weight_scale_factors.view(torch.float8_e4m3fn).unsqueeze(0), + self.ALPHA: alpha, + self.BIAS: bias.view(1, 1, self.n), + self.HIDDEN_GLOBAL_SCALE: hidden_global_scale, + # FROST's generated host ABI uses Int8 as the byte carrier for + # packed FP4 output taps (the logical graph dtype remains E2M1). + self.QH: self.qh.view(torch.int8).unsqueeze(0), + self.SH: self.sh.view(torch.float8_e4m3fn).unsqueeze(0), + }, + stream=stream, + ) + return _QuantizedActivation(self.qh, self.sh, hidden_global_scale, source=None, group="fused_hidden") + + +@dataclass +class _QuantizedActivation: + packed: object + scale_factors: object + global_scale: object + source: object + group: str + + +class QwenImageModelOptNvfp4Adapter: + """Installed, fail-closed dispatcher for the three benchmark arms.""" + + def __init__(self, model, shape, calibration): + import torch + import cudnn + import cudnn.gemm.frost # noqa: F401 -- installs graph recording hooks + from cudnn.gemm.frost.graph_analyzer import ( + build_gemm_plan, + resolve_variant_pack, + ) + from cudnn.gemm.ops._nvfp4_quantize import nvfp4_quantize + from cudnn.gemm.ops import gelu_mlp + + self.torch = torch + self.cudnn = cudnn + self.nvfp4_quantize = nvfp4_quantize + self.gelu_mlp = gelu_mlp + self.shape = dict(shape) + self.layers = shape["layers"] + self.entries = _collect_entries(model, shape, torch) + self.by_name = {entry.qualified_name: entry for entry in self.entries} + if set(calibration) != {"amax", "calls", "metadata"} or set(calibration["amax"]) != set(self.by_name): + raise ValueError("calibration does not match the exact Qwen-Image Linear role set") + self.calibration_metadata = copy.deepcopy(calibration["metadata"]) + self.counters = { + "bf16_linear_calls": 0, + "nvfp4_linear_calls": 0, + "activation_quant_logical": 0, + "activation_quant_physical": 0, + "activation_quant_standalone": 0, + "activation_quant_fused": 0, + "activation_cache_hits": 0, + "weight_pack_calls": 0, + "plan_build_calls": 0, + "fallback_calls": 0, + "forward_scopes": 0, + "mlp_calls": {"torch": 0, "cudnn_bf16": 0, "nvfp4": 0}, + "nvfp4_linear_by_role": {entry.qualified_name: 0 for entry in self.entries}, + } + self._selected = None + self._active = False + self._active_role_order = [] + self._activation_cache = {} + self._installed_generic = {} + self._installed_mod = {} + self._installed_mlp = {} + self._activation_buffers = {} + self._linear_plans = {} + self._fused_fc1_plans = {} + + group_scale = {} + device = self.entries[0].module.weight.device + self._device = device + self._stream = torch.cuda.current_stream(device).cuda_stream + for entry in self.entries: + entry.activation_amax = calibration["amax"][entry.qualified_name] + scale = (_MODELOPT_FP8_MAX * _NVFP4_E2M1_MAX / entry.activation_amax).reshape(1).contiguous() + existing = group_scale.get(entry.activation_group) + if existing is None: + group_scale[entry.activation_group] = scale + elif float(existing.item()) != float(scale.item()): + raise RuntimeError(f"activation group {entry.activation_group} has unequal frozen global scales") + entry.activation_global_scale = group_scale[entry.activation_group] + + entry.weight_amax = entry.module.weight.detach().abs().amax().float() + if not bool(torch.isfinite(entry.weight_amax)) or not bool(entry.weight_amax > 0): + raise RuntimeError(f"invalid weight amax for {entry.qualified_name}: {entry.weight_amax}") + entry.weight_global_scale = (_MODELOPT_FP8_MAX * _NVFP4_E2M1_MAX / entry.weight_amax).reshape(1).contiguous() + entry.alpha = (1.0 / (entry.activation_global_scale * entry.weight_global_scale)).reshape(1, 1, 1).contiguous() + entry.packed_weight, entry.weight_scale_factors = nvfp4_quantize( + entry.module.weight.detach(), + entry.weight_global_scale, + pre_quant_scale=None, + enable_pdl=True, + ) + self._validate_quantized_buffers( + entry.packed_weight, + entry.weight_scale_factors, + entry.n, + entry.k, + f"{entry.qualified_name} weight", + ) + self.counters["weight_pack_calls"] += 1 + + for entry in self.entries: + key = (entry.m, entry.k) + if key not in self._activation_buffers: + self._activation_buffers[key] = ( + torch.empty((entry.m, entry.k // 2), dtype=torch.uint8, device=device), + torch.empty( + (_ceil_to(entry.m, 128), _ceil_to(entry.k // _BLOCK_SIZE, 4)), + dtype=torch.uint8, + device=device, + ), + ) + + linear_shapes = {(entry.m, entry.n, entry.k) for entry in self.entries if entry.role not in MLP_ROLES or entry.role.endswith("net.2")} + fused_shapes = {(entry.m, entry.n, entry.k) for entry in self.entries if entry.role.endswith("net.0.proj")} + with torch.cuda.device(device): + for m, n, k in sorted(linear_shapes): + key = (device.index, self._stream, m, n, k, "linear_bias") + self._linear_plans[key] = _Nvfp4LinearPlan(torch, cudnn, build_gemm_plan, m=m, n=n, k=k, device=device) + self.counters["plan_build_calls"] += 1 + for m, n, k in sorted(fused_shapes): + key = (device.index, self._stream, m, n, k, "linear_bias_gelu_nvfp4") + self._fused_fc1_plans[key] = _Nvfp4FusedFc1Plan(torch, cudnn, build_gemm_plan, m=m, n=n, k=k, device=device) + self.counters["plan_build_calls"] += 1 + actual_contracts = {(m, n, k, epilogue) for _, _, m, n, k, epilogue in (*self._linear_plans, *self._fused_fc1_plans)} + if actual_contracts != set(expected_plan_contracts(shape)): + raise RuntimeError(f"NVFP4 plan contracts changed: got={sorted(actual_contracts)}, expected={sorted(expected_plan_contracts(shape))}") + if self.counters["weight_pack_calls"] != 14 * self.layers or self.counters["plan_build_calls"] != 7: + raise RuntimeError(f"unexpected setup counts: {self.counters}") + + # Resolve stable typed views and graph bindings once per logical Linear. + # Plans remain shared by shape, but weights/biases are per-entry. FC2's + # activation storage is the fixed output tap of its matching FC1 plan. + for entry in self.entries: + if entry.role.endswith("net.0.proj"): + second_name = entry.qualified_name.removesuffix("net.0.proj") + "net.2" + second = self.by_name.get(second_name) + if second is None: + raise RuntimeError(f"cannot prepare fused FC1 without {second_name}") + self._fused_fc1_plan(entry).prepare( + resolve_variant_pack, + self._activation_buffers[(entry.m, entry.k)], + entry, + second.activation_global_scale, + ) + elif entry.role.endswith("net.2"): + first_name = entry.qualified_name.removesuffix("net.2") + "net.0.proj" + first = self.by_name.get(first_name) + if first is None: + raise RuntimeError(f"cannot prepare FC2 without {first_name}") + first_plan = self._fused_fc1_plan(first) + self._linear_plan(entry).prepare( + resolve_variant_pack, + (first_plan.qh, first_plan.sh), + entry, + ) + else: + self._linear_plan(entry).prepare( + resolve_variant_pack, + self._activation_buffers[(entry.m, entry.k)], + entry, + ) + prepared_linear = sum(len(plan._prepared) for plan in self._linear_plans.values()) + prepared_fused = sum(len(plan._prepared) for plan in self._fused_fc1_plans.values()) + if (prepared_linear, prepared_fused) != (12 * self.layers, 2 * self.layers): + raise RuntimeError("unexpected prepared NVFP4 binding counts: " f"linear={prepared_linear}, fused_fc1={prepared_fused}") + self.binding_mode = { + "name": "pre_resolved_run_resolved", + "scope": "per_linear_entry", + "stable_typed_views_cached": True, + "stable_binding_validation": "full reference/signature preflight in select('C'); hot path object identity only", + "linear_output": "fresh allocation; resolved Y slot cleared in finally", + "fused_fc1_outputs": "plan-owned fixed QH/SH consumed immediately by FC2 on the same guarded stream", + "concurrency": "single-thread, non-reentrant, one prepared CUDA stream; mutation after select is unsupported", + "private_lowered_call": False, + "prepared_linear_entries": prepared_linear, + "prepared_fused_fc1_entries": prepared_fused, + } + + self._install_generic_linears() + self._install_modulations(model) + self._install_mlps(model) + + def _validate_quantized_buffers(self, packed, scale_factors, rows, k, label): + torch = self.torch + expected_packed = (rows, k // 2) + expected_scale = (_ceil_to(rows, 128), _ceil_to(k // _BLOCK_SIZE, 4)) + if packed.dtype != torch.uint8 or not packed.is_contiguous() or tuple(packed.shape) != expected_packed: + raise RuntimeError(f"{label} packed output changed: got {packed.dtype}/{tuple(packed.shape)}/{packed.stride()}, expected uint8/{expected_packed}") + if scale_factors.dtype != torch.uint8 or not scale_factors.is_contiguous() or tuple(scale_factors.shape) != expected_scale: + raise RuntimeError( + f"{label} scale-factor output changed: got {scale_factors.dtype}/{tuple(scale_factors.shape)}/{scale_factors.stride()}, " + f"expected uint8/{expected_scale}" + ) + + @staticmethod + def _sample_indices(extent): + return tuple(sorted({0, int(extent) // 2, int(extent) - 1})) + + def _dequantize_rows(self, packed, scale_factors, logical_rows, k, row_indices, label): + """Decode selected rows without using a cuDNN/FROST layout helper.""" + torch = self.torch + rows = torch.tensor(tuple(row_indices), dtype=torch.long, device=self._device) + if packed.dtype != torch.uint8 or tuple(packed.shape) != (logical_rows, k // 2): + raise RuntimeError(f"{label} packed reference input changed: {packed.dtype}/{tuple(packed.shape)}") + expected_sf = (_ceil_to(logical_rows, 128), _ceil_to(k // _BLOCK_SIZE, 4)) + if scale_factors.dtype != torch.uint8 or tuple(scale_factors.shape) != expected_sf: + raise RuntimeError(f"{label} scale reference input changed: {scale_factors.dtype}/{tuple(scale_factors.shape)}") + + selected = packed.index_select(0, rows) + lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=self._device) + low = lut[(selected & 0xF).long()] + high = lut[(selected >> 4).long()] + decoded = torch.stack((low, high), dim=-1).flatten(-2) + + sf_columns = k // _BLOCK_SIZE + column_groups = _ceil_to(sf_columns, 4) // 4 + blocks = torch.arange(sf_columns, dtype=torch.long, device=self._device).unsqueeze(0) + row_grid = rows.unsqueeze(1) + addresses = ((row_grid // 128) * column_groups + (blocks // 4)) * 512 + (row_grid % 32) * 16 + ((row_grid % 128) // 32) * 4 + (blocks % 4) + flat_scales = scale_factors.view(torch.float8_e4m3fn).flatten() + if int(addresses.max().item()) >= flat_scales.numel(): + raise RuntimeError(f"{label} F8_128x4 reference address exceeds the physical blob") + logical_scales = flat_scales[addresses].float() + return decoded * logical_scales.repeat_interleave(_BLOCK_SIZE, dim=1) + + def _reference_linear_samples(self, activation, entry, row_indices, column_indices): + torch = self.torch + x = self._dequantize_rows( + activation.packed, + activation.scale_factors, + entry.m, + entry.k, + row_indices, + f"{entry.qualified_name} activation", + ) + weight = self._dequantize_rows( + entry.packed_weight, + entry.weight_scale_factors, + entry.n, + entry.k, + column_indices, + f"{entry.qualified_name} weight", + ) + corrected = (x @ weight.t()) * entry.alpha.float().reshape(()) + corrected = corrected.to(torch.bfloat16) + columns = torch.tensor(tuple(column_indices), dtype=torch.long, device=self._device) + return (corrected + entry.module.bias.index_select(0, columns)).to(torch.bfloat16) + + def _check_linear_samples(self, actual, expected, *, label): + torch = self.torch + actual_f, expected_f = actual.float(), expected.float() + if not bool(torch.isfinite(actual_f).all()) or not bool(torch.isfinite(expected_f).all()): + raise RuntimeError(f"{label} numerical reference contains a non-finite value") + absolute = (actual_f - expected_f).abs() + allowed = _LINEAR_REFERENCE_ATOL + _LINEAR_REFERENCE_RTOL * expected_f.abs() + violations = int((absolute > allowed).sum().item()) + rel_l2 = float((actual_f - expected_f).norm() / expected_f.norm().clamp_min(1.0e-12)) + maximum = float(absolute.max()) + result = { + "rel_l2": rel_l2, + "max_abs": maximum, + "rtol": _LINEAR_REFERENCE_RTOL, + "atol": _LINEAR_REFERENCE_ATOL, + "elements": actual.numel(), + "violations": violations, + } + if violations: + raise RuntimeError(f"{label} failed independent dequantized-operand reference: {result}") + return result + + def _quantize_gate_input(self, entry, value, label): + packed, scale_factors = self.nvfp4_quantize( + value, + entry.activation_global_scale, + pre_quant_scale=None, + enable_pdl=True, + ) + self._validate_quantized_buffers(packed, scale_factors, entry.m, entry.k, label) + return _QuantizedActivation( + packed, + scale_factors, + entry.activation_global_scale, + source=value, + group=f"gate:{label}", + ) + + def _make_gate_input(self, entry, generator): + torch = self.torch + # Keep the deterministic probe near the frozen calibration range. The + # reference consumes the packed values, so clipping cannot mask a graph + # layout/orientation/alpha defect. + amplitude = max(float(entry.activation_amax.item()) / 8.0, 2.0**-12) + return ( + torch.randn( + (entry.m, entry.k), + dtype=torch.bfloat16, + device=self._device, + generator=generator, + ) + * amplitude + ).contiguous() + + def _stage_prepared_gate_activation(self, plan, entry, activation): + """Copy a gate probe into the exact buffers cached by one entry.""" + prepared = plan._prepared.get(id(entry)) + if prepared is None or prepared.entry is not entry: + raise RuntimeError(f"numerical gate cannot resolve prepared binding for {entry.qualified_name}") + if tuple(prepared.activation_packed.shape) != tuple(activation.packed.shape) or tuple(prepared.activation_scale_factors.shape) != tuple( + activation.scale_factors.shape + ): + raise RuntimeError(f"numerical gate prepared buffers changed for {entry.qualified_name}") + if float(prepared.activation_global_scale.item()) != float(activation.global_scale.item()): + raise RuntimeError(f"numerical gate global scale changed for {entry.qualified_name}") + prepared.activation_packed.copy_(activation.packed) + prepared.activation_scale_factors.copy_(activation.scale_factors) + return _QuantizedActivation( + prepared.activation_packed, + prepared.activation_scale_factors, + prepared.activation_global_scale, + source=None, + group=f"prepared_gate:{entry.qualified_name}", + ) + + def _run_prepared_binding_parity(self, generator): + """Exercise every per-entry resolved map against the dynamic gate path.""" + torch = self.torch + direct_checked = 0 + fused_checked = 0 + sequential_fc2_checked = 0 + qkv_outputs = {} + + for entry in self.entries: + value = self._make_gate_input(entry, generator) + activation = self._quantize_gate_input(entry, value, f"prepared parity {entry.qualified_name}") + if entry.role.endswith("net.0.proj"): + second_name = entry.qualified_name.removesuffix("net.0.proj") + "net.2" + second = self.by_name.get(second_name) + if second is None: + raise RuntimeError(f"prepared parity cannot resolve FC2 for {entry.qualified_name}") + plan = self._fused_fc1_plan(entry) + hidden_scale = plan._prepared[id(entry)].hidden_global_scale + dynamic_hidden = plan.run_unprepared( + activation, + entry, + entry.alpha, + entry.module.bias, + hidden_scale, + ) + dynamic_packed = dynamic_hidden.packed.clone() + dynamic_scales = dynamic_hidden.scale_factors.clone() + prepared_activation = self._stage_prepared_gate_activation(plan, entry, activation) + prepared_hidden = plan( + prepared_activation, + entry, + entry.alpha, + entry.module.bias, + ) + if not torch.equal(dynamic_packed, prepared_hidden.packed) or not torch.equal(dynamic_scales, prepared_hidden.scale_factors): + raise RuntimeError(f"{entry.qualified_name} prepared fused QH/SH differs from dynamic binding") + + # The fixed QH/SH taps are benchmark-safe only because FC2 + # consumes them immediately on the same prepared stream. + second_plan = self._linear_plan(second) + dynamic_output = second_plan.run_unprepared(prepared_hidden, second, second.alpha, second.module.bias) + prepared_output = second_plan(prepared_hidden, second, second.alpha, second.module.bias) + if not torch.equal(dynamic_output, prepared_output): + raise RuntimeError(f"{entry.qualified_name} -> {second.qualified_name} prepared FC2 differs from dynamic binding") + fused_checked += 1 + sequential_fc2_checked += 1 + continue + + plan = self._linear_plan(entry) + dynamic_output = plan.run_unprepared(activation, entry, entry.alpha, entry.module.bias) + prepared_activation = self._stage_prepared_gate_activation(plan, entry, activation) + prepared_output = plan(prepared_activation, entry, entry.alpha, entry.module.bias) + if not torch.equal(dynamic_output, prepared_output): + raise RuntimeError(f"{entry.qualified_name} prepared output differs from dynamic binding") + direct_checked += 1 + if entry.role in ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.add_q_proj", + "attn.add_k_proj", + "attn.add_v_proj", + ): + qkv_outputs.setdefault(entry.activation_group, []).append(prepared_output) + + for group, outputs in qkv_outputs.items(): + pointers = [output.data_ptr() for output in outputs] + if len(outputs) != 3 or len(set(pointers)) != 3: + raise RuntimeError(f"{group} prepared Q/K/V outputs do not have three independent Y allocations: {pointers}") + if ( + direct_checked != 12 * self.layers + or fused_checked != 2 * self.layers + or sequential_fc2_checked != 2 * self.layers + or len(qkv_outputs) != 2 * self.layers + ): + raise RuntimeError( + "prepared parity coverage changed: " + f"direct={direct_checked}, fused={fused_checked}, " + f"sequential_fc2={sequential_fc2_checked}, qkv_groups={len(qkv_outputs)}" + ) + return { + "status": "passed", + "comparison": "bitwise prepared run_resolved vs dynamic variant-pack binding", + "logical_entries_checked": direct_checked + fused_checked, + "direct_entries_checked": direct_checked, + "fused_fc1_entries_checked": fused_checked, + "sequential_fc1_fc2_pairs_checked": sequential_fc2_checked, + "qkv_groups_with_three_live_distinct_outputs": len(qkv_outputs), + "fused_output_lifetime": "plan-owned QH/SH consumed immediately by FC2 on the same guarded stream", + } + + def _reference_fused_hidden(self, activation, entry, row_indices): + torch = self.torch + x = self._dequantize_rows( + activation.packed, + activation.scale_factors, + entry.m, + entry.k, + row_indices, + f"{entry.qualified_name} activation", + ) + chunks = [] + # Chunk the independent decode so the image-token contract does not + # materialize the full 12288x3072 weight as FP32 at once. + for begin in range(0, entry.n, 1024): + end = min(begin + 1024, entry.n) + columns = tuple(range(begin, end)) + weight = self._dequantize_rows( + entry.packed_weight, + entry.weight_scale_factors, + entry.n, + entry.k, + columns, + f"{entry.qualified_name} weight", + ) + corrected = ((x @ weight.t()) * entry.alpha.float().reshape(())).to(torch.bfloat16) + pre = (corrected + entry.module.bias[begin:end]).to(torch.bfloat16) + chunks.append(torch.nn.functional.gelu(pre, approximate="tanh").to(torch.bfloat16)) + return torch.cat(chunks, dim=1) + + @staticmethod + def _relative_l2(torch, actual, expected): + return float((actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1.0e-12)) + + def run_focused_numerical_gate(self): + """Validate all seven plans once, outside route counts and timed work.""" + torch = self.torch + if self._active: + raise RuntimeError("cannot run the NVFP4 numerical gate during a model forward") + self._validate_prepared_bindings() + before = self.snapshot() + expected_contracts = expected_plan_contracts(self.shape) + generator = torch.Generator(device=self._device).manual_seed(0x4E56465034) + results = [] + + def representative(m, n, k, *, fused): + for candidate in self.entries: + candidate_fused = candidate.role.endswith("net.0.proj") + candidate_linear = candidate.role not in MLP_ROLES or candidate.role.endswith("net.2") + if (candidate_fused if fused else candidate_linear) and ( + candidate.m, + candidate.n, + candidate.k, + ) == (m, n, k): + return candidate + raise RuntimeError(f"no representative Linear for numerical contract {(m, n, k, fused)}") + + with torch.inference_mode(), torch.cuda.device(self._device): + prepared_parity = self._run_prepared_binding_parity(generator) + for m, n, k, epilogue in expected_contracts: + fused = epilogue == "linear_bias_gelu_nvfp4" + first = representative(m, n, k, fused=fused) + value = self._make_gate_input(first, generator) + activation = self._quantize_gate_input(first, value, f"gate {first.qualified_name}") + row_indices = self._sample_indices(m) + + if not fused: + column_indices = self._sample_indices(n) + actual_full = self._linear_plan(first).run_unprepared(activation, first, first.alpha, first.module.bias) + rows = torch.tensor(row_indices, dtype=torch.long, device=self._device) + columns = torch.tensor(column_indices, dtype=torch.long, device=self._device) + actual = actual_full.index_select(0, rows).index_select(1, columns) + expected = self._reference_linear_samples(activation, first, row_indices, column_indices) + results.append( + { + "contract": {"m": m, "n": n, "k": k, "epilogue": epilogue}, + "representative_role": first.qualified_name, + "sample_rows": list(row_indices), + "sample_columns": list(column_indices), + "output": self._check_linear_samples(actual, expected, label=f"gate {first.qualified_name}"), + } + ) + continue + + second_name = first.qualified_name.removesuffix("net.0.proj") + "net.2" + second = self.by_name.get(second_name) + if second is None or second.m != first.m or second.k != first.n: + raise RuntimeError(f"fused numerical gate cannot resolve FC2 for {first.qualified_name}") + hidden = self._fused_fc1_plan(first).run_unprepared( + activation, + first, + first.alpha, + first.module.bias, + second.activation_global_scale.reshape(1, 1, 1), + ) + expected_hidden = self._reference_fused_hidden(activation, first, row_indices) + fused_hidden_scaled = self._dequantize_rows( + hidden.packed, + hidden.scale_factors, + first.m, + first.n, + row_indices, + f"gate {first.qualified_name} fused hidden", + ) + hidden_scale = second.activation_global_scale.float().reshape(()) + hidden_rel_l2 = self._relative_l2(torch, fused_hidden_scaled / hidden_scale, expected_hidden) + + oracle_packed, oracle_scale = self.nvfp4_quantize( + expected_hidden.contiguous(), + second.activation_global_scale, + pre_quant_scale=None, + enable_pdl=True, + ) + oracle_scaled = self._dequantize_rows( + oracle_packed, + oracle_scale, + len(row_indices), + first.n, + tuple(range(len(row_indices))), + f"gate {first.qualified_name} standalone hidden oracle", + ) + oracle_rel_l2 = self._relative_l2(torch, fused_hidden_scaled, oracle_scaled) + hidden_metrics = { + "rel_l2_vs_bf16_gelu_reference": hidden_rel_l2, + "limit_vs_bf16_gelu_reference": _FUSED_HIDDEN_REFERENCE_REL_L2, + "rel_l2_vs_standalone_quant_oracle": oracle_rel_l2, + "limit_vs_standalone_quant_oracle": _FUSED_HIDDEN_ORACLE_REL_L2, + } + if not all(math_value == math_value and abs(math_value) != float("inf") for math_value in (hidden_rel_l2, oracle_rel_l2)): + raise RuntimeError(f"gate {first.qualified_name} fused hidden produced non-finite metrics: {hidden_metrics}") + if hidden_rel_l2 > _FUSED_HIDDEN_REFERENCE_REL_L2 or oracle_rel_l2 > _FUSED_HIDDEN_ORACLE_REL_L2: + raise RuntimeError(f"gate {first.qualified_name} fused hidden failed reference: {hidden_metrics}") + + column_indices = self._sample_indices(second.n) + actual_full = self._linear_plan(second).run_unprepared(hidden, second, second.alpha, second.module.bias) + rows = torch.tensor(row_indices, dtype=torch.long, device=self._device) + columns = torch.tensor(column_indices, dtype=torch.long, device=self._device) + actual = actual_full.index_select(0, rows).index_select(1, columns) + expected = self._reference_linear_samples(hidden, second, row_indices, column_indices) + results.append( + { + "contract": {"m": m, "n": n, "k": k, "epilogue": epilogue}, + "representative_role": first.qualified_name, + "sample_rows": list(row_indices), + "sample_columns": list(column_indices), + "fused_hidden": hidden_metrics, + "fc2_output": self._check_linear_samples( + actual, + expected, + label=f"gate {first.qualified_name} -> {second.qualified_name}", + ), + } + ) + torch.cuda.synchronize(self._device) + + if self.snapshot() != before: + raise RuntimeError(f"focused numerical gate contaminated route counters: before={before}, after={self.snapshot()}") + if len(results) != 7 or [tuple((*item["contract"].values(),)) for item in results] != list(expected_contracts): + raise RuntimeError(f"focused numerical gate did not cover the exact seven contracts: {results}") + return { + "status": "passed", + "contracts_checked": len(results), + "prepared_binding": prepared_parity, + "reference": "independent E2M1 decode plus F8_128x4 address map and sampled FP32 matmul; fused hidden also checked against BF16 tanh-GELU and standalone quantization", + "sampling": "first/middle/last rows and output columns; every full FROST plan still executes its exact M/N/K", + "results": results, + } + + def plan_provenance(self): + records = [] + for key, plan in sorted( + (*self._linear_plans.items(), *self._fused_fc1_plans.items()), + key=lambda item: item[0][2:], + ): + device, stream, m, n, k, epilogue = key + records.append( + { + "device": device, + "stream": stream, + "m": m, + "n": n, + "k": k, + "epilogue": epilogue, + "tile_config": plan.compiled.config.name, + "generated_path": str(plan.compiled.generated_path.resolve()), + } + ) + if len(records) != 7: + raise RuntimeError(f"expected seven low-precision plans, got {records}") + return records + + def _install_generic_linears(self): + for entry in self.entries: + if entry.role in MLP_ROLES or entry.role in ("img_mod.1", "txt_mod.1"): + continue + + def bf16_forward(_module, hidden_states, *args, _entry=entry, **kwargs): + self._validate_call(_entry, hidden_states, args, kwargs) + self.counters["bf16_linear_calls"] += 1 + return _entry.original_forward(hidden_states) + + def nvfp4_forward(_module, hidden_states, *args, _entry=entry, **kwargs): + self._validate_call(_entry, hidden_states, args, kwargs) + activation = self._quantize_activation(_entry, hidden_states) + self._record_nvfp4_role(_entry) + output = self._linear_plan(_entry)(activation, _entry, _entry.alpha, _entry.module.bias) + return output.reshape(*hidden_states.shape[:-1], _entry.n) + + self._installed_generic[entry.module] = { + "original": entry.original_forward, + "bf16": types.MethodType(bf16_forward, entry.module), + "nvfp4": types.MethodType(nvfp4_forward, entry.module), + } + + def _install_modulations(self, model): + """Own the raw-temb boundary so shared quantization is identity-guarded. + + The Linear inputs themselves are distinct outputs of two identical SiLU + modules, so comparing their Tensor identities would reject valid sharing. + The pinned block forward does pass the exact same raw ``temb`` object to + every img/txt modulation Sequential. Intercepting that boundary both + proves the sharing key and lets C compute SiLU only for the first use. + """ + torch = self.torch + for block_index, block in enumerate(model.transformer_blocks): + for stream in ("img", "txt"): + module = getattr(block, f"{stream}_mod") + name = f"transformer_blocks.{block_index}.{stream}_mod" + entry = self.by_name[f"{name}.1"] + if type(module) is not torch.nn.Sequential or len(module) != 2: + raise TypeError(f"{name} must be the pinned two-node Sequential") + if type(module[0]) is not torch.nn.SiLU or module[0].inplace or module[1] is not entry.module: + raise TypeError(f"{name} must be non-inplace SiLU followed by its exact Linear") + if "forward" in module.__dict__: + raise NotImplementedError(f"{name} already has an instance-level forward override") + original = module.forward + + def bf16_forward( + _module, + temb, + *args, + _name=name, + _entry=entry, + _original=original, + **kwargs, + ): + self._validate_mod_call(_name, _entry, temb, args, kwargs) + self.counters["bf16_linear_calls"] += 1 + return _original(temb) + + def nvfp4_forward( + _module, + temb, + *args, + _name=name, + _entry=entry, + _silu=module[0], + **kwargs, + ): + self._validate_mod_call(_name, _entry, temb, args, kwargs) + activation = self._quantize_modulation(_entry, temb, _silu) + self._record_nvfp4_role(_entry) + output = self._linear_plan(_entry)(activation, _entry, _entry.alpha, _entry.module.bias) + return output.reshape(*temb.shape[:-1], _entry.n) + + self._installed_mod[module] = { + "original": original, + "bf16": types.MethodType(bf16_forward, module), + "nvfp4": types.MethodType(nvfp4_forward, module), + } + + def _install_mlps(self, model): + torch = self.torch + for block_index, block in enumerate(model.transformer_blocks): + for stream in ("img", "txt"): + module = getattr(block, f"{stream}_mlp") + name = f"transformer_blocks.{block_index}.{stream}_mlp" + if module.__class__.__module__ != "diffusers.models.attention" or module.__class__.__name__ != "FeedForward": + raise TypeError(f"{name} is not the pinned Diffusers FeedForward: {type(module)!r}") + if "forward" in module.__dict__: + raise NotImplementedError(f"{name} already has an instance-level forward override") + first = self.by_name[f"transformer_blocks.{block_index}.{stream}_mlp.net.0.proj"] + second = self.by_name[f"transformer_blocks.{block_index}.{stream}_mlp.net.2"] + net = getattr(module, "net", None) + if not isinstance(net, torch.nn.ModuleList) or len(net) != 3: + raise TypeError(f"{name}.net must contain exactly GELU, Dropout, and output projection") + activation, dropout, output = net + if ( + activation.__class__.__module__ != "diffusers.models.activations" + or activation.__class__.__name__ != "GELU" + or getattr(activation, "approximate", None) != "tanh" + or getattr(activation, "proj", None) is not first.module + or type(dropout) is not torch.nn.Dropout + or dropout.p != 0.0 + or dropout.inplace + or output is not second.module + ): + raise TypeError(f"{name} must be the pinned Linear -> GELU(tanh) -> non-inplace Dropout(0) -> Linear") + original = module.forward + + def torch_forward( + _module, + hidden_states, + *args, + _name=name, + _original=original, + **kwargs, + ): + self._validate_mlp_call(_name, hidden_states, args, kwargs) + self.counters["mlp_calls"]["torch"] += 1 + self.counters["bf16_linear_calls"] += 2 + return _original(hidden_states) + + def cudnn_forward( + _module, + hidden_states, + *args, + _name=name, + _first=first, + _second=second, + **kwargs, + ): + self._validate_mlp_call(_name, hidden_states, args, kwargs) + self.counters["mlp_calls"]["cudnn_bf16"] += 1 + self.counters["bf16_linear_calls"] += 2 + return self.gelu_mlp( + hidden_states, + _first.module.weight, + _first.module.bias, + _second.module.weight, + _second.module.bias, + ) + + def nvfp4_forward( + _module, + hidden_states, + *args, + _name=name, + _first=first, + _second=second, + **kwargs, + ): + self._validate_mlp_call(_name, hidden_states, args, kwargs) + self.counters["mlp_calls"]["nvfp4"] += 1 + activation = self._quantize_activation(_first, hidden_states) + self._record_nvfp4_role(_first) + hidden = self._fused_fc1_plan(_first)( + activation, + _first, + _first.alpha, + _first.module.bias, + ) + self.counters["activation_quant_logical"] += 1 + self.counters["activation_quant_physical"] += 1 + self.counters["activation_quant_fused"] += 1 + self._record_nvfp4_role(_second) + output = self._linear_plan(_second)(hidden, _second, _second.alpha, _second.module.bias) + return output.reshape(*hidden_states.shape[:-1], _second.n) + + self._installed_mlp[module] = { + "original": original, + "torch": types.MethodType(torch_forward, module), + "cudnn_bf16": types.MethodType(cudnn_forward, module), + "nvfp4": types.MethodType(nvfp4_forward, module), + } + + def _validate_call(self, entry, hidden_states, args, kwargs): + if not self._active: + raise RuntimeError(f"{entry.qualified_name} called outside adapter.forward_scope()") + if args or kwargs: + raise NotImplementedError(f"{entry.qualified_name} accepts only hidden_states in this benchmark") + if tuple(hidden_states.shape) != entry.input_shape: + raise ValueError(f"{entry.qualified_name} expected {entry.input_shape}, got {tuple(hidden_states.shape)}") + if hidden_states.dtype != self.torch.bfloat16 or hidden_states.device.type != "cuda" or not hidden_states.is_contiguous(): + raise NotImplementedError(f"{entry.qualified_name} requires contiguous bf16 CUDA activation") + self._validate_stream(hidden_states.device) + + def _validate_mlp_call(self, name, hidden_states, args, kwargs): + if not self._active: + raise RuntimeError(f"{name} called outside adapter.forward_scope()") + if args or kwargs: + raise NotImplementedError(f"{name} accepts only hidden_states in this benchmark") + if hidden_states.dtype != self.torch.bfloat16 or hidden_states.device.type != "cuda" or not hidden_states.is_contiguous(): + raise NotImplementedError(f"{name} requires contiguous bf16 CUDA activation") + self._validate_stream(hidden_states.device) + + def _validate_mod_call(self, name, entry, temb, args, kwargs): + if not self._active: + raise RuntimeError(f"{name} called outside adapter.forward_scope()") + if args or kwargs: + raise NotImplementedError(f"{name} accepts only temb in this benchmark") + if tuple(temb.shape) != entry.input_shape: + raise ValueError(f"{name} expected raw temb shape {entry.input_shape}, got {tuple(temb.shape)}") + if temb.dtype != self.torch.bfloat16 or temb.device.type != "cuda" or not temb.is_contiguous(): + raise NotImplementedError(f"{name} requires contiguous bf16 CUDA temb") + self._validate_stream(temb.device) + + def _validate_stream(self, device): + current = self.torch.cuda.current_stream(device).cuda_stream + if device != self._device or current != self._stream: + raise NotImplementedError( + "the benchmark-local NVFP4 adapter is prepared for exactly one (device, stream); " + f"prepared={(self._device, self._stream)}, current={(device, current)}" + ) + + def _linear_plan(self, entry): + return self._linear_plans[(self._device.index, self._stream, entry.m, entry.n, entry.k, "linear_bias")] + + def _fused_fc1_plan(self, entry): + return self._fused_fc1_plans[ + ( + self._device.index, + self._stream, + entry.m, + entry.n, + entry.k, + "linear_bias_gelu_nvfp4", + ) + ] + + def _validate_prepared_bindings(self): + """Full stable-binding preflight, called outside the timed forward.""" + self._validate_stream(self._device) + for entry in self.entries: + if entry.role.endswith("net.0.proj"): + second_name = entry.qualified_name.removesuffix("net.0.proj") + "net.2" + second = self.by_name.get(second_name) + if second is None: + raise RuntimeError(f"cannot validate fused FC1 without {second_name}") + self._fused_fc1_plan(entry).validate_prepared( + self._activation_buffers[(entry.m, entry.k)], + entry, + second.activation_global_scale, + ) + elif entry.role.endswith("net.2"): + first_name = entry.qualified_name.removesuffix("net.2") + "net.0.proj" + first = self.by_name.get(first_name) + if first is None: + raise RuntimeError(f"cannot validate FC2 without {first_name}") + first_plan = self._fused_fc1_plan(first) + self._linear_plan(entry).validate_prepared((first_plan.qh, first_plan.sh), entry) + else: + self._linear_plan(entry).validate_prepared(self._activation_buffers[(entry.m, entry.k)], entry) + + def _quantize_modulation(self, entry, temb, silu): + key = (entry.m, entry.k) + cached = self._activation_cache.get(key) + if cached is not None and cached.group == entry.activation_group: + # Unlike equal amax values, this proves every reuse came from the + # exact raw conditioning Tensor in the pinned call graph. + if cached.source is not temb: + raise RuntimeError("Qwen-Image modulation stopped sharing the exact raw temb Tensor") + if cached.global_scale is not entry.activation_global_scale: + raise RuntimeError("Qwen-Image modulation reused a different frozen global-scale tensor") + self.counters["activation_quant_logical"] += 1 + self.counters["activation_cache_hits"] += 1 + return cached + activated = silu(temb) + self._validate_call(entry, activated, (), {}) + return self._quantize_activation(entry, activated, cache_source=temb) + + def _quantize_activation(self, entry, hidden_states, *, cache_source=None): + self.counters["activation_quant_logical"] += 1 + source = hidden_states if cache_source is None else cache_source + key = (entry.m, entry.k) + cached = self._activation_cache.get(key) + if cached is not None and cached.group == entry.activation_group: + if cached.global_scale is not entry.activation_global_scale: + raise RuntimeError(f"{entry.activation_group} reused with a different frozen global-scale tensor") + if cached.source is not source: + raise RuntimeError(f"{entry.activation_group} inputs stopped sharing the exact Tensor object") + self.counters["activation_cache_hits"] += 1 + return cached + + packed, scale_factors = self._activation_buffers[key] + result_packed, result_scale = self.nvfp4_quantize( + hidden_states.view(entry.m, entry.k), + entry.activation_global_scale, + pre_quant_scale=None, + out=packed, + scale_factors=scale_factors, + enable_pdl=True, + ) + if result_packed.data_ptr() != packed.data_ptr() or result_scale.data_ptr() != scale_factors.data_ptr(): + raise RuntimeError("nvfp4_quantize ignored the caller-provided activation output buffers") + self._validate_quantized_buffers( + result_packed, + result_scale, + entry.m, + entry.k, + f"{entry.qualified_name} activation", + ) + result = _QuantizedActivation( + packed, + scale_factors, + entry.activation_global_scale, + source, + entry.activation_group, + ) + self._activation_cache[key] = result + self.counters["activation_quant_physical"] += 1 + self.counters["activation_quant_standalone"] += 1 + return result + + def _record_nvfp4_role(self, entry): + self.counters["nvfp4_linear_calls"] += 1 + self.counters["nvfp4_linear_by_role"][entry.qualified_name] += 1 + self._active_role_order.append(entry.qualified_name) + + def select(self, arm): + if arm not in ARM_CONFIGS: + raise ValueError(f"unknown ModelOpt NVFP4 benchmark arm {arm!r}") + if self._active: + raise RuntimeError("cannot switch treatments during a model forward") + if arm == "C": + self._validate_prepared_bindings() + valid_generic = {value for forwards in self._installed_generic.values() for key, value in forwards.items() if key != "original"} + valid_mod = {value for forwards in self._installed_mod.values() for key, value in forwards.items() if key != "original"} + valid_mlp = {value for forwards in self._installed_mlp.values() for key, value in forwards.items() if key != "original"} + if any(module.forward not in valid_generic and "forward" in module.__dict__ for module in self._installed_generic): + raise RuntimeError("a Qwen-Image generic Linear was modified after adapter installation") + if any(module.forward not in valid_mlp and "forward" in module.__dict__ for module in self._installed_mlp): + raise RuntimeError("a Qwen-Image FeedForward was modified after adapter installation") + if any(module.forward not in valid_mod and "forward" in module.__dict__ for module in self._installed_mod): + raise RuntimeError("a Qwen-Image modulation Sequential was modified after adapter installation") + config = ARM_CONFIGS[arm] + for module, forwards in self._installed_generic.items(): + module.forward = forwards[config["generic_linear"]] + for module, forwards in self._installed_mlp.items(): + module.forward = forwards[config["mlp"]] + for module, forwards in self._installed_mod.items(): + module.forward = forwards[config["generic_linear"]] + self._selected = arm + + @contextmanager + def forward_scope(self, arm): + if arm != self._selected: + raise RuntimeError(f"forward scope arm {arm} does not match selected treatment {self._selected}") + if self._active: + raise RuntimeError("nested model forwards are outside the benchmark treatment") + before = self.snapshot() + self._active = True + self._active_role_order = [] + self._activation_cache = {} + self.counters["forward_scopes"] += 1 + completed = False + try: + yield + completed = True + finally: + self._active = False + self._activation_cache = {} + if completed: + delta = counter_delta(self.snapshot(), before) + expected = expected_route_delta(arm, self.layers) + if delta != expected: + raise RuntimeError(f"{arm} adapter route mismatch: delta={delta}, expected={expected}") + if arm == "C": + expected_order = [f"transformer_blocks.{block}.{role}" for block in range(self.layers) for role in ROLE_ORDER] + if self._active_role_order != expected_order: + raise RuntimeError(f"NVFP4 Linear call order changed: got {self._active_role_order}, expected {expected_order}") + + def snapshot(self): + return copy.deepcopy(self.counters) + + def metadata(self): + return { + "recipe": copy.deepcopy(MODELOPT_RECIPE), + "representative_full_blocks": representative_middle_blocks(self.layers), + "role_order": list(ROLE_ORDER), + "role_shapes": { + entry.qualified_name: { + "input": list(entry.input_shape), + "weight": list(entry.weight_shape), + "gemm_mkn": [entry.m, entry.k, entry.n], + "activation_group": entry.activation_group, + "activation_amax": float(entry.activation_amax.item()), + "weight_amax": float(entry.weight_amax.item()), + "activation_global_scale": float(entry.activation_global_scale.item()), + "weight_global_scale": float(entry.weight_global_scale.item()), + "alpha": float(entry.alpha.item()), + } + for entry in self.entries + }, + "calibration": copy.deepcopy(self.calibration_metadata), + "setup_counts": { + "weight_pack_calls": 14 * self.layers, + "plan_build_calls": 7, + }, + "binding_mode": copy.deepcopy(self.binding_mode), + "plan_provenance": self.plan_provenance(), + "expected_c_forward": expected_route_delta("C", self.layers), + } + + def restore(self): + if self._active: + raise RuntimeError("cannot restore adapter during a model forward") + for module, forwards in self._installed_generic.items(): + if module.forward in forwards.values(): + module.__dict__.pop("forward", None) + for module, forwards in self._installed_mlp.items(): + if module.forward in forwards.values(): + module.__dict__.pop("forward", None) + for module, forwards in self._installed_mod.items(): + if module.forward in forwards.values(): + module.__dict__.pop("forward", None) + self._selected = None + + +def install_modelopt_nvfp4_dispatch(model, shape, calibration): + return QwenImageModelOptNvfp4Adapter(model, shape, calibration) diff --git a/benchmark/e2e/Qwen-Image/run_nvfp4.py b/benchmark/e2e/Qwen-Image/run_nvfp4.py new file mode 100644 index 000000000..db9f0df9b --- /dev/null +++ b/benchmark/e2e/Qwen-Image/run_nvfp4.py @@ -0,0 +1,681 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Balanced Qwen-Image BF16/cuDNN/ModelOpt-NVFP4 three-arm benchmark. + +The low-precision arm aligns its Linear placement and NVFP4/max policy with +ModelOpt 0.46.0's Qwen-Image recipe: all fourteen Linear roles are NVFP4 while +the attention core stays BF16 (``quantize_mha=False``). Explicit BF16, +synthetic-calibration, and depth-reduction overrides make this performance +evidence, not an exact ModelOpt state or image-quality claim. +""" + +from __future__ import annotations + +import argparse +from datetime import datetime, timezone +import hashlib +import importlib +import importlib.util +import json +import math +import os +from pathlib import Path +import platform +import statistics +import subprocess +import sys +import time + +THIS_DIR = Path(__file__).resolve().parent +E2E_DIR = THIS_DIR.parent +REPO_ROOT = THIS_DIR.parents[2] +MODEL_PATH = THIS_DIR / "run_model.py" +ADAPTER_PATH = THIS_DIR / "modelopt_nvfp4.py" +FACTORIAL_PATH = E2E_DIR / "_factorial.py" +sys.path.insert(0, str(E2E_DIR)) + +from _factorial import config_fingerprint, paired_stats, percentile # noqa: E402 + +PROTOCOL_DEFAULTS = { + "smoke": {"warmup": 1, "rounds": 6, "repeats": 1}, + "formal": {"warmup": 3, "rounds": 42, "repeats": 3}, +} + +ARMS = ( + { + "id": "A", + "name": "bf16_off", + "linears": "Torch BF16", + "mlp": "pinned Diffusers Torch GELU FFN", + "attention": "forced PyTorch Flash BF16", + "attention_route": "torch_flash", + }, + { + "id": "B", + "name": "bf16_cudnn", + "linears": "Torch BF16 outside FFN", + "mlp": "cudnn.gemm.ops.gelu_mlp BF16", + "attention": "cuDNN BF16", + "attention_route": "cudnn", + }, + { + "id": "C", + "name": "modelopt046_nvfp4_cudnn", + "linears": "cuDNN FROST NVFP4, all 14 roles", + "mlp": "NVFP4 FC1 + BF16 bias/GELU + fused NVFP4 hidden requant + NVFP4 FC2", + "attention": "cuDNN BF16 (ModelOpt quantize_mha=False)", + "attention_route": "cudnn", + }, +) + + +def _utc_now(): + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +def _sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as source: + for chunk in iter(lambda: source.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _source_record(path): + path = Path(path).resolve() + try: + display = str(path.relative_to(REPO_ROOT)) + except ValueError: + display = str(path) + return {"path": display, "sha256": _sha256(path)} + + +def _load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load module from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _git_provenance(): + def run(*arguments): + return subprocess.run(arguments, cwd=REPO_ROOT, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True).stdout.strip() + + return { + "commit": run("git", "rev-parse", "HEAD"), + "branch": run("git", "branch", "--show-current") or "detached", + "dirty": bool(run("git", "status", "--porcelain")), + } + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=PROTOCOL_DEFAULTS, default="smoke") + parser.add_argument("--layers", type=int) + parser.add_argument("--image-tokens", type=int) + parser.add_argument("--text-tokens", type=int) + parser.add_argument("--warmup", type=int) + parser.add_argument("--rounds", type=int) + parser.add_argument("--repeats", type=int) + parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "qwen-image-nvfp4-results") + parser.add_argument("--tag", default="") + parser.add_argument("--compare", type=Path) + return parser.parse_args() + + +def _resolve_protocol(args): + protocol = dict(PROTOCOL_DEFAULTS[args.mode]) + for name in protocol: + value = getattr(args, name) + if value is not None: + protocol[name] = value + if any(not isinstance(value, int) or value <= 0 for value in protocol.values()): + raise ValueError(f"protocol values must be positive integers, got {protocol}") + if protocol["rounds"] % 6: + raise ValueError("rounds must be a multiple of 6 for the balanced three-treatment design") + return protocol + + +def _pick_device(torch, mode): + candidates = [] + for index in range(torch.cuda.device_count()): + properties = torch.cuda.get_device_properties(index) + candidates.append(f"cuda:{index}={properties.name}/{properties.multi_processor_count}SM") + if (properties.major, properties.minor) == (10, 0) and ( + mode != "formal" or (properties.name == "NVIDIA B200" and properties.multi_processor_count == 148) + ): + return torch.device(f"cuda:{index}"), properties + requirement = "a full 148-SM NVIDIA B200" if mode == "formal" else "an SM100 GPU" + raise RuntimeError(f"{mode} mode requires {requirement}; visible: {', '.join(candidates)}") + + +def _rel_l2(actual, expected): + return float((actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1e-12)) + + +def _scale_counter_tree(value, factor): + return {key: _scale_counter_tree(item, factor) if isinstance(item, dict) else int(item) * factor for key, item in value.items()} + + +def _add_counter_trees(*values): + if not values: + return {} + keys = set(values[0]) + if any(set(value) != keys for value in values[1:]): + raise ValueError("counter trees have different keys") + result = {} + for key in values[0]: + items = [value[key] for value in values] + result[key] = _add_counter_trees(*items) if isinstance(items[0], dict) else sum(int(item) for item in items) + return result + + +def _load_comparison(path): + def reject(value): + raise ValueError(f"non-finite JSON constant {value}") + + return json.loads(Path(path).read_text(encoding="utf-8"), parse_constant=reject) + + +def _compare(current, previous): + current_fp = current["config"]["comparability_fingerprint"]["sha256"] + previous_fp = previous["config"]["comparability_fingerprint"]["sha256"] + if current_fp != previous_fp: + raise ValueError(f"comparison fingerprint mismatch: current={current_fp}, previous={previous_fp}") + arms = {} + for arm in ARMS: + key = arm["id"] + new = float(current["summary"][key]["p50_ms"]) + old = float(previous["summary"][key]["p50_ms"]) + if not all(math.isfinite(value) and value > 0 for value in (new, old)): + raise ValueError("comparison p50 values must be finite and positive") + arms[key] = { + "previous_p50_ms": old, + "current_p50_ms": new, + "change_ms": new - old, + "change_percent": (new / old - 1.0) * 100.0, + } + return {"paired_across_runs": False, "arms": arms} + + +def _format_elapsed_effect(ratio): + """Describe a paired elapsed-time ratio without calling regressions wins.""" + ratio = float(ratio) + if not math.isfinite(ratio) or ratio <= 0: + raise ValueError(f"elapsed-time ratio must be finite and positive, got {ratio!r}") + if ratio < 1: + return f"{1 / ratio:.3f}x speedup ({(1 - ratio) * 100:.2f}% lower elapsed time)" + if ratio > 1: + return f"{ratio:.3f}x slower ({(ratio - 1) * 100:.2f}% higher elapsed time)" + return "1.000x (no elapsed-time change)" + + +def _render_markdown(metadata, raw_name, raw_hash): + config = metadata["config"] + comparisons = metadata["comparisons"] + smoke = config["mode"] == "smoke" + lines = [ + "# Qwen-Image ModelOpt 0.46 NVFP4/max benchmark", + "", + f"Generated: `{metadata['completed_utc']}` ", + f"Mode: `{config['mode']}` ", + f"Comparability fingerprint: `{config['comparability_fingerprint']['sha256']}` ", + f"Build/provenance fingerprint: `{config['build_fingerprint']['sha256']}` ", + f"Raw JSON: [`{raw_name}`]({raw_name}) (`sha256:{raw_hash}`)", + "", + ] + if smoke: + lines += [ + "## Smoke validation", + "", + "**Validation only; these reduced-token timings are not a performance headline.**", + "", + ] + else: + ba = comparisons["B_vs_A"] + cb = comparisons["C_vs_B"] + ca = comparisons["C_vs_A"] + lines += [ + "## Result", + "", + f"- BF16 cuDNN effect (B/A): {_format_elapsed_effect(ba['paired_ratio_p50'])}.", + f"- ModelOpt NVFP4 increment (C/B): {_format_elapsed_effect(cb['paired_ratio_p50'])}.", + f"- Total cuDNN platform impact (C/A): {_format_elapsed_effect(ca['paired_ratio_p50'])}.", + "", + ] + lines += [ + "| arm | linears | GELU FFN | attention | p10 | p50 | p90 | paired ratio vs A |", + "|---|---|---|---|---:|---:|---:|---:|", + ] + for arm in ARMS: + value = metadata["summary"][arm["id"]] + lines.append( + f"| `{arm['id']}` | {arm['linears']} | {arm['mlp']} | {arm['attention']} | " + f"{value['p10_ms']:.3f} ms | {value['p50_ms']:.3f} ms | {value['p90_ms']:.3f} ms | " + f"{value['paired_ratio_p50']:.5f} |" + ) + if not smoke: + lines += [ + "", + "## Paired contrasts", + "", + "| contrast | meaning | paired ratio (p50) | elapsed-time effect | delta (p50) | wins |", + "|---|---|---:|---:|---:|---:|", + ] + labels = { + "B_vs_A": "fixed-precision cuDNN MLP + attention effect", + "C_vs_B": "NVFP4 recipe increment; precision and expanded cuDNN Linear coverage", + "C_vs_A": "total cuDNN platform impact", + } + for key, label in labels.items(): + value = comparisons[key] + lines.append( + f"| {key.replace('_vs_', '/')} | {label} | {value['paired_ratio_p50']:.5f} | " + f"{_format_elapsed_effect(value['paired_ratio_p50'])} | {value['paired_delta_p50_ms']:+.3f} ms | " + f"{value['wins']}/{value['batches']} |" + ) + shape = config["shape"] + recipe = config["numerical_recipe"] + route = metadata["route"] + correctness = metadata["correctness"]["model_output_rel_l2_vs_A"] + contract_gate = metadata["correctness"].get("low_precision_contract_gate") + lines += [ + "", + "## Scope and gates", + "", + f"- Shape: B={shape['bs']}, image/text/joint tokens={shape['image_tokens']}/{shape['text_tokens']}/{shape['joint_tokens']}, " + f"H={shape['hidden']}, heads={shape['heads']}x{shape['head_dim']}, FFN={shape['ffn']}, repeated layers={shape['layers']}/60.", + f"- Proxy block mapping: `{config['representative_full_blocks']}` within ModelOpt's quantized full-model range 2..57.", + f"- Recipe placement/format anchor: `{recipe['id']}` at ModelOpt `{recipe['release']}@{recipe['commit']}`; " + "all fourteen Linear roles use NVFP4 E2M1/block-16 with E4M3 block scales.", + f"- Proxy overrides: `{json.dumps(recipe['proxy_overrides'], sort_keys=True)}`. This is not an exact upstream dtype, calibration-state, or workload reproduction.", + "- Attention remains BF16 in C. ModelOpt 0.46 `quantize_mha` defaults to false; this result is not MXFP8 or per-tensor FP8 attention.", + "- Calibration is one deterministic synthetic BF16 max pass and is frozen before timing. Random weights/inputs make this quality-ineligible.", + "- B/A isolates the existing BF16 cuDNN treatments. C/B is intentionally not a kernel-only contrast: it changes precision and moves all " + "fourteen block Linear roles onto cuDNN FROST.", + "- A turns cuDNN off only for the measured FFN/attention treatments; unrelated framework operators are unchanged.", + f"- PyTorch SDPA probe: natural `{route['torch_probe']['natural_choice_name']}`, A forced `{route['torch_probe']['forced_choice_name']}`.", + f"- Model output relative L2 vs A: `{json.dumps(correctness, sort_keys=True)}`. C is recorded as a finite sanity signal, not a quality gate.", + f"- C per-forward activation quantizations: logical `{route['expected_C_per_forward']['activation_quant_logical']}`, " + f"physical `{route['expected_C_per_forward']['activation_quant_physical']}` " + f"(`{route['expected_C_per_forward']['activation_quant_standalone']}` standalone + " + f"`{route['expected_C_per_forward']['activation_quant_fused']}` fused).", + "", + "## Provenance", + "", + "| source | path | sha256 |", + "|---|---|---|", + ] + if contract_gate is not None: + lines.insert( + lines.index("## Provenance") - 1, + f"- Low-precision implementation contract gate: `{contract_gate['status']}` across " + f"`{contract_gate['contracts_checked']}` exact plan shapes/epilogues; this is a kernel-contract gate, not an image-quality claim.", + ) + for name, source in sorted(metadata["provenance"]["sources"].items()): + lines.append(f"| {name} | `{source['path']}` | `{source['sha256']}` |") + comparison = metadata.get("comparison_across_runs") + if comparison is not None: + lines += [ + "", + "## Cross-run comparison", + "", + "**Not paired across runs.** Each row compares independent p50 estimates.", + "", + "| arm | previous p50 | current p50 | non-paired change | change |", + "|---|---:|---:|---:|---:|", + ] + for arm, value in sorted(comparison["arms"].items()): + lines.append( + f"| `{arm}` | {value['previous_p50_ms']:.3f} ms | {value['current_p50_ms']:.3f} ms | " + f"{value['change_ms']:+.3f} ms | {value['change_percent']:+.2f}% |" + ) + lines += [ + "", + "This is a random-weight, depth-reduced transformer-backbone proxy: no text encoder, VAE, scheduler, denoising loop, or image-quality claim.", + "", + ] + return "\n".join(lines) + + +def main(): + started_utc = _utc_now() + args = _parse_args() + protocol = _resolve_protocol(args) + model_api = _load_module("qwen_image_model_nvfp4", MODEL_PATH) + lowp = _load_module("qwen_image_modelopt_nvfp4", ADAPTER_PATH) + shape = model_api.resolve_shape(args.mode, layers=args.layers, image_tokens=args.image_tokens, text_tokens=args.text_tokens) + representative = lowp.representative_middle_blocks(shape["layers"]) + orders = lowp.three_arm_orders() + + # The low-precision adapter calls FROST directly. Keeping the global engine + # opt-in disabled prevents arm B's public BF16 GELU op from acquiring a + # different plan population than the existing BF16 benchmark. + if os.environ.get("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "0").lower() in ("1", "true", "yes", "on"): + raise RuntimeError("disable global FROST engines; this runner invokes only its private NVFP4 FROST plans directly") + + torch, _, cudnn, sdpamod, diffusers, qwen_module = model_api.load_runtime() + diffusers_attention_module = importlib.import_module("diffusers.models.attention") + diffusers_activations_module = importlib.import_module("diffusers.models.activations") + loaded_diffusers_source = _source_record(qwen_module.__file__) + if loaded_diffusers_source["sha256"] != model_api.DIFFUSERS_ANCHOR["source_sha256"]: + raise RuntimeError( + "loaded Diffusers Qwen-Image source does not match the pin: " + f"got {loaded_diffusers_source['sha256']}, expected {model_api.DIFFUSERS_ANCHOR['source_sha256']}" + ) + loaded_supporting_sources = { + "attention": _source_record(diffusers_attention_module.__file__), + "activations": _source_record(diffusers_activations_module.__file__), + } + for name, source in loaded_supporting_sources.items(): + expected = model_api.DIFFUSERS_ANCHOR["supporting_sources"][name]["source_sha256"] + if source["sha256"] != expected: + raise RuntimeError(f"loaded Diffusers {name} source does not match the pin: got {source['sha256']}, expected {expected}") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not hasattr(torch, "float4_e2m1fn_x2"): + raise RuntimeError("this PyTorch build has no packed float4_e2m1fn_x2 dtype") + + device, properties = _pick_device(torch, args.mode) + with torch.cuda.device(device): + if any(hasattr(sdpamod, name) for name in ("sdpa_fwd_d256", "sdpa_bwd_d256")): + raise RuntimeError("loaded FE SDPA predates backend-only #682") + padding_check = None + # Use the same focused mask coverage as the BF16 leaf without importing + # or changing it: dense formal timings remain mask-free. + text_tokens, image_tokens, heads, head_dim, batch = 64, 128, 4, 128, 2 + generator = torch.Generator(device=device).manual_seed(2026) + qkv_shape = (batch, text_tokens + image_tokens, heads, head_dim) + q, k, v = (torch.randn(qkv_shape, device=device, dtype=torch.bfloat16, generator=generator) for _ in range(3)) + text_mask = torch.arange(text_tokens, device=device).unsqueeze(0) < torch.tensor([64, 37], device=device).unsqueeze(1) + mask = torch.cat([text_mask, torch.ones(batch, image_tokens, dtype=torch.bool, device=device)], dim=1)[:, None, None] + pad_select, pad_restore, _, _ = model_api.install_joint_attention_dispatch(qwen_module, text_tokens=text_tokens) + try: + with torch.inference_mode(): + pad_select("torch_reference") + expected = qwen_module.dispatch_attention_fn(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) + pad_select("cudnn") + actual = qwen_module.dispatch_attention_fn(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) + torch.cuda.synchronize(device) + finally: + pad_restore() + padding_rel_l2 = _rel_l2(actual, expected) + if not math.isfinite(padding_rel_l2) or padding_rel_l2 > 0.01: + raise RuntimeError(f"right-padded joint-attention adapter mismatch: rel_l2={padding_rel_l2}") + padding_check = {"shape": list(qkv_shape), "text_valid_lengths": [64, 37], "rel_l2": padding_rel_l2} + + model = model_api.build_model(torch, qwen_module, device, layers=shape["layers"]) + inputs = model_api.make_inputs(torch, device, shape) + calibration = lowp.collect_max_calibration(model, shape, lambda: model_api.forward(model, inputs)) + torch.cuda.synchronize(device) + + attention_calls = {"torch_flash": 0, "cudnn": 0} + torch_probe = {} + select_attention, restore_attention, attention_calls, torch_probe = model_api.install_joint_attention_dispatch( + qwen_module, text_tokens=shape["text_tokens"], counters=attention_calls, torch_probe=torch_probe + ) + adapter = lowp.install_modelopt_nvfp4_dispatch(model, shape, calibration) + numerical_gate = adapter.run_focused_numerical_gate() + + arm_by_id = {arm["id"]: arm for arm in ARMS} + + def configure(arm_id): + adapter.select(arm_id) + select_attention(arm_by_id[arm_id]["attention_route"]) + + def invoke(arm_id): + with adapter.forward_scope(arm_id): + with torch.inference_mode(): + return model_api.forward(model, inputs) + + try: + for arm in ARMS: + arm_id = arm["id"] + configure(arm_id) + adapter_before = adapter.snapshot() + attention_before = dict(attention_calls) + for _ in range(protocol["warmup"]): + invoke(arm_id) + torch.cuda.synchronize(device) + adapter_delta = lowp.counter_delta(adapter.snapshot(), adapter_before) + expected_adapter = _scale_counter_tree(lowp.expected_route_delta(arm_id, shape["layers"]), protocol["warmup"]) + if adapter_delta != expected_adapter: + raise RuntimeError(f"{arm_id} warm adapter route mismatch: delta={adapter_delta}, expected={expected_adapter}") + attention_delta = {name: attention_calls[name] - attention_before[name] for name in attention_calls} + expected_attention = protocol["warmup"] * shape["layers"] + expected_attention_delta = { + "torch_reference": 0, + "torch_flash": expected_attention if arm["attention_route"] == "torch_flash" else 0, + "cudnn": expected_attention if arm["attention_route"] == "cudnn" else 0, + } + if attention_delta != expected_attention_delta: + raise RuntimeError(f"{arm_id} warm attention route mismatch: delta={attention_delta}, expected={expected_attention_delta}") + + outputs = {} + for arm in ARMS: + configure(arm["id"]) + outputs[arm["id"]] = invoke(arm["id"]) + torch.cuda.synchronize(device) + if not bool(torch.isfinite(outputs[arm["id"]]).all()): + raise RuntimeError(f"non-finite model output in arm {arm['id']}") + model_rel_l2 = {arm["id"]: _rel_l2(outputs[arm["id"]], outputs["A"]) for arm in ARMS} + if model_rel_l2["B"] > 0.02: + raise RuntimeError(f"BF16 cuDNN arm B mismatch: rel_l2={model_rel_l2['B']}") + + raw = {arm["id"]: [] for arm in ARMS} + batches = {arm["id"]: [] for arm in ARMS} + timing_started = time.time() + for batch_index in range(protocol["rounds"]): + for arm_index in orders[batch_index % len(orders)]: + arm_id = ARMS[arm_index]["id"] + samples = [] + for _ in range(protocol["repeats"]): + configure(arm_id) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + output = invoke(arm_id) + end.record() + end.synchronize() + elapsed = start.elapsed_time(end) + if not math.isfinite(elapsed) or elapsed <= 0: + raise RuntimeError(f"invalid timing {elapsed}") + samples.append(elapsed) + raw[arm_id].append(samples) + batches[arm_id].append(statistics.median(samples)) + print(f"BATCH {batch_index + 1}/{protocol['rounds']} elapsed_s={time.time() - timing_started:.1f}", flush=True) + + summary = {} + for arm in ARMS: + arm_id = arm["id"] + values = batches[arm_id] + summary[arm_id] = { + "p10_ms": percentile(values, 0.1), + "p50_ms": percentile(values, 0.5), + "p90_ms": percentile(values, 0.9), + "mean_ms": statistics.mean(values), + "batches": len(values), + **paired_stats(values, batches["A"]), + } + comparisons = { + "B_vs_A": paired_stats(batches["B"], batches["A"]), + "C_vs_B": paired_stats(batches["C"], batches["B"]), + "C_vs_A": paired_stats(batches["C"], batches["A"]), + } + + calls_per_arm = protocol["warmup"] + 1 + protocol["rounds"] * protocol["repeats"] + expected_adapter_totals = _add_counter_trees( + *[_scale_counter_tree(lowp.expected_route_delta(arm["id"], shape["layers"]), calls_per_arm) for arm in ARMS] + ) + expected_adapter_totals["weight_pack_calls"] = 14 * shape["layers"] + expected_adapter_totals["plan_build_calls"] = 7 + if adapter.snapshot() != expected_adapter_totals: + raise RuntimeError(f"final adapter route mismatch: got={adapter.snapshot()}, expected={expected_adapter_totals}") + expected_attention_calls = { + "torch_reference": 0, + "torch_flash": calls_per_arm * shape["layers"], + "cudnn": 2 * calls_per_arm * shape["layers"], + } + if attention_calls != expected_attention_calls: + raise RuntimeError(f"final attention route mismatch: got={attention_calls}, expected={expected_attention_calls}") + if torch_probe.get("forced_choice_name") != "FLASH_ATTENTION": + raise RuntimeError(f"forced PyTorch FlashAttention treatment route changed: {torch_probe}") + finally: + adapter.restore() + restore_attention() + + quant_module = importlib.import_module("cudnn.gemm.ops._nvfp4_quantize") + gelu_module = importlib.import_module("cudnn.gemm.ops._gelu_mlp") + frost_module_names = ( + "compiler", + "dtypes", + "epilogue_codegen", + "fusion_ir", + "graph_analyzer", + "kernel_registry", + "recipe", + "tile_config", + ) + frost_modules = {name: importlib.import_module(f"cudnn.gemm.frost.{name}") for name in frost_module_names} + quant_csrc = Path(quant_module.__file__).with_name("csrc") + sources = { + "runner": _source_record(Path(__file__)), + "model_adapter": _source_record(MODEL_PATH), + "modelopt_adapter": _source_record(ADAPTER_PATH), + "statistics": _source_record(FACTORIAL_PATH), + "diffusers_qwen_image": loaded_diffusers_source, + "diffusers_attention": loaded_supporting_sources["attention"], + "diffusers_activations": loaded_supporting_sources["activations"], + "cudnn_sdpa": _source_record(sdpamod.__file__), + "cudnn_gelu_mlp": _source_record(gelu_module.__file__), + "cudnn_nvfp4_quantize": _source_record(quant_module.__file__), + "cudnn_nvfp4_quantize_cuda": _source_record(quant_csrc / "nvfp4_quantize_sm100.cu"), + "cudnn_nvfp4_quantize_header": _source_record(quant_csrc / "nvfp4_smooth_quantize_sm100.cuh"), + } + sources.update({f"cudnn_frost_{name}": _source_record(module.__file__) for name, module in frost_modules.items()}) + for index, plan in enumerate(adapter.plan_provenance()): + sources[f"cudnn_frost_generated_plan_{index:02d}"] = _source_record(plan["generated_path"]) + config = { + "schema_version": 1, + "mode": args.mode, + "timing_role": "validation_only" if args.mode == "smoke" else "formal_performance", + "performance_claim_eligible": args.mode == "formal", + "numerical_claim_eligible": False, + "device": properties.name, + "device_id": str(device), + "sm_arch": f"sm_{properties.major}{properties.minor}", + "sm_count": properties.multi_processor_count, + "python": platform.python_version(), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda or "unknown", + "cudnn_frontend": cudnn.__version__, + "cudnn_backend": cudnn.backend_version(), + "diffusers": getattr(diffusers, "__version__", "unknown"), + "shape": shape, + "representative_full_blocks": representative, + "protocol": protocol, + "workload": "single_conditional_transformer_forward_no_checkpoint", + "weights": "random_init_seed_0; bf16 originals retained; NVFP4 weights prepacked before timing", + "inputs": "random_latents_and_precomputed_text_embeddings_seed_1234", + "numerical_recipe": lowp.MODELOPT_RECIPE, + "model_anchor": dict(model_api.OFFICIAL_MODEL), + "diffusers_anchor": dict(model_api.DIFFUSERS_ANCHOR), + "arms": [{key: value for key, value in arm.items() if key != "attention_route"} for arm in ARMS], + "balanced_orders": orders, + } + comparable_keys = ( + "schema_version", + "mode", + "timing_role", + "performance_claim_eligible", + "numerical_claim_eligible", + "device", + "sm_arch", + "sm_count", + "python", + "torch", + "torch_cuda", + "cudnn_backend", + "diffusers", + "shape", + "representative_full_blocks", + "protocol", + "workload", + "weights", + "inputs", + "numerical_recipe", + "model_anchor", + "diffusers_anchor", + "arms", + "balanced_orders", + ) + comparable = {key: config[key] for key in comparable_keys} + build = {"schema_version": 1, "git": _git_provenance(), "sources": {name: value["sha256"] for name, value in sorted(sources.items())}} + config["comparability_fingerprint"] = {"inputs": comparable, "sha256": config_fingerprint(comparable)} + config["build_fingerprint"] = {"inputs": build, "sha256": config_fingerprint(build)} + metadata = { + "schema_version": 1, + "started_utc": started_utc, + "completed_utc": _utc_now(), + "arguments": {name: str(value) if isinstance(value, Path) else value for name, value in vars(args).items()}, + "config": config, + "correctness": { + "model_output_rel_l2_vs_A": model_rel_l2, + "low_precision_quality_gate": False, + "low_precision_contract_gate": numerical_gate, + "padding_adapter": padding_check, + }, + "summary": summary, + "comparisons": comparisons, + "batch_medians_ms": batches, + "raw_ms": raw, + "route": { + "attention_calls": attention_calls, + "adapter_calls": adapter.snapshot(), + "expected_C_per_forward": lowp.expected_route_delta("C", shape["layers"]), + "torch_probe": torch_probe, + }, + "quantization": adapter.metadata(), + "provenance": {"git": build["git"], "sources": sources}, + } + if args.compare is not None: + if args.mode != "formal": + raise ValueError("--compare is formal-only") + metadata["comparison_across_runs"] = _compare(metadata, _load_comparison(args.compare)) + + args.output_dir.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + suffix = f"-{args.tag}" if args.tag else "" + raw_path = args.output_dir / f"qwen-image-modelopt-nvfp4-{args.mode}-{stamp}{suffix}.json" + report_path = raw_path.with_suffix(".md") + if raw_path.exists() or report_path.exists(): + raise FileExistsError(f"refusing to overwrite existing artifact {raw_path} or {report_path}") + metadata["artifacts"] = {"raw_json": str(raw_path), "markdown": str(report_path)} + raw_path.write_text(json.dumps(metadata, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8") + raw_hash = _sha256(raw_path) + report_path.write_text(_render_markdown(metadata, raw_path.name, raw_hash), encoding="utf-8") + print( + "RESULT " + + json.dumps( + { + "A_p50_ms": summary["A"]["p50_ms"], + "B_p50_ms": summary["B"]["p50_ms"], + "C_p50_ms": summary["C"]["p50_ms"], + "B_over_A": comparisons["B_vs_A"]["paired_ratio_p50"], + "C_over_B": comparisons["C_vs_B"]["paired_ratio_p50"], + "C_over_A": comparisons["C_vs_A"]["paired_ratio_p50"], + }, + sort_keys=True, + ) + ) + print(f"RAW_JSON {raw_path} sha256={raw_hash}") + print(f"MARKDOWN {report_path} sha256={_sha256(report_path)}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/e2e/README.md b/benchmark/e2e/README.md index 5867788fc..475aeb448 100644 --- a/benchmark/e2e/README.md +++ b/benchmark/e2e/README.md @@ -11,6 +11,18 @@ training step cuDNN already owns, and which un-owned op is next. Layer count and be reduced to avoid repeating identical work; every such reduction or stand-in is printed and documented rather than described as a full-model benchmark. +Every model leaf is a controlled **cuDNN-off versus cuDNN-on** experiment. The +off treatment uses the credible non-cuDNN framework route, even when the stock +framework dispatcher already selects cuDNN by default. The broader program has +five purposes: + +1. expose cuDNN technology and its user-visible impact; +2. measure the return from fusions and specialized kernels; +3. surface integration gaps before users hit them; +4. identify the next high-value kernel opportunity; and +5. take missed opportunities back upstream through Megatron, vLLM, SGLang, and + related integrations. + The current dense preset follows Qwen3.8-27B (and the kernel-equivalent Qwen3.5/3.6-27B): H=5120, I=17408, GDN 16 QK / 48 V heads at head-dim 128, and a 3:1 GDN/full-attention period. It keeps four layers and scales the vocabulary by the @@ -101,32 +113,125 @@ four-layer shape proxy, not full 64-layer Qwen throughput. On the same full B200, the Qwen-Image proxy uses the published H=3072, 24x128-head and FFN=12288 dimensions, B=1, 4096 image plus 512 text tokens, and -four of the 60 repeated transformer blocks. Forty balanced batches with three -repeats compare an explicitly forced PyTorch FlashAttention treatment with the -FE public cuDNN backend graph; all projections, QK norm, RoPE, AdaLN, biased -GELU FFNs, residuals, and output work are common and remain inside the timed -transformer forward. - -| SDPA treatment | p50 transformer forward | paired ratio | latency reduction | speedup | wins | -|---|---:|---:|---:|---:|---:| -| forced PyTorch FlashAttention | 9.943 ms | 1.00000 | -- | -- | -- | -| direct FE/cuDNN backend | 7.883 ms | **0.79640** | **20.36%** | **1.256x** | **40/40** | - -The unforced public Torch call already selects `CUDNN_ATTENTION` at this d128 -shape, so this is an implementation A/B rather than a claim of an additional -dispatcher-level user speedup. The complete four-block output matched within -0.141% relative L2. The focused B=2 unequal-text mask case matched within 0.300% -relative L2. Raw artifact SHA-256: `288ce0415c0cfd6564fde99debe0273a2304c07e7810547bd5da8b25cda0fbba`. - -The shared, model-agnostic harness lives in [`_perfshare.py`](_perfshare.py); a -model file only builds its model, applies the swaps, and calls `profile_and_report`. +four of the 60 repeated transformer blocks. Each block has one image-stream and +one text-stream dense biased GELU MLP; this model has no MoE/router. Forty +Williams-balanced batches with three repeats measure two independent treatments: + +- `M`: stock Diffusers GELU MLP versus public + `cudnn.gemm.ops.gelu_mlp(x, w1, b1, w2, b2)`; and +- `A`: explicitly forced PyTorch FlashAttention versus the FE public cuDNN + backend graph for joint SDPA. + +| bits (M/A) | GELU MLP | joint SDPA | p50 transformer forward | paired ratio vs `00` | wins vs `00` | +|---|---|---|---:|---:|---:| +| `00` | Torch | forced PyTorch Flash | 9.978 ms | 1.00000 | -- | +| `01` | Torch | cuDNN backend | 7.921 ms | 0.79444 | 40/40 | +| `10` | cuDNN | forced PyTorch Flash | 9.785 ms | 0.98292 | 34/40 | +| `11` | cuDNN | cuDNN backend | 7.770 ms | **0.78121** | **40/40** | + +The directly paired `11/00` result is 21.88% lower elapsed time, or **1.280x**. +The conditional attention ratio is 0.79285 (**1.261x**, 2.045 ms median +saving); the conditional GELU-MLP ratio is 0.98403 (**1.016x**, 0.142 ms median +saving). The MLP win is directionally consistent in both contexts: 34/40 wins +with Flash attention and 38/40 with cuDNN attention. + +`gelu_mlp` implements the same published biased +`Linear -> GELU(approximate="tanh") -> Linear` FFN semantics and matches Torch +within the documented BF16 tolerance. Its forward fuses the first matmul, bias, +BF16 boundary, and GELU into one cuDNN graph launch, then runs the output linear +as a second graph; its autograd path also fuses +`dout @ w2` with GELU backward. It is not the SwiGLU op used by Qwen3.8. + +A separate CUDA-event diagnostic explains where the E2E gain comes from. These +are mutually exclusive module regions from that diagnostic run, not the formal +factorial samples or FLOP shares: + +Diagnostic artifact SHA-256: `fb3cb4d9381b400fec80d2635cf32c1d19926c3933a6c0356fe2120cd4aa3ef1`. + +| four-block region | cuDNN-off time | off share | cuDNN-on time | +|---|---:|---:|---:| +| two GELU MLPs per block | 2.030 ms | 21.0% | 1.861 ms | +| joint SDPA core | 2.749 ms | 28.4% | 0.677 ms | +| attention projections, QK norm, RoPE, and surrounding work | 3.307 ms | 34.2% | 3.493 ms | +| AdaLN, residuals, output, and other work | 1.589 ms | 16.4% | 1.626 ms | + +The stock unforced Torch call already selects `CUDNN_ATTENTION` at this d128 +shape. That is successful cuDNN adoption, not a reason to discard the result: +the controlled experiment explicitly disables cuDNN SDPA by forcing Flash in +the off arm and quantifies the full-transformer impact of turning cuDNN back on. +It does not claim a further 1.280x from changing Torch's current dispatcher. + +The complete four-block outputs matched the `00` baseline within 0.142% +relative L2; the focused B=2 unequal-text mask case matched within 0.300%. +Raw artifact SHA-256: `63274d0602fe0582088f5241e0dcddcaac244c1426c955bc8e979c4a09fb55d3`. + +## Qwen-Image ModelOpt NVFP4 proxy result + +The sibling [`Qwen-Image/run_nvfp4.py`](Qwen-Image/run_nvfp4.py) leaf anchors +its placement and quantization policy to NVIDIA ModelOpt 0.46.0 commit +`43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a`: Qwen-Image's middle +transformer blocks use NVFP4 E2M1/block-16 Linears with E4M3 block scales and +`max` calibration. ModelOpt's Qwen invocation does not enable +`quantize_mha`, so joint attention remains BF16; this is not an all-FP4 model +or an MXFP8-attention result. + +The four proxy blocks represent full-model blocks `[2, 20, 39, 57]` from the +official quantized range 2..57. Every block routes all 14 Linear roles through +NVFP4, including the two M=1 modulation projections. The proxy deliberately +uses BF16 as its high-precision dtype and one synthetic frozen max-calibration +pass instead of ModelOpt's default FP16/calibration workload. It therefore +claims recipe-policy alignment and kernel-plumbing/performance evidence, not +official calibration state or image quality. All 56 weights are prepacked once +during setup and excluded from event timing; this differs from the quoted bare +ModelOpt CLI's default `compress=false` execution state. + +On the same full B200 and formal four-block shape, 42 position- and +carryover-balanced batches with three repeats measured: + +| arm | Linear / FFN treatment | joint SDPA | p50 transformer forward | paired ratio vs A | wins vs A | +|---|---|---|---:|---:|---:| +| `A` | Torch BF16 / Diffusers GELU FFN | forced PyTorch Flash BF16 | 9.852 ms | 1.00000 | -- | +| `B` | Torch BF16 + cuDNN BF16 `gelu_mlp` | cuDNN BF16 | 7.782 ms | **0.78997** | **42/42** | +| `C` | cuDNN FROST NVFP4 for all 14 Linears | cuDNN BF16 | 7.646 ms | **0.77567** | **42/42** | + +At fixed BF16, B/A is 21.00% lower elapsed time, or **1.266x**. The complete +cuDNN-enabled low-precision stack C is 22.43% lower than A, or **1.289x**. C is +also 1.73% lower than the already-optimized BF16 cuDNN arm B +(`C/B=0.98267`, **1.018x**, 37/42 wins). This incremental low-precision win is +modest: its paired p10--p90 ratio spans 0.97600--1.00452, so a few batches +slightly favor B. + +The final C path caches typed views and resolved FROST bindings per logical +Linear, validates every stable buffer in `select("C")` outside the CUDA-event +region, and uses the public `run_resolved` entry point. Direct Linear outputs +remain fresh allocations whose temporary binding slot is always cleared. This +removes the repeated host binding work exposed by the earlier diagnostic while +retaining strict route and lifetime guards; it does not use a private lowered +launcher. + +The run requires exact successful routes for all 56 NVFP4 Linears, with 56 +logical activation quantizations reduced to 33 physical operations (25 +standalone, eight fused, and 23 cache hits). A setup-only numerical gate +executes all seven distinct M/N/K/epilogue contracts against an independent +E2M1/F8_128x4 dequantized reference. The four-block C output differs from A by +0.852% relative L2, which is recorded only as a finite diagnostic because the +proxy uses random weights and synthetic calibration. + +Raw artifact SHA-256: `7af126f91ea958a8912e611168136afc2241fbc79e9d74d4a26ace907648f7e6`. +The benchmark-private BF16-to-NVFP4 kernel is derived from FlashInfer commit +`f212ec8230486e3615502b8af75fe7022c60b2f3`, retaining its Apache-2.0 notice +and its TensorRT-LLM provenance; FROST folds dequantization into the MMA. + +For Qwen3.8, the shared, model-agnostic harness lives in +[`_perfshare.py`](_perfshare.py); its model file only builds the model, applies +the swaps, and calls `profile_and_report`. ## Models | folder | proxy of | current precision leaf | notes | |---|---|---|---| | [`Qwen3.8/`](Qwen3.8/) | Qwen3.8/3.6/3.5-27B dense hybrid Gated DeltaNet LM | BF16 fwd+CE+bwd, 2^3 GDN/MLP/SDPA | exact MLP/GDN dimensions; 4-layer period; selectable Torch FlashAttention or cuDNN-backend d256 GQA at 20Q/4KV instead of gated 24Q/4KV | -| [`Qwen-Image/`](Qwen-Image/) | Qwen-Image diffusion transformer | BF16 transformer forward, forced PyTorch Flash-vs-cuDNN joint SDPA | exact H=3072, 24x128 and FFN=12288; 4/60 repeated blocks; 4096 image + 512 text tokens | +| [`Qwen-Image/`](Qwen-Image/) | Qwen-Image diffusion transformer | BF16 2^2 GELU-MLP/SDPA plus ModelOpt-anchored NVFP4 three-arm leaf | exact H=3072, 24x128 and FFN=12288; 4/60 repeated blocks; 4096 image + 512 text tokens; no MoE | Planned: Kimi Linear (KDA), DeepSeek-V3. @@ -173,14 +278,25 @@ python benchmark/e2e/Qwen3.8/run_model.py --preset qwen3.5-27b --inspect # equi # Qwen-Image uses the pinned benchmark-only Diffusers implementation. python -m pip install -r benchmark/e2e/Qwen-Image/requirements.txt -# Validation-only reduced-token mask/route/correctness smoke. +# Validation-only reduced-token mask/route/correctness smoke for all four +# Torch/cuDNN MLP x attention treatments. python benchmark/e2e/Qwen-Image/run_bf16.py \ --mode smoke --output-dir qwen-image-bf16-results/smoke -# Formal one-forward BF16 transformer proxy: B=1, 4096 image + 512 text tokens, -# four real-shape blocks, 40 balanced batches x 3 repeats on a full B200. +# Formal four-arm BF16 transformer proxy: B=1, 4096 image + 512 text tokens, +# four real-shape blocks, 40 Williams-balanced batches x 3 repeats on a full B200. python benchmark/e2e/Qwen-Image/run_bf16.py \ --mode formal --output-dir qwen-image-bf16-results/formal + +# ModelOpt-anchored NVFP4 validation. The benchmark-private quantizer lazily +# compiles one SM100 CUDA extension, so CUDA_HOME, an sm_100a-capable nvcc, +# a host C++ compiler, Ninja, and a writable Torch extension cache are required. +python benchmark/e2e/Qwen-Image/run_nvfp4.py \ + --mode smoke --output-dir qwen-image-nvfp4-results/smoke + +# Formal three-arm A/B/C run: BF16 off, BF16 cuDNN, then all-Linear NVFP4 cuDNN. +python benchmark/e2e/Qwen-Image/run_nvfp4.py \ + --mode formal --output-dir qwen-image-nvfp4-results/formal ``` The Qwen3.8 runner requires a cuDNN build with the fused GEMM engine and the @@ -214,9 +330,10 @@ complete denoising loop and is not image-quality evidence. The model/config and Diffusers implementation are pinned in the artifact. The A/B explicitly forces PyTorch FlashAttention versus direct FE/cuDNN; the artifact also reports the unforced public Torch dispatch choice (Torch 2.13 on B200 already selects cuDNN -for this d128 shape). This BF16 leaf has no -ModelOpt claim; a future ModelOpt-anchored NVFP4+FP8 experiment belongs in a -separate `run_nvfp4_fp8.py`. +for this d128 shape). The orthogonal MLP axis compares the stock Diffusers FFN +with the public cuDNN GELU-MLP op. The BF16 leaf itself has no ModelOpt claim; +the separate `run_nvfp4.py` leaf owns the pinned low-precision recipe, synthetic +calibration disclosure, route gates, and low-precision artifact. ## Add a model diff --git a/benchmark/e2e/tests/test_qwen_image_nvfp4_spec.py b/benchmark/e2e/tests/test_qwen_image_nvfp4_spec.py new file mode 100644 index 000000000..5680529c8 --- /dev/null +++ b/benchmark/e2e/tests/test_qwen_image_nvfp4_spec.py @@ -0,0 +1,497 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections import Counter +import importlib.util +import inspect +from pathlib import Path +import sys +from types import SimpleNamespace +import unittest + +E2E_DIR = Path(__file__).resolve().parents[1] +QWEN_DIR = E2E_DIR / "Qwen-Image" + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +LOWP = load("qwen_image_nvfp4_spec", QWEN_DIR / "modelopt_nvfp4.py") +RUNNER = load("qwen_image_nvfp4_runner_spec", QWEN_DIR / "run_nvfp4.py") + + +FORMAL_SHAPE = { + "layers": 4, + "bs": 1, + "image_tokens": 4096, + "text_tokens": 512, + "hidden": 3072, + "ffn": 12288, +} + + +class QwenImageNvfp4SpecTest(unittest.TestCase): + def test_recipe_is_exact_modelopt_046_fp4_max_without_mha_quantization(self): + recipe = LOWP.MODELOPT_RECIPE + self.assertEqual(recipe["release"], "0.46.0") + self.assertEqual(recipe["commit"], "43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a") + self.assertEqual( + recipe["upstream_anchor_args"], + "--model qwen-image --format fp4 --quant-algo max", + ) + self.assertEqual(recipe["repo_url"], "https://github.com/NVIDIA/Model-Optimizer") + for key in ( + "selection", + "qwen_defaults", + "preset", + "numerics", + "mha_policy", + "real_backend", + ): + self.assertIn(key, recipe["sources"]) + self.assertIn(recipe["commit"], recipe["source_permalinks"][key]) + self.assertTrue(recipe["source_permalinks"][key].endswith(recipe["sources"][key])) + self.assertIn("BFloat16", recipe["proxy_overrides"]["model_dtype"]) + self.assertIn("prepacking", recipe["proxy_overrides"]["weights"]) + self.assertIn("placement", recipe["alignment_scope"]) + self.assertFalse(recipe["quantize_mha"]) + self.assertFalse(recipe["numerical_claim_eligible"]) + self.assertIn("synthetic", recipe["calibration"]) + + def test_exact_fourteen_roles_include_both_m1_modulations(self): + self.assertEqual(len(LOWP.ROLE_ORDER), 14) + self.assertEqual(LOWP.ROLE_ORDER[:2], ("img_mod.1", "txt_mod.1")) + self.assertEqual(len(LOWP.MLP_ROLES), 4) + inputs = LOWP.expected_input_shapes(FORMAL_SHAPE) + weights = LOWP.expected_weight_shapes(FORMAL_SHAPE) + self.assertEqual(set(inputs), set(LOWP.ROLE_ORDER)) + self.assertEqual(set(weights), set(LOWP.ROLE_ORDER)) + for role in ("img_mod.1", "txt_mod.1"): + self.assertEqual(inputs[role], (1, 3072)) + self.assertEqual(weights[role], (18432, 3072)) + + def test_formal_proxy_maps_to_reviewed_interior_blocks(self): + self.assertEqual(LOWP.representative_middle_blocks(4), [2, 20, 39, 57]) + self.assertEqual(LOWP.representative_middle_blocks(1), [30]) + self.assertEqual(LOWP.representative_middle_blocks(56), list(range(2, 58))) + for invalid in (0, 57, None): + with self.assertRaises(ValueError): + LOWP.representative_middle_blocks(invalid) + + def test_exact_seven_low_precision_plan_contracts(self): + contracts = LOWP.expected_plan_contracts(FORMAL_SHAPE) + self.assertEqual(len(contracts), 7) + self.assertEqual(len(set(contracts)), 7) + self.assertIn((1, 18432, 3072, "linear_bias"), contracts) + self.assertIn((512, 12288, 3072, "linear_bias_gelu_nvfp4"), contracts) + self.assertIn((4096, 12288, 3072, "linear_bias_gelu_nvfp4"), contracts) + + def test_three_arm_design_balances_positions_and_ordered_carryover(self): + orders = LOWP.three_arm_orders() + self.assertEqual(len(orders), 6) + positions = Counter((position, treatment) for order in orders for position, treatment in enumerate(order)) + carryover = Counter(pair for order in orders for pair in zip(order, order[1:])) + self.assertEqual(set(positions.values()), {2}) + self.assertEqual(set(carryover.values()), {2}) + self.assertEqual([arm["id"] for arm in RUNNER.ARMS], ["A", "B", "C"]) + self.assertEqual(RUNNER.ARMS[2]["attention_route"], "cudnn") + self.assertIn("quantize_mha=False", RUNNER.ARMS[2]["attention"]) + + def test_protocol_requires_complete_six_order_cycles(self): + args = SimpleNamespace(mode="formal", warmup=None, rounds=None, repeats=None) + self.assertEqual(RUNNER._resolve_protocol(args), {"warmup": 3, "rounds": 42, "repeats": 3}) + args.rounds = 40 + with self.assertRaisesRegex(ValueError, "multiple of 6"): + RUNNER._resolve_protocol(args) + + def test_c_route_counts_all_roles_and_shared_quantization(self): + route = LOWP.expected_route_delta("C", 4) + self.assertEqual(route["nvfp4_linear_calls"], 56) + self.assertEqual(route["activation_quant_logical"], 56) + self.assertEqual(route["activation_quant_physical"], 33) + self.assertEqual(route["activation_quant_standalone"], 25) + self.assertEqual(route["activation_quant_fused"], 8) + self.assertEqual(route["activation_cache_hits"], 23) + self.assertEqual(len(route["nvfp4_linear_by_role"]), 56) + self.assertEqual(set(route["nvfp4_linear_by_role"].values()), {1}) + self.assertEqual(route["fallback_calls"], 0) + self.assertEqual(route["weight_pack_calls"], 0) + self.assertEqual(route["plan_build_calls"], 0) + + def test_a_and_b_keep_all_fourteen_logical_linears_bf16(self): + a = LOWP.expected_route_delta("A", 4) + b = LOWP.expected_route_delta("B", 4) + self.assertEqual(a["bf16_linear_calls"], 56) + self.assertEqual(b["bf16_linear_calls"], 56) + self.assertEqual(a["mlp_calls"], {"torch": 8, "cudnn_bf16": 0, "nvfp4": 0}) + self.assertEqual(b["mlp_calls"], {"torch": 0, "cudnn_bf16": 8, "nvfp4": 0}) + self.assertEqual(a["nvfp4_linear_calls"], 0) + self.assertEqual(b["nvfp4_linear_calls"], 0) + + def test_counter_tree_helpers_preserve_strict_nested_shape(self): + one = LOWP.expected_route_delta("C", 1) + two = RUNNER._scale_counter_tree(one, 2) + self.assertEqual(two["nvfp4_linear_calls"], 28) + self.assertEqual(set(two["nvfp4_linear_by_role"].values()), {2}) + total = RUNNER._add_counter_trees(one, two) + self.assertEqual(total["nvfp4_linear_calls"], 42) + self.assertEqual(LOWP.counter_delta(total, two), one) + + def test_pre_resolved_dynamic_output_is_never_retained(self): + output_id = 123 + output = object() + resolved = {456: object()} + + class Compiled: + def __init__(self, fail=False): + self.fail = fail + self.seen = None + + def run_resolved(self, actual, *, stream): + self.seen = (actual[output_id], stream) + if self.fail: + raise ValueError("launch failed") + return "ok" + + compiled = Compiled() + self.assertEqual( + LOWP._run_resolved_with_temporary_output(compiled, resolved, output_id, output, stream=789), + "ok", + ) + self.assertEqual(compiled.seen, (output, 789)) + self.assertNotIn(output_id, resolved) + + compiled = Compiled(fail=True) + with self.assertRaisesRegex(ValueError, "launch failed"): + LOWP._run_resolved_with_temporary_output(compiled, resolved, output_id, output, stream=987) + self.assertEqual(compiled.seen, (output, 987)) + self.assertNotIn(output_id, resolved) + + resolved[output_id] = object() + with self.assertRaisesRegex(RuntimeError, "already occupied"): + LOWP._run_resolved_with_temporary_output(Compiled(), resolved, output_id, output, stream=0) + + def test_timed_plan_paths_use_public_pre_resolved_entrypoint(self): + linear = inspect.getsource(LOWP._Nvfp4LinearPlan.__call__) + fused = inspect.getsource(LOWP._Nvfp4FusedFc1Plan.__call__) + self.assertIn("_run_resolved_with_temporary_output", linear) + self.assertIn("compiled.run_resolved", fused) + self.assertNotIn(".lowered", linear) + self.assertNotIn(".lowered", fused) + + def test_prepared_binding_tracks_every_stable_runtime_buffer(self): + fields = set(LOWP._PreparedNvfp4Binding.__dataclass_fields__) + self.assertTrue( + { + "activation_packed_signature", + "activation_scale_signature", + "activation_global_scale_signature", + "packed_weight_signature", + "weight_scale_signature", + "alpha_signature", + "bias_signature", + } + <= fields + ) + linear_guard = inspect.getsource(LOWP._Nvfp4LinearPlan.validate_prepared) + fused_guard = inspect.getsource(LOWP._Nvfp4FusedFc1Plan.validate_prepared) + for signature in ( + "packed_weight_signature", + "weight_scale_signature", + "alpha_signature", + "bias_signature", + ): + self.assertIn(signature, linear_guard) + self.assertIn(signature, fused_guard) + self.assertNotIn("_tensor_signature", inspect.getsource(LOWP._Nvfp4LinearPlan._binding)) + self.assertNotIn("_tensor_signature", inspect.getsource(LOWP._Nvfp4FusedFc1Plan._binding)) + + def test_prepared_binding_rejects_every_replaced_stable_buffer(self): + class FakeTensor: + next_pointer = 1000 + + def __init__(self, shape=(1,)): + self.pointer = FakeTensor.next_pointer + FakeTensor.next_pointer += 1 + self.shape = shape + self.dtype = "fake_dtype" + self.device = "cuda:0" + + def data_ptr(self): + return self.pointer + + def stride(self): + return tuple(1 for _ in self.shape) + + activation = SimpleNamespace( + packed=FakeTensor((2, 4)), + scale_factors=FakeTensor((128, 4)), + global_scale=FakeTensor(), + ) + entry = SimpleNamespace( + qualified_name="transformer_blocks.0.attn.to_q", + role="attn.to_q", + m=2, + n=8, + k=8, + activation_global_scale=activation.global_scale, + packed_weight=FakeTensor((8, 4)), + weight_scale_factors=FakeTensor((128, 4)), + alpha=FakeTensor((1, 1, 1)), + module=SimpleNamespace(bias=FakeTensor((8,))), + ) + binding = LOWP._PreparedNvfp4Binding( + entry=entry, + resolved={111: FakeTensor()}, + resolved_refs={}, + resolved_signatures={}, + activation_packed=activation.packed, + activation_scale_factors=activation.scale_factors, + activation_global_scale=activation.global_scale, + packed_weight=entry.packed_weight, + weight_scale_factors=entry.weight_scale_factors, + alpha=entry.alpha, + bias=entry.module.bias, + activation_packed_signature=LOWP._tensor_signature(activation.packed), + activation_scale_signature=LOWP._tensor_signature(activation.scale_factors), + activation_global_scale_signature=LOWP._tensor_signature(activation.global_scale), + packed_weight_signature=LOWP._tensor_signature(entry.packed_weight), + weight_scale_signature=LOWP._tensor_signature(entry.weight_scale_factors), + alpha_signature=LOWP._tensor_signature(entry.alpha), + bias_signature=LOWP._tensor_signature(entry.module.bias), + output_id=222, + ) + binding.resolved_refs = dict(binding.resolved) + binding.resolved_signatures = {key: LOWP._tensor_signature(value) for key, value in binding.resolved_refs.items()} + plan = object.__new__(LOWP._Nvfp4LinearPlan) + plan._prepared = {id(entry): binding} + self.assertIs(plan._binding(activation, entry, entry.alpha, entry.module.bias), binding) + + adapter = object.__new__(LOWP.QwenImageModelOptNvfp4Adapter) + adapter._active = False + adapter._device = SimpleNamespace(index=0) + adapter._stream = 333 + adapter._validate_stream = lambda device: None + adapter.entries = [entry] + adapter.by_name = {entry.qualified_name: entry} + adapter._activation_buffers = {(entry.m, entry.k): (activation.packed, activation.scale_factors)} + adapter._linear_plans = {(0, 333, entry.m, entry.n, entry.k, "linear_bias"): plan} + adapter._fused_fc1_plans = {} + adapter._installed_generic = {} + adapter._installed_mod = {} + adapter._installed_mlp = {} + adapter.select("C") + + for owner, attribute in ( + (entry, "activation_global_scale"), + (entry, "packed_weight"), + (entry, "weight_scale_factors"), + (entry, "alpha"), + (entry.module, "bias"), + ): + original = getattr(owner, attribute) + replacement = FakeTensor(original.shape) + setattr(owner, attribute, replacement) + with self.assertRaisesRegex(RuntimeError, "changed after NVFP4 preparation"): + adapter.select("C") + setattr(owner, attribute, original) + + original_buffers = adapter._activation_buffers[(entry.m, entry.k)] + for index in (0, 1): + changed = list(original_buffers) + changed[index] = FakeTensor(changed[index].shape) + adapter._activation_buffers[(entry.m, entry.k)] = tuple(changed) + with self.assertRaisesRegex(RuntimeError, "changed after NVFP4 preparation"): + adapter.select("C") + adapter._activation_buffers[(entry.m, entry.k)] = original_buffers + + original_resolved = binding.resolved[111] + binding.resolved[111] = FakeTensor(original_resolved.shape) + with self.assertRaisesRegex(RuntimeError, "resolved binding changed"): + adapter.select("C") + binding.resolved[111] = original_resolved + + def test_fused_prepared_binding_rejects_hgs_and_fixed_output_replacement(self): + class FakeTensor: + next_pointer = 2000 + + def __init__(self, shape=(1,)): + self.pointer = FakeTensor.next_pointer + FakeTensor.next_pointer += 1 + self.shape = shape + self.dtype = "fake_dtype" + self.device = "cuda:0" + + def data_ptr(self): + return self.pointer + + def stride(self): + return tuple(1 for _ in self.shape) + + activation = SimpleNamespace( + packed=FakeTensor((2, 4)), + scale_factors=FakeTensor((128, 4)), + global_scale=FakeTensor(), + ) + first = SimpleNamespace( + qualified_name="transformer_blocks.0.img_mlp.net.0.proj", + role="img_mlp.net.0.proj", + m=2, + n=32, + k=8, + activation_global_scale=activation.global_scale, + packed_weight=FakeTensor((32, 4)), + weight_scale_factors=FakeTensor((128, 4)), + alpha=FakeTensor((1, 1, 1)), + module=SimpleNamespace(bias=FakeTensor((32,))), + ) + second = SimpleNamespace( + qualified_name="transformer_blocks.0.img_mlp.net.2", + role="img_mlp.net.2", + m=2, + n=8, + k=32, + activation_global_scale=FakeTensor(), + packed_weight=FakeTensor((8, 16)), + weight_scale_factors=FakeTensor((128, 4)), + alpha=FakeTensor((1, 1, 1)), + module=SimpleNamespace(bias=FakeTensor((8,))), + ) + qh, sh = FakeTensor((2, 16)), FakeTensor((128, 4)) + hidden_view = FakeTensor((1, 1, 1)) + resolved_tensor = FakeTensor() + binding = LOWP._PreparedNvfp4Binding( + entry=first, + resolved={111: resolved_tensor}, + resolved_refs={111: resolved_tensor}, + resolved_signatures={111: LOWP._tensor_signature(resolved_tensor)}, + activation_packed=activation.packed, + activation_scale_factors=activation.scale_factors, + activation_global_scale=activation.global_scale, + packed_weight=first.packed_weight, + weight_scale_factors=first.weight_scale_factors, + alpha=first.alpha, + bias=first.module.bias, + activation_packed_signature=LOWP._tensor_signature(activation.packed), + activation_scale_signature=LOWP._tensor_signature(activation.scale_factors), + activation_global_scale_signature=LOWP._tensor_signature(activation.global_scale), + packed_weight_signature=LOWP._tensor_signature(first.packed_weight), + weight_scale_signature=LOWP._tensor_signature(first.weight_scale_factors), + alpha_signature=LOWP._tensor_signature(first.alpha), + bias_signature=LOWP._tensor_signature(first.module.bias), + hidden_global_scale=hidden_view, + hidden_global_scale_source=second.activation_global_scale, + hidden_global_scale_signature=LOWP._tensor_signature(hidden_view), + hidden_global_scale_source_signature=LOWP._tensor_signature(second.activation_global_scale), + output_packed=qh, + output_scale_factors=sh, + output_packed_signature=LOWP._tensor_signature(qh), + output_scale_signature=LOWP._tensor_signature(sh), + ) + plan = object.__new__(LOWP._Nvfp4FusedFc1Plan) + plan._prepared = {id(first): binding} + plan.qh, plan.sh = qh, sh + plan.validate_prepared( + (activation.packed, activation.scale_factors), + first, + second.activation_global_scale, + ) + self.assertIs(plan._binding(activation, first, first.alpha, first.module.bias), binding) + + original_hgs = second.activation_global_scale + second.activation_global_scale = FakeTensor() + with self.assertRaisesRegex(RuntimeError, "changed after NVFP4 preparation"): + plan.validate_prepared( + (activation.packed, activation.scale_factors), + first, + second.activation_global_scale, + ) + second.activation_global_scale = original_hgs + + for attribute in ("qh", "sh"): + original = getattr(plan, attribute) + setattr(plan, attribute, FakeTensor(original.shape)) + with self.assertRaisesRegex(RuntimeError, "changed after NVFP4 preparation"): + plan.validate_prepared( + (activation.packed, activation.scale_factors), + first, + second.activation_global_scale, + ) + with self.assertRaisesRegex(RuntimeError, "changed after NVFP4 preparation"): + plan._binding(activation, first, first.alpha, first.module.bias) + setattr(plan, attribute, original) + + def test_report_separates_all_three_claims(self): + batches = {"A": [10.0, 10.2], "B": [8.0, 8.1], "C": [6.0, 6.1]} + summary = {} + for arm, values in batches.items(): + summary[arm] = { + "p10_ms": min(values), + "p50_ms": sum(values) / len(values), + "p90_ms": max(values), + **RUNNER.paired_stats(values, batches["A"]), + } + comparisons = { + "B_vs_A": RUNNER.paired_stats(batches["B"], batches["A"]), + "C_vs_B": RUNNER.paired_stats(batches["C"], batches["B"]), + "C_vs_A": RUNNER.paired_stats(batches["C"], batches["A"]), + } + metadata = { + "completed_utc": "2026-08-21T00:00:00Z", + "config": { + "mode": "formal", + "comparability_fingerprint": {"sha256": "comparable"}, + "build_fingerprint": {"sha256": "build"}, + "shape": { + **FORMAL_SHAPE, + "joint_tokens": 4608, + "heads": 24, + "head_dim": 128, + }, + "representative_full_blocks": [2, 20, 39, 57], + "numerical_recipe": dict(LOWP.MODELOPT_RECIPE), + }, + "summary": summary, + "comparisons": comparisons, + "correctness": {"model_output_rel_l2_vs_A": {"A": 0.0, "B": 0.01, "C": 0.1}}, + "route": { + "torch_probe": { + "natural_choice_name": "CUDNN_ATTENTION", + "forced_choice_name": "FLASH_ATTENTION", + }, + "expected_C_per_forward": LOWP.expected_route_delta("C", 4), + }, + "provenance": {"sources": {"adapter": {"path": "adapter.py", "sha256": "face"}}}, + } + report = RUNNER._render_markdown(metadata, "result.json", "cafe") + self.assertIn("BF16 cuDNN effect (B/A)", report) + self.assertIn("ModelOpt NVFP4 increment (C/B)", report) + self.assertIn("Total cuDNN platform impact (C/A)", report) + self.assertIn("quantize_mha` defaults to false", report) + self.assertIn("all fourteen Linear roles", report) + self.assertIn("quality-ineligible", report) + + def test_report_labels_elapsed_time_regressions_as_slower(self): + self.assertEqual( + RUNNER._format_elapsed_effect(0.8), + "1.250x speedup (20.00% lower elapsed time)", + ) + self.assertEqual( + RUNNER._format_elapsed_effect(1.0898134), + "1.090x slower (8.98% higher elapsed time)", + ) + self.assertEqual(RUNNER._format_elapsed_effect(1.0), "1.000x (no elapsed-time change)") + for invalid in (0.0, -1.0, float("nan"), float("inf")): + with self.assertRaises(ValueError): + RUNNER._format_elapsed_effect(invalid) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 66d51bdd0..4e188d53b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -106,16 +106,18 @@ jax = [ ] [build-system] -requires = ["setuptools>=64", "cmake>=3.18", "ninja==1.11.1.1", "pybind11[global]>=2.13,<3"] +requires = ["setuptools>=64", "wheel>=0.38.4", "cmake>=3.18", "ninja==1.11.1.1", "pybind11[global]>=2.13,<3"] build-backend = "setuptools.build_meta" [tool.setuptools] packages = {find = {where = ["python", "."], include = ["cudnn*", "include"], namespaces = true}} package-dir = {"" = "python", "include" = "include"} include-package-data = true +license-files = ["LICENSE.txt", "LICENSE-MIT.txt", "NOTICE", "THIRD_PARTY_LICENSES.txt"] [tool.setuptools.dynamic] version = {attr = "cudnn.__version__"} [tool.setuptools.package-data] include = ["**/*"] +"cudnn.gemm.ops" = ["csrc/*.cu", "csrc/*.cuh"] diff --git a/python/cudnn/gemm/ops/_nvfp4_quantize.py b/python/cudnn/gemm/ops/_nvfp4_quantize.py new file mode 100644 index 000000000..446c8dfea --- /dev/null +++ b/python/cudnn/gemm/ops/_nvfp4_quantize.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark-private SM100 BF16-to-NVFP4 activation quantization. + +This module intentionally is not re-exported from :mod:`cudnn.gemm`. Its +packed E2M1 and F8_128x4 scale-factor layouts are an internal bridge for the +Qwen-Image low-precision benchmark, not a stable public API. +""" + +from __future__ import annotations + +import os +from pathlib import Path +import threading +from typing import Optional, Tuple + +import torch + +_EXTENSION = None +_EXTENSION_LOCK = threading.Lock() +_ONES_CACHE = {} +_ONES_LOCK = threading.Lock() +_INT32_MAX = (1 << 31) - 1 + + +def _scale_factor_shape(m: int, k: int) -> Tuple[int, int]: + """Physical byte shape of a row-major logical ``[M, K/16]`` F8_128x4 tensor.""" + return ((m + 127) // 128 * 128, (k // 16 + 3) // 4 * 4) + + +def _load_extension(): + """Build once per process; cpp_extension also caches the binary on disk.""" + global _EXTENSION + if _EXTENSION is not None: + return _EXTENSION + + with _EXTENSION_LOCK: + if _EXTENSION is None: + # Keep cpp_extension and nvcc completely off the import path. In + # particular, unsupported devices fail validation before reaching + # this function. + from torch.utils.cpp_extension import load + + source_dir = Path(__file__).with_name("csrc") + _EXTENSION = load( + name="_cudnn_fe_nvfp4_quantize_sm100_f212ec82", + sources=[str(source_dir / "nvfp4_quantize_sm100.cu")], + extra_include_paths=[str(source_dir)], + extra_cuda_cflags=[ + "-O3", + "-gencode=arch=compute_100a,code=sm_100a", + "-DFLASHINFER_ENABLE_FP8_E8M0", + "-DFLASHINFER_ENABLE_FP4_E2M1", + ], + with_cuda=True, + verbose=os.environ.get("CUDNN_NVFP4_BUILD_VERBOSE", "0") == "1", + ) + return _EXTENSION + + +def _check_tensor( + tensor: torch.Tensor, + *, + name: str, + dtype: torch.dtype, + device: torch.device, + shape: Tuple[int, ...], + alignment: int, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tensor.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.data_ptr() % alignment: + raise ValueError(f"{name} data pointer must be {alignment}-byte aligned") + + +def _ones_pre_quant_scale(device: torch.device, k: int, stream: torch.cuda.Stream) -> torch.Tensor: + # Initialization and first consumption are ordered on the same stream. A + # stream is part of the key so another stream cannot observe an unfinished + # asynchronous fill from the first call. + key = (device.index, stream.cuda_stream, k) + value = _ONES_CACHE.get(key) + if value is not None: + return value + with _ONES_LOCK: + value = _ONES_CACHE.get(key) + if value is None: + value = torch.ones(k, dtype=torch.bfloat16, device=device) + _ONES_CACHE[key] = value + return value + + +def nvfp4_quantize( + x: torch.Tensor, + global_scale: torch.Tensor, + pre_quant_scale: Optional[torch.Tensor] = None, + *, + out: Optional[torch.Tensor] = None, + scale_factors: Optional[torch.Tensor] = None, + enable_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a contiguous BF16 matrix into ModelOpt-style NVFP4 buffers. + + ``x`` has shape ``[M, K]`` and ``pre_quant_scale`` has shape ``[K]``. + The kernel first performs the BF16 multiply ``x * pre_quant_scale`` and + then block-quantizes each 16-value K block. ``global_scale`` is a + same-device contiguous one-element FP32 tensor, conventionally + ``448 * 6 / amax``. + + Returns packed E2M1 bytes ``[M, K/2]`` and F8_128x4 scale-factor bytes + ``[ceil(M/128)*128, ceil((K/16)/4)*4]``. Callers may provide either output + buffer. Execution is asynchronous on the current PyTorch stream. + """ + if not isinstance(x, torch.Tensor): + raise TypeError(f"x must be a torch.Tensor, got {type(x).__name__}") + if not x.is_cuda: + raise ValueError("x must be a CUDA tensor") + if x.dtype is not torch.bfloat16: + raise TypeError(f"x must have dtype {torch.bfloat16}, got {x.dtype}") + if x.ndim != 2: + raise ValueError(f"x must have rank 2, got rank {x.ndim}") + if not x.is_contiguous(): + raise ValueError("x must be contiguous") + + m, k = x.shape + if m <= 0 or k <= 0: + raise ValueError(f"x dimensions must be positive, got {(m, k)}") + if k % 16: + raise ValueError(f"x.shape[1] must be divisible by 16, got {k}") + if m > _INT32_MAX or k > _INT32_MAX: + raise ValueError(f"x dimensions must fit signed 32-bit integers, got {(m, k)}") + if x.data_ptr() % 16: + raise ValueError("x data pointer must be 16-byte aligned") + + device = x.device + capability = torch.cuda.get_device_capability(device) + if capability != (10, 0): + raise RuntimeError(f"nvfp4_quantize requires SM100, got SM{capability[0]}{capability[1]}") + if type(enable_pdl) is not bool: + raise TypeError(f"enable_pdl must be bool, got {type(enable_pdl).__name__}") + + if not isinstance(global_scale, torch.Tensor): + raise TypeError(f"global_scale must be a torch.Tensor, got {type(global_scale).__name__}") + if global_scale.device != device: + raise ValueError(f"global_scale must be on {device}, got {global_scale.device}") + if global_scale.dtype is not torch.float32: + raise TypeError(f"global_scale must have dtype {torch.float32}, got {global_scale.dtype}") + if tuple(global_scale.shape) not in ((), (1,)): + raise ValueError("global_scale must have shape () or (1,), " f"got shape {tuple(global_scale.shape)}") + if not global_scale.is_contiguous(): + raise ValueError("global_scale must be contiguous") + if global_scale.data_ptr() % 4: + raise ValueError("global_scale data pointer must be 4-byte aligned") + + output_shape = (m, k // 2) + sf_shape = _scale_factor_shape(m, k) + with torch.cuda.device(device): + stream = torch.cuda.current_stream(device) + if pre_quant_scale is None: + pre_quant_scale = _ones_pre_quant_scale(device, k, stream) + else: + _check_tensor( + pre_quant_scale, + name="pre_quant_scale", + dtype=torch.bfloat16, + device=device, + shape=(k,), + alignment=16, + ) + + if out is None: + out = torch.empty(output_shape, dtype=torch.uint8, device=device) + else: + _check_tensor( + out, + name="out", + dtype=torch.uint8, + device=device, + shape=output_shape, + alignment=8, + ) + + if scale_factors is None: + scale_factors = torch.empty(sf_shape, dtype=torch.uint8, device=device) + else: + _check_tensor( + scale_factors, + name="scale_factors", + dtype=torch.uint8, + device=device, + shape=sf_shape, + alignment=4, + ) + + extension = _load_extension() + multiprocessor_count = torch.cuda.get_device_properties(device).multi_processor_count + extension.launch( + x.data_ptr(), + pre_quant_scale.data_ptr(), + global_scale.data_ptr(), + out.data_ptr(), + scale_factors.data_ptr(), + m, + k, + multiprocessor_count, + stream.cuda_stream, + enable_pdl, + ) + # The binding receives raw addresses. Explicitly tell the caching + # allocator that every backing allocation remains in use until work on + # the launch stream has completed. + for tensor in (x, pre_quant_scale, global_scale, out, scale_factors): + tensor.record_stream(stream) + return out, scale_factors + + +__all__ = ["nvfp4_quantize"] diff --git a/python/cudnn/gemm/ops/csrc/nvfp4_quantize_sm100.cu b/python/cudnn/gemm/ops/csrc/nvfp4_quantize_sm100.cu new file mode 100644 index 000000000..341f4b9ed --- /dev/null +++ b/python/cudnn/gemm/ops/csrc/nvfp4_quantize_sm100.cu @@ -0,0 +1,73 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include +#include + +#include "nvfp4_smooth_quantize_sm100.cuh" + +namespace { + +void +launch(std::uintptr_t x, + std::uintptr_t pre_quant_scale, + std::uintptr_t global_scale, + std::uintptr_t output, + std::uintptr_t scale_factors, + int m, + int k, + int multiprocessor_count, + std::uintptr_t stream, + bool enable_pdl) { + auto cuda_stream = reinterpret_cast(stream); + flashinfer::gemm::nvfp4_smooth_quantize(reinterpret_cast(output), + reinterpret_cast(scale_factors), + reinterpret_cast(x), + reinterpret_cast(pre_quant_scale), + reinterpret_cast(global_scale), + m, + k, + multiprocessor_count, + cuda_stream, + enable_pdl); + + cudaError_t status = cudaGetLastError(); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("nvfp4 quantize kernel launch failed: ") + cudaGetErrorString(status)); + } +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("launch", + &launch, + pybind11::arg("x"), + pybind11::arg("pre_quant_scale"), + pybind11::arg("global_scale"), + pybind11::arg("output"), + pybind11::arg("scale_factors"), + pybind11::arg("m"), + pybind11::arg("k"), + pybind11::arg("multiprocessor_count"), + pybind11::arg("stream"), + pybind11::arg("enable_pdl")); +} diff --git a/python/cudnn/gemm/ops/csrc/nvfp4_smooth_quantize_sm100.cuh b/python/cudnn/gemm/ops/csrc/nvfp4_smooth_quantize_sm100.cuh new file mode 100644 index 000000000..aa21b7956 --- /dev/null +++ b/python/cudnn/gemm/ops/csrc/nvfp4_smooth_quantize_sm100.cuh @@ -0,0 +1,605 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Derived from FlashInfer's include/flashinfer/gemm/nvfp4_smooth_quantize_sm100.cuh at +// https://github.com/flashinfer-ai/flashinfer/blob/f212ec8230486e3615502b8af75fe7022c60b2f3/include/flashinfer/gemm/nvfp4_smooth_quantize_sm100.cuh +// (introduced by FlashInfer commit f212ec8230486e3615502b8af75fe7022c60b2f3). The upstream +// source credits TensorRT-LLM's kernels/nvfp4SmoothQuantize.cu and quantization helpers below. +// clang-format off +// +// Fused smooth + NVFP4 quantize: apply the per-input-channel pre_quant_scale AND NVFP4-quantize in +// ONE pass over the input, eliminating the separate x_hat = x*s elementwise pass. Self-contained +// port of TensorRT-LLM's kernels/nvfp4SmoothQuantize.cu: the device helpers it reused from +// trtllm's kernels/quantization.cuh (cvt_warp_fp16_to_fp4, PackedVec, cvt_quant_get_sf_out_offset, +// get_sf_out_offset_128x4, ...) are copied verbatim below so the xq+SF output stays byte-identical +// to fp4_quantize(x*s); the only addition is the per-channel multiply before the block-amax + +// quantize. Differences vs the TRT-LLM original: PDL comes in as a function parameter (instead of +// an env probe) and the *_SMOOTH_QUANT_THREADS/*_SMOOTH_QUANT_BLOCKS_PER_SM env overrides are +// dropped (the defaults they fell back to are hardcoded). +#pragma once + +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_BFLOAT16_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT162_OPERATORS__ + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace flashinfer { +namespace gemm { + +namespace smooth_quantize_detail { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Helpers copied verbatim from TensorRT-LLM common/cudaTypeUtils.cuh (only the variants the +// kernels below instantiate: bfloat16/bfloat162 plus the generic templates they specialize). + +// Get type2 from type or vice versa (applied to half and bfloat16) +template +struct TypeConverter { + using Type = half2; +}; // keep for generality + +template <> +struct TypeConverter { + using Type = half; +}; + +template <> +struct TypeConverter { + using Type = half2; +}; + +template <> +struct TypeConverter<__nv_bfloat162> { + using Type = __nv_bfloat16; +}; + +template <> +struct TypeConverter<__nv_bfloat16> { + using Type = __nv_bfloat162; +}; + +template +__device__ inline T cuda_abs(T val) { + assert(false); + return {}; +} + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) { + return __habs(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) { + return __habs2(val); +} +#endif + +// Binary maximum: compute the max of two values. +template +__device__ inline T cuda_max(T val1, T val2) { + return (val1 > val2) ? val1 : val2; +} + +template <> +__device__ inline __nv_bfloat162 cuda_max(__nv_bfloat162 val1, __nv_bfloat162 val2) { + return __hmax2(val1, val2); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Helpers copied verbatim from TensorRT-LLM kernels/quantization.cuh (FP4 quantization section). + +constexpr int CVT_ELTS_PER_THREAD = 8; + +// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). +inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), "f"(array[2].x), + "f"(array[2].y), "f"(array[3].x), "f"(array[3].y)); + return val; +#else + // static_assert(false, "not supported."); + return 0; +#endif +} + +// Convert 8 float2 values into 16 e2m1 values (represented as one uint64_t). +inline __device__ uint64_t fp32_vec_to_e2m1(float2 (&array)[8]) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint64_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + ".reg .b8 byte4;\n" + ".reg .b8 byte5;\n" + ".reg .b8 byte6;\n" + ".reg .b8 byte7;\n" + ".reg .b32 val0;\n" + ".reg .b32 val1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte4, %10, %9;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte5, %12, %11;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte6, %14, %13;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte7, %16, %15;\n" + "mov.b32 val0, {byte0, byte1, byte2, byte3};\n" + "mov.b32 val1, {byte4, byte5, byte6, byte7};\n" + "mov.b64 %0, {val0, val1};\n" + "}" + : "=l"(val) + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), "f"(array[2].x), + "f"(array[2].y), "f"(array[3].x), "f"(array[3].y), "f"(array[4].x), "f"(array[4].y), + "f"(array[5].x), "f"(array[5].y), "f"(array[6].x), "f"(array[6].y), "f"(array[7].x), + "f"(array[7].y)); + return val; +#else + // static_assert(false, "not supported."); + return 0; +#endif +} + +// Fast reciprocal. +inline __device__ float reciprocal_approximate_ftz(float a) { + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// Define a 16 bytes packed data type. +template +struct PackedVec { + typename TypeConverter::Type elts[4]; + static_assert(sizeof(elts) == sizeof(Type) * CVT_ELTS_PER_THREAD, + "Vector size should match the number of elements per thread."); +}; + +// Quantizes the provided PackedVec into the uint32_t output. +// Port note: only the UE4M3 scale-factor path (UE8M0_SF == false) is kept; the removed +// "if constexpr (UE8M0_SF)" branch emitted no code for the instantiation this port uses. +template +__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { + static_assert(!UE8M0_SF, "this port only supports the UE4M3 scale-factor path"); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Get absolute maximum values among the local 8 values. + auto localMax = cuda_abs(vec.elts[0]); + +// Local maximum value. +#pragma unroll + for (int i = 1; i < CVT_ELTS_PER_THREAD / 2; i++) { + localMax = cuda_max(localMax, cuda_abs(vec.elts[i])); + } + + constexpr int CVT_NUM_THREADS_PER_SF = SF_VEC_SIZE / CVT_ELTS_PER_THREAD; + // Get the absolute maximum among all 16 values (two threads for 16, four threads for 32). + localMax = cuda_max(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); + if constexpr (CVT_NUM_THREADS_PER_SF == 4) { + localMax = cuda_max(__shfl_xor_sync(uint32_t(-1), localMax, 2), localMax); + } + // Get the final absolute maximum values. + float vecMax = float(cuda_max(localMax.x, localMax.y)); + + // 8 bits representation of the SF. + uint8_t fp8SFVal; + float outputScale; + // Get the SF (max value of the vector / max value of e2m1). + // maximum value of e2m1 = 6.0. + // TODO: use half as compute data type. + auto SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); + // Here SFValue is always positive, so E4M3 is the same as UE4M3. + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); + fp8SFVal = tmp.__x; + SFValue = static_cast(tmp); + // Get the output scale. + // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal)) * reciprocal(SFScaleVal)) + outputScale = vecMax != 0 + ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) + : 0.0f; + + if (SFout) { + // Write the SF to global memory (STG.8). + *SFout = fp8SFVal; + } + + // Convert the input to float. + float2 fp2Vals[CVT_ELTS_PER_THREAD / 2]; + +#pragma unroll + for (int i = 0; i < CVT_ELTS_PER_THREAD / 2; i++) { + if constexpr (std::is_same_v) { + fp2Vals[i] = __half22float2(vec.elts[i]); + } else { + fp2Vals[i] = __bfloat1622float2(vec.elts[i]); + } + fp2Vals[i].x *= outputScale; + fp2Vals[i].y *= outputScale; + } + + // Convert to e2m1 values. + uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); + + // Write the e2m1 values to global memory. + return e2m1Vec; +#else + return 0; +#endif +} + +// Port note: batch support dropped -- the original takes std::optional batchIdx/numRows, but +// every call site of this port passes std::nullopt/0, making the batch term (batchIdx * +// bTileStride) identically zero. Byte-identical for the single-batch case this port serves. +inline __host__ __device__ int64_t get_sf_out_offset_128x4(int mIdx, int kIdx, int numColVecs) { + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] + // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] + + int32_t innerKIdx = (kIdx % 4); + int64_t innerKStride = 1; + + int32_t innerMIdx = (mIdx % (32 * 4)) / 32; + int64_t innerMStride = 4 * innerKStride; // 4 + + // M tile layout [32, 4] is column-major. + int32_t outerMIdx = (mIdx % 32); + int64_t outerMStride = 4 * innerMStride; // 16 + + int32_t kTileIdx = (kIdx / 4); + int64_t kTileStride = 32 * outerMStride; // 512 + + // SF vector size 16 or 32. We round the "numCols" up to a multiple of 64 or 128. + // It is the same as rounding the "numColVecs" up to a multiple of 4. + int32_t numKTiles = (numColVecs + 4 - 1) / 4; + int32_t mTileIdx = mIdx / (32 * 4); + int64_t mTileStride = numKTiles * kTileStride; + + // Compute the global offset. + int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride + outerMIdx * outerMStride + + innerMIdx * innerMStride + innerKIdx * innerKStride; + + return SFOffset; +} + +// Port note: hardcoded to the SWIZZLED 128x4 layout (the only layout this port dispatches on) and +// batch support dropped as above; otherwise verbatim. +template +__device__ uint8_t* cvt_quant_get_sf_out_offset(int rowIdx, int colVecIdx, int numColVecs, + SFType* SFout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Each thread holds one vector. + static_assert(CVT_NUM_THREADS_PER_SF == 1 || CVT_NUM_THREADS_PER_SF == 2 || + CVT_NUM_THREADS_PER_SF == 4); + + // One pair of threads write one SF to global memory. + // TODO: stage through smem for packed STG.32 + // is it better than STG.8 from 4 threads ? + if (threadIdx.x % CVT_NUM_THREADS_PER_SF == 0) { + // SF vector index (16 elements share one SF in the K dimension). + // numRows and numCols are unpadded. + int32_t kIdx = colVecIdx / CVT_NUM_THREADS_PER_SF; + int32_t mIdx = rowIdx; + + auto SFOffset = get_sf_out_offset_128x4(mIdx, kIdx, numColVecs); + return reinterpret_cast(SFout) + SFOffset; + } +#endif + return nullptr; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// The fused smooth-quantize kernels, ported from TensorRT-LLM kernels/nvfp4SmoothQuantize.cu. + +// bf16, NVFP4 (UE4M3 SF, SF_VEC_SIZE=16), swizzled layout, single batch. +constexpr int SF_VEC_SIZE = 16; +using Type = __nv_bfloat16; +constexpr int ELTS_PER_THREAD = CVT_ELTS_PER_THREAD; +using SmoothPackedVec = PackedVec; +constexpr int CVT_NUM_THREADS_PER_SF = SF_VEC_SIZE / ELTS_PER_THREAD; +constexpr int FAST_ELTS_PER_THREAD = SF_VEC_SIZE; + +// Two of these make one complete 16-element NVFP4 scale block. Keeping the load granularity at +// 128 bits avoids imposing a stronger alignment requirement than the stock quantizer. +union alignas(16) Bf16x8 { + uint4 bits; + __nv_bfloat162 elts[4]; +}; + +static_assert(sizeof(Bf16x8) == 16); + +// trtllm's PadUpFn is a function-like macro (quantization.h); use a plain +// host+device helper here instead. +__host__ __device__ inline int padUp(int x, int y) { return (x + y - 1) / y * y; } + +__device__ __forceinline__ void loadBf16x8(Type const* ptr, Bf16x8& result) { + result.bits = *reinterpret_cast(ptr); +} + +__device__ __forceinline__ uint64_t quantizeSmoothed16(Bf16x8& lo, Bf16x8& hi, Bf16x8 const& pqsLo, + Bf16x8 const& pqsHi, float SFScaleVal, + uint8_t* SFout) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + lo.elts[i] = __hmul2(lo.elts[i], pqsLo.elts[i]); + hi.elts[i] = __hmul2(hi.elts[i], pqsHi.elts[i]); + } + + // Match the legacy even lane's reduction order: reduce each 8-element half independently, then + // merge the high half into the low half. For finite BF16 values this produces the same scale for + // both halves. + auto loMax = cuda_abs(lo.elts[0]); + auto hiMax = cuda_abs(hi.elts[0]); +#pragma unroll + for (int i = 1; i < 4; ++i) { + loMax = cuda_max(loMax, cuda_abs(lo.elts[i])); + hiMax = cuda_max(hiMax, cuda_abs(hi.elts[i])); + } + auto const localMax = cuda_max(hiMax, loMax); + float const vecMax = float(cuda_max(localMax.x, localMax.y)); + + // This is deliberately kept instruction-for-instruction equivalent to cvt_warp_fp16_to_fp4's + // UE4M3 path. + auto SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); + uint8_t const fp8SFVal = tmp.__x; + SFValue = static_cast(tmp); + float const outputScale = + vecMax != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) + : 0.0f; + + if (SFout != nullptr) *SFout = fp8SFVal; + + float2 fp2Vals[FAST_ELTS_PER_THREAD / 2]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + fp2Vals[i] = __bfloat1622float2(lo.elts[i]); + fp2Vals[i + 4] = __bfloat1622float2(hi.elts[i]); + fp2Vals[i].x *= outputScale; + fp2Vals[i].y *= outputScale; + fp2Vals[i + 4].x *= outputScale; + fp2Vals[i + 4].y *= outputScale; + } + return fp32_vec_to_e2m1(fp2Vals); +} + +template +__global__ void __launch_bounds__(512, 4) + smooth_quantize_fast_kernel(int numRows, Type const* __restrict__ in, + Type const* __restrict__ pqs, float const* __restrict__ SFScale, + uint64_t* __restrict__ out, uint8_t* __restrict__ SFout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(NumCols % (4 * SF_VEC_SIZE) == 0); + constexpr int NumSfCols = NumCols / SF_VEC_SIZE; + constexpr int NumSfGroups = NumSfCols / 4; + float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0]; + + cudaGridDependencySynchronize(); + + // Hot path: map a fixed number of complete rows onto each CTA. K=3072 uses two rows and 384 + // threads, so every thread owns exactly one 16-value scale block without a warp shuffle. + for (int rowBase = blockIdx.x * RowsPerCta; rowBase < numRows; + rowBase += gridDim.x * RowsPerCta) { + int const rowsRemaining = numRows - rowBase; + int const rowsThisCta = rowsRemaining < RowsPerCta ? rowsRemaining : RowsPerCta; + int const workItems = rowsThisCta * NumSfCols; + for (int item = threadIdx.x; item < workItems; item += blockDim.x) { + int const rowOffset = item / NumSfCols; + int const sfCol = item - rowOffset * NumSfCols; + int const row = rowBase + rowOffset; + int64_t const vecOffset = static_cast(row) * NumSfCols + sfCol; + + Type const* xPtr = in + vecOffset * FAST_ELTS_PER_THREAD; + Type const* pqsPtr = pqs + sfCol * FAST_ELTS_PER_THREAD; + Bf16x8 xLo; + Bf16x8 xHi; + Bf16x8 pqsLo; + Bf16x8 pqsHi; + loadBf16x8(xPtr, xLo); + loadBf16x8(xPtr + 8, xHi); + loadBf16x8(pqsPtr, pqsLo); + loadBf16x8(pqsPtr + 8, pqsHi); + + int64_t const sfOffset = get_sf_out_offset_128x4(row, sfCol, NumSfCols); + out[vecOffset] = quantizeSmoothed16(xLo, xHi, pqsLo, pqsHi, SFScaleVal, SFout + sfOffset); + } + } + + // Cold path: only scale factors have padded rows. Four consecutive SF columns are contiguous in + // the 128x4 layout, so initialize them with one aligned 32-bit store instead of four byte stores. + int const numPaddedRows = padUp(numRows, 128); + int64_t const numPaddingStores = static_cast(numPaddedRows - numRows) * NumSfGroups; + for (int64_t item = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + item < numPaddingStores; item += static_cast(gridDim.x) * blockDim.x) { + int const paddingRow = static_cast(item / NumSfGroups); + int const sfGroup = static_cast(item - static_cast(paddingRow) * NumSfGroups); + int const row = numRows + paddingRow; + int const sfCol = sfGroup * 4; + int64_t const sfOffset = get_sf_out_offset_128x4(row, sfCol, NumSfCols); + *reinterpret_cast(SFout + sfOffset) = 0u; + } + + cudaTriggerProgrammaticLaunchCompletion(); +#else + // Fail loudly instead of silently leaving out/SFout uninitialized if a build for an + // unsupported architecture is ever launched. + if (threadIdx.x == 0 && blockIdx.x == 0) { + printf("nvfp4_smooth_quantize requires SM100 or newer\n"); + __trap(); + } +#endif +} + +__global__ void __launch_bounds__(512, 4) + smooth_quantize_legacy_kernel(int numRows, int numCols, int numPaddedCols, Type const* in, + Type const* pqs, float const* SFScale, uint32_t* out, + uint32_t* SFout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0]; + int const numPaddedRowsForSf = padUp(numRows, 128); + int const numColsForSf = padUp(numPaddedCols, 4 * SF_VEC_SIZE); + int const numColThreads = numCols / ELTS_PER_THREAD; + int const numPaddedColThreads = numPaddedCols / ELTS_PER_THREAD; + int const numColThreadsForSf = numColsForSf / ELTS_PER_THREAD; + + cudaGridDependencySynchronize(); + for (int rowIdx = blockIdx.x; rowIdx < numPaddedRowsForSf; rowIdx += gridDim.x) { + bool const isRowPadding = (rowIdx >= numRows); + for (int colIdx = threadIdx.x; colIdx < numColThreadsForSf; colIdx += blockDim.x) { + auto sf_out = cvt_quant_get_sf_out_offset( + rowIdx, colIdx, numPaddedCols / SF_VEC_SIZE, SFout); + + if (isRowPadding || colIdx >= numColThreads) { + if (sf_out != nullptr) sf_out[0] = 0x00; + if (!isRowPadding && colIdx >= numColThreads && colIdx < numPaddedColThreads) + reinterpret_cast( + out)[static_cast(rowIdx) * numPaddedColThreads + colIdx] = 0u; + continue; + } + + int64_t const inOffset = static_cast(rowIdx) * numColThreads + colIdx; + int64_t const outOffset = static_cast(rowIdx) * numPaddedColThreads + colIdx; + SmoothPackedVec in_vec = reinterpret_cast(in)[inOffset]; + // --- the fusion: smooth by the per-channel pre_quant_scale (broadcast over rows) --- + SmoothPackedVec p_vec = reinterpret_cast(pqs)[colIdx]; +#pragma unroll + for (int i = 0; i < ELTS_PER_THREAD / 2; i++) + in_vec.elts[i] = __hmul2(in_vec.elts[i], p_vec.elts[i]); + reinterpret_cast(out)[outOffset] = + cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); + } + } + cudaTriggerProgrammaticLaunchCompletion(); +#else + if (threadIdx.x == 0 && blockIdx.x == 0) { + printf("nvfp4_smooth_quantize requires SM100 or newer\n"); + __trap(); + } +#endif +} + +template +void launchSmoothQuantizeFast(void* out, void* sfOut, void const* in, void const* pqs, + float const* sfScale, int numRows, int multiProcessorCount, + int blockThreads, int blocksPerSm, bool enablePDL, + cudaStream_t stream) { + constexpr int NumSfGroups = NumCols / (4 * SF_VEC_SIZE); + int const numPaddedRows = padUp(numRows, 128); + int64_t const hotCtas = (static_cast(numRows) + RowsPerCta - 1) / RowsPerCta; + int64_t const numPaddingStores = static_cast(numPaddedRows - numRows) * NumSfGroups; + int64_t const paddingCtas = (numPaddingStores + blockThreads - 1) / blockThreads; + int64_t const wantedCtas = std::max(hotCtas, paddingCtas); + int64_t const maxCtas = static_cast(multiProcessorCount) * blocksPerSm; + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(std::min(wantedCtas, maxCtas))); + cfg.blockDim = dim3(blockThreads); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePDL ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 1; + + auto* kernel = &smooth_quantize_fast_kernel; + cudaLaunchKernelEx(&cfg, kernel, numRows, reinterpret_cast(in), + reinterpret_cast(pqs), sfScale, reinterpret_cast(out), + reinterpret_cast(sfOut)); +} + +} // namespace smooth_quantize_detail + +// Fused smooth + NVFP4 quantize: (out, sf_out) = NVFP4-quantize(in * pqs) in a single pass over +// in, folding the per-input-channel pre_quant_scale smoothing into the quantize. Byte-identical to +// fp4_quantize(in * pqs) (same cvt_warp_fp16_to_fp4 + swizzled SF layout), so the residual GEMM +// consumes the output unchanged. in [m, n] bf16, pqs [n] bf16, sf_scale f32[1] (the per-tensor +// global scale). out [m, n/2] uint8 (packed e2m1), sf_out swizzled UE4M3 block scales (vec size +// 16). SM100+ only. +inline void nvfp4_smooth_quantize(void* out, void* sf_out, void const* in, void const* pqs, + float const* sf_scale, int m, int n, int multiProcessorCount, + cudaStream_t stream, bool enable_pdl) { + using namespace smooth_quantize_detail; + + if (m == 0 || n == 0) return; + + bool const enablePDL = enable_pdl; + bool const useFastPath = (n == 3072 || n == 12288); + if (useFastPath) { + // Same-node SM100 sweeps over the Qwen image-token M values select 192 threads for K=3072 + // and 256 for K=12288. A grid cap of eight CTAs per SM is best for both. + int const blockThreads = n == 3072 ? 192 : 256; + int const blocksPerSm = 8; + + if (n == 3072) + launchSmoothQuantizeFast<3072, 2>(out, sf_out, in, pqs, sf_scale, m, multiProcessorCount, + blockThreads, blocksPerSm, enablePDL, stream); + else + launchSmoothQuantizeFast<12288, 1>(out, sf_out, in, pqs, sf_scale, m, multiProcessorCount, + blockThreads, blocksPerSm, enablePDL, stream); + return; + } + + dim3 block(std::min(n / ELTS_PER_THREAD, 512)); + int const numBlocksPerSM = std::max(1, 2048 / int(block.x)); + dim3 grid(std::min(padUp(m, 128), multiProcessorCount * numBlocksPerSM)); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = grid; + cfg.blockDim = block; + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePDL ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 1; + // No column padding here (n is the padded width); the residual GEMM and the SF layout use n. + cudaLaunchKernelEx(&cfg, smooth_quantize_legacy_kernel, m, n, n, + reinterpret_cast(in), reinterpret_cast(pqs), + sf_scale, reinterpret_cast(out), + reinterpret_cast(sf_out)); +} + +} // namespace gemm +} // namespace flashinfer +// clang-format on diff --git a/test/python/gemm/test_nvfp4_quantize.py b/test/python/gemm/test_nvfp4_quantize.py new file mode 100644 index 000000000..d3f71113b --- /dev/null +++ b/test/python/gemm/test_nvfp4_quantize.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Focused tests for the benchmark-private SM100 NVFP4 quantizer.""" + +import pytest +import torch + +from cudnn.gemm.ops._nvfp4_quantize import nvfp4_quantize + + +def _is_sm100(): + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + + +def _swizzle_128x4(logical, m, k): + """Independent implementation of the documented F8_128x4 byte mapping.""" + rows = (m + 127) // 128 * 128 + cols = (k // 16 + 3) // 4 * 4 + padded = torch.zeros((rows, cols), dtype=torch.uint8, device=logical.device) + padded[:m, : k // 16] = logical + return padded.view(rows // 128, 4, 32, cols // 4, 4).transpose(1, 3).reshape(rows, cols) + + +def _known_byte_problem(m, k): + # After the exact BF16 pqs multiply, every 16-value block is the complete + # E2M1 codebook at scale 1. This stays away from reciprocal/tie ambiguity + # and gives a literal expected packed-byte sequence. + half_values = torch.tensor( + [ + -3.0, + -2.0, + -1.5, + -1.0, + -0.75, + -0.5, + -0.25, + 0.0, + 0.0, + 0.25, + 0.5, + 0.75, + 1.0, + 1.5, + 2.0, + 3.0, + ], + device="cuda", + dtype=torch.bfloat16, + ) + x = half_values.repeat(m, k // 16).contiguous() + pqs = torch.full((k,), 2.0, device="cuda", dtype=torch.bfloat16) + global_scale = torch.ones(1, device="cuda", dtype=torch.float32) + + packed_block = torch.tensor( + [0xEF, 0xCD, 0xAB, 0x09, 0x10, 0x32, 0x54, 0x76], + device="cuda", + dtype=torch.uint8, + ) + expected = packed_block.repeat(m, k // 16) + one_e4m3 = torch.ones((), device="cuda", dtype=torch.float8_e4m3fn).view(torch.uint8) + logical_sf = one_e4m3.expand(m, k // 16) + expected_sf = _swizzle_128x4(logical_sf, m, k) + return x, pqs, global_scale, expected, expected_sf + + +@pytest.mark.L0 +@pytest.mark.skipif(not _is_sm100(), reason="NVFP4 quantization requires SM100") +def test_nvfp4_quantize_generic_exact_bytes(): + x, pqs, global_scale, expected, expected_sf = _known_byte_problem(37, 256) + out, sf = nvfp4_quantize(x, global_scale, pqs) + + assert torch.equal(out, expected) + assert torch.equal(sf, expected_sf) + + +@pytest.mark.L1 +@pytest.mark.skipif(not _is_sm100(), reason="NVFP4 quantization requires SM100") +@pytest.mark.parametrize("m", [1, 512, 4096]) +@pytest.mark.parametrize("k", [3072, 12288]) +def test_nvfp4_quantize_qwen_shapes_exact_bytes(m, k): + x, pqs, global_scale, expected, expected_sf = _known_byte_problem(m, k) + out, sf = nvfp4_quantize(x, global_scale, pqs) + + assert out.shape == (m, k // 2) + assert sf.shape == ((m + 127) // 128 * 128, (k // 16 + 3) // 4 * 4) + assert torch.equal(out, expected) + assert torch.equal(sf, expected_sf) + + +@pytest.mark.L0 +@pytest.mark.skipif(not _is_sm100(), reason="NVFP4 quantization requires SM100") +def test_nvfp4_quantize_optional_pqs_and_caller_outputs(): + torch.manual_seed(7) + x = torch.randn((37, 256), device="cuda", dtype=torch.bfloat16) + global_scale = ((448.0 * 6.0) / x.float().abs().amax()).reshape(()) + + expected, expected_sf = nvfp4_quantize(x, global_scale) + out = torch.empty_like(expected) + sf = torch.empty_like(expected_sf) + got, got_sf = nvfp4_quantize( + x, + global_scale, + torch.ones(256, device="cuda", dtype=torch.bfloat16), + out=out, + scale_factors=sf, + ) + + assert got is out + assert got_sf is sf + assert torch.equal(got, expected) + assert torch.equal(got_sf, expected_sf) + + with pytest.raises(ValueError, match=r"shape \(\) or \(1,\)"): + nvfp4_quantize(x, global_scale.reshape(1, 1)) + + +@pytest.mark.L0 +@pytest.mark.skipif(not _is_sm100(), reason="NVFP4 quantization requires SM100") +def test_nvfp4_quantize_uses_current_stream(monkeypatch): + import cudnn.gemm.ops._nvfp4_quantize as module + + launches = [] + + class FakeExtension: + @staticmethod + def launch(*args): + launches.append(args) + + monkeypatch.setattr(module, "_load_extension", lambda: FakeExtension) + x = torch.empty((37, 256), device="cuda", dtype=torch.bfloat16) + global_scale = torch.ones(1, device="cuda", dtype=torch.float32) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + out, sf = module.nvfp4_quantize(x, global_scale) + + assert launches[-1][-2] == stream.cuda_stream + assert launches[-1][0] == x.data_ptr() + assert launches[-1][3] == out.data_ptr() + assert launches[-1][4] == sf.data_ptr() + + +@pytest.mark.L0 +def test_nvfp4_quantize_rejects_before_lazy_build(monkeypatch): + import cudnn.gemm.ops._nvfp4_quantize as module + + def unexpected_build(): + raise AssertionError("unsupported input reached the lazy compiler") + + monkeypatch.setattr(module, "_load_extension", unexpected_build) + with pytest.raises(ValueError, match="CUDA tensor"): + module.nvfp4_quantize( + torch.empty((1, 256), dtype=torch.bfloat16), + torch.ones(1, dtype=torch.float32), + ) From 50b534238bf0cae8fc7051eec442e92c4f2ddd00 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 9 Sep 2026 13:34:49 -0700 Subject: [PATCH 4/4] fix(bench): separate sampling protocols from GPU product policy --- AGENTS.md | 1 + benchmark/e2e/Qwen-Image/run_bf16.py | 7 +- benchmark/e2e/Qwen-Image/run_nvfp4.py | 7 +- benchmark/e2e/Qwen3.8/run_matrix.py | 18 +--- benchmark/e2e/README.md | 13 ++- benchmark/e2e/tests/test_device_selection.py | 99 ++++++++++++++++++++ 6 files changed, 117 insertions(+), 28 deletions(-) create mode 100644 benchmark/e2e/tests/test_device_selection.py diff --git a/AGENTS.md b/AGENTS.md index 8fad43f9e..b0c08d426 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -82,6 +82,7 @@ First invocation builds the hook environments and can take >5 minutes; later run ## Conventions +- Keep benchmark collection policy separate from runtime support checks: a request to measure on a particular full-SM GPU constrains the agent's measurement/reporting, not user-facing scripts. Gate on actual kernel capabilities, record device/SM-count metadata, and cover a supported non-matching GPU name and SM count in CPU-only device-selection tests. - `include/` is header-only: no `.cpp` files, no new required dependencies. Vendored third-party code lives in `include/cudnn_frontend/thirdparty/`. - Every new frontend-only Python API needs: `APIBase` subclass + wrapper, lazy export in `python/cudnn/__init__.py`, docs under `docs/fe-oss-apis/`, and pytest coverage under `test/python/fe_api/`. Full recipe: [python/cudnn/AGENTS.md](python/cudnn/AGENTS.md) and the `cutedsl-kernel-integration` skill. - Frontend-only OSS APIs are experimental; keep the lazy-import boundary intact (no eager `torch`/`cutlass` imports at `cudnn` import time). CuTeDSL is a required dependency now, but a tensor framework is not, and `import cudnn` still has to stay cheap. diff --git a/benchmark/e2e/Qwen-Image/run_bf16.py b/benchmark/e2e/Qwen-Image/run_bf16.py index 2f78639d5..b0df9a103 100644 --- a/benchmark/e2e/Qwen-Image/run_bf16.py +++ b/benchmark/e2e/Qwen-Image/run_bf16.py @@ -123,12 +123,9 @@ def _pick_device(torch, mode): for index in range(torch.cuda.device_count()): properties = torch.cuda.get_device_properties(index) candidates.append(f"cuda:{index}={properties.name}/{properties.multi_processor_count}SM") - if (properties.major, properties.minor) == (10, 0) and ( - mode != "formal" or (properties.name == "NVIDIA B200" and properties.multi_processor_count == 148) - ): + if (properties.major, properties.minor) == (10, 0): return torch.device(f"cuda:{index}"), properties - requirement = "a full 148-SM NVIDIA B200" if mode == "formal" else "an SM100 GPU" - raise RuntimeError(f"{mode} mode requires {requirement}; visible: {', '.join(candidates)}") + raise RuntimeError(f"{mode} mode requires an SM100 GPU; visible: {', '.join(candidates)}") def _rel_l2(actual, expected): diff --git a/benchmark/e2e/Qwen-Image/run_nvfp4.py b/benchmark/e2e/Qwen-Image/run_nvfp4.py index db9f0df9b..15d7a7604 100644 --- a/benchmark/e2e/Qwen-Image/run_nvfp4.py +++ b/benchmark/e2e/Qwen-Image/run_nvfp4.py @@ -146,12 +146,9 @@ def _pick_device(torch, mode): for index in range(torch.cuda.device_count()): properties = torch.cuda.get_device_properties(index) candidates.append(f"cuda:{index}={properties.name}/{properties.multi_processor_count}SM") - if (properties.major, properties.minor) == (10, 0) and ( - mode != "formal" or (properties.name == "NVIDIA B200" and properties.multi_processor_count == 148) - ): + if (properties.major, properties.minor) == (10, 0): return torch.device(f"cuda:{index}"), properties - requirement = "a full 148-SM NVIDIA B200" if mode == "formal" else "an SM100 GPU" - raise RuntimeError(f"{mode} mode requires {requirement}; visible: {', '.join(candidates)}") + raise RuntimeError(f"{mode} mode requires an SM100 GPU; visible: {', '.join(candidates)}") def _rel_l2(actual, expected): diff --git a/benchmark/e2e/Qwen3.8/run_matrix.py b/benchmark/e2e/Qwen3.8/run_matrix.py index a22b4e55b..f251208de 100644 --- a/benchmark/e2e/Qwen3.8/run_matrix.py +++ b/benchmark/e2e/Qwen3.8/run_matrix.py @@ -241,7 +241,7 @@ def _parse_args(): "--mode", choices=MODE_DEFAULTS, default="smoke", - help="smoke keeps Qwen kernel dimensions but reduces tokens; formal requires a full 148-SM B200", + help="smoke keeps Qwen kernel dimensions but reduces tokens; formal uses the full workload and sampling protocol", ) parser.add_argument("--preset", default=DEFAULT_PRESET) parser.add_argument("--layers", type=int) @@ -366,20 +366,12 @@ def _serializable_args(args): def _pick_device(mode): visible = [] - smoke_candidate = None for index in range(torch.cuda.device_count()): properties = torch.cuda.get_device_properties(index) visible.append(f"cuda:{index}={properties.name}/{properties.multi_processor_count}SM") - if properties.name == "NVIDIA B200" and properties.multi_processor_count == 148: - if mode == "formal": - return torch.device(f"cuda:{index}") - smoke_candidate = smoke_candidate or torch.device(f"cuda:{index}") - elif 100 <= properties.major * 10 + properties.minor < 120 and smoke_candidate is None: - smoke_candidate = torch.device(f"cuda:{index}") - if mode == "smoke" and smoke_candidate is not None: - return smoke_candidate - requirement = "a full 148-SM NVIDIA B200" if mode == "formal" else "an SM100-family Blackwell GPU" - raise RuntimeError(f"{mode} mode requires {requirement}; visible devices: " + ", ".join(visible)) + if 100 <= properties.major * 10 + properties.minor < 120: + return torch.device(f"cuda:{index}") + raise RuntimeError(f"{mode} mode requires an SM100-family Blackwell GPU; visible devices: " + ", ".join(visible)) def _run_experiment(args, qwen, device, properties, orders, started_utc): @@ -939,8 +931,6 @@ def main(): # Torch-versus-FE selector used by this matrix. qwen = _load_run_model() properties = torch.cuda.get_device_properties(device) - if args.mode == "formal" and (properties.name != "NVIDIA B200" or properties.multi_processor_count != 148): - raise RuntimeError("formal mode requires a full 148-SM NVIDIA B200, got " f"{properties.name}, {properties.multi_processor_count} SMs on {device}") _run_experiment(args, qwen, device, properties, orders, _utc_now()) diff --git a/benchmark/e2e/README.md b/benchmark/e2e/README.md index c1b954667..2d5b491a1 100644 --- a/benchmark/e2e/README.md +++ b/benchmark/e2e/README.md @@ -258,6 +258,12 @@ Planned: Kimi Linear (KDA), DeepSeek-V3. ## Run +`--mode` selects the workload and sampling protocol, not a GPU product or SM +count. The runners retain their kernel architecture/backend requirements and +record the selected device and SM count in each artifact. Performance results +apply to that recorded hardware; selecting `formal` does not certify a B200 run. +Use `CUDA_VISIBLE_DEVICES` to select the GPU for a controlled measurement. + ```bash # The factorial math/reporting tests are CPU-only and do not import Torch. python -m unittest discover -s benchmark/e2e/tests -v @@ -268,7 +274,7 @@ python -m unittest discover -s benchmark/e2e/tests -v python benchmark/e2e/Qwen3.8/run_matrix.py \ --mode smoke --output-dir qwen3.8-factorial-results/smoke -# Formal 8-arm Williams run: full 148-SM B200, bs=4, seq=2048, +# Formal 8-arm Williams run: bs=4, seq=2048, # 40 balanced batches, 3 repeats. Produces timestamped .json and .md artifacts. python benchmark/e2e/Qwen3.8/run_matrix.py \ --mode formal --output-dir qwen3.8-factorial-results/formal @@ -305,7 +311,7 @@ python benchmark/e2e/Qwen-Image/run_bf16.py \ --mode smoke --output-dir qwen-image-bf16-results/smoke # Formal four-arm BF16 transformer proxy: B=1, 4096 image + 512 text tokens, -# four real-shape blocks, 40 Williams-balanced batches x 3 repeats on a full B200. +# four real-shape blocks, 40 Williams-balanced batches x 3 repeats. python benchmark/e2e/Qwen-Image/run_bf16.py \ --mode formal --output-dir qwen-image-bf16-results/formal @@ -326,8 +332,7 @@ cuDNN >= 9.23 backend d256 SDPA path on an SM100 (Blackwell) device; and admits its validated plain, local, bias-free BF16 `swish` module; unsupported runtime configurations fall back to FLA. The default `short_conv=true` GDN axis uses the packed-QKV support landed in #685. After -#682, the FE public d256 SDPA operator is backend-graph-only. Formal mode -additionally rejects anything other than a full 148-SM NVIDIA B200. Both modes verify exact +#682, the FE public d256 SDPA operator is backend-graph-only. Both modes verify exact GDN/MLP/full-attention routes and explicit finite correctness before reporting. Formal shape/timing overrides are allowed but are listed prominently and included in the comparability fingerprint. The default diff --git a/benchmark/e2e/tests/test_device_selection.py b/benchmark/e2e/tests/test_device_selection.py new file mode 100644 index 000000000..7a7770d20 --- /dev/null +++ b/benchmark/e2e/tests/test_device_selection.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only regressions separating benchmark protocols from GPU admission.""" + +from contextlib import nullcontext +import importlib.util +from pathlib import Path +import sys +from types import SimpleNamespace +import unittest +from unittest.mock import patch + +E2E_DIR = Path(__file__).resolve().parents[1] + + +def load(relative_path): + name = "device_selection_" + relative_path.replace("/", "_").replace(".", "_").replace("-", "_") + spec = importlib.util.spec_from_file_location(name, E2E_DIR / relative_path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +MATRIX = load("Qwen3.8/run_matrix.py") +RUNNERS = (MATRIX, load("Qwen-Image/run_bf16.py"), load("Qwen-Image/run_nvfp4.py")) + + +def gpu(name="SM100 test device", sm_count=42, capability=(10, 0)): + return SimpleNamespace(name=name, multi_processor_count=sm_count, major=capability[0], minor=capability[1]) + + +def fake_torch(devices): + return SimpleNamespace( + device=str, + cuda=SimpleNamespace( + is_available=lambda: bool(devices), + device_count=lambda: len(devices), + get_device_properties=lambda index: devices[int(str(index).removeprefix("cuda:"))], + device=lambda _device: nullcontext(), + ), + ) + + +def pick(runner, devices, mode): + torch = fake_torch(devices) + if runner is MATRIX: + with patch.object(runner, "torch", torch): + return runner._pick_device(mode) + return runner._pick_device(torch, mode)[0] + + +class DeviceSelectionTest(unittest.TestCase): + def test_both_modes_accept_supported_architecture_without_sku_or_sm_count_gate(self): + for runner in RUNNERS: + for mode in ("smoke", "formal"): + for properties in (gpu(), gpu("NVIDIA B200", 42), gpu("NVIDIA B200", 148)): + with self.subTest(runner=runner.__name__, mode=mode, gpu=properties): + self.assertEqual(pick(runner, [properties], mode), "cuda:0") + + def test_product_name_cannot_override_unsupported_architecture(self): + for runner in RUNNERS: + for mode in ("smoke", "formal"): + for capability in ((8, 0), (9, 0), (12, 0)): + with self.subTest(runner=runner.__name__, mode=mode, capability=capability): + with self.assertRaisesRegex(RuntimeError, "SM100"): + pick(runner, [gpu("NVIDIA B200", 148, capability)], mode) + + def test_selects_supported_device_in_mixed_inventory(self): + for runner in RUNNERS: + for mode in ("smoke", "formal"): + with self.subTest(runner=runner.__name__, mode=mode): + self.assertEqual(pick(runner, [gpu(capability=(9, 0)), gpu()], mode), "cuda:1") + + def test_no_visible_gpu_has_actionable_error(self): + for runner in RUNNERS: + for mode in ("smoke", "formal"): + with self.subTest(runner=runner.__name__, mode=mode): + with self.assertRaisesRegex(RuntimeError, "SM100"): + pick(runner, [], mode) + + def test_matrix_main_has_no_second_formal_hardware_policy_gate(self): + torch = fake_torch([gpu()]) + args = SimpleNamespace(mode="formal") + with ( + patch.object(MATRIX, "_parse_args", return_value=args), + patch.object(MATRIX.importlib, "import_module", return_value=torch), + patch.dict(MATRIX.os.environ, {"CUDNN_FRONTEND_ENABLE_FROST_ENGINES": "0"}), + patch.object(MATRIX, "_load_run_model", return_value=object()), + patch.object(MATRIX, "_run_experiment") as run, + ): + MATRIX.main() + self.assertEqual(run.call_count, 1) + self.assertEqual(run.call_args.args[2], "cuda:0") + + +if __name__ == "__main__": + unittest.main()