Let the CUDA allocator manage cuBLAS workspaces

qa/L0_pytorch_unittest/test.sh

@@ -34,6 +34,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_optional_flash_attn_import.xml $TE_PATH/tests/pytorch/test_optional_flash_attn_import.py || test_fail "test_optional_flash_attn_import.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gemm_workspace.xml $TE_PATH/tests/pytorch/test_gemm_workspace.py || test_fail "test_gemm_workspace.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py"

tests/pytorch/test_torch_compile.py

@@ -1933,6 +1933,52 @@ def test_to_tensor_spec_quantized(factory, shape):
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+def test_te_linear_cublas_workspace_graph_generations(monkeypatch):
+    """Compiled forward/backward scratch follows graph generations, not a Python cache."""
+    from transformer_engine.pytorch.cpp_extensions import gemm
+
+    original = gemm.tex.generic_gemm
+    captured_workspace_requests = []
+
+    def observe_native_gemm(*args, **kwargs):
+        if torch.cuda.is_current_stream_capturing():
+            # Do not retain tensor arguments or graph-pool owners.
+            captured_workspace_requests.append((args[12] is None, args[13]))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(gemm.tex, "generic_gemm", observe_native_gemm)
+    model = te.Linear(4096, 128, bias=False, params_dtype=torch.bfloat16, device="cuda")
+    with torch.no_grad():
+        model.weight.fill_(1)
+    torch._dynamo.reset()
+    compiled = torch.compile(model, fullgraph=True, mode="reduce-overhead", dynamic=False)
+    with _assert_no_cudagraph_skips(True):
+        for batch, value in [(128, 1), (256, 2), (128, 3)]:
+            for _ in range(3):
+                torch.compiler.cudagraph_mark_step_begin()
+                model.zero_grad(set_to_none=True)
+                inp = torch.full(
+                    (batch, 4096), value, dtype=torch.bfloat16, device="cuda", requires_grad=True
+                )
+                out = compiled(inp)
+                out.sum().backward()
+                torch.testing.assert_close(out, torch.full_like(out, 4096 * value), rtol=0, atol=0)
+                torch.testing.assert_close(inp.grad, torch.full_like(inp, 128), rtol=0, atol=0)
+                torch.testing.assert_close(
+                    model.weight.grad, torch.full_like(model.weight, batch * value), rtol=0, atol=0
+                )
+                del out, inp
+    assert (
+        len(captured_workspace_requests) >= 2
+    ), "Expected CUDA graph captures for different shapes"
+    assert all(
+        native and size == 0 for native, size in captured_workspace_requests
+    ), "Captured GEMMs must request invocation-local native scratch with workspace=None, size=0"
+    del compiled
+    torch._dynamo.reset()
+
+
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.parametrize("compile_mode", _compile_modes)
 @pytest.mark.parametrize(
@@ -3046,16 +3092,17 @@ def run(module, inp):
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_te_ops_linear_workspace_fake_mode(monkeypatch):
     def unexpected_allocation(*args, **kwargs):
-        pytest.fail("Fake initialization must not populate the real workspace cache")
+        pytest.fail("Fake initialization must not allocate a real cuBLAS workspace")
 
     monkeypatch.setattr(
-        "transformer_engine.pytorch.ops.basic.basic_linear.get_cublas_workspace",
+        "transformer_engine.pytorch.cpp_extensions.gemm.get_cublas_workspace",
         unexpected_allocation,
     )
     with FakeTensorMode():
         te.ops.BasicLinear(32, 64, device="cuda", dtype=torch.bfloat16)
 
 
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("training", [False, True])
 @pytest.mark.parametrize(
@@ -3074,11 +3121,27 @@ def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device, quantizati
         import torch
         import transformer_engine.pytorch as te
         from transformer_engine.common import recipe
-        from transformer_engine.pytorch.cpp_extensions.gemm import get_cublas_workspace
+        from transformer_engine.pytorch.cpp_extensions import gemm
         from torch._dynamo.utils import counters
 
         training, initial_device, quantization = sys.argv[1:]
         training = training == "True"
+        workspace_requests, captured_workspace_requests = [], []
+        original = gemm.tex.generic_gemm
+
+        def observe_native_gemm(*args, **kwargs):
+            # Retain only scalar metadata, never tensors or graph-pool owners.
+            request = (args[12] is None, args[13])
+            workspace_requests.append(request)
+            if torch.cuda.is_current_stream_capturing():
+                captured_workspace_requests.append(request)
+            return original(*args, **kwargs)
+
+        def unexpected_python_workspace(*args, **kwargs):
+            raise AssertionError("Native GEMMs must not allocate Python workspace tensors")
+
+        gemm.tex.generic_gemm = observe_native_gemm
+        gemm.get_cublas_workspace = unexpected_python_workspace
         is_nvfp4 = quantization != "None"
         quant_recipe = (
             recipe.NVFP4BlockScaling(
@@ -3091,18 +3154,15 @@ def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device, quantizati
         in_features, out_features = (128, 128) if is_nvfp4 else (32, 64)
         batch_size = 32
         eager_tols = {"rtol": 0, "atol": 0} if is_nvfp4 else {}
-        assert get_cublas_workspace.cache_info().currsize == 0
         linear = te.ops.BasicLinear(
             in_features, out_features, device=initial_device, dtype=torch.bfloat16
         )
         if initial_device == "cpu":
-            assert get_cublas_workspace.cache_info().currsize == 0
             linear.cuda()
         elif initial_device == "meta":
-            assert get_cublas_workspace.cache_info().currsize == 0
             linear.to_empty(device="cuda")
             linear.reset_parameters()
-        assert get_cublas_workspace.cache_info().currsize > 0
+        assert not workspace_requests, "Module initialization must not allocate GEMM scratch"
         model = te.ops.Sequential(linear, te.ops.Bias(out_features, dtype=torch.bfloat16))
         reference_model = copy.deepcopy(model)
         x = torch.randn(
@@ -3135,6 +3195,10 @@ def forward(x, module):
                     del grads, expected_grads
                 del actual, expected
         torch.cuda.synchronize()
+        assert captured_workspace_requests, "Expected native workspace requests during graph capture"
+        assert all(native and size == 0 for native, size in workspace_requests), (
+            "GEMMs must request invocation-local native scratch with workspace=None, size=0"
+        )
         assert not counters["inductor"]["cudagraph_skips"], counters["inductor"]
         assert counters["inductor"]["cudagraph_recorded_non_static_inputs"] > 0, dict(
             counters["inductor"]

transformer_engine/common/gemm/cublaslt_gemm.cu

@@ -14,6 +14,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <exception>
 #include <mutex>
 #include <vector>
 
@@ -1117,31 +1118,38 @@ void multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETens
         cudaStreamWaitEvent(detail::get_compute_stream(s), detail::get_compute_stream_event(0)));
   }
 
-  for (int i = 0; i < num_gemms; i++) {
-    // Check whether GELU or dGELU epilogue is requested
-    Tensor *pre_gelu_tensor = convertNVTETensor(pre_gelu_out[i]);
-    bool with_gelu_dgelu_epilogue =
-        (pre_gelu_tensor != nullptr && pre_gelu_tensor->data.dptr != nullptr);
-
-    // Construct config
-    MatmulConfig config;
-    if (grad) {
-      config.dbias_tensor = bias[i];
-      config.with_dgelu_epilogue = with_gelu_dgelu_epilogue;
-    } else {
-      config.bias_tensor = bias[i];
-      config.with_gelu_epilogue = with_gelu_dgelu_epilogue;
+  std::exception_ptr launch_error;
+  try {
+    for (int i = 0; i < num_gemms; i++) {
+      // Check whether GELU or dGELU epilogue is requested
+      Tensor *pre_gelu_tensor = convertNVTETensor(pre_gelu_out[i]);
+      bool with_gelu_dgelu_epilogue =
+          (pre_gelu_tensor != nullptr && pre_gelu_tensor->data.dptr != nullptr);
+
+      // Construct config
+      MatmulConfig config;
+      if (grad) {
+        config.dbias_tensor = bias[i];
+        config.with_dgelu_epilogue = with_gelu_dgelu_epilogue;
+      } else {
+        config.bias_tensor = bias[i];
+        config.with_gelu_epilogue = with_gelu_dgelu_epilogue;
+      }
+      config.epilogue_aux_tensor = pre_gelu_out[i];
+      config.use_split_accumulator = use_split_accumulator;
+      config.sm_count = math_sm_count;
+
+      // Launch GEMM
+      const float alpha = 1.f;
+      const float beta = accumulate ? 1.f : 0.f;
+      nvte_cublas_gemm_v2(transa, transb, &alpha, A[i], B[i], &beta, D[i], D[i],
+                          workspace[i % num_streams], &config,
+                          detail::get_compute_stream(i % num_streams));
     }
-    config.epilogue_aux_tensor = pre_gelu_out[i];
-    config.use_split_accumulator = use_split_accumulator;
-    config.sm_count = math_sm_count;
-
-    // Launch GEMM
-    const float alpha = 1.f;
-    const float beta = accumulate ? 1.f : 0.f;
-    nvte_cublas_gemm_v2(transa, transb, &alpha, A[i], B[i], &beta, D[i], D[i],
-                        workspace[i % num_streams], &config,
-                        detail::get_compute_stream(i % num_streams));
+  } catch (...) {
+    // Earlier GEMMs may still use caller-owned inputs, outputs, and scratch.
+    // Rejoin the helper streams before propagating the failure to the caller.
+    launch_error = std::current_exception();
   }
 
   // record events on compute streams
@@ -1153,6 +1161,9 @@ void multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETens
   for (int s = 0; s < num_stream_used; s++) {
     NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream, detail::get_compute_stream_event(s)));
   }
+  if (launch_error) {
+    std::rethrow_exception(launch_error);
+  }
 }
 
 void nvte_multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -49,34 +49,40 @@
 _NUM_MAX_UB_STREAMS = 3
 
 
-def get_cublas_workspace_size_bytes() -> None:
-    """Return 32 MiB if using hopper, 4 MiB for all other architectures."""
-    if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 9:
+@functools.lru_cache(maxsize=None)
+def _get_cublas_workspace_size_bytes_for_device(device: int) -> int:
+    """Cached per-device workspace size; the device properties are fixed per process."""
+    if torch.cuda.get_device_properties(device).major >= 9:
         # 32 MiB for NVFP4 GEMM, plus additional 1024 B for alignment and misc scales
         return 32 * 1024 * 1024 + 1024
     return 4_194_304
 
 
-@functools.lru_cache(maxsize=None)
+def get_cublas_workspace_size_bytes() -> int:
+    """Return 32 MiB if using hopper, 4 MiB for all other architectures."""
+    return _get_cublas_workspace_size_bytes_for_device(torch.cuda.current_device())
+
+
 def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Tensor:
-    """Returns workspace for cublas GEMM."""
+    """Allocate cuBLAS scratch for one invocation on the current CUDA stream.
+
+    The caching allocator safely reuses eager allocations on their stream. During
+    capture, the graph's private pool owns the allocation until the graph is
+    destroyed. Keeping the tensor in a Python cache would bypass that ownership
+    on later captures, including graphs replayed on different streams.
+    """
     assert not (ub and grouped_gemm), "UB is unsupported for grouped GEMM."
 
+    size = get_cublas_workspace_size_bytes()
     if ub:
-        return torch.empty(
-            get_cublas_workspace_size_bytes() * _NUM_MAX_UB_STREAMS,
-            dtype=torch.uint8,
-            device=device,
-        )
+        return torch.empty(size * _NUM_MAX_UB_STREAMS, dtype=torch.uint8, device=device)
     if grouped_gemm:
-        _multi_stream_cublas_workspace = []
-        for _ in range(tex.get_num_cublas_streams()):
-            _multi_stream_cublas_workspace.append(
-                torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device)
-            )
-        return _multi_stream_cublas_workspace
+        return [
+            torch.empty(size, dtype=torch.uint8, device=device)
+            for _ in range(tex.get_num_cublas_streams())
+        ]
 
-    return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device)
+    return torch.empty(size, dtype=torch.uint8, device=device)
 
 
 def validate_gemm_scale(scale: Optional[float], required: bool) -> float:
@@ -771,7 +777,6 @@ def general_gemm(
 
     alpha = validate_gemm_scale(alpha, True)
     beta = validate_gemm_scale(beta, accumulate)
-    workspace = get_cublas_workspace(A.device.index, ub is not None, False)
 
     # Temporary hack to route NVFP4 GEMM with UE5M3 scale factors to
     # cuDNN Frontend kernels. UE5M3-specific logic should be removed
@@ -820,6 +825,7 @@ def general_gemm(
 
     # If A or B are custom tensors -> dispatch to quantizers's qgemm implementation
     if is_custom(A) or is_custom(B):
+        workspace = get_cublas_workspace(A.device.index, ub is not None, False)
         return custom_gemm(
             A,
             B,
@@ -845,6 +851,10 @@ def general_gemm(
         # FP8 block-scaling requires split accumulator
         use_split_accumulator = True
 
+    # Native GEMMs allocate scratch through the framework allocator in C++.
+    # Communication overlap still requires an explicit Tensor workspace.
+    workspace = get_cublas_workspace(A.device.index, True, False) if ub is not None else None
+
     args = (
         A,
         transa,  # transa
@@ -859,7 +869,7 @@ def general_gemm(
         gelu_in,
         grad,  # grad
         workspace,
-        workspace.shape[0],
+        workspace.shape[0] if workspace is not None else 0,
         accumulate,
         use_split_accumulator,
     )
@@ -1076,7 +1086,6 @@ def general_grouped_gemm(
     out_dtype = TE_DType[out[0].dtype] if D_dtype is None else D_dtype
 
     sm_count = get_sm_count()
-    workspaces = get_cublas_workspace(A[0].device.index, False, True)
 
     if grad and use_bias:
         grad_bias = [
@@ -1153,8 +1162,8 @@ def general_grouped_gemm(
         single_output,
         gelu_input,  # this is pre_gelu_out
         grad,  # grad
-        workspaces,
-        workspaces[0].shape[0],
+        None,  # Allocate invocation-local scratch in the native binding.
+        0,
         accumulate,
         use_split_accumulator,
         sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))),

transformer_engine/pytorch/csrc/extensions.h

@@ -152,7 +152,7 @@ using MaybeTensor = std::optional<at::Tensor>;
 std::vector<py::object> gemm(py::handle A, bool transa, py::handle B, bool transb, py::object D,
                              py::handle quantizer, std::optional<DType> out_dtype, MaybeTensor bias,
                              DType bias_type, bool gelu, MaybeTensor gelu_in, bool grad,
-                             at::Tensor workspace, size_t workspaceSize, bool accumulate,
+                             MaybeTensor workspace, size_t workspaceSize, bool accumulate,
                              bool use_split_accumulator, CommOverlapCore *comm_overlap = nullptr,
                              std::optional<CommOverlapType> comm_type = std::nullopt,
                              MaybeTensor extra_output = std::nullopt, bool bulk_overlap = false,
@@ -177,8 +177,9 @@ std::optional<std::vector<at::Tensor>> te_general_grouped_gemm(
     std::vector<py::handle> A, bool transa, std::vector<py::handle> B, bool transb,
     std::optional<std::vector<at::Tensor>> D, DType D_type, std::vector<int64_t> m_splits,
     std::vector<at::Tensor> bias, DType bias_type, bool single_output,
-    std::vector<at::Tensor> pre_gelu_out, bool grad, std::vector<at::Tensor> workspace,
-    size_t workspaceSize, bool accumulate, bool use_split_accumulator, int math_sm_count);
+    std::vector<at::Tensor> pre_gelu_out, bool grad,
+    std::optional<std::vector<at::Tensor>> workspace, size_t workspaceSize, bool accumulate,
+    bool use_split_accumulator, int math_sm_count);
 
 py::object te_general_grouped_gemm_for_grouped_tensor(
     py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias,