Fix Cuda Graph based MOE Tests Hang in CI

[vthumbe1503]

Description

make_graphed_callables function returns the module replaced with the graphed module. Recent pytorch upgrade caused the graph replay to hang due to insufficient cleanup.

Irresspective, we were comparing the forward and backward results of the same graphed module being run twice, instead of actually running the non graphed and graphed module and comparing their results numerically.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[vthumbe1503]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR updates grouped linear CUDA graph tests to avoid CI hangs and unstable padded-row checks. The main changes are:

• Separate eager reference modules for grouped linear and grouped MLP comparisons.
• State copied from graphed modules into the reference modules before expected-value checks.
• Active-token-only assertions for padded grouped GEMM inputs.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
tests/pytorch/test_grouped_linear.py	Creates an independent grouped linear reference module for eager CUDA graph safety comparisons.
tests/pytorch/test_grouped_mlp.py	Uses independent reference grouped linear and MLP modules, and compares only active token rows for padded inputs.

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into grouped_linear_..." | Re-trigger Greptile}

[vthumbe1503]

Pipeline: 58029914

[vthumbe1503]

/te-ci pytorch

[timmoon10]

LGTM. Can you elaborate more on how "insufficient cleanup" was causing the hang?

Good catch that we weren't running the non-graphed implementation. It is awkward that make_graphed_callables changes the module forward so it always replays the graph, but this is also the behavior in torch.cuda.make_graphed_callables.

[vthumbe1503]

@timmoon10 I dug a bit deeper and actually it turns out the hang was due to CUDA graph replay compatibilty with CUBLAS GGEMM which exists currently. And it specifically happens on the second graph replay which was happening unintentionally in the test