[gurpreet-dhami]

Motivation

The fused QK-norm + RoPE and fused LayerNorm + modulate kernels were previously gated to CUDA only. This PR enables these fused paths on ROCm, giving FLUX.2 the same eager-mode optimizations on AMD GPUs that already exist on NVIDIA.

Additionally, the single transformer block's packed-QKV unflatten produced strided views that the in-place fused kernel couldn't consume without implicit copies. Instead of forcing .contiguous() copies, we plumb allow_strided_qk=True so the kernel handles strided views natively.

Modifications

• Enable fused QK-norm + RoPE on ROCm: Gate the fused in-place path on _is_cuda or _is_rocm instead of _is_cuda alone in apply_qk_norm and apply_qk_norm_rope (layernorm.py).
• Enable fused LayerNorm + modulate on ROCm: The _flux2_norm_modulate helper (already platform-agnostic via .is_cuda tensor check) now runs on ROCm for both double and single transformer blocks (flux_2.py).
• Support strided QKV views: Add allow_strided_qk parameter to apply_qk_norm_with_optional_rope and pass allow_strided_qk=True in Flux2ParallelSelfAttention and Flux2SingleTransformerBlock, avoiding unnecessary .contiguous() copies.

Results

Benchmarked on MI350 (ROCm) with FLUX.2 at two resolutions for TP=1, ULLYSSES=8. Times are per-step per-layer GPU kernel time in microseconds. Estimated E2E uplift is extrapolated from per-layer kernel savings; measured E2E is in Accuracy Tests below.

Kernel microbenchmarks — 1024×1024 (~10% estimated E2E uplift)

Transformer Component	Q/K Shape	Unfused (us)	Fused (us)	Speedup	Total savings (50 steps, us)
Double Stream (Text)	[1, 64, 48, 128]	101.7	8.4	12.11x	37,320
Double Stream (Image)	[1, 512, 48, 128]	131.6	14.7	8.95x	46,760
Single Stream (Joint)	[1, 576, 48, 128]	153.5	15.8	9.72x	330,480

Kernel microbenchmarks — 2048×2048 (~6% estimated E2E uplift)

Transformer Component	Q/K Shape	Unfused (us)	Fused (us)	Speedup	Total savings (50 steps, us)
Double Stream (Text)	[1, 64, 48, 128]	104.9	8.5	12.34x	38,560
Double Stream (Image)	[1, 2048, 48, 128]	257.1	40.0	6.43x	86,840
Single Stream (Joint)	[1, 2112, 48, 128]	339.1	44.4	7.64x	707,280

Accuracy tests

Verified that the fused ROCm paths (QK-norm + RoPE, LayerNorm + modulate, strided QKV) produce identical accuracy to the unfused baseline on AMD MI350X using SGLang's built-in diffusion consistency test suite (test_server_1_gpu.py -k flux_2_image_t2i).

Setup: FLUX.2-dev, 1 GPU, 1024×1024, 50 denoising steps, seed=42, prompt: "Doraemon is eating dorayaki"

Results

Build	CLIP Score	SSIM	PSNR	Mean Abs Diff	E2E (ms)	Avg Denoise (ms)
Unfused (baseline)	0.9837	0.7330	10.0566	33.7260	16331.11	283.85
Fused (this PR)	0.9843	0.7329	10.0570	33.7424	15179.06	248.37

• No accuracy regression. CLIP score (semantic similarity) is identical or marginally better with fused paths (0.9843 vs 0.9837). All pixel-level metrics (SSIM, PSNR, mean abs diff) are within noise across builds.
• Performance improvement. Fused paths reduce average denoise step time from ~284 ms to ~248 ms (~12.5% faster), and E2E latency from ~16.3 s to ~15.2 s (~6% faster).

Note: SSIM/PSNR/mean_abs_diff values reflect AMD-vs-H100 ground truth comparison. The relevant comparison is across builds on the same platform, where all metrics are consistent.

---

CI States

Latest PR Test (Base): ✅ Run #35280442691
Latest PR Test (Extra): ❌ Run #35280442529
Latest PR Test (AMD ROCm 10): ❌ Run #35280442520

[sushildubey171]

remove contiguous

[gurpreet-dhami]

Added requested changes.

[gurpreet-dhami]

@sushildubey171 @HaiShaw : Can you please review this PR.

[yichiche]

LGTM