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 function

shrink_qkv_permute_3mn_sm100

def shrink_qkv_permute_3mn_sm100(c_lora: TileTensor[Storage=c_lora.Storage, linear_idx_type=c_lora.linear_idx_type, element_size=c_lora.element_size], a: TileTensor[Storage=a.Storage, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[Storage=b.Storage, linear_idx_type=b.linear_idx_type, element_size=b.element_size], a_offsets: TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type, element_size=a_offsets.element_size], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type, element_size=expert_ids.element_size], max_num_tokens_per_expert: Int, num_active_experts: Int, ctx: DeviceContext)

TileTensor primary implementation of shrink_qkv_permute_3mn_sm100.

LoRA shrink GMM with planar Q/K/V output on SM100.

Performs the LoRA 'shrink' grouped matmul for routed tokens: computes [M, K] @ [G, 3N, K]^T per active expert, then permutes the flat [M, 3N] result into a planar layout [3, M, N] (Q, K, V) using an elementwise epilogue, while reusing the same storage.

Constraints:

c_lora must be rank 3 with static first dimension B == 3.
a must be rank 2 with trailing dimension K that matches b[..., K].
b must be rank 3 with shape (G, 3N, K).
The temporary 2D view of c_lora is (M, 3N) in row-major order and aliases the same storage as c_lora.
a_offsets is non-decreasing with a_offsets[0] == 0 and a_offsets[num_active_experts] == M.
expert_ids[i] in [0, G) for valid experts; kernel may treat -1 as inactive.
The epilogue assumes N % vector_width == 0 for aligned vector stores.

Args:

c_lora (TileTensor[Storage=c_lora.Storage, linear_idx_type=c_lora.linear_idx_type, element_size=c_lora.element_size]): Output tensor with planar Q/K/V layout, shape (3, M, N). Backed by row-major storage, used both as a 3D view and as a temporary 2D view (M, 3N) during compute.
a (TileTensor[Storage=a.Storage, linear_idx_type=a.linear_idx_type, element_size=a.element_size]): Routed activation matrix, shape (M, K).
b (TileTensor[Storage=b.Storage, linear_idx_type=b.linear_idx_type, element_size=b.element_size]): Shrink weights per expert, shape (G, 3N, K).
a_offsets (TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type, element_size=a_offsets.element_size]): Inclusive prefix sums of tokens per (active) expert, length (num_experts + 1). Defines per-expert [start, end) in A/C.
expert_ids (TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type, element_size=expert_ids.element_size]): Expert indices for the active groups, length >= num_active_experts.
max_num_tokens_per_expert (Int): Upper bound on tokens for any active expert.
num_active_experts (Int): Number of experts participating in this call.
ctx (DeviceContext): DeviceContext used for enqueues and synchronization.