One inference API, two platform runtimes

The shared declaration is the easy part. Deciding what it is allowed to promise is the whole design.

Draw that line wrong and the abstraction leaks in the direction that hurts most: shared code that compiles everywhere and behaves differently on one platform, with nothing in the type system saying so.

The case this post is about#

One declaration, drawn where it felt natural:

expect class InferenceSession {
    suspend fun run(input: Tensor): Result<Tensor>
    fun close()
}

That looks minimal and it is not. By the time it shipped it had grown an encoder, a decoder, a sampler and a token counter — each added because it sat next to the platform call, and each now a promise two independently written implementations have to keep in step with nothing checking that they do.

Three of those four are arithmetic. This post follows that one declaration down to two methods, because everything that moved above the line became testable on a laptop, and the sampler that was duplicated per platform had already diverged.

The symptom#

A platform capability — inference, biometrics, notifications, a share sheet — gets a shared declaration and two implementations. The structure is obvious and the compiler enforces it.

What the compiler does not enforce is semantic parity. Two implementations can satisfy the same signature while differing in ways that matter:

• one is synchronous and fast, the other takes seconds on first call
• one throws on an unavailable resource, the other returns an empty result
• one supports cancellation, the other ignores it
• one is safe to call concurrently, the other serialises internally

Every one of those type-checks. Shared code written against the declaration will be correct on the platform its author tested and subtly wrong on the other, and there is no signal until someone runs it on a device.

Why the obvious fix failed#

The obvious fix is a richer interface: add methods for the differences, surface the capabilities, expose the knobs.

It migrates the platform difference into the shared declaration, which is where it does the most damage. Now every caller has to reason about both runtimes, and the seam has stopped separating anything — it is a union of two APIs with a Kotlin keyword in front of it.

The second attempt is the opposite: make the declaration minimal and let each side do what it can. That pushes the difference into behaviour, which is worse, because behaviour is not inspectable at the call site.

The actual mechanism#

Put everything that is not the platform call itself in common code, and make the declaration promise only what both sides can genuinely honour.

Concretely, the platform side should own almost nothing:

expect class InferenceSession {
    suspend fun run(input: Tensor): Result<Tensor>
    fun close()
}

Everything else — building the input, validating shape, timing, retry policy, mapping to domain types, the entire state machine around availability — lives above the line, written once, tested once, identical on both platforms.

The two properties that make the declaration honest:

It returns a result type rather than throwing. Failure is a value both sides can produce, so the difference between throws and returns empty never reaches a caller.

It makes no promise about timing. A suspend function says "this may take a while" on both sides, which is the only timing claim that survives two very different runtimes.

Even close() is a promise, and a two-word one hides three questions: is it idempotent, is it safe after a failed run, and does it block. Two implementations will answer those differently unless somebody wrote the answers down, and the shared code that calls it in a finally block depends on all three. The narrowest declaration is still a contract — narrow means fewer clauses, not fewer obligations per clause.

The fix#

The rule that generalises past inference: the expect declaration should be the narrowest thing that still lets common code do its job.

Every method you add is a promise two independently-written implementations must keep in step, and nothing checks that they did. A seam with two methods has two opportunities to diverge; one with ten has ten, and the divergences are invisible.

The test I apply when adding anything to the shared declaration: could I write this above the line instead? If the answer is yes — and it usually is, because most of what surrounds a platform call is arithmetic and state — it goes above. What remains below is the irreducible call, which is a much smaller surface to keep honest.

Applied to a tokenizer-plus-sampler-plus-session design, the seam collapses from six declarations to two, and everything that moves up is testable without a device:

// BELOW the line. Two declarations, both genuinely platform.
expect class NativeSession(modelPath: String) {
    fun forward(tokens: IntArray): FloatArray      // logits for the next position
    fun release()
}

// ABOVE the line. Shared, written once, and the only place these decisions exist.
class TextGenerator(
    private val session: NativeSession,
    private val vocab: Vocab,
    private val rng: () -> Float,                  // injected, so a test is deterministic
) {
    data class Budget(val maxNewTokens: Int, val maxTotalTokens: Int)

    fun generate(prompt: String, budget: Budget, temperature: Float): Result {
        var tokens = vocab.encode(prompt)
        if (tokens.size >= budget.maxTotalTokens) {
            // A refusal, not a truncation. Silently dropping the start of a prompt
            // produces a confident answer to a question nobody asked.
            return Result.Refused("prompt is ${tokens.size} tokens, ceiling is ${budget.maxTotalTokens}")
        }
        val out = StringBuilder()
        var produced = 0

        while (produced < budget.maxNewTokens && tokens.size < budget.maxTotalTokens) {
            val next = sample(session.forward(tokens), temperature)
            if (next == vocab.eos) return Result.Complete(out.toString(), produced)
            out.append(vocab.decode(next))
            tokens += next
            produced++
        }
        return Result.Truncated(out.toString(), produced,
            if (produced >= budget.maxNewTokens) "hit maxNewTokens" else "hit maxTotalTokens")
    }

    private fun sample(logits: FloatArray, temperature: Float): Int {
        if (temperature <= 0f) return logits.indices.maxBy { logits[it] }
        val max = logits.max()
        val weights = logits.map { kotlin.math.exp((it - max) / temperature) }
        var r = rng() * weights.sum()
        weights.forEachIndexed { i, w -> r -= w; if (r <= 0f) return i }
        return weights.lastIndex
    }

    sealed interface Result {
        data class Complete(val text: String, val tokens: Int) : Result
        data class Truncated(val text: String, val tokens: Int, val why: String) : Result
        data class Refused(val why: String) : Result
    }
}

The three outcomes, driven by a fake session with no device involved:

val gen = TextGenerator(FakeSession(emitsEosAt = 12), vocab, rng = { 0.5f })
println(gen.generate("summarise this", Budget(maxNewTokens = 64, maxTotalTokens = 512), 0.7f))
println(gen.generate("summarise this", Budget(maxNewTokens = 5,  maxTotalTokens = 512), 0.7f))
println(gen.generate(longPrompt, Budget(maxNewTokens = 64, maxTotalTokens = 512), 0.7f))

Complete(text=..., tokens=12)
Truncated(text=..., tokens=5, why=hit maxNewTokens)
Refused(why=prompt is 1204 tokens, ceiling is 512)

Every one of those three paths is above the line, which means all three are covered by tests that run on a laptop. The six-declaration version of this seam had encode, decode, sample, tokenCount and the budget logic below it, duplicated per platform — so the Refused branch existed twice, and on one of the two it truncated instead.

How we proved it#

This is where I owe you a clear statement of what is measured and what is not.

Measured: the seam pattern itself, across a codebase where platform bridges exist for permission requests, biometrics, deep links, share sheets, calendar events and notification scheduling, with implementations across six source sets. The narrow-declaration discipline is what keeps those maintainable, and the failures I have seen are exactly the ones above — a method added to the shared declaration because one platform made it easy, and a divergence nobody caught until a device ran it.

Also measured, and relevant: where these modules live and how they are wired is load-bearing in a way that has bitten me. A fork's own module declarations belong in a file the upstream sync does not overwrite; thirty-six of them were sitting in one it replaces wholesale, one sync away from silent deletion. That is not an inference problem, it is a platform-module problem, and anything you build behind an expect/actual seam inherits it.

Not measured: I have not shipped this seam against Core ML and LiteRT. The declaration above is what I would write, and the reasoning behind each property is from the general pattern rather than from two inference runtimes disagreeing with me in production.

So treat the seam discipline as tested and the inference-specific shape as a design I am confident in and have not yet been corrected by.

Where I expect it to leak#

Since I am reasoning rather than reporting here, the useful thing is to be specific about where I expect the design to fail, so it is checkable rather than vague.

Model lifecycle. A session is not just a call — the model is loaded, possibly evicted under memory pressure, and reloaded. If eviction behaves differently across platforms, run returning a failure value is not enough; callers need to distinguish this input failed from the model is gone, and that distinction will want to be in the declaration.

Streaming. A single suspend fun run returning one result is the wrong shape for token streaming, and streaming is the common case for anything text-generating. The honest version is probably a Flow, and I do not know whether both runtimes can honour the same cancellation semantics on it.

First-call cost. suspend says "may take a while" and does not distinguish a 20ms call from a multi-second first load. Shared code that cannot tell those apart will make bad decisions about what to show the user, and the fix is probably an explicit warm-up in the declaration — which is exactly the kind of addition the narrowness rule argues against.

That last one is a genuine tension in the design rather than a detail, and I would rather flag it than present a seam that has not met the problem.

What this generalizes to#

When one API has two platform implementations, decide what the shared declaration is allowed to promise before writing either side. Keep it to the irreducible call; put the arithmetic, the state machine and the policy above the line where they are written and tested once.

Prefer values over exceptions and make no claims about timing, because those are the two places two runtimes diverge without the compiler noticing. Then, when you do meet the platform, expect the pressure to come from lifecycle, streaming and first-call cost — and treat each addition to the shared declaration as a promise you have just made twice with nothing checking that both were kept.