Repository: vllm-project/vllm
File: vllm/lora/ops/triton_ops/fused_moe_lora_op.py (Python)

Code before:
            )
            a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak

        accumulator += tl.dot(a, b)

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
        accumulator = accumulator * moe_weight[:, None]
    accumulator = accumulator.to(c_ptr.dtype.element_ty)
    # Write back the block of the output
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    # When sort_c is true, store the output in c_ptr using token order defined
    # in sorted_token_ids_ptr; otherwise, use the original token order from the prompt
    if sort_c:
        offs_token_id = pid_m * BLOCK_SIZE_M + offs
        c_ptrs = (
            cur_c_ptr
            + lora_id * tl.cdiv(EM, top_k_num) * top_k_num * stride_cm
            + stride_cm * offs_token_id[:, None]
            + stride_cn * offs_cn[None, :]
        )
    else:
        c_ptrs = (
            cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
        )
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)

    if SPLIT_K == 1:
        if ADD_INPUTS:
            prev = tl.load(c_ptrs, mask=c_mask, other=0.0)
            tl.store(c_ptrs, prev + accumulator, mask=c_mask)
        else:
            tl.store(c_ptrs, accumulator, mask=c_mask)
    else:
        tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")


@torch.inference_mode()
def _fused_moe_lora_shrink(
    a_intermediate_cache1: torch.Tensor,
    # (num_slices, num_tokens, top_k_num, max_lora_rank)
    qcurr_hidden_states: torch.Tensor,  # (num_tokens, K,)
    lora_a_stacked: list[
        torch.Tensor
    ],  # [(max_loras, num_experts, max_lora_rank, K,),...]
    topk_weights: torch.Tensor,  # (num_tokens, top_k_num)
    sorted_token_ids: torch.Tensor | None,  # (max_loras, _)
    expert_ids: torch.Tensor,  # (max_loras, _ ,) or (num_tokens * top_k,)
    num_tokens_post_padded: torch.Tensor | None,  # (max_loras, )
    token_lora_mapping: torch.Tensor,

Review comment (suggestion):
would be nice to factor this out to something like, 
```
 a_ptrs, a = load_a()
 b_ptrs, b = load_b()
```

Code after:
        k_remaining = K - cur_k_offset
        # pre-fetch lora weight
        if b_desc is not None:
            b = (
                b_desc.load([lora_id, expert_id, offs_bn, offs_bk + cur_k_offset])
                .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K)
                .T
            )
        else:
            # add (offs_bn < N) mask; optional .ca for B
            b_mask = (offs_k[:, None] < k_remaining) & (offs_bn[None, :] < N)
            if USE_B_L2_CACHE:
                b = tl.load(b_ptrs, mask=b_mask, other=0.0, cache_modifier=".ca")
            else:
                b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk

        if a_desc is not None:
            a = a_desc.load([offs_am, offs_ak + cur_k_offset])
        else:
            a = tl.load(
                a_ptrs,
                mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
                other=0.0,
            )
            a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak

        accumulator += tl.dot(a, b)

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
        accumulator = accumulator * moe_weight[:, None]
    accumulator = accumulator.to(c_ptr.dtype.element_ty)
    # Write back the block of the output
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = _get_c_ptrs(
        cur_c_ptr,
        lora_id,
        pid_m,
        offs,
        offs_token,
        offs_cn,
        stride_cm,
        stride_cn,
        EM,
        BLOCK_SIZE_M,
        sort_c,
    )
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)

    if SPLIT_K == 1: