Let the CUDA allocator manage cuBLAS workspaces

[Vtmpas]

Description

Concurrent GEMMs on independent CUDA streams receive the same cuBLAS workspace from the device/mode cache and can overwrite each other's partial sums. On one H200, two FP32 NT GEMMs with a reduction dimension of 16384 produced incorrect results in 35 of 48 calls; separate scratch produced zero mismatches against serial execution.

Remove the Python workspace cache and let the PyTorch caching allocator manage scratch for each invocation. Eager allocations are reused safely on their stream. Captured allocations belong to their graph's private pool, so a later capture cannot borrow a buffer from an earlier graph generation. Scratch is returned to the allocator when the invocation ends, without retaining tensors for temporary streams.

This also removes Linear's preallocation workaround for the process-global cache. The workspace signature and sizes, GEMM dtypes and algorithm selection are unchanged. No new device synchronization is added.

Type of change

• Documentation change
• Bug fix (non-breaking change which fixes an issue)
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Allocate scratch per invocation instead of caching it in Python.
• Remove the preallocation workaround for globally cached scratch.
• Test eager ownership, temporary-stream cleanup, graph-pool isolation across capture generations, concurrent FP32/BF16 graph replay, grouped GEMM scratch reuse, and compiled Linear forward/backward across graph generations.
• Cover router dgrad NN with weight[128,4096], grad_output[8192,128], output [8192,4096], and reduction 128 against exact analytic references in FP32 and BF16. Preserve NT coverage and check every eager invocation while retaining only one pair of large outputs.
• Include the workspace regression in the PyTorch L0 launcher.

Validation

• One NVIDIA H200, CUDA 13.0, PyTorch 2.13.0+cu130.
• With the exact workspace functions from main (d0b4b32) loaded into native TE 2.19.0.dev0+b5599209: the first 13 regression cases fail before the fix and pass after it. The previous per-stream-cache revision fails 11 of those 13 cases, including both concurrent graph-replay numerics cases.
• The integration patch passes all 17 workspace cases, including grouped GEMM, and is idempotent.
• Rebuilt current main for SM90 and ran NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python -m pytest -v tests/pytorch/test_gemm_workspace.py: all 22 cases pass in 15.98s, including all four new FP32/BF16 NN eager/independent-graph cases, the existing NT cases, and compiled GEMM graph generations. The compiled Linear cases are capability-skipped because this PyTorch build lacks the required opaque-object API.
• Scoped forward/backward timing used benchmarks/linear/benchmark_linear.py::benchmark_linear with BF16/FP8 block scaling on shapes (8192,4096,128) and (128,4096,16384). The repeated timings were unstable, so they do not establish a performance improvement or a reliable regression bound.
• Repository pre-commit checks, changed-file pylint and the L0 license checker pass on macOS.
• Full current-main L0 validation is still pending.

The original workspace factory is identical in main, stable, release_v2.19 and release_v2.20.

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

[Medium risk] Changes how cuBLAS workspace memory is allocated and cached.

The PR appears safe to merge; no new actionable issue or outstanding prior finding remains.

Summary

This PR replaces the Python cuBLAS workspace cache with invocation-local allocations and removes Linear’s preallocation workaround.

• Adds ownership, graph-replay, concurrent GEMM, and compiled-execution regression tests.
• Adds the workspace tests to the PyTorch L0 launcher.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  G[GEMM invocation] --> A[Allocate workspace]
  A --> E[Eager: caching allocator manages reuse]
  A --> C[Capture: graph-private pool owns allocation]
  E --> R[Invocation releases tensor]
  C --> R

Loading

_{Reviews (4) · Last reviewed commit: "Merge branch 'main' into fix/cublas-work..."}