Kina.ai

Specification

Phase 1, Section 5: Operation Set Extensions

Phase 1: Implementation — the shape-typed core compiled end-to-end. Section 5 of 8: Operation Set Extensions.

1. What this section commits to

Four new primitive operations that extend Phase 0's ten-operation core (Section 2: reduce has standard and keepdims typing rules). Each gets:

The four operations:

  1. Variadic-batch matmul — Tensor[T, σ ++ [M, K]] × Tensor[T, σ ++ [K, N]] → Tensor[T, σ ++ [M, N]] for any leading shape σ.
  2. Gather — Tensor[T, [V, D]] × Tensor[I64, σ] → Tensor[T, σ ++ [D]] with a bounds refinement.
  3. Softmax as a primitive — one operation that decomposes internally to the six-op composite from real-world examples §3.4.
  4. repeat_interleave — Tensor[T, σ ++ [N] ++ τ] × n_rep → Tensor[T, σ ++ [N * n_rep] ++ τ] for grouped-query attention.

These four operations close the "expressibility gap" the real-world examples surfaced: each one corresponds to a pattern that's either painfully verbose or impossible without them. After this section, the language can express transformer block forward passes idiomatically rather than awkwardly.

2. Why these four

The real-world examples walkthrough (real-world-examples.md) classified patterns into three groups: clean fits, verbose-but-workable, and outright limits. The four operations here address the third group:

Pattern Without extension With extension
Multi-head attention QK^T Flatten/unflatten dance: 3 reshapes per matmul Direct matmul Q K^T
Token embedding lookup No primitive; impossible without an explicit indexing operation gather embeddings indices
Softmax in attention Six explicit operations; type inference cascades Single softmax operation
Grouped-query attention K/V expansion unsqueeze + expand + reshape with axiomatic shape assertions repeat_interleave k axis n_rep

These four are the minimum set that converts the language from "possible but painful" to "natural." Other operations (cumsum, topk, sort) were considered and deferred — they're useful but not universally load-bearing for transformer architectures, and Phase 1 keeps the operation set tight.

A consistent property: each operation matches a pattern that PyTorch users write idiomatically. No new vocabulary the user has to learn; just the operations they already know, with refinement types added. This is deliberate — Phase 1's job is to validate that the type system is expressive enough for what users want to write, not to invent new primitives that bend programs to fit the system.

3. Variadic-batch matmul

Phase 0's matmul typed [M, K] × [K, N] → [M, N] (2D) and [B, M, K] × [B, K, N] → [B, M, N] (3D). The variadic-batch extension generalizes to any leading shape.

3.1 Typing rule

Extending Phase 0 §2's matmul rule:

Γ ⊢ a : Tensor[T, σₐ ++ [M, K]]
Γ ⊢ b : Tensor[T, σᵦ ++ [K, N]]
σₐ ≡ σᵦ                              (shape unification or SMT obligation)
T ∈ {F16, BF16, F32, F64, I32, I64}   (matmul only on numeric types)
─────────────────────────────────────────────────────
Γ ⊢ matmul a b : Tensor[T, σₐ ++ [M, N]]

The leading shape σₐ (equivalently σᵦ, by the unification premise) is propagated through to the result. No implicit broadcasting (per Phase 0 §2.11); the leading shapes must be equal.

The shape-equality premise discharges via the unification path from Section 2 §5.3 — most cases are syntactically identical or differ only by unbound shape variables. Only when leading dimensions are arithmetically related (e.g., B*H vs. X) does an SMT obligation arise.

3.2 CPU lowering

The strategy: reshape both inputs to 3D [prod(σ), M, K] and [prod(σ), K, N], dispatch standard batched matmul, reshape back.

In MLIR linalg:

// Input: a : tensor<?x?x?x?xf32>  (4D for example)
// Input: b : tensor<?x?x?x?xf32>
// Output: c : tensor<?x?x?x?xf32>

%a_collapsed = tensor.collapse_shape %a [[0, 1], [2], [3]] 
    : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
%b_collapsed = tensor.collapse_shape %b [[0, 1], [2], [3]] 
    : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>

%c_init = tensor.empty(%batch, %m, %n) : tensor<?x?x?xf32>
%c_collapsed = linalg.batch_matmul ins(%a_collapsed, %b_collapsed) outs(%c_init)
    : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

%c = tensor.expand_shape %c_collapsed [[0, 1], [2], [3]] 
    : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>

The collapse/expand operations are zero-cost in MLIR — they're metadata-only when the underlying memory layout permits. The linalg.batch_matmul is the real work.

3.3 GPU lowering

Via Section 3's runtime_gpu_matmul. The runtime function already takes a batch parameter; the compiler passes prod(σ):

// Compiler-emitted code (sketch):
int64_t prod = 1;
for (int i = 0; i < ndim_leading; i++) prod *= leading_shape[i];
runtime_gpu_matmul(out, a, b, prod, m, k, n, dtype_code);

The runtime calls cublasSgemmStridedBatched (or the dtype-appropriate variant). cuBLAS handles the strided-batched form natively — no reshape work needed at the runtime layer.

3.4 Test cases

The 4D case is the load-bearing one: a transformer block uses Q @ K^T where Q : [B, H, S, D] and K^T : [B, H, D, S]. Variadic-batch matmul makes this direct.

4. Gather

Required for embedding lookup, cross-entropy targets, MoE routing, and a dozen other patterns. Without it, these patterns are inexpressible.

4.1 Typing rule

Phase 1 commits to the 2D-values + arbitrary-shape-indices form, which covers embeddings cleanly:

Γ ⊢ values : Tensor[T, [V, D]]
Γ ⊢ indices : Tensor[I64, σ]
∀i ∈ shape(indices). 0 ≤ indices[i] < V    (bounds refinement; runtime-checkable)
─────────────────────────────────────
Γ ⊢ gather values indices : Tensor[T, σ ++ [D]]

The bounds refinement 0 ≤ indices[i] < V cannot in general be discharged at type-check time — it requires reasoning about tensor contents, which is beyond Phase 0 §1.3's quantifier-free Presburger fragment.

Three discharge paths:

  1. Static guarantee — the user's code masks indices into [0, V) via mod (e.g., indices_safe = indices mod V) and the type-checker verifies the masking. Requires a refinement-aware mod operation; in scope for Phase 1.

  2. Runtime check — the gather kernel checks bounds and traps on violation. Default behavior; cost is one comparison per index.

  3. Unchecked — explicit gather_unchecked opt-out for users who can prove externally. Phase 1 ships this only in --unchecked-bounds mode for benchmarking; not for production.

4.2 CPU lowering

Via linalg.generic with an indexed iteration:

%result = linalg.generic
  { indexing_maps = [
      affine_map<(d0, d1, ..., dn, dD) -> (d0, d1, ..., dn)>,           // indices
      affine_map<(d0, d1, ..., dn, dD) -> (dD)>],                        // output dim D
    iterator_types = [...] }
  ins(%indices : tensor<...xi64>)
  outs(%output : tensor<...xf32>)
{
^bb0(%idx: i64, %out: f32):
    %row = affine.apply ... %idx
    %val = tensor.extract %values[%row, %dD] : tensor<?x?xf32>
    linalg.yield %val : f32
}

The linalg.generic form expresses the indexed-read pattern. MLIR's vectorization passes can further optimize when access patterns are coalesced.

4.3 GPU lowering

Via runtime_gpu_gather. The in-house CUDA kernel:

__global__ void gather_f32(
    const float* __restrict__ values,    // [V, D]
    const int64_t* __restrict__ indices, // [N] (flattened)
    float* __restrict__ output,           // [N, D]
    int64_t V, int64_t D, int64_t N)
{
    int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N * D) return;

    int64_t n = i / D;
    int64_t d = i % D;
    int64_t idx = indices[n];

    // Bounds check (compiled out in --unchecked-bounds mode)
    if (idx < 0 || idx >= V) {
        // Set NaN as a soft signal; in checked mode the runtime traps separately
        output[i] = NAN;
        return;
    }

    output[i] = values[idx * D + d];
}

Coalesced when D is a multiple of warp size and D >= 32 — typical for embedding tables. For very small D, the kernel runs but isn't optimal; not a Phase 1 concern.

4.4 Test cases

5. Softmax as a primitive

Phase 0's spec showed softmax decomposed into reduce(max) + broadcast + sub + exp + reduce(sum) + broadcast + div. That works but is verbose; it cascades type-inference work; it produces six MLIR operations where one would do.

5.1 Typing rule

Γ ⊢ x : Tensor[T, σ]
T ∈ {F16, BF16, F32, F64}              (only floating point)
0 ≤ axis < |σ|                         (compile-time integer in valid range)
─────────────────────────────
Γ ⊢ softmax x axis : Tensor[T, σ]

The output shape is identical to the input; the output dtype is preserved. The axis is a compile-time integer; runtime axes are not supported in Phase 1 (would require runtime kernel selection, deferred).

5.2 CPU lowering

The lowering decomposes internally to the six-op composite. The user sees one operation; the compiler emits six MLIR primitives. This is genuine primitive-as-operation — the user-facing surface is one operation, the implementation is the obvious decomposition.

let lower_softmax x axis =
  let max_along = reduce_max x axis in
  let max_broadcast = broadcast max_along (shape_of x) in
  let shifted = binary x max_broadcast Sub in
  let exp_shifted = unary shifted Exp in
  let sum_along = reduce_sum exp_shifted axis in
  let sum_broadcast = broadcast sum_along (shape_of x) in
  binary exp_shifted sum_broadcast Div

The decomposition uses Phase 0 primitives. Numerical stability via the max-shift is preserved in the lowering — that's why this is a primitive: users get the numerically-stable form by default rather than having to remember to subtract the max.

5.3 GPU lowering

Via runtime_gpu_softmax. The dispatch logic:

The in-house kernel implements the standard parallel softmax with a two-pass strategy:

__global__ void softmax_f32_lastdim(
    const float* __restrict__ in,
    float* __restrict__ out,
    int64_t outer, int64_t inner)
{
    // Each block handles one row (outer dimension)
    int64_t row = blockIdx.x;
    if (row >= outer) return;

    const float* in_row = in + row * inner;
    float* out_row = out + row * inner;

    // Pass 1: reduce-max via warp/block reduction
    float thread_max = -INFINITY;
    for (int64_t j = threadIdx.x; j < inner; j += blockDim.x) {
        thread_max = fmaxf(thread_max, in_row[j]);
    }
    __shared__ float row_max;
    /* block-reduce thread_max into row_max ... */

    // Pass 2: exp and reduce-sum
    float thread_sum = 0.0f;
    for (int64_t j = threadIdx.x; j < inner; j += blockDim.x) {
        float e = expf(in_row[j] - row_max);
        out_row[j] = e;  // store unnormalized; normalize in pass 3
        thread_sum += e;
    }
    __shared__ float row_sum;
    /* block-reduce thread_sum into row_sum ... */

    // Pass 3: normalize
    for (int64_t j = threadIdx.x; j < inner; j += blockDim.x) {
        out_row[j] /= row_sum;
    }
}

Three passes; standard pattern. Phase 5's tile-IR work would replace this with a more aggressive fused implementation; Phase 1 ships the straightforward correct version.

5.4 Test cases

6. repeat_interleave

For grouped-query attention (GQA), where K/V heads are repeated to match Q heads.

6.1 Typing rule

Γ ⊢ x : Tensor[T, σ ++ [N] ++ τ]
0 ≤ axis < |σ ++ [N] ++ τ|             (compile-time integer)
n_rep : Int    n_rep ≥ 1               (compile-time integer)
─────────────────────────────────────────
Γ ⊢ repeat_interleave x axis n_rep : Tensor[T, σ ++ [N * n_rep] ++ τ]

The dimension at axis is multiplied by n_rep in the output. Other dimensions are preserved. The axis decomposes the shape into σ ++ [N] ++ τ where σ is everything before the axis and τ is everything after.

6.2 CPU lowering

Via linalg.generic with custom indexing that reads each input element n_rep times:

%result = linalg.generic
  { indexing_maps = [
      affine_map<(d0, ..., dax, ..., dn) -> (d0, ..., dax floordiv n_rep, ..., dn)>,
      affine_map<(d0, ..., dax, ..., dn) -> (d0, ..., dax, ..., dn)>],
    iterator_types = [...] }
  ins(%input : tensor<...>)
  outs(%output : tensor<...>)
{
^bb0(%v: f32, %_: f32):
    linalg.yield %v : f32
}

The floordiv n_rep in the input map is what produces the interleave pattern: output index dax reads input index dax / n_rep.

6.3 GPU lowering

Via runtime_gpu_repeat_interleave. In-house kernel:

__global__ void repeat_interleave_f32(
    const float* __restrict__ in,
    float* __restrict__ out,
    int64_t outer, int64_t N, int64_t inner, int n_rep)
{
    int64_t total = outer * (N * n_rep) * inner;
    int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;

    // Decompose i into (o, n_out, in_idx)
    int64_t in_idx = i % inner;
    int64_t n_out = (i / inner) % (N * n_rep);
    int64_t o = i / (inner * N * n_rep);

    int64_t n_in = n_out / n_rep;
    int64_t in_offset = ((o * N) + n_in) * inner + in_idx;

    out[i] = in[in_offset];
}

Memory access on the input side is non-coalesced for n_rep > 1 — n_rep consecutive output threads read the same input element. The hardware caches help significantly here, but it's not optimal. Phase 5's tile IR work would optimize this with shared-memory staging.

6.4 Test cases

7. Implementation milestones (within Phase 1's three milestones)

These four operations land progressively across Phase 1's existing milestones:

7.1 Milestone 1 (months 1–2): foundation

No operation extensions yet. Phase 0's ten operations are sufficient to demonstrate the CPU-only end-to-end pipeline (per Section 1 §3.1's MLP forward pass acceptance criterion).

7.2 Milestone 2 (months 3–4): GPU + first extensions

These two are sufficient to express attention's Q @ K^T → softmax → @ V pattern idiomatically. Combined with Phase 0's primitives, multi-head attention becomes expressible.

7.3 Milestone 3 (months 5–6): full operation set

By end of Milestone 3, all four extensions are in. The transformer-block forward pass (Section 1's exit criterion) uses all of them.

8. Acceptance criteria specific to operations

Criterion Target
Variadic-batch matmul 4D × 4D works; 5D × 5D works; mismatched leading shapes fail with useful errors
Gather Embedding-lookup pattern compiles cleanly; bounds refinement via mod works
Softmax Numerically stable on large inputs; matches PyTorch within FP tolerance
repeat_interleave GQA pattern (H_kv=4, n_rep=2 → H_q=8) compiles cleanly
Idiomatic transformer block The full block expressible in <200 lines of language code
MLIR linalg lowerings All four operations produce valid MLIR; mlir-opt accepts the output
Runtime FFI dispatch All GPU operations route through the FFI from Section 3 §6.1
Tracing integration All four operations emit spans (per Section 4 §5.2)

The "idiomatic transformer block" criterion is the load-bearing one. After this section's work lands, a full transformer block (RMSNorm + multi-head attention with GQA + softmax + SwiGLU FFN + residuals) should be expressible in language code that's roughly comparable to the PyTorch equivalent in line count and clarity.

9. Open issues

  1. Gather bounds-refinement via mod. Phase 1 commits to refinement-aware mod so users can write indices_safe = indices mod V and have the type-checker discharge the bounds obligation. The implementation: mod V produces output with refinement [0, V). Worth flagging that the type-checker has to track this through subsequent operations (e.g., a slice of a mod V-bounded tensor preserves the bound). Engineering work; not architectural.

  2. Softmax with runtime axis. Phase 1 requires the axis to be compile-time. Some workloads (e.g., dynamic axis selection in MoE) want runtime axis. The path forward: emit a kernel-selection runtime function that dispatches based on the runtime axis value. Phase 2 work; out of Phase 1 scope.

  3. 3D-values gather. The Phase 1 gather requires values to be 2D. Higher-rank values (e.g., gather rows from a 3D tensor) would require a more general typing rule. Pattern: the indices select a leading dimension; remaining dimensions are preserved. Phase 2 likely.

  4. Variadic-batch matmul with mixed precision. The current rule requires both inputs to have the same dtype. cuBLAS supports mixed-precision GEMM (e.g., F16 inputs, F32 accumulator). Adding this requires extending the typing rule with an output-dtype parameter. Phase 2 work; the design is straightforward.

  5. A flatten_pair primitive. Originally on the candidate list (real-world examples §6) but not in this section. Pattern: pair two adjacent dimensions into one ([B, H, S, D] → [B*H, S, D]) and unpair. Useful but not strictly necessary given variadic-batch matmul handles the common cases. Defer.

  6. The softmax primitive's numerical-stability guarantee. The lowering uses max-shift, which is numerically stable. The typing rule doesn't require the implementation to use max-shift — it just types the operation. A future hardening pass should encode the numerical-stability requirement in the operational semantics (Phase 0 §4 work). Not Phase 1; flagged.

10. Section 6 preview

Section 6 covers the verification track — the parallel Lean 4 work that runs alongside the OCaml/Rust implementation:

The verification track has its own discipline and its own acceptance criteria, distinct from the implementation track. Section 6 commits to what Phase 1 delivers on the verification side without overpromising.

End of Phase 1, Section 5.