feat(skippy): execute mixed prefill and decode iterations safely

crates/skippy-correctness/tests/parity_models.rs

@@ -5,7 +5,7 @@ mod harness;
 
 use harness::{
     activation_handoff_matches_full_model, assert_manifest_row_complete,
-    cache_state_restore_matches_recompute, p0_p1_manifest_rows,
+    cache_state_restore_matches_recompute, mixed_iteration_matches_serial, p0_p1_manifest_rows,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
@@ -40,6 +40,12 @@ macro_rules! family_module {
             fn cache_state_restore_matches_recompute() {
                 super::cache_state_restore_matches_recompute(SPEC).unwrap();
             }
+
+            #[test]
+            #[ignore = "downloads and loads model-family GGUF/package artifacts"]
+            fn mixed_iteration_matches_serial() {
+                super::mixed_iteration_matches_serial(SPEC).unwrap();
+            }
         }
     };
 }

crates/skippy-correctness/tests/parity_models/mod.rs

@@ -11,8 +11,8 @@ use model_ref::split_gguf_shard_info;
 use serde::Deserialize;
 use serde_json::Value;
 use skippy_runtime::{
-    ActivationFrame, GGML_TYPE_F16, MtpSource, RuntimeConfig, RuntimeKvPage, RuntimeLoadMode,
-    StageModel, StageSession,
+    ActivationFrame, GGML_TYPE_F16, IterationBatchRequest, MtpSource, RuntimeConfig, RuntimeKvPage,
+    RuntimeLoadMode, StageModel, StageSession,
     package::{PackageStageRequest, inspect_layer_package, materialize_layer_package_details},
     redirect_native_logs_to_file,
 };
@@ -94,6 +94,233 @@ pub(crate) fn cache_state_restore_matches_recompute(spec: FamilySpec) -> Result<
     run_stage_state_restore(&layout, spec, stage_start, stage_end, state_kind)
 }
 
+pub(crate) fn mixed_iteration_matches_serial(spec: FamilySpec) -> Result<()> {
+    prepare_native_logs()?;
+    let Some(case) = resolve_case_for_ignored_test(spec)? else {
+        return Ok(());
+    };
+    let layout = case.layout()?;
+    if case.row.is_package_only() {
+        eprintln!(
+            "skipping mixed-iteration parity for package-only {} / {}",
+            spec.llama_model, spec.family
+        );
+        return Ok(());
+    }
+    let (split_layer, _) = split_layers_for(case.row, layout.layer_count)?;
+    run_mixed_iteration_split(&layout, spec, split_layer)
+}
+
+fn run_mixed_iteration_split(
+    layout: &TestLayout,
+    spec: FamilySpec,
+    split_layer: u32,
+) -> Result<()> {
+    let n_gpu_layers = case_n_gpu_layers(manifest_row(spec)?);
+    let stage0_shape = StageShape {
+        stage_index: 0,
+        layer_start: 0,
+        layer_end: split_layer,
+        include_embeddings: true,
+        include_output: false,
+    };
+    let stage1_shape = StageShape {
+        stage_index: 1,
+        layer_start: split_layer,
+        layer_end: layout.layer_count,
+        include_embeddings: false,
+        include_output: true,
+    };
+    let stage0 = open_stage_model(
+        &stage_path(layout, spec, stage0_shape)?,
+        stage0_shape,
+        n_gpu_layers,
+    )?;
+    let stage1 = open_stage_model(
+        &stage_path(layout, spec, stage1_shape)?,
+        stage1_shape,
+        n_gpu_layers,
+    )?;
+    let tokens = stage0.tokenize(case_prompt(manifest_row(spec)?), true)?;
+    if tokens.len() < 3 {
+        bail!(
+            "{} / {} mixed-iteration prompt produced fewer than three tokens",
+            spec.llama_model,
+            spec.family
+        );
+    }
+    let prefix = &tokens[..2];
+    let long_prefill = &tokens[..tokens.len().min(6)];
+    let short_prefill = &tokens[..3];
+
+    let mut mixed_stage0 = [
+        stage0.create_session()?,
+        stage0.create_session()?,
+        stage0.create_session()?,
+    ];
+    let mut mixed_stage1 = [
+        stage1.create_session()?,
+        stage1.create_session()?,
+        stage1.create_session()?,
+    ];
+    let mut serial_stage0 = [
+        stage0.create_session()?,
+        stage0.create_session()?,
+        stage0.create_session()?,
+    ];
+    let mut serial_stage1 = [
+        stage1.create_session()?,
+        stage1.create_session()?,
+        stage1.create_session()?,
+    ];
+
+    let mixed_prefix_frame = mixed_stage0[0].prefill_chunk_frame(prefix, None, 0)?;
+    let (mixed_decode_token, _) =
+        mixed_stage1[0].prefill_chunk_frame_sampled(prefix, None, Some(&mixed_prefix_frame), 0)?;
+    let serial_prefix_frame = serial_stage0[0].prefill_chunk_frame(prefix, None, 0)?;
+    let (serial_decode_token, _) = serial_stage1[0].prefill_chunk_frame_sampled(
+        prefix,
+        None,
+        Some(&serial_prefix_frame),
+        0,
+    )?;
+    if mixed_decode_token != serial_decode_token {
+        bail!(
+            "{} / {} mixed-iteration setup token {mixed_decode_token} did not match serial {serial_decode_token}",
+            spec.llama_model,
+            spec.family
+        );
+    }
+
+    let decode_tokens = [mixed_decode_token];
+    let mixed_stage0_output = {
+        let [decode, long, short] = &mut mixed_stage0;
+        let mut requests = [
+            IterationBatchRequest {
+                session: decode,
+                token_ids: &decode_tokens,
+                positions: &[],
+                sampling: None,
+                input: None,
+                sample_last: true,
+            },
+            IterationBatchRequest {
+                session: long,
+                token_ids: long_prefill,
+                positions: &[],
+                sampling: None,
+                input: None,
+                sample_last: false,
+            },
+            IterationBatchRequest {
+                session: short,
+                token_ids: short_prefill,
+                positions: &[],
+                sampling: None,
+                input: None,
+                sample_last: true,
+            },
+        ];
+        StageSession::iteration_batch_sampled(&mut requests)?
+    };
+    if !mixed_stage0_output.samples.is_empty() {
+        bail!(
+            "{} / {} intermediate stage unexpectedly sampled tokens",
+            spec.llama_model,
+            spec.family
+        );
+    }
+
+    let (_, serial_decode_frame) =
+        serial_stage0[0].decode_step_frame(serial_decode_token, None, 0)?;
+    let serial_long_frame = serial_stage0[1].prefill_chunk_frame(long_prefill, None, 0)?;
+    let serial_short_frame = serial_stage0[2].prefill_chunk_frame(short_prefill, None, 0)?;
+    let serial_frames = [serial_decode_frame, serial_long_frame, serial_short_frame];
+    for (request_index, (mixed, serial)) in mixed_stage0_output
+        .request_outputs
+        .iter()
+        .zip(&serial_frames)
+        .enumerate()
+    {
+        if !activation_frames_match(mixed, serial) {
+            bail!(
+                "{} / {} mixed intermediate activation for request {request_index} did not match serial",
+                spec.llama_model,
+                spec.family
+            );
+        }
+    }
+
+    let mixed_stage1_output = {
+        let [decode, long, short] = &mut mixed_stage1;
+        let mut requests = [
+            IterationBatchRequest {
+                session: decode,
+                token_ids: &decode_tokens,
+                positions: &[],
+                sampling: None,
+                input: Some(&mixed_stage0_output.request_outputs[0]),
+                sample_last: true,
+            },
+            IterationBatchRequest {
+                session: long,
+                token_ids: long_prefill,
+                positions: &[],
+                sampling: None,
+                input: Some(&mixed_stage0_output.request_outputs[1]),
+                sample_last: false,
+            },
+            IterationBatchRequest {
+                session: short,
+                token_ids: short_prefill,
+                positions: &[],
+                sampling: None,
+                input: Some(&mixed_stage0_output.request_outputs[2]),
+                sample_last: true,
+            },
+        ];
+        StageSession::iteration_batch_sampled(&mut requests)?
+    };
+    let (serial_decode_prediction, _) = serial_stage1[0].decode_step_frame_sampled(
+        serial_decode_token,
+        None,
+        Some(&serial_frames[0]),
+        0,
+    )?;
+    serial_stage1[1].prefill_chunk_frame(long_prefill, Some(&serial_frames[1]), 0)?;
+    let (serial_short_prediction, _) = serial_stage1[2].prefill_chunk_frame_sampled(
+        short_prefill,
+        None,
+        Some(&serial_frames[2]),
+        0,
+    )?;
+    let mixed_samples = mixed_stage1_output
+        .samples
+        .iter()
+        .map(|sample| (sample.request_index, sample.predicted_token))
+        .collect::<Vec<_>>();
+    let expected_samples = vec![(0, serial_decode_prediction), (2, serial_short_prediction)];
+    if mixed_samples != expected_samples {
+        bail!(
+            "{} / {} mixed samples {mixed_samples:?} did not match serial {expected_samples:?}",
+            spec.llama_model,
+            spec.family
+        );
+    }
+    for index in 0..3 {
+        if mixed_stage0[index].token_count() != serial_stage0[index].token_count()
+            || mixed_stage1[index].token_count() != serial_stage1[index].token_count()
+        {
+            bail!(
+                "{} / {} mixed request {index} advanced session state differently from serial",
+                spec.llama_model,
+                spec.family
+            );
+        }
+    }
+    Ok(())
+}
+
 fn resolve_case_for_ignored_test(spec: FamilySpec) -> Result<Option<ResolvedCase>> {
     match ResolvedCase::resolve(spec) {
         Ok(case) => Ok(Some(case)),

crates/skippy-ffi/src/dynamic.rs

@@ -193,7 +193,7 @@ dynamic_symbols! {
     skippy_verify_tokens(session: *mut Session, token_ids: *const i32, token_count: usize, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status;
     skippy_decode_step_sampled(session: *mut Session, token_id: i32, sampling: *const SamplingConfig, input_activation: *const c_void, input_activation_bytes: usize, output_activation: *mut c_void, output_activation_capacity: usize, out_output_activation_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status;
     skippy_decode_batch_sampled(sessions: *const *mut Session, token_ids: *const i32, sampling: *const *const SamplingConfig, request_count: usize, out_predicted_tokens: *mut i32, predicted_token_capacity: usize, out_error: *mut *mut Error) -> Status;
-    skippy_iteration_batch_sampled(requests: *const IterationRequest, request_count: usize, output_descs: *mut ActivationDesc, output_payloads: *const *mut c_void, output_payload_capacities: *const usize, out_output_payload_bytes: *mut usize, out_predicted_tokens: *mut i32, predicted_token_capacity: usize, out_error: *mut *mut Error) -> Status;
+    skippy_iteration_batch_sampled(requests: *const IterationRequest, request_count: usize, output_descs: *mut ActivationDesc, output_payloads: *const *mut c_void, output_payload_capacities: *const usize, out_output_payload_bytes: *mut usize, out_sampled_request_indexes: *mut usize, out_predicted_tokens: *mut i32, sampled_output_capacity: usize, out_sampled_output_count: *mut usize, out_error: *mut *mut Error) -> Status;
     skippy_prefill_chunk_frame(session: *mut Session, token_ids: *const i32, token_count: usize, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_error: *mut *mut Error) -> Status;
     skippy_prefill_chunk_frame_sampled(session: *mut Session, token_ids: *const i32, token_count: usize, sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status;
     skippy_prefill_chunk_frame_with_positions(session: *mut Session, token_ids: *const i32, token_count: usize, positions: *const i32, position_count: usize, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_error: *mut *mut Error) -> Status;

crates/skippy-ffi/src/static_bindings.rs

@@ -212,8 +212,10 @@ unsafe extern "C" {
         output_payloads: *const *mut c_void,
         output_payload_capacities: *const usize,
         out_output_payload_bytes: *mut usize,
+        out_sampled_request_indexes: *mut usize,
         out_predicted_tokens: *mut i32,
-        predicted_token_capacity: usize,
+        sampled_output_capacity: usize,
+        out_sampled_output_count: *mut usize,
         out_error: *mut *mut Error,
     ) -> Status;