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
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], page_stride: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=page_stride.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,gammaand the output (inferred). - freq_dtype (
DType): Element dtype offreqs_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. - page_stride (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=page_stride.static_spec]): Page-to-page distance in elements; -1 when packed. - 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.