Stiletto

compile(f, args...; io_dtype=nothing, intermediate_dtype=Float32, compute_dtype=Float32,
        allocator=(T, dims) -> CuArray{T}(undef, dims),
        max_workspace=nothing, deterministic=false, heuristics=nothing)

Trace f applied to symbolic stand-ins for args, build the resulting cuDNN graph, and return a callable. Calling the result with arrays of the same shapes as args executes the graph and returns outputs obtained from allocator; in-place assignments (mul!, y .= ...) write into the caller's buffers and allocate nothing.

Output element types follow Julia's promotion of the traced computation, so compiled code returns what the plain code would (explicit casts like Float16.(x) included). Passing io_dtype overrides that as a precision policy: every non-cast output is declared and allocated at io_dtype, regardless of what scalar constants promote the trace to.

Engine selection is steerable: max_workspace (bytes) caps the workspace a plan may demand, deterministic=true rejects engines flagged numerically nondeterministic, and heuristics overrides the heuristic modes consulted (a tuple of cuDNN.cudnnBackendHeurMode_t values, tried in order). When no engine satisfies the constraints, compilation throws cuDNN.UnsupportedGraphError rather than silently relaxing them.

g = @compile f(args...; kwargs...)

Expand to compile(f, args...), returning the compiled callable without executing it. Call keywords are closed over as trace-time constants.

jit(f, args...; kwargs...)

Compile f for this argument signature on first use and execute the cached graph, so jit can be called inline in hot code (a forward pass). Entries live in cuDNN's per-handle plan cache. Under CUDA graph capture the call replays as recorded stream work; the in-place forms (mul!, y .= ...) are the natural fit there, binding stable caller buffers and allocating nothing.

Arguments are runtime inputs — what the graph sees: arrays are keyed by shape/dtype/strides and rebind their pointers per call, scalars become by-value tensors rebound per call. Captured values are trace-time constants — what the tracer sees: isbits captures (flags, dims) specialize the graph per value and may be branched on; captured arrays are baked by identity. Pass data as arguments; capture configuration.

Keyword arguments forward to compile — including the engine-selection knobs (max_workspace, deterministic, heuristics) — and participate in the plan key: the same call with a different knob is a different cached plan.

d = @jit f(args...; kwargs...)

Expand to jit(f, args...): compile on first use, execute the cached graph. Call keywords are closed over as trace-time constants.

TracedArray{T,N}

The symbolic stand-in compile and jit run the traced function on: shape and element type are real, values are not. Subtyping AbstractArray admits it into generic array code; interface promises a symbolic value cannot keep (element access, allocation) throw informative errors instead of tracing incorrectly.

Stiletto.graph(c) -> cuDNN.Graph

The built cuDNN graph behind a compiled callable, for callers that schedule execution themselves: execute!(Stiletto.graph(c), b) with b from Stiletto.bindings is exactly what calling c does.

Stiletto.bindings(c, args...) -> (bindings, outputs)

Prepare the tensor bindings a call of c with args would execute with: every graph tensor mapped to its backing (argument storage — reshaped, viewed, or split into components as the trace declared it — captured arrays, by-value scalars, and freshly allocated output buffers). outputs are the arrays the call would return, already present in bindings.

Together with Stiletto.graph this splits c(args...) into its two halves, so execution can be scheduled by the caller — under CUDA graph capture, with rebinding between launches, or inside a larger execution scheme.

Stiletto.workspace(c) -> Int

Bytes of scratch workspace the compiled graph's selected execution plan demands per call. compile's max_workspace caps this at plan selection; querying it tells you what the chosen engine actually asked for.

Stiletto.mul!(C, A, B[, α, β])
Stiletto.mul!(y, A, x[, α, β])

LinearAlgebra.mul! executed as a cuDNN graph: C = α·A·B + β·C with the scaling and accumulation fused as a matmul epilogue; β = 0 leaves the destination unread. Matrix-matrix and matrix-vector, matching Base; batching is spelled batched_mul!.

On traced values this is what LinearAlgebra.mul! dispatches to, so both spellings trace identically. On plain arrays it jit-compiles the graph — the eager mul! for operands no BLAS method claims, such as block-scaled composites. α and β are trace-time constants — each coefficient pair compiles its own plan.

Stiletto.batched_mul(a, b)
Stiletto.batched_mul!(c, a, b[, α, β])
a ⊠ b

Batched matrix multiplication (M, K, batch...) × (K, N, batch...) → (M, N, batch...). Trailing batch extents broadcast pairwise: each pair must match or be 1 — missing dimensions count as 1 — and the result takes the larger. Unlike NNlib.batched_mul, any number of batch dimensions is accepted. ⊠ (\boxtimes) is the infix spelling.

The 5-arg form computes α·A·B + β·C with the scaling and accumulation fused as a matmul epilogue; β = 0 leaves the destination unread. α and β are trace-time constants — each pair compiles its own plan.

Stiletto.attention(q, k, v; stats=false, scale=1/√size(q, 1), causal=false,
                   seq_len_q=nothing, seq_len_kv=nothing)
Stiletto.attention!(o, q, k, v; ...)

Scaled dot-product attention softmax(scale ⋅ qᵀk) ⋅ v over rank-4 operands laid out (head_dim, heads, seq_len, batch). Grouped-query attention gives k/v fewer heads (a divisor of q's). scale binds by value at execution; causal bakes a causal mask into the graph.

stats=true additionally returns the softmax log-sum-exp as a (1, seq_len, heads, batch) Float32 array — the saved state attention_backward consumes.

seq_len_q/seq_len_kv (passed together) are Int32 vectors of per-batch valid lengths — runtime inputs masking each sequence's tail, so one graph built at maximum length serves ragged batches by rebinding lengths per call. Output rows past seq_len_q are undefined.

Stiletto.attention_backward(dO, q, k, v, o, stats; scale=1/√size(q, 1),
                            causal=false, seq_len_q=nothing,
                            seq_len_kv=nothing) -> (dQ, dK, dV)
Stiletto.attention_backward!(dQ, dK, dV, dO, q, k, v, o, stats; ...)

Gradients of attention with respect to q, k, and v, from the forward pass's output o and saved softmax stats (attention(...; stats=true)). scale/causal/sequence lengths must match the forward call. The operands cannot feed other operations in the same trace — the composite pattern owns its graph.

Stiletto.rmsnorm(x, scale; bias=nothing, epsilon=1f-5)
Stiletto.layernorm(x, scale, bias; epsilon=1f-5)
Stiletto.rmsnorm!(y, x, scale; ...) / layernorm!(y, x, scale, bias; ...)

RMS and layer normalization over the dimensions scale spans: a rank-1 scale of length size(x, 1) normalizes over the first dimension. scale and bias must be inputs of the trace (arguments or captures), not computed values — precompute them before compiling. epsilon is a trace-time constant.

Stiletto.batchnorm(x, scale, bias, mean, inv_variance)
Stiletto.batchnorm!(y, x, scale, bias, mean, inv_variance)

Inference-phase batch normalization over the channel dimension — the second-to-last of x, which must be at least 3-dimensional ((spatial..., channels, batch)). Parameters and statistics are per-channel vectors of length size(x, ndims(x) - 1). inv_variance is the folded inverse standard deviation 1 ./ sqrt.(running_var .+ epsilon); fold it once when loading parameters.

Stiletto.conv(x, w; stride=1, dilation=1, pre_padding=0, post_padding=pre_padding)
Stiletto.conv!(y, x, w; ...)

Cross-correlation of x :: (spatial..., channels, batch) with filters w :: (spatial..., channels ÷ groups, out_channels), producing (out_spatial..., out_channels, batch). The group count is inferred from the channel ratio (depthwise: filter channel extent 1). Spatial keywords take a scalar or one value per spatial dimension; asymmetric padding expresses causal windows (pre_padding=K-1, post_padding=0).

Stiletto.maxpool(x, window; stride=window, pre_padding=0, post_padding=pre_padding)
Stiletto.meanpool(x, window; ...; include_padding=false)
Stiletto.maxpool!(y, x, window; ...) / meanpool!(y, x, window; ...)

Pool x :: (spatial..., channels, batch) over window, producing (out_spatial..., channels, batch). stride defaults to the window (non-overlapping). meanpool divides by the number of contributing elements; include_padding=true divides by the full window size instead.

Stiletto.quantize!(dest, x)

Quantize x into dest::BlockscaledArray, writing quantized elements and block scales into dest's storage components. Block size, element type, and scale type are read from dest; the blocked dimension is the one dest's block spans. Inside a trace, dest must be an argument and x is fused into the graph — cuDNN's engines support quantize fused after a matmul (quantize!(dest, w' * x)), the full narrow-precision pipeline in one kernel.

Assignment into a block-scaled destination is the same store — quantization is how the destination stores values — so traced dest .= x and mul!(dest, a, b) route here. Partial writes (view(dest, ...) .= x) are refused: block scales couple elements, so only whole-array stores have a meaning.

Requires the Microscaling extension (using Microscaling).