mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-04-11 05:07:02 +08:00
[Doc] Fixed shape description for fused_batched_moe.py (#25668)
Signed-off-by: Egor <e.a.krivov@gmail.com> Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
parent
2bcc745042
commit
564233d550
@ -355,7 +355,7 @@ def batched_triton_kernel(
|
||||
|
||||
def invoke_moe_batched_triton_kernel(
|
||||
A: torch.Tensor, # [E, max_tokens, K]
|
||||
B: torch.Tensor, # [E, K, N]
|
||||
B: torch.Tensor, # [E, N, K]
|
||||
C: torch.Tensor, # [E, max_tokens, N]
|
||||
expert_num_tokens: torch.Tensor, # [E]
|
||||
compute_type: tl.dtype,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user