feat(MoE): Refactor cuda_graph_scope

[buptzyb]

dev branch PR #1917 & #2353 & #2694 .

With this PR, --cuda-graph-scope in --cuda-graph-impl=transformer_engine mode now supports combinations of the six values:

1. attn: captures operations in TransformerLayer._forward_attention().
2. mlp: captures operations in TransformerLayer._forward_mlp() for a dense layer.
3. moe: captures operations in TransformerLayer._forward_mlp() for a MoE layer.
4. moe_router: captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(). Note that if shared experts overlap is not used, it also captures the shared experts.
5. moe_preprocess: captures operations in MoELayer.preprocess(). Must be used together with moe_router.
6. mamba: captures the mamba layer.

• Example 1:

For a dense model, set --cuda-graph-scope attn mlp to capture the whole Transformer layer, or set --cuda-graph-scope attn to capture the attention part, or set --cuda-graph-scope mlp to capture the mlp part. The non-graphed part will go to the normal pass.

• Example 2:

For a moe model, set --cuda-graph-scope attn moe_router moe_preprocess to capture operations from the beginning of the Transformer layer to the preprocess method in the moe token dispatcher. However, if you are using alltoall dispatcher with drop-no-padding or router-padding-for-fp8 options, you can only set --cuda-graph-scope attn moe_router to capture up to the moe router because there will be cuda synchronization in the preprocess method. Finally, if you are using alltoall dispatcher with drop-padding, then you can directly set --cuda-graph-scope attn moe to capture the whole layer as its sync-free.

• Example 3:

For a model that contains both dense and moe layers like DeepSeek, the dense layer will check for "mlp" and the moe layer will check for "moe*". For example, setting --cuda-graph-scope attn mlp moe_router moe_preprocess captures the whole (attn+mlp) dense layer, and the attn+router+preprocess part of the moe layer.

• Example 4:

Mamba model is different from a traditional TransformerLayer-based model. The mamba, attention, mlp, moe are all in different layer objects. For a mamba+moe model, you can set --cuda-graph-scope mamba attn moe_router to capture the corresponding layers. Or you can also set --cuda-graph-scope attn moe_router if you don't want the mamba layers to be graphed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Skylion007]

Suggested change

	hier_attr_name = attr_name.split('.')
	hier_attr_name = attr_name.rsplit('.', maxsplit=1)

[buptzyb]

I want it to be compatible with multiple levels of hierarchy. Your suggestion works for now because we currently only have up to two levels like the _comm_manager.token_probs. I assume in the future we may need to capture attributes at a deeper level when the code becomes more and more complex.

[buptzyb]

/ok to test 3b85592

[buptzyb]

Hi @rogerwaleffe @duncanriach @JRD971000 could you help review on behalf of NVIDIA/hybrid-mamba? @kvareddy @santhnm2 could you help review? Thanks!

[kvareddy]

@lmcafee-nvidia @mathemakitten can you please sign off on this MR?

[buptzyb]

Hi @rogerwaleffe @duncanriach @JRD971000 could you help review on behalf of hybrid-mamba? @jaredcasper @deepakn94 @santhnm2 @ericharper could you help review on behalf of gpt? Thanks!

[buptzyb]

Hi @rogerwaleffe @duncanriach @JRD971000 could you help review on behalf of hybrid-mamba? @jaredcasper @deepakn94 @santhnm2 @ericharper could you help review on behalf of gpt? @yanring could you help review on behalf of mixture-of-experts-devtech? Thanks!

[buptzyb]

/ok to test 0698a59

[buptzyb]

Hi @rogerwaleffe @duncanriach @JRD971000 could you help review on behalf of hybrid-mamba? @jaredcasper @deepakn94 @santhnm2 @ericharper could you help review on behalf of gpt? @yanring could you help review on behalf of mixture-of-experts-devtech? Thanks!

[Phlip79]

/ok to test 0698a59