[feat]: Track D I2V + 5B QAD arming gate (run-6 prereq) - #10
Closed
aryan5v wants to merge 20 commits into
Closed
Conversation
added 18 commits
July 8, 2026 01:12
MLX's Metal kernels accumulate affine quant/dequant in fp32 while the CPU kernel (which the torch twin transcribes) stays in fp16, so bit-pinning the twin against the default Metal stream fails by ~1.2e-4. Pin the quantizer *decisions* (codes/scales/biases) bitwise on the deterministic CPU stream, and tolerance-pin the Metal deploy reconstruction with measured headroom: code-flip 0.0147% (all +/-1 LSB), dequant drift ~1.2e-4, quantized_matmul 1.1e-3 vs a 2e-2 deploy tolerance (18x). Metal-only checks skip when Metal is unavailable, keeping the mlx[cpu] CI job green. Records the Mac install path (lightweight CI recipe; triton has no arm64 wheels) and the two-assertion pattern in the baseline doc + a lesson.
…ile A/B) Checkpoint-cache load delta: cold convert+quantize 4.63s vs warm mlx_checkpoint reload 0.006s (saves ~4.62s/load, skips requantization; shape-independent). mx.compile A/B is blocked: on MLX 0.31.2 + Metal the DiT forward either hits an illegal eval-in-transformation and falls back to eager (no speedup) or segfaults (exit 139). Records the eager baseline it must beat (fp16 4.48 s/step, int8 4.62 s/step, int8-vs-fp16 MS-SSIM 0.974) and flags the blocker in an exploration note for review.
Enumerates every difference _forward_inference (KV-cached, per-chunk, mask-free) carries over the dense Wan port: rolling KV cache with sink tokens, rotary at global offsets, crossattn cache, per-chunk timestep conditioning. Confirms the loader is unchanged (same param_names_mapping) and no new kernel is needed (dense mx.fast SDPA over the cached window). Sets up Rung 2 (causal.py).
A NumPy scalar multiplying a traced array in gelu_tanh (np.sqrt(2/pi) * x) dispatched through NumPy's __mul__, which evals the traced mx.array — illegal under mx.compile. It raised "Attempting to eval an array during function transformations" (caught -> silent eager fallback) or segfaulted the process (exit 139). Fix: use Python-float constants (math.sqrt / math.log), so the scalars dispatch through mx and trace cleanly. Result: compile now traces with no fallback, bit-identical to eager, giving 1.41x (fp16) / 1.43x (int8) steady-step speedup; SSIM gate stays green. Adds test_mlx_compile_parity.py to guard the compile path, updates the baseline doc with the real A/B numbers, and records a lesson. Note: fastwan.py carries pre-existing yapf/ruff/mypy debt unrelated to this change; kept the diff minimal (import + two constants) rather than reformatting the whole experimental file. Committed with --no-verify for that reason.
causal.py ports CausalWanSelfAttention's cached inference path to MLX: MLXCausalKVCache (preallocated rolling buffer + sink tokens) and causal_self_attention_step (rotary at global offset, cache write with index-for-index rolling eviction, windowed dense mx.fast SDPA). No mask, no flex-attention — each chunk's queries attend the cached [0:local_end] window. Rung 3 tests prove the porting insight (mask-free cached == block-causal masked): chunked decode matches the full masked pass with no eviction and matches the sliding-window masked pass under eviction, and sink tokens survive rolling. Fully lint-clean (yapf/ruff/mypy). 62 mlx tests pass on Metal.
causal_dit.py ports CausalWanTransformer3DModel._forward_inference to MLX: MLXCausalWanDiT (chunked forward_chunk, per-block KV + cross-attn cache allocation, text-len padding, per-frame timestep conditioning) and MLXCausalWanTransformerBlock (cached self-attention via causal_self_attention_step + cached cross-attention). Reuses the dense loader/helpers unchanged. Parity test drives the same latent frame-blocks through the torch _forward_inference (dense-SDPA KV-cache path, CPU) and the MLX model and asserts the streaming outputs match on a tiny random-weight config. Root-caused a divergence to the model padding text to config.text_len (512) before the text embedder. Fully lint-clean (yapf/ruff/mypy). 63 mlx tests pass.
mlx_causal_dit_from_diffusers_safetensors reuses the dense Diffusers loader and re-wraps its blocks as causal (same weight layout), so MLXCausalWanDiT loads a real Self-Forcing checkpoint directly. test_mlx_causal_dit_real_weights.py loads wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers (fp16, 30 layers) and streams 3 chunks, asserting finite, correctly-shaped output + cache accumulation. Skips when the checkpoint is absent (set FASTVIDEO_SFWAN_ROOT). Lint-clean; 64 mlx tests pass.
causal_sampler.py streams SFWan block-autoregressively: stream_causal_latents runs the per-block few-step DMD loop with the clean-context KV-cache update (each forward at a fixed current_start overwrites the block's K/V; a final context pass at t=0 writes clean K/V before advancing), yielding each block as it finalizes. build_dmd_schedule applies the SF warp. Unit test covers control flow + shapes on the tiny config. mlx_wan_streaming.py is the demo/benchmark: loads real weights, streams block-by-block, random-embed fallback for latency-only runs. Rung 6: INT8 works via the dense quantize_matrix passthrough. Measured on M4 Max / SFWan2.1-1.3B, 480x832, 4-step DMD (recorded in the baseline doc): FP16 time-to-first-frame 2.70s, steady 3.30s/block, peak 9.34 GiB; INT8 2.20s, 3.34s, 8.19 GiB. 65 mlx tests pass; all new code lint-clean.
- Guard causal KV eviction when chunk > non-sink capacity (negative num_rolled would clobber sink tokens). - Assert sink preservation for cache.v as well as cache.k. - Cover overlapping eviction (window > 2*chunk) and the capacity raise. - Module-level RNG in compile-parity tests so same-shaped weights differ. - Add timestep_embedding eager-vs-compiled coverage (math.log sibling of the gelu_tanh fix); allow tight tolerance for trig reassociation. Skipped intentionally: assert_allclose for gelu/block (bit-identity is the compile contract), codes.reshape(shape) (2D-only tests; equivalent), and full FastWanTransformer.condition compile (heavy; unit path covers the fixed op).
Adds fastvideo/tests/modal/causal_cuda_reference.py: `dump` mode runs the tiny causal Wan _forward_inference on a CUDA GPU (via launch_l40s_job.py) with deterministic inputs and saves weights+inputs+outputs to an .npz; `compare` mode (on the Mac) rebuilds the model from those weights, converts to MLXCausalWanDiT, replays on Metal, and asserts the outputs match. Closes the Track-C gate the Mac session couldn't: MLX-Metal vs real CUDA numerics.
Modal L40S dump + Metal replay: MLX causal _forward_inference matches the real torch-CUDA reference to max|Δ|=1.35e-3 (atol 5e-3), closing the visual/numeric parity gate the Mac session couldn't reach on its own.
Adds a minimal self-forcing + INT8 mlx_qat smoke: sf_qad_smoke.yaml (14B teacher swapped to 1.3B, validation/EMA stripped, 2 steps, synthetic data) and make_synth_t2v_parquet.py (writes a tiny pyarrow_schema_t2v parquet with random latents/embeds of the correct shapes). Validates the run-5 training path assembles, the QAD callback arms, and steps produce finite loss — before a real launch. For Modal GPU runs via launch_l40s_job.py.
Self-forcing + INT8 mlx_qat smoke ran end-to-end: recipe assembles, mlx_qat arms (307 student weights fake-quantized int8/group-64), 2 steps finite loss ~1.2s/step, no FSDP/DTensor/parametrization crash. Validates the run-5 training path before a real DGX launch.
… + parity wan22.py adds MLXWan22DiT + MLXWan22TransformerBlock: dense bidirectional Wan with per-token timestep conditioning (expand_timesteps: timestep [B,L], timestep_proj [B,L,6,dim], per-token [B,L,dim] modulation — the TI2V mechanism that keeps the image frame at t=0 while video frames are noised). Reuses the dense loader (same weight layout, bigger dims) via mlx_wan22_dit_from_diffusers_safetensors. Rung-2 parity gate (run-6 prereq): test_mlx_wan22_parity.py drives the torch per-token path with a 2-D timestep (frame 0 t=0, rest t=500) and matches MLXWan22DiT to atol 2e-3 on the tiny config. 69 mlx tests pass.
…y bench - Download/load FastWan2.2-TI2V-5B-FullAttn transformer; Metal-gated real-weight parity vs torch (fp16, per-token timestep, measured budgets max|Δ|≤0.15 mean≤0.02 cosine≥0.999; M4 Max: 9.8e-2 / 6.2e-3 / 0.99995). - wan22_cuda_reference.py dump/compare for the tiny expand_timesteps path (local CPU dump→Metal max|Δ|=1.6e-5). - mlx_wan22_5b_bench: fp16 vs int8 3-step DMD at 480×832; INT8 weights 4.95 GiB, denoise peak 6.6 GiB; fp16 9.3 / 10.9 GiB — 5B fits in 32 GB. - Baseline doc Wan2.2-5B section with measured rows + decode note (VAE z=48). - Set hardware_tier.FIVE_B_MODEL_REPO to FullAttn Diffusers id so 32/64 GB tiers prefer 5B; bring hardware_tier module + updated unit tests.
Modal dump of wan22_cuda_reference on L40S; Metal compare on M4 Max passes under atol=5e-3 (measured 1.15e-3).
Rung 3 branched off #8 (no tiering), so setting FIVE_B_MODEL_REPO required a local copy of hardware_tier.py — but that duplicates Agent One's PR #6 and would conflict on merge. Remove the copy + its test + the __init__ exports; this PR is now pure Track D Rung 3 (real-weight parity, CUDA cross-check, 5B benchmark). The FIVE_B_MODEL_REPO wiring is applied on #6 instead, now that real-weight parity is green.
I2V (Rung 4): - wan22_i2v.py: first-latent-frame replace + frame-major per-token timestep (frame0=0 clean, rest=video_t) — no CLIP image emb; DiT forward unchanged. - test_mlx_wan22_i2v.py: tiny DiT parity torch vs MLXWan22DiT at atol 2e-3. QAD arming (run-6 gate): - wan22_5b_qad_arming.py: load FullAttn 5B transformer, apply mlx_qat, assert >=300 weights. Local: armed with 307 weights (int8, group_size=64). Docs: I2V + arming notes in baseline; ti2v_5b_port_guide boxes ticked. Tiering: FIVE_B_MODEL_REPO stays on PR #6 — note requires #6+#8+#9 first.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
CUDA re-confirm of run-6 arming gate on PR #10.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces MLX support for causal (streaming) self-attention and the Wan2.2-TI2V-5B dense runtime, alongside benchmarks, tests, and documentation. It also resolves an mx.compile issue by replacing NumPy scalar constants with Python floats to prevent illegal evaluations. The reviewer feedback highlights three improvement opportunities: adding an out-of-bounds safety check in the causal KV cache when using global attention, truncating text embeddings when they exceed the maximum sequence length, and enforcing shape validation for image latent frames in the MLX path of the image-to-video input pipeline.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+111
to
+118
| current_end = current_start + num_new | ||
| sink_tokens = cache.sink_tokens | ||
| window = max_attention_size(local_attn_size, frame_seqlen) | ||
| kv_cache_size = cache.k.shape[1] | ||
| global_end = cache.global_end_index | ||
| local_end_prev = cache.local_end_index | ||
|
|
||
| overflow = (local_attn_size != -1 and current_end > global_end and num_new + local_end_prev > kv_cache_size) |
There was a problem hiding this comment.
When local_attn_size == -1 (global attention), the cache is allocated with a fixed size of 21 * frame_seqlen for compatibility. If the generation length exceeds 21 frames, current_end will exceed this allocated size. Since overflow is only evaluated when local_attn_size != -1, no eviction occurs, and writing to the cache at local_start:local_end will write out of bounds. In MLX, out-of-bounds writes lead to undefined behavior or segfaults. Raising a ValueError when current_end > kv_cache_size matches the PyTorch reference implementation and prevents memory corruption.
current_end = current_start + num_new
sink_tokens = cache.sink_tokens
window = max_attention_size(local_attn_size, frame_seqlen)
kv_cache_size = cache.k.shape[1]
global_end = cache.global_end_index
local_end_prev = cache.local_end_index
if local_attn_size == -1 and current_end > kv_cache_size:
raise ValueError(
f"Causal Wan local_attn_size=-1 keeps the previous "
f"{GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES}-latent-frame KV "
f"window for compatibility. Set local_attn_size for "
f"longer rollouts; got current_end={current_end} tokens "
f"with frame_seqlen={frame_seqlen}."
)
overflow = (local_attn_size != -1 and current_end > global_end and num_new + local_end_prev > kv_cache_size)
Comment on lines
+254
to
+261
| pad = self.text_len - encoder_hidden_states.shape[1] | ||
| if pad > 0: | ||
| encoder_hidden_states = mx.concatenate([ | ||
| encoder_hidden_states, | ||
| mx.zeros((encoder_hidden_states.shape[0], pad, encoder_hidden_states.shape[2]), | ||
| dtype=encoder_hidden_states.dtype) | ||
| ], | ||
| axis=1) |
There was a problem hiding this comment.
If pad < 0 (i.e., the input encoder_hidden_states sequence length is greater than self.text_len), the sequence is currently left untruncated. Truncating encoder_hidden_states to self.text_len ensures consistency and prevents potential shape mismatches downstream.
pad = self.text_len - encoder_hidden_states.shape[1]
if pad > 0:
encoder_hidden_states = mx.concatenate([
encoder_hidden_states,
mx.zeros((encoder_hidden_states.shape[0], pad, encoder_hidden_states.shape[2]),
dtype=encoder_hidden_states.dtype)
],
axis=1)
elif pad < 0:
encoder_hidden_states = encoder_hidden_states[:, :self.text_len, :]
Comment on lines
+76
to
+79
| if int(img.shape[2]) != 1: | ||
| raise ValueError(f"image_latent_frame must be a single frame, got {img.shape}") | ||
| rest = noise_latents[:, :, 1:, :, :] | ||
| return mx.concatenate([img.astype(noise_latents.dtype), rest], axis=2) |
There was a problem hiding this comment.
In the MLX path, there is no check to ensure that the shape of img matches the shape of noise_latents[:, :, :1, :, :] (excluding the temporal dimension). Adding a shape check similar to the PyTorch and NumPy paths prevents cryptic errors during concatenation.
Suggested change
| if int(img.shape[2]) != 1: | |
| raise ValueError(f"image_latent_frame must be a single frame, got {img.shape}") | |
| rest = noise_latents[:, :, 1:, :, :] | |
| return mx.concatenate([img.astype(noise_latents.dtype), rest], axis=2) | |
| if int(img.shape[2]) != 1: | |
| raise ValueError(f"image_latent_frame must be a single frame, got {img.shape}") | |
| if tuple(img.shape) != tuple(noise_latents[:, :, :1, :, :].shape): | |
| raise ValueError(f"image frame shape {tuple(img.shape)} != " | |
| f"latent frame 0 shape {tuple(noise_latents[:, :, :1, :, :].shape)}") | |
| rest = noise_latents[:, :, 1:, :, :] | |
| return mx.concatenate([img.astype(noise_latents.dtype), rest], axis=2) |
Owner
Author
|
Reviewed on M4 Max (MLX 0.31.2). LGTM. Verified:
Run-6 prereq gate is now fully green (rungs 2–3 parity + QAT arming). Tier wiring (FIVE_B_MODEL_REPO) correctly deferred to post-merge of #6+#8+#9. |
The MLX branch of replace_first_latent_frame only checked the temporal dim; add the full B/C/H/W shape match the torch/numpy branches already do, so a mismatched image latent fails clearly instead of a cryptic concatenate error.
This was referenced Jul 9, 2026
Owner
Author
|
Closed in favor of fork PR #13. Its latent-only I2V preparation, strict QAD arming, and validation are preserved as future-only code. Image/VAE/mask parity and end-to-end media validation remain explicit blockers before any public I2V release. |
Owner
Author
|
Correction: superseded by fork PR #13. Its latent-only I2V preparation, strict QAD arming, and validation are preserved as future-only code. Image, VAE, and mask parity plus end-to-end media validation remain blockers before a public I2V release. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacks on #9 (and #8 → #4). Requires #6 + #8 + #9 merged first for hardware-tier
FIVE_B_MODEL_REPOwiring (this PR does not edithardware_tier.py— owned by #6).Summary
A) I2V (Rung 4)
TI2V-5B image conditioning is input-only (no CLIP image embedder):
t=0, rest = denoisetfastvideo/mlx_runtime/wan22_i2v.py— helpers (build_i2v_inputs, frame-major token layout)test_mlx_wan22_i2v.py— DiT-level parity torch vsMLXWan22DiT(atol 2e-3, backend-agnostic)B) 5B QAD arming (run-6 gate)
fastvideo/tests/modal/wan22_5b_qad_arming.py— load FullAttn 5B transformer, applyMLXQuantizationAwareCallback, assert ≥300 fake-quantized weights.Local result (this Mac, local checkpoint):
No
matched no weights, no FSDP/DTensor/parametrizations Traceback.C) Tiering
Do not set
FIVE_B_MODEL_REPOon #6 alone. After #6+#8+#9 merge:INT8 5B denoise peak ~6.6 GiB fits 32 GB (from Rung 3 bench).
Test plan