[andrewliu2001]

Description

The non-eager EP path crashes on BF16 with IndexError: tuple index out of range from
torch.ops.transformer_engine_ep.prepare_and_dispatch (#3814). dispatch and
combine_bwd have the same problem.

These ops list their optional MXFP8 scale buffers in mutates_args and give them a
None default. When the dispatcher calls a Python kernel it drops trailing arguments
that equal their default, so on the BF16 path, where the scales are None, they are
missing from the boxed args. torch's ADInplaceOrView kernel for custom ops then
indexes args with the schema position of each mutated argument and runs past the
end. PyTorch main now guards that index, but no release through 2.13 does.

This removes the defaults from those arguments so they are always present, and passes
None explicitly in the one call that omitted them (BF16 combine_bwd). The scale
buffers stay in mutates_args, so their version counters are still bumped and
torch.compile still sees the mutation when they are present.

Fixes #3814

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

• Drop the = None defaults from the optional scale arguments of the dispatch,
  prepare_and_dispatch, and combine_bwd custom ops, with a comment explaining why.
• Pass None, None explicitly in the BF16 combine_bwd call.
• Add test_ep_custom_ops_optional_mutated_args to test_permutation.py. It calls the
  three ops with the backend stubbed, so it runs on one GPU without NCCL EP: with None
  scales (the BF16 path) and with scale tensors, where it also checks that the mutated
  scale buffer's version counter is bumped once per op.

Testing

Single GPU, 1x A100 80GB PCIe, driver 570.195.03, CUDA 12.8, PyTorch 2.11.0+cu128,
python -m pytest -v tests/pytorch/test_permutation.py -k ep_custom_ops: both cases pass.
With only ep.py reset to 720b252 (current main), the bf16 case fails with
IndexError: tuple index out of range and the mxfp8_scales case passes, which is
the reported behavior.

EP suite, 8x H100 80GB HBM3 (SXM5, NVLink), same driver and PyTorch, NCCL 2.30.4,
built with NVTE_CUDA_ARCHS=90:

• With ep.py from main, run_ep.py non-eager BF16 (NVTE_EP_EAGER=0): 37 tests,
  14 errors, 20 skipped, all IndexError from prepare_and_dispatch, matching the
  issue. The fused pass (NVTE_EP_FUSED_PREPARE_DISPATCH=1) fails the same way.
• With this change, tests/pytorch/distributed/run_test_ep.sh: no IndexError in any
  pass, and every TestEP test passes or skips. Two things still fail in this
  environment and look unrelated:
  • 5 TestMoeEpSequential.test_megamoe_* tests fail in every pass, eager included,
    with "Single grouped parameters require the native grouped-tensor GroupedLinear
    path, which is unavailable for the current device, dtype, or quantization recipe."
    My build used CUDA 12.8; the issue reports these passing on CUDA 13.4 with the
    workaround applied.
  • The zero-copy passes need NCCL symm-mem support in the torch build, which the
    cu128 wheel I used does not have (same as the side issue noted in [Bug][PyTorch] EP non-eager dispatch crashes: IndexError in torch.ops.transformer_engine_ep.prepare_and_dispatch (mutates_args includes optional trailing None arg) #3814).

I also checked the change on CPU with torch 2.8.0 and 2.11.0: the old schemas raise
the IndexError for None scales and the new ones do not, the mutated scale buffer's
version counter is bumped when present, and the ops trace under
torch.compile(fullgraph=True).

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 (all EP dispatch/combine tests pass; see Testing for the unrelated MoE and zero-copy failures in my environment)

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge, with non-blocking feedback to restore the focused regression test.

Summary

Requires explicit scale arguments for the three affected EP custom ops, preventing trailing None values from being dropped.

• Non-eager EP calls keep optional scale arguments in place.

_{Reviews (2) · Last reviewed commit: "[PyTorch] Remove EP regression test and ..." · Reviewed by Greptile}

[ptrendx]

@pggPL Please take a look

[pggPL]

LGTM

[pggPL]

@andrewliu2001 thanks for the contribution. Note that for PyTorch 2.13 this is no more bug. Since we support torch >= 2.8 I'lll merge this.