[diffusion] Fix native FP8 format handling for FLUX 3 rowwise linears

[Tokha233]

Motivation

The MI300 diffusion unit run on #42121 fails the FLUX 3 FP8r checkpoint test with Expected b.dtype() == at::kFloat8_e4m3fnuz, got Float8_e4m3fn (job). The rowwise linear currently passes serialized E4M3FN weights and activations directly to _scaled_mm, although gfx94 requires FNUZ. ModelOpt already normalizes that format, but two of its regression tests still assume FN and unadjusted scales.

Modifications

• Normalize FLUX 3 FP8r weights at load time on FNUZ platforms and match activation format before _scaled_mm, reusing the existing FN-to-FNUZ helper. Both scales are adjusted; negative zero is normalized to avoid NaN. Clone checkpoint weights before the helper's in-place byte normalization.
• Preserve the CUDA FN path and the serialized FP8r format.
• Cover both native formats, padded zero rows, tensor/tuple outputs, exact dequantized values and checkpoint immutability. Update ModelOpt tests to verify native dtype/scales and dequantized equivalence, including negative zero and CPU offload.

Accuracy Tests

Tested on an isolated RTX 5090 D v2 host, Python 3.12 / PyTorch 2.13.0+cu130, based on upstream f6fcda8.

• New FLUX format regressions against unchanged production code: 2 failed / 2 passed; only the FNUZ contract cases fail.
• Candidate: test_flux3_action.py, test_modelopt_fp8_layerwise_offload_load.py, test_transformer_quant.py: 107 passed, 24 subtests passed. Includes the real CUDA FP8 _scaled_mm checkpoint forward.
• FNUZ contract tests use CPU tensors and a checked/dequantized GEMM substitute; they are not MI300 GPU execution. ModelOpt tests parameterize format conversion and check exact dequantized values.

python -m pytest -q \
  python/sglang/multimodal_gen/test/unit/test_flux3_action.py \
  python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py \
  python/sglang/multimodal_gen/test/unit/test_transformer_quant.py

Please authorize AMD CI to verify the actual MI300 GEMM/offload path. This fixes the concrete format mismatch; it does not yet establish that the separate full FLUX Action non-finite-output failure is resolved. No thresholds or tolerances were relaxed.

Speed Tests and Profiling

Compatibility fix; no throughput claim. CUDA uses its existing FP8 path. The FNUZ path adds format normalization to make the previously invalid GEMM operands usable.

Checklist

• Changed-file pre-commit checks pass.
• Add unit tests.
• Accuracy results and platform limits documented above.
• Follow SGLang code style.
• AMD hardware CI (requires maintainer authorization).

---

CI States

Latest PR Test (Base): ✅ Run #37052849608
Latest PR Test (Extra): ❌ Run #37052849224
Latest PR Test (AMD ROCm 10): ❌ Run #37052849660

[mickqian]

/tag-and-rerun-ci

[Tokha233]

Native AMD validation is now available for this exact head (12bf2a8). In MI300 unit job 111112847081:

• test_fp8r_checkpoint_loads_fused_rowwise_linears passes with the actual MI300 _scaled_mm, replacing the earlier FN/FNUZ operand-type failure.
• All four new native-format contract cases pass.
• The ModelOpt serialized-checkpoint CPU-offload and packed-shard-scale tests pass, including their subtests. The previously identified four FP8 failure entries are resolved.
• The full unit sweep finishes with 23 failed / 3519 passed / 92 skipped / 5 deselected / 3 xfailed / 351 passing subtests. The remaining 23 failures are the 16 performance-policy assertions addressed separately in test(diffusion): cover CUDA and HIP latency guard policies #42250 and seven Qwen LoRA/fusion failures investigated in test(diffusion): cover residual dispatch and Qwen CUDA fallbacks #42256/[diffusion] Preserve eager base dispatch for disabled and merged Linear LoRA #42257. Those independent patches are not included in this branch.

Main CUDA CI and XPU diffusion CI pass. AMD 1-GPU partitions 2 and 3 also pass. This is still not an all-green AMD run: the separate FLUX Action full request continues to contain non-finite values, JoyEcho output fails, H3 encounters its existing HIP allowlist/media admission errors, and LTX-2.5 audio VAE access returns HTTP 403. NPU again fails installing the Triton Ascend wheel (403).

The targeted FP8 fix is now supported by real AMD execution as well as CUDA; no new change or relaxed tolerance was needed. Could the independent policy/fallback PRs be authorized for CI next, and the remaining platform/artifact failures be assessed separately for merge eligibility?