Struct TwoWayTransformerCompiled 

pub struct TwoWayTransformerCompiled {
    pub max_q_n: usize,
    pub k_n: usize,
    pub embed_dim: usize,
    pub num_heads: usize,
    pub num_layers: usize,
    pub masked: bool,
    /* private fields */
}

Fields

max_q_n: usize

Compiled query-token slots (base + MAX_SPARSE when masked).

k_n: usize

embed_dim: usize

num_heads: usize

num_layers: usize

masked: bool

When true, pass per-attention masks and pad queries to max_q_n.

Implementations

impl TwoWayTransformerCompiled

pub fn compile( spec: &TwoWayTransformerSpec, q_n: usize, k_n: usize, device: Device, ) -> Result<TwoWayTransformerCompiled, Error>

pub fn compile_with_profile( spec: &TwoWayTransformerSpec, q_n: usize, k_n: usize, device: Device, masked: bool, profile: &CompileProfile, ) -> Result<TwoWayTransformerCompiled, Error>

pub fn compile_with_sparse_slots( spec: &TwoWayTransformerSpec, base_q_n: usize, k_n: usize, device: Device, ) -> Result<TwoWayTransformerCompiled, Error>

base_q_n + up to MAX_SPARSE_PROMPT_TOKENS padded query slots (masked attention).

pub fn compile_with_sparse_slots_profile( spec: &TwoWayTransformerSpec, base_q_n: usize, k_n: usize, device: Device, profile: &CompileProfile, ) -> Result<TwoWayTransformerCompiled, Error>

pub fn fill_attn_mask( out: &mut [f32], num_heads: usize, max_q: usize, max_k: usize, active_q: usize, active_k: usize, )

Fill [1, H, max_q, max_k] mask (1 = attend, 0 = mask out).

pub fn nchw_to_seq(nchw: &[f32], e: usize, h: usize, w: usize) -> Vec<f32>

NCHW [E, H, W] → sequence [H*W, E] (same layout as host two_way_transformer_forward).

pub fn run_nchw( &mut self, tokens: &[f32], image_nchw: &[f32], image_pe_nchw: &[f32], grid: usize, ) -> Result<(Vec<f32>, Vec<f32>), Error>

tokens: [q_n, E]; image tensors NCHW [E, g, g].

pub fn run_nchw_masked( &mut self, tokens: &[f32], active_q_n: usize, image_nchw: &[f32], image_pe_nchw: &[f32], grid: usize, ) -> Result<(Vec<f32>, Vec<f32>), Error>

Padded-query path: active_q_n real tokens, rest masked out.

pub fn run( &mut self, tokens: &[f32], image_seq: &[f32], image_pe_seq: &[f32], ) -> Result<(Vec<f32>, Vec<f32>), Error>

tokens / query_pe: [q_n, E]; image_seq / image_pe_seq: [k_n, E] row-major.

pub fn run_masked( &mut self, tokens_padded: &[f32], active_q_n: usize, image_seq: &[f32], image_pe_seq: &[f32], ) -> Result<(Vec<f32>, Vec<f32>), Error>