From 1351a4fb72dff668755ac4b3f9f7fc510bf8835e Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 7 Oct 2026 11:26:39 +0800 Subject: [PATCH 1/5] fix(triton): validate grid-pair launches before graph capture Cold compiled FP32 force training can cache a grid-pair launch that exceeds per-block shared memory. wrap_triton records a call during fake-tensor tracing without loading the kernel, so the eager launch loop mistakes a successful trace for a resource-validated configuration. The actual OutOfResources exception arrives later in Inductor's backward compilation, outside the retry loop. Reducing the training batch cannot fix the tile's fixed resource footprint. For preset Pro on RTX PRO 6000 Blackwell, the second-order kernel initially requests 118784 bytes against a 101376-byte device limit. Eager execution finds a smaller grid tile requiring 79872 bytes, but cold graph capture cannot take that retry. Prewarming eagerly can consequently hide the bug. Resolve launch feasibility before recording or executing a kernel: - Separate candidate selection from the runtime launch and use the same specialization parameters for validation and execution. - Compile candidates with Triton warmup using operand dtypes instead of tensor allocations, then load the compiled handles to enforce the target device's resource limits without executing the kernel. - Represent runtime pair counts and physical strides by non-unit scalar placeholders during validation. Keep their actual values in the launch ABI so dynamic batches and degree-major views retain one compiled graph. - Cache validated choices by kernel, device, static geometry, complete operand signature and differentiation-order accumulator count. Cache an unsuccessful search as well, avoiding repeated compilation for layouts served by the reference composition. - Retain exact-device tuned configurations as preferred candidates, while subjecting them to the same validation and deduplicating fallback choices. Forward, first backward and force-loss second backward share the selector. - Preserve compiler-visible Triton operators and the reference einsum path when no candidate fits. Only resource-limit errors reject a candidate; other compilation failures remain visible. Add cold compiled force-loss regressions for the two production Pro grid layouts in FP32 and BF16 autocast. Exercise batch-dependent physical strides with changing node counts, compare outputs and second-order gradients with eager execution, and verify that only one graph is compiled. Clear launch selection state and disable the FX graph cache so prior eager tests cannot mask an invalid first choice. Cover reference fallback through compilation when no fused launch is available. Validation on RTX PRO 6000 Blackwell Server Edition with PyTorch 2.13.0+cu130 and Triton 3.7.1: - All 20 grid-pair tests pass, including existing FP32/BF16 forward and first/second-gradient checks against the float64 reference, dynamic layout tests, and the five cold-compile/fallback regressions. - A cold FP32 second-order reproducer with one graph node passes without an eager warmup. The validated and final compiled backward kernels agree on their resource footprint; the Pro second-order tile uses 79872 bytes. - A preset Pro fine-tuning smoke run completes distributed precompilation and two optimizer steps with one DDP rank, ZeRO Stage 1, a 300-atom batch budget, AMP/TF32 disabled, and both CUDA/Triton training gates enabled. All 278 floating-point model tensors in the saved checkpoint are finite. - An independent review compares 120 actual layout signatures across five geometries, two dtypes, three differentiation orders and four storage layouts. Their shared-memory usage matches the validation signatures. - Ruff and whitespace checks pass. Other GPU architectures and the project-pinned PyTorch 2.10 environment were not exercised. --- .../pt_expt/kernels/triton/sezm/grid_pair.py | 193 +++++++++++------- .../pt_expt/kernels/test_grid_pair_train.py | 110 ++++++++++ 2 files changed, 226 insertions(+), 77 deletions(-) diff --git a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py index a11f5d49d6..37a8c263dc 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py +++ b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py @@ -44,7 +44,11 @@ annotations, ) +from functools import ( + cache, +) from typing import ( + TYPE_CHECKING, Any, ) @@ -60,6 +64,13 @@ ) _Shape = tuple[int | torch.SymInt, ...] +_ShapeKey = tuple[str, int, int, int, int, bool, int, torch.dtype] +_LaunchConfig = tuple[int, int, int] + +if TYPE_CHECKING: + from triton.runtime.jit import ( + JITFunction, + ) __all__ = [ "GRID_PAIR_TRITON_AVAILABLE", @@ -879,8 +890,6 @@ def _unpack(value: Tensor, shape: _Shape, n_frames: int) -> Tensor: ) -_LAUNCH_CACHE: dict[tuple, tuple[int, int, int]] = {} - # Exact production-shape winners on RTX PRO 6000 Blackwell. The grid-pair # kernels hold one, two or three ``(P_PAD, C_BLK)`` accumulator tiles, so the # best channel width and grid tile change independently across differentiation @@ -1002,54 +1011,119 @@ def _unpack(value: Tensor, shape: _Shape, n_frames: int) -> Tensor: def _built_in_launch_config( - device_name: str, shape_key: tuple -) -> tuple[int, int, int] | None: + device_name: str, shape_key: _ShapeKey +) -> _LaunchConfig | None: """Return the built-in launch for one exact grid-pair shape.""" if device_name.startswith("NVIDIA RTX PRO 6000 Blackwell"): return _BLACKWELL_LAUNCH_CONFIGS.get(shape_key) return None +def _slot_partition(p_dim: int) -> tuple[int, int]: + """Split the coefficient axis into power-of-two contraction tiles.""" + # The slot axis is covered by the largest power of two below the count + # plus an optional low segment for the remainder, so 147 (degree six) + # pads to 128 + 32 and 75 (degree four) to 64 + 16 instead of the next + # power of two. The high segment keeps the tensor-core minimum of 16. + p_hi = max(16, 1 << (p_dim.bit_length() - 1)) + p_lo = max(16, _next_pow2(p_dim - p_hi)) if p_dim > p_hi else 0 + return p_hi, p_lo + + +def _launch_kwargs( + shape_key: _ShapeKey, config: _LaunchConfig +) -> dict[str, int | bool]: + """Use identical specialization parameters for validation and execution.""" + _, p_dim, c_per, n_focus, n_frames, packed, n_grid, _ = shape_key + p_hi, p_lo = _slot_partition(p_dim) + c_blk, block_g, stages = config + return { + "n_grid": n_grid, + "PACKED": packed, + "N_FOCUS": n_focus, + "N_FRAMES": n_frames, + "P_DIM": p_dim, + "P_HI": p_hi, + "P_LO": p_lo, + "C_ALL": c_per, + "C_BLK": c_blk, + "ALLOW_TF32": False, + "BLOCK_G": block_g, + "num_warps": 8 if c_blk >= 64 else 4, + "num_stages": stages, + } + + +@cache +def _select_launch_config( + kernel: JITFunction, + device: torch.device, + shape_key: _ShapeKey, + signature: tuple[torch.dtype | int, ...], + n_acc: int, +) -> _LaunchConfig | None: + """Compile and load candidates without allocating or executing their inputs. + + Tile resources depend on the static coefficient geometry and operand + dtypes. The leading pair count and physical strides only address rows; + their non-unit placeholders avoid specializing a dynamic batch size. + Triton's ``warmup`` accepts dtypes in place of tensor operands and + compiles the real device binary without dereferencing a pointer. Loading + its handles checks the device's resource limits before a configuration + enters the cache or a traced graph. Fake-tensor tracing therefore makes + the same choice as eager execution. + + The channel cap keeps the ``n_acc`` fp32 accumulator tiles register + resident. Exact shared-memory requirements are checked by Triton, whose + staging and transpose storage depend on the compiler and target device. + A swept configuration has priority but must pass the same validation. + ``None`` records that the reference composition must serve this layout. + """ + _, p_dim, c_per, *_ = shape_key + p_hi, p_lo = _slot_partition(p_dim) + p_eff = p_hi + p_lo + c_top = min(64, _next_pow2(c_per), max(16, _next_pow2(4096 // (n_acc * p_eff)))) + preferred = _built_in_launch_config(torch.cuda.get_device_name(device), shape_key) + candidates = [ + (c_blk, block_g, stages) + for c_blk in (c_top, 32, 16) + if c_blk <= c_top + for block_g, stages in ((64, 2), (32, 2), (32, 1), (16, 1)) + ] + if preferred is not None: + candidates.insert(0, preferred) + with torch.cuda.device(device): + for config in dict.fromkeys(candidates): + try: + compiled = kernel.warmup( + *signature, + 2, + grid=(1, 1), + **_launch_kwargs(shape_key, config), + ) + compiled._init_handles() + except triton.runtime.errors.OutOfResources: + continue + return config + return None + + def _launch( - kernel, + kernel: JITFunction, value: Tensor, n_grid: int, n_frames: int, packed: bool, - args: tuple, + args: tuple[Tensor | int | torch.SymInt, ...], n_acc: int, ) -> None: - """Launch with the largest tile the register and shared budgets admit. - - The channel block is capped so the ``n_acc`` fp32 accumulator tiles - ``(P_PAD, C_BLK)`` stay register-resident (a spilled accumulator is - read-modify-written through local memory once per grid block, which - dominated the three-accumulator second-order kernel before the cap). - The exact shared footprint additionally depends on Triton's internal - staging (dot operand buffers, transpose scratch), so candidates are - tried from the most to the least aggressive and the first that compiles - is cached per device and exact operator shape. Swept built-in launches - take precedence where available and fall back to the same compile search - if a later Triton version rejects one. - """ - n_batch = value.shape[0] - # The leading pair count follows the workload. The trailing coefficient - # layout and projector size define the Triton tile geometry and remain - # kernel-specialized. + """Record or execute a launch whose tile has passed device validation.""" n_grid = guard_int(n_grid) n_frames = guard_int(n_frames) coeff_dim, n_focus, packed_channels = (guard_int(size) for size in value.shape[1:]) - n_pair = n_batch * n_focus + n_pair = value.shape[0] * n_focus p_dim = coeff_dim * n_frames c_per = packed_channels // n_frames - # The slot axis is covered by the largest power of two below the count - # plus an optional low segment for the remainder, so 147 (degree six) - # pads to 128 + 32 and 75 (degree four) to 64 + 16 instead of the next - # power of two. The high segment keeps the tensor-core minimum of 16. - p_hi = max(16, 1 << (p_dim.bit_length() - 1)) - p_lo = max(16, _next_pow2(p_dim - p_hi)) if p_dim > p_hi else 0 - p_eff = p_hi + p_lo - c_top = min(64, _next_pow2(c_per), max(16, _next_pow2(4096 // (n_acc * p_eff)))) shape_key = ( kernel.fn.__name__, p_dim, @@ -1060,48 +1134,13 @@ def _launch( n_grid, value.dtype, ) - device_name = torch.cuda.get_device_name(value.device) - key = (device_name, *shape_key) - generic_candidates = [ - (c_blk, block_g, stages) - for c_blk in (c_top, 32, 16) - if c_blk <= c_top - for block_g, stages in ((64, 2), (32, 2), (32, 1), (16, 1)) - ] - if key in _LAUNCH_CACHE: - candidates = [_LAUNCH_CACHE[key]] - else: - built_in = _built_in_launch_config(device_name, shape_key) - candidates = ( - generic_candidates - if built_in is None - else [built_in, *(cfg for cfg in generic_candidates if cfg != built_in)] - ) - for c_blk, block_g, stages in candidates: - grid = (n_pair, (c_per + c_blk - 1) // c_blk) - try: - wrap_triton(kernel)[grid]( - *args, - n_pair, - n_grid=n_grid, - PACKED=packed, - N_FOCUS=n_focus, - N_FRAMES=n_frames, - P_DIM=p_dim, - P_HI=p_hi, - P_LO=p_lo, - C_ALL=c_per, - C_BLK=c_blk, - ALLOW_TF32=False, - BLOCK_G=block_g, - num_warps=8 if c_blk >= 64 else 4, - num_stages=stages, - ) - except triton.runtime.errors.OutOfResources: - continue - _LAUNCH_CACHE[key] = (c_blk, block_g, stages) - return - raise _NoViableConfig(p_dim, c_per) + signature = tuple(arg.dtype if isinstance(arg, Tensor) else 2 for arg in args) + config = _select_launch_config(kernel, value.device, shape_key, signature, n_acc) + if config is None: + raise _NoViableConfig(p_dim, c_per) + c_blk = config[0] + grid = (n_pair, (c_per + c_blk - 1) // c_blk) + wrap_triton(kernel)[grid](*args, n_pair, **_launch_kwargs(shape_key, config)) def _strides(*values: Tensor) -> _Shape: @@ -1156,9 +1195,9 @@ def _eager_packed( ) -> tuple[Tensor, ...]: """Eager einsum fallback on the packed layout. - Serves the shapes whose padded tiles exceed the shared-memory budget -- - the wide-slot fp32 regime of the exact-precision harnesses; the - production bf16 shapes never reach it. + Serves layouts for which no candidate tile fits the target device's + resources. The composition remains visible to the compiler when resource + validation rejects every fused configuration during tracing. """ tg = to_grid fg = from_grid diff --git a/source/tests/pt_expt/kernels/test_grid_pair_train.py b/source/tests/pt_expt/kernels/test_grid_pair_train.py index 5d8d6b9223..795bb6cb9f 100644 --- a/source/tests/pt_expt/kernels/test_grid_pair_train.py +++ b/source/tests/pt_expt/kernels/test_grid_pair_train.py @@ -26,6 +26,7 @@ make_fx, ) +from deepmd.pt_expt.kernels.triton.sezm import grid_pair as grid_pair_module from deepmd.pt_expt.kernels.triton.sezm.grid_pair import ( GRID_PAIR_TRITON_AVAILABLE, _built_in_launch_config, @@ -391,3 +392,112 @@ def evaluate( torch.testing.assert_close(got, want) assert compile_counter.frame_count == 1 + + +@pytest.mark.parametrize( + ("n_focus", "channels", "amp", "force_reference"), + [ + (2, 64, False, False), + (1, 192, False, False), + (2, 64, True, False), + (1, 192, True, False), + (2, 16, False, True), + ], +) +def test_cold_compiled_force_backward( + n_focus: int, + channels: int, + amp: bool, + force_reference: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Resolve device resources before tracing a production grid's second order. + + The first execution is compiled, so an eager pass cannot populate the + launch cache first. The two production layouts cover multi-focus messages + and the single-focus feed-forward grid. A device with no viable fused + configuration must retain the reference composition through compilation. + """ + torch._dynamo.reset() + grid_pair_module._select_launch_config.cache_clear() + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + if force_reference: + + def reject_configs(*_args: object, **_kwargs: object) -> None: + return None + + monkeypatch.setattr(grid_pair_module, "_select_launch_config", reject_configs) + + n_frames, coeff_dim, n_grid = 3, 36, 344 + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(40529) + projectors = tuple( + torch.randn( + n_grid, + coeff_dim * n_frames, + device=device, + dtype=torch.float32, + generator=generator, + ) + / n_grid**0.5 + for _ in range(2) + ) + + def make_inputs(n_node: int) -> tuple[torch.Tensor, ...]: + # Degree-major storage gives the coefficient axis a batch-dependent + # stride; the logical entry remains (N, D, F, K * C). + shape = (coeff_dim, n_node, n_focus, n_frames * channels) + return tuple( + torch.randn( + shape, + device=device, + dtype=torch.float32, + generator=generator, + requires_grad=True, + ) + for _ in range(2) + ) + + def force_forward( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): + out = grid_pair_train( + left.permute(1, 0, 2, 3), + right.permute(1, 0, 2, 3), + to_grid, + from_grid, + n_frames, + ) + gradients = torch.autograd.grad(out.sum(), (left, right), create_graph=True) + return out, *gradients + + def force_loss_gradient( + outputs: tuple[torch.Tensor, ...], + inputs: tuple[torch.Tensor, ...], + ) -> tuple[torch.Tensor, ...]: + loss = sum(value.float().square().mean() for value in outputs) + return torch.autograd.grad(loss, inputs) + + traced_inputs = make_inputs(7) + graph = make_fx(force_forward, tracing_mode="symbolic")(*traced_inputs, *projectors) + counter = CompileCounterWithBackend("inductor") + with torch._inductor.config.patch(fx_graph_cache=False): + compiled = torch.compile(graph, backend=counter, dynamic=True, fullgraph=True) + for n_node in (7, 11): + inputs = traced_inputs if n_node == 7 else make_inputs(n_node) + actual = compiled(*inputs, *projectors) + actual_gradients = force_loss_gradient(actual, inputs) + expected = force_forward(*inputs, *projectors) + expected_gradients = force_loss_gradient(expected, inputs) + for got, want in zip( + (*actual, *actual_gradients), + (*expected, *expected_gradients), + strict=True, + ): + torch.testing.assert_close(got, want) + assert counter.frame_count == 1 From e85e510d4dbce82eee973bccd99c716812293bce Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 23 Oct 2026 15:38:52 +0800 Subject: [PATCH 2/5] test(triton): verify grid-pair rejection and cold fused launches --- .../pt_expt/kernels/test_grid_pair_train.py | 152 ++++++++++++++++-- 1 file changed, 142 insertions(+), 10 deletions(-) diff --git a/source/tests/pt_expt/kernels/test_grid_pair_train.py b/source/tests/pt_expt/kernels/test_grid_pair_train.py index 795bb6cb9f..f25cf266d7 100644 --- a/source/tests/pt_expt/kernels/test_grid_pair_train.py +++ b/source/tests/pt_expt/kernels/test_grid_pair_train.py @@ -17,6 +17,16 @@ annotations, ) +from contextlib import ( + nullcontext, +) +from types import ( + SimpleNamespace, +) +from typing import ( + TYPE_CHECKING, +) + import pytest import torch from torch._dynamo.testing import ( @@ -40,10 +50,17 @@ median_deviations, ) +if TYPE_CHECKING: + from collections.abc import ( + Iterator, + ) + pytestmark = [ - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), pytest.mark.skipif(not GRID_PAIR_TRITON_AVAILABLE, reason="Triton is unavailable"), ] +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA is required" +) # ``(lmax, n_frames, n_focus, channels, n_grid)`` spanning the deployed grid # shapes. The slot count ``(lmax + 1)^2 * n_frames`` drives the operator's @@ -63,6 +80,93 @@ DRAW_SEEDS = (11, 2027, 40529) +@pytest.fixture +def isolated_launch_selection(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + """Keep selector tests independent of CUDA hardware and prior cache entries.""" + monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext()) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda device: "test device") + grid_pair_module._select_launch_config.cache_clear() + yield + grid_pair_module._select_launch_config.cache_clear() + + +@pytest.mark.usefixtures("isolated_launch_selection") +@pytest.mark.parametrize("failure_stage", ["warmup", "handles"]) +def test_launch_selection_retries_and_caches( + failure_stage: str, +) -> None: + """Resource rejection tries the next candidate and caches its admission.""" + attempts = [] + + class Kernel: + def warmup(self, *args: object, **kwargs: object) -> SimpleNamespace: + attempts.append((kwargs["C_BLK"], kwargs["BLOCK_G"], kwargs["num_stages"])) + + def init_handles() -> None: + if len(attempts) <= 2: + raise grid_pair_module.triton.runtime.errors.OutOfResources( + 2, 1, "test" + ) + + if failure_stage == "warmup": + init_handles() + return SimpleNamespace(_init_handles=lambda: None) + return SimpleNamespace(_init_handles=init_handles) + + kernel = Kernel() + shape = ("test_kernel", 147, 64, 2, 3, False, 587, torch.float32) + select = grid_pair_module._select_launch_config + result = select(kernel, torch.device("cuda:0"), shape, (torch.float32, 2), 1) + assert result == (32, 32, 1) + assert attempts == [(32, 64, 2), (32, 32, 2), (32, 32, 1)] + assert ( + select(kernel, torch.device("cuda:0"), shape, (torch.float32, 2), 1) == result + ) + assert len(attempts) == 3 + + +@pytest.mark.usefixtures("isolated_launch_selection") +def test_launch_rejection_is_cached_and_reported() -> None: + """An exhausted selector prevents execution and reuses its cached miss.""" + attempts = [] + + class Kernel: + fn = SimpleNamespace(__name__="test_kernel") + + def warmup(self, *args: object, **kwargs: object) -> None: + attempts.append((kwargs["C_BLK"], kwargs["BLOCK_G"], kwargs["num_stages"])) + raise grid_pair_module.triton.runtime.errors.OutOfResources(2, 1, "test") + + kernel = Kernel() + select = grid_pair_module._select_launch_config + value = torch.empty(2, 49, 2, 192, device="cpu") + for _ in range(2): + with pytest.raises(grid_pair_module._NoViableConfig, match="P=147, C=64"): + grid_pair_module._launch(kernel, value, 587, 3, False, (value,), 1) + assert len(attempts) == len(set(attempts)) == 8 + assert select.cache_info().hits == 1 + + +@pytest.mark.usefixtures("isolated_launch_selection") +def test_launch_selection_preserves_compiler_errors() -> None: + """A compiler defect must neither select a fallback nor enter the cache.""" + attempts = [] + + class Kernel: + def warmup(self, *args: object, **kwargs: object) -> None: + attempts.append(kwargs) + raise RuntimeError("compiler defect") + + kernel = Kernel() + select = grid_pair_module._select_launch_config + shape = ("test_kernel", 147, 64, 2, 3, False, 587, torch.float32) + for _ in range(2): + with pytest.raises(RuntimeError, match="compiler defect"): + select(kernel, torch.device("cuda:0"), shape, (torch.float32, 2), 1) + assert len(attempts) == 2 + assert select.cache_info().currsize == 0 + + def test_blackwell_launch_table_uses_exact_grid_shape() -> None: """Keep production launch pins scoped to the swept device and grid.""" shape_key = ( @@ -251,6 +355,7 @@ def _compare( @pytest.mark.parametrize( ("lmax", "n_frames", "n_focus", "channels", "n_grid"), GRID_SHAPES ) +@requires_cuda def test_float32_matches_eager_conditioning( lmax: int, n_frames: int, n_focus: int, channels: int, n_grid: int ) -> None: @@ -261,6 +366,7 @@ def test_float32_matches_eager_conditioning( @pytest.mark.parametrize( ("lmax", "n_frames", "n_focus", "channels", "n_grid"), GRID_SHAPES ) +@requires_cuda def test_autocast_bfloat16_matches_eager_conditioning( lmax: int, n_frames: int, n_focus: int, channels: int, n_grid: int ) -> None: @@ -268,6 +374,7 @@ def test_autocast_bfloat16_matches_eager_conditioning( _compare((lmax, n_frames, n_focus, channels, n_grid), amp=True) +@requires_cuda def test_noncontiguous_operands_match_eager_conditioning() -> None: """Cover the channel-slice strides supplied by production grid nets.""" _compare(GRID_SHAPES[1], amp=True, strided=True) @@ -277,6 +384,7 @@ def test_noncontiguous_operands_match_eager_conditioning() -> None: ("n_focus", "degree_major"), [(1, False), (2, True)], ) +@requires_cuda def test_symbolic_graph_reuses_layout_across_batch_shapes( n_focus: int, degree_major: bool ) -> None: @@ -355,6 +463,7 @@ def evaluate( torch.testing.assert_close(got, want) +@requires_cuda def test_aotautograd_reuses_multifocus_stride_across_batch_shapes() -> None: """Keep a producer's batch-dependent coefficient stride out of guards.""" n_frames, coeff_dim, n_focus, channels, n_grid = 3, 4, 2, 16, 32 @@ -397,13 +506,21 @@ def evaluate( @pytest.mark.parametrize( ("n_focus", "channels", "amp", "force_reference"), [ - (2, 64, False, False), - (1, 192, False, False), - (2, 64, True, False), - (1, 192, True, False), + (2, 96, False, False), + (1, 256, False, False), + (2, 96, True, False), + (1, 256, True, False), (2, 16, False, True), ], + ids=[ + "multifocus-fp32", + "singlefocus-fp32", + "multifocus-amp", + "singlefocus-amp", + "reference", + ], ) +@requires_cuda def test_cold_compiled_force_backward( n_focus: int, channels: int, @@ -419,17 +536,31 @@ def test_cold_compiled_force_backward( configuration must retain the reference composition through compilation. """ torch._dynamo.reset() - grid_pair_module._select_launch_config.cache_clear() monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + fallback_ops = set() + reference = grid_pair_module._eager_packed + + def record_fallback(op: str, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: + fallback_ops.add(op) + return reference(op, *args) + + monkeypatch.setattr(grid_pair_module, "_eager_packed", record_fallback) if force_reference: - def reject_configs(*_args: object, **_kwargs: object) -> None: - return None + def reject_candidate(*args: object, **kwargs: object) -> None: + raise grid_pair_module.triton.runtime.errors.OutOfResources(2, 1, "test") - monkeypatch.setattr(grid_pair_module, "_select_launch_config", reject_configs) + for kernel in ( + grid_pair_module._grid_pair_fwd_kernel, + grid_pair_module._grid_pair_bwd_kernel, + grid_pair_module._grid_pair_bwd2_kernel, + ): + monkeypatch.setattr(kernel, "warmup", reject_candidate) - n_frames, coeff_dim, n_grid = 3, 36, 344 + # Degree six reproduces the P=147 resource pressure. The distinct grid + # count prevents eager production-shape tests from warming this launch key. + n_frames, coeff_dim, n_grid = 3, 49, 587 device = torch.device("cuda") generator = torch.Generator(device=device).manual_seed(40529) projectors = tuple( @@ -501,3 +632,4 @@ def force_loss_gradient( ): torch.testing.assert_close(got, want) assert counter.frame_count == 1 + assert fallback_ops == ({"fwd", "bwd", "bwd2"} if force_reference else set()) From 68c57e07884b1500aced353c525f7fc467b45376 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 23 Oct 2026 16:40:01 +0800 Subject: [PATCH 3/5] test(triton): distinguish device admission from reference fallback --- .../pt_expt/kernels/test_grid_pair_train.py | 48 ++++++++++++++----- 1 file changed, 36 insertions(+), 12 deletions(-) diff --git a/source/tests/pt_expt/kernels/test_grid_pair_train.py b/source/tests/pt_expt/kernels/test_grid_pair_train.py index f25cf266d7..d1190374b2 100644 --- a/source/tests/pt_expt/kernels/test_grid_pair_train.py +++ b/source/tests/pt_expt/kernels/test_grid_pair_train.py @@ -52,9 +52,14 @@ if TYPE_CHECKING: from collections.abc import ( + Callable, Iterator, ) + from triton.compiler.compiler import ( + CompiledKernel, + ) + pytestmark = [ pytest.mark.skipif(not GRID_PAIR_TRITON_AVAILABLE, reason="Triton is unavailable"), ] @@ -539,6 +544,8 @@ def test_cold_compiled_force_backward( monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) fallback_ops = set() + probed_ops = set() + admitted_ops = set() reference = grid_pair_module._eager_packed def record_fallback(op: str, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: @@ -546,17 +553,29 @@ def record_fallback(op: str, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: return reference(op, *args) monkeypatch.setattr(grid_pair_module, "_eager_packed", record_fallback) - if force_reference: - - def reject_candidate(*args: object, **kwargs: object) -> None: - raise grid_pair_module.triton.runtime.errors.OutOfResources(2, 1, "test") - - for kernel in ( - grid_pair_module._grid_pair_fwd_kernel, - grid_pair_module._grid_pair_bwd_kernel, - grid_pair_module._grid_pair_bwd2_kernel, - ): - monkeypatch.setattr(kernel, "warmup", reject_candidate) + for op, kernel in ( + ("fwd", grid_pair_module._grid_pair_fwd_kernel), + ("bwd", grid_pair_module._grid_pair_bwd_kernel), + ("bwd2", grid_pair_module._grid_pair_bwd2_kernel), + ): + + def validate_candidate( + *args: object, + _op: str = op, + _warmup: Callable[..., CompiledKernel] = kernel.warmup, + **kwargs: object, + ) -> CompiledKernel: + probed_ops.add(_op) + if force_reference: + raise grid_pair_module.triton.runtime.errors.OutOfResources( + 2, 1, "test" + ) + candidate = _warmup(*args, **kwargs) + candidate._init_handles() + admitted_ops.add(_op) + return candidate + + monkeypatch.setattr(kernel, "warmup", validate_candidate) # Degree six reproduces the P=147 resource pressure. The distinct grid # count prevents eager production-shape tests from warming this launch key. @@ -623,6 +642,12 @@ def force_loss_gradient( inputs = traced_inputs if n_node == 7 else make_inputs(n_node) actual = compiled(*inputs, *projectors) actual_gradients = force_loss_gradient(actual, inputs) + # Observe compiled selection before eager evaluation can add events. + # A cached/unconditional rejection without validation cannot pass. + assert probed_ops == {"fwd", "bwd", "bwd2"} + assert fallback_ops == probed_ops - admitted_ops + if force_reference: + assert not admitted_ops expected = force_forward(*inputs, *projectors) expected_gradients = force_loss_gradient(expected, inputs) for got, want in zip( @@ -632,4 +657,3 @@ def force_loss_gradient( ): torch.testing.assert_close(got, want) assert counter.frame_count == 1 - assert fallback_ops == ({"fwd", "bwd", "bwd2"} if force_reference else set()) From 6b7707356ffbdde01923fd5581bdb8fb611b12a7 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 26 Oct 2026 10:39:21 +0800 Subject: [PATCH 4/5] test(triton): require admission when a viable tile is available --- .../pt_expt/kernels/triton/sezm/grid_pair.py | 6 +- .../pt_expt/kernels/test_grid_pair_train.py | 69 +++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py index 37a8c263dc..d67b27eb6c 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py +++ b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py @@ -1033,7 +1033,7 @@ def _slot_partition(p_dim: int) -> tuple[int, int]: def _launch_kwargs( shape_key: _ShapeKey, config: _LaunchConfig ) -> dict[str, int | bool]: - """Use identical specialization parameters for validation and execution.""" + """Share static geometry and tile parameters between validation and execution.""" _, p_dim, c_per, n_focus, n_frames, packed, n_grid, _ = shape_key p_hi, p_lo = _slot_partition(p_dim) c_blk, block_g, stages = config @@ -1073,6 +1073,10 @@ def _select_launch_config( enters the cache or a traced graph. Fake-tensor tracing therefore makes the same choice as eager execution. + The probe admits a tile configuration, not a byte-identical runtime binary. + Scalar and pointer-alignment specialization of the actual launch can still + produce a separate compiled artifact. + The channel cap keeps the ``n_acc`` fp32 accumulator tiles register resident. Exact shared-memory requirements are checked by Triton, whose staging and transpose storage depend on the compiler and target device. diff --git a/source/tests/pt_expt/kernels/test_grid_pair_train.py b/source/tests/pt_expt/kernels/test_grid_pair_train.py index d1190374b2..b90102ffc9 100644 --- a/source/tests/pt_expt/kernels/test_grid_pair_train.py +++ b/source/tests/pt_expt/kernels/test_grid_pair_train.py @@ -130,6 +130,48 @@ def init_handles() -> None: assert len(attempts) == 3 +@pytest.mark.usefixtures("isolated_launch_selection") +def test_admitted_launch_preserves_runtime_arguments_and_reuses_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An admitted tile executes with each input's actual layout and pair count.""" + probes = [] + loaded = [] + launches = [] + + class Kernel: + fn = SimpleNamespace(__name__="test_kernel") + + def warmup(self, *args: object, **kwargs: object) -> SimpleNamespace: + probes.append((args, kwargs)) + return SimpleNamespace(_init_handles=lambda: loaded.append(True)) + + class Launcher: + def __getitem__(self, grid: tuple[int, int]) -> Callable[..., None]: + def execute(*args: object, **kwargs: object) -> None: + launches.append((grid, args, kwargs)) + + return execute + + kernel = Kernel() + monkeypatch.setattr(grid_pair_module, "wrap_triton", lambda candidate: Launcher()) + for n_node in (2, 3): + value = torch.empty(49, n_node, 2, 192, device="cpu").permute(1, 0, 2, 3) + strides = value.stride() + grid_pair_module._launch(kernel, value, 587, 3, False, (value, *strides), 1) + grid, args, kwargs = launches[-1] + assert grid == (2 * n_node, 2) + assert args[0] is value + assert args[1:-1] == strides + assert args[-1] == 2 * n_node + assert kwargs == { + key: val for key, val in probes[0][1].items() if key != "grid" + } + assert len(probes) == len(loaded) == 1 + assert len(launches) == 2 + assert grid_pair_module._select_launch_config.cache_info().hits == 1 + + @pytest.mark.usefixtures("isolated_launch_selection") def test_launch_rejection_is_cached_and_reported() -> None: """An exhausted selector prevents execution and reuses its cached miss.""" @@ -546,6 +588,7 @@ def test_cold_compiled_force_backward( fallback_ops = set() probed_ops = set() admitted_ops = set() + probe_inputs = {} reference = grid_pair_module._eager_packed def record_fallback(op: str, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: @@ -566,6 +609,7 @@ def validate_candidate( **kwargs: object, ) -> CompiledKernel: probed_ops.add(_op) + probe_inputs[_op] = (_warmup, args, kwargs) if force_reference: raise grid_pair_module.triton.runtime.errors.OutOfResources( 2, 1, "test" @@ -648,6 +692,31 @@ def force_loss_gradient( assert fallback_ops == probed_ops - admitted_ops if force_reference: assert not admitted_ops + else: + for op in probed_ops - admitted_ops: + warmup, args, kwargs = probe_inputs[op] + # Probe only after the cold compiled execution. This uses + # the original kernel, independently of the selector's + # admission decision, without warming the traced path. + try: + minimum = warmup( + *args, + **( + kwargs + | { + "C_BLK": 16, + "BLOCK_G": 16, + "num_stages": 1, + "num_warps": 4, + } + ), + ) + minimum._init_handles() + except grid_pair_module.triton.runtime.errors.OutOfResources: + continue + pytest.fail( + f"{op} rejected every tile although its minimum tile loads" + ) expected = force_forward(*inputs, *projectors) expected_gradients = force_loss_gradient(expected, inputs) for got, want in zip( From 629f64ec1bdd6eb4d62b57fde2f20595bfe5bb87 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 27 Oct 2026 10:18:36 +0800 Subject: [PATCH 5/5] test(triton): require native BF16 for compiled AMP cases --- source/tests/pt_expt/kernels/test_grid_pair_train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/source/tests/pt_expt/kernels/test_grid_pair_train.py b/source/tests/pt_expt/kernels/test_grid_pair_train.py index b90102ffc9..1d1290dc38 100644 --- a/source/tests/pt_expt/kernels/test_grid_pair_train.py +++ b/source/tests/pt_expt/kernels/test_grid_pair_train.py @@ -582,6 +582,9 @@ def test_cold_compiled_force_backward( and the single-focus feed-forward grid. A device with no viable fused configuration must retain the reference composition through compilation. """ + # Inductor requires native BF16 support, not autocast's software emulation. + if amp and not torch.cuda.is_bf16_supported(including_emulation=False): + pytest.skip("Compiled BF16 requires native device support") torch._dynamo.reset() monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False)