[PyTorch] Retain NVFP4 RNG tensors through quantization

[Connor-XY]

NVFP4 stochastic quantization stores non-owning TE wrappers for temporary rowwise and columnwise RNG tensors. Their ATen owners went out of scope before quantization dispatch, allowing allocator reuse to change seeded output; disabling caching exposed invalid global-memory reads.

Keep both owners alive through dispatch. The regression restores the real CUDA RNG, compares logical packed outputs/scales across ordinary and fragmented allocation patterns, and verifies seed sensitivity. Comments now explain the load-bearing deterministic fill and alternating-free anchors. The 128-column fused shape is parity coverage on SM100/SM110; allocator-based detection is heuristic. The independent no-caching Compute Sanitizer check remains essential evidence.

Validation and scope

Current head: e1cc25658e62dd392c5e059cdb4a3284fc133376. This follow-up changes comments/docstrings only; executable Python AST after removing docstrings is identical to the previously validated test, and the C++ ownership fix is unchanged. No new TE wheel/GPU acceptance run is claimed for the comment-only commit.

At previous head 17c95e206159ec476d88241343a6fa29a21e08d4, exact-source GB200 acceptance built the full PyTorch wheel and CUDA core. Independent verification checked all 10,794 tracked source files, installed binary identities and per-rank artifacts:

• All four regression cases passed on each of four GPUs.
• Stochastic-rounding references passed 28 cases per GPU; four unsupported RHT/dtype combinations skipped.
• Both no-caching Compute Sanitizer controls completed with zero errors.
• MCore eba245d420ec822117445f253d4cae6ab7380c65 passed 31 model and 83 kernel pytest cases per rank. Selected reports were D31/N0/U0 and D49/N0/U0; all six strict author reference/sensitivity requirements passed with actual dependency identities.

Separate pinned-release controls at TE 27486e03cfc1fa41f6932dcecdc47c71c47eac3e established causality: both installed and rebuilt-unmodified controls failed three of four new cases while the ownership patch passed all four. No-caching sanitizer controls reported 136/69 errors for unpatched variants and zero with the patch. The 128-column parity case is not claimed as an old-binary regression failure.

Raw validation data remains on the GPU hosts. Full TE required GPU CI and maintainer review remain pending. A shared RNG owner/view abstraction across other call sites is a separate follow-up; this PR stays focused on ownership and its regression.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, security, or repository-rule issue remains.

Summary

Retains the owning ATen tensors behind NVFP4 rowwise and columnwise RNG wrappers until quantization has been dispatched, preventing temporary RNG storage from being reclaimed prematurely. Adds a GPU regression that restores CUDA RNG state and compares packed outputs and scales across ordinary and fragmented allocation histories.

• Moves both RNG tensor owners to quantize_impl scope while preserving separate rowwise and columnwise RNG states.
• Tests deterministic replay, allocation-history independence, and sensitivity to different seeds across four NVFP4 shapes.
• Documents the allocator-dependent test mechanism and the fused-path parity case.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Initialize Philox states] --> B[Owning ATen RNG tensors]
  B --> C[Non-owning TensorWrapper views]
  C --> D[NVFP4 quantization configuration]
  D --> E[Launch rowwise and columnwise kernels]
  E --> F[Return after dispatch]
  B -->|Owners retained through launch| F

Loading

_{Reviews (4) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..."}