[Bug][AutoDeploy]: Investigate rank drift for SuperV3 + MTP + CudaGraph

System Info

Hardware: 8 x H100

We have observed rank drift for SuperV3 + MTP when torch-cudagraph is enabled, where the target model one rank will occasionally sample a different token than the others. This can lead to crashes due to different acceptance patterns and input_pos across ranks.

This issue seems to stabilize when the models have TP sharding. Since low-latency configurations may lack TP sharding, we have put a broadcast from rank 0 for the sampled logits to ensure consistency.

Who can help?

@govind-ramnarayan

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

• Remove broadcast in EagleWrapper::sample_greedy()
• Run accuracy test
  tests/integration/defs/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp
  while logging the target logits and the implied greedily-sampled tokens per rank

Observe that occasionally the ranks disagree on what the target logits are, and what token should be sampled from them.

This may also just result in a crash in tests/integration/defs/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp
when this leads to different token acceptance across ranks.

Expected behavior

The broadcast should be useless - if we remove it, we should still see the same tokens sampled on each rank.

actual behavior

The broadcast is needed - tokens occasionally differ across ranks during this accuracy test which occasionally leads to a crash.

additional notes

May be related to #13133 - the broadcast, when needed, ends up changing the token a rank samples, which affects the future runs of the model. In particular, the token that we end up forcing upon a rank due to broadcast ends up disagreeing with the data in its KV cache, which leads to a different sampled token and creates an internal inconsistency.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.