Skip to main content

compile_model_to_rdu_streamed

Function compile_model_to_rdu_streamed 

Source
pub fn compile_model_to_rdu_streamed(
    rdu: &RduModel<'_>,
    chip: RduChip,
    weights_dir: &str,
    decode: Option<StepShape>,
    sink: impl FnMut(RduProgramFiles) -> Result<()>,
) -> Result<RduBundle, String>
Expand description

Lower a model to the RDU as the chip would run it: ONE program per layer kind plus the head, each lowered, certified and handed to sink in turn, and a manifest binding every layer to its kind. Returns the manifest JSON and the model constants the bundle ships: a decode deployment’s rotary caches ({weights_dir}/Rcos.f, {weights_dir}/Rsin.f, one pair per rotation the layers turn by — rope_tables at ctx positions, little-endian f32), which every step reads at its positions so the host hands in only tokens and positions.

Layers with equal specs lower to one dataflow (spatial::streamed::layer_kinds), so the whole forward — however deep — is a program per kind and the head. Programs are rendered and sunk one at a time, so the caller never holds more than one program’s JSON. The weights stay outside the programs: the manifest names, for every layer, which file each of its kind’s tensors is read from, and each tensor’s role.

decode: Some(StepShape { ctx, batch, tp }) lowers the DECODE deployment instead of the authored forward: an embedding program, layer kinds whose attention appends to and reads a KV cache of ctx positions, and a head that emits the next token (see streamed::lower_decode_streamed) — for batch sequences per step, every weight read once for all of them, when batch > 1 (streamed::lower_batched_decode_streamed); over tp chips, each layer kind one program with every chip its own heads and a slice of every weight, when tp > 1 (streamed::lower_tp_decode_streamed, the manifest naming each chip tensor’s slice).