From 93a80c7908eea41426ad2c6044f246bb76d3fd04 Mon Sep 17 00:00:00 2001 From: Matvey Saprykin Date: Mon, 5 Oct 2026 12:59:10 +0300 Subject: [PATCH 1/7] Cache cuBLAS workspace by CUDA stream The workspace cache currently ignores the calling CUDA stream. Overlapping GEMMs can therefore overwrite the same scratch buffer and return incorrect results. Include the stream in the cache key, keep reuse on each stream, and add ownership and FP32/BF16 numerical regression tests. Signed-off-by: Matvey Saprykin --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_gemm_workspace.py | 55 +++++++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 12 +++- 3 files changed, 66 insertions(+), 2 deletions(-) create mode 100644 tests/pytorch/test_gemm_workspace.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index b3b6ccacac7..457fb5b798d 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -31,6 +31,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" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 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" diff --git a/tests/pytorch/test_gemm_workspace.py b/tests/pytorch/test_gemm_workspace.py new file mode 100644 index 00000000000..28e265c696e --- /dev/null +++ b/tests/pytorch/test_gemm_workspace.py @@ -0,0 +1,55 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuBLAS workspace ownership and numerical checks for independent CUDA streams.""" + +import pytest +import torch + +from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm, get_cublas_workspace + + +@pytest.mark.parametrize("ub,grouped_gemm", [(False, False), (True, False), (False, True)]) +def test_workspace_is_reused_only_on_its_stream(ub, grouped_gemm): + """A workspace remains reusable without aliasing another stream's active scratch.""" + device = torch.cuda.current_device() + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + workspaces = [] + for stream in streams: + with torch.cuda.stream(stream): + first = get_cublas_workspace(device, ub, grouped_gemm) + assert get_cublas_workspace(device, ub, grouped_gemm) is first + workspaces.append(first if grouped_gemm else [first]) + assert len(workspaces[0]) == len(workspaces[1]) + assert {w.data_ptr() for w in workspaces[0]}.isdisjoint(w.data_ptr() for w in workspaces[1]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_concurrent_gemms_match_exact_reference(dtype): + """Large reductions exercise cuBLAS algorithms that use scratch for partial sums.""" + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + reduction, hidden, experts = 16384, 4096, 128 + inputs = [ + torch.full((reduction, hidden), value, dtype=dtype, device="cuda") for value in (1.0, 2.0) + ] + gradients = [ + torch.full((reduction, experts), value, dtype=dtype, device="cuda") for value in (1.0, 3.0) + ] + ready = torch.cuda.Event() + ready.record() + for stream in streams: + stream.wait_event(ready) + + outputs = [] + for _ in range(8): + for stream, inp, grad in zip(streams, inputs, gradients): + with torch.cuda.stream(stream): + out, *_ = general_gemm(inp, grad, dtype, layout="NT", grad=True) + outputs.append(out) + for stream in streams: + stream.synchronize() + # Both operands and the result are exactly representable in either dtype. + for index, out in enumerate(outputs): + expected = torch.full_like(out, reduction * (1 if index % 2 == 0 else 6)) + torch.testing.assert_close(out, expected, rtol=0, atol=0) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 6939847c8dc..56f23b872f5 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -48,9 +48,17 @@ def get_cublas_workspace_size_bytes() -> None: return 4_194_304 -@functools.lru_cache(maxsize=None) def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Tensor: - """Returns workspace for cublas GEMM.""" + """Returns workspace for cublas GEMM on the current CUDA stream.""" + stream = torch.cuda.current_stream(device).cuda_stream + return _get_cublas_workspace_for_stream(device, ub, grouped_gemm, stream) + + +@functools.lru_cache(maxsize=None) +def _get_cublas_workspace_for_stream( + device: int, ub: bool, grouped_gemm: bool, _stream: int +) -> torch.Tensor: + """Cache by stream so concurrent GEMMs do not overwrite each other's scratch space.""" assert not (ub and grouped_gemm), "UB is unsupported for grouped GEMM." if ub: From e5b55e918ffdbc1118a7dd0e1effbe731c55cbbd Mon Sep 17 00:00:00 2001 From: Matvey Saprykin Date: Mon, 5 Oct 2026 13:41:11 +0300 Subject: [PATCH 2/7] Let the CUDA allocator own cuBLAS workspaces Keep scratch local to each invocation so graph captures use their own private pools and temporary streams do not retain tensors in Python. Remove Linear's preallocation workaround for the global workspace cache. Cover concurrent graph replays, graph generations, scratch reclamation, grouped GEMM stream joins, and compiled Linear forward and backward. Signed-off-by: Matvey Saprykin --- tests/pytorch/test_gemm_workspace.py | 207 +++++++++++++++++- tests/pytorch/test_torch_compile.py | 43 ++++ .../pytorch/cpp_extensions/gemm.py | 15 +- transformer_engine/pytorch/module/linear.py | 35 --- 4 files changed, 246 insertions(+), 54 deletions(-) diff --git a/tests/pytorch/test_gemm_workspace.py b/tests/pytorch/test_gemm_workspace.py index 28e265c696e..013c126ea41 100644 --- a/tests/pytorch/test_gemm_workspace.py +++ b/tests/pytorch/test_gemm_workspace.py @@ -4,29 +4,109 @@ """cuBLAS workspace ownership and numerical checks for independent CUDA streams.""" +import gc +import weakref + import pytest import torch +from transformer_engine.pytorch.cpp_extensions import gemm -from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm, get_cublas_workspace +from transformer_engine.pytorch.cpp_extensions.gemm import ( + general_gemm, + general_grouped_gemm, + get_cublas_workspace, +) @pytest.mark.parametrize("ub,grouped_gemm", [(False, False), (True, False), (False, True)]) -def test_workspace_is_reused_only_on_its_stream(ub, grouped_gemm): - """A workspace remains reusable without aliasing another stream's active scratch.""" +def test_live_workspaces_do_not_alias(ub, grouped_gemm): + """Live invocations own distinct scratch, even on the same capture stream.""" device = torch.cuda.current_device() streams = [torch.cuda.Stream(), torch.cuda.Stream()] workspaces = [] for stream in streams: with torch.cuda.stream(stream): first = get_cublas_workspace(device, ub, grouped_gemm) - assert get_cublas_workspace(device, ub, grouped_gemm) is first workspaces.append(first if grouped_gemm else [first]) + second = get_cublas_workspace(device, ub, grouped_gemm) + workspaces.append(second if grouped_gemm else [second]) assert len(workspaces[0]) == len(workspaces[1]) - assert {w.data_ptr() for w in workspaces[0]}.isdisjoint(w.data_ptr() for w in workspaces[1]) + pointers = [w.data_ptr() for invocation in workspaces for w in invocation] + assert len(pointers) == len(set(pointers)) + + +@pytest.mark.parametrize("ub,grouped_gemm", [(False, False), (True, False), (False, True)]) +def test_temporary_streams_release_workspaces(ub, grouped_gemm): + """Scratch is returned to the allocator when the invocation goes out of scope.""" + device = torch.cuda.current_device() + torch.cuda.synchronize() + gc.collect() + allocated = torch.cuda.memory_allocated() + for _ in range(8): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + workspace = get_cublas_workspace(device, ub, grouped_gemm) + tensors = workspace if grouped_gemm else [workspace] + refs = [weakref.ref(tensor) for tensor in tensors] + for tensor in tensors: + tensor.fill_(1) + del tensor, tensors, workspace + stream.synchronize() + del stream + assert all(ref() is None for ref in refs) + assert torch.cuda.memory_allocated() == allocated + + +@pytest.mark.parametrize("ub,grouped_gemm", [(False, False), (True, False), (False, True)]) +def test_graph_workspaces_survive_replay_and_are_released(ub, grouped_gemm): + """Different captures on one stream must not borrow an earlier graph's scratch.""" + device = torch.cuda.current_device() + capture_stream = torch.cuda.Stream() + replay_streams = [torch.cuda.Stream(), torch.cuda.Stream()] + # Initialize the allocator/library state outside the graph, on a different stream. + workspace = get_cublas_workspace(device, ub, grouped_gemm) + count = len(workspace) if grouped_gemm else 1 + del workspace + outputs = [torch.empty(count, dtype=torch.uint8, device=device) for _ in range(2)] + torch.cuda.synchronize() + torch.cuda.empty_cache() + reserved = torch.cuda.memory_reserved() + allocated = torch.cuda.memory_allocated() + for _ in range(3): + graphs, pointers = [], [] + for value, output in enumerate(outputs, start=1): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + workspace = get_cublas_workspace(device, ub, grouped_gemm) + tensors = workspace if grouped_gemm else [workspace] + pointers.append({tensor.data_ptr() for tensor in tensors}) + refs = [weakref.ref(tensor) for tensor in tensors] + for index, tensor in enumerate(tensors): + tensor.fill_(value) + output[index : index + 1].copy_(tensor[:1]) + del tensor, tensors, workspace + assert all(ref() is None for ref in refs) + graphs.append(graph) + assert pointers[0].isdisjoint(pointers[1]) + for _ in range(8): + for graph, stream in zip(graphs, replay_streams): + with torch.cuda.stream(stream): + graph.replay() + torch.cuda.synchronize() + for value, output in enumerate(outputs, start=1): + torch.testing.assert_close(output, torch.full_like(output, value), rtol=0, atol=0) + for graph in graphs: + graph.reset() + del graph, graphs + gc.collect() + torch.cuda.empty_cache() + assert torch.cuda.memory_allocated() == allocated + assert torch.cuda.memory_reserved() == reserved @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_concurrent_gemms_match_exact_reference(dtype): +@pytest.mark.parametrize("execution", ["eager", "graph"]) +def test_concurrent_gemms_match_exact_reference(dtype, execution): """Large reductions exercise cuBLAS algorithms that use scratch for partial sums.""" streams = [torch.cuda.Stream(), torch.cuda.Stream()] reduction, hidden, experts = 16384, 4096, 128 @@ -41,15 +121,122 @@ def test_concurrent_gemms_match_exact_reference(dtype): for stream in streams: stream.wait_event(ready) - outputs = [] + outputs, graphs = [], [] + if execution == "graph": + capture_stream = torch.cuda.Stream() + capture_stream.wait_event(ready) + for inp, grad in zip(inputs, gradients): + with torch.cuda.stream(capture_stream): + general_gemm(inp, grad, dtype, layout="NT", grad=True) + capture_stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + out, *_ = general_gemm(inp, grad, dtype, layout="NT", grad=True) + graphs.append(graph) + outputs.append(out) for _ in range(8): - for stream, inp, grad in zip(streams, inputs, gradients): + for index, (stream, inp, grad) in enumerate(zip(streams, inputs, gradients)): with torch.cuda.stream(stream): - out, *_ = general_gemm(inp, grad, dtype, layout="NT", grad=True) - outputs.append(out) + if execution == "graph": + graphs[index].replay() + else: + out, *_ = general_gemm(inp, grad, dtype, layout="NT", grad=True) + outputs.append(out) for stream in streams: stream.synchronize() # Both operands and the result are exactly representable in either dtype. for index, out in enumerate(outputs): expected = torch.full_like(out, reduction * (1 if index % 2 == 0 else 6)) torch.testing.assert_close(out, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("execution", ["eager", "graph"]) +def test_grouped_gemms_match_exact_reference(dtype, execution): + """Internal GEMM streams must join the caller before its scratch is recycled.""" + reduction, hidden, experts, count = 8192, 512, 64, 4 + inputs = [torch.ones((reduction, hidden), dtype=dtype, device="cuda") for _ in range(count)] + gradients = [ + torch.full((reduction, experts), index + 1, dtype=dtype, device="cuda") + for index in range(count) + ] + outputs = [torch.empty((experts, hidden), dtype=dtype, device="cuda") for _ in range(count)] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + + def run(): + general_grouped_gemm( + inputs, + gradients, + outputs, + [None] * count, + dtype, + layout="NT", + m_splits=[experts] * count, + grad=True, + ) + + with torch.cuda.stream(stream): + run() + stream.synchronize() + if execution == "graph": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run() + for _ in range(8): + with torch.cuda.stream(stream): + if execution == "graph": + graph.replay() + else: + run() + # Recycle same-size allocations immediately after the internal-stream join. + scratch = get_cublas_workspace(torch.cuda.current_device(), False, True) + for tensor in scratch: + tensor.fill_(0) + del tensor, scratch + stream.synchronize() + for index, output in enumerate(outputs): + expected = torch.full_like(output, reduction * (index + 1)) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + not hasattr(torch.library, "custom_op") + or not hasattr(torch.compiler, "cudagraph_mark_step_begin"), + reason="Custom operators and CUDA graph trees require newer PyTorch", +) +def test_compiled_gemm_graph_generations(monkeypatch): + """An opaque GEMM's scratch must remain valid when graph trees are rerecorded.""" + original = gemm.get_cublas_workspace + captured_workspaces = [] + + def allocate(*args, **kwargs): + workspace = original(*args, **kwargs) + if torch.cuda.is_current_stream_capturing(): + captured_workspaces.append(weakref.ref(workspace)) + return workspace + + monkeypatch.setattr(gemm, "get_cublas_workspace", allocate) + + @torch.library.custom_op("te_workspace_test::gemm", mutates_args=()) + def run(inp: torch.Tensor, grad: torch.Tensor) -> torch.Tensor: + return general_gemm(inp, grad, inp.dtype, layout="NT", grad=True)[0] + + @run.register_fake + def fake(inp, grad): + return inp.new_empty((grad.shape[1], inp.shape[1])) + + torch._dynamo.reset() + compiled = torch.compile(run, fullgraph=True, mode="reduce-overhead", dynamic=False) + for reduction, value in [(8192, 1), (16384, 2), (8192, 3)]: + inp = torch.full((reduction, 512), value, dtype=torch.bfloat16, device="cuda") + grad = torch.ones((reduction, 64), dtype=torch.bfloat16, device="cuda") + for _ in range(3): + torch.compiler.cudagraph_mark_step_begin() + out = compiled(inp, grad) + torch.testing.assert_close(out, torch.full_like(out, reduction * value), rtol=0, atol=0) + del out + assert len(captured_workspaces) >= 2, "Expected CUDA graph captures for different shapes" + assert all(ref() is None for ref in captured_workspaces) + del compiled + torch._dynamo.reset() diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index eae6f0a8a24..e0e28f9b882 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1892,6 +1892,49 @@ 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.""" + import weakref + from transformer_engine.pytorch.cpp_extensions import gemm + + original = gemm.get_cublas_workspace + captured_workspaces = [] + + def allocate(*args, **kwargs): + workspace = original(*args, **kwargs) + if torch.cuda.is_current_stream_capturing(): + captured_workspaces.append(weakref.ref(workspace)) + return workspace + + monkeypatch.setattr(gemm, "get_cublas_workspace", allocate) + 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_workspaces) >= 2, "Expected CUDA graph captures for different shapes" + assert all(ref() is None for ref in captured_workspaces) + 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( diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 56f23b872f5..e7a696cfa4b 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -49,16 +49,13 @@ def get_cublas_workspace_size_bytes() -> None: def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Tensor: - """Returns workspace for cublas GEMM on the current CUDA stream.""" - stream = torch.cuda.current_stream(device).cuda_stream - return _get_cublas_workspace_for_stream(device, ub, grouped_gemm, stream) + """Allocate cuBLAS scratch for one invocation on the current CUDA stream. - -@functools.lru_cache(maxsize=None) -def _get_cublas_workspace_for_stream( - device: int, ub: bool, grouped_gemm: bool, _stream: int -) -> torch.Tensor: - """Cache by stream so concurrent GEMMs do not overwrite each other's scratch space.""" + 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." if ub: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 7f00e19644a..4e4c8d138d4 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -72,7 +72,6 @@ ) from ..cpp_extensions import ( general_gemm, - get_cublas_workspace, ) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..graph import is_graph_capturing @@ -2302,40 +2301,6 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) - self._prealloc_cublas_workspace() - - def _apply(self, *args, **kwargs): - out = super()._apply(*args, **kwargs) - # Params may have just landed on CUDA (.to()/.cuda()/to_empty()). - self._prealloc_cublas_workspace() - return out - - def _prealloc_cublas_workspace(self) -> None: - """Allocate the process-global cuBLAS workspaces before the first GEMM: - under torch.compile it can run inside a cudagraph pool, and a workspace - first allocated there would not survive across graph generations.""" - # pylint: disable=import-outside-toplevel - from torch._guards import detect_fake_mode - from torch._subclasses.fake_tensor import is_fake - - weight = getattr(self, self.weight_names[0], None) - if weight is None or weight.device.type != "cuda": - return - if is_fake(weight) or detect_fake_mode(): - return - get_cublas_workspace(weight.device.index, False, False) - if self.ub_name is not None and any( - ( - self.ub_overlap_rs_fprop, - self.ub_overlap_ag_dgrad, - self.ub_overlap_ag_fprop, - self.ub_overlap_rs_dgrad, - self.ub_bulk_dgrad, - self.ub_bulk_wgrad, - ) - ): - get_cublas_workspace(weight.device.index, True, False) - def forward( self, inp: torch.Tensor, From ec5f5f9ce277bf0e92c9833f73597a123c9cb59c Mon Sep 17 00:00:00 2001 From: Matvey Saprykin Date: Mon, 5 Oct 2026 14:08:47 +0300 Subject: [PATCH 3/7] Test concurrent router dgrad workspace ownership Cover NN router dgrad with FP32 and BF16 operands in concurrent eager execution and independent CUDA graphs. Preserve NT wgrad coverage and exact references while retaining only one pair of large eager outputs. All 22 workspace cases pass on native SM90 H200 with CUDA 13 and PyTorch 2.13, including all four new NN cases. Signed-off-by: Matvey Saprykin --- tests/pytorch/test_gemm_workspace.py | 55 +++++++++++++++++++++------- 1 file changed, 41 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_gemm_workspace.py b/tests/pytorch/test_gemm_workspace.py index 013c126ea41..17bceddaa90 100644 --- a/tests/pytorch/test_gemm_workspace.py +++ b/tests/pytorch/test_gemm_workspace.py @@ -106,48 +106,75 @@ def test_graph_workspaces_survive_replay_and_are_released(ub, grouped_gemm): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("execution", ["eager", "graph"]) -def test_concurrent_gemms_match_exact_reference(dtype, execution): - """Large reductions exercise cuBLAS algorithms that use scratch for partial sums.""" +@pytest.mark.parametrize("layout", ["NT", "NN"]) +def test_concurrent_gemms_match_exact_reference(dtype, execution, layout): + """Concurrent router wgrad and dgrad must own scratch in eager and graph execution.""" streams = [torch.cuda.Stream(), torch.cuda.Stream()] - reduction, hidden, experts = 16384, 4096, 128 - inputs = [ - torch.full((reduction, hidden), value, dtype=dtype, device="cuda") for value in (1.0, 2.0) - ] + hidden, experts = 4096, 128 + if layout == "NT": + # wgrad = grad_output.T @ input; reduce over tokens. + tokens = reduction = 16384 + input_shape, output_shape = (tokens, hidden), (experts, hidden) + else: + # dgrad = grad_output @ weight; match the router's NN [8192, 128, 4096] GEMM. + tokens, reduction = 8192, experts + input_shape, output_shape = (experts, hidden), (tokens, hidden) + inputs = [torch.full(input_shape, value, dtype=dtype, device="cuda") for value in (1.0, 2.0)] gradients = [ - torch.full((reduction, experts), value, dtype=dtype, device="cuda") for value in (1.0, 3.0) + torch.full((tokens, experts), value, dtype=dtype, device="cuda") for value in (1.0, 3.0) ] + # Both operands and these analytic results are exactly representable in either dtype. + expected = [reduction, reduction * 6] ready = torch.cuda.Event() ready.record() for stream in streams: stream.wait_event(ready) - outputs, graphs = [], [] + outputs, graphs, checks = [], [], [] if execution == "graph": capture_stream = torch.cuda.Stream() capture_stream.wait_event(ready) for inp, grad in zip(inputs, gradients): with torch.cuda.stream(capture_stream): - general_gemm(inp, grad, dtype, layout="NT", grad=True) + general_gemm(inp, grad, dtype, layout=layout, grad=True) capture_stream.synchronize() + # Separate private pools, captured on one stream and replayed on two others. graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=capture_stream): - out, *_ = general_gemm(inp, grad, dtype, layout="NT", grad=True) + out, *_ = general_gemm(inp, grad, dtype, layout=layout, grad=True) + assert out.shape == output_shape + assert out.dtype == dtype graphs.append(graph) outputs.append(out) + del out for _ in range(8): for index, (stream, inp, grad) in enumerate(zip(streams, inputs, gradients)): with torch.cuda.stream(stream): if execution == "graph": graphs[index].replay() else: - out, *_ = general_gemm(inp, grad, dtype, layout="NT", grad=True) + out, *_ = general_gemm(inp, grad, dtype, layout=layout, grad=True) + assert out.shape == output_shape + assert out.dtype == dtype outputs.append(out) + del out + if execution == "eager": + # Submit both GEMMs before the checks. Keep only this pair of large NN outputs. + for index, (stream, out) in enumerate(zip(streams, outputs)): + with torch.cuda.stream(stream): + checks.append(torch.all(out == expected[index])) + del out + outputs.clear() for stream in streams: stream.synchronize() - # Both operands and the result are exactly representable in either dtype. for index, out in enumerate(outputs): - expected = torch.full_like(out, reduction * (1 if index % 2 == 0 else 6)) - torch.testing.assert_close(out, expected, rtol=0, atol=0) + checks.append(torch.all(out == expected[index])) + torch.testing.assert_close( + torch.stack(checks), + torch.ones(len(checks), dtype=torch.bool, device="cuda"), + rtol=0, + atol=0, + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) From f01b35ebc9590d88a9bc6c50169d1bf4a40368bf Mon Sep 17 00:00:00 2001 From: Matvey Saprykin Date: Fri, 9 Oct 2026 00:45:04 +0300 Subject: [PATCH 4/7] Adapt BasicLinear graph tests to invocation-local workspaces Replace the cold-start test's cache_info assertions with checks for capture-time allocation and release of workspace tensors. Keep output and gradient comparisons across CUDA, CPU and meta initialization. Remove BasicLinear's obsolete cache preallocation and use the existing opaque-object capability guard for compiled execution. Signed-off-by: Matvey Saprykin --- tests/pytorch/test_torch_compile.py | 27 ++++++++++++++----- .../pytorch/ops/basic/basic_linear.py | 19 +------------ 2 files changed, 21 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 9d1334a9d49..e0e956f974a 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -2959,16 +2959,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("initial_device", ["cuda", "cpu", "meta"]) @@ -2978,23 +2979,33 @@ def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device): """ import contextlib import sys + import weakref import torch import transformer_engine.pytorch as te - 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 = sys.argv[1:] training = training == "True" - assert get_cublas_workspace.cache_info().currsize == 0 + workspaces, captured_workspaces = [], [] + allocate_workspace = gemm.get_cublas_workspace + + def allocate(*args, **kwargs): + workspace = allocate_workspace(*args, **kwargs) + ref = weakref.ref(workspace) + workspaces.append(ref) + if torch.cuda.is_current_stream_capturing(): + captured_workspaces.append(ref) + return workspace + + gemm.get_cublas_workspace = allocate linear = te.ops.BasicLinear(32, 64, 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 workspaces, "Module initialization must not allocate GEMM scratch" model = te.ops.Sequential(linear, te.ops.Bias(64, dtype=torch.bfloat16)) x = torch.randn(32, 32, device="cuda", dtype=torch.bfloat16, requires_grad=training) targets = (x, *model.parameters()) @@ -3020,6 +3031,8 @@ def forward(x): del grads, expected_grads del actual, expected torch.cuda.synchronize() + assert captured_workspaces, "Expected workspace allocations during CUDA graph capture" + assert all(ref() is None for ref in workspaces), "Invocations must release workspace tensors" assert not counters["inductor"]["cudagraph_skips"], counters["inductor"] assert counters["inductor"]["cudagraph_recorded_non_static_inputs"] > 0, dict( counters["inductor"] diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index a7b09198250..81e127fb96e 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -13,7 +13,7 @@ import torch -from ...cpp_extensions import general_gemm, get_cublas_workspace +from ...cpp_extensions import general_gemm from ...dynamo import TensorOrQuantized, TensorSpec from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import ( @@ -411,23 +411,6 @@ def reset_parameters(self) -> None: if not isinstance(weight, torch.nn.Parameter): weight = torch.nn.Parameter(weight) self.weight = weight - self._prealloc_cublas_workspace() - - def _apply(self, *args, **kwargs): - out = super()._apply(*args, **kwargs) - self._prealloc_cublas_workspace() - return out - - def _prealloc_cublas_workspace(self) -> None: - """Keep the cached GEMM workspace outside CUDA graph private pools.""" - # pylint: disable=import-outside-toplevel - from torch._guards import detect_fake_mode - from torch._subclasses.fake_tensor import is_fake - - weight = self.weight - if weight.device.type != "cuda" or is_fake(weight) or detect_fake_mode(): - return - get_cublas_workspace(weight.device.index, False, False) def pre_first_fuser_forward(self) -> None: super().pre_first_fuser_forward() From 8f6ed7805ed8e2e02c0bb3a652a147b9d03e1171 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 20 Oct 2026 10:09:15 -0700 Subject: [PATCH 5/7] Don't allocate workspace for E5M3 path and cache the workspace size Signed-off-by: Kirthi Shankar Sivamani --- .../pytorch/cpp_extensions/gemm.py | 35 ++++++++++--------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index d2e83e54acc..c6faae5161d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -40,14 +40,20 @@ _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 +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: """Allocate cuBLAS scratch for one invocation on the current CUDA stream. @@ -58,21 +64,16 @@ def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Ten """ 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: @@ -767,7 +768,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 @@ -798,6 +798,9 @@ def general_gemm( bulk_overlap, ) + # Allocated after the UE5M3 path above, which does not use cuBLAS scratch. + workspace = get_cublas_workspace(A.device.index, ub is not None, False) + if ub_type is not None: assert ub is not None, ( f"{'AG+GEMM' if ub_type == tex.CommOverlapType.AG else 'GEMM+RS'} overlap requires" From 10f6369f2a82eb5d65bfc5ce52ca41cb1ceec2ab Mon Sep 17 00:00:00 2001 From: Matvey Saprykin Date: Thu, 22 Oct 2026 23:34:55 +0300 Subject: [PATCH 6/7] Allocate cuBLAS workspaces in the native binding Use an invocation-local DataPtr from the PyTorch CUDA allocator instead of creating Python workspace tensors for ordinary GEMMs. Allocate one buffer with aligned, disjoint slices for the grouped multistream path. Keep stream-ordered reuse and graph-pool ownership without adding a persistent cache. Preserve explicit workspaces for custom recipes and communication overlap, and observe native requests in the compiled graph-generation regression. Signed-off-by: Matvey Saprykin --- tests/pytorch/test_gemm_workspace.py | 22 +++-- .../pytorch/cpp_extensions/gemm.py | 15 ++-- transformer_engine/pytorch/csrc/extensions.h | 7 +- .../pytorch/csrc/extensions/gemm.cpp | 84 ++++++++++++++++--- 4 files changed, 99 insertions(+), 29 deletions(-) diff --git a/tests/pytorch/test_gemm_workspace.py b/tests/pytorch/test_gemm_workspace.py index 17bceddaa90..49613058857 100644 --- a/tests/pytorch/test_gemm_workspace.py +++ b/tests/pytorch/test_gemm_workspace.py @@ -234,16 +234,16 @@ def run(): ) def test_compiled_gemm_graph_generations(monkeypatch): """An opaque GEMM's scratch must remain valid when graph trees are rerecorded.""" - original = gemm.get_cublas_workspace - captured_workspaces = [] + original = gemm.tex.generic_gemm + captured_workspace_requests = [] - def allocate(*args, **kwargs): - workspace = original(*args, **kwargs) + def observe_native_gemm(*args, **kwargs): if torch.cuda.is_current_stream_capturing(): - captured_workspaces.append(weakref.ref(workspace)) - return workspace + # Record only scalar metadata, never tensors or capture-pool owners. + captured_workspace_requests.append((args[12] is None, args[13])) + return original(*args, **kwargs) - monkeypatch.setattr(gemm, "get_cublas_workspace", allocate) + monkeypatch.setattr(gemm.tex, "generic_gemm", observe_native_gemm) @torch.library.custom_op("te_workspace_test::gemm", mutates_args=()) def run(inp: torch.Tensor, grad: torch.Tensor) -> torch.Tensor: @@ -263,7 +263,11 @@ def fake(inp, grad): out = compiled(inp, grad) torch.testing.assert_close(out, torch.full_like(out, reduction * value), rtol=0, atol=0) del out - assert len(captured_workspaces) >= 2, "Expected CUDA graph captures for different shapes" - assert all(ref() is None for ref in captured_workspaces) + 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() diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index c6faae5161d..24fa995b142 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -798,9 +798,6 @@ def general_gemm( bulk_overlap, ) - # Allocated after the UE5M3 path above, which does not use cuBLAS scratch. - workspace = get_cublas_workspace(A.device.index, ub is not None, False) - if ub_type is not None: assert ub is not None, ( f"{'AG+GEMM' if ub_type == tex.CommOverlapType.AG else 'GEMM+RS'} overlap requires" @@ -819,6 +816,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, @@ -844,6 +842,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 @@ -858,7 +860,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, ) @@ -987,7 +989,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 = [ @@ -1063,8 +1064,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))), diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index d3e7763a7b5..10c94974668 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -151,7 +151,7 @@ using MaybeTensor = std::optional; std::vector gemm(py::handle A, bool transa, py::handle B, bool transb, py::object D, py::handle quantizer, std::optional 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 comm_type = std::nullopt, MaybeTensor extra_output = std::nullopt, bool bulk_overlap = false, @@ -170,8 +170,9 @@ std::optional> te_general_grouped_gemm( std::vector A, bool transa, std::vector B, bool transb, std::optional> D, DType D_type, std::vector m_splits, std::vector bias, DType bias_type, bool single_output, - std::vector pre_gelu_out, bool grad, std::vector workspace, - size_t workspaceSize, bool accumulate, bool use_split_accumulator, int math_sm_count); + std::vector pre_gelu_out, bool grad, + std::optional> 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, diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index c3484802246..dd41ff75565 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -4,7 +4,9 @@ * See LICENSE for license information. ************************************************************************/ +#include #include +#include #include #include @@ -18,6 +20,28 @@ namespace { +at::Device gemm_device(pybind11::handle tensor) { + NVTE_CHECK(!tensor.is_none(), "Tensor A has not been provided"); + // Avoid a Python property lookup for ordinary tensors. Quantized storage + // objects expose the device of their underlying data through .device. + const auto device = THPVariable_CheckExact(tensor.ptr()) + ? THPVariable_Unpack(tensor.ptr()).device() + : tensor.attr("device").cast(); + NVTE_CHECK(device.is_cuda(), "GEMM expects CUDA operands"); + return device; +} + +size_t cublas_workspace_bytes(int device_id) { + // Match the Tensor-workspace path, including NVFP4 alignment/scales storage. + return transformer_engine::cuda::sm_arch(device_id) >= 90 ? 32 * 1024 * 1024 + 1024 : 4194304; +} + +void check_workspace(const at::Tensor& workspace, size_t size) { + NVTE_CHECK(workspace.is_cuda() && workspace.scalar_type() == at::kByte && + workspace.is_contiguous() && size <= static_cast(workspace.numel()), + "Workspace must be a contiguous CUDA byte tensor with at least workspace_size bytes"); +} + void* get_data_ptr(transformer_engine::pytorch::MaybeTensor tensor) { if (tensor.has_value()) return tensor->data_ptr(); return nullptr; @@ -160,7 +184,7 @@ std::pair createOutputTensor(const std::vector gemm(py::handle A, bool transa, py::handle B, bool transb, py::object D, py::handle quantizer, std::optional 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, std::optional comm_type, MaybeTensor extra_output, bool bulk_overlap, float alpha, std::optional beta) { @@ -169,7 +193,14 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans // Ensure that cublasLt handle is created on the correct device, // overriding torch.cuda.set_device calls from user side. // Assumes all tensors passed are on the same device. - at::cuda::CUDAGuard device_guard(workspace.device()); + at::cuda::CUDAGuard device_guard(workspace ? workspace->device() : gemm_device(A)); + if (workspace) { + check_workspace(*workspace, workspaceSize); + } else { + NVTE_CHECK(comm_overlap == nullptr, "Communication overlap requires an explicit workspace"); + NVTE_CHECK(workspaceSize == 0, "workspace_size must be zero when workspace is None"); + workspaceSize = cublas_workspace_bytes(device_guard.current_device().index()); + } // Input tensors NVTE_CHECK(!A.is_none(), "Tensor A has not been provided"); @@ -304,8 +335,16 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans makeTransformerEngineTensor(get_data_ptr(pre_gelu_out), gelu_shape, gelu_type); // Workspace - auto te_workspace = makeTransformerEngineTensor(workspace.data_ptr(), - std::vector{workspaceSize}, DType::kByte); + // Keep scratch alive until all work has been submitted. The CUDA allocator + // orders eager reuse on this stream and owns capture allocations in the + // graph's private pool. No Python Tensor or persistent workspace cache is needed. + c10::DataPtr workspace_owner; + if (!workspace) { + workspace_owner = c10::cuda::CUDACachingAllocator::get()->allocate(workspaceSize); + } + auto te_workspace = + makeTransformerEngineTensor(workspace ? workspace->data_ptr() : workspace_owner.get(), + std::vector{workspaceSize}, DType::kByte); // Set an external SM Margin to all the GEMMs. // This comes in handy when DP is overlapped with GEMMs @@ -498,8 +537,9 @@ std::optional> te_general_grouped_gemm( std::vector A, bool transa, std::vector B, bool transb, std::optional> D, DType D_type, std::vector m_splits, std::vector bias, DType bias_type, bool single_output, - std::vector pre_gelu_out, bool grad, std::vector workspace, - size_t workspaceSize, bool accumulate, bool use_split_accumulator, int math_sm_count) { + std::vector pre_gelu_out, bool grad, + std::optional> workspace, size_t workspaceSize, bool accumulate, + bool use_split_accumulator, int math_sm_count) { if (single_output && D == std::nullopt) { NVTE_ERROR("not implemented, D should be allocated for single output case."); } @@ -507,7 +547,21 @@ std::optional> te_general_grouped_gemm( // Ensure that cublasLt handle is created on the correct device, // overriding torch.cuda.set_device calls from user side. // Assumes all tensors passed are on the same device. - at::cuda::CUDAGuard device_guard(workspace[0].device()); + NVTE_CHECK(!A.empty(), "Grouped GEMM expects at least one operand"); + const size_t num_workspaces = nvte_get_num_compute_streams(); + NVTE_CHECK(!workspace || workspace->size() == num_workspaces, + "Grouped GEMM expects one workspace per compute stream"); + at::cuda::CUDAGuard device_guard(workspace ? (*workspace)[0].device() : gemm_device(A[0])); + if (workspace) { + for (const auto& buffer : *workspace) { + check_workspace(buffer, workspaceSize); + NVTE_CHECK(buffer.device() == (*workspace)[0].device(), + "Grouped GEMM workspaces must be on the same device"); + } + } else { + NVTE_CHECK(workspaceSize == 0, "workspace_size must be zero when workspace is None"); + workspaceSize = cublas_workspace_bytes(device_guard.current_device().index()); + } void* output_data_ptr = nullptr; if (single_output) { @@ -644,9 +698,19 @@ std::optional> te_general_grouped_gemm( std::vector te_workspace_vector; std::vector te_workspace_wrappers; - for (size_t i = 0; i < workspace.size(); i++) { - auto wsp = makeTransformerEngineTensor(workspace[i].data_ptr(), - std::vector{workspaceSize}, DType::kByte); + // cuBLASLt requires 256-byte alignment for every workspace. One allocation + // supplies disjoint slices; the native grouped call joins its helper streams + // back to the caller before this owner is released. + const size_t workspace_stride = (workspaceSize + 255) / 256 * 256; + c10::DataPtr workspace_owner; + if (!workspace) { + workspace_owner = + c10::cuda::CUDACachingAllocator::get()->allocate(workspace_stride * num_workspaces); + } + for (size_t i = 0; i < num_workspaces; i++) { + void* data = workspace ? (*workspace)[i].data_ptr() + : static_cast(workspace_owner.get()) + i * workspace_stride; + auto wsp = makeTransformerEngineTensor(data, std::vector{workspaceSize}, DType::kByte); te_workspace_vector.emplace_back(wsp.data()); te_workspace_wrappers.emplace_back(std::move(wsp)); } From 48a48c91896a4b51f2f77c9c675159b749ca5e78 Mon Sep 17 00:00:00 2001 From: Matvey Saprykin Date: Fri, 23 Oct 2026 00:06:32 +0300 Subject: [PATCH 7/7] Update compiled workspace tests for native allocation Observe generic_gemm requests during capture instead of the Python workspace factory, which ordinary GEMMs no longer call. Retain only scalar metadata and require the native None/zero workspace contract. Keep the cold-start initialization check by rejecting Python workspace allocation, and preserve the existing fullgraph, capture, and numerical assertions for every initialization and training case. Signed-off-by: Matvey Saprykin --- tests/pytorch/test_torch_compile.py | 54 +++++++++++++++++------------ 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 2a2a51dd61b..4a8fff11f2b 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1936,19 +1936,18 @@ 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.""" - import weakref from transformer_engine.pytorch.cpp_extensions import gemm - original = gemm.get_cublas_workspace - captured_workspaces = [] + original = gemm.tex.generic_gemm + captured_workspace_requests = [] - def allocate(*args, **kwargs): - workspace = original(*args, **kwargs) + def observe_native_gemm(*args, **kwargs): if torch.cuda.is_current_stream_capturing(): - captured_workspaces.append(weakref.ref(workspace)) - return workspace + # 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, "get_cublas_workspace", allocate) + 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) @@ -1970,8 +1969,12 @@ def allocate(*args, **kwargs): model.weight.grad, torch.full_like(model.weight, batch * value), rtol=0, atol=0 ) del out, inp - assert len(captured_workspaces) >= 2, "Expected CUDA graph captures for different shapes" - assert all(ref() is None for ref in captured_workspaces) + 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() @@ -3115,7 +3118,6 @@ def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device, quantizati import contextlib import copy import sys - import weakref import torch import transformer_engine.pytorch as te from transformer_engine.common import recipe @@ -3124,18 +3126,22 @@ def test_te_ops_linear_bias_cold_cudagraphs(training, initial_device, quantizati training, initial_device, quantization = sys.argv[1:] training = training == "True" - workspaces, captured_workspaces = [], [] - allocate_workspace = gemm.get_cublas_workspace + workspace_requests, captured_workspace_requests = [], [] + original = gemm.tex.generic_gemm - def allocate(*args, **kwargs): - workspace = allocate_workspace(*args, **kwargs) - ref = weakref.ref(workspace) - workspaces.append(ref) + 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_workspaces.append(ref) - return workspace + captured_workspace_requests.append(request) + return original(*args, **kwargs) - gemm.get_cublas_workspace = allocate + 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( @@ -3156,7 +3162,7 @@ def allocate(*args, **kwargs): elif initial_device == "meta": linear.to_empty(device="cuda") linear.reset_parameters() - assert not workspaces, "Module initialization must not allocate GEMM scratch" + 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( @@ -3189,8 +3195,10 @@ def forward(x, module): del grads, expected_grads del actual, expected torch.cuda.synchronize() - assert captured_workspaces, "Expected workspace allocations during CUDA graph capture" - assert all(ref() is None for ref in workspaces), "Invocations must release workspace tensors" + 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"]