saaivignesh20/MAGI-2-preview-FP8

MAGI-2 Preview, FP8 routed experts

The preview transformer of sand-ai/MAGI-2-preview (revision 2dea51b6) with its routed-expert weights (*.moe_mlp.W_gate / W_up / W_down, about 95% of the parameters) stored as FP8 E4M3 with one fp32 scale per expert per output channel. Every other tensor is unchanged. preview/ here drops in for SandAI's ckpt/preview (212 GiB -> 112 GiB). This repo also mirrors SandAI's other components unchanged (text_encoder/, vae/, turbo_vae/, stable-audio-open-1.0/, refiner/, same revision), so downloading it gives the whole model in SandAI's ckpt/ layout.

Two expert kernels (MAGI2_FP8_MOE):

• w8a8 (default): the expert inputs are quantized to FP8 too (one scale per token per head; the SwiGLU7 output with a fixed scale of 8, from its clamp bound of 56), so both expert GEMMs run FP8 x FP8 on Hopper's tensor cores. The weights are transposed to K-major at load time, the layout those instructions read from shared memory.
• w8a16: weight-only; FP8 tiles widened to bf16 in the loop, SandAI's bf16 math. Less error, slower.

This is an unofficial conversion; it needs the code in magi2_fp8/ on top of SandAI's MAGI-2-preview code. Loading this checkpoint without the patch fails (missing *_scale keys under strict=True), never silently.

Use

git clone https://github.com/SandAI-org/MAGI-2-preview && cd MAGI-2-preview && git checkout f68a0f9
hf download saaivignesh20/MAGI-2-preview-FP8 --local-dir ckpt         # FP8 preview + SandAI's components, ~185 GB
cp -r ckpt/magi2_fp8 ./ && python3 -m magi2_fp8.configs --gpus 1 --fp8 "$PWD/ckpt/preview"
PYTHONPATH="$PWD" torchrun --nproc_per_node=1 -m magi2_fp8.entry --config configs/run_fp8_refiner.json \
  --resolution 1080p --seconds 10 --num-inference-steps 32 --prompt-file PROMPT.md --output out

Or keep it warm: torchrun --nproc_per_node=1 -m magi2_fp8.serve --config configs/run_fp8_refiner.json loads and compiles once and renders JSON jobs dropped into /root/queue (see magi2_fp8/serve.py). A ready-made RunPod image (ghcr.io/saaivignesh20/magi2-fp8) downloads this repo into the volume and starts that worker.

Prompts: the JSON caption from SandAI's prompt-enhancement templates rendered with inference/prompt_enhancement/pe_markdown.render_markdown (as enhancer.py does), under 7,000 tokens.

Hardware:

• Hopper: SandAI's own stack (their Docker image). With 2 GPUs (cp_size/ep_size 2) 540p needs more than 80 GB per GPU (2x H200 works; 2x H100 runs out of memory). 4x H100 with cp_size/ep_size 4 should fit (untested with FP8; the bf16 model runs there at about 66 GB per GPU).
• Blackwell (B300 tested, one GPU, cp_size/ep_size 1): SandAI's preview model calls FA3, which has no Blackwell build. magi2_fp8/blackwell_attn.py computes the same attention with PyTorch's cuDNN SDPA (which also returns the log-sum-exp) and applies SandAI's attention-sink correction; it installs itself only on sm_10x. Without SandAI's image, magi2_fp8/compat.py stands in for the two Hopper-only imports (flash_attn_interface, magi_attn_extensions) at load time. Tested runtime: torch 2.14.1+cu130, Triton 3.8.0, transformers 5.5.0, MagiCompiler 5950612, plus SandAI's requirements.txt. On one B300 everything fits at once: set MAGI2_TEXT_ENC_OFFLOAD_MODE=gpu MAGI2_PREVIEW_OFFLOAD_MODE=gpu (peak 199 GB of 275).

Files

• preview/: FP8 shards and model.safetensors.index.json (scale keys <weight>_scale).
• text_encoder/, vae/, turbo_vae/, stable-audio-open-1.0/, refiner/: SandAI's, unchanged (mirror).
• quantization_config.json: scheme, tensors, base revision.
• magi2_fp8/quantize.py: the converter (bf16 ckpt/preview -> this layout).
• magi2_fp8/fp8_moe.py: SandAI's multi-head MoE Triton kernel with FP8 weight tiles and scales.
• magi2_fp8/patch.py: FP8 expert parameters, scale sharding across expert-parallel ranks, the forward hook.
• magi2_fp8/entry.py: SandAI's entry.py with the patch applied.
• magi2_fp8/blackwell_attn.py, magi2_fp8/compat.py, magi2_fp8/stubs/: the Blackwell attention path (above).
• magi2_fp8/fa2_compat.py, magi2_fp8/refiner_attn.py, magi2_fp8/refiner_fix.py: the 1080p refiner without FlashAttention-2 / MagiAttention (FlexAttention with a block mask built from SandAI's window ranges) and with its outer forward uncompiled (torch 2.14 breaks its modality masks).
• magi2_fp8/serve.py, magi2_fp8/configs.py: warm worker and the run configs it uses.
• magi2_fp8/fast_sampler.py: opt-in MAGI2_CFG_STOP (guidance on the early steps only). Not recommended: SandAI's guidance pass carries a negative prompt, and dropping it from step 60% on brought in a broadcast logo and lost a line in our test.
• magi2_fp8/check_kernel.py: one-layer check of both kernels against SandAI's bf16 kernel (--sweep times W8A8 tile configs; per-GPU defaults H100 64,64,32,3,4, B300 128,128,32,3,8; MAGI2_FP8_W8A8_CONFIG overrides).
• magi2_fp8/check_attn.py: the Blackwell attention path against an exact fp32 reference with sinks.

Results (2026-10-19)

• Conversion: 108 expert tensors, 161 s on one H100; 112.4 GiB on disk (bf16: 212 GiB). Weight error (relative Frobenius) 2.62% mean, 2.64% max, uniform across layers.

• Kernel check, layer 2, 32,768 random tokens, one H100 (error = relative Frobenius vs SandAI's kernel on the bf16 weights):

  Expert kernel	ms per call	Layer output error
  SandAI bf16	17.7	—
  w8a16	28.3	3.0% (kernel vs dequantized weights 4.6e-3: bf16 rounding)
  w8a8 (tile config 64,64,32,3,4)	11.4	4.6%

  One activation scale per 64 values instead of per row measured 4.56% vs 4.62%, so row scales stay.

• Two 540p, 10 s clips at 50 steps on 2x H200 (cp 2, ep 2): peak 83.2 GB reserved per GPU, so 540p does not fit 2x H100 80GB (out of memory in the first denoise step). Same pod, same seeds, both kernels back to back:

  	Denoise, first clip (incl. compile)	Denoise, second clip
  w8a16	7 min 03 s	5 min 32 s
  w8a8	6 min 45 s	4 min 39 s (16% less)

  Text encoding on GPU (MAGI2_TEXT_ENC_OFFLOAD_MODE=roundtrip) 50-85 s per clip. Both kernels gave the same shots beat for beat at the same seed, and every scripted line came out verbatim and in order (Whisper large-v3-turbo) in all four clips (a photoreal one and a 2D cartoon one per kernel).

• One B300 (cp 1, ep 1, text encoder and preview resident on the GPU), same two clips and seeds:

  Setting	Denoise, second clip	Wall per clip	Lines verbatim
  50 steps	5 min 28 s	5 min 33 s	3/3, 3/3
  32 steps	3 min 30 s	3 min 35 s	3/3, 3/3
  32 steps, guidance on the first 60% (MAGI2_CFG_STOP=0.6)	2 min 53 s	2 min 59 s	3/3, 2/3 (and a logo)

  Kernels on B300 (layer 2, 32,768 tokens): SandAI bf16 with a 64-token tile 24.8 ms (its 128-token tile overflows Blackwell tensor memory in Triton 3.8), w8a8 11.1 ms, w8a16 52.7 ms. Attention check vs exact fp32 with sinks that move the output by 54%: 3.25e-3. Profile of one 50-step denoising step (6.37 s): attention 30%, experts 23%, router top-k 11%, SandAI's RMSNorm kernel 8%, bf16 GEMMs 7%, fp32 SIMT GEMMs 6%.

Status

Unofficial and lightly tested: the clips above. No side-by-side against SandAI's bf16 checkpoint at the same seed yet.

License

Apache-2.0, as the base model; LICENSE is SandAI's. Changes: expert weights quantized; the kernel in magi2_fp8/fp8_moe.py is derived from SandAI's inference/flash_mh_moe/triton/mh_moe_fwd.py.