[Common] Default EP comm kernels to 32 SMs; keep shuffle/preprocess on all SMs

[phu0ngng]

Description

Set a sane default SM budget for EP dispatch/combine (32 SMs) while leaving the local shuffle/preprocess kernels on all device SMs.

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

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

[phu0ngng]

/te-ci L1

[greptile-apps]

Greptile Summary

The PR changes the default NCCL EP communication-kernel budget to 32 SMs, clamped to the device SM count.

• Leaves shuffle and preprocessing kernels on all available SMs unless their environment variables are already configured.
• Updates the public C and JAX API documentation for the new communication-kernel default.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/ep/ep_backend.cpp	Applies the clamped 32-SM communication default and initializes shuffle and preprocessing SM environment defaults.
transformer_engine/common/include/transformer_engine/ep.h	Documents the revised semantics of a zero communication-SM budget in the public C API.
transformer_engine/jax/ep.py	Updates the JAX bootstrap documentation to describe the 32-SM communication default.

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into ep-sm-defaults" | Re-trigger Greptile}

[greptile-apps]

Stale zero-budget diagnostic

The implementation now maps num_comm_sms == 0 to min(32, device_sms), but the native validation message still describes zero as automatic sizing, misleading users who are diagnosing or selecting the communication-kernel SM budget.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[jberchtold-nvidia]

LGTM, thanks!