Cutedsl Tunable Kernel Template

Tunable CuTeDSL Kernel Template

Use this after cutedsl-kernel-template establishes the core kernel. This skill owns the Attention Gym adapter and parent/compiler-process boundary. Load cutedsl-performance separately to design the search space, measure candidates, and accept a selector.

The canonical Attention Gym helpers are:

from attn_gym._backends.cute import benchmark_gpu, compile_tvm_ffi, jit_cache, run_tunable, tune

If another repo has equivalent helpers, reuse them. Do not rebuild cache/process orchestration inside each op.

Ownership Boundary

Keep each value in one domain:

Never pass real tensors to compile(...) or compiler workers. compile_call(...) is the explicit projection from runtime inputs to static compile arguments.

Minimal Kernel Convention

Use a module-scope NamedTuple or frozen dataclass for codegen choices. Module scope makes the config importable and pickleable by fresh compiler processes.

When warps have named protocol responsibilities, represent them with a module-scope IntEnum such as WarpRole.TMA_PRODUCER; do not compare warp indices with unexplained integer literals. Name the responsibility precisely when one warp performs more than one role.

Make approximation choices such as fastmath explicit compile-time arguments. Default them to False unless the public numerical contract deliberately chooses approximate math, encode them in cache and profiler names, and correctness-test every exposed mode.

Put ordinary specialization defaults directly in the owning constructor or public function signature. Avoid module constants that only alias those defaults or a one-kernel policy limit; keep fixed compile-time expressions next to the device code or schedule that consumes them. Derive one-use schedule values in the owning op rather than adding free helper functions.

Do not repeat dtype assertions inside a CuTeDSL entrypoint when the cached compiler boundary already constructs an exact fake-tensor ABI and the runtime operator validates that ABI. Keep small helpers only when they are reused or mark a real protocol, cache, or runtime-ABI boundary; use established integer utilities such as ceildiv directly instead of wrapping one arithmetic expression.

from typing import NamedTuple

from attn_gym._backends.cute.target import CompileTarget


class MyConfig(NamedTuple):
    threads: int
    tile_size: int

A tunable adapter supplies seven distinct responsibilities. When its launch ABI has several values, make that ABI an Args type owned by the adapter rather than a separate module-level private type:

class MyOp:
    class Args(NamedTuple):
        q: torch.Tensor
        k: torch.Tensor
        output: torch.Tensor
        workspace: torch.Tensor

    @staticmethod
    def default_config(args: Args, *, target: CompileTarget) -> MyConfig:
        """Return a deterministic valid config without benchmarking."""
        return MyConfig(threads=128, tile_size=256)

    @staticmethod
    def tuning_key(args: Args, *, target: CompileTarget) -> tuple[int]:
        """Identify workloads that may safely reuse one benchmark winner."""
        return (args.q.shape[1],)

    @staticmethod
    def configs(args: Args) -> tuple[MyConfig, ...]:
        """Return valid candidates derived from actual runtime inputs."""
        ...

    @staticmethod
    @jit_cache
    def compile(*static_args):
        """Construct the fake ABI and compile one specialization."""
        ...

    @staticmethod
    def compile_call(config: MyConfig, args: Args) -> tuple:
        """Project launch arguments to static arguments for compile(...)."""
        ...

    @staticmethod
    def launch(compiled, config: MyConfig, args: Args):
        """Launch in the parent with real tensors and return the public result."""
        ...

Keep these responsibilities distinct, but they need not all live on the CuTeDSL op class. A constructor-heavy DSL op may stay focused on layouts and device code while a small adapter owns the seven static/class methods above.

• default_config(*runtime_args, target=...) owns deterministic no-tune selection. It may inspect tensor metadata, static operation arguments, and target facts, and must return a valid config for every input accepted by the adapter without benchmarking or reading device-resident tensor values. Keep this policy on the adapter rather than selecting a default in the parent wrapper. A constant fallback simply ignores its arguments. Name canonical mode configs after the mode rather than calling an architecture-specific config the default.
• tuning_key(*runtime_args, target=...) defines when an autotuned winner may be reused. Return host-visible tensor metadata and target facts only; never read device values or synchronize, because the key must remain valid during CUDA Graph replay. Use () when compile_call(...) already distinguishes every relevant workload. Exact keys are the safe starting point; introduce buckets only after measurements establish stable winner regions.
• configs(...) may inspect runtime shape, stride, dtype, alignment, or device facts.
• compile(...) owns fake tensors and the cached artifact boundary.
• compile_call(...) prevents runtime tensors from leaking into cache keys or subprocess payloads.
• launch(...) binds the compiled callable to real tensors for correctness checks and timing.

The protocol accepts positional runtime arguments for small ABIs. Keep one or two naturally named values positional; wrapping them in Args would add ceremony without clarifying ownership. Do not, however, grow a positional list of inputs, outputs, and static semantics indefinitely. Once several values form one launch ABI, pass one parent-only Args value. Nest Args in its adapter when no other component owns that ABI; use a descriptive top-level type when multiple components genuinely share it. Do not create a module-level _Runtime or _Args type merely to signal privacy. Reserve “runtime” for an execution environment or lifecycle-bearing object, not a passive argument tuple. compile_call(...) projects Args to static values, while launch(...) binds its real tensors to the compiled callable.

Nesting communicates ownership, not public API status. If an adapter has a normal class name but is an implementation detail, define the module's __all__ explicitly and list only the supported wrapper and configuration types. Do not use leading underscores as a substitute for deciding and documenting the module's actual public surface.

Public Wrapper

Users call one ordinary PyTorch function; they do not construct op objects or compiled callables.

def my_op(
    q: torch.Tensor,
    k: torch.Tensor,
    *,
    config: MyConfig | None = None,
    tune: bool = False,
    configs: Iterable[MyConfig] | None = None,
) -> torch.Tensor:
    args = MyOp.Args(
        q,
        k,
        torch.empty_like(q),
        torch.empty_like(q, dtype=torch.float32),
    )
    return run_tunable(
        MyOp,
        args,
        config=config,
        autotune=tune,
        configs=configs,
    )[0]

This gives three intentional modes:

my_op(q, k)  # conservative default
my_op(q, k, config=MyConfig(64, 128))  # force one specialization
my_op(q, k, tune=True)  # input-aware candidate method
my_op(q, k, tune=True, configs=(cfg_a, cfg_b))  # explicit candidate override

Reject config= with tuning and reject configs= without tuning rather than silently ignoring either argument. If target metadata is supplied explicitly, install it before generating candidates so configs(...) and compilation observe the same target. On heterogeneous multi-GPU hosts, derive that target from the runtime tensor's device rather than the process's current device. The installed target is process-global and sticky; callers that temporarily override it must restore the previous target.

Compile Contract

Inside compile(...):

1. Accept only static values that completely determine generated code and the fake ABI. Keep them structurally stable so jit_cache can encode them for warm process-local lookups before constructing a persistent hash or cache path. When the call tuple is not the right specialization identity, pass cache_key= a pure function returning a complete hashable tuple or frozen config. Return the structural key itself, never hash(key); jit_cache adds function, target, and source identity and uses the same structural key for memory and disk caching.
2. Keep batch-, token-, and head-count-like extents symbolic with cute.sym_int()/cute.sym_int64() when their dependent strides also remain dynamic. Bake a dimension into compile_call(...) only when it changes layout/stride address arithmetic, tiling, vectorization, block shape, shared-memory sizing, compile-time control flow, or another generated-code decision.
3. Build fake compact tensors with runtime-compatible shapes, strides, dtype, and assumed alignment.
4. Instantiate any internal CuTeDSL op object there; never require public callers to construct it.
5. Give compile_tvm_ffi(...) a stable lowercase name encoding every static compile argument. Class entrypoints may expose get_name() instead of passing name= explicitly.
6. Return the compiled TVM-FFI callable directly.

@staticmethod
@jit_cache
def compile(config: MyConfig):
    num_elements = cute.sym_int()
    source = cute.runtime.make_fake_compact_tensor(
        ..., (num_elements,), stride_order=(0,), assumed_align=16
    )
    destination = cute.runtime.make_fake_compact_tensor(
        ..., (num_elements,), stride_order=(0,), assumed_align=16
    )
    op = _MyOp()
    return compile_tvm_ffi(
        op._jit_entrypoint,
        source,
        destination,
        config.threads,
        config.tile_size,
        name=op.get_name(config),
    )

compile_tvm_ffi owns the typed TVM-FFI option and fake environment stream. Do not append another stream or pass string compiler options at call sites. compile_call(...) returns exactly one tuple of positional static arguments; use (config,) when compile(...) accepts only a config.