Fix distributed Newton-Schulz correctness

[vcherepanov-nv]

Description

Fix a number of numerical correctness issues in Newton-Schulz

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:

• Explicitly enable cuSOLVERMp input normalization and compute-type Gram reduction.
• Validate distributed dimensions and reject unsupported tall low-level inputs.
• Reduce device workspace requirements to the grid-wide maximum.
• Synchronize workspace cleanup and retain host workspace for asynchronous operations.
• Normalize the test reference and evaluate Gram and reference errors in FP32.
• Use rank-independent aligned test shapes and strict absolute tolerances.

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

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding findings or new correctness issues introduced since the previous review.

Summary

This PR fixes numerical correctness and workspace lifetime issues in distributed Newton–Schulz operations.

• Enables input normalization and compute-type Gram reduction in cuSOLVERMp.
• Validates distributed dimensions and rejects unsupported tall, column-sharded low-level inputs.
• Caches grid-wide workspace requirements and preserves asynchronous host workspace.
• Expands tall TP coverage and evaluates numerical results in FP32 with strict absolute tolerances.
• Changes since the previous review are formatting-only and introduce no new behavioral issues.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Newton–Schulz input] --> B{TP mode}
  B -->|duplicated| C[Gather full tensor]
  B -->|distributed| D{Partitioned along larger dimension?}
  D -->|No| E[Reject unsupported distribution]
  D -->|Yes| F[Map to wide column-sharded operation]
  C --> F
  F --> G[Normalize input]
  G --> H[Query or reuse cached workspace]
  H --> I[Grid-wide maximum device workspace]
  I --> J[cuSOLVERMp Newton–Schulz]
  J --> K[Synchronize output with caller stream]

Loading

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

[cyanguwa]

It generally looks good to me. Did this PR fix the two bugs we have on L40s then?

Could you address the first P2 comment from Greptile? I feel the second P2 can be explained by the new n <= m requirement. Does the requirement come from cusolvermp? Should we note it somewhere in the docstring?

Also, do we need to apply these same fixes on the Jax side? (Not for this PR, but as a followup if necessary.)

Please run all the necessary CI as well (especially L1). Thanks!