Repository navigation
Emit VisualGen-compatible sparse_attention_config for diffusers skip-softmax export - #2868
Edwardf0t1 wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation |
|---|---|---|
| Description Check |  |
Check skipped - CodeRabbit’s high-level summary is enabled. |
| Title check | |
The title clearly and specifically describes the main change: generating VisualGen-compatible sparse attention configuration for diffusers skip-softmax exports. |
| Docstring Coverage | |
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 4 files. (1 skipped: 1 … |
| Linked Issues check | |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
| Security Anti-Patterns | |
PASS. The PR adds no security-sensitive loading, evaluation, execution, # nosec, or dependency patterns. The only changed package Python code adds a for_diffusion argument and transforms sparse-at… |
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 4 files. (1 skipped: 1 unsupported.)
Full details: Security Anti-Patterns
-
Explanation
PASS. The PR adds no security-sensitive loading, evaluation, execution,
# nosec, or dependency patterns. The only changed package Python code adds afor_diffusionargument and transforms sparse-attention dictionaries. No example Python files or dependency manifests changed.
✨ Finishing Touches
-
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2868 +/- ##
==========================================
+ Coverage 70.16% 78.91% +8.75%
==========================================
Files 659 659
Lines 74081 74092 +11
==========================================
+ Hits 51979 58473 +6494
+ Misses 22102 15619 -6483
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Nudge: the fix is small and looks correct, but nobody has run the tests or tried loading the checkpoint in VisualGen, so a human should check before merge.
Needs action:
- Confirm CI passes on
test_sparse_attention_conversion.pyandtest_wan22_calibrated_export. The PR says pytest was never run. - Confirm with the TRT-LLM VisualGen owners (TRT-LLM PR #14687) that a Wan2.2 checkpoint with
coefficientsand a scalartarget_sparsityloads. - Fix the
next(iter(ratio.values()))default inconversion.py. Python evaluates it even whenprefillis present, and it raisesStopIterationon an empty dict (see inline comment).
No action needed:
- The design review fired only because the PR touches 6 directories. The change is just a keyword-only
for_diffusionflag on the existing exporter, not a new system. The LLM path and the vLLM path throughexport_threshold_scale_factorare unchanged. - The edit to the existing test in
tests/examples/diffusers/sparsity/test_sparsity.pyis justified. Its assertions now checkcoefficientsinstead ofprefill, matching the schema change, and it gains a check thattarget_sparsityis a scalar. - The CodeRabbit summary in the PR body lists unrelated features (GGML, ONNX Runtime). It looks stale, not like an injection attempt.
| 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()))) |
There was a problem hiding this comment.
Bot comment.
Python evaluates the dict.get default eagerly, so next(iter(ratio.values())) runs even when prefill exists, and an empty dict raises StopIteration. Suggest ratio["prefill"] if "prefill" in ratio else next(iter(ratio.values()), None). If the result is None, drop the key, since VisualGen needs a scalar.
There was a problem hiding this comment.
Good catch, verified: the eager default evaluated next(iter(...)) even when prefill was present and would raise on an empty dict. Fixed in 98511f6: uses ratio["prefill"] when present, else the first value, and drops target_sparsity if there is none.
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Nudge: the eager-default bug is fixed and the code looks correct, but the tests still haven't been run and nobody has confirmed the checkpoint loads in VisualGen.
Needs action:
- Confirm CI passes on
test_sparse_attention_conversion.pyandtest_wan22_calibrated_export. The PR says pytest was never run locally. - Confirm with the TRT-LLM VisualGen owners (TRT-LLM PR #14687) that a Wan2.2 checkpoint with
threshold_scale_factor.coefficientsand a scalartarget_sparsityloads.
No action needed:
- ✔️ Resolved since the last review: the eagerly evaluated
next(iter(ratio.values()))default inconversion.py. It now usesratio["prefill"]when present, falls back to the first value, and dropstarget_sparsitywhen the dict is empty. - Design review: the PR only adds a keyword-only
for_diffusionflag to the existing exporter, not a new system. The other callers inplugins/sparse_attn_config.pyandplugins/sparse_attn_calibration.pyuse the default and are unchanged. - The edit to the existing test in
tests/examples/diffusers/sparsity/test_sparsity.pyis justified. It now checkscoefficientsinstead ofprefill, which matches the schema change. - The CodeRabbit summary in the PR body lists unrelated features. It looks stale, not like an injection attempt.
…softmax export Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
98511f6 to
cb0cc62
Compare
This branch has not been deployed
2 participants
What does this PR do?
Type of change: Bug fix
The diffusers skip-softmax HF export wrote the LLM phase-keyed schema (
threshold_scale_factor.prefilland a{phase: ratio}target_sparsity). TRT-LLM VisualGen (PR #14687) expectsthreshold_scale_factor.coefficientsand a scalartarget_sparsity, so the exported Wan2.2 checkpoint fails to load (float() ... not 'dict'). NVBug 6927804.export_sparse_attention_config(model, *, for_diffusion=False): withfor_diffusion=True, emit onecoefficientsblock and a scalartarget_sparsity. The LLM path and the vLLM calibration plugin are unchanged._export_diffusers_checkpointpassesfor_diffusion=True.resolve_for_target_sparsity.Based on #1816 (closed, approved). Deliberately dropped its
initial_disabled_steps->disabled_until_timesteprename:initial_disabled_stepsshipped in 0.46 and the vLLM calibration path now depends on it, so the rename is a separate, breaking change.Usage
Testing
test_sparse_attention_conversion.py(diffusion schema) and the Wan calibrated-export example test.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
NVBug 6927804; supersedes #1816.
🤖 Generated with Claude Code
Summary by CodeRabbit