[PyTorch] Fix mutable QB bounds in CUDA graphs

[harryzhou2000]

Description

Quantile Balancing keeps qb_bin_bounds as a persistent CUDA tensor. In MCore's full-iteration CUDA graph, the router reads the current bounds near the start of an iteration and QB finalization updates the same storage near the end of that iteration. The graph executable is stable, but the state behind its captured device pointer is intentionally mutated once per replayed global batch.

The host validation cache introduced in #3395 keys eager validation to the tensor's PyTorch version. Each eager bounds.copy_() advances that version while preserving the pointer, so the last warmup leaves a valid but stale cached version and capture attempts a forbidden device-to-host validation.

Graph replay has a second important property:

• dispatching the in-place update while capturing advances _version once;
• each later cudaGraphLaunch executes the captured update and changes device contents;
• replay does not re-enter Python, so _version remains fixed.

The fix therefore uses an explicit trusted-producer handshake rather than pretending host version bookkeeping can inspect replayed data:

• eager router calls retain recoverable host validation and exact-version caching;
• mark_qb_bin_bounds_validated() records the current version without synchronizing or inspecting values;
• the marker is allowed immediately after a trusted in-place update is dispatched during capture on the same stream;
• capture still rejects a never-validated tensor or any stale version that was not explicitly produced by the trusted updater;
• replay reads later values through the stable pointer under the updater's validity contract.

MCore's QB updater is such a trusted producer: starting from finite ordered bounds, it computes finite expert biases and constructs the next range as [bias_min - 1, bias_max + 1]. The companion feature-detected MCore call is harryzhou2000/Megatron-LM@78e971efa.

This PR intentionally does not use a device assertion. Invalid replay-time data cannot become a recoverable Python exception without a caller-owned device status tensor and a later host observation point; an assertion or trap would poison the CUDA context. If an untrusted graph-time producer is added later, it should use that explicit status-buffer design. The current trusted updater needs no extra kernel, global write, or replay-time host synchronization.

The final Transformer Engine diff is Python and tests only; it does not change a CUDA kernel, C++ binding, public C ABI, or fused-router launch sequence. Kimi K3's Quantile Balancing algorithm is described in https://github.com/MoonshotAI/Kimi-K3/blob/main/k3_tech_report.pdf.

Tests

Validated on one NVIDIA B200 using the dev_2604 container:

• Focused QB bounds/CUDA-graph selection: 22 passed.
• Complete fused-router QB selection: 47 passed, 2 skipped; the skips require multiple GPUs.
• Both two_kernel and fused_atomic histogram modes.
• Sparse routing-map and caller-provided dense Top-k output paths.
• Exact-version rejection for never-validated and stale pre-capture bounds.
• Capture with a trusted bounds update after router execution, matching full-iteration ordering.
• Production-shaped MCore integration captured routing and QB finalization in one graph, then replayed four times. The bounds pointer stayed stable; bounds changed on every replay and remained finite/ordered; _version advanced once during capture and remained fixed across all replays.
• Targeted pre-commit hooks pass for all changed files.

[greptile-apps]

Greptile Summary

The PR adds an explicit validation-version handshake for mutable Quantile Balancing bounds used during CUDA graph capture.

• Adds mark_qb_bin_bounds_validated for trusted in-place device updates.
• Requires the current tensor version to be validated before graph capture selects the unchecked native path.
• Expands tests across both histogram modes, dense and sparse routing outputs, mutable graph inputs, captured updates, and stale or absent validation.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/router.py	Adds the trusted bounds-version marker, enforces exact-version validation at capture, and documents replay-time caller responsibilities.
tests/pytorch/test_fused_router.py	Extends eager and CUDA-graph QB coverage for mutable bounds, captured updates, stale validation, and both histogram implementations.

_{Reviews (5): Last reviewed commit: "[PyTorch] Allow trusted QB updates durin..." | Re-trigger Greptile}