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

Code before:
                top_k_num,
                max_lora_rank,
            )

    a_intermediate_cache1 = torch.zeros(
        intermediate_cache_shape,
        dtype=output.dtype,
        device=device,
    )

    use_gdc = supports_pdl(device) and not fully_sharded
    _fused_moe_lora_shrink(
        a_intermediate_cache1,
        qcurr_hidden_states,
        lora_a_stacked,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        token_lora_mapping,
        top_k_num,
        lora_ids,
        adapter_enabled,
        ## adding for kernel
        device,
        N,
        M,
        EM,
        K,
        num_tokens,
        num_experts,
        num_slices,
        shrink_block_size_m,
        shrink_block_size_n,
        shrink_block_size_k,
        shrink_group_size_m,
        shrink_num_warps,
        shrink_num_stages,
        shrink_split_k,
        num_active_loras,
        mul_routed_weight,
        use_gdc=use_gdc,
        use_tma=use_tma,
    )

    if fully_sharded:
        if max_lora_rank == w1_lora_b_stacked.shape[-1]:
            a_intermediate_cache1 = tensor_model_parallel_all_reduce(
                a_intermediate_cache1
            )
        else:

Review comment (performance):
Question about intermediate_cache_shape. Focusing on the `use_tma` and `sorted_tokens_ids != None` case IIUC,
this is an intermediate tensor to store the outputs of the shrink operation, without permuting it (as being done in the non-tma case). The idea here is to simply store the outputs in the matching sorted_token_ids block.

For example, 
let num_slices = 1, max_loras = 1, block_size_m = 16, 
let `sorted_token_ids[0]` be  ` 0, 0, 0, ..... | 1, 1, 1, ... | 2, 2, 2 .... | ... | 9, 9, 9 ... | 10, 10, 10 ... | 11, 11, 11 ... | ` where each `|` signifies the end of a block.

In the shrink kernel - for `pid_m`=10, 
the outputs will be stored in `a_intermediate_cache[0, 0, 160: 176, :]` and 

in the expand kernel - for the same `pid_m`=10, we read from `a_intermediate_cache[0, 0, 160:176, :]`
am I missing something ? 

And if that is the case, can we define the intermediate_cache shape as , 
  `(num_slices, sorted_token_ids.shape[0] /* num_active_loras */, EM, max_lora_rank)` 

Code after:

    # TMA is not currently compatiple with fully_sharded due to the non-determinism
    # of token id sorting across ranks.
    use_tma = supports_tma(device) and not fully_sharded

    intermediate_cache_shape = (
        num_slices,
        M,
        top_k_num,
        max_lora_rank,
    )
    if use_tma:
        if num_slices > 1:
            # if num_slices > 1, we construct TMA descriptors for LoRA
            # weights within the kernel, which requires us to first set an allocator
            set_triton_allocator(device)

        # When storing intermediate data in sorted order for TMA, we
        # need an extra 'num_active_loras' dim in the cache to avoid conflicts
        if sorted_token_ids is not None:
            intermediate_cache_shape = (
                num_slices,
                sorted_token_ids.shape[0],
                EM,
                max_lora_rank,
            )

    a_intermediate_cache1 = torch.zeros(
        intermediate_cache_shape,
        dtype=output.dtype,
        device=device,
    )

    use_gdc = supports_pdl(device) and not fully_sharded
    _fused_moe_lora_shrink(
        a_intermediate_cache1,
        qcurr_hidden_states,
        lora_a_stacked,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        token_lora_mapping,
        top_k_num,
        lora_ids,
        adapter_enabled,
        ## adding for kernel
        device,
        N,
        M,
        EM,