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

[Phil-amd]

Motivation

Fixes #33142.

fused_residual_norm.py does not import against FlyDSL 0.3.0: flydsl.expr.buffer_ops
moved to kernels.common.buffer_ops in FlyDSL 7dac082 (FlyDSL #880), and that package
is in the FlyDSL source tree rather than the wheel. Because both call sites in
multimodal_gen/runtime/layers/layernorm.py guard the import with except ImportError,
the ROCm fused-norm path is silently disabled and every call falls back to native, even
with SGLANG_USE_ROCM_FLYDSL=1.

Restoring that one import would not hold for long. 278 of the module's 902 lines
reference something outside FlyDSL's public surface (158 raw upstream MLIR dialect
builders, 122 unstable FlyDSL paths, 2 deprecated with a removal release already
announced). None of those carry a compatibility statement, so the same thing can happen
again at any FlyDSL minor release. Details and the full audit table are in #33142.

Modifications

Both kernel builders are rewritten against the FlyDSL 0.3.0 stable public API only:

• fx.rocdl.make_buffer_tensor + fx.slice + fx.logical_divide + fx.make_copy_atom /
  fx.copy_atom_call over fx.make_rmem_tensor, replacing buffer_ops resources and
  flat element-offset arithmetic
• fx.SharedAllocator over an @fx.struct layout, replacing the hand-built dynamic LDS
  memref and the explicit smem= launch argument
• fx.gpu.shuffle_xor / fx.gpu.barrier, fx.math.rsqrt, fx.memref_load /
  fx.memref_store, typed values and Python control flow, replacing the arith / gpu /
  memref / scf / vector / math dialect builders and flydsl._mlir.ir
• an import-time capability check that raises ImportError (not AttributeError) when a
  required stable symbol is missing, so the existing native fallback still engages against
  an older FlyDSL

Deliberately unchanged: both torch.library.custom_op names, signatures and return arity;
WARP_SIZE/_NUM_WAVES/_VEC/FLYDSL_NORM_MIN_ALIGNED_DIM; one block per row at 640
threads; the two-stage reduction order; the phase-2 register cache; _prep_slices; the
compile-cache key; and the hardcoded eps = 1e-6. layernorm.py is untouched.

Broadcast operands previously addressed via a runtime row * stride; they now select row 0
through (stride != 0).select(row, 0), which is equivalent because a broadcast operand is
always passed as a dense (1, C) tensor.

Tests go from 8 to 22 cases, adding the gate-free and no-affine specializations (both
previously untested), per-row scale/shift, a 4D frame-gate layout, D=10240 for a second
register-cached iteration, a cross-L/layout compile-cache reuse guard, and a wheel-only
import check that fails if the module ever depends on the FlyDSL source tree.

Accuracy Tests

Baseline captured by running the unmodified module under the same FlyDSL 0.3.0 build
via a validation-only buffer_ops alias, so this is a same-compiler comparison.

61 cases covering every specialization (is_rms x has_gate x has_weight), both
scale/shift layouts, D in {5120, 10240}, L in {16, 90000}, and the 4D frame-gate layout:

	gfx942 / MI308X	gfx950 / MI355X
output tensors bit-exact (torch.equal)	100 / 100	100 / 100
test_flydsl_fused_norm.py	22 passed	22 passed

Speed Tests and Profiling

Median over 300 iterations after 200 warmup, baseline and candidate alternated across 3
independent process pairs.

shape	gfx942	gfx950
fused B1-L16-D5120	+1.07%	-0.98%
fused B2-L16-D5120	+1.14%	-0.42%
fused B1-L90000-D5120	+1.85%	+1.09%
nss B1-L16-D5120	+1.54%	+2.96%
nss B2-L16-D5120	+0.77%	+1.93%
nss B1-L90000-D5120	-0.07%	-0.03%

Final ISA is unchanged in every invariant that matters: 7 buffer_load_dwordx4,
2 buffer_store_dwordx4, zero non-128-bit or global_*/flat_* accesses, 2 ds_read /
2 ds_write, 2 s_barrier, and 12 cross-lane shuffles confirming the two-stage x six-step
reduction tree is preserved. Total instruction count 377 to 389.

Resources improve slightly: VGPR 38 to 32, no spills or scratch in either version, and LDS
goes from 88 B dynamic to 56 B static for RMSNorm (the unused accumulator is dead-code
eliminated) or 104 B for LayerNorm.

Known gaps

Two things a reviewer should know rather than discover:

1. CI may not go green on its own. I could not find any FlyDSL provisioning in
   .github/workflows/ or scripts/, and the AMD runner image is resolved dynamically, so
   I cannot tell whether it ships FlyDSL 0.3.0. If it ships an older FlyDSL or none, these
   tests fail rather than skip. That is not a regression (they fail on main today for the
   same underlying reason), but it likely needs a maintainer to confirm the image.
2. The layer-level path was not exercised end to end. Importing
   sglang.multimodal_gen pulls in sglang.srt, which does not import in a plain
   rocm/pytorch container at this revision: transformers==5.12.1 (the pyproject pin)
   lacks PreTrainedConfig required by srt/configs/cohere2_moe.py, while 5.14.1 fails
   on a duplicate qwen3_asr registration. This reproduces on both GPUs and is unrelated to
   this change. I substituted a check on exactly what layernorm.py depends on (exported
   names, both custom ops registered, the guard raising ImportError verified by deleting
   fx.SharedAllocator, both call sites intact, and numerics at the dispatch shape), which
   passes on gfx942 and gfx950 — but the real layernorm.py path itself was never run.

Out of scope, preserved exactly as-is and worth separate fixes: eps is ignored, FP32
affine weights are cast to BF16, 2D [B, C] broadcasting is incorrect, there is no
empty-tensor guard, and fused vs norm-only output dtype handling differs.

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. — no user-facing behavior change, so no docs touched
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #30682896481
Latest PR Test (Extra): ❌ Run #30682896419

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.