Kernel Layer — CUDA, JIT, and the Functional API¶

OASR exposes custom CUDA / CUTLASS kernels to Python through TVM-FFI JIT compilation, in the style of FlashInfer. Nothing is linked ahead of time in the default build: a kernel is compiled by nvcc on its first call and cached, so import oasr works on a machine with no compiled extension at all.

This document covers the C++/CUDA layer, the JIT pipeline, and the Python functional API. For the nn.Module layer that models are built from, see architecture.md § The layer waist. For a step-by-step walkthrough of adding a kernel, use the /add-cuda-kernel skill.

Layered design¶

Python functional API (oasr/<family>.py)  — @oasr_api decorated
    └── JIT generator (oasr/jit/<family>.py) → JitSpec / JinjaJitSpec
            └── TVM-FFI JIT binding (csrc/<family>_jit_binding.cu)
                    └── TVM-FFI launcher (csrc/<family>.cu)
                            └── Pure CUDA kernels (include/oasr/<family>.cuh)  — facade
                                    └── Config    (cutlass_*_configs.h)
                                    └── Template  (*_cutlass_template.h)
                                    └── Dispatch  (*_cutlass.h / *_dispatch.inc)

The C++/CUDA layer¶

Directory map¶

Path	Contents
`include/oasr/common/`	Shared types (`types.h`), scalar/vector dtype conversion (`vec_dtypes.h`), warp/block reduction (`reduction.h`), SM dispatch (`arch_dispatch.h`), per-SM hardware facts (`arch_facts.h`), epilogue functors, and math utilities; plus the Ampere-class CuTe toolkit the three CuTe families share (`cute_sm80.h`) and its CuTe-free tile rules (`tile_rules.h`) — see The C++ CuTe families
`include/oasr/activation.cuh`	Vectorized exact GELU, sigmoid, tanh, ReLU, GLU, Swish, and Swoosh activations; unary sigmoid/tanh/ReLU also consume regular padded row strides such as channel chunks without a copy
`include/oasr/norm.cuh` + `norm_dispatch.inc`	LayerNorm, RMSNorm, fused add+LayerNorm/RMSNorm (with optional residual passthrough), BatchNorm1d, GroupNorm, fused norm+activation
`include/oasr/conv/`	`conv1d.cuh` + `conv1d_dispatch.inc` (depthwise with asymmetric padding and optional FSMN mask/residual fusion, pointwise, causal); dense BTC Conv1D is the height-one specialization of the `conv2d.cuh` CUTLASS facade
`include/oasr/pooling.cuh`	BTC AvgPool1D and MaxPool1D; vectorized 2×2 specialization, generic padding/ceil/count semantics, and a flat narrow-channel path for one-channel traces
`include/oasr/recurrent/`	LSTM and tanh/ReLU RNN inference: fused GEMV/cohort kernels (`recurrent.cuh`), CUTLASS 2.x recurrent GEMM and state epilogues (`recurrent_cutlass.cuh`) with Stream-K/Split-K candidates, the CUTLASS 3.x TMA warp-specialized path for SM90/SM100 (`recurrent_cutlass_sm90.cuh`), and the CuTe fused recurrent step (`recurrent_step_*.h`)
`include/oasr/attention/`	The CuTe fused multi-head attention (`fmha_*.h`, `cutlass_fmha_configs.h`)
`include/oasr/mlp/`	The CuTe fused gated MLP (`gated_mlp_*.h`, `cutlass_gated_mlp_configs.h`)
`include/oasr/gemm/`	`gemm.cuh` facade, `bmm.cuh`, `group_gemm.cuh`
`include/oasr/quant/`, `include/oasr/scaled_mm/`	Per-token and static 8-bit activation quantization, int8 and fp8 e4m3 (`quantize.cuh`'s `QuantTraits`; pinned numerics, bit-exact with the torch twins) and the 8-bit scaled GEMM: CUTLASS 2.x int8 `mma.sync` (int32 accumulator) or Ada e4m3 `mma.sync` (fp32, fast accumulation, one summation order for every tile) with an Sm80 EVT epilogue that dequantizes (per-token × per-channel), adds the bias and applies the activation — in the three-header layout, one JIT module per format. On sm90 / sm100 the native lanes (`*_sm90.h`): CUTLASS 3.x collective builders over TMA and wgmma (ping-pong or cooperative warpgroups) or tcgen05 (1-SM / 2-SM, a double-buffered tensor-memory accumulator), with the same epilogue as an Sm90 EVT tree whose activation is a runtime argument (one kernel per tile and output dtype, not five) and whose activation-scale leaf is a dynamic per-token / per-tensor broadcast; one JIT module per format and lane, int8 on sm100 built for `sm_100a` (`sm_100f` has no tcgen05 `kind::i8`). The element functors both generations share are `scaled_mm_functors.h`; the launcher checks, `csrc/scaled_mm_launcher.h`. See quantization.md
`include/oasr/scaled_conv/`	The int8 implicit-GEMM forward convolution of a quantized `Conv2d` / `Conv1d` (NHWC int8 input, KRSC int8 filter, exact int32 accumulator): CUTLASS 2.x has no EVT convolution, so it is `DefaultConv2dFpropWithBroadcast` with `ScaledConvOutputOp` -- the per-channel `input_scale * weight_scale` as the broadcast vector, the bias as a zero-stride source, the activation fused. JIT module `scaled_conv` (`oasr/jit/scaled_conv.py`), functional `oasr.scaled_conv2d` / `scaled_conv1d`. See quantization.md
`include/oasr/{softmax,topk,fft,features}.cuh`, `sort/`	The remaining families
`include/oasr/ctc_decoder.cuh`, `include/oasr/wfst/`	GPU decoder kernels
`csrc/<family>.cu`	TVM-FFI launcher
`csrc/<family>_jit_binding.cu`	JIT binding exports
`csrc/tvm_ffi_utils.h`	DLPack dtype dispatch, validation macros (`CHECK_GEMM_ALIGNMENT`, `CHECK_CONTIGUOUS_INPUT`, `CHECK_LAST_DIM_CONTIGUOUS_INPUT`, `CHECK_VECTOR_ALIGNED`, `FLATTENED_ROWS`), `OptionalDataPtr`
`csrc/templates/`	Jinja2 templates for config-specific CUTLASS instantiations (`gemm_cutlass_template.cu.jinja`, `bmm_cutlass_template.cu.jinja`, `group_gemm_cutlass_template.cu.jinja`, `scaled_mm_cutlass_template.cu.jinja`, `scaled_conv2d_cutlass_template.cu.jinja`)
`csrc/decoder/ctc/`	GPU CTC launcher + binding (`ctc_decoder.cu`, `ctc_decoder_jit_binding.cu`); `ctc/cpu/` holds the CPU-side C++ decoders compiled into `_C.so` — greedy search, prefix beam search, WFST beam search (via k2), the streaming WFST decoder, `ContextGraph` for phrase boosting, and shared `common/utils`
`csrc/decoder/wfst/`	In-tree GPU WFST decoder (TVM-FFI JIT). Its exact-semantics CPU reference oracle is test-only and lives separately under `csrc/tests/wfst/`, out of the production decoder library.
`csrc/pybind/`	pybind11 module for the CPU decoder bindings, the alignment bindings and legacy enums (`pybind_main.cpp`, `pybind_decoder.h`, `pybind_alignment.h`)
`csrc/alignment/`	The post-decode alignment pass in C++ — emission frames → per-token spans → words, plus `extract_beam_tokens` / `extract_beam_row`, which turn a padded `[batch, beam, max_len]` host tensor into nested lists. Not a kernel: plain host data shuffling, compiled into `_C.so`
`csrc/tokenizers/`	The rendering half of the `symbol_table` tokenizer kind (`token_pieces`), which the word grouping calls once per finished hypothesis

The GPU CTC launcher/binding pair is the one that does not live at the csrc/ root (its JIT generator is oasr/jit/ctc_decoder.py); everything else follows the flat convention.

The three-header CUTLASS pattern¶

Each CUTLASS kernel family splits config, template, and dispatch:

Header	Purpose	Example (GEMM)
`cutlass_*_configs.h`	Config structs (`GemmConfig`), per-SM MMA traits (`SmMMATraits`), default configs (`DefaultGemmConfig`)	`gemm/cutlass_gemm_configs.h`
`*_cutlass_template.h`	CUTLASS kernel template parameterized by Config + MMATraits	`gemm/gemm_cutlass_template.h`
`*_cutlass.h`	Public dispatch interface (JIT mode via `OASR_TARGET_SM`, AOT mode via `OASR_DISPATCH_SM`)	`gemm/gemm_cutlass.h`

Non-CUTLASS kernels (Conv1D, Norm, Activation) use *_dispatch.inc files with VecSize / block_size dispatch macros instead.

The three CuTe families (attention, the gated MLP, the recurrent step) keep the same three roles but split the template into CUTLASS 3.x-style collectives: the config is cutlass_*_configs.h, the template is a mainloop + epilogue composed by an arch-free kernel shell, and the dispatch is *_launch_template.h — see The C++ CuTe families.

Dispatch modes¶

Kernel family	Mode	Config source	Source generation
GEMM, BMM, GroupGEMM	jinja	`cutlass_gemm_configs.h`	Jinja renders `.cu` with baked-in config
Dense Conv1D / Conv2D	jinja	`cutlass_conv2d_configs.h`	Jinja renders `.cu` with baked-in config; each tactic exports strict BTC/KSC Conv1D and NHWC/KRSC Conv2D entry points. SM75–89 and SM120 use the CUTLASS 2.x implicit GEMM; SM90/SM100 use the 3.x `conv::CollectiveBuilder` — see Conv2D on SM90/SM100 for why that config is not a copy of the GEMM one
Depthwise / causal Conv1D	dispatch	`conv1d_dispatch.inc`	Direct compilation, VecSize macro
Grouped / depthwise Conv2D	direct	`grouped_conv2d.cuh`	NHWC 3×3/7×7 specializations; bias and optional activation share the convolution launch
Norm	dispatch	`norm_dispatch.inc`	Direct compilation, block/vec macro
Activation	dispatch	`activation_dispatch.inc`	Direct compilation, VecSize macro
Pooling	direct	`pooling.cuh`	128-bit channel vectors in BTC layout; specialized 2×2, generic, and narrow-channel launches
Recurrent	direct + CUTLASS	`recurrent/recurrent.cuh` + `recurrent/recurrent_cutlass{,_sm90}.cuh`	fused low-latency GEMV at small batch, shared-weight batch warps for cohorts, sequence-wide input projection, state epilogues, autotuned Stream-K/Split-K for wide states, and TMA warp-specialized collectives on SM90/SM100
Fused attention, gated MLP, recurrent step (C++ CuTe)	jinja (per variant)	`cutlass_{fmha,gated_mlp,recurrent_step}_configs.h` + the CuTe-free tile headers	Jinja renders one `.cu` per variant of a cell; the tile is derived (attention) or chosen from a compiled table (MLP, recurrent step) — see The C++ CuTe families

JIT mode (OASR_TARGET_SM defined): a single SM instantiation, with an optional JitGemmConfig / JitConv2dConfig passed via -D flags.
AOT mode (no OASR_TARGET_SM): the OASR_DISPATCH_SM macro switches on the runtime SM version.

SM targets default to 75, 80, 86, 89, 90, 100, 120 in CMakeLists.txt; setup.py queries the host GPUs and falls back to 80;86;89;90. Override either with CUDA_ARCHITECTURES.

Conv2D on SM90 / SM100¶

The CUTLASS 3.x implicit-GEMM conv builder is close enough to the GEMM one to invite a copy, and far enough away that the copy does not compile. Four differences, each load-bearing:

	GEMM	Conv
K mode of `TileShape`	flat `Int<BK>`	nested `Shape<Int<BK>>` — implicit GEMM's K axis is the filter's (C, S, R) modes
Mainloop schedule	pingpong / cooperative	`conv::KernelScheduleAuto`; SM90 conv has one schedule (CUTLASS's own auto-selector has the cooperative branch commented out, and `conv/dispatch_policy.hpp` `static_assert`s on the persistent tags)
2-SM SM100 atom	M tile doubled	M tile passed as-is; the atom comes from the cluster, and `BM * 2` is `"Invalid TileShape_M."`
`Arguments`	leading `GemmUniversalMode`	the mode-less constructor — SM90's `ConvUniversal` inherits `GemmUniversal` and has the mode, SM100's does not

Two tile rules bound the space, both measured by compiling the grid rather than documented upstream, and both enforced as filters in oasr/jit/conv.py because a single unbuildable variant fails the whole JIT module:

SM90 — (BM + BN) * BK * 2 ≤ 106,496 B. StageCountAutoCarveout has to fit two stages plus the epilogue into Hopper's 227 KiB; above that it resolves to zero stages and the build stops.
SM100 — a 256-row tile requires cluster_m == 2.

The tile ladders themselves are unmeasured: they span the M range at two N widths. They are also sized against build cost — each 3.x conv translation unit peaks near 3.7 GB in cicc, and run_ninja only bounds parallelism when MAX_JOBS is set.

Recurrent execution paths¶

The CUDA FP16/BF16 recurrent waist retains two complementary implementations:

recurrent/recurrent.cuh owns the latency path. A CTA computes a complete state row at batch 1, while the cohort specialization stages one output row's weights in shared memory for reuse by several batch warps. Affine accumulation, bias, gate activation, and state writes share a launch.
recurrent/recurrent_cutlass.cuh owns the throughput path. The input affine is one sequence-wide OASR GEMM. LSTM weights are cached in [hidden, gate, K] order, so the four gates for a cell are adjacent; the recurrent CUTLASS epilogue can apply i/f/g/o and write h/c directly. Vanilla RNN uses the common tanh/ReLU CUTLASS epilogue and writes the hidden state directly. Matrix-C and output leading dimensions are independent, so BTC input projections feed the recurrence as strided batch rows without a transpose or intermediate copy.

The recurrent autotuner compares 16/32/64-row M tiles, Stream-K, and parallel Split-K. LSTM also exposes serial Split-K: its custom epilogue delays the nonlinear state transition until the final K partition and reuses GEMM's self-restoring semaphore path, avoiding a workspace clear per timestep. Parallel Split-K and Stream-K materialize one interleaved LSTM gate tile before the state finalizer; this costs one extra launch but is mathematically safe and can expose more SM parallelism for a thin M and large K. Vanilla RNN excludes serial Split-K because applying tanh/ReLU to an intermediate K partition is incorrect.

The new state is never written over the old. At one timestep a single launch reads initial_h while other CTAs of it already store final_h, so an in-place call (final_h=initial_h) returns half-updated state: on an RTX 5090 the scalar LSTM and RNN, the tensor-core LSTM and the fused step all did. Every recurrent launcher therefore refuses an output that shares storage with an input (oasr::StorageOverlaps in csrc/tvm_ffi_utils.h, exact for row-strided views, so two column halves of one buffer are fine). Outputs may share storage with each other: the single-timestep cell ring is final_c. The CuTeDSL step checks nothing on the host, so the fused path leaves it for the decomposed one when the state overlaps.

The direct kernel remains the default for decode-sized and moderate hidden states. This is intentional: a tensor-core GEMM can underfill the device at small M, and packing/projection plus dependent GEMM launches may cost more than the work saved. benchmarks/kernels/recurrent.py exposes native, cutlass16, cutlass32, cutlass64, streamk, splitk, and serial_splitk (LSTM only) arms so architecture-specific crossover changes are measured instead of inferred. The focused matrix is benchmarks/testlists/recurrent_tactics.txt.

Architecture mapping. Every target gets a CUTLASS 2.x composition, and SM90/SM100 additionally get a 3.x one.

recurrent_cutlass.cuh holds the 2.x side. For FP16/BF16 that API specialises exactly two arch tags, so every target maps onto one of them:

JIT target	CUTLASS tag	MMA	Stages
75 (Turing)	`Sm75`	`m16n8k8`	2
80, 86, 89, 90, 100, 103, 120	`Sm80`	`m16n8k16`	3

Turing needs its own row twice over: a narrower MMA, and a kernel::DefaultGemm that is specialised for a two-stage pipeline and no other stage count. Nothing between the two rows is usable — Sm86 has no 2.x tensor-op specialisation at all, and Sm89/Sm90 have one whose DefaultGemmConfiguration covers FP8 only — so Ada, Hopper and Blackwell all run the SM80 composition, which is forward compatible and keeps one epilogue output-thread mapping across them. Two static_asserts pin the arch/stage and arch/instruction-shape pairings, so a remap that would issue an instruction the target lacks fails at compile time instead of at decode.

recurrent_cutlass_sm90.cuh adds the CUTLASS 3.x collective path — TMA plus wgmma on Hopper, tcgen05 on Blackwell datacenter — as two extra tactics, tma_64 (id 6) and tma_128 (id 7). It is compiled only for targets 90 and 100, the two whose 3.x OpClassTensorOp builders accept FP16/BF16; SM120's is restricted to F8/F6/F4, and no CutlassArch entry exists for 103. Everywhere else the two ids are refused rather than rerouted, and the autotuner does not offer them. Hopper's cooperative schedule static_asserts on an M tile below 128 rows, so the 64-row tile takes the pingpong schedule; SM100 selects from kSMs and ignores the flag, which is why one pair of configs covers both.

On the 3.x path the LSTM is decomposed: the collective GEMM materialises one gate-interleaved tile and the existing finalizer applies the state transition. The fused custom epilogue cannot come along — it reconstructs logical coordinates from a PredicatedTileIterator thread map, and 3.x replaced that with cute layouts and an epilogue visitor tree, where a four-gates-to-one-cell column reduction is not an elementwise node. The vanilla RNN has one gate, so its nonlinearity stays fused in the collective epilogue.

The layer's automatic tensor-core selection still requires compute capability 8.0 (oasr/layers/recurrent.py), because the crossover was measured on Ampere and later; on Turing the path is reachable through the functional API and the autotuner.

Conventions¶

The output tensor is the first parameter of every TVM-FFI launcher.
Launchers take N-D tensors and flatten with FLATTENED_ROWS rather than making Python call reshape(-1, K).
Output allocation stays in Python (new_empty varargs) — allocating in the C++ launcher was measured and is slower.
Every GEMM-family launcher enforces the CUTLASS alignment-8 rule uniformly via CHECK_GEMM_ALIGNMENT, with a message naming the fix.
One convolution launch addresses at most 2 GiB per tensor. CUTLASS's implicit-GEMM convolutions use 32-bit byte offsets and refuse anything larger; the launchers say so via CHECK_CONV_LAUNCH_BYTES. Every oasr.conv* entry point therefore splits an oversized batch into launches over row slices (_row_slices in oasr/functionals/conv.py) — exact, because convolution rows are independent — before choosing a backend, and refuses a single row that is already too large. The in-tree grouped, depthwise and depthwise-1D kernels offset whole rows in 64 bits and positions within one in 32, and their launchers refuse a row past that; 64-bit arithmetic in the inner loop measured 1.8–4.6% slower on depthwise shapes.

The JIT pipeline¶

Module	Role
`oasr/jit/core.py`	`JitSpec` (static sources) and `JinjaJitSpec` (Jinja-rendered), `gen_jit_spec()`, `gen_jinja_jit_spec()`, `build_and_load()`
`oasr/jit/templates.py`	Jinja2 rendering (`get_template_env()`, `render_template()`)
`oasr/jit/env.py`	Path constants (`OASR_TEMPLATE_DIR`, `OASR_JIT_DIR`), `cutlass_version_stamp`
`oasr/jit/cpp_ext.py`	The compile and link commands (`build_commands`), the `build.ninja` made from them, and their part of the cache key (`build_identity`)
`oasr/jit/<family>.py`	Per-family generators: `gemm`, `conv`, `norm`, `activation`, `pooling`, `recurrent`, `softmax`, `topk`, `fft`, `features`, `ctc_decoder`, `wfst_decoder`
`oasr/jit/fmha.py`	Fused attention, C++ CUTLASS/CuTe lane — one module per `(sm, dtype, head_dim)` cell, holding all 17 feature variants
`oasr/jit/gated_mlp.py`, `oasr/jit/recurrent_step.py`	The gated-MLP and recurrent-step C++ CUTLASS/CuTe lanes — one module per `(sm, dtype, activation)` / `(sm, dtype, kind)` cell
`oasr/jit/arch_facts.py`	The Python mirror of `include/oasr/common/{arch_facts,tile_rules}.h` the three CuTe lanes' tile mirrors are built on
`oasr/jit/attention.py`	The backend arbiter, and the CuTeDSL lane — different model, see below
`oasr/compilation_context.py`	`CompilationContext` detects GPU SMs at import time; pass `supported_major_versions=[...]` to `get_nvcc_flags_list()` for arch-restricted kernels

Compiled modules are cached in ~/.cache/oasr/jit/<name>/<key>/ (OASR_JIT_DIR). The key covers what is compiled, how, and with which tools:

What: the sources (a rendered one by its text), the include/ tree and the csrc/ headers by relative path and content, and the CUTLASS version stamp.
How: the expanded compile and link commands, whatever put each flag there: the module, OASR_EXTRA_{CFLAGS,CUDAFLAGS,LDFLAGS}, CC, CXX, OASR_NVCC or the arch list.
Which tools: nvcc's build and the host compilers' and TVM-FFI's versions.

Where the checkout sits is not in the key. Its root is replaced by a placeholder in the commands, so worktrees with the same sources share libraries. The tools' own paths stay in, because they select the tools.

A rendered source stays in memory until its module is built. The build writes it into the key's directory, under gen/, so the file a build compiles is always the text its key hashed. Renderers used to write a shared path named after the variant, and another checkout rendering different text could replace that file between the hash and the compile.

The key has no device in it. A process compiles for one architecture: the current device's, or the one OASR_CUDA_ARCH_LIST names. It keeps one module per family, so ASREngine refuses a device of another architecture at construction (jit.core.require_device_family). On a host with mixed GPUs, run one process per architecture and select its GPUs with CUDA_VISIBLE_DEVICES.

A module is built under a file lock to <name>.so.partial and renamed into place only once the link has finished (cubin_loader.locked_compile), so the cached name never refers to a partly written library. That matters more than it looks: dlopen of a truncated library does not raise, it kills the process with SIGBUS. A process that finds a module mid-build therefore waits on the lock instead of loading it, and a build killed mid-link leaves only the .partial, which the next build discards. Before loading, build_and_load also checks that every section and segment the ELF headers describe is in the file (is_complete_shared_object) and rebuilds a library that fails.

A failed build is not repeated: - In one process: once a module's build fails, later calls raise oasr.jit.errors.JitBuildError with the original error, without building. The failure belongs to the key, meaning the same sources, flags and tools. The module getters are functools.cache-wrapped, and that caches results but not exceptions, so each call used to pay for the whole compile again. oasr.jit.core.forget_failed_builds() re-arms the build once the cause is fixed, and clear_cache() does the same. - Across processes: a process that was waiting on the build lock when the build failed gets that failure too, instead of running the same build after it. The builder records its failure in <name>.so.failed before releasing the lock. A process that arrives after the failure builds again, so a fixed cause is picked up. - While waiting: a waiter logs every 30 s, naming the process that holds the lock. After lock_timeout (600 s) it raises JitBuildLockTimeout. A timeout is not a build failure and is not remembered.

Nothing removes old keys automatically. Every change to a header, a flag or the toolchain leaves the previous key behind, and so does every one-off module. oasr jit prune removes them explicitly:

oasr jit prune --dry-run                 # list what would go
oasr jit prune --unused-days 30          # remove keys nothing has loaded in 30 days
oasr jit prune --unused-days 7 --keep 0  # also let a whole unused module go

How the prune decides: - When a key was last used: every load touches last_used in its key's directory. A key built before that stamp existed has only its library's build time to go on. - What it always keeps: each module's --keep most recently used keys (default 1), however old. - What it never touches: a key whose build holds its lock. - When to run it: when no OASR process is starting up. A process that has found a library but not yet loaded it would fail if the prune removed it in between. A process that already loaded a library keeps running on it.

The default flags (core._default_cuda_cflags) include -DNDEBUG, which is how CUTLASS's own release builds compile. Without it, CuTe's device-side assert()s become __assertfail calls inside the GEMM mainloop, and ptxas serialises every wgmma across a function call. A build log shows that as C7510: wgmma pipeline crossing function boundary; every SM90 GEMM kernel had it. OASR's own sources contain no assert().

Fused attention has two kernel lanes and one arbiter over them.

oasr/jit/attention.py is the arbiter: select_backend(), set_backend_mode(), fmha_config_supported(...), fmha_backend_for(...), warmup_fmha(...). It is also the CuTeDSL lane — a functools.cache-keyed wrapper around cutlass.cute.compile() rather than a Ninja JIT spec, which is why it sits apart from the table above. select_backend() probes the device capability eagerly at module load and resolves on sm_80 / 86 / 89 / 120, otherwise "sdpa".

oasr/jit/fmha.py is the C++ lane and is an ordinary Ninja spec. Jinja renders one translation unit per feature variant and ninja builds them into one .so per cell = (target_sm, dtype, padded_head_dim) — the axes that change the shared-memory layouts and so cannot share a binary. Inside a cell:

3 (none / causal / local) × 2 (bias) × 2 (dense / paged)   = 12
+ 4 split-KV (unmasked only) + 1 static binding             = 17

Everything else is runtime, not a compile axis: head counts, page size, window bounds, cache_seqlens / cache_seqstarts, packed-vs-dense input, and whether the bias is vectorisable. nvcc parallelises across translation units but not within one, so a cell costs about one variant's wall time on a many-core box and then covers every shape that cell will ever be asked for.

select_backend() answers a mode question and cannot see the shape. fmha_backend_for(...) answers the rest — whether the preferred lane can serve this call, and what to degrade to when it cannot. The case that needs it is the sliding window, which the C++ lane has and the CuTeDSL one has no argument for.

CUTLASS¶

CUTLASS is the 3rdparty/cutlass git submodule, pinned to v4.6.1. git submodule update --init is what provides it — CMake fetches only pybind11. Nothing links it: every CUTLASS kernel is JIT-compiled, so oasr/jit/env.py hands the include directories to nvcc at runtime.

Its version.h is folded into the JIT cache key. Without that, build_and_load short-circuits on an existing .so and a submodule bump keeps silently loading binaries built against the old headers. Editing a vendored CUTLASS header without bumping the version still needs rm -rf ~/.cache/oasr/jit. The stale libraries also reach every other checkout that shares the cache, so give an experiment on vendored headers its own OASR_JIT_DIR.

The CuTeDSL half of CUTLASS is the separate nvidia-cutlass-dsl wheel (pip install -e .[attention], floor in oasr/jit/attention.py::MIN_CUTEDSL_VERSION), kept at the same 4.6.1 release. It is optional: OASR_ATTN_BACKEND=auto degrades oasr.fmha to SDPA when it is absent.

Evaluation of the 4.4.2 → 4.6.1 move: .artifacts/cutlass_upgrade.md.

The Python functional API¶

Every entry point is @oasr_api-decorated (oasr/api_logging.py — debug logging plus exception context), JIT-compiles on first call via @functools.cache, allocates its output tensor, and calls into the compiled module.

Module	Exposes
`oasr/functionals/gemm.py`	`gemm`, `bmm`, `group_gemm`, and the fused epilogues `gemm_activation` (RELU/tanh-GELU/exact-erf GELU/SWISH) and `gemm_log_softmax` (the CTC head fast path)
`oasr/functionals/gemm_torch.py`	Torch/cuBLAS runners — `torch_gemm`, `torch_gemm_activation`, `torch_bmm`, `torch_gemm_log_softmax` — mirroring the CUTLASS launcher contract exactly (output-first, in-place / CUDA-graph-safe, `D = A @ Bᵀ`). Doubles as a `Tactic("torch")` autotuner candidate and as the production dispatch target. Deliberately free of any `oasr.tune` import.
`oasr/functionals/norm.py`	`layer_norm`, `rms_norm`, `batch_norm1d`, `group_norm`, fused norm+activation
`oasr/functionals/conv.py`	dense / depthwise / pointwise / causal Conv1D; dense, grouped and depthwise Conv2D. Dense NHWC 1×1 Conv2D dispatches as GEMM; Conv1D depthwise padding may be an integer or `(left, right)` pair
`oasr/functionals/activation.py`	standalone exact-erf `gelu`, `sigmoid`, `tanh`, `relu`, `glu`, `swish`, `swoosh_l`, `swoosh_r`
`oasr/functionals/pooling.py`	BTC/TC `avg_pool1d` and `max_pool1d`, including symmetric padding and ceil mode (`count_include_pad` for avg; `dilation` and `return_indices` are refused, not ignored)
`oasr/functionals/softmax.py`, `oasr/functionals/topk.py`, `oasr/functionals/fft.py`	`softmax`, `log_softmax`, `masked_softmax` (one pass over an attention score tensor: an additive bias and two boolean masks, each broadcast against the scores through its own strides, so a shifted `as_strided` relative-position window or a `[..., ::ds]` mask slice is consumed where it is), `topk`, `rfft` / `rfft_power`
`oasr/functionals/feature.py`	`stft_frame`, `dct_lifter`, `fbank_preprocess`, `mel_log`, `whisper_logmel`, `lfr_gather` — see features.md
`oasr/functionals/attention.py`	`fmha(...)` and `fmha.persistent_inputs(...)`
`oasr/functionals/mlp.py`	`gated_mlp` — a whole SwiGLU/GeGLU gate+up+activation+multiply in one dual-B GEMM, plus `gated_mlp_available` (the routing question, so a layer can ask before building anything) and `gated_mlp_backend` (which of the two lanes answered it). Refuses rather than falling back
`oasr/functionals/ctc_decode.py`	`ctc_beam_search_decode`, `GpuStreamingDecoder` — see ctc_decoder_gpu.md
`oasr/decode.py`	Thin helpers over the CPU-side `oasr.decoder` decoders

oasr/decoder/ holds the Python wrappers for the CPU-side C++ decoders — CtcGreedySearch, CtcPrefixBeamSearch, CtcWfstBeamSearch (requires k2), and ContextGraph (a phrase-boosting trie) — plus the k2_available flag. Each lazily imports the compiled _C extension and delegates to a _*Core C++ object.

LayerNorm with a per-token int8 output¶

layer_norm_quant and add_layer_norm_residual_quant (oasr/functionals/norm.py) are LayerNorm and the fused add + LayerNorm with residual passthrough, emitting the normalized output as per-token int8 (or fp8 e4m3, by the output tensor's dtype -- a runtime flag in the store, not a second instantiation) plus one fp32 scale per row for an 8-bit GEMM consumer (quantization.md). That replaces two launches (the norm storing T, then quantize_per_token reading it back) with one. The contract is bit-exactness with that pair, torch.equal on q, scale and residual_out:

layer_norm_quant(x, w, b, eps, qeps)  ==  quantize_per_token(layer_norm(x, w, b, eps), None, qeps)

One kernel template, two sinks. layerNormKernel and addLayerNormResidualKernel take a kQuantOut parameter. The mean and variance passes, the block size and the vector/scalar choice are shared. Only the normalize pass's store differs. The plain kernels compile to the same instructions as before.
What is quantized is the T value. Each normalized value is rounded to the input dtype exactly as the store would round it. The row's amax is taken over those values, and quant::quantizeInt8 (include/oasr/quant/quantize.cuh) quantizes the held values rather than recomputing them. A thread's first vector stays in registers. Past 1024 vectors a row, the rest are staged in dynamic shared memory, with an opt-in past 40 KiB.
The launcher's vector/scalar choice must be LayerNorm's, decided from the input (and residual) pointers alone. It fixes each thread's accumulation order and hence the bits of the mean. The int8 tensor's alignment only decides whether the int8 store is 8 bytes wide.
The affine's rounding is spelled out in source. nvcc contracts n * w + b into an FMA only when it unswitches the bias != nullptr branch, and that is a per-instantiation decision. The plain LayerNorm contracted, the first kQuantOut build did not, and the two then differed by one ulp in a few elements per million. Each kernel now writes the rounding it compiled to before the change: an fmaf in layerNormKernel and in the scalar addLayerNormResidualKernel, __fadd_rn in its vector path. tests/kernels/test_norm_quant.py::TestBitExactVsUnfused::test_volume_catches_a_one_ulp_drift is the config that sees such a drift every time; a rows=1000 grid sees it in about one config in seventy.

In the layer, LayerNorm.quantize_output is armed by oasr.quantization.bind_quantization, never by hand. When armed, forward returns a QuantizedActivation, and so does the first element of forward_add_residual. That happens only where the fused kernel applies: CUDA, fp16/bf16, a contiguous input (and residual), and the oasr backend. Everywhere else the layer returns the plain tensor. reference_layer_norm_quant and reference_add_layer_norm_residual_quant are the torch twins: bit-exact in the quantize, within one int8 step overall, because torch's LayerNorm reduction is not the kernel's. RMSNorm has no quant-out variant yet; the same kQuantOut templating applies to rmsNormKernel / addRmsNormResidualKernel unchanged.

Shape-aware backend selection¶

gemm, gemm_activation, bmm and gemm_log_softmax route per shape. jit.gemm.select_default_config(op, M, N, K) picks:

a CUTLASS variant — default tile, serial split-K, parallel split-K (pk), or Stream-K;
on SM90 and SM100, from either of two lanes. The native one is CUTLASS 3.x: TMA plus wgmma (SM90) or tcgen05 (SM100). It includes small 64-row and 64-wide tiles at K-tile 64 with 1×1 clusters. The other is the CUTLASS 2.x mma.sync space, with its split-K and Stream-K, rendered with the Sm80 arch tag (jit.gemm._MIXED_LANE_SMS). Only the GEMM family takes both lanes; BMM and grouped GEMM stay native (is_native_lane). CUTLASS's 3.x Stream-K scheduler was built and measured, and it won nowhere: its deterministic split-K ran about 1.3× behind the unsplit kernel;
the torch/cuBLAS backend (oasr/functionals/gemm_torch.py);
or, for the CTC head only, the legacy single-call fused launcher.

The rules are data — the tuning database (oasr/tune/database.py, format and policy in autotuning.md). One file per (arch family, kernel family): the shipped system tier (oasr/tune/db/sm120/gemm.json, …/conv1d.json), and a user tier in ~/.cache/oasr/tune/v2/ written by tuning runs on the machine that measured them, which wins per signature. An entry is keyed by the static signature (gemm|op=gemm|dt=half|N=512|K=256) and its ascending (m_hi, config) regions are looked up by rounding M up, so a config only serves sizes at or below the largest one it was measured at.

The lookup chain is user → system → cost model → GEMM_DEFAULT, each tier counted (oasr.tune.database.tier_counts(), format_gap_report()). The model tier ranks the compiled configs with the analytic model the arch's tuning file ships (oasr/tune/cost_model.py) for an aligned gemm/gemm_activation shape no entry covers, and only replaces the default when it predicts at least a 5% gain (OASR_TUNE_MODEL_FALLBACK=0 disables it). A shape with no entry is still a rule miss — counted by jit.gemm.rule_miss_report(), exported for re-tuning by oasr tune export-misses.

Per architecture. jit.gemm._GEMM_HEURISTIC_RULES is a view of the system tier keyed by SM family (and jit.conv._CONV1D_HEURISTIC_RULES likewise); an arch with no file is tuned by oasr tune build on that card, not by an edit to the selector. Only SM120 is measured today.

An architecture with no table is not an error — GEMM_DEFAULT computes the right answer — but it is the largest gap the heuristic can have, because it is every shape rather than one width, and it is the one gap that used to be invisible: the arch fall-through returned before recording anything, so rule_miss_report() printed "every shape this process issued had a tuned rule" on a box that had never opened the table. It is now counted per architecture (not per shape, which would name every GEMM the process issued) and reported by both rule_miss_report() and oasr.layers.format_gap_report():

GEMM heuristic inactive on sm80: no tuned rule table (tuned: sm120), so all 8
shape lookup(s) used GEMM_DEFAULT. Tune this card with scripts/tune_asr_gemm.py.

The production module compiles the tuned subset. get_production_configs(sm) is what the tuning DB references for the arch (both tiers), the default and a small coverage basis — 20 of 35 variants on sm120, 9 of 43 on sm80; the tuning modules (gemm_tune, bmm_tune, group_gemm_tune) compile the rest of the space (get_unique_compile_configs minus production) and are only built by a process that tunes or sweeps. Never a variant in both: two libraries instantiating one CUTLASS kernel share its one-time initialisation (STB_GNU_UNIQUE statics), so the second skips its own shared-memory opt-in and fails to launch. Each family's production set is therefore frozen at first use (jit.gemm._production_set) — a tuning-DB reload cannot move it under a loaded library, and a config the reload adds is compiled by the tuning module instead. OASR_GEMM_COMPILE_SET=all restores the old single-module behaviour.

What a failed launch does. A variant whose can_implement refuses the problem raises oasr.jit.errors.VariantUnsupportedError before anything launches (CHECK_GEMM_STATUS, csrc/tvm_ffi_utils.h). The GEMM, BMM and Conv1D dispatch answer it with the default variant, remember it per shape and log it once. Every other failure — a CUDA error, a failed shared-memory opt-in, initialize() — is a plain RuntimeError and reaches the caller: relaunching the default over it hid a sticky CUDA error until a later call tripped over it, and a variant that could never launch until someone profiled. The candidate space itself is one generator per MMA lane: every CUTLASS 2.x family (sm_75/80/86/89/120) gets the lane's tiles at every pipeline depth its template can take and its shared memory fits — which is how sm_86/sm_89, with sm_120's 100 KiB budget, came to be offered the 4-stage tiles sm_120's table uses.

Adding a table is a measurement. Rules that were reasoned about rather than timed have shipped a 4.6x regression and an empty transcript here; the entries a new table may name are held to its own architecture's emitted variant set by tests/kernels/test_gemm_heuristic.py::TestPerArchRuleRegistry.

Three rules are structural rather than tuned:

GEMM_MIN_ROWS — a row floor below which CUTLASS's M-tiling leaves most of every tile empty and cuBLAS's GEMV-shaped kernel wins.
The candidate space is constrained before it is tuned (jit.gemm._epilogue_covers_warp). CUTLASS's tensor-op epilogue folds a warp's 32 lanes into a grid derived from the tile, and asserts nothing about that grid covering a warp; a tile where it does not compiles, launches at full speed and returns wrong numbers. At the 8-element access width half precision uses, that is every block_n < 32 tile. Such a tile is refused by the per-SM builders — for GEMM, BMM, grouped GEMM and Conv2D alike, since all four render from the same list — and the refusal is readable through jit.gemm.rejected_tiles(). This is not a tuning preference: a rule naming such a tile is a correctness bug that neither dispatch nor timing can see, and one shipped (.artifacts/gemm_thin_n_tile_epilogue.md).
The dispatch decision is a pure function of the call and is deliberately not relaxed under CUDA-graph capture, even though dispatch cost is free there: a capture-dependent branch makes the graph pick a different kernel than eager, and the resulting one-ulp fp16 difference has produced different tokens.

OASR_GEMM_HEURISTIC=0 disables the whole thing. Tuning a card: oasr tune census + oasr tune build (autotuning.md).

Batch-invariant GEMM (`batch_invariant=True`)¶

Every choice above depends on M, and with it the order a row is summed in (cuBLAS's own kernel and split, split-K, Stream-K, and F.linear below GEMM_MIN_ROWS). The result is an ulp of drift per row as the batch changes. An int8 quantizer downstream can turn that drift into a changed word (quantization.md). gemm(..., batch_invariant=True) and gemm_activation(..., batch_invariant=True) pin the call: every output row is bitwise independent of M. They run only the CUTLASS 2.x data-parallel variants that jit.gemm.is_batch_invariant defines: split_k == 1, no Stream-K, no parallel split-K, not sliced-K (warp_k == block_k), at the shared 64-wide K tile. In these variants one thread owns each output element and walks all of K in the same mma.sync order, whatever the tile, stage count or grid. So every member of the class returns the same bits for a row, and the selector may still pick the tile from M (tests/kernels/test_gemm_batch_invariant.py).

jit.gemm.select_batch_invariant_config picks the config:

It keeps a measured rule's choice when that choice is already a member.
Otherwise it applies an occupancy rule: among grids that fit in one wave, the smallest warp tile; past one wave, the least operand traffic.
It is a pure function of the shape and the SM count.
It never picks cuBLAS, split-K, Stream-K or the autotuner, even with tuning enabled.
It falls back only to a GEMM_DEFAULT that is itself a member, and raises if no member is compiled.

On sm_120 at embed.out (N=256, K=4864), the pin costs 2.7-3.0× at M=64 and 2.1× at 256. Each CTA walks all 4,864 of K alone, which takes about 15 µs. It costs 1.25× at 1024 and 1.02-1.20× from 2048 to 8192.

oasr.layers.Linear and LinearActivation expose the pin as batch_invariant = False, set per instance. A pinned projection takes the pinned call at every M on the CUDA fp16/bf16 path, row floor included. CPU, fp32 and OASR_LAYERS_BACKEND=torch keep their routes, and quant_method runs first.

The BMM general lane¶

bmm has a second lane that the shape rules do not reach. The rendered tile variants are alignment-8 iterators over contiguous 3-D operands; a decomposed attention block does not have that shape. Zipformer's is the case that forced the lane: five products per layer, 4-D permuted views of a (time, batch, head, dim) activation, one operand broadcast over the request batch, head dims of 32 / 4 / 12 and a relative-position extent that is always odd. FMHA cannot substitute — the probabilities are materialized once and shared by SelfAttention and NonlinAttention, and a fused kernel never materializes them.

oasr.bmm therefore accepts one or two broadcasting batch dimensions, arbitrary N and K, and — the part that saves a copy at every call site — a B operand contiguous along either trailing axis, so [..., N, K] and its transposed view are both legal and neither needs .contiguous(). A contiguous 3-D alignment-8 call still takes the tuned lane; everything else lands in include/oasr/gemm/bmm.cuh, which turns three run-time facts into one instantiation:

Chosen from	Values	Why it is compile-time
B's memory layout	CUTLASS `ColumnMajor` / `RowMajor`	the iterator's contiguous axis
operand alignment along K	8 / 4 / 2 elements	`cp.async` cannot issue a 2-byte copy, so alignment 1 is not a tensor-op case at all
alignment along N (epilogue, and B when RowMajor)	8 / 4 / 2 / 1	an odd N has to be able to store one element at a time
threadblock tile	128×128 / 64×64 / 32×32 / 64×16 (thin-N)	see below

The grid is 124 instantiations (62 per dtype), rendered one translation unit per (layout, dtype) from csrc/templates/bmm_general_template.cu.jinja — the same template-per-configuration pattern as the tile variants, and for a measured reason: nvcc parallelizes across translation units but not within one, so the module's cold build is set by its largest TU. Inlining the grid into bmm.cu costs 112 s; this split costs 40 s.

Anything outside that grid runs on CUTLASS SIMT GemmBatched, which constrains nothing. Nothing falls back to PyTorch: a shape either has a kernel or the call raises.

Two properties of the lane are worth knowing before changing it:

The tile ladder is the difference between the lane and a regression. These problems are 2–70 MFLOP, so they are latency-bound and the only thing that matters is how much of the device the grid covers. A single 128×128 tile put 8–64 CTAs on an RTX 5090's 170 SMs and burned 246 registers/thread for 8.3% achieved occupancy — 1.74× slower than cuBLAS. Selecting the largest tile that still fills one wave brought that to 1.07×. Note the ceiling: CUTLASS's tensor-op epilogue divides the tile's rows by the warp count in M before it computes an iteration count, so a 32-row tile is a single-warp shape and nothing smaller than 64 rows can use 128 threads. The same ceiling has a column half, and that one is not an assert: the epilogue derives its lane width from the tile's columns and its lane rows from the tile's rows independently, and never checks that the two cover a warp — so the 16-column tile is addressable at 4 elements per access and silently wrong at 8. epilogueAlignment in bmm.cuh caps the store width for that tile alone; the cp.async load keeps its 8 elements.
GemmBatched advances every operand by one constant stride, so two batch axes are one launch only when all three tensors are affine in the flattened index. Both flattening orders are tried, because a contiguous output satisfies one and a head-major view of a (time, batch, head, dim) activation satisfies the other. A broadcast inner axis satisfies neither and costs min(batch0, batch1) launches.

Measurements, including the profile that produced the ladder: .artifacts/kg5_strided_bmm.md.

Fused attention¶

oasr.fmha(q, k, v, *, softmax_scale, attn_bias, cache_seqlens, cache_seqstarts,
          block_table, causal, window_left, window_right, out, backend)

Three cache modes share one signature:

Mode	`block_table`	`cache_seqlens`
Offline	`None`	`None`
Dense streaming (caller concatenated old + new K/V)	`None`	set
Paged streaming (K/V are pool views)	set	required

Three backends share that signature — cxx (C++ CUTLASS/CuTe), cute (CuTeDSL) and sdpa (PyTorch, fp32-friendly). OASR_ATTN_BACKEND selects process-wide; the per-call backend= overrides it without touching the global, which matters because set_backend_mode() clears three compile caches and an in-process A/B driven through the environment variable recompiles the CuTeDSL kernel on every flip. validate=False skips checks for proven inputs.

Naming a backend requires it: a request it cannot serve raises rather than quietly computing something else. backend=None (the default) is the only mode that may degrade, and it degrades through fmha_backend_for(...), which asks whether the preferred lane can serve this shape rather than only which lane the process prefers.

Two capabilities exist on the cxx lane only:

Capability	Why not on `cute`
`window_left` / `window_right` — a per-row sliding window, top-left aligned	its `local` axis is the per-stream `[seqstart_k, seqlen_k)` pair, a different thing with no argument for this
causal and windowed varlen, and varlen at all without a second kernel	packed input there is a separate ~600-line `kernel_varlen`; here it is the same instantiation with a zero batch stride

The cxx lane also splits the K range when the grid cannot fill the machine (flash-decoding). num_splits is a pure function of (shape, SM count) — never of CUDA-graph capture state, because splitting changes the order the fp32 partials are summed (rule 11). The workspace is allocated in Python, from torch's caching allocator, deliberately not from oasr::getCachedWorkspace, which branches on cudaStreamIsCapturing.

The kernel moves q, k, v and the output in 128-bit vectors along the head dimension: - Every operand: the base must be 16-byte aligned, and the stride of every axis longer than one must be a whole vector. oasr.fmha and oasr.fmha_varlen copy an operand that breaks this. - head_dim: it must be a multiple of 8, or 16 for an fp8 KV cache. A backend left to choose picks another lane for any other head_dim; naming cxx refuses it. - Integer inputs: the wrappers convert lengths, cu_seqlens and block tables to int32, and packed-bias offsets to int64. - Direct callers: the launchers check all of this again (csrc/fmha_launcher.h), along with dtypes and shapes. Calling the module directly therefore raises an error instead of stopping the CUDA context with misaligned address.

Packed-bias offsets are int64 because the bias buffer holds sum_s H * T_s^2 elements, which passes 2**31 at 256 segments of 1024 frames and 8 heads. The CuTeDSL varlen kernel still indexes the buffer with 32-bit indices. A larger packed bias therefore runs on cxx under auto, and backend="cute" refuses it. A packed bias may start at any element; when its base is not 4-byte aligned, the kernel reads it element by element instead of in pairs.

Routing policy and measurements: .artifacts/fmha_tuning.md (CuTeDSL) and .artifacts/fmha_cpp_validation.md (the C++ lane's falsifications and the A/B).

The C++ CuTe families¶

Three kernel families are written in C++ CUTLASS/CuTe rather than CUTLASS 2.x device templates: fused attention, the fused gated MLP and the fused recurrent step. They share one layout, file for file, so a reader who knows one knows where everything is in the other two:

Role	attention	gated MLP	recurrent step
Which SMs are served, and the family's ceilings	`cutlass_fmha_configs.h`	`cutlass_gated_mlp_configs.h`	`cutlass_recurrent_step_configs.h`
Tile choice, CuTe-free	`cutlass_fmha_configs.h` (derived)	`gated_mlp_tiles.h` (table)	`recurrent_step_tiles.h` (ladder)
Arguments / params	`fmha_params.h`	`gated_mlp_params.h`	`recurrent_step_params.h`
Family-only CuTe helpers	`fmha_utils.h`, plus softmax, mask, bias, block, seqlen, paged-KV, combine	`gated_mlp_utils.h` (the dual-B gemm)	—
Mainloop collective	`fmha_mainloop_sm80.h`	`gated_mlp_mainloop_sm80.h`	`recurrent_step_mainloop_sm80.h`
Epilogue collective	`fmha_epilogue.h`	`gated_mlp_epilogue.h`	`recurrent_step_epilogue.h`
Kernel shell (arch-free)	`fmha_kernel.h` + `fmha_tile_scheduler.h`	`gated_mlp_kernel.h`	`recurrent_step_kernel.h`
Launch template (the one file naming an arch)	`fmha_launch_template.h`	`gated_mlp_launch_template.h`	`recurrent_step_launch_template.h`
TVM-FFI launcher, one TU per variant	`csrc/templates/fmha{,_split,_kv8}_template.cu.jinja`, checks in `csrc/fmha_launcher.h`	`gated_mlp_template.cu.jinja`	`recurrent_step_template.cu.jinja`
Capability oracle (the static TU)	`csrc/fmha_jit_binding.cu`	`csrc/gated_mlp_jit_binding.cu`	`csrc/recurrent_step_jit_binding.cu`
JIT generator + Python tile mirror	`oasr/jit/fmha.py`	`oasr/jit/gated_mlp.py`	`oasr/jit/recurrent_step.py`
Arbiter over the C++ and CuTeDSL lanes	`oasr/jit/attention.py`	`oasr/jit/mlp.py`	`oasr/jit/recurrent_cute.py`

Names follow the files: Collective<Family>MainloopSm80, Collective<Family>Epilogue, <Family>Kernel, <Family>Arch<SM>, run_<family> — so a profile's kernel name says which family it is.

What the three have in common lives in include/oasr/common/, not in three copies:

cute_sm80.h — the Ampere-class toolkit: ArchAmpere<SM>, the traits base every <Family>Arch<SM> derives from (the mma.sync atom, the ZFILL cp.async, ldmatrix, the budgets and the capability bools); the swizzled shared-memory atom; the accumulator re-views; the packed down-conversion; the zero-filling (copy_zfill) and skipping (copy_if) tiled copies; the two warp-level gemm shapes with their refill hook; and launch_kernel, the one launcher. The copies this replaced had drifted: the gated MLP found that branching around an asynchronous copy costs three dead LDS per copy and that attention's down-conversion returned a view of a dead local, and neither fix had reached the other family.
tile_rules.h — the CuTe-free half: the swizzle-row rule, the warp tiling and tiled-copy pass geometry every tile validator composes, and the occupancy estimate. CuTe-free so the capability oracles build in seconds; mirrored in Python by oasr/jit/arch_facts.py.

What stays in a family directory is what is actually that family's: the online softmax, the masks, the bias read, the paged gather, the split-KV combine, the dual-B gemm, the recurrent state transition, and each family's decision of which architectures it serves.

Three rules hold for all of them. Dispatch on ArchTraits::kIsWarpSpecialized / kHasTma, never on Tag::kMinComputeCapability >= 90 (true for sm_120, which has neither). Every tile choice is a pure function of shape and machine, never of CUDA-graph capture state (rule 11). And a new architecture is a new <Family>Arch<SM>, a new collective and one conditional_t arm in the launch template — nothing else moves.

The C++ attention kernel (`include/oasr/attention/`)¶

Fifteen headers, structured the way CUTLASS 3.x structures a collective: CollectiveFmhaMainloopSm80 + CollectiveFmhaEpilogue + SingleTileScheduler, composed by FmhaKernel — a shell that names no architecture and touches no layout. fmha_launch_template.h is the one file that names one, so a Hopper or Blackwell lane is a second FmhaArch<SM>, a second collective and one more conditional_t arm; the arguments, the FFI signature and the whole Python side do not move.

Two rules in that header are load-bearing rather than stylistic:

ArchTag selects instructions; a separate kIsSm86Or89 selects tuning. cutlass::arch::Sm86 exists, but the collective is the same one sm_80 uses.
Never branch on ArchTag::kMinComputeCapability >= 90. Sm120's is 120, so FlashAttention's Use_TMA_O test is true on a part with no TMA at all. FmhaArch<SM>::kHasTma / ::kIsWarpSpecialized say what they mean.

fmha_softmax.h's header comment carries the numerics contract — twelve numbered items, each one a place where FlashAttention's reference is wrong for OASR's masking semantics, with the failure each one prevents.

Two implementation notes, because nothing in the C++ says them:

Every load is branch-free but one. Q, every deeper K/V tile and the paged gather fold their predicates into the ZFILL cp.async (cute_sm80::copy_zfill); the branching form had left 144–168 dead LDS and a BSSY/BSYNC pair per copy in every variant. The exception is the dense first K tile, issued once per CTA: its branches keep ptxas from hoisting later stages' address math above the Q and first-K requests, which is the critical path of a CTA that walks one or two K tiles (copy_rows_or_clear records the measurement).
The bias is read at its point of use, in the right shape. Rows past the plane are branched around (they are register loads, so there is no asynchronous copy to order against) and boundary tiles read column pairs. A one-tile-ahead register prefetch was built and measured: it bought 3–4% on streaming chunks and cost up to 6% on interior 2–4-tile CTAs, including a realistic Conformer batch, so it is not shipped (fmha_bias.h).

FP8 KV variants (`fmha_kv8.h`)¶

An autoregressive decoder's KV cache can be stored as float8_e4m3fn with one fp32 scale per (token, kv-head) (EngineConfig.quantization.kv_cache = "fp8_e4m3"; see cache_manager.md). oasr.fmha reads such a cache when k/v are fp8 and k_scale/v_scale are given:

Dequantized in shared memory, inside the tile loop. mma.sync has no fp16 x fp8 shape and quantizing Q to use the e4m3 MMA would spend the precision the format keeps, so the cp.async ring carries the fp8 bytes (and each row's scale, in the same commit group) and every stage is converted into one served-dtype tile right before the gemm that reads it -- K, then V, one tile for both. Each thread converts the 16-byte chunks it issued itself, so only the scales need the barrier the ring already takes; the conversion costs one extra __syncthreads per gemm. Nothing past the conversion changes.
Bit-identical to the served-dtype kernel on the dequantized cache. The scales are powers of two (include/oasr/kv_quant/kv_quant.cuh), so q * scale rounded to fp16/bf16 is exact, and fmhaResolveTileKv8 walks the same ladder as fmhaResolveTile, so a shipped head dim lands on the same K tile (the split count and the online-softmax tiling both follow it). tests/kernels/ test_fmha_kv8.py holds the lane to that at atol=0, dense and paged.
Own cell, own variants. fmha_kv8_sm{SM}_{dtype}_d{D} -- dense / paged x unsplit / split, rendered from csrc/templates/fmha_kv8_template.cu.jinja -- so a process that never reads fp8 never compiles it. Rendered for unmasked, bias-free reads only (a decode step, a cross-attention read): a prefill attends its own fresh K/V in the served dtype, so causal / windowed / biased fp8 reads have no caller, and asking for one raises with the reason (oasr.jit.attention.fmha_kv8_refusal), as does a padded head dim the 16-byte gather cannot tile exactly (it covers padded 32 / 64 / 128 / 256).
Routing is a declared rule (fmha_kv8_backend): fp8 goes to this lane under auto and cxx; OASR_ATTN_BACKEND=cute raises (that lane cannot read fp8) and sdpa runs the dequantizing reference, which is also what CPU and an fp32 query get. The engine refuses kv_cache under cute / sdpa.

Measured on an RTX 5090, graph-replayed decode steps (median of 15 interleaved rounds): Whisper-large-v3 cross-attention, B16 x H20 x 1500 keys, 51 us vs 82 us fp16 (1.61x); Qwen2-Audio-7B self-attention, B16 x H32 x ~800 keys, d128, bf16, 114 us vs 120-126 us (1.05-1.10x); at B4 the same step is 34 us vs 26-27 us (0.78x) -- with little K/V to read, the in-smem conversion is the critical path. The quantize-and-scatter that writes the cache (oasr.functionals.kv_quant.quantize_kv_store) is one launch for K and V per layer, replacing the two index_put_ the fp16 path issues.

The C++ gated-MLP kernel (`include/oasr/mlp/`)¶

The same collective decomposition, one family over: CollectiveGatedMlpMainloopSm80 + CollectiveGatedMlpEpilogue, composed by GatedMlpKernel, with gated_mlp_launch_template.h the one file that names an architecture. gated_mlp_tiles.h is deliberately CuTe-free — it holds the CTA tiles and the arithmetic that picks one, as plain constexpr integers, so csrc/gated_mlp_jit_binding.cu can export them as a cheap capability oracle and tests/kernels/test_gated_mlp_cpp.py can hold the Python mirror to them over four architectures from one box.

Where it differs from the attention family, and why:

	attention	gated MLP
tile	resolved: `fmhaResolveTile` derives it from the architecture, because the question is "what fits"	tabled: the question is a wave count, so it depends on `rows`, `N` and the SM count, and all six tiles are compiled as siblings in one module
cell	`(sm, dtype, head_dim)` — the axes that change the smem layouts	`(sm, dtype, activation)` — the activation is the one axis that changes the arithmetic, and a checkpoint means one of them
variants in a cell	3 masks x 2 bias x 2 paged = 12	6 tiles x 2 bias = 12

Two things it does that the CuTeDSL lane cannot:

Capability	Why not on `cute`
a K residue — `K` need not be a whole number of K tiles	there the mainloop predicates only the row axis, so a partial K tile reads the next row of `x`; `gated_mlp_shape_supported` has to refuse it
arbitrary row strides on every operand	there the compiled signature marks the tensors compact, so a row-slice of a wider buffer has to be copied first

One implementation note is worth carrying, because nothing in the C++ says it. The K residue is predicated into the ZFILL cp.async's own src_size operand (cute_sm80::copy_zfill), not into control flow. Written the obvious way — if (pred) copy(...) else clear(...) — ptxas has to order a synchronous STS against an asynchronous LDGSTS to the same shared address, and it does that by bracketing every copy in BSSY/BSYNC and padding it with three dead @!PT LDS RZ, [RZ]. Measured on the 64x64x32 tile: the K loop's load section went from ~60 instructions to ~12, the LSU pipe from 2.34M instructions to 1.2M, and the kernel from 0.89x of the CuTeDSL lane to ahead of it everywhere.

The C++ fused recurrent step (`include/oasr/recurrent/recurrent_step_*.h`)¶

The same collective decomposition again, and the second lane for the kernel oasr/kernels/cute/recurrent/step.py already implements in CuTeDSL: CollectiveRecurrentStepMainloopSm80 + CollectiveRecurrentStepEpilogue, composed by RecurrentStepKernel, with recurrent_step_launch_template.h the one file that names an architecture. recurrent_step_tiles.h is CuTe-free, so csrc/recurrent_step_jit_binding.cu exports it as a capability oracle and tests/kernels/test_recurrent_cpp.py holds the Python mirror to it over six architectures from one box.

One fused timestep, with the gate dimension interleaved so column n is (hidden n / G, gate n % G):

gates[m, n] = sum_k previous_h[m, k] * weight_hh[n, k] + input_gates[m, n]
c[m, i]     = sigmoid(g1) * previous_c[m, i] + sigmoid(g0) * tanh(g2)
h[m, i]     = sigmoid(g3) * tanh(c[m, i])

The interleaving is what keeps the epilogue inside one CTA tile — a hidden unit's gates are adjacent columns, so no cross-tile reduction is needed, and recurrentStepTileValid refuses any block_n that would straddle one.

	gated MLP	recurrent step
tile	wave count: is the last wave still saturating DRAM?	tabled ladder on `(hidden, batch)`: a dependent launch has no next kernel to overlap its tail with, so what the measurements found was a boundary between two regimes, not a wave count
cell	`(sm, dtype, activation)`	`(sm, dtype, kind)` — `lstm` / `rnn_tanh` / `rnn_relu`
variants in a cell	6 tiles x 2 bias = 12	8 tiles

The ladder is inherited from the CuTeDSL lane rather than re-derived, so both lanes run the same tile on every shape and an A/B between them compares kernels. test_the_ladder_reproduces_the_cutedsl_lane pins that.

Three things it does that the CuTeDSL lane cannot:

Capability	Why not on `cute`
a K residue — the hidden width need not be a whole number of K tiles	there the mainloop loops `ceil_div(K, k_block)` and predicates only the row axis, so a partial K tile reads past the end of both operands
arbitrary row strides on every operand	there the compiled signature marks the tensors compact, so a row-slice has to be copied first
a staging buffer larger than the ring	there the epilogue aliases the ring; here it is a `union`, so `can_implement` needs no such clause

Two epilogue details are worth carrying, because nothing in the C++ says them:

The global loads are issued before the mainloop. The transition needs input_gates and previous_c from global memory and the affine from shared; the first two depend on nothing the K loop computes, so the kernel shell issues them first (load_inputs) and their latency hides behind the whole loop — 1.04–1.10x on the latency-bound steps. The two 128x128 tiles, where the carried values would be 16–20 registers against a 512-thread CTA's 128-register cap, load them after the loop instead (kPrefetchInputs).
The accumulator staging is padded by 8 floats per row, not 1. Both avoid the MMA-C store's bank conflict; only 8 also keeps every row 16-byte aligned, so a cell's four gates come back in one LDS.128 instead of four LDS.32. stride % 32 == 8 puts a 16-lane STS.64 phase on 32 distinct banks and an 8-lane LDS.128 phase likewise. cuobjdump -sass confirms both.

One declared limitation: the widest one-gate variant (tile 7, 32 epilogue slots on a 512-thread CTA) sits at REG:128 with a 32-byte stack frame, where every LSTM variant is at REG:54-120 with no spill. Nothing reaches it — a vanilla RNN is never routed under auto, and tile 7 needs hidden > 1536 and batch > 128 — so it is declared rather than fixed. The header comment records the two causes that were tested and refuted, so a later attempt starts past them.

Routing is oasr/jit/recurrent_cute.py — the arbiter, and the CuTeDSL lane's compile cache. OASR_RECURRENT_CUTE decides whether to fuse; OASR_RECURRENT_BACKEND decides which lane. Measurements: .artifacts/recurrent_cpp_validation.md.

CuteDSL kernels (`oasr/kernels/`)¶

oasr/kernels/ holds low-level implementations that do not use the TVM-FFI / Ninja pipeline.

kernels/cute/attention/base.py — abstract FmhaBase + pick_arch_cls(major, minor)
kernels/cute/attention/fmha_sm80.py — FmhaSm80, the mainloop (sm_80: A100, A30)
kernels/cute/attention/fmha_sm{86,89,120}.py — thin subclasses for Ampere consumer, Ada and consumer Blackwell.
kernels/cute/recurrent/step.py — RecurrentStepCute, one fused LSTM/RNN timestep as a tensor-core GEMM with the state transition in the epilogue
kernels/cute/mlp/gated.py — GatedMlpCute, a dual-B GEMM: one A tile in shared memory feeding two mma.sync chains against two B tiles, with activation(gate) * up applied to the FP32 accumulators. The shape of CUTLASS's examples/45_dual_gemm, and the reason a gated MLP needs no intermediate tensor
kernels/cute/ — FlashAttention-style helpers: block_info.py, seqlen_info.py, mask.py, softmax.py, tile_scheduler.py, pack_gqa.py, paged_kv.py, named_barrier.py, copy_utils.py, layout_utils.py, ampere_helpers.py, utils.py

Each is compiled via cutlass.cute.compile() into a Python callable and cached per config — oasr/jit/attention.py::_compiled_fmha, oasr/jit/recurrent_cute.py::_compiled_step, oasr/jit/mlp.py::_compiled_gated_mlp — always with options="--enable-tvm-ffi", which is what lets the callable take torch tensors directly and be captured into a CUDA graph.

Each of the three also owns a routing module beside its compile cache, because a fused kernel that is faster on some shapes and slower on others has to say which: OASR_ATTN_BACKEND, OASR_RECURRENT_CUTE, OASR_GATED_MLP_CUTE, each with the measured band in its module docstring. For the gated MLP the band is one m-tile (M <= 64): with one m-tile every weight element is read from DRAM once, which is the bandwidth argument the fusion rests on; with two it is an ordinary GEMM reading its operands twice and cuBLAS wins. The tile inside the band is chosen by N rather than by M — see select_gated_mlp_tile, and the ten lines of it that exist because a rows-keyed table lost 10% at one model width.

Attention and the gated MLP each have two fused lanes, cute and cxx, and they spell the choice differently on purpose. Attention folds both questions into OASR_ATTN_BACKEND, because its fallback (sdpa) is itself a backend and sits in the same list. The gated MLP cannot: its fallback is two GEMMs, which is not a lane of this kernel at all, so "fuse or not" (OASR_GATED_MLP_CUTE) and "which lane" (OASR_GATED_MLP_BACKEND, auto / cute / cxx) are separate switches. Collapsing them would make a rollback and an A/B the same knob, and they are different decisions: the first is an operator's call about a shape, the second is a claim about two implementations of the same thing.

auto resolves to cxx for the MLP — 1.00x-1.11x of the CuTeDSL lane over four interleaved, graph-replayed sweeps with no row regressing, plus a strictly wider shape contract — and to cute for attention, where the C++ lane's paged mode still reads 0.85-0.94x.

Which routing decisions travel, and which are extrapolations¶

The routing splits cleanly in two, and only one half follows the card it runs on:

	reads the machine	example
derived	yes	`gated_mlp_ctas_per_sm` takes `multi_processor_count`, the real opt-in shared memory and `max_threads_per_multi_processor`, and `select_gated_mlp_tile` does wave arithmetic against them; `selectBmmTile` takes `getDeviceMultiProcessorCount()`
measured	no	`_LSTM_BANDS`'s `(hidden, batch)` cut-offs, the `_TILES` ranking, the gated-MLP candidate list — timed once and written down

A measured cut-off is neither wrong nor a guess; it is the best available estimate. But the crossover it encodes is a function of SM count and memory bandwidth, and the supported set — sm_80, sm_86, sm_89, sm_120 — spans an A30 (56 SMs, 933 GB/s), an A100 (108, 1555), an L40S (142, 864) and an RTX 5090 (170, 1792). All four got the same numbers, with nothing saying whose they were.

oasr/jit/measured.py is where each table now names the GPU it was timed on, the note holding the protocol, and what moves the crossover. Off that card the table still applies — which side wins at each extreme is a property of the algorithm, only the boundary between them is a property of the machine — but the fact is logged once and reported by oasr.layers.format_gap_report():

  routing tables measured on another GPU (applied as an extrapolation):
    jit.recurrent_cute._LSTM_BANDS / _TILES
        measured on NVIDIA GeForce RTX 5090 (sm_120, 170 SMs), 1792 GB/s; running on NVIDIA A30 (sm_80, 56 SMs)

Identity is (compute capability, SM count), both queried exactly. Bandwidth is deliberately not computed from memory_clock_rate × bus_width: that formula is right for HBM and GDDR6 and wrong for GDDR7, so on an RTX 50-series card it would report the measured machine as an extrapolation of itself.

Utilities¶

oasr/utils/:

Module	Contents
`validation.py`	`@supported_compute_capability([80, 86, ...])` marks a check function with the SMs it supports; `@backend_requirement(backend_checks={...}, common_check=fn)` wires validation into the public API function and adds `.is_backend_supported()` / `.is_compute_capability_supported()` helpers
`mappings.py`	dtype and enum helpers
`timer.py`	timing helpers

oasr/testing/bench_gpu_time(fn, args, ...) is the measurement primitive: CUDA event timing with an optional CUPTI fallback via triton.testing.do_bench, returning (median_s, std_s).

Ahead-of-time compilation¶

oasr/aot.py registers every kernel family for AOT builds, including gen_all_gemm_variants() for systematic variant enumeration. AOT is optional — the default path is JIT-on-first-call.

Autotuning¶

oasr/tune/ produces what the shape-aware selection reads: the census of shapes a deployment runs, one benchmark protocol (oasr/tune/bench.py), the cost model, the region partitioner and the tuning database, driven by oasr tune census|build|evaluate|diff|status and by the oasr.autotune() context manager, whose GEMM winners land in the user tier. See autotuning.md.

Adding a kernel family¶

Seven steps, in order:

Kernel header in include/oasr/<family>.cuh
TVM-FFI launcher in csrc/<family>.cu
TVM-FFI JIT binding in csrc/<family>_jit_binding.cu
JIT generator in oasr/jit/<family>.py
Python functional API in oasr/<family>.py
nn.Module wrapper in oasr/layers/<family>.py
AOT registration in oasr/aot.py

The /add-cuda-kernel skill (.claude/skills/add-cuda-kernel/SKILL.md) walks through each with worked code. /benchmark-kernel covers measuring the result.

pybind11 bindings (csrc/pybind/) remain for the work that is not a kernel: the CPU-side CTC/WFST decoders and the post-decode alignment pass. New kernels do not use them — TVM-FFI JIT is the route for anything that runs on the device.

The distinction is about where the work runs, not how fast it is. csrc/alignment/ holds no CUDA at all; it is there because the pass runs on the engine's step-loop thread, which holds the GIL for every request the engine finishes, and in Python it cost more than the CTC decode it decorated. The beam read-back is the same story on the untimed path: out_lengths[b, k] is a 0-d tensor plus an item(), and a slice is another tensor, so materialising a 16-beam block row by row cost more than the decode's own device→host copy.

Neither has a Python twin in the package — oasr/engine/decode/alignment.py is marshalling only, because a fallback here is a slow path a deployment lands on silently. Both files are in OASR_SOURCES, so a successful build always has them and no call site checks. So these are the one part of the tree where test-cpu.yml, which compiles nothing, cannot cover the implementation: the rule is checked against a Python oracle kept inside tests/decoders/test_alignment.py (exact agreement over randomised input and the whole Unicode plane), and that file skips without the extension.