Struct Predictor 

pub struct Predictor { /* private fields */ }

Implementations

impl Predictor

pub fn new( ctx: &CudaCtx, dims: PredictorDims, projection_seed: u64, ) -> Result<Self>

pub fn new_on_stream( stream: u64, dims: PredictorDims, projection_seed: u64, ) -> Result<Self>

pub fn project_q(&self, ctx: &CudaCtx, q: u64, q_proj: u64) -> Result<()>

q device pointer to [num_q_heads, head_dim] BF16. q_proj device pointer to [num_q_heads, r] BF16 output.

pub fn project_q_on_stream( &self, stream: u64, q: u64, q_proj: u64, ) -> Result<()>

Stream-only variant for production callers that already own a CUDA context (and therefore don’t need the test-side CudaCtx wrapper).

pub fn project_kv_block( &self, ctx: &CudaCtx, layer: usize, block_id: usize, k_block: u64, ) -> Result<()>

Update A_g for a single (layer, block_id) slot from the K data the caller just wrote into the KV cache. k_block device ptr to [block_size, num_kv_heads, head_dim] BF16.

pub fn project_kv_block_on_stream( &self, stream: u64, layer: usize, block_id: usize, k_block: u64, ) -> Result<()>

pub fn score_blocks( &self, ctx: &CudaCtx, q_proj: u64, k_lr_seq: u64, scores_out: u64, num_active_blocks: usize, ) -> Result<()>

Score num_active_blocks (already-laid-out) blocks for the current layer. q_proj device ptr to [num_q_heads, r] BF16. a_g_seq is a device ptr to the active sequence’s per-block anchors at this layer ([num_active_blocks, num_kv_heads, r] BF16). scores_out is a device ptr to a [num_active_blocks] f32 buffer.

pub fn score_blocks_on_stream( &self, stream: u64, q_proj: u64, k_lr_seq: u64, scores_out: u64, num_active_blocks: usize, ) -> Result<()>

pub fn dims(&self) -> PredictorDims

pub fn a_g_dev_ptr(&self) -> u64