From 07040ffe88ff1181091701e3b3bfaaf2bcc07196 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:11:16 +0000 Subject: [PATCH 1/2] Fix FP8 QDQ for delegated attention Co-authored-by: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/onnx/export/fp8_exporter.py | 167 ++++++++- .../torch/quantization/plugins/attention.py | 171 +++++---- .../plugins/diffusion/diffusers.py | 40 ++- .../quantization/test_fp8_mha_exporter.py | 140 +++++++- .../plugins/test_diffusers_attention.py | 332 ++++++++++++++++++ 6 files changed, 775 insertions(+), 76 deletions(-) create mode 100644 tests/unit/torch/quantization/plugins/test_diffusers_attention.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3e10c786753..da0c237f239 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Changelog **Bug Fixes** +- Fix FP8 MHA quantization for diffusion attention modules that delegate to a module-level helper, and preserve fusion-ready Q/DQ placement around attention scaling, key transposes, and softmax outputs in exported ONNX graphs. - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 8e94f40f332..0661d101330 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -32,6 +32,8 @@ # when using 1/448 as the Q scale (single fixed value — softmax range is data-independent). _FP8_E4M3_MAX = 448.0 _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX +_ELEMENTWISE_SCALAR_OPS = {"Add", "Div", "Mul", "Pow", "Sub"} +_UNARY_SCALAR_OPS = {"Cast", "Reciprocal", "Sqrt"} def _torch_from_numpy_for_fp8(array: np.ndarray) -> torch.Tensor: @@ -41,6 +43,144 @@ def _torch_from_numpy_for_fp8(array: np.ndarray) -> torch.Tensor: return torch.from_numpy(array) +def _constant_values(tensor: gs.Tensor) -> np.ndarray | None: + if isinstance(tensor, gs.Constant): + return tensor.values + if not isinstance(tensor, gs.Variable) or len(tensor.inputs) != 1: + return None + producer = tensor.inputs[0] + if producer.op != "Constant": + return None + value = producer.attrs.get("value") + return value.values if isinstance(value, gs.Constant) else None + + +def _shape_vector_length(tensor: gs.Tensor) -> int | None: + if not isinstance(tensor, gs.Variable) or len(tensor.inputs) != 1: + return None + shape_node = tensor.inputs[0] + if shape_node.op != "Shape" or len(shape_node.inputs) != 1: + return None + source_shape = shape_node.inputs[0].shape + if source_shape is None: + return None + rank = len(source_shape) + start = shape_node.attrs.get("start", 0) + end = shape_node.attrs.get("end", rank) + return len(range(*slice(start, end).indices(rank))) + + +def _is_data_independent_scale(tensor: gs.Tensor, memo=None) -> bool: + """Return whether a tensor is provably scalar and independent of data values.""" + if memo is None: + memo = {} + tensor_id = id(tensor) + if tensor_id in memo: + return memo[tensor_id] + memo[tensor_id] = False + + values = _constant_values(tensor) + if values is not None: + memo[tensor_id] = values.size == 1 + return memo[tensor_id] + if not isinstance(tensor, gs.Variable) or len(tensor.inputs) != 1: + return False + + producer = tensor.inputs[0] + inputs = producer.inputs + result = False + if producer.op == "Gather" and len(inputs) == 2: + indices = _constant_values(inputs[1]) + if ( + _shape_vector_length(inputs[0]) is not None + and producer.attrs.get("axis", 0) in (0, -1) + and indices is not None + and indices.size == 1 + ): + result = True + elif producer.op == "Slice" and 3 <= len(inputs) <= 5: + vector_length = _shape_vector_length(inputs[0]) + starts = _constant_values(inputs[1]) + ends = _constant_values(inputs[2]) + axes = _constant_values(inputs[3]) if len(inputs) >= 4 else np.array([0]) + steps = _constant_values(inputs[4]) if len(inputs) == 5 else np.array([1]) + if ( + vector_length is not None + and starts is not None + and ends is not None + and axes is not None + and steps is not None + and all(value.size == 1 for value in (starts, ends, axes, steps)) + and axes.item() in (0, -1) + and steps.item() != 0 + ): + result = ( + len(range(*slice(starts.item(), ends.item(), steps.item()).indices(vector_length))) + == 1 + ) + elif producer.op in _UNARY_SCALAR_OPS and len(inputs) == 1: + result = _is_data_independent_scale(inputs[0], memo) + elif producer.op in _ELEMENTWISE_SCALAR_OPS and inputs: + result = all(_is_data_independent_scale(input_tensor, memo) for input_tensor in inputs) + + memo[tensor_id] = result + return result + + +def _rebase_shape_sources(tensor, old_tensors, new_tensor): + """Rebase direct ``Shape(Q/DQ)`` inputs, rejecting dependencies that would cycle.""" + dependency_memo = {} + + def is_old_tensor(candidate): + return any(candidate is old_tensor for old_tensor in old_tensors) + + def depends_on_old_tensor(candidate): + if is_old_tensor(candidate): + return True + if not isinstance(candidate, gs.Variable): + return False + candidate_id = id(candidate) + if candidate_id in dependency_memo: + return dependency_memo[candidate_id] + dependency_memo[candidate_id] = False + dependency_memo[candidate_id] = any( + depends_on_old_tensor(producer_input) + for producer in candidate.inputs + for producer_input in producer.inputs + ) + return dependency_memo[candidate_id] + + shape_nodes = [] + visited = set() + + def collect_shape_nodes(candidate): + if not isinstance(candidate, gs.Variable) or id(candidate) in visited: + return True + visited.add(id(candidate)) + for producer in candidate.inputs: + if producer.op == "Shape": + if any( + not is_old_tensor(input_tensor) and depends_on_old_tensor(input_tensor) + for input_tensor in producer.inputs + ): + return False + if any(is_old_tensor(input_tensor) for input_tensor in producer.inputs): + shape_nodes.append(producer) + continue + if not all(collect_shape_nodes(producer_input) for producer_input in producer.inputs): + return False + return True + + if not collect_shape_nodes(tensor): + return False + for shape_node in shape_nodes: + shape_node.inputs = [ + new_tensor if is_old_tensor(input_tensor) else input_tensor + for input_tensor in shape_node.inputs + ] + return True + + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -253,23 +393,24 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: @staticmethod def _move_mul_before_qdq(graph: gs.Graph) -> int: - """Move attention-scaling Mul(const) from after DQ to before Q for TRT MatMul fusion. + """Move attention scaling from after DQ to before Q for TRT MatMul fusion. - Handles both ``DQ → Mul → MatMul`` and ``DQ → Transpose → Mul → MatMul`` (K path). + The scale must be scalar and data-independent. Handles both ``DQ → Mul → MatMul`` + and ``DQ → Transpose → Mul → MatMul`` (K path). """ count = 0 for mul_node in list(graph.nodes): if mul_node.op != "Mul": continue - const_input = next( - (i for i in mul_node.inputs if isinstance(i, gs.Constant) and i.values.size == 1), + scale_input = next( + (tensor for tensor in mul_node.inputs if _is_data_independent_scale(tensor)), None, ) tensor_input = next( - (i for i in mul_node.inputs if not isinstance(i, gs.Constant)), None + (tensor for tensor in mul_node.inputs if tensor is not scale_input), None ) - if const_input is None or tensor_input is None: + if scale_input is None or tensor_input is None: continue if not (isinstance(tensor_input, gs.Variable) and len(tensor_input.inputs) == 1): continue @@ -289,6 +430,7 @@ def _move_mul_before_qdq(graph: gs.Graph) -> int: continue q_output = dq_node.inputs[0] + dq_output = dq_node.outputs[0] if ( not isinstance(q_output, gs.Variable) or len(q_output.inputs) != 1 @@ -307,6 +449,13 @@ def _move_mul_before_qdq(graph: gs.Graph) -> int: if not mul_consumers or not all(c.op == "MatMul" for c in mul_consumers): continue + if not _rebase_shape_sources(scale_input, (q_output, dq_output), q_input): + continue + if q_input.shape is None: + q_input.shape = dq_output.shape + if q_input.dtype is None: + q_input.dtype = dq_output.dtype + new_mul_output = gs.Variable( q_input.name + "_scaled", dtype=q_input.dtype, shape=q_input.shape ) @@ -314,15 +463,13 @@ def _move_mul_before_qdq(graph: gs.Graph) -> int: gs.Node( op="Mul", name=mul_node.name + "_moved", - inputs=[q_input, const_input], + inputs=[q_input, scale_input], outputs=[new_mul_output], ) ) q_node.inputs[0] = new_mul_output - replacement = ( - transpose_node.outputs[0] if transpose_node is not None else dq_node.outputs[0] - ) + replacement = transpose_node.outputs[0] if transpose_node is not None else dq_output for consumer in mul_consumers: for i, inp in enumerate(consumer.inputs): if inp is mul_output: diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 2113edea8a7..235e36dc367 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -34,6 +34,7 @@ import inspect import tempfile import types +from typing import cast from warnings import warn from ..conversion import register @@ -50,6 +51,8 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: """ source_code = inspect.getsource(attention_cls) model_module = inspect.getmodule(attention_cls) + if model_module is None: + return False head = ast.parse(source_code) bmm_ops = ("matmul", "bmm", "baddbmm") @@ -72,6 +75,48 @@ def is_sdpa(node): def is_bin_matmul(node): return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) + def get_attention_helper_parameters(node): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and callable(helper := model_module.__dict__.get(node.func.id)) + ): + return None + + helper_name = getattr(helper, "__name__", node.func.id).lower() + if "attention" not in helper_name and "attn" not in helper_name: + return None + + try: + positional_parameters = [ + parameter.name + for parameter in inspect.signature(helper).parameters.values() + if parameter.kind + in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + ] + except (TypeError, ValueError): + return None + + keyword_names = {keyword.arg for keyword in node.keywords} + for parameter_names in (("q", "k", "v"), ("query", "key", "value")): + if tuple(positional_parameters[:3]) == parameter_names and all( + name in keyword_names or index < len(node.args) + for index, name in enumerate(parameter_names) + ): + return parameter_names + return None + + def quantize_argument(argument, quantizer_name): + return ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_name, + ctx=ast.Load(), + ), + args=[argument], + keywords=[], + ) + def patch(node, quantizer_names, transpose=False): for index, quantizer_name in enumerate(quantizer_names): if quantizer_name is None: @@ -79,33 +124,16 @@ def patch(node, quantizer_names, transpose=False): arg = node.args[index] if not transpose: - node.args[index] = ast.Call( - func=ast.Attribute( - value=ast.Name(id="self", ctx=ast.Load()), - attr=quantizer_name, - ctx=ast.Load(), - ), - args=[arg], + node.args[index] = quantize_argument(arg, quantizer_name) + else: + transposed_arg = ast.Call( + func=ast.Attribute(value=arg, attr="transpose", ctx=ast.Load()), + args=[ast.Constant(value=-1), ast.Constant(value=-2)], keywords=[], ) - else: node.args[index] = ast.Call( func=ast.Attribute( - value=ast.Call( - func=ast.Attribute( - value=ast.Name(id="self", ctx=ast.Load()), - attr=quantizer_name, - ctx=ast.Load(), - ), - args=[ - ast.Call( - func=ast.Attribute(value=arg, attr="transpose", ctx=ast.Load()), - args=[ast.Constant(value=-1), ast.Constant(value=-2)], - keywords=[], - ) - ], - keywords=[], - ), + value=quantize_argument(transposed_arg, quantizer_name), attr="transpose", ctx=ast.Load(), ), @@ -113,18 +141,22 @@ def patch(node, quantizer_names, transpose=False): keywords=[], ) + def patch_attention_helper(node, parameter_names): + keyword_by_name = { + keyword.arg: keyword for keyword in node.keywords if keyword.arg is not None + } + for index, (parameter_name, quantizer_name) in enumerate( + zip(parameter_names, ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer")) + ): + if keyword := keyword_by_name.get(parameter_name): + keyword.value = quantize_argument(keyword.value, quantizer_name) + else: + node.args[index] = quantize_argument(node.args[index], quantizer_name) + def patch_binop(node, quantizer_names, transpose=False): assert len(quantizer_names) == 2 if quantizer_names[0] is not None: - node.left = ast.Call( - func=ast.Attribute( - value=ast.Name(id="self", ctx=ast.Load()), - attr=quantizer_names[0], - ctx=ast.Load(), - ), - args=[node.left], - keywords=[], - ) + node.left = quantize_argument(node.left, quantizer_names[0]) if quantizer_names[1] is not None: arg = node.right if transpose: @@ -137,15 +169,7 @@ def patch_binop(node, quantizer_names, transpose=False): args=[arg, ast.Constant(value=-1), ast.Constant(value=-2)], keywords=[], ) - quant_arg = ast.Call( - func=ast.Attribute( - value=ast.Name(id="self", ctx=ast.Load()), - attr=quantizer_names[1], - ctx=ast.Load(), - ), - args=[arg], - keywords=[], - ) + quant_arg = quantize_argument(arg, quantizer_names[1]) if transpose: quant_arg = ast.Call( func=ast.Attribute( @@ -158,28 +182,43 @@ def patch_binop(node, quantizer_names, transpose=False): ) node.right = quant_arg - nodes = list(ast.walk(head)) - org_class_name = nodes[1].name # type: ignore[attr-defined] - new_class_name = nodes[1].name = "_Quant" + nodes[1].name # type: ignore[attr-defined] - - bmm_nodes = [] - sdpa_nodes = [] - bin_matmul_nodes = [] - for node in ast.walk(head): - if is_bmm(node): - bmm_nodes.append(node) - if is_sdpa(node): - sdpa_nodes.append(node) - if is_bin_matmul(node): - bin_matmul_nodes.append(node) - if len(bmm_nodes) != 2 and len(sdpa_nodes) != 1 and len(bin_matmul_nodes) != 2: + class_node = next(node for node in head.body if isinstance(node, ast.ClassDef)) + org_class_name = class_node.name + new_class_name = class_node.name = "_Quant" + class_node.name + forward_node = next( + ( + node + for node in class_node.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "forward" + ), + None, + ) + attention_helper_nodes = [ + (node, parameter_names) + for node in (ast.walk(forward_node) if forward_node is not None else ()) + if (parameter_names := get_attention_helper_parameters(node)) is not None + ] + + direct_nodes = list(ast.walk(head)) + bmm_nodes = [cast("ast.Call", node) for node in direct_nodes if is_bmm(node)] + sdpa_nodes = [cast("ast.Call", node) for node in direct_nodes if is_sdpa(node)] + bin_matmul_nodes = [cast("ast.BinOp", node) for node in direct_nodes if is_bin_matmul(node)] + patch_bmm = len(bmm_nodes) == 2 and all(len(node.args) >= 2 for node in bmm_nodes) + patch_sdpa = len(sdpa_nodes) == 1 and len(sdpa_nodes[0].args) >= 3 + patch_bin_matmul = len(bin_matmul_nodes) == 2 + if not attention_helper_nodes and not (patch_bmm or patch_sdpa or patch_bin_matmul): print(f"Expect 2 bmm/matmul op in the {org_class_name}, found {len(bmm_nodes)}") print(f"Or expect 1 sdpa op in the {org_class_name}, found {len(sdpa_nodes)}") print(f"Or expect 2 @ op in the {org_class_name}, found {len(bin_matmul_nodes)}") + print(f"Or expect an attention helper call in the {org_class_name}") print("Auto quantization of KV Cache fails") return False - if len(bmm_nodes) == 2: + if attention_helper_nodes: + for node, parameter_names in attention_helper_nodes: + patch_attention_helper(node, parameter_names) + print(f"Patching {len(attention_helper_nodes)} attention helper call(s) with quantizers") + if patch_bmm: # transpose k cache here to enable per-token quantization # without transpose, the quantization will be per-channel, i.e., # self.k_bmm_quantizer(key_states.transpose(-1, -2)) @@ -187,9 +226,13 @@ def patch_binop(node, quantizer_names, transpose=False): # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) # removing the additional transpose is doable but not trivial patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) - patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch( + bmm_nodes[1], + quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), + transpose=True, + ) print("Patching 2 BMM/Matmul operators with quantizers") - if len(bin_matmul_nodes) == 2: + if patch_bin_matmul: patch_binop( bin_matmul_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), @@ -197,10 +240,14 @@ def patch_binop(node, quantizer_names, transpose=False): ) patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 @ operators with quantizers") - - if len(sdpa_nodes) == 1: + if patch_sdpa: patch( - sdpa_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") + sdpa_nodes[0], + quantizer_names=( + "q_bmm_quantizer", + "k_bmm_quantizer", + "v_bmm_quantizer", + ), ) print("Patching 1 scaled_dot_product_attention operator with quantizers") diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..5b31e355412 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -16,6 +16,7 @@ """Support quantization of diffusers layers.""" from collections.abc import Callable, Iterator +from contextlib import suppress from functools import partial from types import ModuleType from typing import TYPE_CHECKING @@ -25,6 +26,7 @@ import torch from diffusers.models.attention_processor import Attention from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from diffusers.models.modeling_utils import ModelMixin from packaging.version import parse as parse_version if parse_version(diffusers.__version__) >= parse_version("0.35.0"): @@ -64,7 +66,8 @@ TensorQuantizer, ) from ...nn.modules.quant_conv import _QuantConv3d -from ..custom import _QuantFunctionalMixin +from ..attention import register_attention_for_kv_quant +from ..custom import CUSTOM_MODEL_PLUGINS, _QuantFunctionalMixin onnx_dtype_map = { "BFloat16": onnx.TensorProto.BFLOAT16, @@ -208,6 +211,41 @@ def forward(self, *args, **kwargs): ) +def _try_register_attention(attention_cls): + with suppress( + IndentationError, IndexError, KeyError, OSError, SyntaxError, TypeError, ValueError + ): + register_attention_for_kv_quant(attention_cls) + + +def _register_diffusers_attentions_on_the_fly(model): + """Register unrecognized leaf attention modules in diffusers model graphs.""" + if not isinstance(model, ModelMixin): + return + + seen_classes = set() + for module in model.modules(): + module_type = type(module) + if ( + not module_type.__name__.endswith("Attention") + or module_type in seen_classes + or module_type in QuantModuleRegistry + or module_type.__module__.startswith("transformers.") + or hasattr(module_type, "_setup") + or not all(isinstance(getattr(module, name, None), torch.nn.Module) for name in "qkv") + or any( + child is not module and type(child).__name__.endswith("Attention") + for child in module.modules() + ) + ): + continue + seen_classes.add(module_type) + _try_register_attention(module_type) + + +CUSTOM_MODEL_PLUGINS.add(_register_diffusers_attentions_on_the_fly) + + original_scaled_dot_product_attention = F.scaled_dot_product_attention diff --git a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py index 1f7251a9ad9..7edb947fa1f 100644 --- a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py +++ b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py @@ -22,14 +22,15 @@ from modelopt.onnx.export.fp8_exporter import FP8QuantExporter -def _var(name): - return gs.Variable(name, dtype=np.float32) +def _var(name, shape=None): + return gs.Variable(name, dtype=np.float32, shape=shape) def _qdq(src): """Build ``QuantizeLinear → DequantizeLinear`` and return [Q, DQ], dq_out.""" scale = gs.Constant("scale", np.array(0.1, dtype=np.float32)) - q_out, dq_out = _var("q_out"), _var("dq_out") + q_out = _var("q_out", src.shape) + dq_out = _var("dq_out", src.shape) return [ gs.Node(op="QuantizeLinear", inputs=[src, scale], outputs=[q_out]), gs.Node(op="DequantizeLinear", inputs=[q_out, scale], outputs=[dq_out]), @@ -57,6 +58,139 @@ def test_move_mul_before_qdq_rewrites_dq_mul_matmul_pattern(): assert q.inputs[0].inputs[0].op == "Mul" +def test_move_mul_before_qdq_rewrites_constant_expression_scale(): + """The scalar scale may be a constant expression emitted by the Torch exporter.""" + x, k, y, mul_out = _var("x"), _var("k"), _var("y"), _var("mul_out") + constant_out, sqrt_out = _var("constant_out"), _var("sqrt_out") + qdq_nodes, dq_out = _qdq(x) + constant = gs.Node( + op="Constant", + attrs={"value": gs.Constant("constant", np.array(0.25, dtype=np.float32))}, + outputs=[constant_out], + ) + sqrt = gs.Node(op="Sqrt", inputs=[constant_out], outputs=[sqrt_out]) + mul = gs.Node(op="Mul", inputs=[dq_out, sqrt_out], outputs=[mul_out]) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph([*qdq_nodes, constant, sqrt, mul, mm], [x, k], [y]) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert q.inputs[0].inputs[0].op == "Mul" + + +@pytest.mark.parametrize("shape_source", ["q", "dq"]) +@pytest.mark.parametrize("indirect", [False, True]) +def test_move_mul_before_qdq_handles_shape_dependencies(shape_source, indirect): + x, k, y, identity_out, shape_out, slice_out, scale, mul_out = ( + _var("x", [1, 2, 3, 4]), + _var("k"), + _var("y"), + _var("identity_out", [1, 2, 3, 4]), + _var("shape_out"), + _var("slice_out"), + _var("shape_scale"), + _var("mul_out"), + ) + qdq_nodes, dq_out = _qdq(x) + q_out = qdq_nodes[0].outputs[0] + shape_input = q_out if shape_source == "q" else dq_out + pre_shape_nodes = [] + if indirect: + pre_shape_nodes.append(gs.Node(op="Identity", inputs=[shape_input], outputs=[identity_out])) + shape_input = identity_out + shape = gs.Node(op="Shape", inputs=[shape_input], outputs=[shape_out]) + starts, ends = _var("starts"), _var("ends") + slice_constants = [ + gs.Node( + op="Constant", + attrs={"value": gs.Constant("starts_value", np.array([-1], dtype=np.int64))}, + outputs=[starts], + ), + gs.Node( + op="Constant", + attrs={ + "value": gs.Constant( + "ends_value", np.array([np.iinfo(np.int64).max], dtype=np.int64) + ) + }, + outputs=[ends], + ), + ] + shape_slice = gs.Node(op="Slice", inputs=[shape_out, starts, ends], outputs=[slice_out]) + sqrt = gs.Node(op="Sqrt", inputs=[slice_out], outputs=[scale]) + mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph( + [*qdq_nodes, *pre_shape_nodes, shape, *slice_constants, shape_slice, sqrt, mul, mm], + [x, k], + [y], + ) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == (0 if indirect else 1) + assert shape.inputs == [identity_out if indirect else x] + + +def test_move_mul_before_qdq_skips_vector_shape_scale(): + x, k, y, shape_out, scale, mul_out = ( + _var("x", [1, 2, 3, 4]), + _var("k"), + _var("y"), + _var("shape_out"), + _var("shape_scale"), + _var("mul_out"), + ) + qdq_nodes, dq_out = _qdq(x) + shape = gs.Node(op="Shape", inputs=[dq_out], outputs=[shape_out]) + cast = gs.Node(op="Cast", inputs=[shape_out], outputs=[scale], attrs={"to": 1}) + mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph([*qdq_nodes, shape, cast, mul, mm], [x, k], [y]) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == 0 + assert shape.inputs == [dq_out] + + +def test_move_mul_before_qdq_rewrites_gathered_shape_scale(): + x, k, y, shape_out, scale, mul_out = ( + _var("x", [1, 2, 3, 4]), + _var("k"), + _var("y"), + _var("shape_out"), + _var("shape_scale"), + _var("mul_out"), + ) + qdq_nodes, dq_out = _qdq(x) + shape = gs.Node(op="Shape", inputs=[dq_out], outputs=[shape_out]) + gather = gs.Node( + op="Gather", + inputs=[shape_out, gs.Constant("index", np.array(-1, dtype=np.int64))], + outputs=[scale], + attrs={"axis": 0}, + ) + mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph([*qdq_nodes, shape, gather, mul, mm], [x, k], [y]) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == 1 + assert shape.inputs == [x] + + +def test_move_mul_before_qdq_skips_data_dependent_scale(): + x, scale, k, y, mul_out = ( + _var("x"), + _var("dynamic_scale"), + _var("k"), + _var("y"), + _var("mul_out"), + ) + qdq_nodes, dq_out = _qdq(x) + mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph([*qdq_nodes, mul, mm], [x, scale, k], [y]) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == 0 + + def test_move_transpose_before_qdq_rewrites_dq_transpose_matmul_pattern(): """``DQ → Transpose → MatMul`` collapses to ``Transpose → Q → DQ → MatMul``.""" k_in, q_in, scores, t_out = _var("k_in"), _var("q_in"), _var("scores"), _var("t_out") diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_attention.py b/tests/unit/torch/quantization/plugins/test_diffusers_attention.py new file mode 100644 index 00000000000..b6cbffe2de4 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_diffusers_attention.py @@ -0,0 +1,332 @@ +# 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. + +"""Tests for on-the-fly registration of delegated diffusion attention.""" + +import copy +import io +from collections import defaultdict + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +pytest.importorskip("diffusers") +from diffusers import ModelMixin + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.plugins.attention import register_attention_for_kv_quant + +_QKV_QUANTIZER_NAMES = ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") + + +def delegated_attention(q, k, v): + return F.scaled_dot_product_attention(q, k, v) + + +class DelegatedAttention(nn.Module): + def __init__(self, hidden_size=16, num_heads=2): + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.q = nn.Linear(hidden_size, hidden_size) + self.k = nn.Linear(hidden_size, hidden_size) + self.v = nn.Linear(hidden_size, hidden_size) + self.o = nn.Linear(hidden_size, hidden_size) + + def _reshape(self, hidden_states): + batch_size, sequence_length, _ = hidden_states.shape + return hidden_states.view( + batch_size, sequence_length, self.num_heads, self.head_dim + ).transpose(1, 2) + + def _unused_attention(self, q, k, v): + scores = torch.matmul(q, k) + return torch.matmul(scores, v) + + def forward(self, hidden_states): + q = self._reshape(self.q(hidden_states)) + k = self._reshape(self.k(hidden_states)) + v = self._reshape(self.v(hidden_states)) + output = delegated_attention(q=q * 0.5, k=k * 0.25, v=v) + return self.o(output.transpose(1, 2).flatten(2)) + + +class PositionalDelegatedAttention(DelegatedAttention): + def forward(self, hidden_states): + q = self._reshape(self.q(hidden_states)) + k = self._reshape(self.k(hidden_states)) + v = self._reshape(self.v(hidden_states)) + output = delegated_attention(q * 0.5, k * 0.25, v) + return self.o(output.transpose(1, 2).flatten(2)) + + +class RepeatedDelegatedAttention(DelegatedAttention): + def forward(self, hidden_states): + q = self._reshape(self.q(hidden_states)) + k = self._reshape(self.k(hidden_states)) + v = self._reshape(self.v(hidden_states)) + output = delegated_attention(q, k, v) + output = delegated_attention(output, k, v) + return self.o(output.transpose(1, 2).flatten(2)) + + +class HelperWithMethodMatmulAttention(DelegatedAttention): + def forward(self, hidden_states): + q = self._reshape(self.q(hidden_states)) + k = self._reshape(self.k(hidden_states)) + v = self._reshape(self.v(hidden_states)) + if hidden_states.shape[0] == 0: + scores = q.matmul(k.transpose(-2, -1)) + output = scores.matmul(v) + else: + output = delegated_attention(q, k, v) + return self.o(output.transpose(1, 2).flatten(2)) + + +class AuxiliaryHelperAttention(nn.Module): + def forward(self, hidden_states): + return hidden_states + + def _unused_attention(self, q, k, v): + return delegated_attention(q, k, v) + + +class MethodMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = q.matmul(k.transpose(-2, -1)) + return scores.matmul(v) + + +class CustomSetupDelegatedAttention(DelegatedAttention): + def _setup(self): + pass + + def forward(self, hidden_states): + q = self._reshape(self.q(hidden_states)) + k = self._reshape(self.k(hidden_states)) + v = self._reshape(self.v(hidden_states)) + output = delegated_attention(q, k, v) + return self.o(output.transpose(1, 2).flatten(2)) + + +class DelegatedDiffusionModel(ModelMixin): + def __init__(self, attention_cls): + super().__init__() + self.attn = attention_cls() + + def forward(self, hidden_states): + return self.attn(hidden_states) + + +def _get_fp8_attention_config(): + quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + quant_config["quant_cfg"].append( + { + "quantizer_name": "*[qkv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + } + ) + return quant_config + + +@pytest.mark.parametrize( + ("attention_cls", "expected_calls"), + [ + (DelegatedAttention, 1), + (PositionalDelegatedAttention, 1), + (RepeatedDelegatedAttention, 2), + (HelperWithMethodMatmulAttention, 1), + ], +) +def test_quantize_registers_delegated_attention(attention_cls, expected_calls): + model = DelegatedDiffusionModel(attention_cls) + inputs = torch.randn(2, 4, 16) + + try: + mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) + + for name in _QKV_QUANTIZER_NAMES: + quantizer = getattr(model.attn, name) + assert quantizer.is_enabled + assert quantizer.amax is not None + + call_counts = dict.fromkeys(_QKV_QUANTIZER_NAMES, 0) + + def make_count_hook(name): + def count_call(_module, _args, _output): + call_counts[name] += 1 + + return count_call + + handles = [ + getattr(model.attn, name).register_forward_hook(make_count_hook(name)) + for name in _QKV_QUANTIZER_NAMES + ] + model(inputs) + for handle in handles: + handle.remove() + assert call_counts == dict.fromkeys(_QKV_QUANTIZER_NAMES, expected_calls) + finally: + if attention_cls in mtq.QuantModuleRegistry: + mtq.unregister(attention_cls) + + +def test_helper_registration_preserves_direct_attention_methods(): + model = DelegatedDiffusionModel(DelegatedAttention) + inputs = torch.randn(2, 4, 16) + + try: + mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) + call_counts = dict.fromkeys(_QKV_QUANTIZER_NAMES, 0) + + def make_count_hook(name): + def count_call(_module, _args, _output): + call_counts[name] += 1 + + return count_call + + handles = [ + getattr(model.attn, name).register_forward_hook(make_count_hook(name)) + for name in _QKV_QUANTIZER_NAMES + ] + q, k, v = (torch.randn(1, 2, 2) for _ in range(3)) + model.attn._unused_attention(q, k, v) + for handle in handles: + handle.remove() + assert call_counts == dict.fromkeys(_QKV_QUANTIZER_NAMES, 1) + finally: + if DelegatedAttention in mtq.QuantModuleRegistry: + mtq.unregister(DelegatedAttention) + + +def test_default_fp8_config_disables_delegated_attention_quantizers(): + model = DelegatedDiffusionModel(DelegatedAttention) + inputs = torch.randn(2, 4, 16) + + try: + mtq.quantize( + model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda quant_model: quant_model(inputs) + ) + for name in _QKV_QUANTIZER_NAMES: + assert not getattr(model.attn, name).is_enabled + finally: + if DelegatedAttention in mtq.QuantModuleRegistry: + mtq.unregister(DelegatedAttention) + + +@pytest.mark.parametrize("attention_cls", [AuxiliaryHelperAttention, MethodMatmulAttention]) +def test_register_attention_rejects_unsupported_patterns(attention_cls): + try: + assert not register_attention_for_kv_quant(attention_cls) + assert attention_cls not in mtq.QuantModuleRegistry + finally: + if attention_cls in mtq.QuantModuleRegistry: + mtq.unregister(attention_cls) + + +def test_discovery_skips_attention_with_custom_setup(): + model = DelegatedDiffusionModel(CustomSetupDelegatedAttention) + inputs = torch.randn(2, 4, 16) + + try: + mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) + assert not hasattr(model.attn, "q_bmm_quantizer") + finally: + if CustomSetupDelegatedAttention in mtq.QuantModuleRegistry: + mtq.unregister(CustomSetupDelegatedAttention) + + +def test_exported_attention_matmuls_have_fp8_qdq(): + onnx = pytest.importorskip("onnx") + from modelopt.onnx.export import FP8QuantExporter + + model = DelegatedDiffusionModel(DelegatedAttention).eval() + inputs = torch.randn(2, 4, 16) + + try: + mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) + assert model.attn.o.input_quantizer.is_enabled + buffer = io.BytesIO() + torch.onnx.export(model, inputs, buffer, opset_version=20, dynamo=False) + exported_model = onnx.load_model_from_string(buffer.getvalue()) + processed_model = FP8QuantExporter.process_model(exported_model) + onnx.checker.check_model(processed_model) + + producer_by_output = { + output: node for node in processed_model.graph.node for output in node.output + } + consumers_by_input = defaultdict(list) + for node in processed_model.graph.node: + for tensor_name in node.input: + consumers_by_input[tensor_name].append(node) + + def is_activation_dq(tensor_name): + dq_node = producer_by_output.get(tensor_name) + return ( + dq_node is not None + and dq_node.op_type == "DequantizeLinear" + and (q_node := producer_by_output.get(dq_node.input[0])) is not None + and q_node.op_type == "QuantizeLinear" + ) + + def is_softmax_qdq_input(tensor_name): + dq_node = producer_by_output[tensor_name] + q_node = producer_by_output[dq_node.input[0]] + source_node = producer_by_output.get(q_node.input[0]) + return source_node is not None and source_node.op_type == "Softmax" + + attention_matmuls = [ + node + for node in processed_model.graph.node + if node.op_type == "MatMul" and all(is_activation_dq(name) for name in node.input) + ] + assert len(attention_matmuls) == 2 + + value_matmuls = [ + node + for node in attention_matmuls + if any(is_softmax_qdq_input(tensor_name) for tensor_name in node.input) + ] + assert len(value_matmuls) == 1 + + def feeds_dequantized_matmul(q_node): + return any( + dq_node.op_type == "DequantizeLinear" + and any( + consumer.op_type == "MatMul" + for consumer in consumers_by_input[dq_node.output[0]] + ) + for dq_node in consumers_by_input[q_node.output[0]] + ) + + pending = list(value_matmuls[0].output) + visited = set(pending) + output_projection_qdq = False + while pending: + tensor_name = pending.pop() + for node in consumers_by_input.get(tensor_name, []): + if node.op_type == "QuantizeLinear" and feeds_dequantized_matmul(node): + output_projection_qdq = True + for output_name in node.output: + if output_name not in visited: + visited.add(output_name) + pending.append(output_name) + assert output_projection_qdq + finally: + if DelegatedAttention in mtq.QuantModuleRegistry: + mtq.unregister(DelegatedAttention) From e31020d99faac377e10908d727e2fec6698d5dcc Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Sat, 12 Sep 2026 00:03:08 +0000 Subject: [PATCH 2/2] Simplify delegated FP8 attention support Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/export/fp8_exporter.py | 167 +------ .../torch/quantization/plugins/attention.py | 171 +++---- .../plugins/diffusion/diffusers.py | 132 +++++- .../quantization/test_fp8_mha_exporter.py | 140 +----- .../plugins/test_diffusers_attention.py | 423 ++++++------------ 5 files changed, 336 insertions(+), 697 deletions(-) diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 0661d101330..8e94f40f332 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -32,8 +32,6 @@ # when using 1/448 as the Q scale (single fixed value — softmax range is data-independent). _FP8_E4M3_MAX = 448.0 _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX -_ELEMENTWISE_SCALAR_OPS = {"Add", "Div", "Mul", "Pow", "Sub"} -_UNARY_SCALAR_OPS = {"Cast", "Reciprocal", "Sqrt"} def _torch_from_numpy_for_fp8(array: np.ndarray) -> torch.Tensor: @@ -43,144 +41,6 @@ def _torch_from_numpy_for_fp8(array: np.ndarray) -> torch.Tensor: return torch.from_numpy(array) -def _constant_values(tensor: gs.Tensor) -> np.ndarray | None: - if isinstance(tensor, gs.Constant): - return tensor.values - if not isinstance(tensor, gs.Variable) or len(tensor.inputs) != 1: - return None - producer = tensor.inputs[0] - if producer.op != "Constant": - return None - value = producer.attrs.get("value") - return value.values if isinstance(value, gs.Constant) else None - - -def _shape_vector_length(tensor: gs.Tensor) -> int | None: - if not isinstance(tensor, gs.Variable) or len(tensor.inputs) != 1: - return None - shape_node = tensor.inputs[0] - if shape_node.op != "Shape" or len(shape_node.inputs) != 1: - return None - source_shape = shape_node.inputs[0].shape - if source_shape is None: - return None - rank = len(source_shape) - start = shape_node.attrs.get("start", 0) - end = shape_node.attrs.get("end", rank) - return len(range(*slice(start, end).indices(rank))) - - -def _is_data_independent_scale(tensor: gs.Tensor, memo=None) -> bool: - """Return whether a tensor is provably scalar and independent of data values.""" - if memo is None: - memo = {} - tensor_id = id(tensor) - if tensor_id in memo: - return memo[tensor_id] - memo[tensor_id] = False - - values = _constant_values(tensor) - if values is not None: - memo[tensor_id] = values.size == 1 - return memo[tensor_id] - if not isinstance(tensor, gs.Variable) or len(tensor.inputs) != 1: - return False - - producer = tensor.inputs[0] - inputs = producer.inputs - result = False - if producer.op == "Gather" and len(inputs) == 2: - indices = _constant_values(inputs[1]) - if ( - _shape_vector_length(inputs[0]) is not None - and producer.attrs.get("axis", 0) in (0, -1) - and indices is not None - and indices.size == 1 - ): - result = True - elif producer.op == "Slice" and 3 <= len(inputs) <= 5: - vector_length = _shape_vector_length(inputs[0]) - starts = _constant_values(inputs[1]) - ends = _constant_values(inputs[2]) - axes = _constant_values(inputs[3]) if len(inputs) >= 4 else np.array([0]) - steps = _constant_values(inputs[4]) if len(inputs) == 5 else np.array([1]) - if ( - vector_length is not None - and starts is not None - and ends is not None - and axes is not None - and steps is not None - and all(value.size == 1 for value in (starts, ends, axes, steps)) - and axes.item() in (0, -1) - and steps.item() != 0 - ): - result = ( - len(range(*slice(starts.item(), ends.item(), steps.item()).indices(vector_length))) - == 1 - ) - elif producer.op in _UNARY_SCALAR_OPS and len(inputs) == 1: - result = _is_data_independent_scale(inputs[0], memo) - elif producer.op in _ELEMENTWISE_SCALAR_OPS and inputs: - result = all(_is_data_independent_scale(input_tensor, memo) for input_tensor in inputs) - - memo[tensor_id] = result - return result - - -def _rebase_shape_sources(tensor, old_tensors, new_tensor): - """Rebase direct ``Shape(Q/DQ)`` inputs, rejecting dependencies that would cycle.""" - dependency_memo = {} - - def is_old_tensor(candidate): - return any(candidate is old_tensor for old_tensor in old_tensors) - - def depends_on_old_tensor(candidate): - if is_old_tensor(candidate): - return True - if not isinstance(candidate, gs.Variable): - return False - candidate_id = id(candidate) - if candidate_id in dependency_memo: - return dependency_memo[candidate_id] - dependency_memo[candidate_id] = False - dependency_memo[candidate_id] = any( - depends_on_old_tensor(producer_input) - for producer in candidate.inputs - for producer_input in producer.inputs - ) - return dependency_memo[candidate_id] - - shape_nodes = [] - visited = set() - - def collect_shape_nodes(candidate): - if not isinstance(candidate, gs.Variable) or id(candidate) in visited: - return True - visited.add(id(candidate)) - for producer in candidate.inputs: - if producer.op == "Shape": - if any( - not is_old_tensor(input_tensor) and depends_on_old_tensor(input_tensor) - for input_tensor in producer.inputs - ): - return False - if any(is_old_tensor(input_tensor) for input_tensor in producer.inputs): - shape_nodes.append(producer) - continue - if not all(collect_shape_nodes(producer_input) for producer_input in producer.inputs): - return False - return True - - if not collect_shape_nodes(tensor): - return False - for shape_node in shape_nodes: - shape_node.inputs = [ - new_tensor if is_old_tensor(input_tensor) else input_tensor - for input_tensor in shape_node.inputs - ] - return True - - class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -393,24 +253,23 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: @staticmethod def _move_mul_before_qdq(graph: gs.Graph) -> int: - """Move attention scaling from after DQ to before Q for TRT MatMul fusion. + """Move attention-scaling Mul(const) from after DQ to before Q for TRT MatMul fusion. - The scale must be scalar and data-independent. Handles both ``DQ → Mul → MatMul`` - and ``DQ → Transpose → Mul → MatMul`` (K path). + Handles both ``DQ → Mul → MatMul`` and ``DQ → Transpose → Mul → MatMul`` (K path). """ count = 0 for mul_node in list(graph.nodes): if mul_node.op != "Mul": continue - scale_input = next( - (tensor for tensor in mul_node.inputs if _is_data_independent_scale(tensor)), + const_input = next( + (i for i in mul_node.inputs if isinstance(i, gs.Constant) and i.values.size == 1), None, ) tensor_input = next( - (tensor for tensor in mul_node.inputs if tensor is not scale_input), None + (i for i in mul_node.inputs if not isinstance(i, gs.Constant)), None ) - if scale_input is None or tensor_input is None: + if const_input is None or tensor_input is None: continue if not (isinstance(tensor_input, gs.Variable) and len(tensor_input.inputs) == 1): continue @@ -430,7 +289,6 @@ def _move_mul_before_qdq(graph: gs.Graph) -> int: continue q_output = dq_node.inputs[0] - dq_output = dq_node.outputs[0] if ( not isinstance(q_output, gs.Variable) or len(q_output.inputs) != 1 @@ -449,13 +307,6 @@ def _move_mul_before_qdq(graph: gs.Graph) -> int: if not mul_consumers or not all(c.op == "MatMul" for c in mul_consumers): continue - if not _rebase_shape_sources(scale_input, (q_output, dq_output), q_input): - continue - if q_input.shape is None: - q_input.shape = dq_output.shape - if q_input.dtype is None: - q_input.dtype = dq_output.dtype - new_mul_output = gs.Variable( q_input.name + "_scaled", dtype=q_input.dtype, shape=q_input.shape ) @@ -463,13 +314,15 @@ def _move_mul_before_qdq(graph: gs.Graph) -> int: gs.Node( op="Mul", name=mul_node.name + "_moved", - inputs=[q_input, scale_input], + inputs=[q_input, const_input], outputs=[new_mul_output], ) ) q_node.inputs[0] = new_mul_output - replacement = transpose_node.outputs[0] if transpose_node is not None else dq_output + replacement = ( + transpose_node.outputs[0] if transpose_node is not None else dq_node.outputs[0] + ) for consumer in mul_consumers: for i, inp in enumerate(consumer.inputs): if inp is mul_output: diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 235e36dc367..2113edea8a7 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -34,7 +34,6 @@ import inspect import tempfile import types -from typing import cast from warnings import warn from ..conversion import register @@ -51,8 +50,6 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: """ source_code = inspect.getsource(attention_cls) model_module = inspect.getmodule(attention_cls) - if model_module is None: - return False head = ast.parse(source_code) bmm_ops = ("matmul", "bmm", "baddbmm") @@ -75,48 +72,6 @@ def is_sdpa(node): def is_bin_matmul(node): return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) - def get_attention_helper_parameters(node): - if not ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and callable(helper := model_module.__dict__.get(node.func.id)) - ): - return None - - helper_name = getattr(helper, "__name__", node.func.id).lower() - if "attention" not in helper_name and "attn" not in helper_name: - return None - - try: - positional_parameters = [ - parameter.name - for parameter in inspect.signature(helper).parameters.values() - if parameter.kind - in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) - ] - except (TypeError, ValueError): - return None - - keyword_names = {keyword.arg for keyword in node.keywords} - for parameter_names in (("q", "k", "v"), ("query", "key", "value")): - if tuple(positional_parameters[:3]) == parameter_names and all( - name in keyword_names or index < len(node.args) - for index, name in enumerate(parameter_names) - ): - return parameter_names - return None - - def quantize_argument(argument, quantizer_name): - return ast.Call( - func=ast.Attribute( - value=ast.Name(id="self", ctx=ast.Load()), - attr=quantizer_name, - ctx=ast.Load(), - ), - args=[argument], - keywords=[], - ) - def patch(node, quantizer_names, transpose=False): for index, quantizer_name in enumerate(quantizer_names): if quantizer_name is None: @@ -124,16 +79,33 @@ def patch(node, quantizer_names, transpose=False): arg = node.args[index] if not transpose: - node.args[index] = quantize_argument(arg, quantizer_name) - else: - transposed_arg = ast.Call( - func=ast.Attribute(value=arg, attr="transpose", ctx=ast.Load()), - args=[ast.Constant(value=-1), ast.Constant(value=-2)], + node.args[index] = ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_name, + ctx=ast.Load(), + ), + args=[arg], keywords=[], ) + else: node.args[index] = ast.Call( func=ast.Attribute( - value=quantize_argument(transposed_arg, quantizer_name), + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_name, + ctx=ast.Load(), + ), + args=[ + ast.Call( + func=ast.Attribute(value=arg, attr="transpose", ctx=ast.Load()), + args=[ast.Constant(value=-1), ast.Constant(value=-2)], + keywords=[], + ) + ], + keywords=[], + ), attr="transpose", ctx=ast.Load(), ), @@ -141,22 +113,18 @@ def patch(node, quantizer_names, transpose=False): keywords=[], ) - def patch_attention_helper(node, parameter_names): - keyword_by_name = { - keyword.arg: keyword for keyword in node.keywords if keyword.arg is not None - } - for index, (parameter_name, quantizer_name) in enumerate( - zip(parameter_names, ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer")) - ): - if keyword := keyword_by_name.get(parameter_name): - keyword.value = quantize_argument(keyword.value, quantizer_name) - else: - node.args[index] = quantize_argument(node.args[index], quantizer_name) - def patch_binop(node, quantizer_names, transpose=False): assert len(quantizer_names) == 2 if quantizer_names[0] is not None: - node.left = quantize_argument(node.left, quantizer_names[0]) + node.left = ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_names[0], + ctx=ast.Load(), + ), + args=[node.left], + keywords=[], + ) if quantizer_names[1] is not None: arg = node.right if transpose: @@ -169,7 +137,15 @@ def patch_binop(node, quantizer_names, transpose=False): args=[arg, ast.Constant(value=-1), ast.Constant(value=-2)], keywords=[], ) - quant_arg = quantize_argument(arg, quantizer_names[1]) + quant_arg = ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_names[1], + ctx=ast.Load(), + ), + args=[arg], + keywords=[], + ) if transpose: quant_arg = ast.Call( func=ast.Attribute( @@ -182,43 +158,28 @@ def patch_binop(node, quantizer_names, transpose=False): ) node.right = quant_arg - class_node = next(node for node in head.body if isinstance(node, ast.ClassDef)) - org_class_name = class_node.name - new_class_name = class_node.name = "_Quant" + class_node.name - forward_node = next( - ( - node - for node in class_node.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "forward" - ), - None, - ) - attention_helper_nodes = [ - (node, parameter_names) - for node in (ast.walk(forward_node) if forward_node is not None else ()) - if (parameter_names := get_attention_helper_parameters(node)) is not None - ] - - direct_nodes = list(ast.walk(head)) - bmm_nodes = [cast("ast.Call", node) for node in direct_nodes if is_bmm(node)] - sdpa_nodes = [cast("ast.Call", node) for node in direct_nodes if is_sdpa(node)] - bin_matmul_nodes = [cast("ast.BinOp", node) for node in direct_nodes if is_bin_matmul(node)] - patch_bmm = len(bmm_nodes) == 2 and all(len(node.args) >= 2 for node in bmm_nodes) - patch_sdpa = len(sdpa_nodes) == 1 and len(sdpa_nodes[0].args) >= 3 - patch_bin_matmul = len(bin_matmul_nodes) == 2 - if not attention_helper_nodes and not (patch_bmm or patch_sdpa or patch_bin_matmul): + nodes = list(ast.walk(head)) + org_class_name = nodes[1].name # type: ignore[attr-defined] + new_class_name = nodes[1].name = "_Quant" + nodes[1].name # type: ignore[attr-defined] + + bmm_nodes = [] + sdpa_nodes = [] + bin_matmul_nodes = [] + for node in ast.walk(head): + if is_bmm(node): + bmm_nodes.append(node) + if is_sdpa(node): + sdpa_nodes.append(node) + if is_bin_matmul(node): + bin_matmul_nodes.append(node) + if len(bmm_nodes) != 2 and len(sdpa_nodes) != 1 and len(bin_matmul_nodes) != 2: print(f"Expect 2 bmm/matmul op in the {org_class_name}, found {len(bmm_nodes)}") print(f"Or expect 1 sdpa op in the {org_class_name}, found {len(sdpa_nodes)}") print(f"Or expect 2 @ op in the {org_class_name}, found {len(bin_matmul_nodes)}") - print(f"Or expect an attention helper call in the {org_class_name}") print("Auto quantization of KV Cache fails") return False - if attention_helper_nodes: - for node, parameter_names in attention_helper_nodes: - patch_attention_helper(node, parameter_names) - print(f"Patching {len(attention_helper_nodes)} attention helper call(s) with quantizers") - if patch_bmm: + if len(bmm_nodes) == 2: # transpose k cache here to enable per-token quantization # without transpose, the quantization will be per-channel, i.e., # self.k_bmm_quantizer(key_states.transpose(-1, -2)) @@ -226,13 +187,9 @@ def patch_binop(node, quantizer_names, transpose=False): # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) # removing the additional transpose is doable but not trivial patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) - patch( - bmm_nodes[1], - quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), - transpose=True, - ) + patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) print("Patching 2 BMM/Matmul operators with quantizers") - if patch_bin_matmul: + if len(bin_matmul_nodes) == 2: patch_binop( bin_matmul_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), @@ -240,14 +197,10 @@ def patch_binop(node, quantizer_names, transpose=False): ) patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 @ operators with quantizers") - if patch_sdpa: + + if len(sdpa_nodes) == 1: patch( - sdpa_nodes[0], - quantizer_names=( - "q_bmm_quantizer", - "k_bmm_quantizer", - "v_bmm_quantizer", - ), + sdpa_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") ) print("Patching 1 scaled_dot_product_attention operator with quantizers") diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index 5b31e355412..be84c77bba4 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -15,8 +15,8 @@ """Support quantization of diffusers layers.""" +import inspect from collections.abc import Callable, Iterator -from contextlib import suppress from functools import partial from types import ModuleType from typing import TYPE_CHECKING @@ -66,7 +66,6 @@ TensorQuantizer, ) from ...nn.modules.quant_conv import _QuantConv3d -from ..attention import register_attention_for_kv_quant from ..custom import CUSTOM_MODEL_PLUGINS, _QuantFunctionalMixin onnx_dtype_map = { @@ -78,6 +77,7 @@ "UINT8": onnx.TensorProto.UINT8, } mha_valid_precisions = {"Half", "BFloat16"} +_QKV_PARAMETER_NAMES = (("q", "k", "v"), ("query", "key", "value")) class _QuantLoRACompatibleLinearConvBase(QuantLinearConvBase): @@ -111,6 +111,26 @@ def _quantized_baddbmm(self, input, batch1, batch2, *args, **kwargs): return torch._baddbmm(input, self.q_bmm_quantizer(q), self.k_bmm_quantizer(k), *args, **kwargs) +def _fp8_mha_disabled(self, query, key, value): + if hasattr(self, "_disable_fp8_mha") or not getattr(self, "_auto_fp8_mha", False): + return getattr(self, "_disable_fp8_mha", True) + names = ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer", "softmax_quantizer") + quantizers = tuple(getattr(self, name) for name in names) + if any(not quantizer.is_enabled or not quantizer.is_fp8 for quantizer in quantizers): + return True + high_precision = self.q_bmm_quantizer.trt_high_precision_dtype + if high_precision not in mha_valid_precisions or { + quantizer.trt_high_precision_dtype for quantizer in quantizers + } != {high_precision}: + return True + if any( + tensor.dtype != getattr(torch, high_precision.lower()) for tensor in (query, key, value) + ): + return True + head_dims = tuple(int(tensor.shape[-1]) for tensor in (query, key, value)) + return len(set(head_dims)) != 1 or head_dims[0] % 16 != 0 + + def _quantized_sdpa(self, *args, **kwargs): fp8_sdpa = FP8SDPA.apply parameters = [ @@ -135,11 +155,14 @@ def _quantized_sdpa(self, *args, **kwargs): while fp8_sdpa_args and fp8_sdpa_args[-1] is None: fp8_sdpa_args.pop() query, key, value = fp8_sdpa_args[:3] - - if not torch.onnx.is_in_onnx_export(): + exporting = torch.onnx.is_in_onnx_export() + if not exporting: query = self.q_bmm_quantizer(query) key = self.k_bmm_quantizer(key) value = self.v_bmm_quantizer(value) + disable_fp8_mha = _fp8_mha_disabled(self, query, key, value) + if exporting and getattr(self, "_auto_fp8_mha", False) and not disable_fp8_mha: + self._delegated_sdpa_export_calls = getattr(self, "_delegated_sdpa_export_calls", 0) + 1 q_quantized_scale = self.q_bmm_quantizer._get_amax(query) k_quantized_scale = self.k_bmm_quantizer._get_amax(key) @@ -158,7 +181,7 @@ def _quantized_sdpa(self, *args, **kwargs): self.q_bmm_quantizer.trt_high_precision_dtype if hasattr(self.q_bmm_quantizer, "trt_high_precision_dtype") else "Half", - self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, + disable_fp8_mha, ) ) @@ -211,15 +234,103 @@ def forward(self, *args, **kwargs): ) +original_scaled_dot_product_attention = F.scaled_dot_product_attention + + +def _call_delegated_attention(self, helper, parameter_names, *args, **kwargs): + bound = inspect.signature(helper).bind(*args, **kwargs) + qkv = tuple(bound.arguments[name] for name in parameter_names) + if torch.onnx.is_in_onnx_export(): + before = self._delegated_sdpa_export_calls + eligible = not _fp8_mha_disabled(self, *qkv) + output = helper(*bound.args, **bound.kwargs) + self._delegated_missed_sdpa |= eligible and before == self._delegated_sdpa_export_calls + return output + + quantizers = (self.q_bmm_quantizer, self.k_bmm_quantizer, self.v_bmm_quantizer) + for name, quantizer in zip(parameter_names, quantizers): + bound.arguments[name] = quantizer(bound.arguments[name]) + return self.bmm2_output_quantizer(helper(*bound.args, **bound.kwargs)) + + +class _QuantDelegatedAttention(_QuantAttention): + def forward(self, *args, **kwargs): + self._delegated_missed_sdpa = self._delegated_sdpa_export_calls = 0 + output = super().forward(*args, **kwargs) + if torch.onnx.is_in_onnx_export() and self._delegated_missed_sdpa: + raise RuntimeError("Eligible delegated FP8 MHA export did not reach SDPA.") + return output + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + module, name, parameter_names = self._delegated_attention_helper + helper = getattr(module, name, None) + if not callable(helper): + return + quantized_sdpa = partial(_quantized_sdpa, self) + if helper is original_scaled_dot_product_attention: + replacement = quantized_sdpa + else: + replacement = partial(_call_delegated_attention, self, helper, parameter_names) + yield module, name, replacement + if not torch.onnx.is_in_onnx_export(): + return + + yield F, "scaled_dot_product_attention", quantized_sdpa + if helper is original_scaled_dot_product_attention: + return + helper = inspect.unwrap(helper) + helper_globals = getattr(helper, "__globals__", {}) + helper_module = inspect.getmodule(helper) + code = getattr(helper, "__code__", None) + if code is None or not isinstance(helper_module, ModuleType): + return + if helper_module.__dict__ is helper_globals: + for alias in code.co_names: + if helper_globals.get(alias) is original_scaled_dot_product_attention: + yield helper_module, alias, quantized_sdpa + + def _try_register_attention(attention_cls): - with suppress( - IndentationError, IndexError, KeyError, OSError, SyntaxError, TypeError, ValueError + try: + forward = inspect.unwrap(attention_cls.forward) + except (TypeError, ValueError): + return + code = getattr(forward, "__code__", None) + function_globals = getattr(forward, "__globals__", {}) + if ( + code is None + or not isinstance(module := inspect.getmodule(forward), ModuleType) + or module.__dict__ is not function_globals ): - register_attention_for_kv_quant(attention_cls) + return + candidates = [] + for name in dict.fromkeys(code.co_names): + helper = function_globals.get(name) + if not callable(helper): + continue + try: + parameter_names = tuple(inspect.signature(helper).parameters)[:3] + except (TypeError, ValueError): + if helper is not original_scaled_dot_product_attention: + continue + parameter_names = _QKV_PARAMETER_NAMES[1] + helper_name = f"{name} {getattr(helper, '__name__', '')}".lower() + if parameter_names in _QKV_PARAMETER_NAMES and ( + "attention" in helper_name or "attn" in helper_name + ): + candidates.append((module, name, parameter_names)) + if len(candidates) != 1: + return + quantized_cls = type( + f"_Quant{attention_cls.__name__}", + (_QuantDelegatedAttention,), + {"_auto_fp8_mha": True, "_delegated_attention_helper": candidates[0]}, + ) + QuantModuleRegistry.register({attention_cls: attention_cls.__name__})(quantized_cls) def _register_diffusers_attentions_on_the_fly(model): - """Register unrecognized leaf attention modules in diffusers model graphs.""" if not isinstance(model, ModelMixin): return @@ -246,9 +357,6 @@ def _register_diffusers_attentions_on_the_fly(model): CUSTOM_MODEL_PLUGINS.add(_register_diffusers_attentions_on_the_fly) -original_scaled_dot_product_attention = F.scaled_dot_product_attention - - class FP8SDPA(Function): """A customized FP8 SDPA op for the onnx export.""" diff --git a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py index 7edb947fa1f..1f7251a9ad9 100644 --- a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py +++ b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py @@ -22,15 +22,14 @@ from modelopt.onnx.export.fp8_exporter import FP8QuantExporter -def _var(name, shape=None): - return gs.Variable(name, dtype=np.float32, shape=shape) +def _var(name): + return gs.Variable(name, dtype=np.float32) def _qdq(src): """Build ``QuantizeLinear → DequantizeLinear`` and return [Q, DQ], dq_out.""" scale = gs.Constant("scale", np.array(0.1, dtype=np.float32)) - q_out = _var("q_out", src.shape) - dq_out = _var("dq_out", src.shape) + q_out, dq_out = _var("q_out"), _var("dq_out") return [ gs.Node(op="QuantizeLinear", inputs=[src, scale], outputs=[q_out]), gs.Node(op="DequantizeLinear", inputs=[q_out, scale], outputs=[dq_out]), @@ -58,139 +57,6 @@ def test_move_mul_before_qdq_rewrites_dq_mul_matmul_pattern(): assert q.inputs[0].inputs[0].op == "Mul" -def test_move_mul_before_qdq_rewrites_constant_expression_scale(): - """The scalar scale may be a constant expression emitted by the Torch exporter.""" - x, k, y, mul_out = _var("x"), _var("k"), _var("y"), _var("mul_out") - constant_out, sqrt_out = _var("constant_out"), _var("sqrt_out") - qdq_nodes, dq_out = _qdq(x) - constant = gs.Node( - op="Constant", - attrs={"value": gs.Constant("constant", np.array(0.25, dtype=np.float32))}, - outputs=[constant_out], - ) - sqrt = gs.Node(op="Sqrt", inputs=[constant_out], outputs=[sqrt_out]) - mul = gs.Node(op="Mul", inputs=[dq_out, sqrt_out], outputs=[mul_out]) - mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) - graph = _graph([*qdq_nodes, constant, sqrt, mul, mm], [x, k], [y]) - - assert FP8QuantExporter._move_mul_before_qdq(graph) == 1 - q = next(n for n in graph.nodes if n.op == "QuantizeLinear") - assert q.inputs[0].inputs[0].op == "Mul" - - -@pytest.mark.parametrize("shape_source", ["q", "dq"]) -@pytest.mark.parametrize("indirect", [False, True]) -def test_move_mul_before_qdq_handles_shape_dependencies(shape_source, indirect): - x, k, y, identity_out, shape_out, slice_out, scale, mul_out = ( - _var("x", [1, 2, 3, 4]), - _var("k"), - _var("y"), - _var("identity_out", [1, 2, 3, 4]), - _var("shape_out"), - _var("slice_out"), - _var("shape_scale"), - _var("mul_out"), - ) - qdq_nodes, dq_out = _qdq(x) - q_out = qdq_nodes[0].outputs[0] - shape_input = q_out if shape_source == "q" else dq_out - pre_shape_nodes = [] - if indirect: - pre_shape_nodes.append(gs.Node(op="Identity", inputs=[shape_input], outputs=[identity_out])) - shape_input = identity_out - shape = gs.Node(op="Shape", inputs=[shape_input], outputs=[shape_out]) - starts, ends = _var("starts"), _var("ends") - slice_constants = [ - gs.Node( - op="Constant", - attrs={"value": gs.Constant("starts_value", np.array([-1], dtype=np.int64))}, - outputs=[starts], - ), - gs.Node( - op="Constant", - attrs={ - "value": gs.Constant( - "ends_value", np.array([np.iinfo(np.int64).max], dtype=np.int64) - ) - }, - outputs=[ends], - ), - ] - shape_slice = gs.Node(op="Slice", inputs=[shape_out, starts, ends], outputs=[slice_out]) - sqrt = gs.Node(op="Sqrt", inputs=[slice_out], outputs=[scale]) - mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) - mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) - graph = _graph( - [*qdq_nodes, *pre_shape_nodes, shape, *slice_constants, shape_slice, sqrt, mul, mm], - [x, k], - [y], - ) - - assert FP8QuantExporter._move_mul_before_qdq(graph) == (0 if indirect else 1) - assert shape.inputs == [identity_out if indirect else x] - - -def test_move_mul_before_qdq_skips_vector_shape_scale(): - x, k, y, shape_out, scale, mul_out = ( - _var("x", [1, 2, 3, 4]), - _var("k"), - _var("y"), - _var("shape_out"), - _var("shape_scale"), - _var("mul_out"), - ) - qdq_nodes, dq_out = _qdq(x) - shape = gs.Node(op="Shape", inputs=[dq_out], outputs=[shape_out]) - cast = gs.Node(op="Cast", inputs=[shape_out], outputs=[scale], attrs={"to": 1}) - mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) - mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) - graph = _graph([*qdq_nodes, shape, cast, mul, mm], [x, k], [y]) - - assert FP8QuantExporter._move_mul_before_qdq(graph) == 0 - assert shape.inputs == [dq_out] - - -def test_move_mul_before_qdq_rewrites_gathered_shape_scale(): - x, k, y, shape_out, scale, mul_out = ( - _var("x", [1, 2, 3, 4]), - _var("k"), - _var("y"), - _var("shape_out"), - _var("shape_scale"), - _var("mul_out"), - ) - qdq_nodes, dq_out = _qdq(x) - shape = gs.Node(op="Shape", inputs=[dq_out], outputs=[shape_out]) - gather = gs.Node( - op="Gather", - inputs=[shape_out, gs.Constant("index", np.array(-1, dtype=np.int64))], - outputs=[scale], - attrs={"axis": 0}, - ) - mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) - mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) - graph = _graph([*qdq_nodes, shape, gather, mul, mm], [x, k], [y]) - - assert FP8QuantExporter._move_mul_before_qdq(graph) == 1 - assert shape.inputs == [x] - - -def test_move_mul_before_qdq_skips_data_dependent_scale(): - x, scale, k, y, mul_out = ( - _var("x"), - _var("dynamic_scale"), - _var("k"), - _var("y"), - _var("mul_out"), - ) - qdq_nodes, dq_out = _qdq(x) - mul = gs.Node(op="Mul", inputs=[dq_out, scale], outputs=[mul_out]) - mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) - graph = _graph([*qdq_nodes, mul, mm], [x, scale, k], [y]) - - assert FP8QuantExporter._move_mul_before_qdq(graph) == 0 - - def test_move_transpose_before_qdq_rewrites_dq_transpose_matmul_pattern(): """``DQ → Transpose → MatMul`` collapses to ``Transpose → Q → DQ → MatMul``.""" k_in, q_in, scores, t_out = _var("k_in"), _var("q_in"), _var("scores"), _var("t_out") diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_attention.py b/tests/unit/torch/quantization/plugins/test_diffusers_attention.py index b6cbffe2de4..4c7e0be308f 100644 --- a/tests/unit/torch/quantization/plugins/test_diffusers_attention.py +++ b/tests/unit/torch/quantization/plugins/test_diffusers_attention.py @@ -13,320 +13,179 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for on-the-fly registration of delegated diffusion attention.""" - import copy import io -from collections import defaultdict import pytest import torch import torch.nn as nn import torch.nn.functional as F -pytest.importorskip("diffusers") -from diffusers import ModelMixin +ModelMixin = pytest.importorskip("diffusers").ModelMixin import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.plugins.attention import register_attention_for_kv_quant +from modelopt.torch.quantization.plugins.diffusion.diffusers import _fp8_mha_disabled -_QKV_QUANTIZER_NAMES = ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") +_SDPA_ALIAS = F.scaled_dot_product_attention -def delegated_attention(q, k, v): - return F.scaled_dot_product_attention(q, k, v) +def _fake_attention(q, k, v): + return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v -class DelegatedAttention(nn.Module): - def __init__(self, hidden_size=16, num_heads=2): - super().__init__() - self.num_heads = num_heads - self.head_dim = hidden_size // num_heads - self.q = nn.Linear(hidden_size, hidden_size) - self.k = nn.Linear(hidden_size, hidden_size) - self.v = nn.Linear(hidden_size, hidden_size) - self.o = nn.Linear(hidden_size, hidden_size) - - def _reshape(self, hidden_states): - batch_size, sequence_length, _ = hidden_states.shape - return hidden_states.view( - batch_size, sequence_length, self.num_heads, self.head_dim - ).transpose(1, 2) - - def _unused_attention(self, q, k, v): - scores = torch.matmul(q, k) - return torch.matmul(scores, v) - - def forward(self, hidden_states): - q = self._reshape(self.q(hidden_states)) - k = self._reshape(self.k(hidden_states)) - v = self._reshape(self.v(hidden_states)) - output = delegated_attention(q=q * 0.5, k=k * 0.25, v=v) - return self.o(output.transpose(1, 2).flatten(2)) +def _wan_attention(q, k, v): + if torch.onnx.is_in_onnx_export(): + return F.scaled_dot_product_attention(q, k, v) + return _fake_attention(q, k, v) -class PositionalDelegatedAttention(DelegatedAttention): - def forward(self, hidden_states): - q = self._reshape(self.q(hidden_states)) - k = self._reshape(self.k(hidden_states)) - v = self._reshape(self.v(hidden_states)) - output = delegated_attention(q * 0.5, k * 0.25, v) - return self.o(output.transpose(1, 2).flatten(2)) - - -class RepeatedDelegatedAttention(DelegatedAttention): - def forward(self, hidden_states): - q = self._reshape(self.q(hidden_states)) - k = self._reshape(self.k(hidden_states)) - v = self._reshape(self.v(hidden_states)) - output = delegated_attention(q, k, v) - output = delegated_attention(output, k, v) - return self.o(output.transpose(1, 2).flatten(2)) +class DelegatedAttention(nn.Module): + def __init__(self, style="keyword"): + super().__init__() + self.style = style + self.q, self.k, self.v, self.o = (nn.Linear(32, 32) for _ in range(4)) + def _qkv(self, x): + qkv = (self.q(x), self.k(x), self.v(x)) + return tuple(tensor.view(2, 4, 2, 16).transpose(1, 2) for tensor in qkv) -class HelperWithMethodMatmulAttention(DelegatedAttention): - def forward(self, hidden_states): - q = self._reshape(self.q(hidden_states)) - k = self._reshape(self.k(hidden_states)) - v = self._reshape(self.v(hidden_states)) - if hidden_states.shape[0] == 0: - scores = q.matmul(k.transpose(-2, -1)) - output = scores.matmul(v) + def forward(self, x): + q, k, v = self._qkv(x) + if self.style == "keyword": + output = _wan_attention(q=q, k=k, v=v) else: - output = delegated_attention(q, k, v) + output = _wan_attention(q, k, v) + if self.style == "repeated": + output = _wan_attention(output, k, v) return self.o(output.transpose(1, 2).flatten(2)) -class AuxiliaryHelperAttention(nn.Module): - def forward(self, hidden_states): - return hidden_states +class AliasedAttention(DelegatedAttention): + def forward(self, x): + return self.o(_SDPA_ALIAS(*self._qkv(x)).transpose(1, 2).flatten(2)) - def _unused_attention(self, q, k, v): - return delegated_attention(q, k, v) +class NonSDPAAttention(DelegatedAttention): + def forward(self, x): + return self.o(_fake_attention(*self._qkv(x)).transpose(1, 2).flatten(2)) -class MethodMatmulAttention(nn.Module): - def forward(self, q, k, v): - scores = q.matmul(k.transpose(-2, -1)) - return scores.matmul(v) +class DelegatedModel(ModelMixin): + def __init__(self, attention_cls=DelegatedAttention, **kwargs): + super().__init__() + self.attn = attention_cls(**kwargs) -class CustomSetupDelegatedAttention(DelegatedAttention): - def _setup(self): - pass + def forward(self, x): + return self.attn(x) - def forward(self, hidden_states): - q = self._reshape(self.q(hidden_states)) - k = self._reshape(self.k(hidden_states)) - v = self._reshape(self.v(hidden_states)) - output = delegated_attention(q, k, v) - return self.o(output.transpose(1, 2).flatten(2)) +@pytest.fixture(autouse=True) +def _clean_registrations(): + yield + for cls in (DelegatedAttention, AliasedAttention, NonSDPAAttention): + if cls in mtq.QuantModuleRegistry: + mtq.unregister(cls) -class DelegatedDiffusionModel(ModelMixin): - def __init__(self, attention_cls): - super().__init__() - self.attn = attention_cls() - def forward(self, hidden_states): - return self.attn(hidden_states) +def _quantize(model, enabled=True): + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + attrs = {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"} + config["quant_cfg"].extend( + {"quantizer_name": name, "cfg": attrs, "enable": enabled} + for name in ("*[qkv]_bmm_quantizer", "*softmax_quantizer") + ) + inputs = torch.randn(2, 4, 32) + mtq.quantize(model, config, lambda quant_model: quant_model(inputs)) + return inputs -def _get_fp8_attention_config(): - quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - quant_config["quant_cfg"].append( - { - "quantizer_name": "*[qkv]_bmm_quantizer", - "cfg": {"num_bits": (4, 3), "axis": None}, - "enable": True, - } - ) - return quant_config - - -@pytest.mark.parametrize( - ("attention_cls", "expected_calls"), - [ - (DelegatedAttention, 1), - (PositionalDelegatedAttention, 1), - (RepeatedDelegatedAttention, 2), - (HelperWithMethodMatmulAttention, 1), - ], -) -def test_quantize_registers_delegated_attention(attention_cls, expected_calls): - model = DelegatedDiffusionModel(attention_cls) - inputs = torch.randn(2, 4, 16) - - try: - mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) - - for name in _QKV_QUANTIZER_NAMES: - quantizer = getattr(model.attn, name) - assert quantizer.is_enabled - assert quantizer.amax is not None - - call_counts = dict.fromkeys(_QKV_QUANTIZER_NAMES, 0) - - def make_count_hook(name): - def count_call(_module, _args, _output): - call_counts[name] += 1 - - return count_call - - handles = [ - getattr(model.attn, name).register_forward_hook(make_count_hook(name)) - for name in _QKV_QUANTIZER_NAMES - ] - model(inputs) - for handle in handles: - handle.remove() - assert call_counts == dict.fromkeys(_QKV_QUANTIZER_NAMES, expected_calls) - finally: - if attention_cls in mtq.QuantModuleRegistry: - mtq.unregister(attention_cls) - - -def test_helper_registration_preserves_direct_attention_methods(): - model = DelegatedDiffusionModel(DelegatedAttention) - inputs = torch.randn(2, 4, 16) - - try: - mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) - call_counts = dict.fromkeys(_QKV_QUANTIZER_NAMES, 0) - - def make_count_hook(name): - def count_call(_module, _args, _output): - call_counts[name] += 1 - - return count_call - - handles = [ - getattr(model.attn, name).register_forward_hook(make_count_hook(name)) - for name in _QKV_QUANTIZER_NAMES - ] - q, k, v = (torch.randn(1, 2, 2) for _ in range(3)) - model.attn._unused_attention(q, k, v) - for handle in handles: - handle.remove() - assert call_counts == dict.fromkeys(_QKV_QUANTIZER_NAMES, 1) - finally: - if DelegatedAttention in mtq.QuantModuleRegistry: - mtq.unregister(DelegatedAttention) - - -def test_default_fp8_config_disables_delegated_attention_quantizers(): - model = DelegatedDiffusionModel(DelegatedAttention) - inputs = torch.randn(2, 4, 16) - - try: - mtq.quantize( - model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda quant_model: quant_model(inputs) - ) - for name in _QKV_QUANTIZER_NAMES: - assert not getattr(model.attn, name).is_enabled - finally: - if DelegatedAttention in mtq.QuantModuleRegistry: - mtq.unregister(DelegatedAttention) - - -@pytest.mark.parametrize("attention_cls", [AuxiliaryHelperAttention, MethodMatmulAttention]) -def test_register_attention_rejects_unsupported_patterns(attention_cls): - try: - assert not register_attention_for_kv_quant(attention_cls) - assert attention_cls not in mtq.QuantModuleRegistry - finally: - if attention_cls in mtq.QuantModuleRegistry: - mtq.unregister(attention_cls) - - -def test_discovery_skips_attention_with_custom_setup(): - model = DelegatedDiffusionModel(CustomSetupDelegatedAttention) - inputs = torch.randn(2, 4, 16) - - try: - mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) - assert not hasattr(model.attn, "q_bmm_quantizer") - finally: - if CustomSetupDelegatedAttention in mtq.QuantModuleRegistry: - mtq.unregister(CustomSetupDelegatedAttention) - - -def test_exported_attention_matmuls_have_fp8_qdq(): +def _export(attention_cls): onnx = pytest.importorskip("onnx") - from modelopt.onnx.export import FP8QuantExporter - - model = DelegatedDiffusionModel(DelegatedAttention).eval() - inputs = torch.randn(2, 4, 16) - - try: - mtq.quantize(model, _get_fp8_attention_config(), lambda quant_model: quant_model(inputs)) - assert model.attn.o.input_quantizer.is_enabled - buffer = io.BytesIO() - torch.onnx.export(model, inputs, buffer, opset_version=20, dynamo=False) - exported_model = onnx.load_model_from_string(buffer.getvalue()) - processed_model = FP8QuantExporter.process_model(exported_model) - onnx.checker.check_model(processed_model) - - producer_by_output = { - output: node for node in processed_model.graph.node for output in node.output - } - consumers_by_input = defaultdict(list) - for node in processed_model.graph.node: - for tensor_name in node.input: - consumers_by_input[tensor_name].append(node) - - def is_activation_dq(tensor_name): - dq_node = producer_by_output.get(tensor_name) - return ( - dq_node is not None - and dq_node.op_type == "DequantizeLinear" - and (q_node := producer_by_output.get(dq_node.input[0])) is not None - and q_node.op_type == "QuantizeLinear" - ) - - def is_softmax_qdq_input(tensor_name): - dq_node = producer_by_output[tensor_name] - q_node = producer_by_output[dq_node.input[0]] - source_node = producer_by_output.get(q_node.input[0]) - return source_node is not None and source_node.op_type == "Softmax" - - attention_matmuls = [ - node - for node in processed_model.graph.node - if node.op_type == "MatMul" and all(is_activation_dq(name) for name in node.input) - ] - assert len(attention_matmuls) == 2 - - value_matmuls = [ - node - for node in attention_matmuls - if any(is_softmax_qdq_input(tensor_name) for tensor_name in node.input) - ] - assert len(value_matmuls) == 1 - - def feeds_dequantized_matmul(q_node): - return any( - dq_node.op_type == "DequantizeLinear" - and any( - consumer.op_type == "MatMul" - for consumer in consumers_by_input[dq_node.output[0]] - ) - for dq_node in consumers_by_input[q_node.output[0]] - ) - - pending = list(value_matmuls[0].output) - visited = set(pending) - output_projection_qdq = False - while pending: - tensor_name = pending.pop() - for node in consumers_by_input.get(tensor_name, []): - if node.op_type == "QuantizeLinear" and feeds_dequantized_matmul(node): - output_projection_qdq = True - for output_name in node.output: - if output_name not in visited: - visited.add(output_name) - pending.append(output_name) - assert output_projection_qdq - finally: - if DelegatedAttention in mtq.QuantModuleRegistry: - mtq.unregister(DelegatedAttention) + model = DelegatedModel(attention_cls).eval() + inputs, buffer = _quantize(model), io.BytesIO() + torch.onnx.export(model.half(), inputs.half(), buffer, opset_version=20, dynamo=False) + graph = onnx.load_model_from_string(buffer.getvalue()) + onnx.checker.check_model(graph) + return graph + + +def _attention_qdq(graph): + producers = {output: node for node in graph.graph.node for output in node.output} + + def source(name): + node = producers.get(name) + if node is not None and node.op_type == "Cast": + node = producers.get(node.input[0]) + if node is None or node.op_type != "TRT_FP8DequantizeLinear": + return None + quantize = producers.get(node.input[0]) + if quantize is None or quantize.op_type != "TRT_FP8QuantizeLinear": + return None + return quantize, producers.get(quantize.input[0]) + + return producers, [ + (node, inputs) + for node in graph.graph.node + if node.op_type == "MatMul" and all(inputs := [source(name) for name in node.input]) + ] + + +@pytest.mark.parametrize(("style", "calls"), [("keyword", 1), ("positional", 1), ("repeated", 2)]) +def test_delegated_helper_call_styles(style, calls): + model = DelegatedModel(style=style) + inputs = _quantize(model) + seen = [] + with model.attn.q_bmm_quantizer.register_forward_hook(lambda *_: seen.append(None)): + model(inputs) + assert len(seen) == calls + assert all(getattr(model.attn, f"{name}_bmm_quantizer").amax is not None for name in "qkv") + + +@pytest.mark.parametrize("attention_cls", [DelegatedAttention, AliasedAttention]) +def test_export_uses_fp8_sdpa_symbolic(attention_cls): + graph = _export(attention_cls) + producers, attention = _attention_qdq(graph) + qk = next(x for _, x in attention if all(source.op_type == "Mul" for _, source in x)) + assert any(producers[source.input[0]].op_type == "Transpose" for _, source in qk) + pv, inputs = next( + x for x in attention if any(source.op_type == "Softmax" for _, source in x[1]) + ) + softmax_q = next(q for q, source in inputs if source.op_type == "Softmax") + scale = producers[softmax_q.input[1]].attribute[0].t.raw_data + assert scale == torch.tensor(1 / 448, dtype=torch.float16).numpy().tobytes() + reachable = set(pv.output) + for node in graph.graph.node: + if reachable.intersection(node.input): + reachable.update(node.output) + qdq = {node.input[0] for node in graph.graph.node if node.op_type == "TRT_FP8QuantizeLinear"} + assert reachable & qdq + + +def test_export_fails_when_helper_does_not_reach_sdpa(): + helper = _fake_attention + with pytest.raises(RuntimeError, match="did not reach SDPA"): + _export(NonSDPAAttention) + assert _fake_attention is helper and F.scaled_dot_product_attention is _SDPA_ALIAS + + +@pytest.mark.parametrize("case", ["eligible", "fp32", "int8", "misaligned", "explicit"]) +def test_fp8_mha_eligibility(case): + model = DelegatedModel() + _quantize(model) + model.attn.q_bmm_quantizer.num_bits = 8 if case == "int8" else (4, 3) + if case == "explicit": + model.attn._disable_fp8_mha = True + dtype = torch.float32 if case == "fp32" else torch.float16 + head_dim = 15 if case == "misaligned" else 16 + qkv = (torch.randn(2, 2, 4, head_dim, dtype=dtype) for _ in range(3)) + assert _fp8_mha_disabled(model.attn, *qkv) is (case != "eligible") + + +@pytest.mark.parametrize(("model_mixin", "enabled"), [(True, True), (True, False), (False, True)]) +def test_registration_and_enablement_guards(model_mixin, enabled): + model = DelegatedModel() if model_mixin else nn.Sequential(DelegatedAttention()) + _quantize(model, enabled) + attention = model.attn if model_mixin else model[0] + assert bool(getattr(attention, "_auto_fp8_mha", False)) is model_mixin + if model_mixin: + assert attention.q_bmm_quantizer.is_enabled is enabled