Module variant_bake 

Optimize-time kernel-variant baking — the same-device arm selector.

The gap this closes

optimize_graph records Op::Branch decision points whose arms are competing routes. Two kinds of arm coexist under one representation:

• Placement arms — the same op on different (backend, device) pairs (CPU vs. CUDA). These stay live Op::Branch points, resolved at dispatch by the runtime route picker on live load / VRAM (crate::ranker::route_picker, 06-runtime).
• Kernel-variant arms — different implementations of the same op on the SAME device (the decomposed attention region vs. a fused flash kernel, both on CUDA). Per 04-optimization these are “largely baked at optimize time”, NOT deferred to the runtime picker.

The shipped runtime picker resolves only placement arms: it ranks arms by their (backend, device) keyed on arm 0’s op, so two same-device arms map to the same candidate and always resolve to arm 0 (the finding in commit c7bfc1b5). A same-device fused variant — e.g. the CUDA flash decode arm from crate::decode_flash::offer_decode_flash_arm — can be offered but never selected.

What this module does (the bake)

bake_variant_branches runs at optimize time (Picker 1, where costs + caps + the Judge oracle are in hand — NOT the executor). For every reachable same-device variant branch it:

1. costs arm 0 (the decomposed-region oracle) and each variant arm;
2. picks the unique variant that is strictly cheaper than arm 0;
3. collapses the branch to that winner via fuel_graph::Graph::collapse_variant_branch (the merge is rewired to the winner; the branch becomes an inert single arm — “the Branch collapses” of 04-optimization).

Conservative defaults (the oracle stands): a tie, an unknown/sentinel cost, or a variant whose kernel is not available on the placed device all resolve to arm 0 — the decomposed correctness oracle. Placement branches (arms on ≥2 devices) are never touched — they are left live for the runtime picker. A graph with no same-device variant branch (every CPU/Vulkan build today, since no CUDA flash arm is offered there) is a no-op, so the byte-exact suites are unaffected by construction.

Where it rides

The bake mutates the graph in place during optimize_graph, so the baked choice rides crate::optimize::OptimizedGraph (which derives its dispatch order from the graph) with no extra state — plan-once persistent decode bakes once per generation and every realize lowers the collapsed graph.

What the full version adds

This is the smallest constitutional version. The full version adds: Judge-refined variant ranking (a Layer-2 measurement overriding the Layer-1 composed cost — plugs into the cost closure); Pareto retention of both arms (keep the variant as a runtime-selectable path instead of pruning it, once the picker can distinguish same-device variants); and, if it ever becomes constitutional, runtime load-adaptive variant switching. Today variant choice is baked and the loser is pruned.

Functions

bake_variant_branches

Bake same-device kernel-variant branches to their cost winner (see the module docs). Returns the number of branches collapsed to a variant arm (arm > 0) — a diagnostic the caller / tests assert on, also accumulated into the thread-local variant_bakes_thread_local counter. Never panics: a structural surprise from Graph::collapse_variant_branch is ignored rather than crashing (never-panic on a production path).

decode_arm_composite_ns

The production cost provider (04-optimization “cost-from-decompose” + the region composed-Layer-1 sum). Folds a Layer-1 CostEstimate over the arm’s interior and converts it to a composite-ns figure on the placed backend’s throughput roofline. This is the defensible v1 the session prompt blesses: compare the fused arm’s declared cost against the region’s composed Layer-1 sum over its nodes.

decode_arm_composite_ns_judged

The Layer-2-aware cost provider — the slice-4 payoff. Consults the Judge oracle (empirical measured latency, 06-optimization Layer-2) first per interior node, falling back to the decode_arm_composite_ns Layer-1 composite for any node the Judge hasn’t measured. Summed per arm exactly as the Layer-1 provider does, so the comparison stays apples-to-apples:

variant_bakes_thread_local

The current thread’s cumulative variant-bake count. Read it before and after an optimize_graph/realize call and diff to learn how many branches this pass collapsed to a fused variant winner — e.g. a BF16 CUDA decode step should bake one flash-decode arm per layer once the flash-arm capability seam is bridged. Thread-local so a single-threaded realize sequence isn’t polluted by concurrent test threads.

Type Aliases

ArmCostFn

The composite-nanosecond cost of a branch arm, or None when it is unknown or inadmissible — an unpriced/sentinel-cost op, or a variant whose kernel is not available on the placed device. None means “this arm cannot win”, so the conservative arm-0 oracle stands.