Skip to content

[Perf] [AMD] diffusion: VDN-H3 and MiniMax-H3 ROCm MI355X compile path, AITER kernels, MXFP8 fixes - #3

Open
jackyhkust wants to merge 10 commits into
mainfrom
feat/vdn-h3-rocm-mi355x
Open

[Perf] [AMD] diffusion: VDN-H3 and MiniMax-H3 ROCm MI355X compile path, AITER kernels, MXFP8 fixes#3
jackyhkust wants to merge 10 commits into
mainfrom
feat/vdn-h3-rocm-mi355x

Conversation

@jackyhkust

@jackyhkust jackyhkust commented Oct 22, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Owner

Summary

VDN-H3 (OpenVDN/vdn-minimax-h3) and MiniMax-H3 changes from the sgld-vdn-compile branch, rebased onto current main as a linear series. Most changes are ROCm (MI355X, gfx950); the two MXFP8 commits fix Blackwell.

Commit Change
VDN-H3: allow torch.compile on ROCm for the compiled u8 sample ROCm compile path for the VDN-H3 DiT, delta-factor and linear-branch kernels
VDN-H3: use functional collectives in the Ulysses exchange Compile-friendly Ulysses all-to-all; restores usp._maybe_wait
VDN-H3: default the ROCm window softmax to AITER CK varlen on gfx950 hybrid_window_attn_h3 window on AITER CK on gfx950
MiniMax-H3 VAE: use AITER CK flash attention for the ViT on gfx950 Replaces the forced math SDPA in the video-VAE ViT; SGLANG_MINIMAX_H3_VAE_ROCM_ATTN=math is the exact control
MiniMax-H3 audio VAE: dilated fp32 Conv1d as unfold + matmul on ROCm MIOpen runs these with naive_conv (about 50x slower)
FP8: log per-layer FP8 decisions with SGLANG_LOG_FP8_LAYERS=1 Diagnostic only, off by default
MXFP8: honor --quantization-ignored-layers On Blackwell, fp8 is remapped to mxfp8, and the ignore list was dropped, so every linear was quantized
MXFP8: make the online linear compile-safe via custom ops The MXFP8 quantizer and GEMM as opaque custom ops; without this, --enable-torch-compile with MXFP8 fails
VDN-H3: count inference steps as DiT forwards Follows sgl-project#38671: the default is 8 (the distilled schedule, nine sigma grid points); 8 and the 50-step comparison sample are accepted

Two benchmark-only commits from the branch are not in this PR: the SGLANG_VDN_SPD_INVERSE=inv workaround for old rocSOLVER, and the opt-in per-tensor FP8 cuDNN path.

Rebase notes

Testing

  • On 8x MI355X, at the pre-rebase branch (5400d640e5 plus the MXFP8 commits), patched against unpatched: fox (124 frames) +0.0%, ex2 (345 frames) −0.3%, piano (345 frames) −0.2%. Video is bit-identical, and the FP8 layer plan and recompiles are unchanged.

  • B200: the MXFP8 commits were validated on 8x B200 (matched FP8 layer coverage, compile on, no recompiles in the timed request).

  • All 22 changed files parse in the run image.

  • MI355X rerun on this branch (80f81b089f): passed. 8x MI355X, compiled Ulysses 8, num_inference_steps=8. Canonical value is the last of 3 requests:

    Cell Pre-rebase (5400d640e5 + MXFP8) This branch Delta
    fox (124 frames) 4.171 s 4.102 s −1.7%
    ex2 (345 frames) 11.363 s 11.128 s −2.1%
    piano (345 frames) 11.345 s 11.099 s −2.2%
    • 8 DiT forwards per request in every cell, and the step time is about 2% lower. Decode plus the post-decode time is the same (568 against 569 ms on fox); [diffusion] Encode the MP4 while the VAE decodes (MiniMax-H3, Wan) sgl-project/sglang#41819 moves the MP4 encode into the decode stage.
    • The FP8 layer plan is the same. There are no recompiles in the timed request, and there are no MIOpen errors.
    • The output is valid. The scene and subject are the same, with small differences in pose and layout (video PSNR 20–24 dB against the pre-rebase output). The same-code floor is bit-identical, so changes in main give a different sampling path.
    • Before the last commit, the branch ran 9 forwards and was 8.5–8.8% slower (fox 4.537 s, ex2 12.346 s, piano 12.306 s).

CI States

Latest PR Test (Base): ❌ Run #39897757357
Latest PR Test (Extra): ❌ Run #39897757298
Latest PR Test (AMD ROCm 10): ❌ Run #39897757369

@jackyhkust
[diffusion] VDN-H3: allow torch.compile on ROCm for the compiled u8 s…
ff7a5cc 
…ample.

The gfx950 run needs the AITER window path. The linear branch stays on eager Cholesky, and the breakable CUDA graph stays off.
@jackyhkust
[diffusion] VDN-H3: use functional collectives in the Ulysses exchange
16094db 
The VDN attention module's Ulysses all-to-all and frame-sum all-reduce
used raw async_op=True collectives and Work.wait(). Dynamo cannot trace
a Work handle, so every call forced a hard graph break under
torch.compile.

Switch the three call sites (_vdn_a2a_rows_to_heads,
_vdn_a2a_heads_to_rows, and the frame_work all-reduce in
_vdn_ulysses_hybrid_core) to torch.distributed._functional_collectives
and resolve results with the existing _maybe_wait() helper from usp.py
at the point of use, keeping the launch-early/wait-late overlap.

Extend the exchange round-trip test with a 4-rank case (heads=8 so
local heads divide evenly).

MI355X, compiled, 1344x768x124f, 9 steps (last-clip stage sum):
  u2 t2va: 108.66 s -> 65.10 s
  u4 t2va: 32.12 s, u4 fl2va: 32.84 s
@jackyhkust
[diffusion] VDN-H3: default the ROCm window softmax to AITER CK varle…
80351a8 
…n on gfx950

The window passes, the dense leg, the full-cover dense call and the token
refiner ran through AITER Triton varlen on every ROCm arch, because the
grouped-varlen CK/ASM kernel hangs on MiniMax-H3 packed lengths on gfx942.
On gfx950 it does not hang and is faster.

Register aiter.flash_attn_varlen_func (how_v3_bf16_cvt=2 when supported) at
import as the opaque custom op sglang::vdn_aiter_ck_varlen, so torch.compile
treats it as one node, and dispatch _fa_varlen to it. The default is
aiter_ck when is_gfx95_supported(), else triton; SGLANG_VDN_WINDOW_BACKEND
overrides. The text-encoder dense path stays on Triton. If aiter cannot be
imported, the backend logs a warning and keeps Triton.

Kernel check (1 GPU): rel L2 vs Triton 6-8e-5, compiled == eager bit-exact
under fullgraph=True; 1.36-1.63x faster on large segments.

MI355X x8, compiled u8, t2va 768x1344, 9 steps, fp8, total_duration_ms:
  fox   124f: 4.564 s -> 4.227 s (-7.4%)
  ex2   345f: 13.071 s -> 11.450 s (-12.4%)
  piano 345f: 13.028 s -> 11.455 s (-12.1%)
@jackyhkust
[diffusion] MiniMax-H3 VAE: use AITER CK flash attention for the ViT …
ca959bf 
…on gfx950

On gfx950 the video VAE ViT attention was pinned to math SDPA. Fused ROCm
SDPA (AOTriton flash and efficient) returns intermittently wrong outputs at
the ViT shapes (head dim 64, large B x S) with HIP 7.16/7.17 and torch 2.13:
12-31% of trials, rel L2 up to 1e-1. Math is exact but upcasts to fp32 and
runs unfused.

Add SGLANG_MINIMAX_H3_VAE_ROCM_ATTN = aiter | math | sdpa (gfx95 only;
default aiter). aiter calls aiter.flash_attn_func on the [B, S, H, D] tensors
the decoder already holds; fp32 inputs and a missing aiter fall back to
math. sdpa keeps the fused path as an opt-in.

AITER stays at the fp64 reference and is run-to-run deterministic.
MI355X x8, decode vs math: -626 ms (124f), -1792 ms (345f).
End to end, total_duration_ms, math -> aiter:
  fox   124f: 4.842 s -> 4.139 s
  ex2   345f: 13.157 s -> 11.337 s
  piano 345f: 13.125 s -> 11.323 s
@jackyhkust
[diffusion] MiniMax-H3 audio VAE: run dilated fp32 Conv1d as unfold +…
6e588d0 
… matmul on ROCm

MIOpen runs the BigVGAN resblock dilated fp32 Conv1d with naive_conv on
gfx950: 3-9 ms per conv, about 50x the existing _conv1d_as_matmul path.
Route groups=1, dilation>1, fp32 Conv1d to unfold + matmul first on ROCm.
SGLANG_MINIMAX_H3_AUDIO_DILATED_CONV_MATMUL=0 restores the conv path.

MI355X: -248 ms audio decode per 124-frame clip.
@jackyhkust
[diffusion] FP8: log per-layer FP8 decisions with SGLANG_LOG_FP8_LAYE…
0bfba9f 
…RS=1

With SGLANG_LOG_FP8_LAYERS=1, Fp8Config.get_quant_method logs whether each
linear leaf name is quantized or skipped by --quantization-ignored-layers,
once per (name, decision). This verifies the FP8 scope at runtime, e.g. that
an ignore list leaves exactly qkv_proj, out_proj, to_out_linear, fc1 and fc2
in FP8 for VDN-H3. No effect when the variable is unset.
@jackyhkust
[diffusion] MXFP8: honor --quantization-ignored-layers
52fb063 
VDN-H3 remaps --quantization fp8 to mxfp8 on Blackwell (sm100) and sm120.
The loader only forwarded quantization_ignored_layers for fp8/mxfp4/kitchen_int8,
and MXFP8Config.get_quant_method had no skip check, so every linear layer was
quantized. Forward the ignore list for mxfp8, skip ignored prefixes with
is_layer_skipped, and log decisions under SGLANG_LOG_FP8_LAYERS like Fp8Config.
@jackyhkust
[diffusion] MXFP8: make the online linear compile-safe via custom ops
dbcaf2e 
can_use_mxfp8_swizzled() is False under torch.compile, so apply() hit the Triton
quantizer guard (k % 32 ValueError), and Inductor cannot lower _scaled_mm_v2 with
swizzled scales. Register the activation quantizer and the block-scaled GEMM as
opaque custom ops (sglang::mxfp8_quantize_act, sglang::mxfp8_gemm) with fake impls.
Eager kernels and numerics are unchanged.
@jackyhkust
[diffusion] VDN-H3: count inference steps as DiT forwards
80f81b0 
sgl-project#38671 changed MiniMax-H3 num_inference_steps from sigma grid points to
denoise transitions (DiT forwards), and set the VDN-H3 default to 8. The
VDN-H3 sampling params on this branch kept the old count (default 9,
accepts 9 and 50). With the new stages, a request for 9 runs 9 forwards
on a 10-point grid, which is off the distilled 8-NFE schedule and about
9% slower.

Use the upstream count again: the default is 8 (the distilled schedule,
nine sigma grid points), and 8 and the 50-step comparison sample are
accepted.
@jackyhkust
jackyhkust marked this pull request as ready for review October 22, 2026 23:48
@jackyhkust jackyhkust changed the title [diffusion] VDN-H3 / MiniMax-H3: ROCm MI355X compile path, AITER kernels, MXFP8 fixes [Perf] [AMD] diffusion: VDN-H3 and MiniMax-H3 ROCm MI355X compile path, AITER kernels, MXFP8 fixes Oct 22, 2026
@github-actions

This comment has been minimized.

Sign in to view
@jackyhkust
jackyhkust force-pushed the feat/vdn-h3-rocm-mi355x branch from 04967d1 to 80f81b0 Compare October 23, 2026 03:50
@github-actions

This comment has been minimized.

Sign in to view
@jackyhkust jackyhkust self-assigned this Oct 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@jackyhkust jackyhkust

Labels

amd diffusion jit-kernel

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jackyhkust