examples/diffusers/quantization/config.py (1)

39-39: Fix spacing inconsistency in configuration.

The configuration has inconsistent spacing after commas. Line 39 has a missing space after the comma between (4, 3) and "block_sizes".

-        "*[qkv]_bmm_quantizer": {"type": "dynamic", "num_bits": (4, 3),"block_sizes": {-2: 32}},
+        "*[qkv]_bmm_quantizer": {"type": "dynamic", "num_bits": (4, 3), "block_sizes": {-2: 32}},

modelopt/torch/quantization/export_onnx.py (2)

237-263: Consider adding validation for block_sizes parameter.

The new _fp8_block_quantize function should validate the structure and values of the block_sizes parameter to prevent runtime errors during ONNX export.

Add validation at the beginning of the function:

 def _fp8_block_quantize(
     g: torch.onnx._internal.jit_utils.GraphContext,
     inputs: torch.Value,
     trt_high_precision_dtype: str,
     block_sizes: list,
 ):
     """Helper Function for Quantization."""
+    if not isinstance(block_sizes, list) or not block_sizes:
+        raise ValueError(f"block_sizes must be a non-empty list, got {block_sizes}")
+    if not all(isinstance(b, int) and b > 0 for b in block_sizes):
+        raise ValueError(f"All block sizes must be positive integers, got {block_sizes}")
+        
     output_shape = sym_help._get_tensor_sizes(inputs)

---

534-535: Fix typo in comment.

-        # We cannot do block quant for the softmax's output 
+        # We cannot do block quant for the softmax's output

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (2)

653-653: Remove trailing whitespace from blank lines.

These blank lines contain unnecessary whitespace which violates Python style guidelines.

-        
+
         Args:
             shape: The tensor shape to use for conversion (can be tuple or torch.Size)
-            
+
         Returns:
             List of block sizes for each dimension, or None if block_sizes is None
-            
+
         Example:

Also applies to: 656-656, 659-659

---

961-962: Consider thread-safety for the _original_input_shape attribute.

Setting and deleting _original_input_shape as a temporary attribute could cause issues in multi-threaded scenarios where the same quantizer is used by multiple threads simultaneously.

Consider using a context manager or local variable approach instead:

-            setattr(self, "_original_input_shape", inputs.shape)
-            inputs = self._process_for_blockquant(inputs)
+            original_shape = inputs.shape
+            inputs = self._process_for_blockquant(inputs)
+            # Pass original_shape to methods that need it

Alternatively, consider storing it in a thread-local storage if multi-threading support is required.

modelopt/torch/quantization/plugins/diffusers.py (6)

117-124: Fix assertion logic for mixed dynamic/non-dynamic quantizers

The current implementation requires all QKV quantizers to be either dynamic or non-dynamic together. However, the logic flow is problematic - if they're all non-dynamic, scales are computed, but if any is dynamic, it asserts all must be dynamic. This creates a rigid constraint that may not be necessary for all use cases.

Consider refactoring to handle mixed cases more gracefully:

-    if not self.q_bmm_quantizer._dynamic and not self.k_bmm_quantizer._dynamic and not self.v_bmm_quantizer._dynamic:
-        q_quantized_scale = self.q_bmm_quantizer._get_amax(query)
-        k_quantized_scale = self.k_bmm_quantizer._get_amax(key)
-        v_quantized_scale = self.v_bmm_quantizer._get_amax(value)
-    else:
-        assert self.q_bmm_quantizer._dynamic and self.k_bmm_quantizer._dynamic and self.v_bmm_quantizer._dynamic, "QKV QDQS must be in the same type"
-        q_quantized_scale, k_quantized_scale, v_quantized_scale = None, None, None
+    # Compute scales for non-dynamic quantizers, set None for dynamic ones
+    q_quantized_scale = None if self.q_bmm_quantizer._dynamic else self.q_bmm_quantizer._get_amax(query)
+    k_quantized_scale = None if self.k_bmm_quantizer._dynamic else self.k_bmm_quantizer._get_amax(key)
+    v_quantized_scale = None if self.v_bmm_quantizer._dynamic else self.v_bmm_quantizer._get_amax(value)
+    
+    # Optionally validate consistency if needed
+    dynamic_states = [self.q_bmm_quantizer._dynamic, self.k_bmm_quantizer._dynamic, self.v_bmm_quantizer._dynamic]
+    if len(set(dynamic_states)) > 1:
+        # Log warning or handle mixed dynamic states if necessary
+        pass

---

122-122: Fix line length violation

Line 122 exceeds the 120 character limit (149 characters).

-        assert self.q_bmm_quantizer._dynamic and self.k_bmm_quantizer._dynamic and self.v_bmm_quantizer._dynamic, "QKV QDQS must be in the same type"
+        assert (self.q_bmm_quantizer._dynamic and 
+                self.k_bmm_quantizer._dynamic and 
+                self.v_bmm_quantizer._dynamic), "QKV QDQS must be in the same type"

---

144-146: Remove trailing whitespace

Line 145 has trailing whitespace after the comma.

             q_block_sizes,
-            k_block_sizes, 
+            k_block_sizes,
             v_block_sizes,

---

231-233: Inconsistent default values for scale parameters

The scale parameters have inconsistent default values in the symbolic method signature (float | None = 1.0) which doesn't match the forward method where they default to None.

-        q_quantized_scale: float | None = 1.0,
-        k_quantized_scale: float | None = 1.0,
-        v_quantized_scale: float | None = 1.0,
+        q_quantized_scale: float | None = None,
+        k_quantized_scale: float | None = None,
+        v_quantized_scale: float | None = None,

---

200-202: Consider using TypeAlias for block shape type consistency

The block shape parameters use list | None type annotations repeatedly. Consider defining a type alias for better maintainability and consistency.

Add at the top of the file after imports:

from typing import TypeAlias

BlockShape: TypeAlias = list[int] | None

Then update the signatures:

-        q_block_shape: list | None = None,
-        k_block_shape: list | None = None,
-        v_block_shape: list | None = None,
+        q_block_shape: BlockShape = None,
+        k_block_shape: BlockShape = None,
+        v_block_shape: BlockShape = None,

Also applies to: 236-238

---

126-128: Add validation for block sizes consistency

The code retrieves block sizes from quantizers but doesn't validate that they're compatible with the actual tensor shapes or the quantization configuration.

Consider adding validation to ensure block sizes are appropriate:

     # Get block sizes lists for each quantizer if needed
     q_block_sizes = self.q_bmm_quantizer._get_block_sizes_list(query.shape)
     k_block_sizes = self.k_bmm_quantizer._get_block_sizes_list(key.shape)
     v_block_sizes = self.v_bmm_quantizer._get_block_sizes_list(value.shape)
+    
+    # Validate block sizes if dynamic quantization is enabled
+    if self.q_bmm_quantizer._dynamic and q_block_sizes:
+        for dim, block_size in enumerate(q_block_sizes):
+            if block_size > 1 and query.shape[dim] % block_size != 0:
+                raise ValueError(f"Query dimension {dim} (size {query.shape[dim]}) is not divisible by block size {block_size}")
+    # Similar validation for k and v can be added

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

959-963: Ensure _original_input_shape is always cleaned up (wrap in try/finally)

State can leak if an exception occurs between setting and deleting _original_input_shape. Wrap the forward section after setup in try/finally.

-        if (
+        cleanup_original_shape = False
+        if (
             self.block_sizes is not None
             and self.block_sizes.get("type", None) != "dynamic"
             and self._fake_quant
         ):
             # Reshape is required if the logic isnt handled in the simulation kernel
             self._setup_for_blockquant(inputs)
             setattr(self, "_original_input_shape", inputs.shape)
+            cleanup_original_shape = True
             inputs = self._process_for_blockquant(inputs)
 
-        outputs = inputs
+        try:
+            outputs = inputs
             ...
-        if hasattr(self, "_original_input_shape"):
-            delattr(self, "_original_input_shape")
+        finally:
+            if cleanup_original_shape and hasattr(self, "_original_input_shape"):
+                delattr(self, "_original_input_shape")

Also applies to: 1002-1003

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

681-699: Comment is misleading within dynamic-only branch

The comment says “Double scale Block quantization, including dynamic and static block quantization” but this branch executes only when type == "dynamic". Tighten the comment to avoid confusion.

-            # Double scale Block quantization, including dynamic and static block quantization
+            # Dynamic double-scale block quantization path

modelopt/torch/quantization/plugins/diffusers.py (2)

117-132: Align QKV mode and error message; avoid computing scales in export path

• Message “QKV QDQS must be in the same type” is unclear. Make it explicit: “Q, K, and V quantizers must all be dynamic or all be static.”
• Skip _get_amax when exporting; it’s unused at runtime and can be None for dynamic. Guard by torch.onnx.is_in_onnx_export().

-    if (
+    if (
         not self.q_bmm_quantizer._dynamic
         and not self.k_bmm_quantizer._dynamic
         and not self.v_bmm_quantizer._dynamic
     ):
-        q_quantized_scale = self.q_bmm_quantizer._get_amax(query)
-        k_quantized_scale = self.k_bmm_quantizer._get_amax(key)
-        v_quantized_scale = self.v_bmm_quantizer._get_amax(value)
+        if not torch.onnx.is_in_onnx_export():
+            q_quantized_scale = self.q_bmm_quantizer._get_amax(query)
+            k_quantized_scale = self.k_bmm_quantizer._get_amax(key)
+            v_quantized_scale = self.v_bmm_quantizer._get_amax(value)
+        else:
+            q_quantized_scale = k_quantized_scale = v_quantized_scale = None
     else:
         assert (
             self.q_bmm_quantizer._dynamic
             and self.k_bmm_quantizer._dynamic
             and self.v_bmm_quantizer._dynamic
-        ), "QKV QDQS must be in the same type"
+        ), "Q, K, and V quantizers must all be dynamic or all be static."
         q_quantized_scale, k_quantized_scale, v_quantized_scale = None, None, None

---

133-137: Using a private helper across modules; consider promoting to public API

Calling _get_block_sizes_list from another module couples to a private method. Expose it as a public helper (e.g., get_block_sizes_list or a util function) to avoid brittle dependencies.

Would you like me to extract a small utility, e.g., modelopt/torch/quantization/utils/block_sizes.py:get_block_sizes_list(shape, block_sizes)?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

959-963: Guarantee cleanup of _original_input_shape on exceptions.

Deletion is unconditional but not exception-safe; wrap the quantization region in try/finally. Also fix the typo “isnt” → “isn't”.

-        if (
+        cleanup_original_input_shape = False
+        if (
             self.block_sizes is not None
             and self.block_sizes.get("type", None) != "dynamic"
             and self._fake_quant
         ):
-            # Reshape is required if the logic isnt handled in the simulation kernel
+            # Reshape is required if the logic isn't handled in the simulation kernel
             self._setup_for_blockquant(inputs)
             setattr(self, "_original_input_shape", inputs.shape)
+            cleanup_original_input_shape = True
             inputs = self._process_for_blockquant(inputs)
 
-        outputs = inputs
+        try:
+            outputs = inputs
@@
-        if (
+            if (
                 self.block_sizes is not None
                 and self.block_sizes.get("type", None) != "dynamic"
                 and self._fake_quant
-        ):
-            outputs = self._reset_to_original_shape(outputs)
-
-        if hasattr(self, "_original_input_shape"):
-            delattr(self, "_original_input_shape")
-        return outputs
+            ):
+                outputs = self._reset_to_original_shape(outputs)
+            return outputs
+        finally:
+            if cleanup_original_input_shape and hasattr(self, "_original_input_shape"):
+                delattr(self, "_original_input_shape")

Also applies to: 1002-1003

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

651-673: Type-safety and input validation for _get_block_sizes_list.

Add explicit typing, validate keys, and guard length mismatches to avoid silently passing malformed shapes to downstream ONNX ops.

-def _get_block_sizes_list(self, shape):
+from typing import Sequence
+
+def _get_block_sizes_list(self, shape: Sequence[int] | torch.Size) -> list[int] | None:
@@
-        block_sizes_list = []
-        for dim in range(len(shape)):
+        # Only allow integer axes plus known metadata keys.
+        valid_meta = {"type", "scale_bits", "scale_block_sizes"}
+        assert all(
+            isinstance(k, int) or k in valid_meta for k in self.block_sizes.keys()
+        ), f"Invalid block_sizes keys: {list(self.block_sizes.keys())}"
+
+        rank = len(shape)
+        block_sizes_list: list[int] = []
+        for dim in range(rank):
             # Check both positive and negative dimension indices
-            dim_negative = dim - len(shape)
+            dim_negative = dim - rank
             block_size = self.block_sizes.get(dim, None) or self.block_sizes.get(dim_negative, None)
             block_sizes_list.append(block_size if block_size is not None else 1)
         return block_sizes_list

modelopt/torch/quantization/export_onnx.py (2)

238-265: Validate block_shape rank and surface clearer errors.

Add a rank check before emitting TRT_DynamicQuantize; mis-sized block_shapes currently fall through to TRT with cryptic errors.

 def _fp8_block_quantize(
@@
-    input_type = inputs.type().scalarType()
+    input_type = inputs.type().scalarType()
+    rank = symbolic_helper._get_tensor_rank(inputs)
+    assert rank is not None, "Input rank must be known at export time."
+    assert len(block_sizes) == rank, (
+        f"block_shape length ({len(block_sizes)}) must match input rank ({rank})."
+    )
@@
     quantized_output, scales_output = g.op(
         "trt::TRT_DynamicQuantize",
         inputs,
         block_shape_i=block_sizes,

---

503-509: Block-shape consistency in FP8 MHA path.

Validate q/k/v block shapes match input ranks; also ensure softmax path never receives a block shape.

-        query_scaled = export_fp8(
-            g, query_scaled, q_quantized_scale, high_precision_flag, q_block_shape
-        )
+        assert (q_block_shape is None) or (
+            len(q_block_shape) == symbolic_helper._get_tensor_rank(query_scaled)
+        ), "q_block_shape rank mismatch."
+        query_scaled = export_fp8(g, query_scaled, q_quantized_scale, high_precision_flag, q_block_shape)
@@
-        key_transposed_scaled = export_fp8(
-            g, key_transposed_scaled, k_quantized_scale, high_precision_flag, k_block_shape
-        )
+        assert (k_block_shape is None) or (
+            len(k_block_shape) == symbolic_helper._get_tensor_rank(key_transposed_scaled)
+        ), "k_block_shape rank mismatch."
+        key_transposed_scaled = export_fp8(g, key_transposed_scaled, k_quantized_scale, high_precision_flag, k_block_shape)
@@
-        # We cannot do block quant for the softmax's output
-        attn_weight = export_fp8(g, attn_weight, 1.0, high_precision_flag, None)
+        # We cannot do block quant for the softmax's output
+        attn_weight = export_fp8(g, attn_weight, 1.0, high_precision_flag, None)
@@
-        value = export_fp8(g, value, v_quantized_scale, high_precision_flag, v_block_shape)
+        assert (v_block_shape is None) or (
+            len(v_block_shape) == symbolic_helper._get_tensor_rank(value)
+        ), "v_block_shape rank mismatch."
+        value = export_fp8(g, value, v_quantized_scale, high_precision_flag, v_block_shape)

Also applies to: 535-549

modelopt/torch/quantization/plugins/diffusers.py (2)

117-132: Q/K/V quantization mode must match: improve error and skip redundant work.

The assertion is good. Minor: clarify message and avoid computing per-tensor amax if any quantizer is dynamic.

-        ), "QKV QDQS must be in the same type"
+        ), "Q/K/V quantization modes must match: either all dynamic or all static."

---

133-137: Guard block size list creation when block_sizes is None.

_if a quantizer has no block_sizes, _get_block_sizes_list returns None; that’s fine. Add a quick comment to make intent explicit and future-proof.

-    q_block_sizes = self.q_bmm_quantizer._get_block_sizes_list(query.shape)  # type: ignore[union-attr]
+    # Returns None for per-tensor paths; ONNX export handles that by taking the non-block path.
+    q_block_sizes = self.q_bmm_quantizer._get_block_sizes_list(query.shape)  # type: ignore[union-attr]

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (2)

649-671: Return None for “metadata‑only” block_sizes and add typing.

Avoid forcing block mode with [1,…,1] when there are no per‑dim entries; this can route ONNX path incorrectly.

Apply:

-    def _get_block_sizes_list(self, shape):
+    def _get_block_sizes_list(self, shape) -> list[int] | None:
         """Convert block_sizes dict to list format based on tensor shape.
@@
-        if self.block_sizes is None:
+        if self.block_sizes is None:
             return None
-
-        block_sizes_list = []
+        # If there are no integer dimension entries with a meaningful block size, treat as no block quant.
+        has_dim_sizes = any(
+            isinstance(k, int) and (v is not None and v != 1) for k, v in self.block_sizes.items()
+        )
+        if not has_dim_sizes:
+            return None
+
+        block_sizes_list: list[int] = []
         for dim in range(len(shape)):
             # Check both positive and negative dimension indices
             dim_negative = dim - len(shape)
             block_size = self.block_sizes.get(dim, None) or self.block_sizes.get(dim_negative, None)
             block_sizes_list.append(block_size if block_size is not None else 1)
         return block_sizes_list

---

1005-1007: Centralize deletion of _original_input_shape in a finally block.

Move this into a single finally at the end of forward so it runs regardless of success/failure.

Apply:

-        if hasattr(self, "_original_input_shape"):
-            delattr(self, "_original_input_shape")
-        return outputs
+        try:
+            return outputs
+        finally:
+            if hasattr(self, "_original_input_shape"):
+                delattr(self, "_original_input_shape")

modelopt/torch/quantization/export_onnx.py (3)

294-321: Remove brittle default [1,1,128,1] and validate rank in _fp8_block_dequantize.

Defaulting silently is dangerous and shape-dependent. Require explicit block_sizes and assert correctness. (Same concern raised earlier.)

Apply:

-def _fp8_block_dequantize(
+def _fp8_block_dequantize(
     g: torch.onnx._internal.jit_utils.GraphContext,
     inputs: torch.Value,
     scales: torch.Value,
     trt_high_precision_dtype: str,
     otype: str | None = None,
-    block_sizes: list = [1, 1, 128, 1],
+    block_sizes: list,
 ):
     """Helper Function for Dequantization."""
     output_shape = sym_help._get_tensor_sizes(inputs)
+    # Validate block shape
+    rank = sym_help._get_tensor_rank(inputs)
+    assert rank is not None, "Input rank must be known at export time."
+    assert isinstance(block_sizes, (list, tuple)) and len(block_sizes) == rank, (
+        f"block_shape length ({len(block_sizes)}) must match input rank ({rank})."
+    )
+    assert all(isinstance(b, int) and b > 0 for b in block_sizes), (
+        "All entries in block_shape must be positive integers."
+    )
+    if otype is None:
+        otype = inputs.type().scalarType()

---

323-345: Handle non-Python amax safely and validate block_shapes before block Q/DQ path.

float(amax) will break when amax is a graph Value/0‑dim tensor; also assert block_shapes align with input rank before calling block ops. (Echoing prior comment.)

Apply:

 def export_fp8(
     g: torch.onnx._internal.jit_utils.GraphContext,
     inputs: torch.Value,
-    amax: float | None,
+    amax: float | None,
     trt_high_precision_dtype: str | None,
     block_sizes: list | None,
 ):
     """Export quantized model to FP8 ONNX."""
-    scale = 1.0 if amax is None else 448.0 / float(amax)
+    if amax is None:
+        scale = 1.0
+    else:
+        amax_const = sym_help._get_const(amax, "f", "amax")
+        # If not a constant at export time, fall back to neutral scale to avoid exporter errors.
+        scale = 1.0 if (amax_const is None or amax_const == 0) else 448.0 / float(amax_const)
@@
-    if not block_sizes:
+    if not block_sizes:
         q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype)
         return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype, otype)
     else:
+        # Validate block shape early
+        rank = sym_help._get_tensor_rank(inputs)
+        assert rank is not None, "Input rank must be known at export time."
+        assert isinstance(block_sizes, (list, tuple)) and len(block_sizes) == rank, (
+            f"block_shape length ({len(block_sizes)}) must match input rank ({rank})."
+        )
+        assert all(isinstance(b, int) and b > 0 for b in block_sizes), (
+            "All entries in block_shape must be positive integers."
+        )
         q_tensor, scales_output = _fp8_block_quantize(
             g, inputs, trt_high_precision_dtype, block_sizes
         )
         return _fp8_block_dequantize(
             g, q_tensor, scales_output, trt_high_precision_dtype, otype, block_sizes
         )

---

238-265: Validate block_shape against input rank and values in _fp8_block_quantize.

Guard against mismatched ranks and non-positive entries to avoid invalid custom op attributes at export time.

Apply:

 def _fp8_block_quantize(
     g: torch.onnx._internal.jit_utils.GraphContext,
     inputs: torch.Value,
     trt_high_precision_dtype: str,
     block_sizes: list,
 ):
     """Helper Function for Quantization."""
     output_shape = sym_help._get_tensor_sizes(inputs)
 
+    # Validate block shape
+    rank = sym_help._get_tensor_rank(inputs)
+    assert rank is not None, "Input rank must be known at export time."
+    assert isinstance(block_sizes, (list, tuple)) and len(block_sizes) == rank, (
+        f"block_shape length ({len(block_sizes)}) must match input rank ({rank})."
+    )
+    assert all(isinstance(b, int) and b > 0 for b in block_sizes), (
+        "All entries in block_shape must be positive integers."
+    )

modelopt/torch/quantization/export_onnx.py (1)

512-518: Pre-validate q/k block_shapes vs tensor ranks to fail fast.

Catch mismatches early instead of deep inside TRT ops.

Apply:

-        query_scaled = export_fp8(
+        # Sanity-check block shapes
+        for name, t, bs in (("q", query_scaled, q_block_shape), ("k", key_transposed_scaled, k_block_shape)):
+            if bs is not None:
+                r = sym_help._get_tensor_rank(t)
+                assert r is not None and len(bs) == r, f"{name}_block_shape must match rank ({r})."
+        query_scaled = export_fp8(
             g, query_scaled, q_quantized_scale, high_precision_flag, q_block_shape
         )
@@
-        key_transposed_scaled = export_fp8(
+        key_transposed_scaled = export_fp8(
             g, key_transposed_scaled, k_quantized_scale, high_precision_flag, k_block_shape
         )

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

962-965: Guarantee cleanup of _original_input_shape via try/finally; fix typo.

Ensure attribute is deleted even on exceptions; also fix “isnt” -> “isn't”.

Apply:

-            # Reshape is required if the logic isnt handled in the simulation kernel
-            self._setup_for_blockquant(inputs)
-            setattr(self, "_original_input_shape", inputs.shape)
-            inputs = self._process_for_blockquant(inputs)
+            # Reshape is required if the logic isn't handled in the simulation kernel
+            cleanup_original_shape = False
+            try:
+                self._setup_for_blockquant(inputs)
+                setattr(self, "_original_input_shape", inputs.shape)
+                cleanup_original_shape = True
+                inputs = self._process_for_blockquant(inputs)
+            except Exception:
+                # Make sure we don't leak transient attributes on failure
+                if cleanup_original_shape and hasattr(self, "_original_input_shape"):
+                    delattr(self, "_original_input_shape")
+                raise

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (2)

649-671: Don’t force “block mode” when block_sizes has only metadata; return None.

Current logic returns [1, ...] even if block_sizes has no per‑dim sizes, unintentionally enabling block handling downstream. Return None when there are no integer keys with v not in {None, 1}. Add typing while here.

-    def _get_block_sizes_list(self, shape):
+    def _get_block_sizes_list(self, shape) -> list[int] | None:
@@
-        if self.block_sizes is None:
-            return None
-
-        block_sizes_list = []
+        if self.block_sizes is None:
+            return None
+        # If there are no per-dimension entries with a real block size, treat as no block quant.
+        has_dim_sizes = any(
+            isinstance(k, int) and (v is not None and v != 1)
+            for k, v in self.block_sizes.items()
+        )
+        if not has_dim_sizes:
+            return None
+
+        block_sizes_list: list[int] = []
         for dim in range(len(shape)):
             # Check both positive and negative dimension indices
             dim_negative = dim - len(shape)
             block_size = self.block_sizes.get(dim, None) or self.block_sizes.get(dim_negative, None)
             block_sizes_list.append(block_size if block_size is not None else 1)
         return block_sizes_list

---

1006-1007: Guarantee deletion of _original_input_shape with try/finally (prev feedback).

Move deletion into a finally block guarded by the local flag so it always runs.

-        if hasattr(self, "_original_input_shape"):
-            delattr(self, "_original_input_shape")
-        return outputs
+        try:
+            return outputs
+        finally:
+            if cleanup_original_input_shape and hasattr(self, "_original_input_shape"):
+                delattr(self, "_original_input_shape")

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

962-966: Make _original_input_shape cleanup exception‑safe: set flag when set, clean in finally.

Set a local flag when creating the attribute so the finally block can reliably clean it even if an exception occurs later.

         if (
             self.block_sizes is not None
             and self.block_sizes.get("type", None) != "dynamic"
             and self._fake_quant
         ):
             # Reshape is required if the logic is not handled in the simulation kernel
             # Only MX format and NVFP4 reshape are currently supported by the kernel.
             self._setup_for_blockquant(inputs)
-            setattr(self, "_original_input_shape", inputs.shape)
+            setattr(self, "_original_input_shape", inputs.shape)
+            cleanup_original_input_shape = True
             inputs = self._process_for_blockquant(inputs)

Add the flag near the top of forward (before this block):

-        # Rotating the input
+        cleanup_original_input_shape = False
+        # Rotating the input

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/quantization/export_onnx.py (2)

299-325: Remove brittle default block_sizes=[1,1,128,1] and validate rank.

A hardcoded default of [1, 1, 128, 1] is unsafe—it assumes a specific tensor rank and shape. Require callers to pass block_sizes explicitly and validate it against the input tensor rank.

 def _fp8_block_dequantize(
     g: torch.onnx._internal.jit_utils.GraphContext,
     inputs: torch.Value,
     scales: torch.Value,
     trt_high_precision_dtype: str,
     otype: str | None = None,
-    block_sizes: list = [1, 1, 128, 1],
+    block_sizes: list | None = None,
 ):
     """Helper Function for Dequantization."""
+    assert block_sizes is not None, "block_sizes must be provided for block dequantization."
     output_shape = sym_help._get_tensor_sizes(inputs)
+    assert len(block_sizes) == len(output_shape), (
+        f"block_sizes length ({len(block_sizes)}) must match input rank ({len(output_shape)})."
+    )
     assert trt_high_precision_dtype in (otype, "Float"), (
         "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float."
     )
     out = g.op(
         "trt::TRT_BlockDequantize",
         inputs,
         scales,
         block_shape_i=block_sizes,
     ).setType(
         inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape)
     )
 
     # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT
     # custom ops, so cast the output if needed.
     if trt_high_precision_dtype != otype:
         out = g.op("Cast", out, to_i=onnx_dtype_map[otype])  # type: ignore[index]
     return out

Based on learnings

---

328-349: Handle tensor/Value amax correctly during symbolic export.

float(amax) on line 336 will fail if amax is a graph Value or 0-dim tensor. Use _get_const to extract Python floats when possible, or handle symbolic amax.

 def export_fp8(
     g: torch.onnx._internal.jit_utils.GraphContext,
     inputs: torch.Value,
     amax: float | None,
     trt_high_precision_dtype: str | None,
     block_sizes: list | None,
 ):
     """Export quantized model to FP8 ONNX."""
-    scale = 1.0 if amax is None else 448.0 / float(amax)
+    # Accept Python floats or Tensor constants; None indicates dynamic/block path not using amax.
+    if amax is None:
+        scale = 1.0
+    else:
+        amax_const = sym_help._get_const(amax, "f", "amax") if hasattr(amax, "node") else amax
+        # If still not constant, fall back to 1.0 to avoid build-time failures.
+        scale = 1.0 if amax_const is None else 448.0 / float(amax_const)
     otype = inputs.type().scalarType()
     if trt_high_precision_dtype is None:
         trt_high_precision_dtype = otype
     if not block_sizes:
         q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype)
         return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype, otype)
     else:
         q_tensor, scales_output = _fp8_block_quantize(
             g, inputs, trt_high_precision_dtype, block_sizes
         )
         return _fp8_block_dequantize(
             g, q_tensor, scales_output, trt_high_precision_dtype, otype, block_sizes
         )

Based on learnings

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (2)

622-643: Guard against metadata-only block_sizes dicts returning [1,1,1,1].

If block_sizes contains only non-integer keys like {"type": "dynamic", "scale_bits": (4,3)}, you return [1]*rank, which signals "block mode" downstream even though no dimension has an actual block size. This can incorrectly trigger the FP8 block path.

Return None when there are no integer dimension keys with a block size > 1.

 def _get_block_sizes_list(self, shape):
     """Convert block_sizes dict to list format based on tensor shape.
     
     Args:
         shape: The tensor shape to use for conversion (can be tuple or torch.Size)
     
     Returns:
         List of block sizes for each dimension, or None if block_sizes is None
     
     Example:
         block_sizes = {-2: 32} with shape [2, 24, 4608, 128] -> [1, 1, 32, 1]
     """
     if self.block_sizes is None:
         return None
+    
+    # Check if there are any actual per-dimension block sizes (not just metadata)
+    has_dim_sizes = any(
+        isinstance(k, int) and v not in (None, 1) for k, v in self.block_sizes.items()
+    )
+    if not has_dim_sizes:
+        return None
     
     block_sizes_list = []
     for dim in range(len(shape)):
         # Check both positive and negative dimension indices
         dim_negative = dim - len(shape)
         block_size = self.block_sizes.get(dim, None) or self.block_sizes.get(dim_negative, None)
         block_sizes_list.append(block_size if block_size is not None else 1)
     return block_sizes_list

Based on learnings

---

975-976: Cleanup of _original_input_shape should be in a try-finally block.

The deletion of _original_input_shape at the end of forward is not guaranteed if an exception occurs during the forward pass. This can leave stale state that corrupts subsequent calls.

Wrap the forward logic that depends on _original_input_shape in a try-finally block to guarantee cleanup.

         if (
             self.block_sizes is not None
             and self.block_sizes.get("type", None) != "dynamic"
             and self._fake_quant
         ):
             # Reshape is required if the logic is not handled in the simulation kernel
             # Only MX format and NVFP4 reshape are currently supported by the kernel.
             self._setup_for_blockquant(inputs)
             setattr(self, "_original_input_shape", inputs.shape)
             inputs = self._process_for_blockquant(inputs)
 
-        outputs = inputs
-
-        block_size = None
-        if self._if_calib and not self._dynamic:
-            if self._calibrator is None:
-                raise RuntimeError("Calibrator was not created.")
-            # Shape is only known when it sees the first tensor
-            if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic":
-                block_size = self.block_sizes.get(-1, None) or self.block_sizes.get(
-                    inputs.dim() - 1, None
-                )
-                assert block_size is not None, "block size for dynamic quantization not found."
-
-            self.collect(inputs)
-
-        if self._if_quant:
-            # Check if the input tensor is contiguous
-            # Non-contiguous tensors will generate incorrect FP4 quantization results
-            if hasattr(inputs, "is_contiguous") and not inputs.is_contiguous():
-                inputs.data = inputs.data.contiguous()
-            if self.fake_quant:
-                outputs = self._fake_quantize(inputs)
-            elif not self._dequantize:
-                outputs = self._real_quantize(inputs)
-            else:
-                raise ValueError(
-                    "self._dequantize is True and self.fake_quant is False. "
-                    "This case should have been handled."
-                )
-
-        if (
-            self.block_sizes is not None
-            and self.block_sizes.get("type", None) != "dynamic"
-            and self._fake_quant
-        ):
-            outputs = self._reset_to_original_shape(outputs)
-
-        if hasattr(self, "_original_input_shape"):
-            delattr(self, "_original_input_shape")
+        try:
+            outputs = inputs
+
+            block_size = None
+            if self._if_calib and not self._dynamic:
+                if self._calibrator is None:
+                    raise RuntimeError("Calibrator was not created.")
+                # Shape is only known when it sees the first tensor
+                if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic":
+                    block_size = self.block_sizes.get(-1, None) or self.block_sizes.get(
+                        inputs.dim() - 1, None
+                    )
+                    assert block_size is not None, "block size for dynamic quantization not found."
+
+                self.collect(inputs)
+
+            if self._if_quant:
+                # Check if the input tensor is contiguous
+                # Non-contiguous tensors will generate incorrect FP4 quantization results
+                if hasattr(inputs, "is_contiguous") and not inputs.is_contiguous():
+                    inputs.data = inputs.data.contiguous()
+                if self.fake_quant:
+                    outputs = self._fake_quantize(inputs)
+                elif not self._dequantize:
+                    outputs = self._real_quantize(inputs)
+                else:
+                    raise ValueError(
+                        "self._dequantize is True and self.fake_quant is False. "
+                        "This case should have been handled."
+                    )
+
+            if (
+                self.block_sizes is not None
+                and self.block_sizes.get("type", None) != "dynamic"
+                and self._fake_quant
+            ):
+                outputs = self._reset_to_original_shape(outputs)
+        finally:
+            if hasattr(self, "_original_input_shape"):
+                delattr(self, "_original_input_shape")
+                
         return outputs

Based on learnings

modelopt/torch/quantization/plugins/diffusers.py (1)

231-235: Fix ONNX parse_args: allow optional scales and block-shapes (prevents export crash).

• q/k/v scales can be None for dynamic quantization; using 't' rejects None.
• q/k/v block-shapes are optional; using 'is' rejects None.

Update parse_args to accept optionals with 'v' for these 6 entries.

-    @symbolic_helper.parse_args(
-        "v", "v", "v", "v", "f", "b", "v", "t", "t", "t", "s", "b", "is", "is", "is"
-    )
+    @symbolic_helper.parse_args(
+        "v", "v", "v", "v", "f", "b", "v", "v", "v", "v", "s", "b", "v", "v", "v"
+    )

Also normalize the 3 block-shape values to Python lists inside symbolic before forwarding (see next comment). Based on learnings.

Also applies to: 243-251

modelopt/torch/quantization/plugins/diffusers.py (4)

120-135: Clarify assertion message for mixed dynamic/static Q/K/V quantizers.

Message is cryptic. Make intent explicit.

-        ), "QKV QDQS must be in the same type"
+        ), "q/k/v quantizers must all be either dynamic or static"

---

243-251: Adjust symbolic type hints to reflect actual tensor/value types.

Scales are Values/Tensors, not floats; block-shapes are list[int] | None.

-        scale: torch._C.Value | None = None,
-        q_quantized_scale: float | None = 1.0,
-        k_quantized_scale: float | None = 1.0,
-        v_quantized_scale: float | None = 1.0,
+        scale: "torch._C.Value | None" = None,
+        q_quantized_scale: "torch._C.Value | None" = None,
+        k_quantized_scale: "torch._C.Value | None" = None,
+        v_quantized_scale: "torch._C.Value | None" = None,
         high_precision_flag: str = "Half",
         disable_fp8_mha: bool = True,
-        q_block_shape: list | None = None,
-        k_block_shape: list | None = None,
-        v_block_shape: list | None = None,
+        q_block_shape: "list[int] | None" = None,
+        k_block_shape: "list[int] | None" = None,
+        v_block_shape: "list[int] | None" = None,

---

136-139: Block-shape None path is now safe once parse_args is fixed; optional naming nit.

With the parse_args/symbolic normalization above, passing None here is OK. Consider renaming variables to *_block_shape for consistency with FP8SDPA.forward.

-    q_block_sizes = self.q_bmm_quantizer._get_block_sizes_list(query.shape)  # type: ignore[union-attr]
-    k_block_sizes = self.k_bmm_quantizer._get_block_sizes_list(key.shape)  # type: ignore[union-attr]
-    v_block_sizes = self.v_bmm_quantizer._get_block_sizes_list(value.shape)  # type: ignore[union-attr]
+    q_block_shape = self.q_bmm_quantizer._get_block_sizes_list(query.shape)  # type: ignore[union-attr]
+    k_block_shape = self.k_bmm_quantizer._get_block_sizes_list(key.shape)  # type: ignore[union-attr]
+    v_block_shape = self.v_bmm_quantizer._get_block_sizes_list(value.shape)  # type: ignore[union-attr]

Remember to update the call site accordingly (Lines 155-157).

---

155-157: *If you rename to _block_shape, update the call site accordingly.

-            q_block_sizes,
-            k_block_sizes,
-            v_block_sizes,
+            q_block_shape,
+            k_block_shape,
+            v_block_shape,

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro