Struct BatchedAttnMetadata 

pub struct BatchedAttnMetadata {

    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,
}

Q12 batched-prefill device-side metadata.

The single-stream AttnMetadataDev collapses per-stream pointers into concrete device pointers because there’s only one stream. For Q12 we dispatch N concurrent prefilling streams through one batched kernel, and the kernel takes:

• stacked positions / slot tables (one big buffer with all streams’ data concatenated in cu_seqlens order), and
• per-stream pointer arrays for block_table / seq_len / h_state.

Built once per prefill_batch_chunk_dispatch call by stage_batched_attn_metadata; threaded through the model-level per-layer batched dispatch (prefill_attn_batched_layer, prefill_ssm_batched_layer) — see model/trait_impl/prefill_b/batch.rs.

Fields

positions_stacked: DevicePtr

Stacked positions across all streams: [total_tokens] u32 at this address. For MRoPE interleaved this is the temporal (T) stream.

positions_h_stacked: DevicePtr

MRoPE H position stream, stacked. Equal to positions_stacked when MRoPE is disabled.

positions_w_stacked: DevicePtr

MRoPE W position stream, stacked. Equal to positions_stacked when MRoPE is disabled.

slot_stacked: DevicePtr

Stacked slot indices for KV writes: [total_tokens] i64.

block_table_ptrs: DevicePtr

Per-stream block_table pointer array: [batch_size] of DevicePtr, each element pointing to a stream’s chunked-prefill block_table. Used by prefill_attention_paged_*_batched kernels.

seq_len_ptrs: DevicePtr

Per-stream seq_len pointer array: [batch_size] of DevicePtr.

batch_size: u32

Number of batched streams.

chunk_len: u32

Per-stream chunk_len. In the legacy same-length path this is uniform; in the VARLEN path (cu_seqlens populated) it is the MAX per-stream length (retained only for buffer-bound/debug use — per-stream lengths come from cu_seqlens).

total_tokens: u32

Total tokens stacked across streams. Legacy: batch_size * chunk_len. VARLEN: Σ per-stream lengths (= cu_seqlens_host[batch_size]).

cu_seqlens: DevicePtr

VARLEN geometry: [batch_size+1] i32 prefix-sum of per-request token counts, on device (read by the GDN kernel + FlashInfer). DevicePtr::NULL in the legacy same-length path (callers fall back to b*chunk_len).

cu_seqlens_host: Vec<i32>

Host copy of cu_seqlens ([batch_size+1] i32) — FlashInfer’s PrefillPlan dereferences the indptr on the CPU, and per-request slice offsets are computed host-side. Empty in the legacy path.

kv_lens: DevicePtr

VARLEN geometry: per-stream KV length [batch_size] i32 on device, kv_lens[b] = chunk_start + per-stream token count. The batched paged attention kernels need this per stream: a single scalar at the MAX makes short streams index their block_table past the blocks they actually own, and applies the wrong causal bound. DevicePtr::NULL in the legacy same-length path (kernels fall back to the scalar kv_len).

kv_lens_host: Vec<i32>

Host copy of kv_lens. Empty in the legacy path.

max_blocks_per_seq: u32

Maximum block_table length across the batch (kernel uses for bounds checking; per-stream block_table reads via the pointer array dereference).

staged_bytes: usize

Exact byte footprint of this metadata block within the scratch buffer (from scratch_offset_bytes to the end of seq_len_ptrs). SSOT for the caller’s scratch-cursor advance — the per-SSM-layer h_state_ptrs slot is placed at scratch_cursor + staged_bytes, so an under-estimate here would overwrite the live slot_stacked array with device pointers and produce wild KV-cache slots (#110 bug #2).