[feat]: Hardware-adaptive model tiering (auto-select quant+caps by unified memory) by aryan5v · Pull Request #6 · aryan5v/FastVideo

aryan5v · 2026-07-09T05:09:16Z

Stacks on PR #4 — review/merge #4 first.

Summary

North-star feature: automatically pick the best MLX quantization + memory caps for the Mac in use.

New fastvideo/mlx_runtime/hardware_tier.py

detect_unified_memory_gib() — macOS sysctl hw.memsize → MLX device_info → Linux /proc/meminfo → safe 16 GiB default
recommend_tier(memory_gib=None, *, prefer_5b=True) — pure, injectable

Tiers (constants, easy to retune):

Memory	Tier	Today	Cap
≤18 GiB	small	1.3B int8 + TAEHV	12
≤40 GiB	medium	1.3B fp16 + TAEHV (5B int8 when Track D sets `FIVE_B_MODEL_REPO`)	24
>40 GiB	large	1.3B fp16 + Wan-VAE (5B fp16 when Track D lands)	48

Benchmark: --auto-tier / --prefer-5b / --no-prefer-5b; mac-32gb and mac-64gb presets now carry MLX caps + MPS watermarks (additive; mac-16gb unchanged)
Tests: test_hardware_tier.py — injected sizes 8…128 GiB, non-Mac fallback, detect on this host, 5B path via injected repo id
Docs: “Hardware tiers” section in apple_silicon_benchmark_baseline.md

On this M4 Max box (36 GiB): recommend_tier() → medium, 1.3B fp16, 24 GiB MLX cap, mac-32gb.

Test plan

export PYTHONPATH=$PWD FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA TOKENIZERS_PARALLELISM=false
export MASTER_ADDR=localhost MASTER_PORT=29513
python -m pytest fastvideo/tests/mlx/test_hardware_tier.py -q
# 20 passed
python -m pytest fastvideo/tests/mlx/ -q
# 88 passed
pre-commit run --files fastvideo/mlx_runtime/hardware_tier.py \
  fastvideo/mlx_runtime/__init__.py fastvideo/benchmarks/mlx_fastwan_bench.py \
  docs/design/apple_silicon_benchmark_baseline.md
# all hooks passed

Out of scope (other agents)

Does not touch fastwan.py, causal.py, causal_dit.py, causal_sampler.py, tiny_wan.py, or mlx_wan_streaming.py.

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.

Add fastvideo/mlx_runtime/hardware_tier.py to auto-select model repo, quantization, decoder, and MLX allocator cap from detected unified memory (sysctl hw.memsize → MLX device_info → /proc/meminfo → safe 16 GiB default). Tiers (thresholds as constants): ≤18 GiB small INT8+TAEHV/12 GiB cap; ≤40 GiB medium (1.3B fp16 until Track D sets FIVE_B_MODEL_REPO); >40 GiB large fp16+Wan-VAE/48 GiB cap. prefer_5b flag reserved for 5B. Wire --auto-tier / --prefer-5b into mlx_fastwan_bench, fill mac-32gb and mac-64gb presets with memory caps matching the tier table, export the public API from mlx_runtime, document under Hardware tiers in the baseline doc, and cover with backend-agnostic injected-memory unit tests.

coderabbitai · 2026-07-09T05:09:23Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 79daa58b-d756-4ae5-b5fd-440d8e09864a

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch aryan/mac-hardware-tiering

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

gemini-code-assist · 2026-07-09T05:10:52Z

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

aryan5v · 2026-07-09T05:23:02Z

Reviewed on M4 Max (MLX 0.31.2). LGTM — ready to merge (stacked on #4; merge #4 first). Can't formally approve (same-account fork), so noting here.

Verified:

Scope clean: no protected files touched.
pytest fastvideo/tests/mlx/test_hardware_tier.py -q → 20 passed; full mlx suite → 88 passed (no regression).
Lint clean: yapf/ruff/mypy/codespell all pass on the new module + init + benchmark. (Benchmark's pre-existing mypy debt is unrelated.)
Detection on this Mac: 36 GiB → medium tier → 1.3B fp16, cap 24 GiB. Correct.
Additive benchmark changes (--auto-tier/--prefer-5b; existing presets/flags preserved).

Forward-dep (already handled well): FIVE_B_MODEL_REPO=None + TODO — once Track D's 5B port is parity-green I'll set it and the medium/large tiers auto-upgrade to 5B.

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.

aryan5v · 2026-07-09T22:07:03Z

Closed in favor of fork PR #13. Its hardware-tiering implementation and tests are preserved there as future-only code; automatic 5B/16 GB selection remains disabled until the pinned-artifact and clean-machine gates pass. See the preservation map in .

aryan5v · 2026-07-09T22:07:38Z

Correction: superseded by fork PR #13. Its hardware-tiering implementation and tests are preserved there as future-only code. Automatic 5B and 16 GB selection remains disabled until pinned-artifact and clean-machine gates pass. See tests/local_tests/wan2_2_ti2v_5b/README.md on #13.

Close the fused int8xint8 Metal speed gate with the friend M5 24 GB table (correct, 0.01-0.85x fp16). Canonical design receipt, next-wins #6 CLOSED, and a ready-to-ship negative-result section in the FastWan INT8 launch draft.

Aryan Kumar added 14 commits July 8, 2026 01:12

[docs]: record CUDA cross-check result for the MLX causal port

fd15f02

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.

This was referenced Jul 9, 2026

[feat]: Track D — Wan2.2-TI2V-5B MLX port (per-token timestep, Rung 2) #8

Closed

[feat]: Track D Rung 3 — Wan2.2-5B real-weight parity + memory/latency #9

Closed

This was referenced Jul 9, 2026

[feat]: Track D I2V + 5B QAD arming gate (run-6 prereq) #10

Closed

[feat]: preserve Apple MLX future 5B, I2V, and self-forcing tracks #13

Draft

aryan5v closed this Jul 9, 2026

aryan5v deleted the aryan/mac-hardware-tiering branch July 9, 2026 22:08

Navigation Menu

[feat]: Hardware-adaptive model tiering (auto-select quant+caps by unified memory) - #6

[feat]: Hardware-adaptive model tiering (auto-select quant+caps by unified memory)#6
aryan5v wants to merge 14 commits into
aryan/apple-silicon-fastwan-mlxfrom
aryan/mac-hardware-tiering

aryan5v commented Jul 9, 2026

Uh oh!

coderabbitai Bot commented Jul 9, 2026

Review skipped

Uh oh!

gemini-code-assist Bot commented Jul 9, 2026

Uh oh!

aryan5v commented Jul 9, 2026

Uh oh!

aryan5v commented Jul 9, 2026

Uh oh!

aryan5v commented Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

aryan5v commented Jul 9, 2026

Summary

Test plan

Out of scope (other agents)

Uh oh!

coderabbitai Bot commented Jul 9, 2026

Review skipped

Uh oh!

gemini-code-assist Bot commented Jul 9, 2026

Uh oh!

aryan5v commented Jul 9, 2026

Uh oh!

aryan5v commented Jul 9, 2026

Uh oh!

aryan5v commented Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant