Index

_

_data (ejkernel.ops.config.cache.ConfigCache attribute)

_logger (ejkernel.loggings.ProgressLogger attribute)

A

abstractify() (in module ejkernel.ops.utils.fingerprint) algorithm (ejkernel.kernels._registry.KernelSpec attribute), [1] (ejkernel.modules.base.KernelConfig attribute), [1] align_to() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) allow_autotune (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] allow_heuristics (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] ANY (ejkernel.kernels._registry.Backend attribute) apply_attention_sinks_block() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) apply_causal() (ejkernel.types.mask.MaskInfo method) apply_chunked() (ejkernel.types.mask.MaskInfo method) apply_kv_lengths() (ejkernel.types.mask.MaskInfo method) apply_logits_soft_cap() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) apply_native_sparse_attention() (in module ejkernel.kernels._triton.native_sparse_attention._interface) (in module ejkernel.kernels._xla.native_sparse_attention._interface) apply_sliding_window() (ejkernel.types.mask.MaskInfo method) apply_token_type_ids() (ejkernel.types.mask.MaskInfo method) arch_supports_fp8() (in module ejkernel.utils) args (ejkernel.ops.core.kernel.Invocation attribute), [1] array (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.NumpyMask attribute) as_overlay() (ejkernel.ops.execution.tuning.AutotuningResult method) assert_close() (in module ejkernel.utils) assert_is_supported_dtype() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._utils) Attention (class in ejkernel.modules.operations.attention) attention() (in module ejkernel.kernels._xla.attention._interface) (in module ejkernel.modules.operations.attention)

attention_mask (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) (ejkernel.types.mask.MaskInfo attribute) (ejkernel.types.mask.MaskInfo property) (ejkernel.types.mask.MaskSharding attribute), [1] attention_mask_to_qkv_cu_seqlens() (in module ejkernel.types.mask) attention_pack_from_cu_static() (in module ejkernel.kernels._triton.blocksparse_attention._utilities) (in module ejkernel.kernels._triton.flash_attention._utilities) attention_pack_with_static_shape() (in module ejkernel.kernels._triton.blocksparse_attention._utilities) (in module ejkernel.kernels._triton.flash_attention._utilities) attention_reference() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) attention_reference_custom() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) attention_unpack_with_static_shape() (in module ejkernel.kernels._triton.blocksparse_attention._utilities) (in module ejkernel.kernels._triton.flash_attention._utilities) AttentionConfig (class in ejkernel.modules.operations.configs) AttentionConfigError autotune() (ejkernel.ops.config.selection.Tuner method) (ejkernel.ops.execution.tuning.Autotuner method) (in module ejkernel.ops.execution.tuning) autotune_lowered() (in module ejkernel.ops.execution.offline) autotune_recorded() (in module ejkernel.ops.execution.tuning) AutotuneData (class in ejkernel.ops.execution.tuning) AutotunePolicy (class in ejkernel.ops.config.selection) Autotuner (class in ejkernel.ops.execution.tuning) AutotuningResult (class in ejkernel.ops.execution.tuning) aval_size_bytes() (in module ejkernel.callib._triton_call) avals_to_layouts() (in module ejkernel.callib._triton_call)

B

Backend (class in ejkernel.kernels._registry) backend (ejkernel.kernels._registry.KernelSpec attribute), [1] (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.BaseOperationConfig attribute) baked_in_masks (ejkernel.types.mask.MaskInfo property) barrier_sync() (in module ejkernel.utils) BaseOperationConfig (class in ejkernel.modules.operations.configs) basic_attention_refrence() (in module ejkernel.kernels._triton.blocksparse_attention._utilities) (in module ejkernel.kernels._triton.flash_attention._utilities) batch_axes (ejkernel.ops.core.kernel.Invocation attribute), [1] batch_axis_name (ejkernel.types.mask.MaskInfo attribute) batch_size (ejkernel.types.mask.MaskInfo property) below_or_on_diag() (in module ejkernel.kernels._pallas.tpu.flash_attention._utils) (in module ejkernel.kernels._xla.ring_attention._utils) benchmark() (in module ejkernel.ops.execution.tuning) bias (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) (ejkernel.types.mask.MaskInfo property) binary (ejkernel.callib._triton_call.CompilationResult attribute), [1] block_b (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_d (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.GLAttentionConfig attribute) (ejkernel.modules.operations.configs.LightningAttentionConfig attribute) (ejkernel.modules.operations.configs.NativeSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.RecurrentAttentionConfig attribute) block_k (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.AttentionConfig attribute) (ejkernel.modules.operations.configs.FlashMLAConfig attribute) (ejkernel.modules.operations.configs.GLAttentionConfig attribute) (ejkernel.modules.operations.configs.GroupedMatmulConfig attribute) (ejkernel.modules.operations.configs.LightningAttentionConfig attribute) (ejkernel.modules.operations.configs.NativeSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.RecurrentAttentionConfig attribute) block_k_dkv (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_k_dq (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_k_major (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_k_major_dkv (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_k_major_dq (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_kv (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) block_kv_compute (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) block_kv_dkv (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) block_kv_dkv_compute (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) block_kv_dq (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute)

block_m (ejkernel.modules.operations.configs.GroupedMatmulConfig attribute) block_mask (ejkernel.kernels._pallas.tpu.blocksparse_attention._info.MaskInfo attribute), [1] block_n (ejkernel.modules.operations.configs.GroupedMatmulConfig attribute) block_q (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.AttentionConfig attribute) (ejkernel.modules.operations.configs.FlashMLAConfig attribute) (ejkernel.modules.operations.configs.GLAttentionConfig attribute) (ejkernel.modules.operations.configs.LightningAttentionConfig attribute) (ejkernel.modules.operations.configs.NativeSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.RecurrentAttentionConfig attribute) block_q_dkv (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_q_dq (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_q_major_dkv (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes attribute) block_size (ejkernel.modules.operations.configs.MeanPoolingConfig attribute) (ejkernel.modules.operations.configs.NativeSparseAttentionConfig attribute) blocksize_heads (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_k (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_keys (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_m (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_n (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] BlockSizes (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) (class in ejkernel.kernels._pallas.tpu.flash_attention._utils) blocksparse_attention() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) (in module ejkernel.kernels._triton.blocksparse_attention._interface) (in module ejkernel.kernels._xla.blocksparse_attention._interface) (in module ejkernel.modules.operations.blocksparse_attention) BlockSparseAttention (class in ejkernel.modules.operations.blocksparse_attention) BlockSparseAttentionConfig (class in ejkernel.modules.operations.configs) buffered_pallas_call() (in module ejkernel.callib._pallas_call) bwd_params (ejkernel.modules.operations.configs.BlockSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.FlashAttentionConfig attribute) (ejkernel.modules.operations.configs.RingAttentionConfig attribute) bwd_triton_impl() (in module ejkernel.kernels._triton.mean_pooling._triton_impl_bwd) (in module ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd) (in module ejkernel.kernels._triton.recurrent._triton_impl_bwd) BwdParams (class in ejkernel.ops.utils.datacarrier) bypass_xla_tiling (ejkernel.modules.operations.configs.GroupedMatmulConfig attribute)

C

cache (ejkernel.ops.config.selection.ConfigSelectorChain attribute) cache_miss_fallback (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] calc_bias_strides() (in module ejkernel.kernels._triton.blocksparse_attention._utilities) (in module ejkernel.kernels._triton.flash_attention._utilities) calculate_blocksize_and_wraps() (in module ejkernel.utils) call_key (ejkernel.ops.core.kernel.Invocation property) (ejkernel.ops.execution.tuning.Entry attribute), [1] candidate_cfgs() (ejkernel.modules.operations.attention.Attention method) (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.gated_linear_attention.GLAttention method) (ejkernel.modules.operations.grouped_matmul.GroupedMatmul method) (ejkernel.modules.operations.kernel_delta_attention.KernelDeltaAttention method) (ejkernel.modules.operations.lightning_attention.LightningAttention method) (ejkernel.modules.operations.multi_head_latent_attention.FlashMLA method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.page_attention.PageAttention method) (ejkernel.modules.operations.pooling.MeanPooling method) (ejkernel.modules.operations.prefill_page_attention.PrefillPageAttention method) (ejkernel.modules.operations.ragged_decode_attention.RaggedDecodeAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.recurrent.RecurrentAttention method) (ejkernel.modules.operations.ring_attention.RingAttention method) (ejkernel.modules.operations.scaled_dot_product_attention.ScaledDotProductAttention method) (ejkernel.modules.operations.state_space_v1.StateSpaceV1 method) (ejkernel.modules.operations.state_space_v2.StateSpaceV2 method) (ejkernel.modules.operations.unified_attention.UnifiedAttention method) (ejkernel.ops.core.kernel.Kernel method) candidate_cfgs_gpu() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.ragged_decode_attention.RaggedDecodeAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) candidate_cfgs_shard_map_gpu() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) candidate_cfgs_shard_map_tpu() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.ring_attention.RingAttention method) candidate_cfgs_shard_map_xla() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) candidate_cfgs_tpu() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.ring_attention.RingAttention method) candidate_cfgs_xla() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) causal_mask_baked_in (ejkernel.types.mask.MaskInfo attribute), [1] CausalMask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) cdiv() (in module ejkernel.callib._utils) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) (in module ejkernel.utils) (in module ejkernel.xla_utils.utils)

cfg (ejkernel.ops.execution.tuning.Entry attribute), [1] (ejkernel.ops.execution.tuning.Measurement attribute), [1] check_bool_flag() (in module ejkernel.callib._utils) check_vma (ejkernel.ops.core.kernel.Invocation attribute), [1] choose() (ejkernel.ops.config.selection.ConfigSelectorChain method) (ejkernel.ops.execution.executor.ConfigChooser method), [1] choose_config() (ejkernel.ops.execution.executor.Executor method) chooser (ejkernel.ops.execution.executor.Executor attribute) chunk_global_cumsum() (in module ejkernel.xla_utils.cumsum) chunk_global_cumsum_scalar() (in module ejkernel.xla_utils.cumsum) chunk_global_cumsum_vector() (in module ejkernel.xla_utils.cumsum) chunk_local_cumsum() (in module ejkernel.xla_utils.cumsum) chunk_local_cumsum_scalar() (in module ejkernel.xla_utils.cumsum) chunk_local_cumsum_vector() (in module ejkernel.xla_utils.cumsum) chunk_prefill_size (ejkernel.modules.operations.configs.RaggedPageAttentionv3Config attribute) chunk_size (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.ChunkedCausalMask attribute), [1] chunked_mask_baked_in (ejkernel.types.mask.MaskInfo attribute), [1] chunked_prefill_attention_kernel() (in module ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd) ChunkedCausalMask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) clear() (ejkernel.ops.config.cache.ConfigCache method) cluster_dims (ejkernel.callib._triton_call.CompilationResult attribute), [1] ColorFormatter (class in ejkernel.loggings) CompilationResult (class in ejkernel.callib._triton_call) compile() (ejkernel.ops.execution.executor.Executor method) compile_ttir_inplace() (in module ejkernel.callib._triton_call) compile_ttir_to_hsaco_inplace() (in module ejkernel.callib._triton_call) compile_ttir_to_ptx_inplace() (in module ejkernel.callib._triton_call) complete() (ejkernel.loggings.ProgressLogger method) config_prune_kernel() (in module ejkernel.kernels._triton.flash_attention._triton_impl_bwd) (in module ejkernel.kernels._triton.flash_attention._triton_impl_fwd) ConfigCache (class in ejkernel.ops.config.cache) ConfigChooser (class in ejkernel.ops.execution.executor) ConfigSelectorChain (class in ejkernel.ops.config.selection) CPU (ejkernel.kernels._registry.Backend attribute) create_attention_mask() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) create_bias() (ejkernel.types.mask.MaskInfo method) create_chunked_attention_mask() (ejkernel.types.mask.MaskInfo static method) create_default_executor() (in module ejkernel.modules.base) create_shard_map_wrapper() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.grouped_matmul.GroupedMatmul method) (ejkernel.modules.operations.page_attention.PageAttention method) (ejkernel.modules.operations.prefill_page_attention.PrefillPageAttention method) (ejkernel.modules.operations.ragged_decode_attention.RaggedDecodeAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.ring_attention.RingAttention method) (ejkernel.modules.operations.scaled_dot_product_attention.ScaledDotProductAttention method) (ejkernel.ops.core.kernel.Kernel method) create_sparsity_mask() (in module ejkernel.kernels._triton.blocksparse_attention._mask) create_step_profiler() (in module ejkernel.loggings) cu_seqlens_kv (ejkernel.types.mask.MaskInfo attribute) (ejkernel.types.mask.MaskInfo property) (ejkernel.types.mask.MaskSharding attribute), [1] cu_seqlens_q (ejkernel.types.mask.MaskInfo attribute) (ejkernel.types.mask.MaskInfo property) (ejkernel.types.mask.MaskSharding attribute), [1] cu_seqlens_to_mask() (in module ejkernel.types.mask) CUDA (ejkernel.kernels._registry.Platform attribute)

D

data_next (ejkernel.kernels._pallas.tpu.blocksparse_attention._info.MaskInfo attribute), [1] decode_attn_sequence() (in module ejkernel.kernels._pallas.gpu.ragged_decode_attention._pallas_impl_fwd) decorate() (ejkernel.ops.execution.tuning.FNAutotuner method) default_key_builder_with_sharding() (in module ejkernel.ops.utils.fingerprint) define_sparse_mask_fn() (in module ejkernel.kernels._triton.blocksparse_attention._mask) detect_platform() (in module ejkernel.modules.base) device (ejkernel.ops.execution.tuning.AutotuningResult attribute), [1] device_fingerprint() (in module ejkernel.ops.utils.fingerprint) device_kind() (in module ejkernel.ops.utils.fingerprint)

dropout_seed (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) dtype (ejkernel.callib._utils.ShapeDtype property) dtype_index() (in module ejkernel.utils) DummyStream (class in ejkernel.callib._utils) dumper (ejkernel.ops.config.persistent.PersistentCache attribute) dynamic_init() (ejkernel.types.mask.MaskInfo class method) dynamic_validate_inputs() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64)

E

ejit() (in module ejkernel.callib._ejit) ejkernel.callib._ejit module ejkernel.callib._pallas_call module ejkernel.callib._triton_call module ejkernel.callib._utils module ejkernel.kernels._pallas.gpu.ragged_decode_attention._interface module ejkernel.kernels._pallas.gpu.ragged_decode_attention._pallas_impl_fwd module ejkernel.kernels._pallas.gpu.scaled_dot_product_attention._interface module ejkernel.kernels._pallas.tpu.blocksparse_attention._info module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel module ejkernel.kernels._pallas.tpu.blocksparse_attention._masks module ejkernel.kernels._pallas.tpu.flash_attention._interface module ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_bwd module ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_fwd module ejkernel.kernels._pallas.tpu.flash_attention._utils module ejkernel.kernels._pallas.tpu.grouped_matmul._interface module ejkernel.kernels._pallas.tpu.grouped_matmul._pallas_impl module ejkernel.kernels._pallas.tpu.grouped_matmul._utils module ejkernel.kernels._pallas.tpu.grouped_matmulv2._interface module ejkernel.kernels._pallas.tpu.grouped_matmulv2._pallas_impl module ejkernel.kernels._pallas.tpu.page_attention._interface module ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd module ejkernel.kernels._pallas.tpu.prefill_page_attention._interface module ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd module ejkernel.kernels._pallas.tpu.ragged_decode_attention._interface module ejkernel.kernels._pallas.tpu.ragged_decode_attention._pallas_impl_fwd module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._interface module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._interface module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64 module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils module ejkernel.kernels._pallas.tpu.ring_attention._interface module ejkernel.kernels._pallas.tpu.ring_attention._ring_splash module ejkernel.kernels._registry module ejkernel.kernels._triton.blocksparse_attention._interface module ejkernel.kernels._triton.blocksparse_attention._mask module ejkernel.kernels._triton.blocksparse_attention._triton_impl_bwd module ejkernel.kernels._triton.blocksparse_attention._triton_impl_fwd module ejkernel.kernels._triton.blocksparse_attention._utilities module ejkernel.kernels._triton.flash_attention._interface module ejkernel.kernels._triton.flash_attention._triton_impl_bwd module ejkernel.kernels._triton.flash_attention._triton_impl_fwd module ejkernel.kernels._triton.flash_attention._utilities module ejkernel.kernels._triton.flash_mla._interface module ejkernel.kernels._triton.flash_mla._triton_impl_bwd module ejkernel.kernels._triton.flash_mla._triton_impl_fwd module ejkernel.kernels._triton.flash_mla._utilities module ejkernel.kernels._triton.gla._interface module ejkernel.kernels._triton.lightning_attn._interface module ejkernel.kernels._triton.mean_pooling._interface module ejkernel.kernels._triton.mean_pooling._triton_impl_bwd module ejkernel.kernels._triton.mean_pooling._triton_impl_fwd module ejkernel.kernels._triton.native_sparse_attention._compression module ejkernel.kernels._triton.native_sparse_attention._interface module ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd module ejkernel.kernels._triton.native_sparse_attention._triton_impl_fwd module ejkernel.kernels._triton.native_sparse_attention._utilities module ejkernel.kernels._triton.page_attention._interface module ejkernel.kernels._triton.page_attention._triton_impl_fwd module ejkernel.kernels._triton.ragged_decode_attention._interface module ejkernel.kernels._triton.ragged_decode_attention._triton_impl_fwd module ejkernel.kernels._triton.ragged_page_attention_v2._interface module ejkernel.kernels._triton.ragged_page_attention_v2._triton_impl_fwd module ejkernel.kernels._triton.ragged_page_attention_v3._interface module ejkernel.kernels._triton.ragged_page_attention_v3._triton_impl_fwd module ejkernel.kernels._triton.recurrent._interface module ejkernel.kernels._triton.recurrent._triton_impl_bwd module ejkernel.kernels._triton.recurrent._triton_impl_fwd module ejkernel.kernels._triton.ring_attention._interface module ejkernel.kernels._triton.ring_attention._ring_kernel module ejkernel.kernels._triton.unified_attention._interface module ejkernel.kernels._triton.unified_attention._triton_impl_fwd module ejkernel.kernels._xla.attention._interface module ejkernel.kernels._xla.blocksparse_attention._interface module ejkernel.kernels._xla.flash_attention._interface module ejkernel.kernels._xla.flash_attention._xla_impl_bwd module ejkernel.kernels._xla.flash_attention._xla_impl_fwd module ejkernel.kernels._xla.gla._interface module

ejkernel.kernels._xla.grouped_matmul._interface module ejkernel.kernels._xla.kernel_delta_attention._interface module ejkernel.kernels._xla.kernel_delta_attention._xla_impl_fwd module ejkernel.kernels._xla.lightning_attn._interface module ejkernel.kernels._xla.mean_pooling._interface module ejkernel.kernels._xla.native_sparse_attention._interface module ejkernel.kernels._xla.native_sparse_attention._xla_impl_bwd module ejkernel.kernels._xla.native_sparse_attention._xla_impl_fwd module ejkernel.kernels._xla.page_attention._interface module ejkernel.kernels._xla.page_attention._xla_impl_fwd module ejkernel.kernels._xla.prefill_page_attention._impl module ejkernel.kernels._xla.ragged_decode_attention._interface module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd module ejkernel.kernels._xla.ragged_page_attention_v2._interface module ejkernel.kernels._xla.ragged_page_attention_v2._xla_impl_fwd module ejkernel.kernels._xla.ragged_page_attention_v3._interface module ejkernel.kernels._xla.ragged_page_attention_v3._kernel module ejkernel.kernels._xla.recurrent._interface module ejkernel.kernels._xla.recurrent._xla_impl_bwd module ejkernel.kernels._xla.recurrent._xla_impl_fwd module ejkernel.kernels._xla.ring_attention._interface module ejkernel.kernels._xla.ring_attention._utils module ejkernel.kernels._xla.ring_attention._xla_impl_bwd module ejkernel.kernels._xla.ring_attention._xla_impl_fwd module ejkernel.kernels._xla.scaled_dot_product_attention._interface module ejkernel.kernels._xla.state_space_v1._interface module ejkernel.kernels._xla.state_space_v1._xla_impl_bwd module ejkernel.kernels._xla.state_space_v1._xla_impl_fwd module ejkernel.kernels._xla.state_space_v2._interface module ejkernel.kernels._xla.state_space_v2._xla_impl_bwd module ejkernel.kernels._xla.state_space_v2._xla_impl_fwd module ejkernel.kernels._xla.unified_attention._interface module ejkernel.kernels._xla.unified_attention._xla_impl_fwd module ejkernel.loggings module ejkernel.modules.base module ejkernel.modules.operations.attention module ejkernel.modules.operations.blocksparse_attention module ejkernel.modules.operations.configs module ejkernel.modules.operations.flash_attention module ejkernel.modules.operations.gated_linear_attention module ejkernel.modules.operations.grouped_matmul module ejkernel.modules.operations.kernel_delta_attention module ejkernel.modules.operations.lightning_attention module ejkernel.modules.operations.multi_head_latent_attention module ejkernel.modules.operations.native_sparse_attention module ejkernel.modules.operations.page_attention module ejkernel.modules.operations.pooling module ejkernel.modules.operations.prefill_page_attention module ejkernel.modules.operations.ragged_decode_attention module ejkernel.modules.operations.ragged_page_attention_v2 module ejkernel.modules.operations.ragged_page_attention_v3 module ejkernel.modules.operations.recurrent module ejkernel.modules.operations.ring_attention module ejkernel.modules.operations.scaled_dot_product_attention module ejkernel.modules.operations.state_space_v1 module ejkernel.modules.operations.state_space_v2 module ejkernel.modules.operations.unified_attention module ejkernel.ops.config.cache module ejkernel.ops.config.persistent module ejkernel.ops.config.selection module ejkernel.ops.core.kernel module ejkernel.ops.core.types module ejkernel.ops.execution.batch module ejkernel.ops.execution.executor module ejkernel.ops.execution.offline module ejkernel.ops.execution.profiler module ejkernel.ops.execution.tuning module ejkernel.ops.registry module ejkernel.ops.utils.datacarrier module ejkernel.ops.utils.fingerprint module ejkernel.ops.utils.meta module ejkernel.ops.utils.serialize module ejkernel.types.mask module ejkernel.utils module ejkernel.xla_utils.cumsum module ejkernel.xla_utils.shardings module ejkernel.xla_utils.utils module entries (ejkernel.ops.execution.tuning.AutotuningResult attribute), [1] Entry (class in ejkernel.ops.execution.tuning) event_filter_regex (ejkernel.ops.execution.profiler.Profiler attribute) Executor (class in ejkernel.ops.execution.executor) extinguish_profiler() (in module ejkernel.loggings) extract_labels_from_hlo_text() (in module ejkernel.ops.utils.meta)

F

fastest_config (ejkernel.ops.execution.tuning.AutotuneData property) find_device_plane_ids() (ejkernel.ops.execution.profiler.Profiler static method) find_labels_in_lowered() (in module ejkernel.ops.utils.meta) flash_attention() (in module ejkernel.kernels._pallas.tpu.flash_attention._interface) (in module ejkernel.kernels._triton.flash_attention._interface) (in module ejkernel.kernels._xla.flash_attention._interface) (in module ejkernel.modules.operations.flash_attention) flash_attention_block() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) flash_attention_kernel() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) flash_mla_attention() (in module ejkernel.kernels._triton.flash_mla._interface) flash_mla_attention_call() (in module ejkernel.kernels._triton.flash_mla._interface) FlashAttention (class in ejkernel.modules.operations.flash_attention) FlashAttentionConfig (class in ejkernel.modules.operations.configs) FlashMLA (class in ejkernel.modules.operations.multi_head_latent_attention) FlashMLAConfig (class in ejkernel.modules.operations.configs) flush() (ejkernel.callib._utils.DummyStream method) FNAutotuner (class in ejkernel.ops.execution.tuning) forbid_reautotune (ejkernel.ops.config.selection.ConfigSelectorChain attribute) format() (ejkernel.loggings.ColorFormatter method), [1] forward_kernel() (in module ejkernel.kernels._pallas.gpu.ragged_decode_attention._pallas_impl_fwd)

from_attention_mask() (ejkernel.types.mask.MaskInfo class method) from_cu_seqlens() (ejkernel.types.mask.MaskInfo class method) from_head_minor() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) from_inputs() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask class method) from_json() (in module ejkernel.ops.utils.serialize) from_random() (ejkernel.types.mask.MaskInfo class method) from_segments() (ejkernel.types.mask.MaskInfo class method) from_tuple() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask method) FullMask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) fwd_params (ejkernel.modules.operations.configs.BlockSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.FlashAttentionConfig attribute) (ejkernel.modules.operations.configs.RaggedDecodeAttentionConfig attribute) (ejkernel.modules.operations.configs.RingAttentionConfig attribute) fwd_triton_impl() (in module ejkernel.kernels._triton.mean_pooling._triton_impl_fwd) (in module ejkernel.kernels._triton.native_sparse_attention._triton_impl_fwd) (in module ejkernel.kernels._triton.recurrent._triton_impl_fwd) fwd_with_residuals() (ejkernel.ops.core.kernel.Kernel method) fwd_with_residuals_shard_map() (ejkernel.ops.core.kernel.Kernel method) fwd_with_residuals_shard_map_gpu() (ejkernel.ops.core.kernel.Kernel method) FwdParams (class in ejkernel.ops.utils.datacarrier)

G

generate_block_indices() (in module ejkernel.utils) get() (ejkernel.kernels._registry.KernelRegistry method) (ejkernel.ops.config.cache.ConfigCache method) (ejkernel.ops.config.persistent.PersistentCache method) get_abs_err() (in module ejkernel.utils) get_autotune_configs() (in module ejkernel.kernels._triton.page_attention._triton_impl_fwd) get_cache_dir() (in module ejkernel.callib._utils) get_corrected_named_sharding() (in module ejkernel.xla_utils.shardings) get_cuda_backend() (in module ejkernel.callib._triton_call) get_debug_mode() (in module ejkernel.types.mask) get_default() (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes class method) (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes class method) get_device_name() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_device_platform() (in module ejkernel.ops.utils.fingerprint) get_dtype_bitwidth() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) get_dtype_packing() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) get_empty_sharding() (ejkernel.types.mask.MaskInfo static method) get_err_ratio() (in module ejkernel.utils) get_events_from_plane() (ejkernel.ops.execution.profiler.Profiler method) get_gpu_arch() (in module ejkernel.utils) get_hash_of_lowering() (in module ejkernel.callib._ejit) get_hip_backend() (in module ejkernel.callib._triton_call) get_impl() (ejkernel.modules.operations.attention.Attention method) (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.gated_linear_attention.GLAttention method) (ejkernel.modules.operations.grouped_matmul.GroupedMatmul method) (ejkernel.modules.operations.kernel_delta_attention.KernelDeltaAttention method) (ejkernel.modules.operations.lightning_attention.LightningAttention method) (ejkernel.modules.operations.multi_head_latent_attention.FlashMLA method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.page_attention.PageAttention method) (ejkernel.modules.operations.pooling.MeanPooling method) (ejkernel.modules.operations.prefill_page_attention.PrefillPageAttention method) (ejkernel.modules.operations.ragged_decode_attention.RaggedDecodeAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.recurrent.RecurrentAttention method) (ejkernel.modules.operations.ring_attention.RingAttention method) (ejkernel.modules.operations.scaled_dot_product_attention.ScaledDotProductAttention method) (ejkernel.modules.operations.state_space_v1.StateSpaceV1 method) (ejkernel.modules.operations.state_space_v2.StateSpaceV2 method) (ejkernel.modules.operations.unified_attention.UnifiedAttention method) get_input_shapes() (in module ejkernel.utils) get_invocations() (in module ejkernel.ops.registry)

get_kernel_name() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) get_kv_cache_shape() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) get_logger() (in module ejkernel.loggings) get_lookup_keys() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_lookup_keys_h64() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_mha_cost_estimate() (in module ejkernel.kernels._pallas.tpu.ragged_decode_attention._pallas_impl_fwd) get_min_heads_per_blk() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd) get_min_page_size() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils) get_or_compute_attention_mask() (ejkernel.types.mask.MaskInfo method) get_or_compute_positions() (ejkernel.types.mask.MaskInfo method) get_or_compute_qkv_cu_seqlens() (ejkernel.types.mask.MaskInfo method) get_or_compute_segment_ids() (ejkernel.types.mask.MaskInfo method) get_or_create_triton_kernel() (in module ejkernel.callib._triton_call) get_padded_headsize() (in module ejkernel.utils) get_qkv_masks() (ejkernel.types.mask.MaskInfo method) get_qkv_shardings() (in module ejkernel.utils) get_safe_hash_int() (in module ejkernel.callib._ejit) (in module ejkernel.modules.operations.configs) (in module ejkernel.ops.utils.datacarrier) get_segments_shardings() (in module ejkernel.utils) get_sharding() (in module ejkernel.utils) get_shardings() (ejkernel.types.mask.MaskInfo method) get_simplified_raw_key() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_simplified_raw_key_h64() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_smem_estimate_bytes() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) get_stride() (in module ejkernel.utils) get_strides() (in module ejkernel.utils) get_tpu_generation() (in module ejkernel.utils) get_tpu_version() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_triton_type() (in module ejkernel.callib._triton_call) get_tuned_block_sizes() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_tuned_block_sizes_h64() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) get_vmem_estimate_bytes() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) gla_attention() (in module ejkernel.modules.operations.gated_linear_attention) GLAttention (class in ejkernel.modules.operations.gated_linear_attention) GLAttentionConfig (class in ejkernel.modules.operations.configs) GPU (ejkernel.kernels._registry.Backend attribute) grouped_matmul() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._interface) (in module ejkernel.kernels._pallas.tpu.grouped_matmul._pallas_impl) (in module ejkernel.kernels._pallas.tpu.grouped_matmulv2._pallas_impl) (in module ejkernel.kernels._xla.grouped_matmul._interface) (in module ejkernel.modules.operations.grouped_matmul) grouped_matmulv2() (in module ejkernel.kernels._pallas.tpu.grouped_matmulv2._interface) GroupedMatmul (class in ejkernel.modules.operations.grouped_matmul) GroupedMatmulConfig (class in ejkernel.modules.operations.configs)

H

has_backward_blocks (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes property) (ejkernel.kernels._pallas.tpu.flash_attention._utils.BlockSizes property) hash_fn() (in module ejkernel.callib._ejit) (in module ejkernel.modules.operations.configs) (in module ejkernel.ops.utils.datacarrier) HEAD_DIM_MINOR (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.QKVLayout attribute) heuristic_cfg() (ejkernel.modules.operations.attention.Attention method) (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.gated_linear_attention.GLAttention method) (ejkernel.modules.operations.grouped_matmul.GroupedMatmul method) (ejkernel.modules.operations.kernel_delta_attention.KernelDeltaAttention method) (ejkernel.modules.operations.lightning_attention.LightningAttention method) (ejkernel.modules.operations.multi_head_latent_attention.FlashMLA method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.page_attention.PageAttention method) (ejkernel.modules.operations.pooling.MeanPooling method) (ejkernel.modules.operations.prefill_page_attention.PrefillPageAttention method) (ejkernel.modules.operations.ragged_decode_attention.RaggedDecodeAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.recurrent.RecurrentAttention method) (ejkernel.modules.operations.ring_attention.RingAttention method) (ejkernel.modules.operations.scaled_dot_product_attention.ScaledDotProductAttention method) (ejkernel.modules.operations.state_space_v1.StateSpaceV1 method) (ejkernel.modules.operations.state_space_v2.StateSpaceV2 method) (ejkernel.modules.operations.unified_attention.UnifiedAttention method) (ejkernel.ops.core.kernel.Kernel method)

heuristic_cfg_gpu() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) heuristic_cfg_tpu() (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) hyperparams (ejkernel.ops.execution.tuning.TimingResult attribute)

I

identity_dtype_convert() (in module ejkernel.xla_utils.utils) ignite_profiler() (in module ejkernel.loggings) implementation (ejkernel.kernels._registry.KernelSpec attribute), [1] in_specs (ejkernel.ops.core.kernel.Invocation attribute), [1] inner_decode_tpu() (in module ejkernel.kernels._pallas.tpu.ragged_decode_attention._pallas_impl_fwd) inner_decode_triton() (in module ejkernel.kernels._triton.ragged_decode_attention._triton_impl_fwd) inner_decode_xla() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) Invocation (class in ejkernel.ops.core.kernel) is_cdna() (in module ejkernel.utils)

is_dynamic_mask (ejkernel.kernels._pallas.tpu.blocksparse_attention._info.MaskInfo attribute), [1] is_fp8() (in module ejkernel.utils) is_hip() (in module ejkernel.utils) is_multi_sequence (ejkernel.types.mask.MaskInfo property) is_rdna() (in module ejkernel.utils) is_self_attention() (ejkernel.types.mask.MaskInfo method) is_tpu() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._utils) items() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask method) iters (ejkernel.ops.config.selection.Tuner attribute) (ejkernel.ops.execution.tuning.Autotuner attribute)

K

k (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) k_layout (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) kda() (in module ejkernel.kernels._xla.kernel_delta_attention._interface) kda_attention() (in module ejkernel.modules.operations.kernel_delta_attention) kda_decay() (in module ejkernel.kernels._xla.kernel_delta_attention._interface) Kernel (class in ejkernel.ops.core.kernel) kernel_delta_attention() (in module ejkernel.kernels._xla.kernel_delta_attention._interface) (in module ejkernel.modules.operations.kernel_delta_attention) KernelConfig (class in ejkernel.modules.base) KernelDeltaAttention (class in ejkernel.modules.operations.kernel_delta_attention) KernelDeltaAttentionConfig (class in ejkernel.modules.operations.configs) KernelRegistry (class in ejkernel.kernels._registry) KernelSpec (class in ejkernel.kernels._registry) key_builder (ejkernel.ops.core.kernel.Kernel attribute), [1] keys() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask method) (ejkernel.ops.config.cache.ConfigCache method)

kv (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.SegmentIds attribute), [1] (ejkernel.kernels._pallas.tpu.flash_attention._utils.SegmentIds attribute), [1] (ejkernel.kernels._pallas.tpu.ring_attention._ring_splash.SegmentIds attribute) kv_blocksize (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] kv_len (ejkernel.types.mask.MaskInfo property) kv_lens (ejkernel.types.mask.MaskInfo property) kv_positions (ejkernel.types.mask.MaskInfo attribute), [1] (ejkernel.types.mask.MaskSharding attribute), [1] kv_segment_ids (ejkernel.types.mask.MaskInfo attribute) (ejkernel.types.mask.MaskInfo property) (ejkernel.types.mask.MaskSharding attribute), [1] kvheads_axis_name (ejkernel.types.mask.MaskInfo attribute) kw_strides() (in module ejkernel.utils) kwargs (ejkernel.ops.core.kernel.Invocation attribute), [1]

L

label() (in module ejkernel.ops.utils.meta) labels_to_configs() (in module ejkernel.ops.utils.meta) LazyLogger (class in ejkernel.loggings) left (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LogicalAnd attribute) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LogicalOr attribute) level (ejkernel.loggings.LazyLogger attribute) (ejkernel.loggings.LazyLogger property) lightning_attention() (in module ejkernel.modules.operations.lightning_attention) lightning_attn() (in module ejkernel.kernels._triton.lightning_attn._interface) (in module ejkernel.kernels._xla.lightning_attn._interface) LightningAttention (class in ejkernel.modules.operations.lightning_attention) LightningAttentionConfig (class in ejkernel.modules.operations.configs)

list_algorithms() (ejkernel.kernels._registry.KernelRegistry method) list_implementations() (ejkernel.kernels._registry.KernelRegistry method) llir (ejkernel.callib._triton_call.CompilationResult attribute), [1] load_cached_functions() (in module ejkernel.callib._ejit) load_compiled_fn() (in module ejkernel.callib._ejit) loader (ejkernel.ops.config.persistent.PersistentCache attribute) LocalMask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) LogicalAnd (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) LogicalOr (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) lower_bounds (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask attribute), [1] lower_full_bounds (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask attribute), [1] lse (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute)

M

make_attention_reference() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) make_causal_mask() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) make_chunk_attention_mask() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) make_dummy_rpa_inputs() (in module ejkernel.utils) make_group_metadata() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._pallas_impl) (in module ejkernel.kernels._pallas.tpu.grouped_matmulv2._pallas_impl) make_key() (ejkernel.ops.core.kernel.Invocation method) make_local_attention_mask() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) make_masked_mha_reference() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) make_masked_mqa_reference() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) make_mesh() (in module ejkernel.utils) make_random_mask() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) make_ring_attention() (in module ejkernel.kernels._pallas.tpu.ring_attention._ring_splash) make_splash_mha() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) make_splash_mha_single_device() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) make_splash_mqa() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) make_splash_mqa_single_device() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) manual_sharding_spec() (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.SplashAttentionKernel method) Mask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) mask_next (ejkernel.kernels._pallas.tpu.blocksparse_attention._info.MaskInfo attribute), [1] mask_to_segment_ids() (in module ejkernel.types.mask) MaskInfo (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._info) (class in ejkernel.types.mask) masks (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.MultiHeadMask attribute) MaskSharding (class in ejkernel.types.mask) materialize_attention_mask() (ejkernel.types.mask.MaskInfo method) materialize_segment_ids() (ejkernel.types.mask.MaskInfo method) max_events_per_profile (ejkernel.ops.execution.profiler.Profiler attribute) mean_pooling() (in module ejkernel.kernels._triton.mean_pooling._interface) (in module ejkernel.kernels._xla.mean_pooling._interface) (in module ejkernel.modules.operations.pooling) MeanPooling (class in ejkernel.modules.operations.pooling) MeanPoolingConfig (class in ejkernel.modules.operations.configs) measure() (ejkernel.ops.config.selection.Tuner method) Measurement (class in ejkernel.ops.execution.tuning) measurements (ejkernel.ops.execution.tuning.AutotuneData attribute), [1] merge_kv() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) mesh (ejkernel.ops.core.kernel.Invocation attribute), [1] method (ejkernel.ops.core.kernel.Invocation attribute), [1] mha_reference() (in module ejkernel.kernels._pallas.tpu.flash_attention._utils) mha_reference_bwd() (in module ejkernel.kernels._pallas.tpu.flash_attention._utils) mha_reference_no_custom_vjp() (in module ejkernel.kernels._pallas.tpu.flash_attention._utils) min_duration_ns (ejkernel.ops.execution.profiler.Profiler attribute) mla_attention() (in module ejkernel.modules.operations.multi_head_latent_attention) module ejkernel.callib._ejit ejkernel.callib._pallas_call ejkernel.callib._triton_call ejkernel.callib._utils ejkernel.kernels._pallas.gpu.ragged_decode_attention._interface ejkernel.kernels._pallas.gpu.ragged_decode_attention._pallas_impl_fwd ejkernel.kernels._pallas.gpu.scaled_dot_product_attention._interface ejkernel.kernels._pallas.tpu.blocksparse_attention._info ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel ejkernel.kernels._pallas.tpu.blocksparse_attention._masks ejkernel.kernels._pallas.tpu.flash_attention._interface ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_bwd ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_fwd ejkernel.kernels._pallas.tpu.flash_attention._utils ejkernel.kernels._pallas.tpu.grouped_matmul._interface ejkernel.kernels._pallas.tpu.grouped_matmul._pallas_impl ejkernel.kernels._pallas.tpu.grouped_matmul._utils ejkernel.kernels._pallas.tpu.grouped_matmulv2._interface ejkernel.kernels._pallas.tpu.grouped_matmulv2._pallas_impl ejkernel.kernels._pallas.tpu.page_attention._interface ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd ejkernel.kernels._pallas.tpu.prefill_page_attention._interface ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd ejkernel.kernels._pallas.tpu.ragged_decode_attention._interface ejkernel.kernels._pallas.tpu.ragged_decode_attention._pallas_impl_fwd ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._interface ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._interface ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64 ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils ejkernel.kernels._pallas.tpu.ring_attention._interface ejkernel.kernels._pallas.tpu.ring_attention._ring_splash ejkernel.kernels._registry ejkernel.kernels._triton.blocksparse_attention._interface ejkernel.kernels._triton.blocksparse_attention._mask ejkernel.kernels._triton.blocksparse_attention._triton_impl_bwd ejkernel.kernels._triton.blocksparse_attention._triton_impl_fwd ejkernel.kernels._triton.blocksparse_attention._utilities ejkernel.kernels._triton.flash_attention._interface ejkernel.kernels._triton.flash_attention._triton_impl_bwd ejkernel.kernels._triton.flash_attention._triton_impl_fwd ejkernel.kernels._triton.flash_attention._utilities ejkernel.kernels._triton.flash_mla._interface ejkernel.kernels._triton.flash_mla._triton_impl_bwd ejkernel.kernels._triton.flash_mla._triton_impl_fwd ejkernel.kernels._triton.flash_mla._utilities ejkernel.kernels._triton.gla._interface ejkernel.kernels._triton.lightning_attn._interface ejkernel.kernels._triton.mean_pooling._interface ejkernel.kernels._triton.mean_pooling._triton_impl_bwd ejkernel.kernels._triton.mean_pooling._triton_impl_fwd ejkernel.kernels._triton.native_sparse_attention._compression ejkernel.kernels._triton.native_sparse_attention._interface ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd ejkernel.kernels._triton.native_sparse_attention._triton_impl_fwd ejkernel.kernels._triton.native_sparse_attention._utilities ejkernel.kernels._triton.page_attention._interface ejkernel.kernels._triton.page_attention._triton_impl_fwd ejkernel.kernels._triton.ragged_decode_attention._interface ejkernel.kernels._triton.ragged_decode_attention._triton_impl_fwd ejkernel.kernels._triton.ragged_page_attention_v2._interface ejkernel.kernels._triton.ragged_page_attention_v2._triton_impl_fwd ejkernel.kernels._triton.ragged_page_attention_v3._interface ejkernel.kernels._triton.ragged_page_attention_v3._triton_impl_fwd ejkernel.kernels._triton.recurrent._interface ejkernel.kernels._triton.recurrent._triton_impl_bwd ejkernel.kernels._triton.recurrent._triton_impl_fwd ejkernel.kernels._triton.ring_attention._interface ejkernel.kernels._triton.ring_attention._ring_kernel ejkernel.kernels._triton.unified_attention._interface ejkernel.kernels._triton.unified_attention._triton_impl_fwd ejkernel.kernels._xla.attention._interface ejkernel.kernels._xla.blocksparse_attention._interface ejkernel.kernels._xla.flash_attention._interface ejkernel.kernels._xla.flash_attention._xla_impl_bwd ejkernel.kernels._xla.flash_attention._xla_impl_fwd ejkernel.kernels._xla.gla._interface ejkernel.kernels._xla.grouped_matmul._interface ejkernel.kernels._xla.kernel_delta_attention._interface ejkernel.kernels._xla.kernel_delta_attention._xla_impl_fwd ejkernel.kernels._xla.lightning_attn._interface ejkernel.kernels._xla.mean_pooling._interface ejkernel.kernels._xla.native_sparse_attention._interface ejkernel.kernels._xla.native_sparse_attention._xla_impl_bwd ejkernel.kernels._xla.native_sparse_attention._xla_impl_fwd ejkernel.kernels._xla.page_attention._interface ejkernel.kernels._xla.page_attention._xla_impl_fwd ejkernel.kernels._xla.prefill_page_attention._impl ejkernel.kernels._xla.ragged_decode_attention._interface ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd ejkernel.kernels._xla.ragged_page_attention_v2._interface ejkernel.kernels._xla.ragged_page_attention_v2._xla_impl_fwd ejkernel.kernels._xla.ragged_page_attention_v3._interface ejkernel.kernels._xla.ragged_page_attention_v3._kernel ejkernel.kernels._xla.recurrent._interface ejkernel.kernels._xla.recurrent._xla_impl_bwd ejkernel.kernels._xla.recurrent._xla_impl_fwd ejkernel.kernels._xla.ring_attention._interface ejkernel.kernels._xla.ring_attention._utils ejkernel.kernels._xla.ring_attention._xla_impl_bwd ejkernel.kernels._xla.ring_attention._xla_impl_fwd ejkernel.kernels._xla.scaled_dot_product_attention._interface ejkernel.kernels._xla.state_space_v1._interface ejkernel.kernels._xla.state_space_v1._xla_impl_bwd ejkernel.kernels._xla.state_space_v1._xla_impl_fwd ejkernel.kernels._xla.state_space_v2._interface ejkernel.kernels._xla.state_space_v2._xla_impl_bwd ejkernel.kernels._xla.state_space_v2._xla_impl_fwd ejkernel.kernels._xla.unified_attention._interface ejkernel.kernels._xla.unified_attention._xla_impl_fwd ejkernel.loggings ejkernel.modules.base ejkernel.modules.operations.attention ejkernel.modules.operations.blocksparse_attention ejkernel.modules.operations.configs ejkernel.modules.operations.flash_attention ejkernel.modules.operations.gated_linear_attention ejkernel.modules.operations.grouped_matmul ejkernel.modules.operations.kernel_delta_attention ejkernel.modules.operations.lightning_attention ejkernel.modules.operations.multi_head_latent_attention ejkernel.modules.operations.native_sparse_attention ejkernel.modules.operations.page_attention ejkernel.modules.operations.pooling ejkernel.modules.operations.prefill_page_attention ejkernel.modules.operations.ragged_decode_attention ejkernel.modules.operations.ragged_page_attention_v2 ejkernel.modules.operations.ragged_page_attention_v3 ejkernel.modules.operations.recurrent ejkernel.modules.operations.ring_attention ejkernel.modules.operations.scaled_dot_product_attention ejkernel.modules.operations.state_space_v1 ejkernel.modules.operations.state_space_v2 ejkernel.modules.operations.unified_attention ejkernel.ops.config.cache ejkernel.ops.config.persistent ejkernel.ops.config.selection ejkernel.ops.core.kernel ejkernel.ops.core.types ejkernel.ops.execution.batch ejkernel.ops.execution.executor ejkernel.ops.execution.offline ejkernel.ops.execution.profiler ejkernel.ops.execution.tuning ejkernel.ops.registry ejkernel.ops.utils.datacarrier ejkernel.ops.utils.fingerprint ejkernel.ops.utils.meta ejkernel.ops.utils.serialize ejkernel.types.mask ejkernel.utils ejkernel.xla_utils.cumsum ejkernel.xla_utils.shardings ejkernel.xla_utils.utils

MultiHeadMask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks) MultiPageAsyncCopyDescriptor (class in ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd) (class in ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd) (class in ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd)

N

n_groups (ejkernel.modules.operations.configs.StateSpaceV2Config attribute) name (ejkernel.callib._triton_call.CompilationResult attribute), [1] (ejkernel.loggings.LazyLogger attribute) (ejkernel.loggings.LazyLogger property) (ejkernel.loggings.ProgressLogger attribute) narrow() (in module ejkernel.utils) native_sparse_attention() (in module ejkernel.kernels._triton.native_sparse_attention._interface) (in module ejkernel.kernels._xla.native_sparse_attention._interface) (in module ejkernel.modules.operations.native_sparse_attention) NativeSparseAttention (class in ejkernel.modules.operations.native_sparse_attention) NativeSparseAttentionConfig (class in ejkernel.modules.operations.configs) next_power_of_2() (in module ejkernel.callib._utils) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils) (in module ejkernel.utils) normalize_grid() (in module ejkernel.callib._triton_call) nsa_block_mask() (in module ejkernel.kernels._triton.native_sparse_attention._utilities) nsa_compression() (in module ejkernel.kernels._triton.native_sparse_attention._compression) nsa_compression_bwd() (in module ejkernel.kernels._triton.native_sparse_attention._compression) nsa_compression_fwd() (in module ejkernel.kernels._triton.native_sparse_attention._compression) nsa_topk() (in module ejkernel.kernels._triton.native_sparse_attention._triton_impl_fwd) num_key_splits (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] num_kv_pages_per_block (ejkernel.modules.operations.configs.RaggedPageAttentionv2Config attribute) (ejkernel.modules.operations.configs.RaggedPageAttentionv3Config attribute) num_par_softmax_segments (ejkernel.modules.operations.configs.UnifiedAttentionConfig attribute) num_queries_per_block (ejkernel.modules.operations.configs.RaggedPageAttentionv2Config attribute) (ejkernel.modules.operations.configs.RaggedPageAttentionv3Config attribute) num_splits (ejkernel.modules.operations.configs.PageAttentionConfig attribute) num_stages (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.AttentionConfig attribute) (ejkernel.modules.operations.configs.FlashMLAConfig attribute) (ejkernel.modules.operations.configs.GLAttentionConfig attribute) (ejkernel.modules.operations.configs.GroupedMatmulConfig attribute) (ejkernel.modules.operations.configs.LightningAttentionConfig attribute) (ejkernel.modules.operations.configs.MeanPoolingConfig attribute) (ejkernel.modules.operations.configs.NativeSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.PageAttentionConfig attribute) (ejkernel.modules.operations.configs.PrefillPageAttentionConfig attribute) (ejkernel.modules.operations.configs.RaggedPageAttentionv2Config attribute) (ejkernel.modules.operations.configs.RaggedPageAttentionv3Config attribute) (ejkernel.modules.operations.configs.RecurrentAttentionConfig attribute) (ejkernel.modules.operations.configs.UnifiedAttentionConfig attribute) (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1]

num_warps (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.AttentionConfig attribute) (ejkernel.modules.operations.configs.FlashMLAConfig attribute) (ejkernel.modules.operations.configs.GLAttentionConfig attribute) (ejkernel.modules.operations.configs.GroupedMatmulConfig attribute) (ejkernel.modules.operations.configs.LightningAttentionConfig attribute) (ejkernel.modules.operations.configs.MeanPoolingConfig attribute) (ejkernel.modules.operations.configs.NativeSparseAttentionConfig attribute) (ejkernel.modules.operations.configs.PageAttentionConfig attribute) (ejkernel.modules.operations.configs.PrefillPageAttentionConfig attribute) (ejkernel.modules.operations.configs.RaggedPageAttentionv2Config attribute) (ejkernel.modules.operations.configs.RaggedPageAttentionv3Config attribute) (ejkernel.modules.operations.configs.RecurrentAttentionConfig attribute) (ejkernel.modules.operations.configs.UnifiedAttentionConfig attribute) (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] numeric_gen() (in module ejkernel.utils) NumpyMask (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._masks)

O

o (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) offset (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.CausalMask attribute), [1] (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LocalMask attribute), [1] on_event (ejkernel.ops.config.selection.ConfigSelectorChain attribute) op_id (ejkernel.modules.operations.kernel_delta_attention.KernelDeltaAttention attribute) (ejkernel.ops.core.kernel.Invocation attribute), [1] (ejkernel.ops.core.kernel.Kernel attribute), [1]

op_id_v (ejkernel.ops.execution.tuning.Entry attribute), [1] out_specs (ejkernel.ops.core.kernel.Invocation attribute), [1] overlay_cache (class in ejkernel.ops.config.cache) override_cfg (ejkernel.ops.core.kernel.Invocation attribute), [1]

P

pad_to_block_size() (in module ejkernel.kernels._triton.blocksparse_attention._utilities) page_attention() (in module ejkernel.kernels._pallas.tpu.page_attention._interface) (in module ejkernel.kernels._triton.page_attention._interface) (in module ejkernel.kernels._xla.page_attention._interface) (in module ejkernel.modules.operations.page_attention) PageAttention (class in ejkernel.modules.operations.page_attention) PageAttentionConfig (class in ejkernel.modules.operations.configs) paged_flash_attention_kernel() (in module ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd) paged_flash_attention_kernel_inline_seq_dim() (in module ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd) pages_per_compute_block (ejkernel.modules.operations.configs.PageAttentionConfig attribute) PALLAS (ejkernel.kernels._registry.Platform attribute) parse_profile_from_bytes() (ejkernel.ops.execution.profiler.Profiler static method) partial_mask_blocks (ejkernel.kernels._pallas.tpu.blocksparse_attention._info.MaskInfo attribute), [1] path (ejkernel.ops.config.persistent.PersistentCache attribute) persist_autotune (ejkernel.ops.config.selection.ConfigSelectorChain attribute) persistent (ejkernel.ops.config.selection.ConfigSelectorChain attribute) PersistentCache (class in ejkernel.ops.config.persistent) Platform (class in ejkernel.kernels._registry) platform (ejkernel.kernels._registry.KernelSpec attribute), [1] (ejkernel.modules.base.KernelConfig attribute), [1] (ejkernel.modules.operations.configs.BaseOperationConfig attribute) pmap_with_config() (in module ejkernel.ops.execution.batch) policy (ejkernel.ops.config.selection.ConfigSelectorChain attribute) policy_override (class in ejkernel.ops.config.selection) prefill_page_attention() (in module ejkernel.kernels._pallas.tpu.prefill_page_attention._interface) (in module ejkernel.kernels._xla.prefill_page_attention._impl) (in module ejkernel.modules.operations.prefill_page_attention) PrefillPageAttention (class in ejkernel.modules.operations.prefill_page_attention)

PrefillPageAttentionConfig (class in ejkernel.modules.operations.configs) prefix_filter (ejkernel.ops.execution.profiler.Profiler attribute) PREFIX_FN (ejkernel.ops.execution.tuning.FNAutotuner attribute) prepare() (ejkernel.ops.core.kernel.Kernel method) prepare_chunk_indices() (in module ejkernel.xla_utils.utils) prepare_chunk_offsets() (in module ejkernel.xla_utils.utils) prepare_cu_seqlens_from_mask() (in module ejkernel.xla_utils.utils) prepare_inputs() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) prepare_lens() (in module ejkernel.xla_utils.utils) prepare_lens_from_mask() (in module ejkernel.xla_utils.utils) prepare_outputs() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) prepare_position_ids() (in module ejkernel.xla_utils.utils) prepare_sequence_ids() (in module ejkernel.xla_utils.utils) prepare_token_indices() (in module ejkernel.xla_utils.utils) priority (ejkernel.kernels._registry.KernelSpec attribute), [1] (ejkernel.modules.base.KernelConfig attribute), [1] process_dynamic_mask() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._info) process_dynamic_mask_dkv() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._info) process_mask() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._info) process_mask_dkv() (in module ejkernel.kernels._pallas.tpu.blocksparse_attention._info) profile_time_by_function_id() (ejkernel.ops.execution.profiler.Profiler method) Profiler (class in ejkernel.ops.execution.profiler) ProfilingError ProgressLogger (class in ejkernel.loggings) put() (ejkernel.ops.config.cache.ConfigCache method) (ejkernel.ops.config.persistent.PersistentCache method)

Q

q (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.SegmentIds attribute) (ejkernel.kernels._pallas.tpu.flash_attention._utils.SegmentIds attribute), [1] (ejkernel.kernels._pallas.tpu.ring_attention._ring_splash.SegmentIds attribute) (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) q_attention_mask (ejkernel.types.mask.MaskInfo property) q_blocksize (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] q_layout (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) q_len (ejkernel.types.mask.MaskInfo property) q_lens (ejkernel.types.mask.MaskInfo property) q_position_ids (ejkernel.types.mask.MaskInfo property)

q_positions (ejkernel.types.mask.MaskInfo attribute), [1] (ejkernel.types.mask.MaskSharding attribute), [1] q_segment_ids (ejkernel.types.mask.MaskInfo attribute) (ejkernel.types.mask.MaskInfo property) (ejkernel.types.mask.MaskSharding attribute), [1] q_sequence (ejkernel.kernels._pallas.tpu.blocksparse_attention._info.MaskInfo attribute), [1] qheads_axis_name (ejkernel.types.mask.MaskInfo attribute) qkv_cu_seqlens_to_attention_mask() (in module ejkernel.types.mask) qkv_cu_seqlens_to_qkv_masks() (in module ejkernel.types.mask) qkv_masks_to_cu_seqlens() (in module ejkernel.types.mask) QKVLayout (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) quiet() (in module ejkernel.callib._utils)

R

ragged_decode_attention() (in module ejkernel.kernels._pallas.gpu.ragged_decode_attention._interface) (in module ejkernel.kernels._pallas.tpu.ragged_decode_attention._interface) (in module ejkernel.kernels._triton.ragged_decode_attention._interface) (in module ejkernel.kernels._xla.ragged_decode_attention._interface) (in module ejkernel.modules.operations.ragged_decode_attention) ragged_decode_mqa() (in module ejkernel.kernels._pallas.tpu.ragged_decode_attention._pallas_impl_fwd) ragged_decode_mqa_triton() (in module ejkernel.kernels._triton.ragged_decode_attention._triton_impl_fwd) ragged_decode_mqa_xla() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) ragged_flash_attention_kernel() (in module ejkernel.kernels._pallas.tpu.ragged_decode_attention._pallas_impl_fwd) ragged_flash_attention_xla() (in module ejkernel.kernels._xla.ragged_decode_attention._xla_impl_fwd) ragged_page_attention_kernel() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd) ragged_page_attention_v2() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._interface) (in module ejkernel.kernels._triton.ragged_page_attention_v2._interface) (in module ejkernel.kernels._xla.ragged_page_attention_v2._interface) (in module ejkernel.modules.operations.ragged_page_attention_v2) ragged_page_attention_v3() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._interface) (in module ejkernel.kernels._triton.ragged_page_attention_v3._interface) (in module ejkernel.kernels._xla.ragged_page_attention_v3._interface) (in module ejkernel.modules.operations.ragged_page_attention_v3) ragged_paged_attention() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) ragged_paged_attention_triton() (in module ejkernel.kernels._triton.ragged_page_attention_v3._triton_impl_fwd) ragged_paged_attention_triton_call() (in module ejkernel.kernels._triton.ragged_page_attention_v2._triton_impl_fwd) ragged_paged_attention_triton_call_qblock() (in module ejkernel.kernels._triton.ragged_page_attention_v2._triton_impl_fwd) RaggedDecodeAttention (class in ejkernel.modules.operations.ragged_decode_attention) RaggedDecodeAttentionConfig (class in ejkernel.modules.operations.configs) RaggedPageAttentionv2 (class in ejkernel.modules.operations.ragged_page_attention_v2) RaggedPageAttentionv2Config (class in ejkernel.modules.operations.configs) RaggedPageAttentionv3 (class in ejkernel.modules.operations.ragged_page_attention_v3) RaggedPageAttentionv3Config (class in ejkernel.modules.operations.configs) random_dense() (in module ejkernel.utils) record_invocation() (in module ejkernel.ops.registry) recurrent() (in module ejkernel.kernels._triton.recurrent._interface) (in module ejkernel.kernels._xla.recurrent._interface) recurrent_attention() (in module ejkernel.modules.operations.recurrent) recurrent_gla() (in module ejkernel.kernels._triton.gla._interface) (in module ejkernel.kernels._xla.gla._interface) RecurrentAttention (class in ejkernel.modules.operations.recurrent) RecurrentAttentionConfig (class in ejkernel.modules.operations.configs) ref_paged_attention() (in module ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd) ref_prefill_page_attention() (in module ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd) ref_ragged_page_attention() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd)

ref_ragged_paged_attention() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) ref_ragged_paged_attention_hd64() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) register() (ejkernel.kernels._registry.KernelRegistry method) reorder_sequence() (in module ejkernel.xla_utils.shardings) replace() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask method) (ejkernel.types.mask.MaskInfo method) require_tf (ejkernel.ops.execution.profiler.Profiler attribute) right (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LogicalAnd attribute) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LogicalOr attribute) ring_attention() (in module ejkernel.kernels._pallas.tpu.ring_attention._interface) (in module ejkernel.kernels._triton.ring_attention._interface) (in module ejkernel.kernels._xla.ring_attention._interface) (in module ejkernel.modules.operations.ring_attention) ring_splash_attention() (in module ejkernel.kernels._pallas.tpu.ring_attention._ring_splash) RingAttention (class in ejkernel.modules.operations.ring_attention) RingAttentionConfig (class in ejkernel.modules.operations.configs) RingFlashResiduals (class in ejkernel.kernels._triton.ring_attention._ring_kernel) RingSplashAttentionKernel (class in ejkernel.kernels._pallas.tpu.ring_attention._ring_splash) rmsnorm_eps (ejkernel.modules.operations.configs.StateSpaceV2Config attribute) run() (ejkernel.modules.operations.attention.Attention method) (ejkernel.modules.operations.blocksparse_attention.BlockSparseAttention method) (ejkernel.modules.operations.flash_attention.FlashAttention method) (ejkernel.modules.operations.gated_linear_attention.GLAttention method) (ejkernel.modules.operations.grouped_matmul.GroupedMatmul method) (ejkernel.modules.operations.kernel_delta_attention.KernelDeltaAttention method) (ejkernel.modules.operations.lightning_attention.LightningAttention method) (ejkernel.modules.operations.multi_head_latent_attention.FlashMLA method) (ejkernel.modules.operations.native_sparse_attention.NativeSparseAttention method) (ejkernel.modules.operations.page_attention.PageAttention method) (ejkernel.modules.operations.pooling.MeanPooling method) (ejkernel.modules.operations.prefill_page_attention.PrefillPageAttention method) (ejkernel.modules.operations.ragged_decode_attention.RaggedDecodeAttention method) (ejkernel.modules.operations.ragged_page_attention_v2.RaggedPageAttentionv2 method) (ejkernel.modules.operations.ragged_page_attention_v3.RaggedPageAttentionv3 method) (ejkernel.modules.operations.recurrent.RecurrentAttention method) (ejkernel.modules.operations.ring_attention.RingAttention method) (ejkernel.modules.operations.scaled_dot_product_attention.ScaledDotProductAttention method) (ejkernel.modules.operations.state_space_v1.StateSpaceV1 method) (ejkernel.modules.operations.state_space_v2.StateSpaceV2 method) (ejkernel.modules.operations.unified_attention.UnifiedAttention method) (ejkernel.ops.core.kernel.Kernel method) run_shard_map() (ejkernel.ops.core.kernel.Kernel method) run_shard_map_gpu() (ejkernel.ops.core.kernel.Kernel method)

S

safe_autotune() (in module ejkernel.utils) save_compiled_fn() (in module ejkernel.callib._ejit) scaled_dot_product_attention() (in module ejkernel.kernels._pallas.gpu.scaled_dot_product_attention._interface) (in module ejkernel.kernels._xla.scaled_dot_product_attention._interface) (in module ejkernel.modules.operations.scaled_dot_product_attention) ScaledDotProductAttention (class in ejkernel.modules.operations.scaled_dot_product_attention) ScaledDotProductAttentionConfig (class in ejkernel.modules.operations.configs) seconds (ejkernel.ops.execution.tuning.Measurement attribute), [1] segment_ids_to_mask() (in module ejkernel.types.mask) segment_ids_to_qkv_masks() (in module ejkernel.types.mask) SegmentIds (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) (class in ejkernel.kernels._pallas.tpu.flash_attention._utils) (class in ejkernel.kernels._pallas.tpu.ring_attention._ring_splash) select_input_dtype() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._utils) SEQ_MINOR (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.QKVLayout attribute) seq_threshold_3d (ejkernel.modules.operations.configs.UnifiedAttentionConfig attribute) sequence_axis_name (ejkernel.types.mask.MaskInfo attribute) set_debug_mode() (in module ejkernel.types.mask) shape (ejkernel.callib._utils.ShapeDtype property) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.FullMask property) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LogicalAnd property) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LogicalOr property) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.Mask property) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.MultiHeadMask property) (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.NumpyMask property) (ejkernel.types.mask.MaskInfo property) ShapeDtype (class in ejkernel.callib._utils) sharding_fingerprint() (in module ejkernel.ops.utils.fingerprint) shared_mem_bytes (ejkernel.callib._triton_call.CompilationResult attribute), [1]

short_hash() (in module ejkernel.ops.utils.fingerprint) simplify_key() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils) size() (ejkernel.ops.config.cache.ConfigCache method) sliding_window_baked_in (ejkernel.types.mask.MaskInfo attribute), [1] smart_compile() (in module ejkernel.callib._ejit) SparseMask (class in ejkernel.kernels._triton.blocksparse_attention._mask) SplashAttentionKernel (class in ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel) stable_json() (in module ejkernel.ops.utils.fingerprint) stamp (ejkernel.ops.core.kernel.Invocation attribute), [1] stamp_prefix (ejkernel.ops.execution.executor.Executor attribute) start() (ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd.MultiPageAsyncCopyDescriptor method) (ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd.MultiPageAsyncCopyDescriptor method) (ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd.MultiPageAsyncCopyDescriptor method) start_time (ejkernel.loggings.ProgressLogger attribute) state_space_v1() (in module ejkernel.kernels._xla.state_space_v1._interface) (in module ejkernel.modules.operations.state_space_v1) state_space_v2() (in module ejkernel.kernels._xla.state_space_v2._interface) (in module ejkernel.modules.operations.state_space_v2) StateSpaceV1 (class in ejkernel.modules.operations.state_space_v1) StateSpaceV1Config (class in ejkernel.modules.operations.configs) StateSpaceV2 (class in ejkernel.modules.operations.state_space_v2) StateSpaceV2Config (class in ejkernel.modules.operations.configs) static_validate_inputs() (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd) (in module ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64) (in module ejkernel.kernels._xla.ragged_page_attention_v3._kernel) strides_from_shape() (in module ejkernel.callib._utils) (in module ejkernel.utils) supports_bfloat16_matmul() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._utils)

T

t_mean (ejkernel.ops.execution.tuning.TimingResult attribute) t_std (ejkernel.ops.execution.tuning.TimingResult attribute) TimingResult (class in ejkernel.ops.execution.tuning) to_dtype() (ejkernel.types.mask.MaskInfo method) to_json() (in module ejkernel.ops.utils.serialize) to_tuple() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask method) token_type_ids_baked_in (ejkernel.types.mask.MaskInfo attribute), [1] TPU (ejkernel.kernels._registry.Backend attribute) tpu_generation() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._utils) tpu_kind() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._utils) transposed_grouped_matmul() (in module ejkernel.kernels._pallas.tpu.grouped_matmul._pallas_impl) (in module ejkernel.kernels._pallas.tpu.grouped_matmulv2._pallas_impl) tree_flatten() (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.SplashAttentionKernel method) (ejkernel.kernels._pallas.tpu.ring_attention._ring_splash.RingSplashAttentionKernel method) (ejkernel.types.mask.MaskInfo method)

tree_unflatten() (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.SplashAttentionKernel class method) (ejkernel.kernels._pallas.tpu.ring_attention._ring_splash.RingSplashAttentionKernel class method) (ejkernel.types.mask.MaskInfo class method) TRITON (ejkernel.kernels._registry.Platform attribute) triton_call() (in module ejkernel.callib._triton_call) triton_kernel_call_abstract_eval() (in module ejkernel.callib._triton_call) triton_kernel_call_lowering() (in module ejkernel.callib._triton_call) triton_kernel_call_raise_on_jvp() (in module ejkernel.callib._triton_call) triton_kernel_call_raise_on_vmap() (in module ejkernel.callib._triton_call) ttgir (ejkernel.callib._triton_call.CompilationResult attribute), [1] tune() (ejkernel.ops.execution.tuning.FNAutotuner method) Tuner (class in ejkernel.ops.config.selection) tuner (ejkernel.ops.config.selection.ConfigSelectorChain attribute)

U

unified_attention() (in module ejkernel.kernels._triton.unified_attention._interface) (in module ejkernel.kernels._xla.unified_attention._interface) (in module ejkernel.modules.operations.unified_attention) unified_attention_triton() (in module ejkernel.kernels._triton.unified_attention._triton_impl_fwd) UnifiedAttention (class in ejkernel.modules.operations.unified_attention) UnifiedAttentionConfig (class in ejkernel.modules.operations.configs)

update() (ejkernel.loggings.ProgressLogger method) update_simple() (ejkernel.loggings.ProgressLogger method) upper_bounds (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask attribute), [1] upper_full_bounds (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask attribute), [1] use_fused_bwd_kernel (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) use_gated_rmsnorm (ejkernel.modules.operations.configs.StateSpaceV2Config attribute) use_tty (ejkernel.loggings.ProgressLogger attribute)

V

v (ejkernel.kernels._triton.ring_attention._ring_kernel.RingFlashResiduals attribute) v_layout (ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel.BlockSizes attribute) validate_backward (ejkernel.ops.config.selection.AutotunePolicy attribute) validate_signatures() (ejkernel.kernels._registry.KernelRegistry method) values() (ejkernel.kernels._triton.blocksparse_attention._mask.SparseMask method) verbose (ejkernel.ops.execution.profiler.Profiler attribute)

version (ejkernel.modules.operations.flash_attention.FlashAttention attribute) (ejkernel.ops.core.kernel.Kernel attribute), [1] visualize() (ejkernel.types.mask.MaskInfo method) vjp() (ejkernel.ops.core.kernel.Kernel method) vjp_shard_map() (ejkernel.ops.core.kernel.Kernel method) vjp_shard_map_gpu() (ejkernel.ops.core.kernel.Kernel method) vmap_with_config() (in module ejkernel.ops.execution.batch)

W

wait() (ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd.MultiPageAsyncCopyDescriptor method) wait_and_get_loaded() (ejkernel.kernels._pallas.tpu.page_attention._pallas_impl_fwd.MultiPageAsyncCopyDescriptor method) (ejkernel.kernels._pallas.tpu.prefill_page_attention._pallas_impl_fwd.MultiPageAsyncCopyDescriptor method)

warmup (ejkernel.ops.config.selection.Tuner attribute) (ejkernel.ops.execution.tuning.Autotuner attribute) window_size (ejkernel.kernels._pallas.tpu.blocksparse_attention._masks.LocalMask attribute), [1] write() (ejkernel.callib._utils.DummyStream method)

X

XLA (ejkernel.kernels._registry.Platform attribute)