flashinfer.fused_moe.bgmv_moe_gemm2_lora_delta

flashinfer.fused_moe.bgmv_moe_gemm2_lora_delta(gemm1_activation_output: Tensor, expanded_idx_to_permuted_idx: Tensor, w_ptr_a: Tensor, lora_stride_a: int, w_ptr_b: Tensor, lora_stride_b: int, topk_ids: Tensor, topk_weights: Tensor, lora_ids: Tensor, rank: int, hidden_size: int, *, lora_dtype: dtype = torch.bfloat16, scale: float = 1.0, out_dtype: dtype = torch.bfloat16) → Tensor

FC2 (down_proj) LoRA delta for a routed MoE, to be ADDED to the MoE output.

Consumes the post-SwiGLU activation returned by trtllm_*_moe (called with gemm1_lora_delta set and do_finalize=True). For each routed pair (token t, slot j) with expert e and adapter l:

delta[t] = scale * Σ_j  w[t, j] * ( B_down[l,e] @ (A_down[l,e] @ a[t, j]) )

Weighted and combined over experts (added after FC2, post-combine).

Parameters:

  • gemm1_activation_output (torch.Tensor) – [padded_rows, I] PERMUTED, post-SwiGLU activation (trtllm return). I (the down-projection input width) is read from shape[1].

  • expanded_idx_to_permuted_idx (torch.Tensor) – [T*top_k] int — maps expanded index token*top_k+slot to the permuted row (trtllm return); < 0 marks an inactive slot.

  • w_ptr_a (torch.Tensor) – [1, num_experts] int64 base-pointer table for [A_down], from fill_w_ptr().

  • lora_stride_a (int) – Element stride between adapters in the A_down bank.

  • w_ptr_b (torch.Tensor) – [1, num_experts] int64 base-pointer table for [B_down].

  • lora_stride_b (int) – Element stride between adapters in the B_down bank.

  • topk_ids (torch.Tensor) – [T, top_k] routed expert ids (int) and per-expert combine weights (f32).

  • topk_weights (torch.Tensor) – [T, top_k] routed expert ids (int) and per-expert combine weights (f32).

  • lora_ids (torch.Tensor) – [T] int adapter id per token (-1 = none).

  • rank (int) – LoRA rank.

  • hidden_size (int) – H — the down-projection output width.

  • lora_dtype (torch.dtype) – Dtype of the LoRA weights / activation gather buffer (bf16/fp16).

  • scale (float) – LoRA alpha / rank scaling.

  • out_dtype (torch.dtype) – Output dtype (match the MoE output, e.g. bf16).

Returns:

[T, H] in out_dtype. Add to the MoE output.

Return type:

torch.Tensor