From 59d7dc8d326b8f7442c308e10ed16845abd90596 Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Sun, 18 Oct 2026 22:14:06 -0700 Subject: [PATCH 1/2] Emit VisualGen-compatible sparse_attention_config for diffusers skip-softmax export Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Zhiyu Cheng --- CHANGELOG.rst | 1 + examples/diffusers/README.md | 2 +- examples/diffusers/sparsity/README.md | 7 +++---- modelopt/torch/export/unified_export_hf.py | 2 +- .../sparsity/attention_sparsity/conversion.py | 17 ++++++++++++++++- .../diffusers/sparsity/test_sparsity.py | 5 +++-- .../test_sparse_attention_conversion.py | 9 +++++++++ 7 files changed, 34 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7974cad044f..5ac7132fbed 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 8d6e61e9644..bb6330db23c 100644 --- a/examples/diffusers/README.md +++ b/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 diff --git a/examples/diffusers/sparsity/README.md b/examples/diffusers/sparsity/README.md index e102f1b6d5d..bc1188d2ea0 100644 --- a/examples/diffusers/sparsity/README.md +++ b/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 diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a5b09a2f974..14ce84e13d5 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/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(): diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index b2acd137930..7640628dd25 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/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,15 @@ 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): + skip_group["target_sparsity"] = ratio.get("prefill", next(iter(ratio.values()))) config_groups[f"group_{group_idx}"] = skip_group group_idx += 1 diff --git a/tests/examples/diffusers/sparsity/test_sparsity.py b/tests/examples/diffusers/sparsity/test_sparsity.py index 4f57ba8ecda..1727ee39b7f 100644 --- a/tests/examples/diffusers/sparsity/test_sparsity.py +++ b/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 diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py index 6cfdd2775b9..d57e914d641 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py +++ b/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() From cb0cc62feaa3ade8a779a8c56099c66a1db21415 Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Wed, 21 Oct 2026 12:12:27 -0700 Subject: [PATCH 2/2] Avoid eager default and StopIteration when scalarizing target_sparsity Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Zhiyu Cheng --- modelopt/torch/sparsity/attention_sparsity/conversion.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index 7640628dd25..5a3d30b390b 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -509,7 +509,13 @@ def export_sparse_attention_config( tsf["coefficients"] = phases.get("prefill") or next(iter(phases.values())) ratio = skip_group.get("target_sparsity") if isinstance(ratio, dict): - skip_group["target_sparsity"] = ratio.get("prefill", next(iter(ratio.values()))) + 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