flashinfer.fused_moe.bgmv_moe_gemm1_lora_delta

flashinfer.fused_moe.bgmv_moe_gemm1_lora_delta(hidden_states: Tensor, w_ptr_a: Tensor, lora_stride_a: int, w_ptr_b: Tensor, lora_stride_b: int, topk_ids: Tensor, lora_ids: Tensor, rank: int, intermediate_size: int, *, lora_dtype: dtype = torch.bfloat16, scale: float = 1.0, out_dtype: dtype = torch.bfloat16) → Tensor

FC1 (gate_up_proj) LoRA delta for a routed MoE, in the layout consumed by trtllm_*_moe’s gemm1_lora_delta.

For each routed pair (token t, slot j) with expert e = topk_ids[t, j] and adapter l = lora_ids[t] (skipped when l < 0):

delta[t, j] = scale * concat( B_gate[l,e] @ (A_gate[l,e] @ x[t]),
                              B_up[l,e]   @ (A_up[l,e]   @ x[t]) )

Unweighted and kept per-(token, slot) (it is added before the nonlinear SwiGLU, so it must not be summed over experts nor scaled by routing weights).

Performant (shared-A) LoRA is selected by passing 1D [num_experts] pointer tables (no slice dim): a single shared A feeds both slices and B is fused horizontally into one [2I, rank] matrix, so the delta is scale * B_fused @ (A_shared @ x[t]) — one shrink + one fused expand instead of two of each.

Parameters:

  • hidden_states (torch.Tensor) – [T, H] FFN input. Cast to lora_dtype for the LoRA path, independent of the FP8/MXFP8 base weights.

  • w_ptr_a (torch.Tensor) – LoRA-A int64 base-pointer table, from fill_w_ptr(). [2, num_experts] (regular) points to [A_gate, A_up]; [num_experts] 1D (no slice dim) selects performant LoRA (one shared A).

  • lora_stride_a (int) – Element stride between adapters in the A bank(s) (the fill_w_ptr return value).

  • w_ptr_b (torch.Tensor) – LoRA-B int64 base-pointer table. [2, num_experts] (regular) points to [B_gate, B_up]; [num_experts] 1D selects performant LoRA (B fused as [2I, rank]: gate rows 0:I, up rows I:2I).

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

  • topk_ids (torch.Tensor) – [T, top_k] int — UNPACKED routed expert id per (token, slot).

  • lora_ids (torch.Tensor) – [T] int — adapter id per token, -1 for no adapter.

  • rank (int) – LoRA rank (the A/B contraction dim).

  • intermediate_size (int) – I — the per-slice (gate / up) output width; total FC1 width is 2*I.

  • lora_dtype (torch.dtype) – Dtype of the LoRA weights the w_ptr tables point to (bf16/fp16).

  • scale (float) – LoRA alpha / rank scaling applied to the delta.

  • out_dtype (torch.dtype) – Output dtype; bf16 to match gemm1_lora_delta’s required dtype.

Returns:

[T, top_k, 2*I] in out_dtype. Pass as gemm1_lora_delta.

Return type:

torch.Tensor