[RFC] FlyDSL kernels in downstream libraries

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

The answer is now clear: FlyDSL should own the compiler/runtime, while the downstream library owns its production kernels and framework-facing operator contract. For the first Quack integration, no new backend abstraction is needed—one isolated experimental RMSNorm module is enough.

Current status (Jul 24, 2026)

Quack scope is accepted in Dao-AILab/quack#178: a FlyDSL path is in scope if it is feature-gated and existing users do not need to install FlyDSL.

There is now no required FlyDSL-core change blocking a constrained plain-RMSNorm PoC:

• #795 added forward rstd, plain backward, and autograd.
• #800 added fused-add/residual backward.
• #855 added the staged dweight reduction and runtime selector.
• #884 added the Quack-relevant mixed-weight contract: FP16/BF16 activations with FP32 weights, on both gfx942 and gfx950.
• #788 landed the CI rule that broad test runs do not search implicitly.
• #786 remains open for optional offline autotune artifacts; it is not a PoC prerequisite.
• #878 was closed as not planned. Mixed-architecture target resolution is therefore not a supported production claim; the PoC must use a homogeneous target plus explicit/per-device guards and caches.

The latest published FlyDSL wheel is 0.2.4 while main reports 0.3.0. The current plain mixed-weight forward/backward path imports and a focused case runs on MI355X/gfx950 with the 0.2.4 wheel, but Quack still needs a pinned minimum version and CI against the published wheel.

Ownership boundary

FlyDSL owns

• the DSL, compiler, and runtime;
• JIT/cache primitives, stream launch support, diagnostics, and generic device/runtime behavior;
• infrastructure demonstrated to be useful to more than one downstream.

Quack owns

• the production FlyDSL RMSNorm source used by Quack;
• the PyTorch API, validation, autograd, torch.compile integration, feature gates, and fallback/dispatch policy;
• Quack-specific tests, benchmarks, tuning choices, packaging, and CI.

The published flydsl wheel does not ship the repository's top-level kernels/ examples. Quack therefore cannot depend on kernels.norm.rmsnorm_kernel; it must adapt the minimal plain forward/backward code into Quack with Apache-2.0 attribution, as AITER does for its FlyDSL kernels. Quantized, fused-add, and other unrelated variants should not be copied into the first change.

Minimum first-merge contract

Use an explicit module such as quack.rmsnorm_flydsl; do not change quack.rmsnorm dispatch yet.

The first claimed subset should be:

• required 1-D weight;
• FP16, BF16, or FP32 input;
• matching input/weight dtype, or FP32 weight with FP16/BF16 input;
• caller-supplied eps (specializing/cache-keying by value is acceptable);
• positive-size, flattenable inputs with output shape and dtype matching the input;
• eager forward + backward and torch.compile(fullgraph=True) forward + backward;
• gfx942 and gfx950 only, under a homogeneous-process target contract.

Initially validate and claim model-relevant hidden sizes through 8192. Empty rows, larger hidden sizes, optional weight, cross-FP16/BF16 weights, bias, residual/prenorm, per-head mode, output dtype conversion, and weight_offset remain explicitly unsupported until tested. Unsupported combinations should fail clearly rather than silently selecting a different semantic path.

Remaining Quack work

1. Isolate imports before adding a kernel

Today quack/__init__.py eagerly imports quack.dsl and the CuTe RMSNorm module, and pyproject.toml makes nvidia-cutlass-dsl unconditional. Importing Quack without Cutlass currently fails before a FlyDSL submodule can be reached; pytest collection also enters the package through quack.testing.pytest_plugin.

The first Quack PR should be kernel-free:

• preserve the existing top-level API through lazy imports;
• make import quack and the ROCm test module collect without importing Cutlass/CUDA bindings;
• keep FlyDSL optional and add a clear availability/version check;
• add import-smoke tests with Cutlass absent and FlyDSL absent.

For the PoC, the least disruptive packaging choice is a FlyDSL/ROCm optional extra while leaving the current default CUDA dependencies unchanged. That preserves existing installs, but it does not produce a clean NVIDIA-free dependency graph. If a CUDA-free Quack installation is a first-merge requirement, use a separate adapter distribution or make a larger packaging decision; moving the current CUDA dependency to a new extra would break the existing default install contract.

2. Add the Quack-owned plain RMSNorm kernel and wrapper

Extract only the plain weighted forward/backward path and its minimal helpers. The wrapper must own:

• tensor/device/dtype/shape validation using real exceptions rather than optimization-removable assertions;
• output and rstd allocation;
• flatten/restore behavior;
• per-device compiled-launcher caches and current-stream dispatch;
• autograd saved state and dweight dtype semantics.

The mixed FP32-weight kernel work is already complete in FlyDSL #884; it is no longer an integration gap.

3. Make the adapter opaque to Dynamo

The raw FlyDSL Python JIT/launch path is not traceable by torch.compile(fullgraph=True); #596 established the working pattern. Quack should put preallocated-output forward and backward launches behind distinct torch.library.custom_op nodes with fake/meta behavior, then use its autograd wrapper around those nodes.

Quack's existing cute_op helper already implements the useful mutation-only fake plus eager-bypass pattern, but it currently lives under quack.dsl and cannot be imported without CuTe. Either move that small helper to a backend-neutral module and keep a compatibility re-export, or add the isolated equivalent in the FlyDSL module.

torch.compile should be part of the first kernel merge bar because Quack's RMSNorm contract already tests eager and fullgraph paths. If maintainers explicitly accept an eager-only staging PR, the experimental entry must reject/gate compiled use clearly rather than imply parity.

4. Make the target assumptions executable

Because #878 is closed, the initial supported deployment must be explicit:

• homogeneous gfx942 or homogeneous gfx950 per process;
• when overriding detection, keep ARCH and FLYDSL_GPU_ARCH consistent;
• cache compiled/loaded launchers per device and run on the calling tensor's current stream;
• do not claim mixed-architecture safety.

The Quack test matrix should cover non-default streams, warm-cache reuse, and two same-architecture devices where available.

5. Add downstream AMD CI

Quack's current GPU workflow covers H100/B300/SM120, not ROCm. Before the backend can be considered maintained, Quack needs AMD-owned jobs for gfx942 and gfx950 that install a released FlyDSL wheel and run only the isolated FlyDSL tests.

The minimum matrix should cover:

• eager and torch.compile(fullgraph=True);
• forward and numerical gradients for x and weight;
• FP16/BF16/FP32 plus FP32 weights for 16-bit inputs;
• aligned and tail hidden sizes, small/non-power-of-two/large row counts, and multiple eps values;
• import isolation, current-stream behavior, and cache reuse.

Suggested sequence

1. Quack PR: lazy import/package boundary plus optional-dependency smoke tests.
2. Quack PR: experimental rmsnorm_flydsl with Quack-owned kernel, custom-op/autograd adapter, and gfx942/gfx950 correctness CI.
3. Follow-ups only after that contract is stable: wider feature parity, transparent dispatch, Quack-owned tuning artifacts, and performance regression thresholds.

The RFC's compiler/kernel prerequisite phase is complete. The critical path is now entirely the downstream Quack integration and its CI/packaging contract.

Start with quack/init.py, pyproject.toml, and quack.testing.pytest_plugin to trace the eager Cutlass imports and pytest collection path. Run the import-smoke and ROCm collection tests with Cutlass and FlyDSL absent. The first step is complete when existing top-level imports remain compatible, optional dependencies are isolated, and both absent-dependency cases collect cleanly.