From f75e208348c0179bd4e9a1dad4ec45ab56ae00c1 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 22:25:46 -0700 Subject: [PATCH 1/4] Make the public gemm path work for the flavors it claims #547 normalized the variant pack, which changed what an engine is handed: buffers arrive as slots that describe memory rather than as tensors. Two places in the gemm engine were writing THROUGH the caller's buffer with torch methods, which worked only while the buffer happened to be a torch tensor: - a reduction output is seeded with its identity before the kernel runs, via tensor.fill_() - a norm2 reduction is finalized with tensor.sqrt_() The engine owns the first one now. Every seed a reduction uses (0, 1, +-inf) is a 32-bit pattern, so buffers.fill_f32_async drives cuMemsetD32Async and no kernel is needed. The second is unreachable: a norm2 reduction is refused while the graph is lowered, so no plan exists to execute -- recorded as a test rather than left as a live-looking path. A bare device address as the WORKSPACE measured 0 bytes, and Workspace.over read that as "empty" and refused any engine that needs scratch. It means "the pack could not measure it": a raw pointer carries no size and the backend takes one without checking, so refusing here made the same call depend on which plan ran. Zero now skips the size check, here and in the C carve's bounds check. The reason none of this was caught: no test drove a non-trivial gemm flavor through graph.execute(). The direct-call tests construct a fusion chain and invoke the compiled object with torch tensors, which skips operand binding, the pack, and the conversion execute() performs -- exactly the part that changed. test_public_execute_flavors.py covers plain matmul, epilogue fusion, the four reduction modes, norm2's refusal, and the bare-address operand form through the entry point a caller actually has. That last one is xfail: frost reads its extents by axis position, and a bare address describes the operand the way the GRAPH declares it (a matmul's B is [batch, K, N]) rather than the way a caller's buffer reports it. It was broken before this too -- the geometry-less Tensor made it an IndexError instead. The fix is the engine recording which axis is M/N/K at build. Also: two unused imports in engines/base.py that broke the lazy frost boundary, and three kernel docstrings claiming the descriptor builders do not recur. --- python/cudnn/engines/base.py | 2 - python/cudnn/frost/buffers.py | 19 ++ python/cudnn/frost/workspace.py | 8 +- python/cudnn/gemm/frost/compiler.py | 24 ++- .../frost/kernel/gdn2_prefill_f16.py | 5 +- .../frost/kernel/gdn_bprop_f16.py | 5 +- .../frost/kernel/kda_prefill_f16.py | 5 +- python/pygraph/variant_pack.cpp | 7 +- .../gemm/frost/test_public_execute_flavors.py | 186 ++++++++++++++++++ 9 files changed, 240 insertions(+), 21 deletions(-) create mode 100644 test/python/gemm/frost/test_public_execute_flavors.py diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 46c523a19..907c93fa1 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -48,8 +48,6 @@ def execute(self, graph, uid_to_data, ctx): from typing import TYPE_CHECKING, Any, Dict, List from .engine_ids import PYTHON_ENGINE_ID_BASE # noqa: F401 — re-exported for engine authors -from ..datatypes import _CUDNN_TO_FROST_DTYPE_NAME -from ..frost.buffers import DeviceView if TYPE_CHECKING: from ..pygraph import pygraph diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 003b761d3..10f4b7fb2 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -18,6 +18,7 @@ from __future__ import annotations import ctypes +import struct from cudnn import _pybind_module @@ -297,6 +298,24 @@ def memset_zero_async(ptr: int, nbytes: int, stream) -> None: raise RuntimeError(f"cudaMemsetAsync failed: {err}") +def fill_f32_async(ptr: int, count: int, value: float, stream) -> None: + """Stream-ordered fill of ``count`` fp32 elements with ``value``. + + An engine that needs to seed a caller's buffer owns that operation itself: + reaching for ``tensor.fill_()`` works only while the buffer happens to be a + torch tensor, which is the coupling the variant pack exists to remove. + Every seed value a reduction uses (0, 1, +-inf) is a 32-bit pattern, so the + driver's D32 memset covers them without a kernel. + """ + from cuda.bindings import driver as _drv + + pattern = int.from_bytes(struct.pack(" buffers.DeviceView: def _check_span(self, offset: int, span: int) -> None: end = int(offset) + int(span) - if end > self._nbytes: + if self._nbytes and end > self._nbytes: raise ValueError(f"{self._owner}: workspace overrun — region [{offset}, {end}) exceeds the " f"{self._nbytes}-byte buffer (sizing bug)") diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index b44b18ebe..7eb48efb4 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1803,13 +1803,23 @@ def _expected_output_shape(spec, chain: FusionChain, mnk) -> tuple[int, int, int return tuple(1 if spec.dim[i] == 1 else full[i] for i in range(3)) -def _initialize_reduction_outputs(chain: FusionChain, outputs) -> None: +def _initialize_reduction_outputs(chain: FusionChain, outputs, stream=None) -> None: + """Seed each reduction output with its identity, before the kernel runs. + + Through the driver rather than the buffer: an engine that reaches for + ``tensor.fill_()`` works only while the caller happened to pass a torch + tensor, and the variant pack exists so that it does not have to. + """ for spec, tensor in zip(chain.outputs, outputs): if not spec.is_reduction: continue - red_idx = int(spec.source.rsplit("_", 1)[1]) - red = chain.reductions[red_idx] - tensor.fill_(_REDUCTION_INIT_VALUE[red.compute_dtype][red.mode]) + red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] + value = _REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] + fill = getattr(tensor, "fill_", None) + if fill is not None: # a torch tensor from the direct-call entry + fill(value) + else: + buffers.fill_f32_async(tensor.data_ptr(), int(tensor.numel()), value, stream) def _finalize_reductions(chain, out_bufs) -> None: @@ -1999,7 +2009,7 @@ def _call_positional(self, *args, stream=None): f"got A={tuple(a.shape)}, B={tuple(b.shape)}, " f"C={[tuple(ci.shape) for ci in cs]}" ) - _initialize_reduction_outputs(self.chain, cs) + _initialize_reduction_outputs(self.chain, cs, stream) if self.chain.output_specs: base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) @@ -2086,7 +2096,7 @@ def _call_multi_gemm(self, gemm_pairs, c, mnk, *aux, stream=None): for t in slots: if len(t.shape) != 3: raise ValueError(f"multi-GEMM {role} operand must be rank-3; got {tuple(t.shape)}") - _initialize_reduction_outputs(chain, cs) + _initialize_reduction_outputs(chain, cs, stream) base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) a_permuted = [t.permute(1, 2, 0) for t in a_slots] @@ -2158,7 +2168,7 @@ def _call_block_scale_multi_gemm(self, gemm_pairs, c, mnk, *aux, stream=None): raise ValueError( f"block-scale multi-GEMM output {spec.source!r} must have " f"shape {_expected_output_shape(spec, chain, mnk)}; " f"got {tuple(ci.shape)}" ) - _initialize_reduction_outputs(chain, cs) + _initialize_reduction_outputs(chain, cs, stream) base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) # Grouped by kind (all A, all B, all SFA, all SFB); single-GEMM → a,b,sfa,sfb. a_permuted = [d.permute(1, 2, 0) for d, _ in a_slots] diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index 83bf5b409..4d107cc21 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -2618,9 +2618,8 @@ def _build_descs( stream: cuda_driver.CUstream, ): """Build the 8 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - beta, w, o, h) into ``tensormap_workspace``. Compiled + launched - separately from the main kernel and cached by input identity in the host - bridge, so the builder launches do not recur in steady-state replay. + beta, w, o, h) into ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into + GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence length, so the main kernel's coordinates are sequence-relative and tail diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py index b78ebf950..704e7cd3b 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -3958,9 +3958,8 @@ def _build_descs( ): """Build the per-(b,h) TMA-descriptor arrays (Q, K, V, dO, H loads; dQ, dK, dV stores; the io-dtype S0 loads when ``s0`` is given) into - ``tensormap_workspace``. Compiled + launched separately from the main - kernel and cached by input identity in the host bridge, so the builder - launches do not recur in steady-state replay. + ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into + GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. The H descriptor is 3-D ``(dv, dk, h)`` over the packed ``[total_h, HO, DK, DV]`` H tensor; ``build_h_descs_kernel`` derives the diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py index 235d10e38..7184b30a4 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -2541,9 +2541,8 @@ def _build_descs( stream: cuda_driver.CUstream, ): """Build the 6 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - o, h) into ``tensormap_workspace``. Compiled + launched separately from - the main kernel and cached by input identity in the host bridge, so the - builder launches do not recur in steady-state replay. Each descriptor + o, h) into ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into + GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence length, so the main kernel's coordinates are sequence-relative and tail chunks clip in hardware. The diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 8ef85175c..0c25db2b7 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -324,6 +324,11 @@ class VariantPackSlot { // Ownership transfers with the capsule, per DLPack: the struct carries its // own copy of the shape and stride and a deleter that frees them, so it // outlives this slot rather than aliasing storage the slot owns. + // + // Unversioned only: max_version is ignored and the capsule is always + // "dltensor". The consumer this exists for is cute's compile-time + // from_dlpack; tvm-ffi reads a slot through the exchange vtable and never + // gets here. py::capsule dlpack(py::object /*stream*/, py::object /*max_version*/) const { struct Owned { @@ -691,7 +696,7 @@ class WorkspaceCarve { std::vector out; out.reserve(protos_.size()); for (size_t i = 0; i < protos_.size(); i++) { - if (ends_[i] > nbytes) { + if (nbytes != 0 && ends_[i] > nbytes) { // 0 = size unknown (bare address) throw py::value_error(owner_ + ": workspace overrun -- region [" + std::to_string(offsets_[i]) + ", " + std::to_string(ends_[i]) + ") exceeds the " + std::to_string(nbytes) + "-byte buffer (sizing bug)"); diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py new file mode 100644 index 000000000..4e16186b6 --- /dev/null +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Every gemm flavor, driven through the PUBLIC ``graph.execute()``. + +The rest of this directory calls the compiled object directly with torch +tensors. That skips the whole engine wrapper -- operand binding, the variant +pack, and the buffer conversion ``execute()`` performs -- so a break that only +appears when the engine hands the kernel what the pack holds is invisible to +it. A reduction output shipped in exactly that state: the epilogue initializes +it with ``fill_`` and finalizes ``norm2`` with ``sqrt_``, both of which the +compiled object used to receive as torch tensors and now does not. + +So this file exists to exercise the same flavors the direct-call tests cover, +but through the entry point a user actually has. +""" + +import pytest +import torch + +import cudnn +from cudnn.engines import is_python_engine + +pytestmark = pytest.mark.L0 + +_GPU = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="the FROST gemm engine claims SM100", +) + +BF16 = cudnn.data_type.BFLOAT16 +F32 = cudnn.data_type.FLOAT +M = N = 128 +K = 64 + + +@pytest.fixture(autouse=True) +def _frost_opt_in(monkeypatch): + monkeypatch.setenv("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "1") + + +def _pin_frost(g): + """Select the FROST plan, or skip when it does not claim this graph.""" + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + frost = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)] + if not frost: + pytest.skip("no FROST plan for this graph") + g.select_plan(frost[0]) + g.check_support() + g.build_plans() + return g + + +def _operands(): + a = torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda") + return a, b, torch.einsum("bmk,bnk->bmn", a.float(), b.float()) + + +def _run(g, data): + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute(data, ws) + torch.cuda.synchronize() + + +@_GPU +def test_plain_matmul(): + """The control: if this fails the harness is wrong, not the flavor.""" + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, C: c}) + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_epilogue_fusion(): + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, Y: y}) + torch.testing.assert_close(y, torch.relu(ref).to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +@pytest.mark.parametrize( + "mode,reference", + ( + (cudnn.reduction_mode.ADD, lambda t: t.sum().reshape(1, 1, 1)), + (cudnn.reduction_mode.AMAX, lambda t: t.abs().max().reshape(1, 1, 1)), + (cudnn.reduction_mode.MAX, lambda t: t.max().reshape(1, 1, 1)), + (cudnn.reduction_mode.MIN, lambda t: t.min().reshape(1, 1, 1)), + ), + ids=("add", "amax", "max", "min"), +) +def test_reduction_output(mode, reference): + """A reduction tap alongside the epilogue output. + + The epilogue writes the tap's initial value before the kernel runs and, for + ``norm2``, takes a square root after it. Both are device operations on a + caller buffer, which is where a buffer that is only a description rather + than a tensor shows up. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + R = g.reduction(input=Y, mode=mode, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]) + R.set_output(True).set_data_type(F32) + _pin_frost(g) + + a, b, ref = _operands() + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + r = torch.empty(1, 1, 1, dtype=torch.float32, device="cuda") + _run(g, {A: a, B: b, Y: y, R: r}) + + relu = torch.relu(ref) + torch.testing.assert_close(y, relu.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + torch.testing.assert_close(r, reference(relu), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_norm2_reduction_is_refused_at_build(): + """``norm2`` is the one reduction mode with a post-kernel finalize. + + It never reaches one: the backend refuses the reduction descriptor while + the graph is being lowered, so no plan exists and ``execute()`` is never + called. Recorded because the finalize would otherwise look like a live path + that needs a device-side square root. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + R = g.reduction(input=Y, mode=cudnn.reduction_mode.NORM2, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]) + R.set_output(True).set_data_type(F32) + g.validate() + with pytest.raises(RuntimeError, match="NOT_SUPPORTED"): + g.build_operation_graph() + + +@_GPU +@pytest.mark.xfail( + reason=( + "the operand a bare address describes is the GRAPH's declaration, and frost reads its " + "extents by axis position from the layout a caller's buffer would report -- for a matmul " + "B those are [batch, K, N] and (batch, N, K). Broken before this branch too (the " + "geometry-less Tensor made it an IndexError); the fix is the engine recording which axis " + "is M/N/K at build, which belongs with the executor rewrite." + ), + strict=True, +) +def test_bare_address_operands(): + """The backend has always taken a raw device address; so must a python plan.""" + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + _pin_frost(g) + + a, b, ref = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {1: a.data_ptr(), 2: b.data_ptr(), 3: c.data_ptr()}) + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) From 952bb20984bdb291ba4fbcbb66ae389458805ac2 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 23:08:33 -0700 Subject: [PATCH 2/4] Seed a reduction output through the driver, on the kernel's stream Three passes over the same fix, each caught by measuring rather than by the suite going green: 1. tensor.fill_() does not exist on a slot -> use cuMemsetD32Async. 2. That lost the STRIDE. torch's fill_ sets every element; a memset writes a contiguous byte range, and the two agree only for a dense buffer. Six strided reduction tests went NaN -- the elements past the first run were never written. 3. One memset per contiguous run is correct and 572 us for a per-row scalar output, against 3.5 for the single kernel torch launches. Correct is not mergeable; nothing in the suite would have flagged it. So: the driver where it is right, which is every contiguous buffer -- and that also puts the seed on the stream the kernel will run on, where tensor.fill_() queues on torch's current stream and is the same stream only by luck. A padded output keeps the torch path it already had, and a padded output arriving as a slot -- the public path, which is where this was broken and where nothing could seed it -- is refused with the measurement in a TODO. The fix is a fill kernel; this is the last place the engine writes through the caller's buffer. Also from review: remaining() refuses rather than returning a negative extent when the workspace size is unknown, and three cases cover the unknown-capacity path end to end (bare-address workspace, undersized-but-known, and the tail refusal). --- python/cudnn/frost/buffers.py | 16 +++-- python/cudnn/frost/workspace.py | 5 ++ python/cudnn/gemm/frost/compiler.py | 22 ++++-- .../gemm/frost/test_public_execute_flavors.py | 68 +++++++++++++++++++ 4 files changed, 101 insertions(+), 10 deletions(-) diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 10f4b7fb2..0ac9915d2 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -299,13 +299,17 @@ def memset_zero_async(ptr: int, nbytes: int, stream) -> None: def fill_f32_async(ptr: int, count: int, value: float, stream) -> None: - """Stream-ordered fill of ``count`` fp32 elements with ``value``. + """Stream-ordered fill of ``count`` CONTIGUOUS fp32 elements with ``value``. - An engine that needs to seed a caller's buffer owns that operation itself: - reaching for ``tensor.fill_()`` works only while the buffer happens to be a - torch tensor, which is the coupling the variant pack exists to remove. - Every seed value a reduction uses (0, 1, +-inf) is a 32-bit pattern, so the - driver's D32 memset covers them without a kernel. + An engine that seeds a caller's buffer owns that operation itself: reaching + for ``tensor.fill_()`` works only while the buffer happens to be a torch + tensor, and queues on torch's current stream rather than the one the kernel + will run on. Every seed a reduction uses (0, 1, +-inf) is a 32-bit pattern, + so the driver's D32 memset covers them without a kernel. + + Contiguous only, and the caller checks: a strided buffer needs one memset + per run, which for a per-row scalar output is one per row -- measured at + 572 us against 3.5 for the single kernel torch would have launched. """ from cuda.bindings import driver as _drv diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index 193fc947d..1089bfe5f 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -168,6 +168,11 @@ def take(self, numel: int, dtype: str) -> buffers.DeviceView: def remaining(self) -> buffers.DeviceView: """The tail no :meth:`take` has claimed, as uint8 — for a nested carver.""" + if not self._nbytes: + raise ValueError( + f"{self._owner}: the workspace was passed as a bare address, so its size is unknown " + "and the unclaimed tail cannot be measured; pass a sized buffer to execute()" + ) return buffers.DeviceView(self._ptr + self._offset, (self._nbytes - self._offset,), "uint8", self._device) def _check_span(self, offset: int, span: int) -> None: diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 7eb48efb4..ef2f14b69 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1815,11 +1815,25 @@ def _initialize_reduction_outputs(chain: FusionChain, outputs, stream=None) -> N continue red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] value = _REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] - fill = getattr(tensor, "fill_", None) - if fill is not None: # a torch tensor from the direct-call entry - fill(value) - else: + shape, strides = tuple(tensor.shape), tuple(tensor.stride()) + if buffers.is_contiguous(shape, strides): + # The driver, on the stream the kernel will run on. tensor.fill_() + # would queue on torch's current stream instead, which is the same + # stream only by luck. buffers.fill_f32_async(tensor.data_ptr(), int(tensor.numel()), value, stream) + continue + # A padded output needs one memset per run -- for a per-row scalar that + # is one per row, measured at 572 us against 3.5 for a single fill + # kernel. TODO: emit that kernel and delete this branch, which is the + # last place the engine writes through the caller's buffer. + fill = getattr(tensor, "fill_", None) + if fill is None: + raise NotImplementedError( + f"frost_gemm: a strided reduction output (shape {shape} stride {strides}) can only be " + "seeded through a buffer that fills itself; pass a contiguous one, or run this graph " + "on the backend" + ) + fill(value) def _finalize_reductions(chain, out_bufs) -> None: diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index 4e16186b6..db586ea49 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -184,3 +184,71 @@ def test_bare_address_operands(): c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") _run(g, {1: a.data_ptr(), 2: b.data_ptr(), 3: c.data_ptr()}) torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_bare_address_workspace(): + """A workspace passed as a raw address has no measurable size. + + Zero means "the pack could not measure it", not "empty" -- the backend + takes a raw workspace pointer without checking either, so an engine that + needs scratch must not refuse one. This drives the unknown-capacity path + through both `Workspace.over` and the C carve's bounds check. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute({A: a, B: b, C: c}, ws.data_ptr()) + torch.cuda.synchronize() + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_undersized_workspace_still_rejected(): + """A workspace whose size IS known is still bounds-checked.""" + from cudnn.engines.base import VariantPack + from cudnn.frost.workspace import Workspace + + tiny = torch.empty(16, dtype=torch.uint8, device="cuda") + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + _pin_frost(g) + + a, b, _ = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + pack = g._normalize(g._uid_to_data({1: a, 2: b, 3: c}), tiny) + assert pack.workspace_bytes == 16 + with pytest.raises(ValueError, match="needs a .*-byte workspace"): + Workspace.over(pack, 4096, "probe") + + +@_GPU +def test_unknown_size_workspace_refuses_to_measure_its_tail(): + """``remaining()`` cannot answer for a workspace whose size is unknown.""" + from cudnn.frost.workspace import Workspace + + ws = torch.empty(4096, dtype=torch.uint8, device="cuda") + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + _pin_frost(g) + + a, b, _ = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + pack = g._normalize(g._uid_to_data({1: a, 2: b, 3: c}), ws.data_ptr()) + assert pack.workspace_bytes == 0 + carver = Workspace.over(pack, 1024, "probe") + with pytest.raises(ValueError, match="size is unknown"): + carver.remaining() From 605e443f1f5e9848b2d12dd62166341c10efc9a8 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 00:31:27 -0700 Subject: [PATCH 3/4] Pack a reduction identity as the output's dtype, and seed on the MoE stream A memset moves bits, not numbers, so the identity has to be packed as the dtype the kernel reads it back as. int32's identities are the ends of its range and are exactly where that bites: -2**31 packed as float is 0xcf000000, so an int32 MAX reduction returned -822083584 for every input below it. fill_f32_async becomes fill_word_async over a 32-bit pattern, with init_word turning a value into one. The four MoE launchers seeded on the null stream rather than the execute one -- they were the call sites that had no stream to pass before this path moved to the driver, and passing None was not the same thing afterwards. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/frost/buffers.py | 29 +++++++++++++++---- python/cudnn/gemm/frost/compiler.py | 13 +++++---- .../gemm/frost/test_public_execute_flavors.py | 29 +++++++++++++++++++ 3 files changed, 59 insertions(+), 12 deletions(-) diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 0ac9915d2..417c7efea 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -298,14 +298,32 @@ def memset_zero_async(ptr: int, nbytes: int, stream) -> None: raise RuntimeError(f"cudaMemsetAsync failed: {err}") -def fill_f32_async(ptr: int, count: int, value: float, stream) -> None: - """Stream-ordered fill of ``count`` CONTIGUOUS fp32 elements with ``value``. +_WORD_FORMAT = {"fp32": " int: + """The 32-bit pattern that writes ``value`` to a buffer of ``dtype``. + + A memset moves bits, not numbers, so the value has to be packed as the dtype + the kernel will read it back as. int32's reduction identities are the ends + of its range and are exactly where that bites: -2**31 packed as float is + 0xcf000000 where the kernel wants 0x80000000. + """ + fmt = _WORD_FORMAT.get(dtype) + if fmt is None: + raise NotImplementedError(f"no 32-bit fill pattern for dtype {dtype!r}") + return int.from_bytes(struct.pack(fmt, value), "little") + + +def fill_word_async(ptr: int, count: int, word: int, stream) -> None: + """Stream-ordered fill of ``count`` CONTIGUOUS 32-bit words with ``word``. An engine that seeds a caller's buffer owns that operation itself: reaching for ``tensor.fill_()`` works only while the buffer happens to be a torch tensor, and queues on torch's current stream rather than the one the kernel - will run on. Every seed a reduction uses (0, 1, +-inf) is a 32-bit pattern, - so the driver's D32 memset covers them without a kernel. + will run on. Every seed a reduction uses is a 32-bit pattern, so the + driver's D32 memset covers them without a kernel -- see ``init_word`` for + turning a value into one. Contiguous only, and the caller checks: a strided buffer needs one memset per run, which for a per-row scalar output is one per row -- measured at @@ -313,8 +331,7 @@ def fill_f32_async(ptr: int, count: int, value: float, stream) -> None: """ from cuda.bindings import driver as _drv - pattern = int.from_bytes(struct.pack(" N if buffers.is_contiguous(shape, strides): # The driver, on the stream the kernel will run on. tensor.fill_() # would queue on torch's current stream instead, which is the same - # stream only by luck. - buffers.fill_f32_async(tensor.data_ptr(), int(tensor.numel()), value, stream) + # stream only by luck. The pattern is packed as the OUTPUT's dtype, + # not as float: int32's identities are the ends of its range. + buffers.fill_word_async(tensor.data_ptr(), int(tensor.numel()), buffers.init_word(red.compute_dtype, value), stream) continue # A padded output needs one memset per run -- for a per-row scalar that # is one per row, measured at 572 us against 3.5 for a single fill @@ -2989,7 +2990,7 @@ def _launch_single(self, token, weight, first_token_offset, output, snke, worksp raise ValueError( f"MoE output {spec.source!r} must have shape " f"{_expected_output_shape(spec, self.chain, (S, N, K))}; " f"got {tuple(t.shape)}" ) - _initialize_reduction_outputs(self.chain, outputs) + _initialize_reduction_outputs(self.chain, outputs, stream) # num_experts = weight batch (E); num_groups = first_token_offset len # (BxE, may exceed E; group g uses expert g % E). From runtime tensors. num_experts = int(weight.shape[0]) @@ -3115,7 +3116,7 @@ def _call_multi_gemm(self, gemm_pairs, first_token_offset, output, snke, *aux, w raise ValueError( f"multi-GEMM MoE output {spec.source!r} must have shape " f"{_expected_output_shape(spec, chain, (S, N, K))}; got {tuple(ci.shape)}" ) - _initialize_reduction_outputs(chain, outs) + _initialize_reduction_outputs(chain, outs, stream) num_experts = int(b_slots[0].shape[0]) num_groups = int(first_token_offset.shape[0]) a_stride_perms = [t.permute(1, 2, 0) for t in a_slots] @@ -3371,7 +3372,7 @@ def _launch_single(self, token, weight, sfa, sfb, first_token_offset, output, sn f"{_expected_output_shape(spec, self.chain, (S, N, K))}; " f"got {tuple(t.shape)}" ) - _initialize_reduction_outputs(self.chain, outputs) + _initialize_reduction_outputs(self.chain, outputs, stream) # num_experts = weight batch (E); num_groups = first_token_offset len # (BxE, may exceed E; group g uses expert g % E). From runtime tensors. num_experts = int(weight.shape[0]) @@ -3507,7 +3508,7 @@ def _call_multi_gemm(self, gemm_pairs, first_token_offset, output, snke, *aux, w f"multi-GEMM MoE block-scale output {spec.source!r} must have shape " f"{_expected_output_shape(spec, chain, (S, N, K))}; got {tuple(ci.shape)}" ) - _initialize_reduction_outputs(chain, outs) + _initialize_reduction_outputs(chain, outs, stream) num_experts = int(b_slots[0][0].shape[0]) num_groups = int(first_token_offset.shape[0]) a0, b0 = a_slots[0][0], b_slots[0][0] diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index db586ea49..7eda4645b 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -137,6 +137,35 @@ def test_reduction_output(mode, reference): torch.testing.assert_close(r, reference(relu), atol=1e-1, rtol=1e-2) +@_GPU +def test_int32_reduction_seed_is_packed_as_int32(): + """A memset moves bits, so the identity has to be packed as the dtype. + + int32's identities are the ends of its range, which is exactly where the + difference shows: -2**31 packed as float is 0xcf000000, and a MAX reduction + seeded with that returns -822083584 for any input below it. + """ + I32 = cudnn.data_type.INT32 + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + bias = g.tensor(name="bias", dim=[1, 1, N], stride=[N, N, 1], data_type=I32) + C = g.matmul(A=A, B=B, name="mm") + Y = g.add(a=C, b=bias, name="add_i32", compute_data_type=I32) + Y.set_output(True).set_data_type(I32) + R = g.reduction(input=Y, mode=cudnn.reduction_mode.MAX, name="red", compute_data_type=I32) + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]).set_output(True).set_data_type(I32) + _pin_frost(g) + + floor = -2_000_000_000 # below the float-packed seed, above int32's minimum + a = torch.zeros(1, M, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(1, N, K, dtype=torch.bfloat16, device="cuda") + y = torch.empty(1, M, N, dtype=torch.int32, device="cuda") + r = torch.empty(1, 1, 1, dtype=torch.int32, device="cuda") + _run(g, {A: a, B: b, bias: torch.full((1, 1, N), floor, dtype=torch.int32, device="cuda"), Y: y, R: r}) + assert int(r.item()) == floor + + @_GPU def test_norm2_reduction_is_refused_at_build(): """``norm2`` is the one reduction mode with a post-kernel finalize. From 99c44422b4c8e095afafba24d7618c6f0f46db39 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 17:20:07 -0700 Subject: [PATCH 4/4] Seed a padded reduction output through the driver too The contiguous case already went through cuMemsetD32Async on the execute-time stream. The padded one fell back to tensor.fill_(), which queues on torch's CURRENT stream -- the same stream only by luck -- and exists at all only while the caller happened to pass a torch tensor. That contradicted what this branch claims to do, so it is gone. strided_fill_plan collapses the layout to the runs a memset can cover and returns the 2D memsets that cover it exactly once. cuMemsetD2D32Async takes a pitch, so a per-row scalar tap is ONE call rather than one per row -- that reading, 572 us at one memset per row, is why the fallback was there. What remains is one call per point of whatever axis is left outside the 2D region, which for a rank-3 output is the batch and is usually 1. The plan is returned before anything is written, and is None for a layout that would write an element twice: a stride of 0 over a real extent, or an outer stride that does not clear the axis below it. Checking only the innermost pair (pitch >= width) is not enough -- shape (2, 2) stride (2, 2) has width 1 and lands both axes on element 2. Also in this commit: a missing space in three linear-attention kernel docstrings (review caught two of the three), and a raw-string pytest.raises pattern. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/frost/buffers.py | 104 +++++++++++++++++- python/cudnn/gemm/frost/compiler.py | 28 ++--- .../frost/kernel/gdn2_prefill_f16.py | 2 +- .../frost/kernel/gdn_bprop_f16.py | 2 +- .../frost/kernel/kda_prefill_f16.py | 2 +- .../gemm/frost/test_public_execute_flavors.py | 101 ++++++++++++++++- 6 files changed, 214 insertions(+), 25 deletions(-) diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 417c7efea..280928e71 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -325,9 +325,8 @@ def fill_word_async(ptr: int, count: int, word: int, stream) -> None: driver's D32 memset covers them without a kernel -- see ``init_word`` for turning a value into one. - Contiguous only, and the caller checks: a strided buffer needs one memset - per run, which for a per-row scalar output is one per row -- measured at - 572 us against 3.5 for the single kernel torch would have launched. + :func:`fill_word_strided_async` is the same fill for a buffer that is not + one dense run. """ from cuda.bindings import driver as _drv @@ -337,6 +336,105 @@ def fill_word_async(ptr: int, count: int, word: int, stream) -> None: raise RuntimeError(f"cuMemsetD32Async failed: {err}") +def _fill_word_2d_async(ptr: int, pitch_words: int, width: int, height: int, word: int, stream) -> None: + from cuda.bindings import driver as _drv + + res = _drv.cuMemsetD2D32Async(int(ptr), int(pitch_words) * 4, int(word), int(width), int(height), int(stream) if stream is not None else 0) + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"cuMemsetD2D32Async failed: {err}") + + +def collapse_layout(shape, strides) -> list: + """``(extent, stride)`` outermost first, with unit axes dropped and adjacent + axes merged where one exactly fills the other's gap. + + A padded output is usually dense underneath its declared rank -- a rank-3 + ``(1, M, 1)`` tap is one strided run, and a contiguous one is a single dense + run whatever rank it was declared at. Merging first is what keeps the fill + below down to one memset in both cases. + """ + axes = sorted(((int(d), int(s)) for d, s in zip(shape, strides) if int(d) != 1), key=lambda ds: -ds[1]) + out: list = [] + for extent, stride in axes: + if out and out[-1][1] == extent * stride: + out[-1] = (out[-1][0] * extent, stride) + else: + out.append((extent, stride)) + return out + + +def strided_fill_plan(shape, strides) -> "list | None": + """The 2D memsets that cover a strided region exactly once, or None. + + None means the region writes some element twice -- a stride of 0 over a real + extent, or an outer stride that does not clear the axis below it. That is a + write race whichever buffer it is, so it is refused rather than issued; the + caller decides how to say so. + + Returned before anything is written, which is the point: the seed's + preconditions have to be settled while the caller's buffer is still + untouched, and a plan is what lets several outputs all be checked before the + first of them is filled. + + Each entry is ``(offset, pitch, width, height)`` in ELEMENTS. The driver's + 2D memset takes a pitch, so a per-row scalar tap is one entry rather than one + per row (the reading that made this look expensive: 572 us at one memset per + row); what remains is one entry per point of whatever axis is left outside + the 2D region, which for a rank-3 output is the batch and is usually 1. + """ + if any(int(s) == 0 and int(d) != 1 for d, s in zip(shape, strides)): + return None + axes = collapse_layout(shape, strides) + # Non-overlapping iff each axis clears the whole span of the one below it. + # `pitch >= width` is this rule at the innermost pair and misses the rest: + # shape (2, 2) stride (2, 2) has width 1 and passes it, and lands both axes + # on the same element. + for (_outer_extent, outer_stride), (inner_extent, inner_stride) in zip(axes, axes[1:]): + if outer_stride < inner_extent * inner_stride: + return None + if not axes: + return [(0, 1, 1, 1)] + # The innermost run is the memset's width when it is dense; otherwise every + # element stands alone and the width is one. + width, rest = (axes[-1][0], axes[:-1]) if axes[-1][1] == 1 else (1, axes) + if not rest: + return [(0, width, width, 1)] + height, pitch = rest[-1] + offsets = [0] + for extent, stride in reversed(rest[:-1]): + offsets = [base + i * stride for base in offsets for i in range(extent)] + return [(base, pitch, width, height) for base in offsets] + + +def apply_fill_plan(ptr: int, plan, word: int, stream) -> None: + """Issue a plan from :func:`strided_fill_plan`, stream-ordered.""" + for offset, pitch, width, height in plan: + if height == 1: + fill_word_async(ptr + offset * 4, width, word, stream) + else: + _fill_word_2d_async(ptr + offset * 4, pitch, width, height, word, stream) + + +def fill_word_strided_async(ptr: int, shape, strides, elem_bytes: int, word: int, stream) -> None: + """Plan and issue in one call, for a caller with a single region to seed. + + The engine owns seeding a reduction output, and a padded one is the case + that used to send it back to the caller's ``fill_()`` -- the last place + anything here wrote through a buffer it does not own, and the reason a + perfectly legal call had to fall off the fast path. A caller with SEVERAL + regions wants :func:`strided_fill_plan` for all of them first: this one + cannot know whether the next region is refusable, so it would leave the + earlier ones filled. + """ + if elem_bytes != 4: + raise NotImplementedError(f"frost: a reduction seed is a 32-bit pattern; this output stores {elem_bytes}-byte elements") + plan = strided_fill_plan(shape, strides) + if plan is None: + raise ValueError(f"frost: a reduction output cannot write an element twice (shape {tuple(shape)} stride {tuple(strides)})") + apply_fill_plan(ptr, plan, word, stream) + + # The CuTe primitives these engines lower through landed in 4.7.0; older DSLs # fail during codegen with errors that name a missing attribute rather than the # version, so the check belongs where an engine can still decline. diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index d961a7681..011322cf8 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1815,26 +1815,18 @@ def _initialize_reduction_outputs(chain: FusionChain, outputs, stream=None) -> N continue red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] value = _REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] + # The driver, on the stream the kernel will run on -- for a padded output + # too. tensor.fill_() would queue on torch's current stream instead, + # which is the same stream only by luck, and only exists at all while + # the caller happened to pass a torch tensor. The pattern is packed as + # the OUTPUT's dtype, not as float: int32's identities are the ends of + # its range. shape, strides = tuple(tensor.shape), tuple(tensor.stride()) + word = buffers.init_word(red.compute_dtype, value) if buffers.is_contiguous(shape, strides): - # The driver, on the stream the kernel will run on. tensor.fill_() - # would queue on torch's current stream instead, which is the same - # stream only by luck. The pattern is packed as the OUTPUT's dtype, - # not as float: int32's identities are the ends of its range. - buffers.fill_word_async(tensor.data_ptr(), int(tensor.numel()), buffers.init_word(red.compute_dtype, value), stream) - continue - # A padded output needs one memset per run -- for a per-row scalar that - # is one per row, measured at 572 us against 3.5 for a single fill - # kernel. TODO: emit that kernel and delete this branch, which is the - # last place the engine writes through the caller's buffer. - fill = getattr(tensor, "fill_", None) - if fill is None: - raise NotImplementedError( - f"frost_gemm: a strided reduction output (shape {shape} stride {strides}) can only be " - "seeded through a buffer that fills itself; pass a contiguous one, or run this graph " - "on the backend" - ) - fill(value) + buffers.fill_word_async(tensor.data_ptr(), int(tensor.numel()), word, stream) + else: + buffers.fill_word_strided_async(tensor.data_ptr(), shape, strides, tensor.element_size(), word, stream) def _finalize_reductions(chain, out_bufs) -> None: diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index 4d107cc21..ff80aad30 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -2618,7 +2618,7 @@ def _build_descs( stream: cuda_driver.CUstream, ): """Build the 8 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - beta, w, o, h) into ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into + beta, w, o, h) into ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py index 704e7cd3b..f9fd8c670 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -3958,7 +3958,7 @@ def _build_descs( ): """Build the per-(b,h) TMA-descriptor arrays (Q, K, V, dO, H loads; dQ, dK, dV stores; the io-dtype S0 loads when ``s0`` is given) into - ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into + ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. The H descriptor is 3-D ``(dv, dk, h)`` over the packed diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py index 7184b30a4..dfd622955 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -2541,7 +2541,7 @@ def _build_descs( stream: cuda_driver.CUstream, ): """Build the 6 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - o, h) into ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into + o, h) into ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence length, so the main kernel's diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index 7eda4645b..068e159c4 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -257,7 +257,7 @@ def test_undersized_workspace_still_rejected(): c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") pack = g._normalize(g._uid_to_data({1: a, 2: b, 3: c}), tiny) assert pack.workspace_bytes == 16 - with pytest.raises(ValueError, match="needs a .*-byte workspace"): + with pytest.raises(ValueError, match=r"needs a .*-byte workspace"): Workspace.over(pack, 4096, "probe") @@ -281,3 +281,102 @@ def test_unknown_size_workspace_refuses_to_measure_its_tail(): carver = Workspace.over(pack, 1024, "probe") with pytest.raises(ValueError, match="size is unknown"): carver.remaining() + + +# --- seeding a padded reduction output, which the engine owns ---------------- + + +@pytest.mark.parametrize( + "shape,stride,want", + [ + ((1, 8, 4), (32, 4, 1), [(32, 1)]), # dense whatever rank it was declared at + ((1, 8, 1), (32, 4, 1), [(8, 4)]), # a per-row scalar: one strided run + ((2, 8, 4), (64, 8, 1), [(16, 8), (4, 1)]), # padded rows: batch merges into the row count + ((1, 1, 1), (1, 1, 1), []), # one element + ], +) +def test_a_layout_collapses_to_the_runs_a_memset_can_cover(shape, stride, want): + """Unit axes carry no elements and adjacent dense axes are one run. + + Collapsing first is what keeps the seed to a single memset for a contiguous + tap and one per batch for a padded one, rather than one per row. + """ + from cudnn.frost.buffers import collapse_layout + + assert collapse_layout(shape, stride) == want + + +@pytest.mark.parametrize( + "shape,stride", + [ + ((1, 8, 1), (0, 0, 1)), # stride 0 over a real extent + ((1, 4, 4), (16, 2, 1)), # rows closer together than they are wide + ((1, 2, 2), (4, 2, 2)), # two axes landing on the same elements + ], +) +def test_a_reduction_output_that_writes_an_element_twice_is_rejected(shape, stride): + """Two elements at one address is a write race, not a layout to support. + + The rule is that each axis clears the whole span of the one below it. + Checking the innermost pair alone -- pitch against width -- passes the last + case, whose width is 1 and whose two axes both land on element 2. + """ + from cudnn.frost.buffers import fill_word_strided_async, strided_fill_plan + + assert strided_fill_plan(shape, stride) is None + with pytest.raises(ValueError, match="twice"): + fill_word_strided_async(0, shape, stride, 4, 0, None) + + +@_GPU +@pytest.mark.parametrize("shape,stride", [((1, 8, 1), (32, 4, 1)), ((2, 8, 4), (64, 8, 1)), ((3, 5, 1), (7, 1, 1))]) +def test_a_padded_seed_writes_its_own_elements_and_no_others(shape, stride): + """The engine seeds a padded output itself, without the caller's ``fill_()``. + + Borrowing that method worked only while the buffer happened to be a torch + tensor -- and queued on torch's stream, not the one the kernel will run on. + What it has to get right is exactly this: every element the view covers, and + nothing between them. + """ + from cudnn.frost.buffers import fill_word_strided_async, init_word + + span = 1 + sum((d - 1) * s for d, s in zip(shape, stride)) + flat = torch.zeros(span, dtype=torch.float32, device="cuda") + view = torch.as_strided(flat, shape, stride) + fill_word_strided_async(flat.data_ptr(), shape, stride, 4, init_word("fp32", 3.5), None) + torch.cuda.synchronize() + + expected = torch.zeros(span, dtype=torch.float32, device="cuda") + torch.as_strided(expected, shape, stride).fill_(3.5) + assert torch.equal(flat, expected) + assert torch.equal(view, torch.full(shape, 3.5, device="cuda")) + + +@_GPU +def test_a_padded_reduction_output_is_seeded_on_the_kernel_s_stream(): + """The case that used to reach ``tensor.fill_()``. + + A tap declared into a padded buffer is legal and goes through the public + path like any other. Seeding it by calling ``fill_()`` on the caller's + tensor queued on torch's current stream rather than the one the kernel runs + on -- the same stream only by luck -- and only worked at all while that + buffer was a torch tensor. It is the driver's 2D memset now, so this asserts + both halves: the tap is right, and the padding it does not own is untouched. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + R = g.reduction(input=Y, mode=cudnn.reduction_mode.ADD, name="red") + R.set_dim([1, M, 1]).set_stride([M * 4, 4, 1]).set_output(True).set_data_type(F32) + _pin_frost(g) + + a, b, ref = _operands() + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + pad = torch.full((1, M, 4), -7.0, dtype=torch.float32, device="cuda") + _run(g, {A: a, B: b, Y: y, R: pad[:, :, :1]}) + + torch.testing.assert_close(pad[:, :, 0], torch.relu(ref).sum(dim=2), atol=1.0, rtol=2e-2) + assert torch.equal(pad[:, :, 1:], torch.full((1, M, 3), -7.0, device="cuda"))