[RFC] FlyDSL kernels in downstream libraries

[jhinpan]

This RFC started with a simple question: how should a downstream kernel library consume kernels written in FlyDSL?

After implementing RMSNorm backward and studying both Quack and AITER, the main issue is no longer kernel capability. It is ownership.

FlyDSL now has RMSNorm forward + backward, residual backward, autograd coverage, and an optimized staged-reduction path (#795, #800, #855). The current backward implementation is competitive with the tested PyTorch and AIter baselines on MI355X.

The remaining question is where production kernel code and integration logic should live.

Proposed boundary

Follow the model already used by AITER:

FlyDSL owns

• the DSL, compiler, and runtime;
• target selection, JIT cache correctness, device/context safety, streams, and diagnostics;
• generic infrastructure that is useful to more than one downstream.

The downstream library owns

• the production FlyDSL kernel used by that library;
• operator semantics, dtype support, autograd, framework custom ops, fallback behavior, and dispatch;
• operator-specific tuning, tests, benchmarks, packaging, and CI.

This keeps one semantic owner for each operator. It also avoids turning FlyDSL into an operator library just so Quack can import one RMSNorm implementation. Quack can adapt the existing Apache-2.0 FlyDSL kernel with attribution and then evolve it with Quack's API.

Quack RMSNorm PoC

The scope question is now open in Dao-AILab/quack#178. We should not open implementation sub-issues until Quack maintainers confirm that a non-CuTe backend belongs in the project.

If accepted, the work should be split into small Quack-owned changes:

1. Isolate eager CuTe/CUDA imports and decide how CUDA/ROCm dependencies are installed.
2. Add an explicit experimental FlyDSL RMSNorm entry point.
3. Implement plain weighted forward + backward first, including FP16/BF16/FP32 inputs, FP32 weights for 16-bit inputs, runtime eps, and flattenable contiguous inputs.
4. Validate on gfx942 and gfx950 against Quack's references and ROCm PyTorch/AIter.
5. Enable transparent quack.rmsnorm dispatch only after the agreed API and torch.compile contract are covered.

Bias, residual/prenorm, per-head mode, output dtype conversion, optional weight, and weight_offset are follow-ups owned by Quack—not reasons to expand FlyDSL's public operator surface.

FlyDSL work

Only one required FlyDSL-core gap has been identified:

• #878: resolve the JIT target from the invocation device and keep target/cache reuse safe across GPU architectures.

The existing torch.compile discussion in #596 may still lead to a useful generic helper, but it is not a blocker for the Quack PoC. Quack and AITER already own operator-specific custom-op wrappers.

The RMSNorm training tracker #769 is complete and closed. Generic autotune work remains in #770, but Quack-specific search spaces and selected configs should live in Quack.

Decision

For this RFC, the proposed direction is:

• keep FlyDSL focused on reusable compiler/runtime capability;
• let Quack own its FlyDSL RMSNorm implementation and integration;
• treat [Runtime] Resolve the JIT target from the input device #878 as the only current FlyDSL blocker;
• wait for the scope decision in Quack [Feat] hide LLVM symbols in libFlyPythonCAPI.so to avoid conflicts wi… #178 before creating Quack implementation tasks.

[jhinpan]

Supplement: interface-gap analysis — quack rmsnorm ↔ FlyDSL rmsnorm_kernel.py

Concrete findings from reading the actual code on both sides (quack quack/rmsnorm.py @ v0.5.3; FlyDSL kernels/rmsnorm_kernel.py @ main). This grounds PoC step 1 ("interface parity") in real signatures rather than intent.

1. The two sides sit at different abstraction layers

	quack	FlyDSL
Public entry	rmsnorm(x, weight, bias, residual, out_dtype, residual_dtype, eps, prenorm) → autograd-aware Tensor (rmsnorm.py:1432)	build_rmsnorm_module(N, dtype_str) → returns a launch_rmsnorm(Input, Gamma, Output, m_in, stream) closure (rmsnorm_kernel.py:113,294)
Layer	PyTorch op: takes torch.Tensor, ships its own torch.autograd.Function, torch.compile-traceable	kernel launcher: takes fx.Tensor, caller passes grid/block and pre-allocates the output
Compile model	runtime @autotune over shape/dtype	N and dtype are baked at build time (build_rmsnorm_module(N, dtype_str)), one module per N

Core gap: this is not a "swap the kernel backend" change. quack's shim expects a (torch.Tensor) -> torch.Tensor autograd op; FlyDSL today exposes a low-level launcher with fixed N, manual launch, and a caller-allocated buffer. The real PoC deliverable is the adapter layer in between.

2. Per-feature gap (the PoC step-1 checklist)

Capability	quack	FlyDSL rmsnorm_kernel.py	Gap
forward	✅	✅ launch_rmsnorm	aligned
backward	✅ RMSNormFunction + rmsnorm_bwd (dx/dw/db, fp32 partials) (:1252,1368)	❌ none	biggest gap — see note below
weight (gamma)	optional	required Gamma	handle weight=None
bias	✅ optional	❌ no bias input	gap
residual / prenorm	✅ same entry, returns residual_out	✅ but a separate module launch_fused_add_rmsnorm(Input, ResidualIn, Gamma, Output, ResidualOut, ...) (:604)	adapter must route to the right module
rstd output	✅ store_rstd (needed by bwd)	❌ not emitted	bwd dependency
out_dtype / residual_dtype (mixed precision)	✅	⚠️ dtype baked via build_*_module(dtype_str), in==out	needs extension
eps	✅ runtime arg, default 1e-6	⚠️ compile-time constant EPS = 1e-5 (:25)	value mismatch and eps is not runtime-settable
per-head (B,S,H,D)	✅ (weight.dim()==2)	❌ 2D (M,N) only	defer for PoC, but record
fused quant (dynamic/smooth)	❌	✅ FlyDSL has extra build_rmsnorm_dynamicquant_module etc. (:966,979)	FlyDSL is ahead here; unused by PoC

On backward: this isn't a one-file omission — it's structural. The whole FlyDSL norm family is forward-only: layernorm_kernel.py has the same build_*/launch_* shape with no bwd, there are no *bwd*/*backward* files under kernels/, and FlyDSL's own tests/kernels/test_rmsnorm.py (965 lines) exercises forward only (no autograd). So "rmsnorm fwd+bwd, directly comparable to quack" overstates what exists today: forward is the part that's already there; bwd is net-new kernel work.

3. Reusing quack's test suite (PoC step 3) — two concrete blockers

quack's tests/test_rmsnorm.py parametrizes input_dtype ∈ {bf16, fp16, fp32}, weight_dtype ∈ {bf16, fp32, None}, M ∈ {1, 199, 8192}, several N, and use_compile ∈ {False, True}.

• It drives autograd (asserts on .grad) → needs FlyDSL bwd, which doesn't exist yet.
• use_compile=True needs a torch.compile-traceable op → the raw launcher isn't one.

Recommendation: scope the PoC to rmsnorm forward, Option-1 shim, gating quack's forward + use_compile=False subset; treat bwd and compile-compat as fast-follows.

4. Feeds back into the RFC's open questions

• Q4 (which PoC kernel): rmsnorm is still the right pick, but forward-only rmsnorm is the genuinely "already there" slice. Recommend stating the PoC as forward-first.
• Q2 (API stability): the surface worth stabilizing for downstreams isn't the low-level launcher — it's a thin functional adapter with runtime eps, auto-allocated output, torch.Tensor in/out, optional weight. That adapter's signature is what a downstream lib would actually depend on.
• Q3 / step 6 (what FlyDSL is missing): (1) runtime eps (not a compile-time constant), (2) optional rstd output from forward, (3) a functional entry that isn't pinned to a fixed N, (4) a bwd kernel, (5) a torch.compile-compatible op wrapper.

[jhinpan]

Spinning two concrete, actionable sub-issues out of this RFC:

• [Feature] Add rmsnorm backward (+ forward saves rstd) to enable training / downstream (quack) integration #769 — Add rmsnorm backward (+ have forward save rstd) to enable training and a training-capable downstream integration (quack). This addresses the forward-only gap surfaced in the interface analysis above.
• [Feature] Make autotune a safe, adopted path #770 — Make autotune a real, adopted tuning path, synthesizing FlyDSL's existing online @autotune with aiter's offline CSV + CI-regression approach (and aiter's existing ops/flydsl/gemm_tune path). Supersedes and closes the prior discussion-only Discussion: make autotune a supported and broadly usable tuning path #612.

Both are prerequisites/enablers for the downstream-backend contribution path proposed here.

[jhinpan]

Status update (Jul 22, 2026): downstream scope accepted; kernel prerequisites have landed

The earlier forward-only analysis is now stale. The state has materially changed:

Landed in FlyDSL

Training/backward series:

• [rmsnorm] Add backward pass + forward store_rstd for training (PR 1/3, #769) #795 added RMSNorm forward rstd, plain backward, and autograd.
• [rmsnorm] Add fused-add / residual backward (PR 2/3, #769) #800 added fused-add/residual backward.
• [layernorm] Add backward pass for training (PR 3/3, #769) #801 added LayerNorm backward and completed the planned [Feature] Add rmsnorm backward (+ forward saves rstd) to enable training / downstream (quack) integration #769 series.
• [Test] Expand RMSNorm backward configs, cache reuse, and Torch benchmarks #854 expanded backward coverage, cache-reuse checks, and reproducible Torch comparisons.
• [Perf] Optimize RMSNorm backward with staged reduction and vec8 I/O #855 replaced the large-M dweight-atomic bottleneck with a staged reduction/vec8 path. Its direct-wrapper measurements on MI355X are competitive with—and faster than—the tested PyTorch/AIter paths at the reported shapes.

So there is no longer a missing core RMSNorm backward kernel; #769 is complete.

Autotune series:

• [1/5] autotune: harden cache key + add restore_value (#770) #783 landed safe cache identity plus restore_value / reset_to_zero.
• [2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770) #785 landed explicit opt-in search, a zero-search default path, compiler-hint tuning, and the first RMSNorm adopter.
• [3/5] autotune: add offline config artifacts (#770) #786 (offline artifacts) and autotune: keep CI search opt-in (#770) #788 (keep broad CI search disabled while dedicated tests opt in) remain open under [Feature] Make autotune a safe, adopted path #770.

One important distinction: #785 currently tunes the explicit RMSNorm forward entry; #855's backward path uses a measured runtime selector, not the generic autotuner. That is sufficient to start the downstream PoC. Offline artifacts and broader backward tuning are not prerequisites for it.

Quack scope decision

In Dao-AILab/quack#178, @tridao confirmed that this is in scope: Quack can gate the backend to the features FlyDSL implements, provided existing users do not have to install extra dependencies. The suggested precedent is the JAX softmax support living in a separate file.

That resolves the main RFC question. The working boundary is now: a Quack-owned, opt-in FlyDSL module/entry point, without changing the existing CuTe path or imposing FlyDSL on current users.

What is still missing

1. Quack import/package isolation. Today quack/__init__.py eagerly imports quack.dsl, and pyproject.toml makes nvidia-cutlass-dsl unconditional. A ROCm entry must be importable without loading the CuTe/CUDA stack, while FlyDSL remains an opt-in dependency. This should be the first, kernel-free Quack PR.
2. A minimal Quack-owned adapter/kernel. Start with a separate experimental plain weighted RMSNorm forward+backward entry, flattenable contiguous inputs, runtime eps, and explicit feature gates. Do not route quack.rmsnorm transparently yet.
3. Mixed weight dtype parity. The landed FlyDSL wrapper currently requires weight.dtype == x.dtype; Quack supports FP32 weights with FP16/BF16 inputs. Supporting that initial contract requires a separate gamma dtype/load path in the Quack-owned FlyDSL implementation (or temporarily narrowing and documenting the gate).
4. torch.compile contract. Quack's RMSNorm tests exercise eager and torch.compile(fullgraph=True), while the raw FlyDSL launcher/autograd wrapper is not by itself the agreed Quack compile integration. We need to decide whether compile support is a first-merge requirement; if it is, the PoC needs a Quack-owned custom-op/compile wrapper. If not, the experimental entry should fail or gate clearly under compile until the follow-up lands.
5. Device validation and target correctness. Validate the agreed subset on gfx942 and gfx950, including non-default streams, cache reuse, and numerical gradients. [Runtime] Resolve the JIT target from the input device #878 remains the only FlyDSL-core correctness gap: invocation-device target resolution. An explicitly pinned, homogeneous target can unblock the PoC, but [Runtime] Resolve the JIT target from the input device #878 should land before mixed-architecture use is considered production-safe.
6. Downstream tuning ownership. Start from the measured [Perf] Optimize RMSNorm backward with staged reduction and vec8 I/O #855 backward selector and [2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770) #785's no-search/default behavior. Any Quack-specific search space, selected configs, and deployment artifacts should live in Quack; [3/5] autotune: add offline config artifacts (#770) #786/autotune: keep CI search opt-in (#770) #788 do not need to block the first implementation.

Suggested immediate sequence

1. Quack PR: isolate imports/dependencies and add a ROCm-safe module boundary.
2. Quack PR: add the experimental plain RMSNorm forward+backward path plus eager correctness tests on gfx942/gfx950.
3. Quack PR: add the agreed torch.compile contract, then consider transparent dispatch and Quack-owned tuning artifacts.

In short, the kernel capability question is resolved. The remaining work is the downstream package boundary, mixed-dtype/API adaptation, compile contract, and device validation.