Struct DsaSelectInputs 

pub struct DsaSelectInputs {
    pub k_normed: DevicePtr,
    pub gate: DevicePtr,
    pub valid: DevicePtr,
    pub ape: DevicePtr,
    pub q: DevicePtr,
    pub weights: DevicePtr,
    pub q_pos: DevicePtr,
    pub q_mask: DevicePtr,
    pub first_key: i32,
    pub geom_dev: DevicePtr,
}

2026-09-25: Device-side inputs to a selection pass, all owned by the caller.

Fields

k_normed: DevicePtr

2026-09-25: [seq, index_head_dim] BF16 indexer keys, after the k_norm LayerNorm.

gate: DevicePtr

2026-09-25: [seq, index_head_dim] BF16 index_kpool_compress_gate projection.

valid: DevicePtr

2026-09-25: [seq] u8 per-key validity.

ape: DevicePtr

2026-09-25: [index_kpool, index_head_dim] f32 APE table. The checkpoint stores BF16; build_dsa_weights uploads it as f32.

q: DevicePtr

2026-09-25: [q_rows, index_heads, index_head_dim] f32.

weights: DevicePtr

2026-09-25: [q_rows, index_heads] f32, already carrying the index_heads^-0.5 factor, which dsa_index_scores does not apply.

q_pos: DevicePtr

2026-09-25: [q_rows] i32 absolute position of each query.

q_mask: DevicePtr

2026-09-25: [q_rows] u8; a row whose entry is 0 selects nothing and stays all -1.

first_key: i32

2026-09-25: Index of the first valid key; pooling starts here, so left padding is skipped.

geom_dev: DevicePtr

2026-09-25: [5] i32 device geometry written by dsa_write_geom, or NULL for the scalar path.

A captured graph fixes every scalar argument, while S, the pool counts, the tile and select_k grow with the context; the kernels read them from here when it is non-null. Only with q_rows == 1: select_tokens refuses a ceiling launch otherwise.

Trait Implementations

impl Clone for DsaSelectInputs

fn clone(&self) -> DsaSelectInputs

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for DsaSelectInputs

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

Formats the value using the given formatter.

impl Copy for DsaSelectInputs