Make the public gemm path work for the flavors it claims

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

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,143 @@ def memset_zero_async(ptr: int, nbytes: int, stream) -> None:
         raise RuntimeError(f"cudaMemsetAsync failed: {err}")
 
 
+_WORD_FORMAT = {"fp32": "<f", "int32": "<i"}
+
+
+def init_word(dtype: str, value) -> 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 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.
+
+    :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
+
+    res = _drv.cuMemsetD32Async(int(ptr), int(word), int(count), 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"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.

python/cudnn/frost/workspace.py

@@ -125,7 +125,11 @@ def over(cls, variant_pack, required_bytes: int, owner: str, *, align: int = DEF
                 f"{owner} requires a {required_bytes}-byte workspace but execute() received "
                 f"none; allocate graph.get_workspace_size() bytes and pass the buffer to execute()"
             )
-        if nbytes < required_bytes:
+        # 0 means the pack could not measure it, not that it is empty: a bare
+        # device address carries no size, and the backend takes one without
+        # checking either. Refusing here would make the same call depend on
+        # which plan ran.
+        if nbytes and nbytes < required_bytes:
             raise ValueError(f"{owner}: needs a {required_bytes}-byte workspace, got {nbytes} bytes (size it with graph.get_workspace_size())")
         if ptr % align != 0:
             raise ValueError(f"{owner}: the workspace buffer must be {align}-byte aligned; got 0x{ptr:x}")
@@ -164,9 +168,14 @@ 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:
         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)")

python/cudnn/gemm/frost/compiler.py

@@ -1803,13 +1803,30 @@ 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]
+        # 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):
+            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:
@@ -1999,7 +2016,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 +2103,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 +2175,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]
@@ -2965,7 +2982,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])
@@ -3091,7 +3108,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]
@@ -3347,7 +3364,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])
@@ -3483,7 +3500,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]

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

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

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

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<VariantPackSlot *> 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)");