Skip to main content

Mojo function

mla_decoding

mla_decoding[q_type: DType, k_t: MHAOperand, output_type: DType, mask_t: MHAMask, score_mod_t: ScoreModTrait, valid_layout: Layout, BM: Scalar[DType.uint], BN: Scalar[DType.uint], BK: Scalar[DType.uint], WM: Scalar[DType.uint], WN: Scalar[DType.uint], depth: Scalar[DType.uint], num_heads: Scalar[DType.uint], num_threads: Scalar[DType.uint], num_pipeline_stages: Scalar[DType.uint], group: Scalar[DType.uint] = 1, use_score_mod: Bool = False, ragged: Bool = False, _use_valid_length: Bool = False, _is_cache_length_accurate: Bool = False, decoding_warp_split_k: Bool = False](q_ptr: UnsafePointer[Scalar[q_type], MutAnyOrigin], k: k_t, output_ptr: UnsafePointer[Scalar[output_type], MutAnyOrigin], exp_sum_ptr: UnsafePointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin], qk_max_ptr: UnsafePointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin], scale: Float32, batch_size: Int, num_partitions: Int, max_cache_valid_length: Int, valid_length: LayoutTensor[DType.uint32, valid_layout, MutAnyOrigin], mask: mask_t, score_mod: score_mod_t)

Was this page helpful?