Emit VisualGen-compatible sparse_attention_config for diffusers skip-softmax export

CHANGELOG.rst

@@ -39,6 +39,7 @@ Changelog
 
 - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root.
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
+- Fix the diffusers skip-softmax HF export to emit the schema TensorRT-LLM VisualGen loads (``threshold_scale_factor.coefficients`` and a scalar ``target_sparsity``). Re-export checkpoints that failed to load.
 - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
 - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11.
 - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit.

examples/diffusers/README.md

@@ -361,7 +361,7 @@ By following these steps, your PEFT LoRA model should be efficiently quantized u
 
 ## Sparse Attention (Skip-Softmax)
 
-Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. Calibrated coefficients can be exported as a Hugging Face checkpoint (embedded in each component's `config.json` under `sparse_attention_config`) consumed directly by TRT-LLM's `SkipSoftmaxAttentionConfig.resolve_for_target_sparsity` — no extra conversion needed downstream.
+Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. Calibrated coefficients can be exported as a Hugging Face checkpoint (embedded in each component's `config.json` under `sparse_attention_config`) in the schema TRT-LLM VisualGen loads (`threshold_scale_factor.coefficients` plus a scalar `target_sparsity`) — no extra conversion needed downstream.
 
 ### Getting Started
 

examples/diffusers/sparsity/README.md

@@ -65,10 +65,9 @@ python wan22_skip_softmax.py \
 
 `--export-dir` writes a Hugging Face checkpoint with the calibrated
 `threshold_scale_factor` block embedded in each component's `config.json`
-(under the `sparse_attention_config` key). TRT-LLM's
-`SkipSoftmaxAttentionConfig.resolve_for_target_sparsity` reads the
-`(a, b)` directly via `coeffs['a'] * math.exp(coeffs['b'] * sparsity)` —
-no extra conversion needed downstream.
+(under the `sparse_attention_config` key), in the schema TRT-LLM VisualGen
+loads: `(a, b)` under `threshold_scale_factor.coefficients` and a scalar
+`target_sparsity`. No extra conversion is needed downstream.
 
 ## Threshold Modes
 

modelopt/torch/export/unified_export_hf.py

@@ -1376,7 +1376,7 @@ def _export_diffusers_checkpoint(
 
         # Update config.json with sparse attention info (both quantized and non-quantized)
         if export_sparse_attention_config is not None:
-            sparse_attn_config = export_sparse_attention_config(component)
+            sparse_attn_config = export_sparse_attention_config(component, for_diffusion=True)
             if sparse_attn_config is not None:
                 config_path = component_export_dir / "config.json"
                 if config_path.exists():

modelopt/torch/sparsity/attention_sparsity/conversion.py

@@ -394,7 +394,9 @@ def _get_sparse_softmax_export_config(module: SparseAttentionModule) -> dict[str
     return sparse_softmax_config
 
 
-def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None:
+def export_sparse_attention_config(
+    model: nn.Module, *, for_diffusion: bool = False
+) -> dict[str, Any] | None:
     """Extract sparse attention config for export to config.json.
 
     Extracts calibrated skip-softmax parameters and N:M sparse-softmax metadata
@@ -406,6 +408,10 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None:
 
     Args:
         model: Model with sparse attention applied
+        for_diffusion: Emit the single-phase schema read by the TRT-LLM VisualGen
+            runtime: one ``threshold_scale_factor.coefficients`` block and a scalar
+            ``target_sparsity``. The default is the LLM schema, with per-phase
+            (``prefill``/``decode``) blocks and a ``{phase: ratio}`` dict.
 
     Returns:
         Dictionary with sparse attention config for HuggingFace config.json export.
@@ -495,6 +501,21 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None:
         skip_group["threshold_scale_factor"] = export_threshold_scale_factor(calibration_params)
         if target_sparse_ratio is not None:
             skip_group["target_sparsity"] = target_sparse_ratio
+        if for_diffusion:
+            # Diffusion calibrates a single (prefill) phase; VisualGen rejects phase-keyed dicts.
+            tsf = skip_group["threshold_scale_factor"]
+            phases = {p: tsf.pop(p) for p in ("prefill", "decode") if p in tsf}
+            if phases:
+                tsf["coefficients"] = phases.get("prefill") or next(iter(phases.values()))
+            ratio = skip_group.get("target_sparsity")
+            if isinstance(ratio, dict):
+                scalar = (
+                    ratio["prefill"] if "prefill" in ratio else next(iter(ratio.values()), None)
+                )
+                if scalar is None:
+                    del skip_group["target_sparsity"]  # VisualGen needs a scalar
+                else:
+                    skip_group["target_sparsity"] = scalar
         config_groups[f"group_{group_idx}"] = skip_group
         group_idx += 1
 

tests/examples/diffusers/sparsity/test_sparsity.py

@@ -232,8 +232,9 @@ def test_wan22_calibrated_export(tiny_wan22_path, tmp_path):
         # threshold_scale_factor lives inside the skip_softmax group.
         tsf = group_0["threshold_scale_factor"]
         assert tsf["formula"] == "a * exp(b * target_sparsity)"
-        assert tsf["prefill"]["a"] == pytest.approx(math.exp(test_log_a))
-        assert tsf["prefill"]["b"] == pytest.approx(test_b)
+        assert tsf["coefficients"]["a"] == pytest.approx(math.exp(test_log_a))
+        assert tsf["coefficients"]["b"] == pytest.approx(test_b)
+        assert isinstance(group_0["target_sparsity"], float)  # scalar, as VisualGen expects
 
         # Calibrated mode — no raw_threshold.
         assert "raw_threshold" not in group_0

tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py

@@ -387,6 +387,15 @@ def test_exports_when_calibration_present(self):
         assert group_0["target_sparsity"] == {"prefill": 0.4, "decode": 0.6}
         assert out["producer"]["name"] == "modelopt"
 
+        # Diffusion (VisualGen) schema: one coefficient block, scalar target_sparsity.
+        group_0 = export_sparse_attention_config(model, for_diffusion=True)["config_groups"][
+            "group_0"
+        ]
+        tsf = group_0["threshold_scale_factor"]
+        assert tsf["coefficients"] == {"a": 3.14, "b": 7.5}
+        assert "prefill" not in tsf and "decode" not in tsf
+        assert group_0["target_sparsity"] == 0.4
+
     def test_exports_sparse_softmax_metadata(self):
         """N:M sparse-softmax params are exported without calibration metadata."""
         model = SimpleAttentionModel()