diff --git a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py index a11f5d49d6..d67b27eb6c 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,123 @@ 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]: + """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 + 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 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. + 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 +1138,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 +1199,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..1d1290dc38 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 ( @@ -26,6 +36,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, @@ -39,10 +50,22 @@ median_deviations, ) +if TYPE_CHECKING: + from collections.abc import ( + Callable, + Iterator, + ) + + from triton.compiler.compiler import ( + CompiledKernel, + ) + 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 @@ -62,6 +85,135 @@ 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_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.""" + 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 = ( @@ -250,6 +402,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: @@ -260,6 +413,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: @@ -267,6 +421,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) @@ -276,6 +431,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: @@ -354,6 +510,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 @@ -391,3 +548,184 @@ def evaluate( torch.testing.assert_close(got, want) assert compile_counter.frame_count == 1 + + +@pytest.mark.parametrize( + ("n_focus", "channels", "amp", "force_reference"), + [ + (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, + 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. + """ + # 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) + 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, ...]: + fallback_ops.add(op) + return reference(op, *args) + + monkeypatch.setattr(grid_pair_module, "_eager_packed", record_fallback) + 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) + probe_inputs[_op] = (_warmup, args, kwargs) + 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. + n_frames, coeff_dim, n_grid = 3, 49, 587 + 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) + # 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 + 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( + (*actual, *actual_gradients), + (*expected, *expected_gradients), + strict=True, + ): + torch.testing.assert_close(got, want) + assert counter.frame_count == 1