pub struct BatchedAttnMetadata {Show 15 fields
pub positions_stacked: DevicePtr,
pub positions_h_stacked: DevicePtr,
pub positions_w_stacked: DevicePtr,
pub slot_stacked: DevicePtr,
pub block_table_ptrs: DevicePtr,
pub seq_len_ptrs: DevicePtr,
pub batch_size: u32,
pub chunk_len: u32,
pub total_tokens: u32,
pub cu_seqlens: DevicePtr,
pub cu_seqlens_host: Vec<i32>,
pub kv_lens: DevicePtr,
pub kv_lens_host: Vec<i32>,
pub max_blocks_per_seq: u32,
pub staged_bytes: usize,
}Expand description
2026-09-25: Device metadata for one kernel-batched prefill chunk over several streams,
built by stage_batched_attn_metadata (model-engine prefill_b/stage_batched.rs) and
passed to prefill_attn_batched_layer and prefill_ssm_batched_layer. Positions and
KV slots are stacked in cu_seqlens order; block tables and sequence lengths are
per-stream pointer arrays.
Fields§
§positions_stacked: DevicePtr2026-09-25: Stacked positions, [total_tokens] u32.
positions_h_stacked: DevicePtr2026-09-25: Stacked MRoPE H positions; the same pointer as positions_stacked
without MRoPE.
positions_w_stacked: DevicePtr2026-09-25: Stacked MRoPE W positions; the same pointer as positions_stacked
without MRoPE.
slot_stacked: DevicePtr2026-09-25: Stacked KV write slots, [total_tokens] i64.
block_table_ptrs: DevicePtr2026-09-25: [batch_size] device pointers, each to one stream’s block table.
seq_len_ptrs: DevicePtr2026-09-25: [batch_size] device pointers, each to one stream’s seq_len.
batch_size: u322026-09-25: Number of streams.
chunk_len: u322026-09-25: The largest per-stream token count. Per-stream counts are in
cu_seqlens.
total_tokens: u322026-09-25: Tokens stacked across all streams (cu_seqlens_host[batch_size]).
cu_seqlens: DevicePtr2026-09-25: [batch_size + 1] i32 prefix sum of per-stream token counts, on the
device.
cu_seqlens_host: Vec<i32>2026-09-25: Host copy of cu_seqlens; the model slices per-stream token ranges
with it.
kv_lens: DevicePtr2026-09-25: Per-stream KV length, [batch_size] i32 on the device: the stream’s
chunk start plus its token count. The batched attention needs it per stream; one
value at the maximum would give a short stream the longest stream’s causal bound
and block count.
kv_lens_host: Vec<i32>2026-09-25: Host copy of kv_lens.
max_blocks_per_seq: u322026-09-25: The largest block count of any stream.
staged_bytes: usize2026-09-25: Bytes this block occupies in scratch, from the offset it was staged at
through the end of kv_lens. The caller advances its scratch cursor by it before
placing the per-layer h_state pointer array (prefill_b/batch_kernel.rs), so an
undercount would put that array over live metadata.