IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

Struct_mla_prefill_sparse_paged

struct Struct_mla_prefill_sparse_paged

Registers the mo.mla.prefill.sparse.paged graph op with the graph compiler.

Implemented traitsโ€‹

AnyType, ImplicitlyDeletable

Methodsโ€‹

executeโ€‹

static def execute[dtype: DType, cache_dtype: DType, //, target: StringSlice[ImmStaticOrigin], indices_stride: Int](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], q: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.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, input_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec], sparse_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=sparse_indices.static_spec], topk_lengths: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=topk_lengths.static_spec], attn_sink: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=attn_sink.static_spec], scale: Float32, context: DeviceContext)

Was this page helpful?