Mojo function
mla_prefill_single_batch_amd
mla_prefill_single_batch_amd[output_type: DType, q_type: DType, k_t: MHAOperand, v_t: MHAOperand, k_rope_t: MHAOperand, mask_t: MHAMask, *, config: MHAConfig, group: Int, q_depth: Int, cache_depth: Int](output: UnsafePointer[Scalar[output_type]], q: UnsafePointer[Scalar[q_type]], k: k_t, v: v_t, k_rope: k_rope_t, seq_len: Int, num_keys: Int, scale: Float32, batch_idx: Int, start_pos: UInt32, cache_start_pos: UInt32, mask: mask_t)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!