Skip to content

GitLab

Commit a02a5490 authored by silencealiang's avatar silencealiang
Browse files

减少cpu侧开销

parent 1310cbf8
Pipeline #2463 passed with stage
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
megatron/core/transformer/moe/moe_utils.py
View file @ a02a5490
...@@ -312,7 +312,7 @@ def device_limited_topk( ...@@ -312,7 +312,7 @@ def device_limited_topk(
return probs, top_indices return probs, top_indices
@torch.compile(options={"triton.cudagraphs": True, "triton.cudagraph_trees": False})
def topk_softmax_with_capacity( def topk_softmax_with_capacity(
logits: torch.Tensor, logits: torch.Tensor,
topk: int, topk: int,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment