diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d10..4e01d7a7972 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" 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..49613058857 --- /dev/null +++ b/tests/pytorch/test_gemm_workspace.py @@ -0,0 +1,273 @@ +# 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 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, + general_grouped_gemm, + get_cublas_workspace, +) + + +@pytest.mark.parametrize("ub,grouped_gemm", [(False, False), (True, False), (False, True)]) +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) + 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]) + 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]) +@pytest.mark.parametrize("execution", ["eager", "graph"]) +@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()] + 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((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, 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=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=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=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() + for index, out in enumerate(outputs): + 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]) +@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.tex.generic_gemm + captured_workspace_requests = [] + + def observe_native_gemm(*args, **kwargs): + if torch.cuda.is_current_stream_capturing(): + # 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.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: + 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_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/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6868de7ef18..4a8fff11f2b 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/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"] diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f2f8537cc23..24fa995b142 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -40,34 +40,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: @@ -762,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 @@ -811,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, @@ -836,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 @@ -850,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, ) @@ -979,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 = [ @@ -1055,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 f6cf7d2fa77..e0a94774b33 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)); } diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 808de29c134..66c287b478a 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -73,7 +73,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 @@ -2325,40 +2324,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, diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index df8e13604a9..69e1c92d1a8 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 ( @@ -415,23 +415,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_fuser_forward(self, *, requires_grad: bool) -> None: super().pre_fuser_forward(requires_grad=requires_grad)