Skip to content

[refactor] Simplify Wan sampling and tests - #1825

Open
SolitaryThinker wants to merge 3 commits into
hao-ai-lab:mainfrom
SolitaryThinker:refactor/wan-cleanup
Open

[refactor] Simplify Wan sampling and tests#1825
SolitaryThinker wants to merge 3 commits into
hao-ai-lab:mainfrom
SolitaryThinker:refactor/wan-cleanup

Conversation

@SolitaryThinker

@SolitaryThinker SolitaryThinker commented Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Purpose

Keep Wan sampling code with the Wan pipelines. Make the first numerical check small and cheap. Continue #1823 and #1824 in one larger PR.

Changes

  • Move the causal transformer into fastvideo/models/wan/. Keep aliases at the old paths and update internal imports.
  • Put dense Wan, DMD, and causal sampling under fastvideo/pipelines/basic/wan/stages/.
  • Keep shared CFG and scheduler execution in the shared dense loop. Move Wan expert selection, input packing, and timestep rules into the Wan stage.
  • Encode the first frame before sampling. Preserve normalization and offload order. Standard and DMD causal stages share cache allocation, not sampling algorithms.
  • Make DMD scheduler ownership explicit. Its full training-noise table stays separate from the scheduler used for timestep preparation.
  • Share tiny test fixtures and collect the CFG tests. Remove network access from local DreamX config tests. Component parity tests fetch only their pinned component.
  • Add VAE, causal-cache, and three-step denoising goldens. Run component goldens before parity. Make selected integration lanes wait for the golden lane.
  • Add bash scripts/validate_wan.sh all and update contributor guidance.

Shared T5/UMT5 and CLIP stay shared. No model is removed. No training pipeline is migrated. Checkpoint keys and model math stay the same.

Test evidence

Current head: 275cf421beedf2fb538cd6dbe13a4714d6a49202.

  • All-golden canary 5081: 20 passed in 59.37 seconds on GB200, with exact tensor matches. This is pytest time, excluding container setup.
  • Fastcheck 1147: encoder, VAE, transformer, and kernel lanes passed. Component results: 41 encoder tests, 81 VAE tests, and 47 transformer tests passed. Wan's component goldens ran first.
  • Unit lane: 1269 passed, 8 skipped, 5 failed. All Wan sampling, wiring, import, and new isolation contracts passed. The five failures are the unchanged NVFP4 tests listed below.
  • Pre-commit and docs build passed.
  • Before the stage split, 14 sampling contracts passed against unchanged production code. They include the real scheduler's 50-step control flow.
  • Each new tensor reference was generated from unchanged main at 7f8316423355a2f07db533a5a339ac16e0da31ff in two separate processes. The candidate then matched both. The causal gate asserts the actual FlashAttention backend.
  • New references and provenance. No existing tensor or video reference was replaced.
  • Private runner policy: 70 driver tests passed. Bootstrap reject and cancel checks passed. The trusted uploader accepted the complete 20-lane graph.

The first full golden canary caught a test-order bug: the loader left a BF16 compute policy behind and changed the next block test's attention backend. The tensor-golden context now clears and restores that state. Four tests cover absent/existing state and normal/error exits. Canary 5081 verified the fix in the full test order, without changing references or tolerances.

CI blockers outside this diff

These need separate repairs. No unrelated H3 changes or weaker tolerances are included. macOS MLX and DreamVerse checks are still running.

Validation limits

This evidence is from GB200, not L40S. The 50-step tests use tiny fake denoisers and the real scheduler; they are not full video renders. No new default/full-quality SSIM render is claimed. Independent component parity and SSIM coverage remain.

First-frame VAE time now has its own stage timing instead of being counted inside DiT time. Total request latency still includes it. No performance baseline was reseeded.

The six Fastcheck lanes and all twenty full-suite lanes remain. Direct reruns keep their existing meaning when the conditional golden step is skipped. The private uploader accepts the old or complete golden-first graph during rollout, and rejects partial or allow-failure dependency changes.

@mergify mergify Bot added type: refactor Code restructure without behavior change scope: inference Inference pipeline, serving, CLI scope: infra CI, tests, Docker, build scope: docs Documentation scope: model Model architecture (DiTs, encoders, VAEs) labels Sep 6, 2026
@mergify

mergify Bot commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

Protection Waiting on
🔴 PR merge requirements 👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

  • #approved-reviews-by>=1
  • check-success=fastcheck-passed
  • check-success=full-suite-passed
This rule is failing.
  • #approved-reviews-by>=1
  • check-success=fastcheck-passed
  • check-success=full-suite-passed
  • check-success~=pre-commit
  • title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

scope: docs Documentation scope: inference Inference pipeline, serving, CLI scope: infra CI, tests, Docker, build scope: model Model architecture (DiTs, encoders, VAEs) type: refactor Code restructure without behavior change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant