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

Device-side inputs to a selection pass. Every one is owned by the caller; this module allocates nothing but its own scratch.

Fields

k_normed: DevicePtr

[seq, index_head_dim] BF16 — indexer keys, already LayerNorm’d. 🪤 indexer.k_norm is an nn.LayerNorm with a bias, not an RMSNorm.

gate: DevicePtr

[seq, index_head_dim] BF16 — index_kpool_compress_gate projection.

valid: DevicePtr

[seq] u8 — per-key validity.

ape: DevicePtr

[index_kpool, index_head_dim] f32 — the APE table. 🪤 BF16 on disk, f32 to the kernel; the loader must upconvert. This is the #347 dtype-mismatch class, so the width is stated here rather than inferred.

q: DevicePtr

[q_rows, index_heads, index_head_dim] f32.

weights: DevicePtr

[q_rows, index_heads] f32, already carrying the index_heads^-0.5 factor — dsa_index_scores does not apply it.

q_pos: DevicePtr

[q_rows] i32 — absolute position of each query.

q_mask: DevicePtr

[q_rows] u8 — a zero row selects nothing and stays all -1.

first_key: i32

Index of the first valid key; pooling starts here so left padding is skipped.

geom_dev: DevicePtr

[5] i32 DEVICE geometry (dsa_write_geom), or NULL for the scalar path.

🔴 Non-null is what makes a captured decode step replay correctly: S, the pool count, the padded sort axis and select_k all grow with the context, and a graph freezes every scalar it was captured with. Decode-only — see Replay below.

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