Index A | B | C | D | E | F | G | H | I | K | L | M | N | P | R | S | T | U | V A additive_bias_mod() (in module attnax.kernels) alibi_mod() (in module attnax.kernels) alibi_slopes() (in module attnax.kernels) allocate_blocks() (in module attnax) append_kv() (in module attnax) apply_rope() (in module attnax) AttentionFn (class in attnax.kernels) AttentionType (class in attnax) B block_size (attnax.PagedKVCache property) block_table (attnax.PagedKVCache attribute) C causal_mod() (in module attnax.kernels) combine_masks() (in module attnax) compose_score_mods() (in module attnax.kernels) cos_sin_for_positions() (attnax.RotaryEmbedding method) create_norm() (in module attnax) D DecoderBlock (class in attnax) document_mask_mod() (in module attnax.kernels) E EncoderBlock (class in attnax) F FeedForward (class in attnax) FfActivation (in module attnax) FLASH (attnax.AttentionType attribute) flash_attention() (in module attnax.kernels) G gather_kv() (in module attnax) H head_dim (attnax.PagedKVCache property) I init_decoder_kv_caches() (in module attnax) init_decoder_kv_caches_from_config() (in module attnax) init_kv_layer_cache() (in module attnax) init_paged_kv_cache() (in module attnax) K key_pool (attnax.PagedKVCache attribute) KVLayerCache (class in attnax) L LearnedPositionalEmbedding (class in attnax) LINEAR (attnax.AttentionType attribute) linear_attention() (in module attnax.kernels) LITE (attnax.AttentionType attribute) M make_causal_mask() (in module attnax) make_document_mask() (in module attnax) make_padding_mask() (in module attnax) make_sliding_window_mask() (in module attnax) mask_mod_to_boolean_mask() (in module attnax.kernels) mask_mod_to_score_mod() (in module attnax.kernels) MaskMod (class in attnax.kernels) max_len (attnax.KVLayerCache property) MEMORY_EFFICIENT (attnax.AttentionType attribute) memory_efficient_attention() (in module attnax.kernels) MixtureOfExperts (class in attnax) MultiHeadAttention (class in attnax) N NormKind (in module attnax) num_blocks (attnax.PagedKVCache property) num_kv_heads (attnax.PagedKVCache property) P paged_attention() (in module attnax.kernels) PagedKVCache (class in attnax) PALLAS_FLASH (attnax.AttentionType attribute) pallas_flash_attention() (in module attnax.kernels) PatchEmbedding (class in attnax) Pool (in module attnax) PosEmbType (in module attnax) PositionalEncoding (class in attnax) prefix_lm_mod() (in module attnax.kernels) R ring_attention() (in module attnax.kernels) ring_attention_reference() (in module attnax.kernels) RMSNorm (class in attnax) rope_cos_sin_from_positions() (in module attnax) rope_cos_sin_table() (in module attnax) rope_inv_freq() (in module attnax) RotaryEmbedding (class in attnax) S ScoreMod (class in attnax.kernels) seq_lengths (attnax.PagedKVCache attribute) sliding_window_mod() (in module attnax.kernels) STANDARD (attnax.AttentionType attribute) standard_attention() (in module attnax.kernels) T TokenEmbedding (class in attnax) TransformerConfig (class in attnax) TransformerDecoder (class in attnax) TransformerEncoder (class in attnax) U update_kv_layer_cache() (in module attnax) V value_pool (attnax.PagedKVCache attribute) VisionTransformer (class in attnax) VisionTransformerConfig (class in attnax)
