减少cpu侧开销

a02a5490 · silencealiang · 1310cbf8 · a02a5490
Commit a02a5490 authored Mar 10, 2025 by silencealiang
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

megatron/core/transformer/moe/moe_utils.py megatron/core/transformer/moe/moe_utils.py +1 -1

No files found.
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -312,7 +312,7 @@ def device_limited_topk(

    return probs, top_indices

-
+@torch.compile(options={"triton.cudagraphs": True, "triton.cudagraph_trees": False})
 def topk_softmax_with_capacity(
    logits: torch.Tensor,
    topk: int,