From 58233b2f4bf42b1bfab3084335afafc32fceb7cc Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 2 Oct 2026 22:53:41 +0800 Subject: [PATCH 1/5] fix(pt): stabilize full-batch fitting GEMM reductions Port the reviewed opt-in policy onto current master. Select a fixed canonical-shape cuBLASLt algorithm while evaluating complete node matrices in pedantic FP32, with device/thread/stream-local scratch and explicit unsupported-shape errors. Preserve the legacy default and scientific tolerances. Make the CUDA regression run in a fresh subprocess under ordinary GPU CI, and cover minimally aligned features in addition to batch/gradient invariance, independent streams and host threads, and a double-precision autograd reference. The current fitting sources compiled and all focused GPU checks passed on n2 under Slurm job 2030. Coding-Agent: Codex Codex-Version: codex-cli 0.160.0 Model: gpt-6-astra Reasoning-Effort: xhigh --- doc/env.md | 18 ++ source/op/pt/graph_fitting.cu | 14 +- source/op/pt/graph_fitting_gemm.h | 252 ++++++++++++++++++ .../test_graph_fitting_batch_invariant.py | 171 ++++++++++++ 4 files changed, 453 insertions(+), 2 deletions(-) create mode 100644 source/op/pt/graph_fitting_gemm.h create mode 100644 source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py diff --git a/doc/env.md b/doc/env.md index ba08490a06..4e9959f70b 100644 --- a/doc/env.md +++ b/doc/env.md @@ -55,6 +55,24 @@ Control high (double) or low (float) precision of training. {{ tensorflow_icon }} Enable JIT. Note that this option may either improve or decrease the performance. Requires TensorFlow to support JIT. ::: +:::{envvar} DP_GRAPH_FITTING_GEMM_POLICY + +**Default**: `legacy` + +For fused CUDA graph fitting, `shape_stable` selects pedantic-FP32 cuBLASLt +GEMMs with a configuration chosen at a fixed 256-node heuristic shape. +The actual operation evaluates the complete node matrix: it does not split +inference into frame subbatches or change precision. The default `legacy` +retains the existing SGEMM policy. Set the policy before the first CUDA fitting +call; it is immutable within a process. CPU fitting is unaffected. + +The shape-stable path caches launch metadata and 128 MiB of workspace per +thread/device/stream. Unsupported shapes or workspace requirements fail +explicitly rather than selecting a different reduction algorithm. Qualify +energy and force parity for the desired model, shapes, GPU, and cuBLAS version; +this setting is not a cross-hardware or cross-version reproducibility promise. +::: + :::{envvar} DP_INFER_BATCH_SIZE **Default**: `1024` on CPUs; automatically sized on GPUs diff --git a/source/op/pt/graph_fitting.cu b/source/op/pt/graph_fitting.cu index af997a1bf3..468ac77180 100644 --- a/source/op/pt/graph_fitting.cu +++ b/source/op/pt/graph_fitting.cu @@ -36,6 +36,7 @@ #include #include +#include "graph_fitting_gemm.h" #include "graph_ops.h" namespace { @@ -48,8 +49,9 @@ namespace { cublasHandle_t cublas_handle() { // A cuBLAS handle is device-bound and unsafe to share across threads, so - // cache one per device in thread-local storage. Pedantic math keeps the fp32 - // potential-energy surface exact (no TF32, no split-K reordering) for MD. + // cache one per device in thread-local storage. Pedantic math disables TF32, + // but does not guarantee identical reductions across different node counts. + // The optional shape-stable path fixes its full-matrix configuration. thread_local std::unordered_map handles; int device = 0; cudaGetDevice(&device); @@ -70,6 +72,10 @@ void gemm_nn(cudaStream_t stream, int n, int k, float beta = 0.f) { + if (deepmd_fitting::shape_stable_enabled()) { + deepmd_fitting::context(stream).run(stream, a, b, c, m, n, k, false, beta); + return; + } cublasSetStream(cublas_handle(), stream); const float alpha = 1.f; cublasSgemm(cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, b, n, @@ -85,6 +91,10 @@ void gemm_nt(cudaStream_t stream, int n, int k, float beta = 0.f) { + if (deepmd_fitting::shape_stable_enabled()) { + deepmd_fitting::context(stream).run(stream, a, b, c, m, n, k, true, beta); + return; + } cublasSetStream(cublas_handle(), stream); const float alpha = 1.f; cublasSgemm(cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, n, m, k, &alpha, b, k, diff --git a/source/op/pt/graph_fitting_gemm.h b/source/op/pt/graph_fitting_gemm.h new file mode 100644 index 0000000000..df9608e79c --- /dev/null +++ b/source/op/pt/graph_fitting_gemm.h @@ -0,0 +1,252 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace deepmd_fitting { + +inline void check_blas(cublasStatus_t status, const char* operation) { + TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, operation, + " failed with cuBLAS status ", static_cast(status)); +} + +// This is a numerical reproducibility policy, not a precision switch. Read it +// once per process so a live model cannot silently change arithmetic policy. +inline bool shape_stable_enabled() { + static const bool enabled = [] { + const char* value = std::getenv("DP_GRAPH_FITTING_GEMM_POLICY"); + const std::string policy = value ? value : "legacy"; + TORCH_CHECK(policy == "legacy" || policy == "shape_stable", + "DP_GRAPH_FITTING_GEMM_POLICY must be legacy or shape_stable"); + return policy == "shape_stable"; + }(); + return enabled; +} + +// cuBLASLt descriptors are host-side launch metadata. Keep them alive and +// update only the node-axis length; neither selecting an algorithm nor +// constructing descriptors belongs in the steady-state per-step path. +struct Layout { + cublasLtMatrixLayout_t value = nullptr; + Layout(uint64_t rows, uint64_t columns, int64_t leading) { + check_blas( + cublasLtMatrixLayoutCreate(&value, CUDA_R_32F, rows, columns, leading), + "cublasLtMatrixLayoutCreate"); + } + ~Layout() { + if (value) { + cublasLtMatrixLayoutDestroy(value); + } + } + Layout(const Layout&) = delete; + Layout& operator=(const Layout&) = delete; + void columns(uint64_t columns) { + check_blas( + cublasLtMatrixLayoutSetAttribute(value, CUBLASLT_MATRIX_LAYOUT_COLS, + &columns, sizeof(columns)), + "cublasLtMatrixLayoutSetAttribute"); + } +}; + +struct Operation { + cublasLtMatmulDesc_t value = nullptr; + Operation() { + check_blas(cublasLtMatmulDescCreate(&value, CUBLAS_COMPUTE_32F_PEDANTIC, + CUDA_R_32F), + "cublasLtMatmulDescCreate"); + } + ~Operation() { + if (value) { + cublasLtMatmulDescDestroy(value); + } + } + Operation(const Operation&) = delete; + Operation& operator=(const Operation&) = delete; +}; + +struct Preference { + cublasLtMatmulPreference_t value = nullptr; + Preference() { + check_blas(cublasLtMatmulPreferenceCreate(&value), + "cublasLtMatmulPreferenceCreate"); + } + ~Preference() { + if (value) { + cublasLtMatmulPreferenceDestroy(value); + } + } + Preference(const Preference&) = delete; + Preference& operator=(const Preference&) = delete; +}; + +// The fitting API supplies row-major node features. In column-major BLAS +// notation the variable node axis is N, while M and K depend only on weights. +// Pick a configuration using a fixed N=256, then submit the complete actual +// matrix. This does not tile windows or reduce the neural inference batch. +class Plan { + public: + static constexpr uint64_t canonical_nodes = 256; + // Select using the canonical problem's budget, independently of the larger + // scratch reservation needed by the same split-K algorithm on full batches. + static constexpr size_t heuristic_workspace = 32 * 1024 * 1024; + static constexpr size_t workspace_limit = 128 * 1024 * 1024; + + Plan(cublasLtHandle_t handle, int m, int k, bool transpose) + : weight_(transpose ? k : m, transpose ? m : k, transpose ? k : m), + features_(k, canonical_nodes, k), + output_(m, canonical_nodes, m) { + const cublasOperation_t trans = transpose ? CUBLAS_OP_T : CUBLAS_OP_N; + check_blas(cublasLtMatmulDescSetAttribute(operation_.value, + CUBLASLT_MATMUL_DESC_TRANSA, + &trans, sizeof(trans)), + "cublasLtMatmulDescSetAttribute"); + Preference preference; + const size_t bytes = heuristic_workspace; + check_blas(cublasLtMatmulPreferenceSetAttribute( + preference.value, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &bytes, sizeof(bytes)), + "cublasLtMatmulPreferenceSetAttribute workspace"); + // Slices need only satisfy float alignment. Never let allocator addresses + // choose a different reduction algorithm for otherwise identical rows. + const uint32_t alignment = alignof(float); + for (auto attribute : {CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES, + CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES, + CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES, + CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES}) { + check_blas( + cublasLtMatmulPreferenceSetAttribute(preference.value, attribute, + &alignment, sizeof(alignment)), + "cublasLtMatmulPreferenceSetAttribute alignment"); + } + cublasLtMatmulHeuristicResult_t result{}; + int count = 0; + check_blas( + cublasLtMatmulAlgoGetHeuristic( + handle, operation_.value, weight_.value, features_.value, + output_.value, output_.value, preference.value, 1, &result, &count), + "cublasLtMatmulAlgoGetHeuristic"); + TORCH_CHECK(count == 1 && result.state == CUBLAS_STATUS_SUCCESS, + "No shape-stable FP32 fitting GEMM configuration for M=", m, + ", K=", k, ", transpose=", transpose); + algorithm_ = result.algo; + } + + void run(cublasLtHandle_t handle, + cudaStream_t stream, + void* workspace, + const float* weight, + const float* features, + float* output, + int nodes, + float beta) { + if (nodes != last_nodes_) { + features_.columns(nodes); + output_.columns(nodes); + cublasLtMatmulHeuristicResult_t support{}; + check_blas(cublasLtMatmulAlgoCheck( + handle, operation_.value, weight_.value, features_.value, + output_.value, output_.value, &algorithm_, &support), + "cublasLtMatmulAlgoCheck"); + TORCH_CHECK(support.state == CUBLAS_STATUS_SUCCESS && + support.workspaceSize <= workspace_limit, + "Shape-stable fitting GEMM does not support this node count " + "within its 128 MiB workspace: ", + nodes, " (required bytes: ", support.workspaceSize, + ", status: ", static_cast(support.state), + "). Refusing a shape-dependent algorithm fallback."); + last_nodes_ = nodes; + } + const float alpha = 1.f; + check_blas(cublasLtMatmul(handle, operation_.value, &alpha, weight, + weight_.value, features, features_.value, &beta, + output, output_.value, output, output_.value, + &algorithm_, workspace, workspace_limit, stream), + "cublasLtMatmul shape-stable fitting"); + } + + private: + Operation operation_; + Layout weight_, features_, output_; + cublasLtMatmulAlgo_t algorithm_{}; + int last_nodes_ = -1; +}; + +// A separate context per host thread, CUDA device, and stream prevents races +// in both launch metadata and scratch memory. A stream serializes repeated +// uses of its own workspace, with no device-wide synchronization in the loop. +class Context { + public: + explicit Context(int device) : device_(device) { + workspace_ = torch::empty({static_cast(Plan::workspace_limit)}, + torch::TensorOptions() + .dtype(torch::kUInt8) + .device(torch::kCUDA, device)); + check_blas(cublasLtCreate(&handle_), "cublasLtCreate"); + } + ~Context() { + // CUDA may already be shutting down. Destructors must not throw; ordinary + // tensor destruction releases the allocator-owned scratch on its device. + int previous = -1; + if (cudaGetDevice(&previous) == cudaSuccess && + cudaSetDevice(device_) == cudaSuccess) { + plans_.clear(); + if (handle_) { + cublasLtDestroy(handle_); + } + cudaSetDevice(previous); + } + } + Context(const Context&) = delete; + Context& operator=(const Context&) = delete; + + void run(cudaStream_t stream, + const float* features, + const float* weight, + float* output, + int nodes, + int width, + int reduction, + bool transpose, + float beta) { + TORCH_CHECK(nodes >= 0 && width > 0 && reduction > 0, + "Invalid shape-stable fitting GEMM dimensions"); + if (nodes == 0) { + return; + } + const auto key = std::make_tuple(width, reduction, transpose); + auto& plan = plans_[key]; + if (!plan) { + plan = std::make_unique(handle_, width, reduction, transpose); + } + plan->run(handle_, stream, workspace_.data_ptr(), weight, features, output, + nodes, beta); + } + + private: + int device_; + cublasLtHandle_t handle_ = nullptr; + torch::Tensor workspace_; + std::map, std::unique_ptr> plans_; +}; + +inline Context& context(cudaStream_t stream) { + thread_local std::map, + std::unique_ptr> + contexts; + const int device = c10::cuda::current_device(); + auto& result = contexts[std::make_pair( + device, reinterpret_cast(stream))]; + if (!result) { + result = std::make_unique(device); + } + return *result; +} + +} // namespace deepmd_fitting diff --git a/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py b/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py new file mode 100644 index 0000000000..c6087c7bda --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""GPU regression for a fixed full-matrix fitting GEMM configuration. + +Run this focused file in a fresh process with +``DP_GRAPH_FITTING_GEMM_POLICY=shape_stable``. The policy is process-immutable. +The optional library path supports testing an isolated C++ build without +replacing an installed Python environment. +""" + +import importlib +import os +import subprocess +import sys +import unittest +from concurrent.futures import ( + ThreadPoolExecutor, +) +from pathlib import ( + Path, +) + +import torch + +if os.environ.get("DEEPMD_TEST_OP_LIBRARY"): + torch.ops.load_library(os.environ["DEEPMD_TEST_OP_LIBRARY"]) +else: + importlib.import_module("deepmd.pt.cxx_op") + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestShapeStablePolicyProcess(unittest.TestCase): + """Exercise the opt-in path in ordinary GPU CI without changing its policy. + + The operator reads its policy only once. Other tests may already have used + the legacy policy, so changing os.environ in this test process is unsafe. + A child process also preserves the caller's Slurm device visibility. + """ + + def test_shape_stable_in_fresh_process(self): + environment = dict(os.environ, DP_GRAPH_FITTING_GEMM_POLICY="shape_stable") + result = subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "TestShapeStableFitting", + "-v", + ], + env=environment, + capture_output=True, + text=True, + timeout=180, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertIn("Ran 4 tests", result.stderr) + self.assertNotIn("skipped", result.stderr) + + +@unittest.skipUnless( + torch.cuda.is_available() + and os.environ.get("DP_GRAPH_FITTING_GEMM_POLICY") == "shape_stable", + "requires CUDA and a fresh shape_stable fitting process", +) +class TestShapeStableFitting(unittest.TestCase): + """Check the entire feature matrix, residual backward, and stream ownership.""" + + def setUp(self): + torch.set_num_threads(1) + torch.backends.cuda.matmul.allow_tf32 = False + generator = torch.Generator(device="cuda").manual_seed(20260917) + self.x = torch.randn(1087, 144, generator=generator, device="cuda") * 0.2 + self.types = torch.arange(1087, device="cuda") % 6 + self.weights = [ + torch.randn(a, b, generator=generator, device="cuda") * 0.04 + for a, b in [(144, 192), (192, 192), (192, 192)] + ] + self.biases = [ + torch.randn(192, generator=generator, device="cuda") * 0.03 + for _ in range(3) + ] + self.head = torch.randn(192, generator=generator, device="cuda") * 0.04 + self.head_bias = torch.zeros(1, device="cuda") + self.atomic_bias = torch.arange(6, dtype=torch.float64, device="cuda") + + def evaluate(self, x, types, activation=1): + """Invoke the same fitting forward and reverse paths as canonical MD.""" + energy, saved = torch.ops.deepmd.graph_fitting( + x, + types, + self.weights, + self.biases, + [1, 1, 1], + self.head, + self.head_bias, + self.atomic_bias, + activation, + ) + gradient = torch.ops.deepmd.graph_fitting_backward( + torch.ones_like(energy), + saved, + self.weights, + self.biases, + [1, 1, 1], + self.head, + activation, + ) + return energy, gradient + + def test_batch_rows_and_backward_are_identical(self): + for activation in (0, 1): + single = self.evaluate(self.x[:268], self.types[:268], activation) + for repetitions in (1, 2, 4, 8, 16): + x = self.x.repeat(repetitions, 1) + types = self.types.repeat(repetitions) + full = self.evaluate(x, types, activation) + for expected, actual in zip(single, full, strict=True): + torch.testing.assert_close(actual[:268], expected, atol=0, rtol=0) + + def test_independent_streams_and_host_threads(self): + expected = self.evaluate(self.x, self.types) + torch.cuda.synchronize() + + def worker(_): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + outputs = self.evaluate(self.x, self.types) + stream.synchronize() + return outputs + + with ThreadPoolExecutor(max_workers=2) as pool: + outputs = list(pool.map(worker, range(4))) + for output in outputs: + for reference, value in zip(expected, output, strict=True): + torch.testing.assert_close(reference, value, atol=0, rtol=0) + + def test_minimum_aligned_feature_slice(self): + # A contiguous view can start only four bytes into its allocation. + # Pointer alignment must not change the selected reduction schedule. + backing = torch.empty(self.x.numel() + 1, device="cuda", dtype=self.x.dtype) + shifted = backing[1:].view_as(self.x) + shifted.copy_(self.x) + self.assertEqual(shifted.data_ptr() % 16, 4) + for expected, actual in zip( + self.evaluate(self.x, self.types), + self.evaluate(shifted, self.types), + strict=True, + ): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def test_independent_double_precision_reference(self): + # Independent autograd reference: do not obtain it from the candidate's + # own backward implementation or weaken tolerances to hide a mismatch. + x = self.x[:19].double().requires_grad_(True) + current = x + for weight, bias in zip(self.weights, self.biases, strict=True): + value = torch.nn.functional.silu(current @ weight.double() + bias.double()) + if current.shape[1] == value.shape[1]: + value = value + current + current = value + energy = current @ self.head.double() + self.atomic_bias[self.types[:19]] + (gradient,) = torch.autograd.grad(energy.sum(), x) + actual_energy, actual_gradient = self.evaluate(self.x[:19], self.types[:19]) + torch.testing.assert_close( + actual_energy.flatten(), energy, atol=1e-6, rtol=1e-6 + ) + torch.testing.assert_close( + actual_gradient.double(), gradient, atol=1e-6, rtol=1e-6 + ) + + +if __name__ == "__main__": + unittest.main() From 43eeea5c02abd7eefeeac1a4ea1d9a39b6385faf Mon Sep 17 00:00:00 2001 From: "A bot of @njzjz" Date: Fri, 2 Oct 2026 23:28:27 +0800 Subject: [PATCH 2/5] fix(pt): invalidate fitting layout cache before shape checks Prevent a rejected shape transition from leaving last_nodes_ inconsistent with the mutated cuBLASLt layouts. Any later call now restores and revalidates its requested shape. Agent: ChatGPT Model: GPT-5.6 Sol --- source/op/pt/graph_fitting_gemm.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/source/op/pt/graph_fitting_gemm.h b/source/op/pt/graph_fitting_gemm.h index df9608e79c..3a2559e375 100644 --- a/source/op/pt/graph_fitting_gemm.h +++ b/source/op/pt/graph_fitting_gemm.h @@ -147,6 +147,10 @@ class Plan { int nodes, float beta) { if (nodes != last_nodes_) { + // Invalidate before mutating either descriptor. If a layout update or + // support check throws, the next call must restore and revalidate its + // requested node count instead of trusting stale descriptor state. + last_nodes_ = -1; features_.columns(nodes); output_.columns(nodes); cublasLtMatmulHeuristicResult_t support{}; From b7361f05fd616385fe0bc6ecaf8f7af9469e045b Mon Sep 17 00:00:00 2001 From: "A bot of @njzjz" Date: Wed, 4 Nov 2026 02:17:16 +0800 Subject: [PATCH 3/5] test(pt): cover every fitting row and overlapping same-thread CUDA streams Agent: ChatGPT Model: GPT-6 --- .../test_graph_fitting_batch_invariant.py | 32 +++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py b/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py index c6087c7bda..d02babc683 100644 --- a/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py +++ b/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py @@ -51,7 +51,7 @@ def test_shape_stable_in_fresh_process(self): timeout=180, ) self.assertEqual(result.returncode, 0, result.stdout + result.stderr) - self.assertIn("Ran 4 tests", result.stderr) + self.assertIn("Ran 5 tests", result.stderr) self.assertNotIn("skipped", result.stderr) @@ -107,13 +107,22 @@ def evaluate(self, x, types, activation=1): def test_batch_rows_and_backward_are_identical(self): for activation in (0, 1): - single = self.evaluate(self.x[:268], self.types[:268], activation) + single = self.evaluate(self.x, self.types, activation) for repetitions in (1, 2, 4, 8, 16): x = self.x.repeat(repetitions, 1) types = self.types.repeat(repetitions) full = self.evaluate(x, types, activation) for expected, actual in zip(single, full, strict=True): - torch.testing.assert_close(actual[:268], expected, atol=0, rtol=0) + # Check every repeated block, including the final partial + # GEMM tile, not only the leading rows of the result. + torch.testing.assert_close( + actual.reshape(repetitions, *expected.shape), + expected.unsqueeze(0).expand( + repetitions, *expected.shape + ), + atol=0, + rtol=0, + ) def test_independent_streams_and_host_threads(self): expected = self.evaluate(self.x, self.types) @@ -132,6 +141,23 @@ def worker(_): for reference, value in zip(expected, output, strict=True): torch.testing.assert_close(reference, value, atol=0, rtol=0) + def test_concurrent_streams_on_one_host_thread(self): + # Launch both streams before synchronizing either. Unlike worker + # threads, this exercises the stream dimension of the thread-local + # GEMM context key and its private scratch workspace. + expected = self.evaluate(self.x, self.types) + torch.cuda.synchronize() + streams = [torch.cuda.Stream() for _ in range(2)] + outputs = [] + for stream in streams: + with torch.cuda.stream(stream): + outputs.append(self.evaluate(self.x, self.types)) + for stream in streams: + stream.synchronize() + for output in outputs: + for reference, value in zip(expected, output, strict=True): + torch.testing.assert_close(reference, value, atol=0, rtol=0) + def test_minimum_aligned_feature_slice(self): # A contiguous view can start only four bytes into its allocation. # Pointer alignment must not change the selected reduction schedule. From 3dd3702e30afc0d5baa02e80dd1389209cdaf6e0 Mon Sep 17 00:00:00 2001 From: "A bot of @njzjz" Date: Wed, 4 Nov 2026 02:17:40 +0800 Subject: [PATCH 4/5] docs(pt): clarify tiling and CUDA graph capture qualification Agent: ChatGPT Model: GPT-6 --- doc/env.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/doc/env.md b/doc/env.md index 4e9959f70b..7e65b29003 100644 --- a/doc/env.md +++ b/doc/env.md @@ -71,6 +71,16 @@ thread/device/stream. Unsupported shapes or workspace requirements fail explicitly rather than selecting a different reduction algorithm. Qualify energy and force parity for the desired model, shapes, GPU, and cuBLAS version; this setting is not a cross-hardware or cross-version reproducibility promise. + +The full-matrix statement applies to the fused operator's ordinary invocation. +An explicit positive `tile` argument to `graph_fitting_energy_gradient` +partitions the input and is outside the full-batch invariance guarantee. + +CUDA graph capture has not been qualified with `shape_stable`: each new +thread/device/stream context lazily creates cuBLASLt metadata and scratch, +which may not be safe inside capture. Initialize that context outside capture +on the same stream before capturing, and independently validate capture +compatibility for the intended workloads; otherwise use `legacy`. ::: :::{envvar} DP_INFER_BATCH_SIZE From f442d2431a3280d08dc0dcee0fb37c620b05db48 Mon Sep 17 00:00:00 2001 From: "A bot of @njzjz" Date: Wed, 4 Nov 2026 02:57:46 +0800 Subject: [PATCH 5/5] test(pt): strengthen fitting bias and thread coverage Use a nonzero head bias in the independent numerical reference and require two host workers to meet at a bounded barrier before concurrent launches. Agent: ChatGPT Model: GPT-6 Astra Pro --- .../test_graph_fitting_batch_invariant.py | 39 ++++++++++++++----- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py b/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py index d3dcf6bac4..7cef6ab8d7 100644 --- a/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py +++ b/source/tests/pt_expt/descriptor/test_graph_fitting_batch_invariant.py @@ -11,6 +11,9 @@ from concurrent.futures import ( ThreadPoolExecutor, ) +from threading import ( + Barrier, +) import torch @@ -40,7 +43,9 @@ def setUp(self): ] self.biases = [torch.randn(192, **options) * 0.03 for _ in range(3)] self.head = torch.randn(192, **options) * 0.04 - self.head_bias = torch.zeros(1, device=self.device, dtype=torch.float32) + self.head_bias = torch.full( + (1,), 0.125, device=self.device, dtype=torch.float32 + ) self.atomic_bias = torch.arange(6, dtype=torch.float64, device=self.device) def evaluate(self, x, types, activation=1): @@ -116,15 +121,27 @@ def test_independent_streams_and_host_threads(self): for evaluate in (self.evaluate, self.evaluate_fused): - def worker(_): - stream = torch.cuda.Stream(device=self.device) - with torch.cuda.stream(stream): - outputs = evaluate(self.x, self.types) - stream.synchronize() - return outputs + def worker(barrier): + try: + stream = torch.cuda.Stream(device=self.device) + with torch.cuda.stream(stream): + barrier.wait() + outputs = evaluate(self.x, self.types) + stream.synchronize() + return outputs + except BaseException: + # Release the other worker if setup fails before the wait. + barrier.abort() + raise + outputs = [] with ThreadPoolExecutor(max_workers=2) as pool: - outputs = list(pool.map(worker, range(4))) + for _ in range(2): + # Each wave must occupy two host threads, even if one + # submission would otherwise finish before the next starts. + barrier = Barrier(2, timeout=30) + futures = [pool.submit(worker, barrier) for _ in range(2)] + outputs.extend(future.result() for future in futures) for output in outputs: for reference, value in zip(expected, output, strict=True): torch.testing.assert_close(reference, value, atol=0, rtol=0) @@ -186,7 +203,11 @@ def test_independent_double_precision_reference(self): if current.shape[1] == value.shape[1]: value = value + current current = value - energy = current @ self.head.double() + self.atomic_bias[self.types[:19]] + energy = ( + current @ self.head.double() + + self.head_bias.double() + + self.atomic_bias[self.types[:19]] + ) (gradient,) = torch.autograd.grad(energy.sum(), x) for evaluate in (self.evaluate, self.evaluate_fused): actual_energy, actual_gradient = evaluate(self.x[:19], self.types[:19])