[AMD][Diffusion] Migrate FlyDSL fused norm kernels to the v0.3.0 stable API

[yctseng0211]

Motivation

Migrate FlyDSL fused norm kernels to the v0.3.0 stable API, this migration is required for upcoming AITER/FlyDSL versions: the current code depends on non-public FlyDSL internals that already broke once, and when they break again the kernels silently fall back to forward_native.

Rebase of #33144 onto the post-#35114 layout: the kernel now lives at norm/fused_residual_norm_flydsl.py and #34536's import fallback is dropped, since the capability guard supersedes it.

Co-authored-by: Phil Li haicli@amd.com

FlyDSL 0.3 API migration

v0.3.0 both removed flydsl.expr.buffer_ops (the upstream kernel carries an aiter fallback for exactly this) and introduced FlyDSL's stability contract, docs/api_stability.md, which is absent at v0.2.4. Under that contract the rest of the old imports are unstable too: §2.4 covers flydsl._mlir.*, underscore-prefixed names, and direct upstream-MLIR dialect calls; §2.2 excludes deep flydsl.compiler.* paths such as kernel_function.CompilationContext. None of it appears in the §3 deprecation table, so none of it was protected.

Each replacement is documented at the v0.3.0 tag:

before	after	reference
buffer_ops.create_buffer_resource / buffer_load / buffer_store / _unwrap_value	fx.rocdl.make_buffer_tensor, fx.make_copy_atom(fx.rocdl.BufferCopy128b(), ...), fx.memref_load_vec	authoring §4.3, layout §9
LDS via _memref.ViewOp / StoreOp / LoadOp	@fx.struct + fx.SharedAllocator, fx.memref_load / _store	authoring §6.1
_gpu.ShuffleOp(..., mode=ShuffleMode.XOR)	fx.gpu.shuffle_xor(x, offset, width)	api_stability §3
math_ops.RsqrtOp, arith_ops.*, scf.IfOp	fx.math.rsqrt, plain Python operators and if	arithmetic_types
CompilationContext	dropped; @flyc.jit owns compilation	authoring §1.1

Accuracy and Speed Tests

e2e validation on WAN2.2 T2V-A14B

Ran the rebased kernel end-to-end against WAN2.2; no regression.

Environment

• Image: lmsysorg/sglang-rocm:v0.5.18-rocm700-mi35x-20260824
• Single MI355X (gfx950), FlyDSL 0.3.1
• A/B method: only python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py was swapped between the upstream-main version and this PR's version. Everything else in the tree is identical, and both runs used a warm torch.compile cache.

Command

export SGLANG_USE_ROCM_FLYDSL=1
export SGLANG_USE_ROCM_VAE_CONV2D_BF16=1
export SGLANG_DIFFUSION_AITER_FP8_ATTN=1

sglang generate \
  --model-path /models/Wan2.2-T2V-A14B-Diffusers \
  --prompt "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." \
  --negative-prompt " " \
  --720p --num-frames 193 --num-inference-steps 8 --seed 42 \
  --num-gpus 1 \
  --dit-layerwise-offload true --dit-cpu-offload false \
  --vae-cpu-offload false --text-encoder-cpu-offload true \
  --warmup-mode request --enable-torch-compile true

Result (from [DenoisingStage] average time per step)

kernel	s/step	denoising total
upstream main	31.026	260.90 s
this PR	30.922	260.96 s
upstream main (repeat)	30.937	261.12 s

The two upstream runs differ by 0.29%, so this PR's 0.05% delta sits inside run-to-run noise. The full pipeline finished with no GPU faults and the decoded video was written normally.

A microbenchmark at the shapes WAN2.2 actually uses (LayerNorm, seq 176400, dim 5120, elementwise_affine=False for norm1/norm_out) agrees: the total time these kernels spend per DiT forward goes from 147.8 ms to 147.1 ms.

Output quality

Both upstream main and this PR produce a byte-identical video at seed 42 (720p, 193 frames, 8 steps):

run	md5
upstream main	4959325aeb967ce0dcfd6bce405a271f
upstream main (repeat)	4959325aeb967ce0dcfd6bce405a271f
this PR	4959325aeb967ce0dcfd6bce405a271f

The two upstream runs matching each other confirms the pipeline is deterministic under a fixed seed, so this is a meaningful comparison rather than a coincidence: the kernel rewrite changes nothing in the generated output.

This is consistent with the numerics checks at the lower levels — the 22 kernel unit tests pass against the reference implementations, and at the layer level forward_hip matches forward_native to within bf16 tolerance (max abs diff 0.125).

Note: e2e was run on the rocm700 image because torch.compile currently core dumps on the rocm720 image for unrelated reasons.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #33466569118
Latest PR Test (Extra): ❌ Run #33466568860
Latest PR Test (AMD ROCm 7.2): ❌ Run #33466568976

[yctseng0211]

This change is ROCm-only: it is reached solely through forward_hip and is opt-in behind SGLANG_USE_ROCM_FLYDSL, so no CUDA or CPU path is affected.

Verified end-to-end on WAN2.2 T2V-A14B: this PR and upstream main produce a byte-identical video at seed 42. The unit tests in test/registered/kernels/ops/diffusion/test_norm_flydsl.py pass in CI.

[sushildubey171]

  --dit-layerwise-offload true --dit-cpu-offload false \
  --vae-cpu-offload false --text-encoder-cpu-offload true

@yctseng0211, a few things don't look clear to me. Could you clarify why offloading was enabled when MI35X appears to have sufficient memory? It would also be great to include benchmark metrics demonstrating the performance improvements so the benefits are easier to assess.

you can also dump the performance breakup with: --perf-dump-path

For running perf benchmark you can use below command:

  
  sglang serve \
  --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
  --num-gpus 8 \
  --ulysses-degree 8 \
  --port 30000 \
  --attention-backend aiter \
  --performance-mode speed
  --enable-torch-compile false


python3 -m sglang.multimodal_gen.benchmarks.bench_serving \
  --port 30000 \
  --dataset vbench \
  --task text-to-video \
  --num-prompts 20 \
  --max-concurrency 1

cc: @HaiShaw @yichiche

[yctseng0211]

@sushildubey171
Thanks for taking a look. Two separate things here:

1. The offload flags just came from our existing WAN2.2 benchmark recipe so the numbers stay comparable with our earlier runs

2. Also, this PR is a migration to the new FlyDSL API rather than an optimization, there is no speedup to show, the goal is to drop the non-public APIs while keeping output and latency unchanged.

[sushildubey171]

@sushildubey171 Thanks for taking a look. Two separate things here:

1. The offload flags just came from our existing WAN2.2 benchmark recipe so the numbers stay comparable with our earlier runs
2. Also, this PR is a migration to the new FlyDSL API rather than an optimization, there is no speedup to show, the goal is to drop the non-public APIs while keeping output and latency unchanged.

@yctseng0211, thanks for the clarification. I understand this is a migration PR and not a performance optimization. My concern is mainly around regression validation:

• Since MI35X has sufficient memory, offloading should not be required.
• It would be useful to benchmark using the intended configuration (without offloading).
• Establishing before/after numbers would help confirm that the FlyDSL migration does not introduce any performance regressions.
• For an apples-to-apples comparison, we can regenerate the older benchmark using the same configuration and commands.

[yctseng0211]

Environment

• Image: lmsysorg/sglang-rocm:v0.5.18-rocm700-mi35x-20260824
• Single MI355X (gfx950), FlyDSL 0.3.1
• A/B method: only python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py is swapped between the upstream-main version and this PR's version. Everything else is identical.

Command

export SGLANG_USE_ROCM_FLYDSL=1
export SGLANG_USE_ROCM_VAE_CONV2D_BF16=1
export SGLANG_DIFFUSION_AITER_FP8_ATTN=1

sglang generate \
  --model-path /models/Wan2.2-T2V-A14B-Diffusers \
  --prompt "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." \
  --negative-prompt " " \
  --720p --num-frames 193 --num-inference-steps 8 --seed 42 \
  --num-gpus 1 \
  --dit-layerwise-offload <true|false> --text-encoder-cpu-offload <true|false> \
  --dit-cpu-offload false --vae-cpu-offload false \
  --warmup-mode request --enable-torch-compile true

Microbenchmark

Same image and GPU as the e2e runs. Shapes are what WAN2.2 T2V-A14B actually instantiates at 720p / 193 frames: 49 latent frames x 3600 tokens = seq 176400, dim 5120, LayerNorm. Median of 200 iterations after 50 warmup, timed with CUDA events, with baseline and this PR run in alternating processes (they register the same torch.library op names, so only one can live in a process).

kernel config	calls / DiT forward	upstream	this PR	delta
norm_scale_shift, no affine	41	1.0596 ms	1.0185 ms	-3.9%
fused residual, affine	40	1.7101 ms	1.6876 ms	-1.3%
fused residual, no affine	40	1.5740 ms	1.5582 ms	-1.0%
total per DiT forward		174.8 ms	171.6 ms	-1.8%

Reproduction:

import statistics, torch
from sglang.kernels.ops.diffusion import (
    flydsl_fused_residual_norm_scale_shift,
    flydsl_norm_scale_shift,
)
SEQ, DIM = 176400, 5120
mk = lambda *s, d=torch.bfloat16: torch.randn(*s, device="cuda", dtype=d)
x, res = mk(1, SEQ, DIM), mk(1, SEQ, DIM)
gate, scale, shift = mk(1, 1, DIM), mk(1, 1, DIM), mk(1, 1, DIM)
w, b = mk(DIM, d=torch.float32), mk(DIM, d=torch.float32)
cases = {
    "nss no-affine  (norm1 x40, norm_out x1)":
        lambda: flydsl_norm_scale_shift(x, None, None, scale, shift, "layer", 1e-6),
    "fused affine   (self_attn x40)":
        lambda: flydsl_fused_residual_norm_scale_shift(res, x, gate, w, b, scale, shift, "layer", 1e-6),
    "fused no-affine(cross_attn x40)":
        lambda: flydsl_fused_residual_norm_scale_shift(res, x, gate, None, None, scale, shift, "layer", 1e-6),
}
for name, fn in cases.items():
    for _ in range(50):
        fn()
    torch.cuda.synchronize()
    s, e, samples = torch.cuda.Event(True), torch.cuda.Event(True), []
    for _ in range(200):
        s.record(); fn(); e.record(); e.synchronize()
        samples.append(s.elapsed_time(e))
    print(f"{name}: {statistics.median(samples):.4f} ms")

End-to-end

[DenoisingStage] average time per step, 8 steps:

offload	kernel	s/step	denoising total
on	upstream	30.937	261.12 s
on	this PR	30.922	260.96 s
off	upstream	31.201	260.77 s
off	this PR	31.027	261.44 s

Two upstream runs of the same config differ by 0.29%, so every delta above sits inside run-to-run noise. Turning layerwise offload off does not change the step time either, so the offload flags were not distorting the comparison

Output

Byte-identical video in both configurations:

offload	upstream	this PR
on	4959325aeb967ce0dcfd6bce405a271f	4959325aeb967ce0dcfd6bce405a271f
off	ed85dcf0fe721a68318f9d023ca0cef6	ed85dcf0fe721a68318f9d023ca0cef6

At the kernel level the two versions are bit-identical on the same inputs across all three configurations WAN2.2 uses, at both a small shape and the real seq 176400. Dispatch to the FlyDSL custom ops was verified with a call counter, so neither side is silently falling back to forward_native.

[yichiche]

@sushildubey171 The recipe used for offloading is to compare with the B200 side by side for profiling and benchmarking. It's fine here either with or without this flag, considering the intention here is the upgrade of FlyDSL.

[yctseng0211]

@sushildubey171 @yichiche
results with and without offloading are above, both show no regression.

[yichiche]

@yctseng0211 The v0.3.0 rewrite is quite a shift from the previous FlyDSL internals (unified primitives + the new tuning-config path). Could you add a short pointer in the PR body to the AITER FlyDSL 0.3 example/docs this was patterned after? A link plus a brief old-API → new-API mapping would be enough for later kernel work; no need to paste a full kernel here.

[yctseng0211]

@yctseng0211 The v0.3.0 rewrite is quite a shift from the previous FlyDSL internals (unified primitives + the new tuning-config path). Could you add a short pointer in the PR body to the AITER FlyDSL 0.3 example/docs this was patterned after? A link plus a brief old-API → new-API mapping would be enough for later kernel work; no need to paste a full kernel here.

@yichiche thanks for your suggestion, added

[yctseng0211]

This change is ROCm-only: it is reached solely through forward_hip and is opt-in behind SGLANG_USE_ROCM_FLYDSL, so no CUDA or CPU path is affected.

Verified end-to-end on WAN2.2 T2V-A14B: this PR and upstream main produce a byte-identical video at seed 42. The unit tests in test/registered/kernels/ops/diffusion/test_norm_flydsl.py pass in CI.

cc: @HaiShaw