Repository navigation
[Perf] [AMD] Qwen4-Exp HC: fuse decode-sized hyper-connection norm and inject into FlyDSL mix - #48648
amd-cgilli wants to merge 1 commit into
Conversation
… (11 + 16); PR links with state badges in the tables
…d inject into FlyDSL mix On gfx950, behind SGLANG_OPT_QWEN4_FUSED_HC, HC sites with <= 128 rows fold the per-branch Gemma RMSNorm and the combine's inject GEMV into a FlyDSL HC mix (down split-K -> reduce + SiLU -> up + gate). The mix hands its combine the inject logits as gate partials, so the combine is one elementwise kernel that also returns the sums of squares the next mix needs; combine_and_normalize passes them on in place of the next norm, and the next mix takes them as its normalized_input. Qwen K=10240, R=320, mix + combine per site: 33 -> 11 us at M=1, 34 -> 15 us at M=24, 45 -> 25 us at M=96. Not bit-exact (BF16 rounding points of the normed input and (1 + w) move). Co-authored-by: Nikolai Protasov <nprotaso@amd.com>
f6f9c1c to
ac68749
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Required checks failed on ac68749; this PR cannot be merged until they pass. PR title (run)
|
Motivation
On Qwen3.8-Flash-Next (
qwen4_exp_text), each hyper-connection (HC) site at decode runs the per-branch Gemma RMSNorm, the mix and the combine's inject GEMV as separate launches. At decode-sized batches these sites are launch- and memory-bound on MI355X.Modifications
kernels/ops/gemm/hc_fused_flydsl.py(new): FlyDSL kernels that fold the per-branch RMSNorm and the inject GEMV into the HC mix (hc_fused_mix); the combine becomes one elementwise kernel (hc_apply_inject) that also returns the sums of squares of its output. The inject logits travel in the gate-partials slot. Both are registered asgemm.hc_fused_mix/gemm.hc_apply_inject(FlyDSL, HIP).layers/hyperconnection.py:GatedResidualuses them on gfx950 for HC sites with <= 128 rows; larger batches keep the existing path unchanged.combine_and_normalizereturns the combine's sums of squares in place of the next norm, and the next fused mix takes them asnormalized_input. Without them (first layer, PLE layers) the down kernel computes them.hc_fused_shapes_supported); a site with unsupported shapes stays on the existing path.flydsl_hc_mix_supportedfrom [Perf] [AMD] Qwen3.8: AITER FlyDSL hc_mix for Qwen3.8 hyper-connections on gfx950 #41880 (bf16, contiguous, 16-B aligned, 1-128 rows) and skips the fused path undertorch.compile.environ.py:SGLANG_OPT_QWEN4_FUSED_HC(default off).test/registered/kernels/ops/gemm/test_hc_fused_flydsl.py(MI35x suite): kernels vs the unfused reference, including the mix -> combine -> mix chain that passes the sums of squares.Co-authored with Nikolai Protasov.
Accuracy Tests
Not bit-exact: the normed input and
(1 + w)are no longer rounded to BF16 separately. On MI355X, fused vs unfused for M = 1-128: combine outputs within 1 BF16 ulp, with the same error against an fp32 reference as the unfused path; M = 129 falls back and is bit-exact; CUDA-graph replay equals eager.GSM8K (1319, real MTP verification, Qwen3.8-Flash-Next TP1): 96.3% with
SGLANG_OPT_QWEN4_FUSED_HC=1vs 96.6% with it off (same upstream main, same AITER).Speed Tests and Profiling
MI355X, per HC site (K=10240, R=320), mix + combine, CUDA graph over 16 layers:
Related: #46881 (AMD Triton seam fusion, draft), #40552 (SM100 HC), ROCm/aiter#5913 (two-stage FlyDSL gated residual).
Checklist