Composable Kernel Language: Making GPU Tile Interfaces Negotiable

Codebase for CKL

What’s the problem?

Persistent kernel and mega kernel are gaining more and more popularity with projects like Mirage Persistent Kernel and Event Tensor. This is because as the compiler is powerful enough to extract the most capacity out of a standalone kernel, the overhead of kernel scheduling becomes the next critical hotspot to attack. While CUDA graph can potentially help with cpu launching overhead, it does not reduce the time for memory movement between kernels, if the layout of two kernels are incompatible.

Within a kernel, each operation expects a data layout. A producer may leave a tile in one lane/register distribution while its consumer expects another. The mismatch might be free, a register permutation, a warp shuffle, a shared-memory exchange, or a full kernel boundary through global memory. For example, triton uses a specialized convert layout ops to convert data layout within the kernel. This guarantees the correctness and compatibilty across different operations, but it’s not optimal. TileLang, on the other hand, introduces layout inference to produce the optimal layout for the entire kernel, so it does not need explicit layout conversion.

Composable Kernel Language (CKL) is trying to solve a similar issue, but between kernels instead of within a kernel. A task states a logical computation and exposes one or more physical realizations. Each realization describes its input and output data distributions, placement, resource requirements, capabilities, and estimated cost. CKL constructs a graph of task invocations and selects realizations jointly with the conversions between them. Layout is therefore not merely an attribute of a tensor. It is part of a negotiable interface between computations.

Note: Realization just means one possible kernel variation, which can be introduced by using different data layout, different block size/tile size/vectorization size, or different target features.

1. The composition problem

Consider a fused sequence that dequantizes an input tile and feeds it to a matrix multiplication. At the logical level, composition is easy. It is basically quantized tile ── dequantize ──> f16 tile ── matrix multiply ──> accumulator tile.

At the physical level, the boundary carries much more information. To name just a few, which lane owns each element, their placement in registers, alignment and vectorization constraints, etc. Suppose the dequantization implementation naturally emits four adjacent values per lane, while an MMA instruction expects a lane/value mapping fixed by the instruction encoding. Several choices are possible:

1. select a different dequantization schedule that directly produces the MMA fragment;
2. keep both implementations and insert a lane-local permutation;
3. exchange values with subgroup shuffles;
4. stage through shared memory and synchronize;
5. decline fusion and materialize through global memory.

The cheapest decision cannot, in general, be made by either task alone. Nor can every boundary be optimized independently. If a producer fans out to two consumers, changing its output distribution affects both edges. A locally slower producer can therefore be globally cheaper when it avoids two expensive conversions.

For a task graph $G=(V,E)$, where each vertice $v$ is a kernel invocation and each edge $e$ is a producer-consumer dependency, CKL’s planning problem is to choose one realization $r_v$ for every task invocation $v$ and one converstion strategy $c_e$ for every boundary e, that minimizes

$\sum_{rv} \text{ExecutionCost}(rv)+\sum_{c_e} \text{ConversionCost}(c_e)$

This is not a complete model for kernel scheduling, because we still need to consider for example resource pressure if both kernel content on one resource, but I just want to start with a simple model.

2. Prior art

CKL builds on several strong systems.

2.1 CuTe and CUTLASS

CuTe’s decisive simplification is to model a layout as a mapping from coordinates to indices. CuTe’s layout documentation describes a layout as a function from a coordinate space to an index space. Its algebra constructs larger mappings using composition, products, and division.

The recent paper Categorical Foundations for CuTe Layouts makes the algebraic structure more precise. CKL shares the goal of giving layout transformations laws that can be tested and reasoned about.

2.2 AMD Composable Kernel and CK Tile

AMD’s CK Tile provides the most direct baseline for CKL’s distribution model. The important lesson is its separation of coordinate spaces:

• P,Y ─> X: processing coordinate and local value to logical tile
• Y -> D: local value to per-thread register storage
• X: address tensor coordinate to backing storage

In CK source, it retains both the P,Y -> X adaptor and the Y -> D per-thread descriptor.

// include/ck_tile/core/tensor/tile_distribution.hpp
template <typename PsYs2XsAdaptor_,
          typename Ys2DDescriptor_,
          typename StaticTileDistributionEncoding_,
          typename TileDistributionDetail_,
          bool IsWarpLevelParallelOnly_ = false>
struct tile_distribution
{
    using PsYs2XsAdaptor = remove_cvref_t<PsYs2XsAdaptor_>;
    using Ys2DDescriptor = remove_cvref_t<Ys2DDescriptor_>;
    using DstrEncode     = remove_cvref_t<StaticTileDistributionEncoding_>;
    using DstrDetail     = remove_cvref_t<TileDistributionDetail_>;

    static_assert(PsYs2XsAdaptor::is_static() && Ys2DDescriptor::is_static(),
                  "wrong! should be static");

    static constexpr index_t NDimX = PsYs2XsAdaptor::get_num_of_bottom_dimension();
    static constexpr index_t NDimY = Ys2DDescriptor::get_num_of_top_dimension();
    static constexpr index_t NDimP = PsYs2XsAdaptor::get_num_of_top_dimension() - NDimY;
    static constexpr index_t NDimR = StaticTileDistributionEncoding_::NDimR;
    ...

The per-thread descriptor are later used to to index register-backed storage.

// include/ck_tile/core/tensor/static_distributed_tensor.hpp
    template <typename TileDistributedIndices>
    CK_TILE_HOST_DEVICE constexpr const DataType& operator[](TileDistributedIndices) const
    {
        static_assert(is_static_v<TileDistributedIndices>,
                      "wrong! Tile Distributed Indices should be static");

        constexpr auto y_idx = get_tile_distribution().get_y_indices_from_distributed_indices(
            TileDistributedIndices{});

        return thread_buf_[number<ThreadTensorDesc{}.calculate_offset(y_idx) / PackedSize>{}];
    }

CKL preserves this separation in a representation designed for compiler analysis.

3. Layout Abstraction

The initial design considered a single broad “layout” object. Experience with CK and the core validation work showed that this conflates questions with different domains and laws. CKL instead uses four objects: index spaces, index maps, storage layouts, and distributions.

3.1 Index spaces

An IndexSpace is a finite product of named axes. It preserves hierarchical grouping while also providing a flat coordinate domain.

tile(m:16, n-factors(n0:2, n1:4))

We need factors grouping because for example, (6,2) and (4,3) can be related through the common prime-factor refinement (2,2,3).

3.2 Index maps

An IndexMap is a guarded mapping between index spaces:

$f: A \rightarrow B$

Its result expressions currently include input dimensions, constants, addition, multiplication, floor division, modulo, and XOR. A predicate describes where the map is active. These primitives cover reshaping, slicing, padding guards, common strided layouts, distribution formulas, and bounded swizzles without pretending to represent arbitrary programs.

Composition is ordinary function composition:

$(g \circ f)(x) = g(f(x))$

CKL basically uses substitution to realize the composition. It also implemented a naive expression tree for constant propagation and evaluation, so that it supports simple provenance.

3.3 Storage layouts

StorageLayout concerns about logical space, address space, address unit, index-to-address map, and alignment. It does not care about thread ownership.

3.4 Distributions

A Distribution determins thread ownership actually. It answers the question of which logical tile elements does each executor own, and in what local order. It maps executor space (block, warp, lane) and local storage (register) to the abstract tile coordinate. The ownership map and local-storage map must remain separate. Two distributions may give every lane the same logical elements but arrange those elements differently in registers.

The complete path from an executing lane to memory is a composition of semantically distinct maps:

4. Layout Conversion

Once source and target distributions are explicit, CKL classifies the physical difference between them. CKL currently produces one of six classes (which also corresponds to the cases in Triton’s convert layout op implementation):

This part is embarrassingly just a stub at the moment. Concrete lowering has not be implemented yet.

5. Tasks and alternatives

A CKL task is a logical computation. An alternative is one legal realization of that task. Alternatives can come from four sources:

• a user-authored algorithmic variant;
• a compiler-derived schedule;
• a target extension, such as a fixed MMA instruction realization;
• an imported library implementation.

A simplified Python declaration looks like:

lane = Space(lane=32)
lhs_local = Space(kHalf=2, rowHalf=2, pair=2)
lhs_tile_space = Space(m=16, k=16)
lhs_distribution = Distribution(
    lane,
    lhs_local,
    lhs_tile_space,
    IndexMap(
        lane.product(lhs_local),
        lhs_tile_space,
        [dim(0) // 4 + dim(2) * 8, (dim(0) % 4) * 2 + dim(3) + dim(1) * 8],
    ),
    IndexMap(lhs_local, Space(address=8), [0 + dim(0) * 4 + dim(1) * 2 + dim(2) * 1]),
)
@task(
    alternatives=[
        Alternative(
            "mma-sync-m16n8k16-f16-f32",
            {
                "implementation_id":
                    "nvidia.mma.sync.m16n8k16.row.col.f32.f16.f16.f32",
                "estimated_execution_cost": 1,
            },
            inputs=(
                Port("lhs", lhs_distribution),
                Port("rhs", rhs_distribution),
                Port("acc", acc_distribution),
            ),
            outputs=(Port("result", acc_distribution),),
        )
    ]
)
def mma(lhs: lhs_tile, rhs: rhs_tile, acc: acc_tile) -> acc_tile:
    ...

Well, I have to admit this is pretty ugly and not fun to write. In future, the Port and lhs_distribution should be inferred and optimized by the compiler naturally, so we can focus on just high-level implementation, and expose a Triton-like API.

6. Graph-wide realization selection

For a sequence of kernel invocations, CKL models them as graph nodes. Two invocations of the same task may select different alternatives. Vertices are invocations and edges identify producer and consumer dependencies.

As we mentioned in the first section, CKL models kernel alternative selection as a graph optimization problem. The implementation right now brute-force searches all potential combination to find the optimal alternative selection, based on heuristic cost model. Memory boundaries participate in the same graph. A direct ckl.load_tile producer and ckl.store_tile consumer become fixed boundary nodes whose distributions are determined by the memory operations.

As illustrated below, “global optimization” means finding the implementationm (data layout) that reduces the sum of all conversion cost and execution cost.

7. Why explicit conversions matter

Many compilers can insert a conversion after discovering a type mismatch. CKL insists that conversion must remain visible during selection because explicit conversion planning provides four benefits:

1. Optimization: conversion cost participates in graph-wide selection and fusion.
2. Correctness: ownership changes generate concrete movement witnesses and synchronization.
3. Diagnostics: the compiler can explain why two contracts differ and where values move.
4. Provenance: a selected plan can record alternatives considered, constraints applied, and the reason for the final choice.

For example, “layouts differ” is not an actionable diagnostic. A useful report is closer to:

selected producer schedule P2 and consumer instruction C1
total estimated score: 14

boundary dequant.result -> mma.lhs:
  conversion: subgroup exchange
  32 lane-to-lane movements
  rejected P1/C1: lower producer cost, but required shared-memory exchange and barrier

This is obviously not implemented yet, but should be really interesting and helpful for kernel developer.

11. Conclusion

Overall, this is a fun research project. Based on what I’ve experimented so far, composability can be achieved in theory, but global optimization across kernels is very hard. Throughout the project, I used the concepts like “execution cost”, “conversion cost”, “resource pressure penalty”, but they sometimes are not static numbers that we can derive from. Nevertheless, this opens a structured exploration space, and can potentially benefits kernel design agents or maybe ML compilers that models those costs with weights and biases.