Mojo function

flare_mla_prefill_dispatch

flare_mla_prefill_dispatch[k_t: MHAOperand, v_t: MHAOperand, k_rope_t: MHAOperand, mask_t: MHAMask, dtype: DType, output_type: DType, q_layout: Layout, //, kv_num_heads: Int, q_depth: Int = 192, cache_depth: Int = 576, config: MHAConfig[dtype] = MHAConfig(SIMD(Int.__init__[IntTuple](q_layout.shape[(q_layout.rank() - 2)])), SIMD(Int.__init__[IntTuple](q_layout.shape[(q_layout.rank() - 1)])), Optional(None), Optional(None), Optional(None), Optional(None), Optional(None), 4, 1, FlashAttentionAlgorithm(-1), TensorMapSwizzle.SWIZZLE_128B), _ndbuffer_mha_operand: Bool = False](output: LayoutTensor[output_type, output.layout, output.origin, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], q: LayoutTensor[dtype, q_layout, q.origin, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: k_t, v: v_t, k_rope: k_rope_t, mask_functor: mask_t, valid_length: LayoutTensor[DType.uint32, valid_length.layout, valid_length.origin, element_layout=valid_length.element_layout, layout_int_type=valid_length.layout_int_type, linear_idx_type=valid_length.linear_idx_type, masked=valid_length.masked, alignment=valid_length.alignment], max_prompt_len: Int, scale: Float32, ctx: DeviceContext, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(VariadicList(-1)), MutAnyOrigin]] = None)