For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

QSACompressKeysPaged

struct QSACompressKeysPaged

Pools a paged group of raw QSA index keys into dense block keys.

Tensor shapes: - block_keys : [batch * max_blocks, head_dim] (OUT) - gamma : [head_dim] - freqs_cis : [positions, rotary_dim] - key_counts : [batch] (int32) - kv_blocks : [num_pages, 2, num_layers, page_size, 1, head_dim] (mut in) - cache_lengths : [batch] (uint32) - kv_lookup_table : [batch, max_pages] (uint32) - max_prompt_length : [1] (uint32) - max_cache_length : [1] (uint32)

key_counts is each sequence's total raw-key count including this forward's tokens, so it is cache_lengths + this step's row lengths rather than either alone. max_blocks is block_keys.dim(0) // batch, the same per-sequence stride qsa_block_score re-derives.

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[dtype: DType, freq_dtype: DType, target: StringSpan[ImmStaticOrigin], head_dim: Int, rotary_dim: Int, ratio: Int, eps: StringSpan[ImmStaticOrigin]](block_keys: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=block_keys.static_spec], gamma: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gamma.static_spec], freqs_cis: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=freqs_cis.static_spec], key_counts: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=key_counts.static_spec], kv_blocks: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=kv_blocks.static_spec], cache_lengths: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_lengths.static_spec], kv_lookup_table: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=kv_lookup_table.static_spec], max_prompt_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_prompt_length.static_spec], max_cache_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_cache_length.static_spec], layer_idx: UInt32, ctx: DeviceContext)

Runs qsa_compress_keys for one layer.

Parameters:

• ​dtype (DType): Element dtype of the keys, gamma and the output (inferred).
• ​freq_dtype (DType): Element dtype of freqs_cis (inferred).
• ​target (StringSpan[ImmStaticOrigin]): Compilation target.
• ​head_dim (Int): indexer_head_dim.
• ​rotary_dim (Int): Channels the rotation covers.
• ​ratio (Int): indexer_compress_ratio.
• ​eps (StringSpan[ImmStaticOrigin]): rms_norm_eps, string-encoded because the extensibility bridge takes no float parameters.

Args:

• ​block_keys (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=block_keys.static_spec]): Dense block keys, at the per-sequence stride.
• ​gamma (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gamma.static_spec]): k_layernorm.weight.
• ​freqs_cis (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=freqs_cis.static_spec]): Rotary cos/sin-pair table, indexed by position.
• ​key_counts (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=key_counts.static_spec]): Raw keys held per sequence, this step included.
• ​kv_blocks (ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=kv_blocks.static_spec]): Paged blocks of the indexer cache group.
• ​cache_lengths (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_lengths.static_spec]): Per-sequence cached-key count.
• ​kv_lookup_table (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=kv_lookup_table.static_spec]): Per-sequence page table.
• ​max_prompt_length (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_prompt_length.static_spec]): Max new query tokens this step.
• ​max_cache_length (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_cache_length.static_spec]): Max cached context this step.
• ​layer_idx (UInt32): Layer index into the indexer cache.
• ​ctx (DeviceContext): Device context.

Raises:

Error: If the operand shapes disagree with the parameters.