[damien-lejeune]

Motivation

Implement a two-stage Hyper Connection Gated Residual kernel in order to improve performance on the decode case. The current baseline is Triton/Hipblaslt 5-stage kernel.

Architecture: gfx950 and gfx942

Note: this kernel on gfx942 needs a BF16 GEMM, which is not available for now on FlyDSL. An adaptation of the gfx950 GEMM is provided temporarily but we aim to remove this implementation as soon a more efficient GEMM is ready as a drop-in replacement.

Technical Details

Paper: https://arxiv.org/abs/2409.19606

In the context of Hyper-Connections, let R ∈ ℝ^{HC×HS} be the incoming streams (flattened R ∈ ℝ^{K}, K=10240),
y ∈ ℝ^{HS} the block output to inject, inj ∈ ℝ^{HC} the incoming injection logits (from the previous layer), σ sigmoid and ⊙ elementwise product. HC is the number of parallel residual streams and HS is the hidden width of a single stream (2560, in our case). K is the flatten residual width, or equivalently, the down-GEMM contraction dimension

• (COMBINE) Depth re-injection

iw = 2·σ(inj / HC)                     inj = incoming injection logits (input) → iw:[HC]
R2[s,c] = R[s,c] + y[c] · iw[s]        R2 ∈ ℝ^{HC×HS}   (next residual)

• (N) Grouped RMSNorm, normalize each stream independently:

rrms = r_s = rsqrt(mean_c R2[s,c]² + eps)            the reciprocal-RMS scale, per stream s
xn[s,c] = R2[s,c]·r_s·(1+Wnorm)                      xn ∈ ℝ^{HC×HS}  (flattened ℝ^{K})

xn is the single normalized input reused by all three projections below.

• (G1) Down-GEMM, Wd : ℝ^{K} → ℝ^{r}, then SiLU. This GEMM (G1) encodes the low rank space.

lora = xn · Wdᵀ
h = SiLU(lora / HC ),   xn:[K] · Wdᵀ:[K×r] → h:[r]      r = 320

• (G2) Up-GEMM, Wu : ℝ^{r} → ℝ^{K}, then sigmoid. This GEMM (G2) decoder expands the low rank code back to the input vector K dimension size, per stream.

gate = h · Wuᵀ 
G = σ(gate),             h:[r] · Wuᵀ:[r×K] → G:[K]  ≡  G ∈ ℝ^{HC×HS}

• (MIX) Gated mean

x[c] = (1/HC) · Σ_{s=1..HC} G[s,c] · xn[s,c]      x ∈ ℝ^{HS}   = h_0

• (INJ) Inject-GEMM, the injection logits for the NEXT layer's combine:

inj_next = xn · Wiᵀ          for the next layer's combine; xn:[K] · Wiᵀ:[K×HC] → inj_next:[HC]  

The 2-stage kernel is split between:

• Kernel 1:
  • (COMBINE), inject the previous block output
  • (N), grouped RMSNorm
  • (G1), the down gemm h = (xn·Wdᵀ)/HC
  • (INJ)
  • write out the next residual R2, the reciprocal-RMS scale rrms, the down-GEMM output lora (pre-SiLU), and the next layer's injection logits inj_next into HBM
• Kernel 2:
  • (G1), the SiLU operation only
  • (G2), the up-GEMM
  • (MIX), gated mean
  • Rebuild the xn = R2 * rrms * (1+Wnorm) on the fly, such that xn and gate (i.e. up-GEMM output before sigmoid) are never stored in HBM

The 5-stage kernel is split between:

1. COMBINE + N
2. G1 (GEMM only) + INJ
3. G1 (SiLU tail)
4. G2 (GEMM)
5. G2 (sigmoid) + MIX

Every intermediate (R2, xn, gate, etc) incurs an HBM roundtrip

Test Plan

Unit tests have been provided. They should pass on both target architectures

Performance

gfx950 performance vs the vLLM baseline

Setup: MI355X (gfx950), single GPU, ROCm 7.2.4, PyTorch 2.10.0, Triton 3.6.0, FlyDSL 0.4.1, head 03121cd.

• Shape: HC=4, HS=2560 (hidden 10240), lowrank 320, bf16.
• Merged down/inject weight: 336 rows with (1+w) folded in, as shipped by vLLM.
• Baseline: vLLM qwen4_exp hc Triton kernels (_hc_combine_norm, _hc_silu, _hc_gate_mix) plus hipBLASLt F.linear for the down/inject and up projections, i.e. the 5-stage path. vLLM commit 114abd1.
• Ours: the public flydsl_gr_two_stage_combine_and_mix(..., fold_w=True).
• Method: every path is checked against the float32 oracle before timing. Each cell is the median of 60 timed calls, then the median of 3 full sweeps.
• Ratio: vLLM time / our time, so > 1 means the two-stage op is faster.
  HIP_VISIBLE_DEVICES=2 PYTHONPATH=<aiter>:<vllm> python op_tests/flydsl_tests/_bench_vs_triton_graph.py --tokens 1 2 4 8 16 32 512 1024 2048 4096 8192 16384 32768 65536 --mode both --iters 60

M	vLLM, graph (µs)	Two-stage, graph (µs)	Speed-up, graph	vLLM, eager (µs)	Two-stage, eager (µs)	Speed-up, eager
1	31.5	21.0	1.50×	82.3	29.7	2.77×
2	36.1	19.3	1.87×	89.5	28.9	3.10×
4	36.5	26.6	1.37×	88.0	30.4	2.89×
8	35.1	33.0	1.06×	89.1	37.4	2.38×
16	35.0	28.8	1.22×	90.8	50.0	1.82×
32	35.3	28.8	1.23×	89.1	49.3	1.81×
512	58.7	51.0	1.15×	90.4	59.9	1.51×
1024	75.9	76.3	0.99×	95.8	83.5	1.15×
2048	113.4	106.8	1.06×	123.4	116.2	1.06×
4096	194.0	184.5	1.05×	207.6	190.6	1.09×
8192	363.6	339.2	1.07×	383.2	343.6	1.12×
16384	673.5	618.8	1.09×	677.0	622.8	1.09×
32768	1294.7	1324.2	0.98×	1264.7	1318.8	0.96×
65536	2598.1	2636.4	0.99×	2505.7	2613.3	0.96×
Summary (graph replay, the mode vLLM uses for decode):

• 1.06–1.87× faster at decode, M=1–32;
• 1.04–1.15× faster at M=512 and from 2048 to 16384;
• parity (0.98–0.99×) at M=1024, 32768 and 65536.
  In eager mode the decode gain grows to 2.4–3.1×, because the 5-stage baseline is bound by launch overhead there.
  Run-to-run spread over the 3 sweeps:
• graph: under 2.5% for ours, up to 4.5% for vLLM (M=16384);
• eager: up to 5.4% for ours, up to 8.8% for vLLM.
  The GPU was shared with other processes during the runs.

PR Self-Review

Reviewed revision

03121cd

Behavior and intent

Adds a new FlyDSL two-stage Hyper-Connection gated-residual op (aiter/ops/flydsl/kernels/hyper_connection_gated_residual/) for gfx950 (the target) and gfx942 (functional support), bf16. Today this layer runs as a 5-stage Triton/hipBLASLt pipeline that round-trips every intermediate (R2, xn, gate) through HBM.

• K1 fuses combine, grouped RMSNorm, the down/inject GEMM and SiLU. It has three paths: split-K, decoupled async-LDS, and a skinny GEMV decode for M ≤ 8 with a folded weight.
• K2 fuses the up GEMM and the gated mean, re-forming xn from R2 and rrms in registers so that neither xn nor gate is ever stored.
• Public entry points:
  • flydsl_gr_two_stage_combine_and_mix, flydsl_gr_two_stage_mix and flydsl_gr_two_stage_combine;
  • the direct kernels flydsl_k1_combine_norm_down and flydsl_up_gate_mix_norm.
• The public entry points reject unsupported architectures (NotImplementedError), widths the kernels cannot tile, token counts no tile divides (pad with gemm_pad), mismatched streams and devices, stages < 2, row tails on direct K2, and malformed or aliased caller output buffers (ValueError).
• Existing behavior is unchanged: this PR adds a new op and a new config file and modifies no existing kernel. The only edit outside the new package adds four helpers (_llvm_ptr_global, _av_scope, av_store_b128, atomic_add_agent_acq_rel_av_none) and one import to aiter/ops/flydsl/kernels/communication_ops_utils.py. Existing functions there are untouched, and the helpers are used only by the skinny decode.

Validation and unverified scope

gfx950 (target): MI355X, single GPU, ROCm 7.2.4, FlyDSL 0.4.1 (as pinned by main), at head 03121cd.

• HIP_VISIBLE_DEVICES=2 python op_tests/test_hc_gated_residual.py passed: 292 checks, all against an independent float32 PyTorch oracle. Coverage:
  • the token sweep 512/2048/4096, folded and unfolded, for combine_and_mix, mix, combine and the final mixer;
  • decode M=1, 3–8 and 32;
  • n_pad=336 from M=1 to 8192;
  • the tuned range M=16384–65536;
  • direct K1 at M=224 and M=100: M=100 is rejected unpadded and matches the oracle with gemm_pad=128;
  • input, buffer, aliasing, stream, arch and pipeline-depth validation;
  • the skinny counter isolation, cache bound and publication stress tests.
• Ruff 0.16.0 and Black are clean on all changed files.
• Regression check against the pre-review revision 3ab1a23 (FlyDSL 0.4.0, before merging main):
  • the launched kernels are identical for 126 shape/fold/n_pad cases;
  • the outputs are bit-identical for 81 cases;
  • the CUDA-graph device time is within ±1% from M=1 to 24576, except combine-only, which is about 10% faster because it no longer computes the unused RMS.
• Skinny decode ISA on FlyDSL 0.4.1: there are no buffer_wbl2/buffer_inv instructions, and the LLVM IR keeps the acq_rel amdgcn-av:none arrival atomic and the llvm.amdgcn.av.store.b128 partial stores.
• Performance against the vLLM baseline at 03121cd: see the Performance section above.

gfx942 (functional support): AMD Instinct MI300X (304 CUs), single GPU, ROCm 7.2.4, PyTorch 2.10.0, FlyDSL 0.4.1, at head 8393037.

• python op_tests/test_hc_gated_residual.py passed: 265 checks against the float32 oracle. Coverage:
  • the token sweep, folded, unfolded and full-width norm weights, for combine_and_mix, mix, combine and the final mixer;
  • the skinny decode, including the counter isolation, cache bound, side-stream and publication stress tests (AV partial stores, amdgcn-av:none arrival atomic, raw sc1 loads);
  • the wide-stream K1 (HS=40960), strided combine, direct K1/K2 and the skinny n_pad=8 case;
  • the decode sweep and n_pad=336 from M=1 to 8192.

Unverified:

• gfx942 was not re-run at 03121cd. That commit replaces an assert with a ValueError in K1 and adds checks that run only on gfx950.
• gfx942 performance is not measured or claimed; it is functional support only.
• Pipeline-depth validation is gfx950-only: no exact-preload configuration fits gfx942 (block_k=256 exceeds its 64 KiB LDS; block_k=128 needs ≥ 5 stages, beyond the GEMM's vmcnt budget), so _run_pipeline_depth skips there.
• MI308X and other reduced-CU gfx942 parts are not tested.
• Cross-device execution is not tested: _run_stream_device_mismatch skips itself with only one visible GPU.
• Above M=8192, the oracle covers only the unfolded path; the folded path and n_pad=336 are not covered there.

PR-specific question

The only consumer of aiter/configs/model_configs/hc_gated_residual_tuned.json is this op's tuned.py; the tuner op_tests/tuners/tune_hc_gated_residual.py writes it. No other op reads the file, so no unrelated workload's selection can change.

Workload keys (all for HC=4, HS=2560, r=320):

Section	Token keys
gfx950/k1	64 to 65536
gfx950/k1_n336	64 to 65536
gfx950/up_gate_mix	64 to 8192
gfx950/up_gate_mix_n336	64
gfx942/k1	2048, 4096, 8192

How lookups behave:

• Token counts resolve to the nearest key in log space.
• Below the smallest K1 key, and for entries marked method: "heuristic", the kernel keeps its heuristic.
• up_gate_mix_n* entries apply only to exactly measured token counts.
• _k1_tuned_plan discards any plan that cannot tile the actual shape, such as a different hidden size or too few K tiles for the pipeline depth, and falls back to the heuristic. _run_untuned_hidden (HC=2) and _run_wide_stream_lds (HS=8192 and 40960) test this.

Changed op tests

• op_tests/test_hc_gated_residual.py (new) is the correctness suite for the new op; nothing in aiter covers it. It checks every public entry point and every kernel path (split-K, decoupled, skinny decode, gfx942 GEMM) against a float32 oracle, including the tuned token range. It also covers the error contract and the skinny cross-wave counter protocol:
  • malformed, aliased and non-contiguous inputs, row tails, untileable direct token counts, stages, arch and stream device;
  • hc isolation, the bounded cache under CUDA-graph capture, and a publication stress test.
• op_tests/flydsl_tests/bench_hc_gated_residual.py (new) is a standalone latency benchmark for the op. It times combine_and_mix with a pre-folded weight across a token sweep and checks each shape against the oracle before timing. It doesn't duplicate the test, because the test never measures latency. The comparison against vLLM in the Performance section uses a separate script, outside this PR, because it needs a vLLM checkout.
• op_tests/tuners/tune_hc_gated_residual.py (new) generates hc_gated_residual_tuned.json.
  • It sweeps split-K and decoupled K1 candidates per token count and rejects any candidate whose outputs are non-finite or outside tolerance.
  • It records heuristic when the untuned path wins.
  • It refuses to export when any token count is invalid, and writes the table atomically.
  • It is the only way to regenerate the config.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5913 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[junhaha666]

@aiter-bot review

[aiter-bot]

ROCm/aiter PR #5913 — [FlyDSL] feat(2stageGR) Two-stage Hyper-Connection Gated Residual kernel

Adds a new two-stage FlyDSL Hyper-Connection gated-residual op for gfx942/gfx950 bf16 — K1 fuses combine + grouped RMSNorm + down/inject GEMM, K2 fuses SiLU + up GEMM + gated mean with xn re-formed in registers so it never touches HBM — plus a fused skinny decode path, tuned configs, an oracle test, a bench and a tuner.

Review (advisory): ⚠️ NEEDS WORK
Validation (deterministic): NOT RUN — REVIEW_AUTO_VALIDATE=0 (auto-validation disabled in this environment); the single test target op_tests/test_hc_gated_residual.py exists but was not executed — an environment gap, not a PR defect.
Perf (advisory): NOT RUN — the triaged target op_tests/test_hc_gated_residual.py exposes no benchmark entry point, and the PR's own op_tests/flydsl_tests/bench_hc_gated_residual.py is head-only (the op is new, so base has no counterpart to time against); no base-vs-head timing exists for this kernel.

⚠️ [verified] flydsl_gr_two_stage_combine (op.py:272-277) is the only entry point that hands injection to the combine kernel with no contiguity assert, no .contiguous() and no stride handling, while both siblings handle the strided inj_next view this package itself returns (packed[:, lowrank:lowrank+hc_count], row stride n_pad — op.py:140, k1.py:1435, skinny.py:1004-1005), so a caller feeding the previous layer's injection logits into this public export gets silently wrong r2 (the kernel reads at a baked contiguous stride, in-bounds so nothing faults). Author must assert injection.is_contiguous() or thread the stride through _build_combine_rms as the skinny path does, and add a strided-injection case to op_tests/test_hc_gated_residual.py — no test currently passes one.
⚠️ [verified] k1.py:1480-1488 takes the tuned plan's sk_block_m with no divisibility check while the decouple path guards its own tile height (k1.py:1628-1631), so a direct call to the exported flydsl_k1_combine_norm_down with e.g. tokens=224 on gfx950 (nearest entry "256" in aiter/configs/model_configs/hc_gated_residual_tuned.json picks sk_block_m=64) launches grid_m=ceil(224/64)=4 tiles that read past r2's 224 rows and write past the [split_k*224, n_pad] partial buffer — both wrapped in max_size=True descriptors (k1.py:1231-1233), so nothing faults: silent heap corruption. Author must apply the decouple path's divisibility fallback (or assert gemm_tokens % sk_block_m == 0) before building the split-K partial.
⚠️ [verified] The skinny decode path adds three undocumented permanent env knobs — AITER_GR_SKINNY_TWO_KERNEL (k1.py:961, which also silently changes DECODE_MAX_M from 8 to 4), AITER_GR_SKINNY_COLS_PER_WAVE (skinny.py:447) and GR_FORCE_PIPE (k1.py:1456) — and its persistent arrival counters are torch.zeros-allocated on the first call's current stream then read by kernels launched on the caller's stream with no sync (skinny.py:451-457, 516-530), on top of the docstring's own "must not run concurrently on different streams" contract (skinny.py:477-479); a side-stream first call can race the zero-fill and concurrent calls corrupt the arrival detection into lost packed writes. Author must drop the knobs or document each as a temporary rollback switch, and allocate the counters on the caller's stream (or sync once at first use).
⚠️ [verified] The package rolls its own ISA detection instead of the chip_info helpers the FlyDSL policy mandates: common.py:60 reads torch.cuda.get_device_properties(device).gcnArchName inside the runtime dispatch path, gemm_a16w16_gfx942.py:173-187 probes get_rocm_arch() and hand-copies an LDS capacity map {gfx942: 64 KiB, gfx950: 160 KiB} duplicating chip_info._LDS_CAPACITY_BYTES, and k1.py:1456/1632 hardcode arch == "gfx942" literals — a second arch table that can drift; the same file also carries the gfx950 twin's dead split-K/HTI machinery with a reverted fp32-output fix (gemm_a16w16_gfx942.py:140/305/652/772), dead only because K1 always requests split_k=1 and bf16 output. Author must route arch and LDS queries through get_gfx_runtime()/get_lds_capacity_bytes() (fixing chip_info.py if per-device detection is genuinely needed) and trim or fence the unadapted dead paths.
⚠️ [verified] No base-vs-head timing exists for this kernel: the triaged target has no benchmark entry point, the PR's own bench is head-only (the op is new — base has no counterpart in aiter), and the claimed numbers (gfx950 +20% decode / 2-8% slower prefill; gfx942 +20-40% at M 1-8, "far behind" above) trace to no attached output; the tuned entries above 8192 tokens were validated by op_tests/tuners/tune_hc_gated_residual.py against the kernel's own heuristic output rather than the oracle and are never oracle-checked by the test sweep, and no consumer PR wires flydsl_gr_two_stage_combine_and_mix (op.py:168) into any model. Author must attach the bench log behind the claimed numbers, extend an oracle check across the tuned token range, and name the consumer that will activate the op — the untraceable numbers and the half-adapted twin copy warrant thorough manual verification of the dispatch logic and test coverage.

[junhaha666]

@damien-lejeune @mjkvaak-amd

Additional review findings

I reviewed the current head (b864f051b) and independently checked the issues reported by aiter-bot. The bot's strided-injection and split-K findings look valid. I also found two additional correctness issues that are not covered by that review.

1. flydsl_gr_two_stage_combine misreads a strided injection view

combine_and_mix returns inj_next as a view of packed[:, lowrank:lowrank + hc_count], whose row stride is n_pad, not hc_count. The standalone flydsl_gr_two_stage_combine entry point passes this tensor to _build_combine_rms without making it contiguous or passing its row stride, so the kernel reads subsequent rows at the wrong offsets.

I reproduced this on MI300X/gfx942 with an injection tensor of shape [3, 4] and stride (384, 1):

contiguous injection: max error = 0.0
strided injection:    max error = 1.46875

Please either make injection contiguous in this entry point or pass its row stride to _build_combine_rms, as the skinny path already does. A regression test should feed the strided inj_next returned by the previous layer directly into flydsl_gr_two_stage_combine.

2. The tuned split-K tile height can overrun direct K1 calls

For gfx950, an invocation such as tokens=224 selects the nearest 256 tuning entry, which has sk_block_m=64. flydsl_k1_combine_norm_down accepts that value without checking that it divides gemm_tokens, while _build_down_norm_partial launches:

grid_m = (tokens + block_m - 1) // block_m

The final tile therefore accesses rows beyond the 224-row r2/rrms tensors. The partial output is also indexed using bid_k * grid_m + bid_m, although only split_k * tokens rows were allocated. Because the affected buffers use max_size=True, this can become silent out-of-bounds access rather than a clean failure.

The high-level combine_and_mix wrapper pads M, but flydsl_k1_combine_norm_down is publicly exported in __all__, so direct calls remain unsafe. Please apply the same divisibility fallback used by the decoupled path, or reject unsupported token counts before launching the kernel.

3. Persistent skinny counters are not isolated by hc_count or invocation

The counter cache key is currently:

key = (device.index, n_pad, m_groups)

However, the arrival condition wraps modulo hc_count. Two sequential calls with the same n_pad and m_groups, but different hc_count, therefore reuse counters with an incompatible phase. This does not require concurrent streams to fail.

I reproduced the following sequence on MI300X/gfx942 using valid skinny-kernel shapes:

call 1: hc_count = 4
call 2: hc_count = 8

HC=8 injection max error before contamination: ~0.007
HC=8 injection max error after contamination:   5.78

This is in addition to the cross-stream race mentioned by aiter-bot: the counters are initialized on the current stream and are shared by concurrent launches on other streams. Adding hc_count to the cache key fixes the sequential phase mismatch, but it is not sufficient for concurrency. Each active invocation/stream needs isolated counters, or the counter lifecycle must be explicitly synchronized and reset.

4. w_up is silently interpreted as contiguous BF16 without validation

Both the tiled K2 kernel and the skinny up kernel construct a BF16 row-major view of w_up, but the public operator validates only its shape. A float32 or non-contiguous tensor with the expected shape is accepted and then read using the wrong element type or strides, producing silent numerical corruption.

On MI300X/gfx942 at M=8:

contiguous BF16 w_up: max block-input error = ~0.004
float32 w_up:         max block-input error = ~0.158
strided BF16 w_up:    max block-input error = ~0.205

Please validate at the shared public dispatch boundary that w_up is on the same device, has dtype torch.bfloat16, is two-dimensional and row-major contiguous, and has the expected shape.

Validation performed

The PR's correctness test passed on gfx942 for the exercised paths, including:

• M=512 and M=8192;
• decode M=1, 3 and 8;
• folded and unfolded norm weights;
• full-width norm weights;
• caller-padded n_pad=336 for M=1 through M=8192.

Those tests do not cover the failure modes above. I recommend treating findings 1-4 as correctness blockers. I also agree with aiter-bot that the gfx950 tuned range above M=8192 should be checked against an independent oracle rather than only against the kernel's untuned heuristic path.

[aiter-bot]

READY FOR HUMAN REVIEW

PR Self-Review

Revision 03121cda7ded · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

Adds a new FlyDSL two-stage Hyper-Connection gated-residual op (`aiter/ops/flydsl/kernels/hyper_connection_gated_residual/`) for gfx950 (the target) and gfx942 (functional support), bf16. Today this layer runs as a 5-stage Triton/hipBLASLt pipeline that round-trips every intermediate (`R2`, `xn`, `gate`) through HBM.

- **K1** fuses combine, grouped RMSNorm, the down/inject GEMM and SiLU. It has three paths: split-K, decoupled async-LDS, and a skinny GEMV decode for M ≤ 8 with a folded weight.
- **K2** fuses the up GEMM and the gated mean, re-forming `xn` from `R2` and `rrms` in registers so that neither `xn` nor `gate` is ever stored.
- Public entry points:
  - `flydsl_gr_two_stage_combine_and_mix`, `flydsl_gr_two_stage_mix` and `flydsl_gr_two_stage_combine`;
  - the direct kernels `flydsl_k1_combine_norm_down` and `flydsl_up_gate_mix_norm`.
- The public entry points reject unsupported architectures (`NotImplementedError`), widths the kernels cannot tile, token counts no tile divides (pad with `gemm_pad`), mismatched streams and devices, `stages < 2`, row tails on direct K2, and malformed or aliased caller output buffers (`ValueError`).
- Existing behavior is unchanged: this PR adds a new op and a new config file and modifies no existing kernel. The only edit outside the new package adds four helpers (`_llvm_ptr_global`, `_av_scope`, `av_store_b128`, `atomic_add_agent_acq_rel_av_none`) and one import to `aiter/ops/flydsl/kernels/communication_ops_utils.py`. Existing functions there are untouched, and the helpers are used only by the skinny decode.

Validation and unverified scope

**gfx950 (target):** MI355X, single GPU, ROCm 7.2.4, FlyDSL 0.4.1 (as pinned by `main`), at head 03121cd.

- `HIP_VISIBLE_DEVICES=2 python op_tests/test_hc_gated_residual.py` passed: 292 checks, all against an independent float32 PyTorch oracle. Coverage:
  - the token sweep 512/2048/4096, folded and unfolded, for combine_and_mix, mix, combine and the final mixer;
  - decode M=1, 3–8 and 32;
  - `n_pad=336` from M=1 to 8192;
  - the tuned range M=16384–65536;
  - direct K1 at M=224 and M=100: M=100 is rejected unpadded and matches the oracle with `gemm_pad=128`;
  - input, buffer, aliasing, stream, arch and pipeline-depth validation;
  - the skinny counter isolation, cache bound and publication stress tests.
- Ruff 0.16.0 and Black are clean on all changed files.
- Regression check against the pre-review revision 3ab1a23 (FlyDSL 0.4.0, before merging `main`):
  - the launched kernels are identical for 126 shape/fold/`n_pad` cases;
  - the outputs are bit-identical for 81 cases;
  - the CUDA-graph device time is within ±1% from M=1 to 24576, except combine-only, which is about 10% faster because it no longer computes the unused RMS.
- Skinny decode ISA on FlyDSL 0.4.1: there are no `buffer_wbl2`/`buffer_inv` instructions, and the LLVM IR keeps the acq_rel `amdgcn-av:none` arrival atomic and the `llvm.amdgcn.av.store.b128` partial stores.
- Performance against the vLLM baseline at 03121cd: see the Performance section above.

**gfx942 (functional support):** AMD Instinct MI300X (304 CUs), single GPU, ROCm 7.2.4, PyTorch 2.10.0, FlyDSL 0.4.1, at head 8393037.

- `python op_tests/test_hc_gated_residual.py` passed: 265 checks against the float32 oracle. Coverage:
  - the token sweep, folded, unfolded and full-width norm weights, for combine_and_mix, mix, combine and the final mixer;
  - the skinny decode, including the counter isolation, cache bound, side-stream and publication stress tests (AV partial stores, `amdgcn-av:none` arrival atomic, raw `sc1` loads);
  - the wide-stream K1 (HS=40960), strided combine, direct K1/K2 and the skinny `n_pad=8` case;
  - the decode sweep and `n_pad=336` from M=1 to 8192.

Unverified:

- gfx942 was not re-run at 03121cd. That commit replaces an assert with a `ValueError` in K1 and adds checks that run only on gfx950.
- gfx942 performance is not measured or claimed; it is functional support only.
- Pipeline-depth validation is gfx950-only: no exact-preload configuration fits gfx942 (`block_k=256` exceeds its 64 KiB LDS; `block_k=128` needs ≥ 5 stages, beyond the GEMM's vmcnt budget), so `_run_pipeline_depth` skips there.
- MI308X and other reduced-CU gfx942 parts are not tested.
- Cross-device execution is not tested: `_run_stream_device_mismatch` skips itself with only one visible GPU.
- Above M=8192, the oracle covers only the unfolded path; the folded path and `n_pad=336` are not covered there.

PR-specific question

The only consumer of `aiter/configs/model_configs/hc_gated_residual_tuned.json` is this op's `tuned.py`; the tuner `op_tests/tuners/tune_hc_gated_residual.py` writes it. No other op reads the file, so no unrelated workload's selection can change.

Workload keys (all for HC=4, HS=2560, r=320):

| Section | Token keys |
|---|---|
| `gfx950/k1` | 64 to 65536 |
| `gfx950/k1_n336` | 64 to 65536 |
| `gfx950/up_gate_mix` | 64 to 8192 |
| `gfx950/up_gate_mix_n336` | 64 |
| `gfx942/k1` | 2048, 4096, 8192 |

How lookups behave:

- Token counts resolve to the nearest key in log space.
- Below the smallest K1 key, and for entries marked `method: "heuristic"`, the kernel keeps its heuristic.
- `up_gate_mix_n*` entries apply only to exactly measured token counts.
- `_k1_tuned_plan` discards any plan that cannot tile the actual shape, such as a different hidden size or too few K tiles for the pipeline depth, and falls back to the heuristic. `_run_untuned_hidden` (HC=2) and `_run_wide_stream_lds` (HS=8192 and 40960) test this.

Changed op tests

- `op_tests/test_hc_gated_residual.py` (new) is the correctness suite for the new op; nothing in aiter covers it. It checks every public entry point and every kernel path (split-K, decoupled, skinny decode, gfx942 GEMM) against a float32 oracle, including the tuned token range. It also covers the error contract and the skinny cross-wave counter protocol:
  - malformed, aliased and non-contiguous inputs, row tails, untileable direct token counts, `stages`, arch and stream device;
  - hc isolation, the bounded cache under CUDA-graph capture, and a publication stress test.
- `op_tests/flydsl_tests/bench_hc_gated_residual.py` (new) is a standalone latency benchmark for the op. It times `combine_and_mix` with a pre-folded weight across a token sweep and checks each shape against the oracle before timing. It doesn't duplicate the test, because the test never measures latency. The comparison against vLLM in the Performance section uses a separate script, outside this PR, because it needs a vLLM checkout.
- `op_tests/tuners/tune_hc_gated_residual.py` (new) generates `hc_gated_residual_tuned.json`.
  - It sweeps split-K and decoupled K1 candidates per token count and rejects any candidate whose outputs are non-finite or outside tolerance.
  - It records `heuristic` when the untuned path wins.
  - It refuses to export when any token count is invalid, and writes the table atomically.
  - It is the only way to regenerate the config.

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot did not review the current head SHA

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected
• ADVISORY — new op test op_tests/flydsl_tests/bench_hc_gated_residual.py has no obvious assertion/check in the available patch
• ADVISORY — new op test op_tests/tuners/tune_hc_gated_residual.py has no obvious assertion/check in the available patch

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• PASS — Aiter Test Gate: passed

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[damien-lejeune]

@junhaha666 thanks for the detailed review and repros. All four findings, and the oracle-coverage point, are addressed:

Problem	Commit(s)	Comments
1. flydsl_gr_two_stage_combine misreads a strided injection view (the inj_next returned by combine_and_mix, row stride n_pad)	b9dc14d, 4773830	The combine entry point now passes the injection's row stride to _build_combine_rms, as the skinny path does, and accepts any injection with a contiguous inner dimension. _run_strided_combine feeds the strided inj_next returned by the previous layer straight into flydsl_gr_two_stage_combine and checks it against the oracle.
2. Tuned split-K tile height overruns direct K1 calls (e.g. tokens=224 picking sk_block_m=64)	b9dc14d, 03121cd	The split-K tile height now falls back to one that divides the GEMM row count, so a nearest-neighbour plan can't launch past the buffers. _run_direct_k1_untiled_tokens checks M=224 against the oracle. Row counts that no tile divides (e.g. M=100) now raise a ValueError that names gemm_pad, instead of reaching the kernel; the test also checks the padded call (gemm_pad=128) against the oracle for split_k auto, 1 and 4.
3. Persistent skinny counters not isolated by hc_count or invocation; cross-stream race on the zero-fill	b9dc14d, 4773830, 6bfa2a8, 3a693c3, cfada82, 6b47fa8, 2751c40	Counter sets are keyed by device, n_pad, m_groups, hc_count and stream, and are zero-filled on the launch stream (b9dc14d, 4773830). Arrival is now an agent-scope acq_rel atomic tagged amdgcn-av:none, with available-on-write partial stores and sc1 loads (6bfa2a8, 3a693c3, cfada82). Only lane 0 publishes, counter slots are padded to 128 B, and the cache is LRU-bounded (6b47fa8). Sets used by a CUDA graph are pinned and bounded too (2751c40). Tests: _run_counter_hc_isolation (your HC=4 then HC=8 sequence), _run_side_stream_pressure, _run_skinny_publication_stress and _run_counter_cache_bound.
4. w_up silently read as contiguous bf16 without validation	b9dc14d	_validate_w_up runs at the shared public dispatch boundary and on the direct K2 entry point. It checks the device, torch.bfloat16, 2-D shape (hidden, lowrank) and row-major contiguity, and raises ValueError otherwise. _run_w_up_validation covers float32 and non-contiguous w_up on both the skinny and tiled paths.
gfx950 tuned range above M=8192 checked only against the untuned kernel (also raised by aiter-bot)	b9dc14d	On gfx950, every tuned token count above the ordinary sweep (M=16384, 24576, 32768, 49152, 65536) is checked against the independent float32 PyTorch oracle. That coverage is for the unfolded path; the folded path and n_pad=336 above 8192 are not covered.

Validated on gfx950 (full suite, FlyDSL 0.4.1) and on gfx942 (MI300X, FlyDSL 0.4.1); details are in the PR Self-Review section.

[damien-lejeune]

@aiter-bot findings from the Nov 2 review, and how each was addressed:

Problem	Commit(s)	Comments
1. flydsl_gr_two_stage_combine hands a strided injection (inj_next view, row stride n_pad) to the combine kernel at a baked contiguous stride, giving silently wrong r2	b9dc14d, 4773830	The injection's row stride is now passed to _build_combine_rms, as on the skinny path. _run_strided_combine feeds the strided inj_next from combine_and_mix straight into flydsl_gr_two_stage_combine and checks it against the oracle.
2. Tuned sk_block_m taken without a divisibility check, so a direct flydsl_k1_combine_norm_down call (e.g. tokens=224) writes past the r2/partial buffers	b9dc14d, 03121cd	The split-K tile height now falls back to one that divides the GEMM row count (_run_direct_k1_untiled_tokens, M=224, oracle-checked). Row counts that no tile divides now raise a ValueError that names gemm_pad, instead of reaching the kernel; the test also checks the padded call (gemm_pad=128) against the oracle for split_k auto, 1 and 4.
3a. Undocumented permanent env knobs AITER_GR_SKINNY_TWO_KERNEL (also changes DECODE_MAX_M), AITER_GR_SKINNY_COLS_PER_WAVE, GR_FORCE_PIPE	b9dc14d	Each is now documented at its definition as a temporary diagnostic or rollback switch that is not a supported tuning knob and may be removed once the path is proven. The defaults are the validated configuration.
3b. Skinny arrival counters zero-filled on the first call's current stream and shared across streams, so a side-stream call can race the zero-fill and concurrent calls can lose packed writes	b9dc14d, 4773830, 6bfa2a8, 3a693c3, cfada82, 6b47fa8, 2751c40	Counter sets are keyed by device, n_pad, m_groups, hc_count and stream, and are zero-filled on the launch stream, so concurrent streams never share a set (b9dc14d, 4773830). Arrival is now an agent-scope acq_rel atomic tagged amdgcn-av:none, with available-on-write partial stores and sc1 loads (6bfa2a8, 3a693c3, cfada82). Only lane 0 publishes, slots are padded to 128 B, and the cache is LRU-bounded (6b47fa8). Sets used by a CUDA graph are pinned and bounded too (2751c40). Tests: _run_side_stream_pressure, _run_counter_hc_isolation, _run_skinny_publication_stress and _run_counter_cache_bound.
4. Own ISA detection (gcnArchName), a hand-copied LDS capacity map, arch == "gfx942" literals, and dead split-K/HTI paths with a reverted fp32 fix in gemm_a16w16_gfx942.py	b9dc14d, 4773830	Arch and LDS queries now go through chip_info.get_gfx_runtime() and get_lds_capacity_bytes(); the gcnArchName probe and the copied capacity map are gone (b9dc14d). The gfx942 GEMM lost 1,222 lines of the gfx950 twin's dead split-K/HTI machinery, and those modes now raise ValueError (4773830). The arch == "gfx942" checks that remain choose gfx942-specific code (32-bit async load width, register-prefetch split-K), not capacities, so they cannot drift from chip_info.
5. No base-vs-head timing behind the claimed numbers; tuned entries above 8192 validated only against the heuristic; no named consumer	b9dc14d (oracle); benchmark log in this comment / PR description	gfx950 (the target) is now benchmarked against the vLLM baseline (hc Triton kernels + hipBLASLt F.linear) at head 03121cd, on M=1–65536, with oracle validation, the command and versions; the claims in the PR description are updated to match. Graph replay: 1.06–1.87× at decode M=1–32, 1.04–1.15× at M=512 and 2048–16384, parity at 1024/32768/65536. On gfx950, every tuned token count above M=8192 is checked against the independent float32 oracle. gfx942 is provided as functional support only; its performance is not claimed. Consumers: vllm-project/vllm#59442 and vllm-project/vllm#59575 dispatch combine_and_mix to this op.
PR self-review readiness (Nov 3 comment)	n/a	Addressed in the PR description (PR Self-Review section).

The Ruff F821 findings reported by reviewdog on communication_ops_utils.py (after merging main) are fixed in a64d3ce.