Trait ModelForward 

pub trait ModelForward: ModelStreams {
    // Required methods
    fn prefill(
        &self,
        tokens: &[u32],
        seq: &mut SequenceState,
        stream: u64,
    ) -> Result<DevicePtr>;
    fn prefill_chunk(
        &self,
        tokens: &[u32],
        seq: &mut SequenceState,
        chunk_start: usize,
        chunk_len: usize,
        is_last_chunk: bool,
        stream: u64,
    ) -> Result<DevicePtr>;
    fn decode(
        &self,
        token: u32,
        seq: &mut SequenceState,
        stream: u64,
    ) -> Result<DevicePtr>;
    fn decode_batch(
        &self,
        tokens: &[u32],
        seqs: &mut [&mut SequenceState],
        stream: u64,
    ) -> Result<DevicePtr>;

    // Provided methods
    fn supports_beam(&self) -> bool { ... }
    fn generate_beam_batch(&self, _reqs: &[BeamReq]) -> Result<Vec<Vec<u32>>> { ... }
    fn mixed_forward(
        &self,
        decode_tokens: &[u32],
        decode_seqs: &mut [&mut SequenceState],
        prefill_tokens: &[u32],
        prefill_seq: &mut SequenceState,
        prefill_chunk_start: usize,
        prefill_chunk_len: usize,
        prefill_is_last: bool,
        stream: u64,
    ) -> Result<MixedForwardResult> { ... }
    fn prefill_batch_chunk(
        &self,
        streams: &mut [PrefillSlice<'_>],
        stream: u64,
    ) -> Result<Vec<DevicePtr>> { ... }
    fn prefill_batch_chunk_rows(
        &self,
        streams: &mut [PrefillSlice<'_>],
        stream: u64,
        _row_base: usize,
    ) -> Result<Vec<DevicePtr>> { ... }
    fn mixed_forward_batch(
        &self,
        decode_tokens: &[u32],
        decode_seqs: &mut [&mut SequenceState],
        prefill_streams: &mut [PrefillSlice<'_>],
        stream: u64,
    ) -> Result<MixedBatchResult> { ... }
    fn normalize_ssm_states(
        &self,
        _seq: &SequenceState,
        _stream: u64,
    ) -> Result<()> { ... }
    fn prefill_twophase(
        &self,
        tokens: &[u32],
        seq: &mut SequenceState,
        _chunk_size: usize,
        stream: u64,
    ) -> Result<DevicePtr> { ... }
    fn is_mla(&self) -> bool { ... }
    fn hc_mult(&self) -> usize { ... }
    fn kv_block_size(&self) -> Option<usize> { ... }
    fn prefill_tail_split(&self, _tokens: &[u32]) -> Option<usize> { ... }
}

2026-09-26: Prefill, decode and mixed forward passes, beam search, and the model-shape queries is_mla, hc_mult and kv_block_size.

Required Methods

fn prefill( &self, tokens: &[u32], seq: &mut SequenceState, stream: u64, ) -> Result<DevicePtr>

2026-09-25: Prefill the whole prompt and return the logits of its last position.

fn prefill_chunk( &self, tokens: &[u32], seq: &mut SequenceState, chunk_start: usize, chunk_len: usize, is_last_chunk: bool, stream: u64, ) -> Result<DevicePtr>

2026-09-25: Prefill chunk_len prompt tokens from chunk_start. Only a chunk with is_last_chunk returns logits (of the last prompt position); any other chunk returns DevicePtr::NULL.

fn decode( &self, token: u32, seq: &mut SequenceState, stream: u64, ) -> Result<DevicePtr>

2026-09-25: One decode step for one sequence; returns the logits for the next token.

fn decode_batch( &self, tokens: &[u32], seqs: &mut [&mut SequenceState], stream: u64, ) -> Result<DevicePtr>

2026-09-25: One decode step for each sequence; returns logits [seqs.len(), vocab].

Provided Methods

fn supports_beam(&self) -> bool

2026-09-25: Whether this model implements Self::generate_beam_batch. Default false; NllbGpuModel returns true.

fn generate_beam_batch(&self, _reqs: &[BeamReq]) -> Result<Vec<Vec<u32>>>

2026-09-25: Run beam search to completion for each request and return each winning hypothesis. The scheduler calls it at prefill for num_beams > 1, instead of the token-by-token decode loop. Default: an error.

fn mixed_forward( &self, decode_tokens: &[u32], decode_seqs: &mut [&mut SequenceState], prefill_tokens: &[u32], prefill_seq: &mut SequenceState, prefill_chunk_start: usize, prefill_chunk_len: usize, prefill_is_last: bool, stream: u64, ) -> Result<MixedForwardResult>

2026-09-25: One decode step for each of decode_seqs plus one prefill chunk. Default: decode_batch (skipped when there are no decode tokens), then prefill_chunk.

fn prefill_batch_chunk( &self, streams: &mut [PrefillSlice<'_>], stream: u64, ) -> Result<Vec<DevicePtr>>

2026-09-25: Prefill one chunk for each of streams. Returns one entry per stream, in order: the last-token logits when that stream’s chunk is its last, else DevicePtr::NULL. Default: prefill_chunk per stream.

fn prefill_batch_chunk_rows( &self, streams: &mut [PrefillSlice<'_>], stream: u64, _row_base: usize, ) -> Result<Vec<DevicePtr>>

2026-09-25: Self::prefill_batch_chunk, with finishing stream i’s logits written to row row_base + i of the shared logits arena instead of row i.

decode_batch writes lane i’s logits to row i of the same arena, and in a mixed step the caller samples the decode rows after the prefill has run. Without the shift a finishing prefill stream overwrites a decode lane’s logits, and that lane samples another request’s first-token distribution. TransformerModel falls back to row_base = 0 when row_base + streams.len() exceeds the arena’s rows. Default: ignore row_base.

fn mixed_forward_batch( &self, decode_tokens: &[u32], decode_seqs: &mut [&mut SequenceState], prefill_streams: &mut [PrefillSlice<'_>], stream: u64, ) -> Result<MixedBatchResult>

2026-09-25: One decode step for each of decode_seqs plus one prefill chunk for each of prefill_streams. Default: decode_batch, a synchronise of the default stream, then prefill_batch_chunk_rows with the prefill rows placed after the decode rows.

fn normalize_ssm_states(&self, _seq: &SequenceState, _stream: u64) -> Result<()>

2026-09-25: Normalise the sequence’s SSM h_state. The scheduler calls it between prefill chunks. Default: no-op.

fn prefill_twophase( &self, tokens: &[u32], seq: &mut SequenceState, _chunk_size: usize, stream: u64, ) -> Result<DevicePtr>

2026-09-25: Prefill the whole prompt in one call, with chunk_size as the chunk length when an implementation splits it. Returns the last position’s logits. Default: one prefill_chunk over all of tokens.

fn is_mla(&self) -> bool

2026-09-25: Whether chunked prefill must run as one chunk. TransformerModel answers requires_single_chunk_prefill: MLA models (kv_lora_rank > 0) except glm5_next and kimi_k3. Default false.

fn hc_mult(&self) -> usize

2026-09-25: mHC hyper-connection stream count (config.hc_mult in TransformerModel); 0 means no highway. Default 0.

fn kv_block_size(&self) -> Option<usize>

2026-09-25: Tokens per paged-KV block, or None without a paged KV cache. The scheduler uses it to end a prefill chunk at metrale_gpu_runtime::ssm_tail_boundary. Default None.

fn prefill_tail_split(&self, _tokens: &[u32]) -> Option<usize>

2026-09-27: The token at which the prefill of tokens is split so that an SSM snapshot lands at prefill_plan::tail_split_point, or None when this model does not split it. The scheduler ends a non-last chunk there (prefill_plan::plan_chunk_len). Default None.

Implementors

impl ModelForward for NllbGpuModel

impl ModelForward for TransformerModel