Struct RaggedBatch 

pub struct RaggedBatch {
    pub token_ids: Vec<u32>,
    pub cumulative_offsets: Vec<u32>,
    pub max_seq_len: usize,
}

A ragged (unpadded) batch for efficient ModernBERT inference.

ModernBERT achieves its speed advantage by avoiding padding tokens entirely. Instead of [batch, max_seq_len], it uses a single contiguous 1D sequence with offset indices to track document boundaries.

Memory Layout

Traditional (padded):
[doc1_tok1, doc1_tok2, PAD, PAD, PAD]  <- wasted compute
[doc2_tok1, doc2_tok2, doc2_tok3, PAD, PAD]

Ragged (unpadded):
[doc1_tok1, doc1_tok2, doc2_tok1, doc2_tok2, doc2_tok3]
cumulative_offsets: [0, 2, 5]  <- doc1 is [0..2], doc2 is [2..5]

Fields

token_ids: Vec<u32>

Token IDs flattened into a single contiguous array. Shape: [total_tokens] (1D, no padding)

cumulative_offsets: Vec<u32>

Cumulative sequence lengths. Length: batch_size + 1 Document i spans tokens [offsets[i]..offsets[i+1])

max_seq_len: usize

Maximum sequence length in this batch (for kernel bounds).

Implementations

impl RaggedBatch

pub fn from_sequences(sequences: &[Vec<u32>]) -> Self

Create a new ragged batch from sequences.

pub fn batch_size(&self) -> usize

Get the number of documents in this batch.

pub fn total_tokens(&self) -> usize

Get the total number of tokens (no padding).

pub fn doc_range(&self, doc_idx: usize) -> Option<Range<usize>>

Get token range for a specific document.

pub fn doc_tokens(&self, doc_idx: usize) -> Option<&[u32]>

Get tokens for a specific document.

pub fn padding_savings(&self) -> f64

Calculate memory saved vs padded batch.

Trait Implementations

impl Clone for RaggedBatch

fn clone(&self) -> RaggedBatch

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for RaggedBatch

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.