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

expand_qkv_sm100

def expand_qkv_sm100(q_out: TileTensor[Storage=q_out.Storage, linear_idx_type=q_out.linear_idx_type], kv_out: TileTensor[Storage=kv_out.Storage, linear_idx_type=kv_out.linear_idx_type], p: TileTensor[Storage=p.Storage, linear_idx_type=p.linear_idx_type], b: TileTensor[Storage=b.Storage, linear_idx_type=b.linear_idx_type], a_offsets: TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type], max_num_tokens_per_expert: Int, num_active_experts: Int, ctx: DeviceContext)

Single-launch LoRA-B expand grouped GEMM (load/store specialized).

Computes the LoRA 'expand' (up-projection by the fused LoRA-B weight) for routed tokens in ONE grouped-GEMM launch over a single fused weight, GQA-correct (q_dim != kv_dim). For token row m (in adapter group g) and output column j in [0, q_dim + 2*kv_dim):

j < q_dim -> region Q -> dot(B[g, j, :], P[0, m, :]) written to q_out[m, j].
q_dim <= j < q_dim+kv_dim -> region K -> dot(B[g, j, :], P[1, m, :]) written to kv_out[m, j-q_dim].
otherwise -> region V -> dot(B[g, j, :], P[2, m, :]) written to kv_out[M+m, j-q_dim-kv_dim].

Reuses the tuned persistent grouped matmul: the fused weight is the matmul B operand and the planar shrink output P [3, M, R] (viewed [3M, R]) is the A operand. The per-region operand switch (which plane of P) is a load-stage specialization (a_plane_splits, shifting the activation row by plane * M); the Q/K/V output routing is an output-epilogue specialization (elementwise_lambda_fn). K and V are row-stacked in kv_out (K into rows [0, M), V into rows [M, 2M)), matching the layout kv_cache_ragged_2m_iadd consumes downstream.

Constraints:

p must be rank 3 with static shape (3, M, R) (R static).
b must be rank 3 with static shape (G, q_dim + 2*kv_dim, R).
q_out and kv_out are row-major with the same dtype; q_out trailing dim is q_dim, kv_out trailing dim is kv_dim.
q_dim and kv_dim must be statically known and > 0. For the SM100 tensor-core path they must be tile-aligned (multiples of the grouped-matmul BM and the epilogue vector width) so no output tile or store vector straddles a Q/K/V region boundary.
a_offsets is non-decreasing with a_offsets[0] == 0 and a_offsets[num_active_experts] == M.

Args:

q_out (TileTensor[Storage=q_out.Storage, linear_idx_type=q_out.linear_idx_type]): Q output, shape (M, q_dim), region Q.
kv_out (TileTensor[Storage=kv_out.Storage, linear_idx_type=kv_out.linear_idx_type]): K/V output, shape (2M, kv_dim); region K into rows [0, M), region V into rows [M, 2M).
p (TileTensor[Storage=p.Storage, linear_idx_type=p.linear_idx_type]): Shrink output in planar layout, shape (3, M, R) (plane 0 = q-lora, 1 = k-lora, 2 = v-lora). This is exactly the output of shrink_qkv_permute_3mn_sm100.
b (TileTensor[Storage=b.Storage, linear_idx_type=b.linear_idx_type]): Fused LoRA-B weight per adapter, shape (G, q_dim + 2*kv_dim, R). Per adapter the q_dim + 2*kv_dim rows are the output features, partitioned along that axis into Q, then K, then V; R (the LoRA rank, the columns) is shared across all three.
a_offsets (TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type]): Inclusive prefix sums of tokens per active adapter, length (num_experts + 1). The SAME grouping used by the shrink -- no separate K/V offset array.
expert_ids (TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type]): Adapter indices for the active groups, length >= num_active_experts; -1 marks an inactive group.
max_num_tokens_per_expert (Int): Upper bound on tokens for any active adapter.
num_active_experts (Int): Number of adapter groups in this call.
ctx (DeviceContext): DeviceContext used for the enqueue.