From 78c842aa64e086a871a37b6a0b31d08930fd269f Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 23:49:51 -0700 Subject: [PATCH 1/9] Decide a gemm's per-call facts when it compiles, and lower them to one line Which axis of each operand carries M/N/K, each major, the fp4 packing factor, every output's required alignment and shape rule, which outputs are reductions: all settled by the time cute hands back a launchable, and all re-derived on every execute. gemm/frost/recipe.py reads them once, and gives that table two consumers -- run_views interprets it and serves every flavor, _lower emits a straight line for the single-GEMM shape. 256x256x128 bf16, host enqueue, min over 25 reps of a 64-call burst from a drained queue, one process and one plan: 44.13 -> 34.79 interpreted -> 17.54 lowered. The lowered path never raises. Anything it is not certain of it hands to run_views, which owns every rejection message, so it can only accept a subset of what the general path accepts and there is no second set of error strings. The version of this I wrote by hand first had two such divergences -- it lost the operand batch check and pinned an fp4 output at N where the graph says N/2. Also closes the bare-address xfail: cuDNN declares a matmul's B as [batch, K, N] while a caller allocates it (batch, N, K), so reading an extent by axis position answered one of those and not the other. The recipe records both orders and picks by where stride 1 landed. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/frost/README.md | 10 + python/cudnn/gemm/frost/compiler.py | 349 +++++++-------- python/cudnn/gemm/frost/dtypes.py | 18 +- python/cudnn/gemm/frost/engine.py | 26 +- python/cudnn/gemm/frost/recipe.py | 416 ++++++++++++++++++ test/python/gemm/frost/test_execute_recipe.py | 261 +++++++++++ .../gemm/frost/test_public_execute_flavors.py | 18 +- 7 files changed, 889 insertions(+), 209 deletions(-) create mode 100644 python/cudnn/gemm/frost/recipe.py create mode 100644 test/python/gemm/frost/test_execute_recipe.py diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index e76a14c7f..348c8f488 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -235,6 +235,16 @@ CompiledPlan.execute(graph, uid_to_data, ctx) (the hot path) - **`build_plan` runs once per (graph, plan)** at `build_plans()` time and the compiled artifact lives on the graph, so one engine instance is safely reusable across graphs. +- **What a runtime value cannot change belongs in a build-time table.** An + operand's role and major, each output's shape rule and required alignment, + which outputs are reductions -- all settled when the kernel compiled, and + deciding them again per call is most of what a python execute path costs + (measured: 44 -> 18 us for one gemm). `gemm/frost/recipe.py` is the worked + example: one table, read by an interpreter that serves every flavor and by a + straight line lowered for the common one. Emitting a second call path per + flavor instead is how the two disagree, so the lowered one never raises -- + anything it is unsure of it hands back to the interpreter, which owns every + rejection message. - **`ExecutionContext` carries handle, stream and workspace explicitly.** No engine may hard-code a stream, reach into private graph state, or allocate hidden workspace. `uid_to_data` is the caller's variant pack (tensor uid -> diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 011322cf8..2c82bc0c3 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -19,11 +19,12 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Callable, ClassVar +from typing import Any, Callable, ClassVar import cudnn from cuda.bindings import driver as _cuda from cudnn.frost import buffers +from cudnn.frost.device import current_device from cudnn.frost.workspace import Workspace _LOG = logging.getLogger(__name__) @@ -53,6 +54,21 @@ def _as_custream(stream): ) from .epilogue_codegen import EpilogueSnippets, generate from .fusion_ir import ZERO_PRESERVING_OPS, FusionChain, TensorRef +from .recipe import ( + AX_BATCH, + AX_K, + AX_MN, + CONST, + FROM_M, + FROM_N, + KERNEL_AXES, + REDUCTION_INIT_VALUE, + _output_rule, + build as build_recipe, + contiguous_modulus, + expected_shape, + gate as gate_recipe, +) from .graph_analyzer import ( GemmBinding, analyze_with_binding, @@ -711,8 +727,6 @@ def _plan_device() -> int: """The GPU a plan compiled right now targets. Every device-derived constant baked into the kernel (ab_stages, grid_num_clusters, the target SM) is read for this device, so the plan records it and re-checks it at execute time.""" - from cudnn.frost.device import current_device - return current_device() @@ -726,10 +740,9 @@ def _check_plan_device(plan_device: int) -> None: devices either: ``create_variant_pack`` sets only pointers, uids and the workspace, and ``graph_interface.h`` takes its ordinal from ``cuda_get_device``. Reading the current device once is 0.74 us against - 1.45 per operand for the walk this replaces. + 1.45 per operand for the walk this replaces -- and the import is at module + scope because doing it here costs 1.1 us of the 1.7 this function takes. """ - from cudnn.frost.device import current_device - device = current_device() if device != plan_device: raise ValueError( @@ -1765,42 +1778,8 @@ def _reshape_aux_to_fake(t: object, ref: TensorRef) -> object: return t.reshape(shape) -_REDUCTION_INIT_VALUE = { - "fp32": { - "add": 0.0, - "amax": 0.0, - "max": -float("inf"), - "min": float("inf"), - "avg": 0.0, - "norm1": 0.0, - "norm2": 0.0, - "mul": 1.0, - "mul_no_zeros": 1.0, - }, - "int32": { - "add": 0, - "amax": 0, - "max": -(2**31), - "min": 2**31 - 1, - "norm1": 0, - }, -} - - def _expected_output_shape(spec, chain: FusionChain, mnk) -> tuple[int, int, int]: - full = (chain.matmul.batch, int(mnk[0]), int(mnk[1])) - if spec.is_quant_scale: - assert spec.dim is not None - return spec.dim - if not spec.is_reduction: - if spec.dtype == "fp4_e2m1": - return (full[0], full[1], full[2] // 2) - return full - assert spec.dim is not None - red_idx = int(spec.source.rsplit("_", 1)[1]) - if chain.reductions[red_idx].grouped_by_moe: - return spec.dim - return tuple(1 if spec.dim[i] == 1 else full[i] for i in range(3)) + return expected_shape(_output_rule(spec, chain), int(mnk[0]), int(mnk[1])) def _initialize_reduction_outputs(chain: FusionChain, outputs, stream=None) -> None: @@ -1814,7 +1793,7 @@ def _initialize_reduction_outputs(chain: FusionChain, outputs, stream=None) -> N if not spec.is_reduction: continue red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] - value = _REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] + 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 @@ -1873,6 +1852,19 @@ class CompiledFusedGemm: # from the execute-time cuDNN handle and forwards it as `stream=`. Engines # that do not carry the param stay on the default stream (see dispatch). accepts_stream: ClassVar[bool] = True + # Everything the call path needs that a runtime value cannot change, read + # once here rather than rebuilt per execute. None when this object was + # constructed without a binding and so has no call path. + recipe: Any = field(default=None, init=False, repr=False, compare=False) + # The recipe as one straight line, when this graph is a shape it emits. + launch: Any = field(default=None, init=False, repr=False, compare=False) + bound: Any = field(default=(), init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + if self.binding is not None: + self.bound = tuple(self.binding.bound_tensors()) + self.recipe = build_recipe(self) + self.launch = self._lower() def __call__(self, variant_pack, stream=None): # The runtime call is a variant-pack dict keyed by cuDNN tensor object @@ -1886,96 +1878,144 @@ def __call__(self, variant_pack, stream=None): return self.run_resolved(resolve_variant_pack(variant_pack, self.binding), stream=stream) def run_resolved(self, resolved, stream=None): - """Launch over ``{id(bound_tensor): buffer}``, already resolved. + """Launch over ``{id(bound_tensor): buffer}``, already resolved.""" + views = [] + for i, t in enumerate(self.bound): + buf = resolved.get(id(t)) + if buf is None: + raise KeyError(f"variant pack is missing a buffer for {self.recipe.roles[i]}") + views.append(buf) + return (self.launch or self.run_views)(views, stream=stream) + + def _lower(self): + """The recipe as one straight line, or None for a shape it does not emit. + + Every branch the general path takes -- multi-GEMM or not, block scale or + not, how many outputs, which are reductions, which axis of each operand + carries M/N/K -- was settled when the kernel compiled. Taking them again + per call is what stands between this path and its floor: the + load-bearing validation measures 1.1 us against 21 for the walk that + rebuilds it. + + The emitted body never raises. Anything it is not sure of it hands to + ``run_views``, which serves every flavor and owns every rejection + message -- so this can only ever accept a subset of what the general + path accepts, and there is no second set of error strings to drift. + + Emitting another flavor means widening the recipe, not copying this + body: a second body is the drift the recipe exists to prevent. + """ + r = self.recipe + if r is None or not _TVM_FFI_OK or r.multi_gemm or r.block_size or r.aux or r.workspace_bytes: + return None + if len(r.inputs) != 2 or len(r.outputs) != 1: + return None + out = r.outputs[0] + # A raw output is a reduction or a quant scale: both want a pre-kernel + # seed and one wants a post-kernel sqrt, which are device operations the + # engine does not own yet (see _initialize_reduction_outputs). + if out.raw or out.init is not None or out.sqrt: + return None + rule = out.rule + if tuple(s for s, _ in rule) != (CONST, FROM_M, FROM_N): + return None + + a, b = r.a, r.b + ai, bi, ci = a.view, b.view, out.view + # kc is both the axis whose stride must be 1 and the axis whose extent + # enters the TMA rule -- they are the same axis by definition of major. + a_kc, b_kc = a.kc, b.kc + a_mod, b_mod = a.modulus, b.modulus + a_batch, b_batch = a.batch, b.batch + a_kpack, b_kpack = a.kpack, b.kpack + out_batch, out_ndiv, out_align = rule[0][1], rule[2][1], out.align + # Every batch extent the launch could read is pinned by a check above + # it, so the kernel's batch is settled here rather than re-read. + batch = out_batch if r.has_output_specs else max(a_batch, b_batch) + device = self.device + launchable = self._launchable + general = self.run_views + + def launch(views, stream=None): + _check_plan_device(device) + av, bv, cv = views[ai], views[bi], views[ci] + a_sh, b_sh, c_sh = av.shape, bv.shape, cv.shape + if len(a_sh) != 3 or len(b_sh) != 3 or len(c_sh) != 3: + return general(views, stream=stream) + a_st, b_st, c_st = av.stride(), bv.stride(), cv.stride() + m, n, k = a_sh[1], b_sh[1], a_sh[2] * a_kpack + if ( + a_st[a_kc] != 1 + or b_st[b_kc] != 1 + or a_sh[a_kc] % a_mod + or b_sh[b_kc] % b_mod + or a_sh[0] != a_batch + or b_sh[0] != b_batch + or b_sh[2] * b_kpack != k + or c_sh[0] != out_batch + or c_sh[1] != m + or c_sh[2] != n // out_ndiv + or _pow2_floor(av.data_ptr()) < 16 + or _pow2_floor(bv.data_ptr()) < 16 + or tensor_alignment(tuple(c_sh), tuple(c_st), cv.element_size(), ptr=cv.data_ptr()) < out_align + ): + return general(views, stream=stream) + problem_size = ( + m, + n, + k, + batch, + # permute(1, 2, 0) relabels axes, so the kernel's strides are + # that permutation of the ones just read -- no second read. + a_st[1], + a_st[2], + a_st[0], + b_st[1], + b_st[2], + b_st[0], + c_st[1], + c_st[2], + c_st[0], + ) + # One output and no aux: the TMA-store and plain argument orders + # coincide, and a shape where they do not is not emitted here. + return launchable(problem_size, av.permute(1, 2, 0), bv.permute(1, 2, 0), cv.permute(1, 2, 0), stream=_as_custream(stream)) + + return launch + + def run_views(self, views, stream=None): + """Launch over the operand buffers, in bound-tensor order. Which bound tensor holds which operand is fixed when the plan compiles, so a caller that knows it -- the engine, which binds the graph's slots - once -- has no reason to rebuild the by-object / by-uid / by-name tables - resolve_variant_pack needs on every execute. + once -- indexes rather than resolves, and the gate below reads a table + built at that same moment instead of rebuilding one per call. """ _check_plan_device(self.device) - b = self.binding - - def pull(t, role): - if t is None or id(t) not in resolved: - raise KeyError(f"variant pack is missing a buffer for {role}") - return resolved[id(t)] - - a_bufs = [pull(t, "A operand") for t in b.a_operands] - b_bufs = [pull(t, "B operand") for t in b.b_operands] - out_bufs = [pull(t, "output") for t in b.outputs] - aux_bufs = [pull(t, f"aux {self.aux_names[i]!r}") for i, t in enumerate(b.aux)] - # (M, N, K) from buffer shapes: A=(batch,M,K) (FP4-packed A stores K/2 - # elem/byte → scale back up), N from the B operand. - k_factor = 2 if self.chain.matmul.a_dtype == "fp4_e2m1" else 1 - M = a_bufs[0].shape[1] - K = a_bufs[0].shape[2] * k_factor - # N from B: a dense output may be FP4-packed (N/2 bytes). - N = b_bufs[0].shape[1] - mnk = (M, N, K) - c_arg = out_bufs if len(out_bufs) > 1 else out_bufs[0] - - # The kernel is shape-agnostic, so the RUNTIME dims must satisfy the - # TMA 16-byte alignment rule — same check as the graph-time gate. - mm = self.chain.matmul - _align_reason = _tma_alignment_reject(mm.a_dtype, mm.b_dtype, mm.a_major, mm.b_major, M, N, K) - if _align_reason is not None: - raise ValueError(_align_reason) - - # M/N/K came from the FIRST A and B buffer; every other operand must - # agree, and every operand's layout must match the compiled major. - _operands = [(f"A operand[{i}]", x, mm.a_major, M, 2 if mm.a_dtype == "fp4_e2m1" else 1) for i, x in enumerate(a_bufs)] + [ - (f"B operand[{i}]", x, mm.b_major, N, 2 if mm.b_dtype == "fp4_e2m1" else 1) for i, x in enumerate(b_bufs) - ] - for _reject in ( - _operand_shape_reject(_operands, M, N, K), - _operand_layout_reject([(role, x, major) for role, x, major, _d1, _kp in _operands]), - ): - if _reject is not None: - raise ValueError(_reject) + recipe = self.recipe + mnk, axes = recipe.problem(views) + gate_recipe(recipe, views, mnk, axes) - _out_reqs = _output_align_reqs(self.chain, self.use_tma_store, vec_bytes=self.vec_bytes_epi) - _aux_reqs = _aux_align_reqs(self.chain, vec_bytes=self.vec_bytes_epi) - _named = ( - [("A operand", x, 16, "ptr") for x in a_bufs] - + [("B operand", x, 16, "ptr") for x in b_bufs] - + [(f"output[{i}]", x, _out_reqs[i], "full") for i, x in enumerate(out_bufs)] - + [(f"aux {self.aux_names[i]!r}", x, _aux_reqs[self.aux_names[i]], "full") for i, x in enumerate(aux_bufs)] - ) - if self.block_scale: - _named += [("SFA", pull(t, "SFA"), 16, "ptr") for t in b.sfa_operands] - _named += [("SFB", pull(t, "SFB"), 16, "ptr") for t in b.sfb_operands] - _align_reason2 = _alignment_reject(_named) - if _align_reason2 is not None: - raise ValueError(_align_reason2) - - if self.block_scale: - _sf_k4 = ((K // self.chain.block_scale.block_size) + 3) // 4 - _sf_reason = _sf_blob_reject( - [(f"SFA[{i}]", pull(t, "SFA"), 512 * _sf_k4 * ((M + 127) // 128) * int(a_bufs[i].shape[0])) for i, t in enumerate(b.sfa_operands)] - + [(f"SFB[{j}]", pull(t, "SFB"), 512 * _sf_k4 * ((N + 127) // 128) * int(b_bufs[j].shape[0])) for j, t in enumerate(b.sfb_operands)] - ) - if _sf_reason is not None: - raise ValueError(_sf_reason) + out_bufs = [views[o.view] for o in recipe.outputs] + aux_bufs = [views[x.view] for x in recipe.aux] + c_arg = out_bufs if len(out_bufs) > 1 else out_bufs[0] + # Everything below indexes an operand's axes by position, so an operand + # that arrived in the graph's order is re-expressed in the kernel's here + # rather than threaded through every launcher. + operands = [(views[op.view] if ax is KERNEL_AXES else views[op.view].permute(ax[AX_BATCH], ax[AX_MN], ax[AX_K])) for op, ax in zip(recipe.inputs, axes)] + a_bufs, b_bufs = operands[: recipe.b_at], operands[recipe.b_at :] + sf_bufs = [views[s.view] for s in recipe.sf] if self.chain.is_multi_gemm: + sfa, sfb = sf_bufs[: recipe.b_at], sf_bufs[recipe.b_at :] if self.block_scale: - sfa = [pull(t, "SFA") for t in b.sfa_operands] - sfb = [pull(t, "SFB") for t in b.sfb_operands] pairs = [((a_bufs[ai], sfa[ai]), (b_bufs[bi], sfb[bi])) for ai, bi in self.chain.gemm_operands] else: pairs = [(a_bufs[ai], b_bufs[bi]) for ai, bi in self.chain.gemm_operands] - r = self._call_positional(pairs, c_arg, mnk, *aux_bufs, stream=stream) - _finalize_reductions(self.chain, out_bufs) - return r - - if self.block_scale: - sfa = pull(b.sfa_operands[0], "SFA") - sfb = pull(b.sfb_operands[0], "SFB") - r = self._call_positional(a_bufs[0], b_bufs[0], c_arg, mnk, sfa, sfb, *aux_bufs, stream=stream) - _finalize_reductions(self.chain, out_bufs) - return r - r = self._call_positional(a_bufs[0], b_bufs[0], c_arg, mnk, *aux_bufs, stream=stream) + args = (pairs, c_arg, mnk) + else: + args = (a_bufs[0], b_bufs[0], c_arg, mnk, *sf_bufs) + r = self._call_positional(*args, *aux_bufs, stream=stream) _finalize_reductions(self.chain, out_bufs) return r @@ -2362,76 +2402,15 @@ def _tma_alignment_reject( ("A", a_dtype, a_major, K if a_major == "k" else M, "K" if a_major == "k" else "M"), ("B", b_dtype, b_major, K if b_major == "k" else N, "K" if b_major == "k" else "N"), ): - bits = _dtype_bits(dtype) - if (extent * bits) % 128 != 0: - bad.append(f"{name} ({major}-major, {dtype}) requires {dim} % {128 // bits} == 0, " f"got {dim}={extent}") + modulus, pack = contiguous_modulus(dtype, major == "k") + logical = modulus * pack + if extent % logical: + bad.append(f"{name} ({major}-major, {dtype}) requires {dim} % {logical} == 0, " f"got {dim}={extent}") if bad: return "TMA input contiguous dimensions must be 16-byte aligned: " + "; ".join(bad) return None -# A/B are rank-3 (batch, M|N, K); the graph's major names which dim is contiguous. -_MAJOR_CONTIGUOUS_DIM = {"k": 2, "m": 1, "n": 1} - - -def _contiguous_dim(buf) -> "int | None": - """Index of the unambiguously contiguous dim of a rank-3 buffer, else None. - - A size-1 dim can carry stride 1 without meaning anything, so a layout is only - judged when exactly one dim of size > 1 is contiguous.""" - shape = tuple(buf.shape) - strides = tuple(buf.stride()) - unit = [i for i, s in enumerate(strides) if s == 1 and shape[i] > 1] - return unit[0] if len(unit) == 1 else None - - -def _operand_layout_reject(named_operands) -> "str | None": - """Each operand's runtime layout must match the major baked into the kernel. - - The major comes from the GRAPH's declared strides and is compiled into the - TMA descriptor and the MMA operand descriptor, while the launch reads the - RUNTIME buffer's strides — so a buffer whose contiguous dim disagrees with - the declaration computes silently wrong numbers rather than faulting. - Entries are ``(role, buffer, major)``. Returns a reason string, or ``None``.""" - bad = [] - for role, buf, major in named_operands: - if buf is None or len(tuple(buf.shape)) != 3: - continue - got = _contiguous_dim(buf) - want = _MAJOR_CONTIGUOUS_DIM[major] - if got is not None and got != want: - names = {0: "batch", 1: "M/N", 2: "K"} - bad.append( - f"{role}: graph declares {major}-major (dim {want} contiguous) but the buffer has dim {got} ({names[got]}) contiguous, stride={tuple(buf.stride())}" - ) - if bad: - return "runtime operand layout does not match the layout the kernel was compiled for: " + "; ".join(bad) - return None - - -def _operand_shape_reject(named_operands, M: int, N: int, K: int) -> "str | None": - """Every A operand must be (batch, M, K) and every B operand (batch, N, K), - for the SAME K — M/N/K are inferred from the first A and B buffer, so a - disagreeing operand would otherwise be read past its end (the kernel walks - the inferred K on every operand). Entries are ``(role, buffer, major, dim1, - k_pack)``; ``k_pack`` is 2 for FP4-packed data (two elements per byte-slot). - Returns a reason string, or ``None``.""" - bad = [] - for role, buf, _major, dim1, k_pack in named_operands: - if buf is None: - continue - shape = tuple(buf.shape) - if len(shape) != 3: - bad.append(f"{role}: expected a rank-3 buffer, got shape {shape}") - continue - want = (dim1, K // k_pack) - if (shape[1], shape[2]) != want: - bad.append(f"{role}: expected (batch, {want[0]}, {want[1]}), got {shape}") - if bad: - return f"runtime operand shapes disagree with the inferred problem size (M={M}, N={N}, K={K}): " + "; ".join(bad) - return None - - def _alignment_reject(named_buffers) -> "str | None": """Every runtime buffer's alignment must be >= the kernel's compiled requirement for its role. Entries are ``(role, buffer, required_bytes, diff --git a/python/cudnn/gemm/frost/dtypes.py b/python/cudnn/gemm/frost/dtypes.py index dfd9d0644..460792e9f 100644 --- a/python/cudnn/gemm/frost/dtypes.py +++ b/python/cudnn/gemm/frost/dtypes.py @@ -7,6 +7,7 @@ from __future__ import annotations +import functools from typing import Any import cudnn @@ -113,15 +114,25 @@ def tensor_alignment(shape, stride, elem_bytes: int, ptr: "int | None" = None, c one with ``stride==1 and shape!=1``; no such dim -> ``elem_bytes`` (nothing is contiguous, so only a single element can be moved at a time). """ - a = cap + a = _layout_alignment(tuple(shape), tuple(stride), elem_bytes, cap) if ptr is not None: a = min(a, _pow2_floor(int(ptr), cap)) + return a + + +@functools.lru_cache(maxsize=256) +def _layout_alignment(shape: tuple, stride: tuple, elem_bytes: int, cap: int) -> int: + """``min(A_stride, A_shape)`` -- the half the pointer does not enter. + Memoized because it is a function of values, not of objects: the same layout + recurs on every execute, and the set of shapes a caller cycles through under + dynamic shape is small. Only the pointer half is recomputed per call, and + that is one power-of-two floor. + """ stride_align = cap for sh, st in zip(shape, stride): if sh != 1 and st != 1: stride_align = min(stride_align, _pow2_floor(int(st) * elem_bytes, cap)) - a = min(a, stride_align) shape_align = None for sh, st in zip(shape, stride): @@ -130,8 +141,7 @@ def tensor_alignment(shape, stride, elem_bytes: int, ptr: "int | None" = None, c break if shape_align is None: shape_align = min(elem_bytes, cap) - a = min(a, shape_align) - return a + return min(stride_align, shape_align) def allowed_store_vsize(dim, stride, dtype: str) -> int: diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index a30382dfe..f9ffdde54 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -32,6 +32,12 @@ def __init__(self, compiled): # as ambiguous. self._tensors = list(compiled.binding.bound_tensors()) self._slots = None + # Which call path this plan uses is a property of the compiled kernel, + # so it is chosen here and not re-asked per execute. ``launch`` is the + # straight line the recipe lowers to when the kernel is a shape it + # emits; ``run_views`` serves everything else. + self._lowered = getattr(compiled, "launch", None) + self._run_views = self._lowered or getattr(compiled, "run_views", None) def get_workspace_size(self) -> int: return int(getattr(self._compiled, "workspace_bytes", 0) or 0) @@ -47,17 +53,19 @@ def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: # which carry the shape this execute runs, override_shapes included. views = variant_pack.views(slots) required = self.get_workspace_size() - # Scratch is carved from the CALLER's workspace: stable pointers, so a - # plan stays safe to capture in a CUDA graph. - extra = (Workspace.over(variant_pack, required, "frost_gemm"),) if required else () - run_resolved = getattr(self._compiled, "run_resolved", None) - if run_resolved is not None: + if required: + # Scratch is carved from the CALLER's workspace: stable pointers, so + # a plan stays safe to capture in a CUDA graph. Only the MoE + # launchers need it, and they take the variant-pack dict. + self._compiled(dict(zip(self._tensors, views)), Workspace.over(variant_pack, required, "frost_gemm"), stream=ctx.stream) + return + run_views = self._run_views + if run_views is not None: # Which bound tensor holds which operand was settled at build, so - # resolve_variant_pack's by-object / by-uid / by-name tables have - # no question left to answer. - run_resolved({id(t): v for t, v in zip(self._tensors, views)}, *extra, stream=ctx.stream) + # the buffers arrive in that order and the launcher indexes them. + run_views(views, stream=ctx.stream) else: - self._compiled(dict(zip(self._tensors, views)), *extra, stream=ctx.stream) + self._compiled(dict(zip(self._tensors, views)), stream=ctx.stream) class FrostGemmEngine(BaseEngine): diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py new file mode 100644 index 000000000..6a6323411 --- /dev/null +++ b/python/cudnn/gemm/frost/recipe.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""What one compiled gemm needs per call, settled when it compiles. + +Operand roles, majors, packing factors, alignment requirements, output shape +rules and which outputs are reductions are all fixed by the time cute hands back +a launchable. A call carries M, N, K, the strides and the pointers, and nothing +else. This module writes the first set down, so that neither the interpreted +path nor the straight line lowered from it re-derives them per call. + +The two consumers read the same recipe but do not share a body: :func:`gate` +interprets it, and ``CompiledFusedGemm._lower`` emits a straight line with the +constants inlined. That is a compiler beside its interpreter, kept honest the +way those always are -- ``test_execute_recipe.py`` runs both over the same +accepts and rejects and requires the same answer. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from .dtypes import DTYPE_BYTES, _aux_align_reqs, _output_align_reqs, _pow2_floor, tensor_alignment +from .fusion_ir import FusionChain + +# An operand's three axes, named by role rather than by position. +AX_BATCH, AX_MN, AX_K = 0, 1, 2 + +# cuDNN's matmul ABI declares A as [b, m, k] and B as [b, k, n], while a caller +# allocates B the way the kernel reads it -- (b, n, k). Both describe the same +# memory, so an operand arrives in one of exactly two axis orders and its +# strides alone do not say which: a graph-order K-major B and a caller-order +# N-major B carry stride 1 on the same axis. Its SHAPE does say, since only a +# buffer described from the graph's own declaration carries the declared dims. +_DECLARED_AXES = {"a": (0, 1, 2), "b": (0, 2, 1)} +KERNEL_AXES = (0, 1, 2) + +# How an output axis follows from the problem size. +CONST, FROM_M, FROM_N = 0, 1, 2 + +REDUCTION_INIT_VALUE = { + "fp32": { + "add": 0.0, + "amax": 0.0, + "max": -float("inf"), + "min": float("inf"), + "avg": 0.0, + "norm1": 0.0, + "norm2": 0.0, + "mul": 1.0, + "mul_no_zeros": 1.0, + }, + "int32": { + "add": 0, + "amax": 0, + "max": -(2**31), + "min": 2**31 - 1, + "norm1": 0, + }, +} + + +def contiguous_modulus(dtype: str, contiguous_is_k: bool) -> tuple: + """``(stored_modulus, packing)`` for one operand's TMA 16-byte rule. + + TMA encodes the contiguous dimension in 16-byte units, so the LOGICAL extent + must divide ``128 // bits``. A shape carries the STORED count, and fp4 packs + two elements per slot along K -- so a buffer is checked against the quotient + while a user wrote their shape against the product. Both the graph-time gate + and the per-call one read this, which is the one number that could drift. + """ + bits = 4 if dtype == "fp4_e2m1" else DTYPE_BYTES[dtype] * 8 + pack = 2 if (dtype == "fp4_e2m1" and contiguous_is_k) else 1 + return (128 // bits) // pack, pack + + +def expected_shape(rule, m: int, n: int) -> tuple: + """One output's runtime shape, from the rule the build recorded.""" + return tuple(v if s == CONST else (m if s == FROM_M else n // v) for s, v in rule) + + +def _output_rule(spec, chain: FusionChain) -> tuple: + """How one output's shape follows from (M, N), as a per-axis rule. + + The single answer to a question that used to be asked in two places and + disagreed: an fp4 dense output is (batch, M, N/2), and a hand-written launch + path that pinned N instead rejected a legal call. + """ + batch = chain.matmul.batch + if spec.is_quant_scale: + return tuple((CONST, int(d)) for d in spec.dim) + if not spec.is_reduction: + return ((CONST, batch), (FROM_M, 0), (FROM_N, 2 if spec.dtype == "fp4_e2m1" else 1)) + red_idx = int(spec.source.rsplit("_", 1)[1]) + if chain.reductions[red_idx].grouped_by_moe: + return tuple((CONST, int(d)) for d in spec.dim) + # A reduced axis collapses to 1; the rest follow the problem size. + return ( + (CONST, 1 if spec.dim[0] == 1 else batch), + (CONST, 1) if spec.dim[1] == 1 else (FROM_M, 0), + (CONST, 1) if spec.dim[2] == 1 else (FROM_N, 1), + ) + + +@dataclass(frozen=True) +class Operand: + """One A or B operand: what a call must satisfy, with the rest baked.""" + + view: int # position in the view list the engine hands over + role: str # names this operand in a rejection message + major: str + kpack: int # 2 when fp4 stores two elements per slot + batch: int # the extent the graph declares on the batch axis + is_b: bool # its non-K extent is N rather than M + modulus: int # the contiguous STORED extent must divide this + pack: int # ... and multiplying by this recovers the logical extent + contiguous_role: int # AX_K for a k-major operand, else AX_MN + declared: tuple # (batch, mn, k) axis positions, graph order + declared_dim: tuple # the extents that order comes with + dc: int # where stride 1 lands in the graph's order + kc: int # ... and in the caller's + + def axes(self, shape, stride) -> "tuple | None": + """Which axis holds which role, or None if neither order fits.""" + if stride[self.kc] == 1: + return KERNEL_AXES + if stride[self.dc] == 1 and tuple(shape) == self.declared_dim: + return self.declared + return None + + +@dataclass(frozen=True) +class Output: + view: int + role: str + rule: tuple + align: int + raw: bool # the kernel takes only its pointer + init: Any = None # reduction identity, seeded before the kernel runs + sqrt: bool = False # norm2 takes a square root after + + +@dataclass(frozen=True) +class Aux: + view: int + role: str + align: int + ref: Any # TensorRef, for the fake-shape reshape + + +@dataclass(frozen=True) +class ScaleFactor: + view: int + role: str + is_a: bool + operand_at: int # the A or B operand this one scales, as an index into inputs + + +@dataclass(frozen=True) +class GemmRecipe: + inputs: tuple # every A operand then every B operand, in kernel slot order + outputs: tuple + aux: tuple + sf: tuple + roles: tuple # what occupies each view position, for a missing-buffer message + a_at: int # M and K are read off inputs[a_at] + b_at: int # N off inputs[b_at] + has_output_specs: bool + block_size: "int | None" + workspace_bytes: int + multi_gemm: bool + + @property + def a(self) -> Operand: + return self.inputs[self.a_at] + + @property + def b(self) -> Operand: + return self.inputs[self.b_at] + + def problem(self, views) -> tuple: + """``((M, N, K), axes-per-input)``, located rather than assumed. + + Reading M off axis 1 assumes the caller laid the buffer out the way the + kernel reads it, which a bare device address does not: the pack + describes that one from the graph's declaration, which orders B the + other way round. + """ + axes, bad = [], [] + for op in self.inputs: + v = views[op.view] + ax = op.axes(v.shape, v.stride()) + if ax is None: + bad.append( + f"{op.role}: graph declares {op.major}-major (dim {op.kc} contiguous) but the buffer has shape={tuple(v.shape)}, stride={tuple(v.stride())}" + ) + ax = KERNEL_AXES + axes.append(ax) + if bad: + raise ValueError("runtime operand layout does not match the layout the kernel was compiled for: " + "; ".join(bad)) + a, b = self.a, self.b + a_ax, b_ax = axes[self.a_at], axes[self.b_at] + a_shape, b_shape = views[a.view].shape, views[b.view].shape + return (a_shape[a_ax[AX_MN]], b_shape[b_ax[AX_MN]], a_shape[a_ax[AX_K]] * a.kpack), tuple(axes) + + +def _tma_reject(recipe: GemmRecipe, views, axes) -> "str | None": + """TMA encodes the contiguous dimension in 16-byte units; a misaligned + extent silently mis-strides every row past the first.""" + bad = [] + for op, ax in zip(recipe.inputs, axes): + role = op.contiguous_role + extent = views[op.view].shape[ax[role]] + if extent % op.modulus: + # Report the LOGICAL extent and modulus: fp4 stores two elements per + # slot, so "K % 32" is the rule a user wrote their shape against. + name = "K" if role == AX_K else ("N" if op.is_b else "M") + bad.append(f"{op.role} ({op.major}-major) requires {name} % {op.modulus * op.pack} == 0, got {name}={extent * op.pack}") + if bad: + return "TMA input contiguous dimensions must be 16-byte aligned: " + "; ".join(bad) + return None + + +def _shape_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": + """Every operand must agree with the M/N/K read off the first A and B -- + the kernel walks the inferred K on all of them.""" + m, n, k = mnk + bad = [] + for op, ax in zip(recipe.inputs, axes): + shape = views[op.view].shape + want = (op.batch, n if op.is_b else m, k // op.kpack) + got = (shape[ax[AX_BATCH]], shape[ax[AX_MN]], shape[ax[AX_K]]) + if got != want: + bad.append(f"{op.role}: expected (batch, M|N, K) = {want}, got {got}") + if bad: + return f"runtime operand shapes disagree with the inferred problem size (M={m}, N={n}, K={k}): " + "; ".join(bad) + return None + + +def _output_shape_reject(recipe: GemmRecipe, views, mnk) -> "str | None": + m, n, _ = mnk + bad = [] + for out in recipe.outputs: + shape = tuple(views[out.view].shape) + want = expected_shape(out.rule, m, n) + if shape != want: + bad.append(f"{out.role}: expected {want}, got {shape}") + if bad: + return "runtime tensors must be rank-3 with shapes matching the graph: " + "; ".join(bad) + return None + + +def _align_reject(recipe: GemmRecipe, views) -> "str | None": + """Every buffer's alignment must meet the width its role's access uses. + + Inputs and scale factors are TMA-loaded, so only the base pointer is at + stake; outputs and aux are stored and loaded directly, so their strides and + contiguous extent bound the vector too. + """ + bad = [] + for item in recipe.inputs + recipe.sf: + ptr = int(views[item.view].data_ptr()) + align = _pow2_floor(ptr) + if align < 16: + bad.append(f"{item.role}: alignment {align}B < required 16B (ptr=0x{ptr:x})") + for item in recipe.outputs + recipe.aux: + v = views[item.view] + ptr = int(v.data_ptr()) + align = tensor_alignment(tuple(v.shape), tuple(v.stride()), v.element_size(), ptr=ptr) + if align < item.align: + bad.append(f"{item.role}: alignment {align}B < required {item.align}B (ptr=0x{ptr:x})") + if bad: + return "runtime tensor alignment is below the kernel's compiled requirement: " + "; ".join(bad) + return None + + +def _sf_blob_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": + """A block-scale SF reaches the kernel as a base pointer plus a layout the + template re-synthesizes from M/N/K, so a blob that is not one dense byte run + of at least the required size is read out of bounds with no fault.""" + m, n, k = mnk + k4 = ((k // recipe.block_size) + 3) // 4 + bad = [] + for sf in recipe.sf: + v = views[sf.view] + op = recipe.inputs[sf.operand_at] + rows = m if sf.is_a else n + batch = views[op.view].shape[axes[sf.operand_at][AX_BATCH]] + required = 512 * k4 * ((rows + 127) // 128) * int(batch) + span = 1 + sum((int(s) - 1) * int(st) for s, st in zip(v.shape, v.stride())) + if int(v.numel()) != span: + bad.append(f"{sf.role} shape {tuple(v.shape)} stride {tuple(v.stride())} is not a dense byte run") + continue + have = int(v.numel()) * v.element_size() + if have < required: + bad.append( + f"{sf.role} is {have}B but the kernel reads {required}B ({required // 512} atoms of 128 rows x 4 SF-K) — was it produced by to_blocked()?" + ) + if bad: + return "block-scale F8_128x4 scale factors must be a packed blob: " + "; ".join(bad) + return None + + +def gate(recipe: GemmRecipe, views, mnk, axes) -> None: + """Raise on anything about this call the compiled kernel cannot serve.""" + reasons = [ + _tma_reject(recipe, views, axes), + _shape_reject(recipe, views, axes, mnk), + _output_shape_reject(recipe, views, mnk), + _align_reject(recipe, views), + ] + if recipe.block_size: + reasons.append(_sf_blob_reject(recipe, views, axes, mnk)) + for reason in reasons: + if reason is not None: + raise ValueError(reason) + + +def _declared_dim(tensor) -> tuple: + """The dims the graph declared, or ``()`` when the tensor cannot say. + + An empty tuple never matches a runtime shape, so an operand that cannot + report its declaration is simply read in the caller's order -- which is + what every path did before this table existed. + """ + try: + return tuple(int(d) for d in tensor.get_dim()) + except Exception: # noqa: BLE001 — an analyzer-synthesized ref has no dims + return () + + +def _operand(view: int, role: str, tensor, *, major: str, dtype: str, batch: int, is_b: bool) -> Operand: + declared = _DECLARED_AXES["b" if is_b else "a"] + contiguous = AX_K if major == "k" else AX_MN + modulus, pack = contiguous_modulus(dtype, contiguous == AX_K) + return Operand( + view=view, + role=role, + major=major, + kpack=2 if dtype == "fp4_e2m1" else 1, + batch=batch, + is_b=is_b, + modulus=modulus, + pack=pack, + contiguous_role=contiguous, + declared=declared, + declared_dim=_declared_dim(tensor), + dc=declared[contiguous], + kc=KERNEL_AXES[contiguous], + ) + + +def build(compiled) -> GemmRecipe: + """Read one compiled gemm into the table its call path runs off.""" + chain: FusionChain = compiled.chain + mm = chain.matmul + binding = compiled.binding + order = {} + for i, t in enumerate(binding.bound_tensors()): + order.setdefault(id(t), i) + + inputs = [ + _operand(order[id(t)], f"A operand[{i}]", t, major=mm.a_major, dtype=mm.a_dtype, batch=mm.a_batch, is_b=False) for i, t in enumerate(binding.a_operands) + ] + [ + _operand(order[id(t)], f"B operand[{i}]", t, major=mm.b_major, dtype=mm.b_dtype, batch=mm.b_batch, is_b=True) for i, t in enumerate(binding.b_operands) + ] + + out_reqs = _output_align_reqs(chain, compiled.use_tma_store, vec_bytes=compiled.vec_bytes_epi) + aux_reqs = _aux_align_reqs(chain, vec_bytes=compiled.vec_bytes_epi) + outputs = [] + for i, (spec, t) in enumerate(zip(chain.outputs, binding.outputs)): + init, sqrt = None, False + if spec.is_reduction: + red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] + init = REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] + sqrt = red.mode == "norm2" + outputs.append( + Output( + view=order[id(t)], + role=spec.source, + rule=_output_rule(spec, chain), + align=out_reqs[i], + raw=bool(spec.is_reduction or spec.is_quant_scale), + init=init, + sqrt=sqrt, + ) + ) + + aux = tuple( + Aux(view=order[id(t)], role=f"aux {compiled.aux_names[i]!r}", align=aux_reqs[compiled.aux_names[i]], ref=ref) + for i, (t, ref) in enumerate(zip(binding.aux, chain.aux_tensors)) + ) + na = len(binding.a_operands) + sf = tuple( + [ScaleFactor(view=order[id(t)], role=f"SFA[{i}]", is_a=True, operand_at=i) for i, t in enumerate(binding.sfa_operands)] + + [ScaleFactor(view=order[id(t)], role=f"SFB[{j}]", is_a=False, operand_at=na + j) for j, t in enumerate(binding.sfb_operands)] + ) + + roles = [f"bound tensor {i}" for i in range(len(binding.bound_tensors()))] + for item in (*inputs, *outputs, *aux, *sf): + roles[item.view] = item.role + + return GemmRecipe( + inputs=tuple(inputs), + outputs=tuple(outputs), + aux=aux, + sf=sf, + roles=tuple(roles), + a_at=0, + b_at=na, + has_output_specs=bool(chain.output_specs), + block_size=chain.block_scale.block_size if compiled.block_scale else None, + workspace_bytes=int(getattr(compiled, "workspace_bytes", 0) or 0), + multi_gemm=bool(chain.is_multi_gemm), + ) diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py new file mode 100644 index 000000000..6ede445b7 --- /dev/null +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The build-time recipe, and the straight line lowered from it. + +``_lower`` emits a call path with the recipe's constants inlined; ``run_views`` +interprets the same recipe. That is a compiler beside its interpreter, and the +two can drift -- an earlier hand-written version of the emitted path lost the +operand batch check and pinned an fp4 output at N instead of N/2, both of which +made it accept or reject calls the general path did not. + +So the differential below is the point of this file: every case runs through +BOTH entry points and the two must return the same verdict and the same numbers. +The rest asserts that the fast path is the one a public ``execute()`` actually +takes, and that the flavors it declines are declined on purpose. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from gemm_test_utils import requires_sm100, vp + +import cudnn +import cudnn.gemm.frost # noqa: F401 — installs the cudnn.pygraph recorder hook +from cudnn.engines import is_python_engine +from cudnn.gemm.frost.compiler import jit_from_cudnn_graph +from cudnn.gemm.frost.graph_analyzer import resolve_variant_pack +from cudnn.gemm.frost.recipe import _output_rule, expected_shape + +pytestmark = pytest.mark.L0 + +BF16 = cudnn.data_type.BFLOAT16 +F32 = cudnn.data_type.FLOAT +M = N = 256 +K = 128 + + +# --- the output shape rule, which is pure data ------------------------------ + + +def _spec(**kw): + base = dict(source="matmul", dtype="bf16", dim=None, is_reduction=False, is_quant_scale=False) + base.update(kw) + return SimpleNamespace(**base) + + +def _chain(batch=1, reductions=()): + return SimpleNamespace(matmul=SimpleNamespace(batch=batch), reductions=list(reductions)) + + +def test_dense_output_follows_m_and_n(): + assert expected_shape(_output_rule(_spec(), _chain(batch=3)), 128, 256) == (3, 128, 256) + + +def test_fp4_dense_output_is_half_as_wide(): + """fp4 packs two elements per stored slot, so the last axis is N/2. + + The hand-written launch path this replaces pinned it at N, which rejected a + legal call. One rule, read by both paths, is the answer to that. + """ + assert expected_shape(_output_rule(_spec(dtype="fp4_e2m1"), _chain()), 128, 256) == (1, 128, 128) + + +def test_a_reduced_axis_collapses_to_one(): + chain = _chain(batch=2, reductions=[SimpleNamespace(grouped_by_moe=False)]) + rule = _output_rule(_spec(source="reduction_0", is_reduction=True, dim=(1, 128, 1)), chain) + assert expected_shape(rule, 512, 256) == (1, 512, 1) + + +def test_a_quant_scale_output_is_fixed_at_build(): + rule = _output_rule(_spec(source="quant_scale_0", is_quant_scale=True, dim=(1, 128, 4)), _chain()) + assert expected_shape(rule, 999, 999) == (1, 128, 4) + + +# --- which flavors lower, and which decline --------------------------------- + + +def _plain_graph(m=M, n=N, k=K, batch=1): + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[batch, m, k], stride=[m * k, k, 1]) + B = g.tensor(name="B", dim=[batch, k, n], stride=[k * n, 1, k]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + return g + + +def _reduction_graph(): + g = _plain_graph() + Y = [t for t in g._nodes[-1].outputs.values()][0] + R = g.reduction(input=Y, mode=cudnn.reduction_mode.ADD, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]).set_output(True).set_data_type(F32) + return g + + +def _aux_graph(): + 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=F32) + C = g.matmul(A=A, B=B, name="mm") + Y = g.add(a=C, b=bias, name="bias_add") + Y.set_output(True).set_data_type(BF16) + return g + + +@requires_sm100 +@pytest.mark.parametrize( + "build,lowered", + ( + (_plain_graph, True), + (_reduction_graph, False), + (_aux_graph, False), + ), + ids=("plain", "reduction", "aux"), +) +def test_which_flavors_lower(build, lowered): + """A declined flavor is declined by a named recipe field, not by accident. + + Reductions want a pre-kernel seed and aux wants a fake-shape reshape; both + are work the emitted line does not carry yet, so it hands them over rather + than growing a branch per flavor. + """ + compiled = jit_from_cudnn_graph(build()) + assert (compiled.launch is not None) is lowered + + +@requires_sm100 +def test_public_execute_takes_the_lowered_path(): + """The straight line is what a user's ``execute()`` runs, not a side door.""" + g = _plain_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() + assert g._compiled_plans[g._plan_index]._lowered is not None + + +# --- the differential ------------------------------------------------------- + + +def _views(compiled, a, b, c): + resolved = resolve_variant_pack(vp(compiled, a, b, c), compiled.binding) + return [resolved[id(t)] for t in compiled.binding.bound_tensors()] + + +def _verdict(run, views, c): + c.zero_() + try: + run(views, stream=None) + except ValueError: + return "rejected", None + torch.cuda.synchronize() + return "ran", c.clone() + + +def _operands(m=M, n=N, k=K, batch=1): + a = torch.randn(batch, m, k, dtype=torch.bfloat16, device="cuda") + b = torch.randn(batch, n, k, dtype=torch.bfloat16, device="cuda") + c = torch.empty(batch, m, n, dtype=torch.bfloat16, device="cuda") + return a, b, c + + +def _good(): + return _operands() + + +def _short_k(): + a, _, c = _operands() + return a, torch.randn(1, N, K // 2, dtype=torch.bfloat16, device="cuda"), c + + +def _wrong_major(): + a, _, c = _operands() + return a, torch.randn(1, K, N, dtype=torch.bfloat16, device="cuda").transpose(1, 2), c + + +def _wrong_output_shape(): + a, b, _ = _operands() + return a, b, torch.empty(1, M, N // 2, dtype=torch.bfloat16, device="cuda") + + +def _misaligned_output(): + a, b, _ = _operands() + # One element in is one element off every alignment the epilogue wants. + return a, b, torch.empty(1, M, N + 8, dtype=torch.bfloat16, device="cuda")[:, :, 1 : N + 1] + + +def _misaligned_a(): + """One bf16 element in is 2 bytes in, and TMA wants a 16-byte base.""" + _, b, c = _operands() + wide = torch.randn(1, M, K + 8, dtype=torch.bfloat16, device="cuda") + return wide[:, :, 1 : K + 1], b, c + + +def _padded_rows(): + """Legal: the outer stride is free, only the contiguous extent is pinned.""" + a, b, c = _operands() + return a, b, torch.empty(1, M, N * 2, dtype=torch.bfloat16, device="cuda")[:, :, :N] + + +@requires_sm100 +@pytest.mark.parametrize( + "case", + (_good, _short_k, _wrong_major, _wrong_output_shape, _misaligned_output, _misaligned_a, _padded_rows), + ids=lambda f: f.__name__.strip("_"), +) +def test_lowered_and_interpreted_agree(case): + """Same operands, both entry points, one verdict. + + An unsound fast path shows up here as an accept where the general path + rejects -- which is exactly how the two regressions this replaced would have + read. + """ + compiled = jit_from_cudnn_graph(_plain_graph()) + if compiled.launch is None: + pytest.skip("this build does not lower (no tvm-ffi front door)") + + a, b, c = case() + fast, fast_out = _verdict(compiled.launch, _views(compiled, a, b, c), c) + slow, slow_out = _verdict(compiled.run_views, _views(compiled, a, b, c), c) + assert fast == slow, f"lowered says {fast}, interpreted says {slow}" + if fast == "ran": + torch.testing.assert_close(fast_out, slow_out, atol=0, rtol=0) + ref = torch.einsum("bmk,bnk->bmn", a.float(), b.float()) + torch.testing.assert_close(fast_out.float(), ref, atol=2e-1, rtol=2e-2) + + +@requires_sm100 +def test_a_graph_order_operand_falls_back_and_still_runs(): + """A bare device address is described from the graph, which orders B + ``[batch, K, N]``. The emitted line only serves the caller's order, so it + hands this one over rather than reading the wrong axis.""" + compiled = jit_from_cudnn_graph(_plain_graph()) + a, b, c = _operands() + graph_order_b = b.transpose(1, 2) # (1, K, N) over the same memory + views = _views(compiled, a, graph_order_b, c) + compiled.run_views(views, stream=None) + torch.cuda.synchronize() + torch.testing.assert_close(c.float(), torch.einsum("bmk,bnk->bmn", a.float(), b.float()), atol=2e-1, rtol=2e-2) + + +@requires_sm100 +def test_operand_batch_is_checked(): + """The graph pins each operand's batch, and a launch that ignored it read + one batch of A against three of B.""" + compiled = jit_from_cudnn_graph(_plain_graph(batch=2)) + a, b, c = _operands(batch=2) + one_batch_b = b[:1].contiguous() + for run in (compiled.launch, compiled.run_views): + if run is None: + continue + with pytest.raises(ValueError): + run(_views(compiled, a, one_batch_b, c), stream=None) diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index 068e159c4..ae994c2b8 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -190,18 +190,14 @@ def test_norm2_reduction_is_refused_at_build(): @_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.""" + """The backend has always taken a raw device address; so must a python plan. + + The operand a bare address describes is the GRAPH's declaration, which + orders a matmul's B ``[batch, K, N]`` where a caller allocates it + ``(batch, N, K)``. Reading an extent by axis position answers one of those + and not the other, so the recipe records which axis carries M/N/K instead. + """ 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]) From 6d20caed107ab6042184c1584ebbfb63e7e05f77 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 01:05:31 -0700 Subject: [PATCH 2/9] Take an operand's axis order from the graph, not from a guess at its shape A matmul's B arrives either as cuDNN's declared [b, K, N] or as this engine's own direct-call (b, N, K), and at N == K those are the same shape AND the same stride. Inferring which from the description read one as the other and computed a transpose, silently -- and both call paths read that inference, so the differential between them agreed and stayed green. The tie-break is the backend's own rule: the graph's tensor descriptor defines the tensor and the variant pack supplies only a pointer. So a buffer that reports the declared (dim, stride) is read as the declaration. Measured on the ambiguous case: that agrees with the backend to bf16 tolerance, where reading it as (b, N, K) differs by 65. VariantPack.graph_described names the slots the pack described FROM the graph, which a bare address's live shape cannot show once override_shapes has moved it. That closes the other half: a bare pointer running a smaller problem inside its allocation was refused as a layout mismatch. This also fixes a pre-existing divergence, not introduced here: test_every_variant_pack_form_still_works fails all four forms on develop with CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1, because it builds its graph with tensor_like() on a [b, K, N] tensor -- the matmul ABI, literally -- and the engine assumed its own order. Same call, two answers, decided by plan selection, which is what that test exists to deny. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/_pygraph.py | 11 ++- python/cudnn/engines/base.py | 9 +- python/cudnn/gemm/frost/compiler.py | 21 ++++- python/cudnn/gemm/frost/engine.py | 11 ++- python/cudnn/gemm/frost/recipe.py | 69 ++++++++------ test/python/gemm/frost/test_execute_recipe.py | 91 +++++++++++++++++-- 6 files changed, 168 insertions(+), 44 deletions(-) diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 39a7011d1..949e06db0 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -1882,10 +1882,19 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= # caller's mistake. A python-only graph's layout is every wired port, # including optional ones, where a hole means "not requested". strict = self._lowered_graph is not None + from_graph = [] for i in unread: data = uid_to_data.get(order[i]) if data is None: continue # named below if this graph requires it + if type(data) is int: + # A bare address has no geometry of its own, so _describe lends + # it the graph's -- including the graph's AXIS ORDER, which for + # a matmul's B is [batch, K, N] where a caller allocates + # (batch, N, K). Nothing in the resulting description says which + # of the two it is (at N == K the two are bit-identical), so the + # slot that borrowed one is named here. + from_graph.append(i) ptr, tensor = self._describe(data, order[i]) native.set_slot(i, ptr, tuple(tensor.dim), tuple(tensor.stride), *_dlpack_code_bits(tensor.data_type)) if strict: @@ -1924,7 +1933,7 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= workspace_bytes = _byte_size(workspace_tensor) else: workspace_ptr, workspace_bytes = extent - return VariantPack(tuple(order), native, workspace_ptr, workspace_bytes) + return VariantPack(tuple(order), native, workspace_ptr, workspace_bytes, tuple(from_graph)) def _describe(self, data: Any, uid: int): """``(pointer, Tensor)`` for one caller buffer. diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 907c93fa1..532fbd8af 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -133,13 +133,18 @@ class VariantPack: pointers — silently, because every pointer in it is individually valid. """ - __slots__ = ("uids", "native", "_slot_of", "workspace", "workspace_bytes", "_device") + __slots__ = ("uids", "native", "_slot_of", "workspace", "workspace_bytes", "_device", "graph_described") - def __init__(self, uids, native, workspace_ptr: int = 0, workspace_bytes: int = 0): + def __init__(self, uids, native, workspace_ptr: int = 0, workspace_bytes: int = 0, graph_described=()): self.uids = uids self.native = native self.workspace = workspace_ptr self.workspace_bytes = workspace_bytes + # Slots whose dim/stride were lent by the graph because the caller + # passed a bare address. Usually empty. An engine that reads extents by + # axis position needs this: the graph and the caller order a matmul's B + # differently, and the description does not say which one it is. + self.graph_described = graph_described self._slot_of = None # built on first lookup: the backend never does one self._device = None diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 2c82bc0c3..820a65932 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1921,6 +1921,14 @@ def _lower(self): return None a, b = r.a, r.b + # A buffer reporting the declaration is read in the graph's axis order, + # which this body does not serve. The stride guard below tells that + # apart -- unless the declaration itself would satisfy the guard, which + # only a unit extent can arrange, and which is settled here not per call. + for op in (a, b): + declared_stride = op.declared_layout[1] + if op.dc != op.kc and declared_stride and declared_stride[op.kc] == 1: + return None ai, bi, ci = a.view, b.view, out.view # kc is both the axis whose stride must be 1 and the axis whose extent # enters the TMA rule -- they are the same axis by definition of major. @@ -1936,8 +1944,12 @@ def _lower(self): launchable = self._launchable general = self.run_views - def launch(views, stream=None): + def launch(views, graph_order=None, stream=None): _check_plan_device(device) + if graph_order: + # An operand described from the graph is in the graph's axis + # order; this body reads the caller's. + return general(views, graph_order, stream=stream) av, bv, cv = views[ai], views[bi], views[ci] a_sh, b_sh, c_sh = av.shape, bv.shape, cv.shape if len(a_sh) != 3 or len(b_sh) != 3 or len(c_sh) != 3: @@ -1983,17 +1995,20 @@ def launch(views, stream=None): return launch - def run_views(self, views, stream=None): + def run_views(self, views, graph_order=None, stream=None): """Launch over the operand buffers, in bound-tensor order. Which bound tensor holds which operand is fixed when the plan compiles, so a caller that knows it -- the engine, which binds the graph's slots once -- indexes rather than resolves, and the gate below reads a table built at that same moment instead of rebuilding one per call. + + ``graph_order`` is the pack's per-view "this operand's layout is the + graph's, not the caller's", or None when they are all the caller's. """ _check_plan_device(self.device) recipe = self.recipe - mnk, axes = recipe.problem(views) + mnk, axes = recipe.problem(views, graph_order) gate_recipe(recipe, views, mnk, axes) out_bufs = [views[o.view] for o in recipe.outputs] diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index f9ffdde54..33a438b39 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -63,7 +63,16 @@ def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: if run_views is not None: # Which bound tensor holds which operand was settled at build, so # the buffers arrive in that order and the launcher indexes them. - run_views(views, stream=ctx.stream) + # Which AXIS ORDER each one arrived in is a per-call fact only the + # pack knows, since a bare address wears the graph's layout. None + # means every operand here is the caller's own, which is what both + # launchers are written for. + graph_order = None + borrowed = variant_pack.graph_described + if borrowed: + flags = tuple(s in borrowed for s in slots) + graph_order = flags if any(flags) else None + run_views(views, graph_order, stream=ctx.stream) else: self._compiled(dict(zip(self._tensors, views)), stream=ctx.stream) diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py index 6a6323411..98e1f92de 100644 --- a/python/cudnn/gemm/frost/recipe.py +++ b/python/cudnn/gemm/frost/recipe.py @@ -27,12 +27,19 @@ # An operand's three axes, named by role rather than by position. AX_BATCH, AX_MN, AX_K = 0, 1, 2 -# cuDNN's matmul ABI declares A as [b, m, k] and B as [b, k, n], while a caller -# allocates B the way the kernel reads it -- (b, n, k). Both describe the same -# memory, so an operand arrives in one of exactly two axis orders and its -# strides alone do not say which: a graph-order K-major B and a caller-order -# N-major B carry stride 1 on the same axis. Its SHAPE does say, since only a -# buffer described from the graph's own declaration carries the declared dims. +# cuDNN's matmul ABI declares A as [b, m, k] and B as [b, k, n]; this engine's +# own direct-call API takes B the way the kernel reads it, (b, n, k). Both +# describe the same memory, so an operand arrives in one of two axis orders and +# the description alone cannot always say which -- at N == K the two are +# identical tuples. +# +# The tie-break is the backend's own rule: a graph's tensor descriptor DEFINES +# the tensor and the variant pack supplies only a pointer. So a buffer whose +# (shape, stride) is the declaration is read as the declaration, which is what +# the backend computes from it (measured: a matmul whose B matches the declared +# [b, K, N] agrees with the backend to bf16 tolerance, and differs from the +# (b, N, K) reading by 65). Anything else is the caller's own labelling of the +# same memory, which is the direct-call order. _DECLARED_AXES = {"a": (0, 1, 2), "b": (0, 2, 1)} KERNEL_AXES = (0, 1, 2) @@ -117,17 +124,23 @@ class Operand: pack: int # ... and multiplying by this recovers the logical extent contiguous_role: int # AX_K for a k-major operand, else AX_MN declared: tuple # (batch, mn, k) axis positions, graph order - declared_dim: tuple # the extents that order comes with + declared_layout: tuple # the (dim, stride) that order comes with dc: int # where stride 1 lands in the graph's order kc: int # ... and in the caller's - def axes(self, shape, stride) -> "tuple | None": - """Which axis holds which role, or None if neither order fits.""" - if stride[self.kc] == 1: - return KERNEL_AXES - if stride[self.dc] == 1 and tuple(shape) == self.declared_dim: - return self.declared - return None + def axes(self, shape, stride, graph_order: bool) -> "tuple | None": + """Which axis holds which role, or None when the layout is not the major. + + ``graph_order`` says the pack described this slot FROM the graph (a bare + address has no geometry of its own), so it is the declaration by + construction. A buffer that reports the declaration is read as the + declaration too -- that is what the backend computes from it, and a + caller must not get a different answer for having landed on a python + plan. Everything else is the caller's own labelling of the same memory. + """ + if graph_order or (tuple(shape), tuple(stride)) == self.declared_layout: + return self.declared if stride[self.dc] == 1 else None + return KERNEL_AXES if stride[self.kc] == 1 else None @dataclass(frozen=True) @@ -179,21 +192,24 @@ def a(self) -> Operand: def b(self) -> Operand: return self.inputs[self.b_at] - def problem(self, views) -> tuple: + def problem(self, views, graph_order=None) -> tuple: """``((M, N, K), axes-per-input)``, located rather than assumed. Reading M off axis 1 assumes the caller laid the buffer out the way the kernel reads it, which a bare device address does not: the pack describes that one from the graph's declaration, which orders B the - other way round. + other way round. ``graph_order`` is the pack's per-view answer to which + it was, or None when every operand is the caller's own. """ axes, bad = [], [] for op in self.inputs: v = views[op.view] - ax = op.axes(v.shape, v.stride()) + borrowed = bool(graph_order and graph_order[op.view]) + ax = op.axes(v.shape, v.stride(), borrowed) if ax is None: + want = op.dc if borrowed else op.kc bad.append( - f"{op.role}: graph declares {op.major}-major (dim {op.kc} contiguous) but the buffer has shape={tuple(v.shape)}, stride={tuple(v.stride())}" + f"{op.role}: graph declares {op.major}-major (dim {want} contiguous) but the buffer has shape={tuple(v.shape)}, stride={tuple(v.stride())}" ) ax = KERNEL_AXES axes.append(ax) @@ -317,17 +333,16 @@ def gate(recipe: GemmRecipe, views, mnk, axes) -> None: raise ValueError(reason) -def _declared_dim(tensor) -> tuple: - """The dims the graph declared, or ``()`` when the tensor cannot say. +def _declared_layout(tensor) -> tuple: + """The ``(dim, stride)`` the graph declared, or ``((), ())`` when it cannot say. - An empty tuple never matches a runtime shape, so an operand that cannot - report its declaration is simply read in the caller's order -- which is - what every path did before this table existed. + An empty pair never equals a runtime description, so an operand whose + declaration is unreadable is simply read in the direct-call order. """ try: - return tuple(int(d) for d in tensor.get_dim()) - except Exception: # noqa: BLE001 — an analyzer-synthesized ref has no dims - return () + return tuple(int(d) for d in tensor.get_dim()), tuple(int(s) for s in tensor.get_stride()) + except Exception: # noqa: BLE001 -- an analyzer-synthesized ref has no dims + return (), () def _operand(view: int, role: str, tensor, *, major: str, dtype: str, batch: int, is_b: bool) -> Operand: @@ -344,8 +359,8 @@ def _operand(view: int, role: str, tensor, *, major: str, dtype: str, batch: int modulus=modulus, pack=pack, contiguous_role=contiguous, + declared_layout=_declared_layout(tensor), declared=declared, - declared_dim=_declared_dim(tensor), dc=declared[contiguous], kc=KERNEL_AXES[contiguous], ) diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index 6ede445b7..86534aced 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -234,17 +234,88 @@ def test_lowered_and_interpreted_agree(case): @requires_sm100 -def test_a_graph_order_operand_falls_back_and_still_runs(): - """A bare device address is described from the graph, which orders B - ``[batch, K, N]``. The emitted line only serves the caller's order, so it - hands this one over rather than reading the wrong axis.""" - compiled = jit_from_cudnn_graph(_plain_graph()) - a, b, c = _operands() - graph_order_b = b.transpose(1, 2) # (1, K, N) over the same memory - views = _views(compiled, a, graph_order_b, c) - compiled.run_views(views, stream=None) +def test_an_operand_reporting_the_declaration_agrees_with_the_backend(monkeypatch): + """At N == K the two axis orders are the SAME shape and the SAME stride. + + The graph declares B ``[batch, K, N]`` k-major, stride ``[K*N, 1, K]``; this + engine's direct-call API takes ``(batch, N, K)``, stride ``(K*N, 1, N)``. + Identical tuples when N == K, so the description cannot say which the caller + meant -- and the backend does not ask, it computes from the descriptor. A + python plan has to reach the same answer, because the caller does not choose + which plan the heuristics land on. + """ + monkeypatch.setenv("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "1") + torch.manual_seed(0) + d = 128 + a = torch.randn(1, d, d, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, d, d, dtype=torch.bfloat16, device="cuda").transpose(1, 2) + assert (tuple(b.shape), tuple(b.stride())) == ((1, d, d), (d * d, 1, d)) # == the declaration + + out = {} + for want_frost in (False, True): + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, d, d], stride=[d * d, d, 1]) + B = g.tensor(name="B", dim=[1, d, d], stride=[d * d, 1, d]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + hits = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id) == want_frost] + if not hits: + pytest.skip("no plan of the requested kind for this graph") + g.select_plan(hits[0]) + g.check_support() + g.build_plans() + c = torch.zeros(1, d, d, 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) + torch.cuda.synchronize() + out[want_frost] = c + torch.testing.assert_close(out[True].float(), out[False].float(), atol=2e-1, rtol=2e-2) + + +@requires_sm100 +@pytest.mark.parametrize("n,k", ((256, 128), (128, 128)), ids=("n!=k", "n==k")) +def test_bare_addresses_run_at_a_smaller_live_shape(n, k): + """A bare address wears the graph's layout AND the graph's allocation size. + + ``override_shapes`` then says the call runs a smaller problem inside it, so + the live shape matches neither the declaration nor a caller's buffer. The + backend takes this; a python plan that inferred the axis order from the + shape refused it. + """ + MB, NB, KB = 256, n, 128 + m, live_k = 128, k // 2 + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, MB, KB], stride=[MB * KB, KB, 1]) + B = g.tensor(name="B", uid=2, dim=[1, KB, NB], stride=[KB * NB, 1, KB]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + 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() + + a = torch.randn(1, MB, KB, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, NB, KB, dtype=torch.bfloat16, device="cuda") + c = torch.zeros(1, MB, NB, dtype=torch.bfloat16, device="cuda") + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute( + {1: a.data_ptr(), 2: b.data_ptr(), 3: c.data_ptr()}, + ws, + override_uids=[1, 2, 3], + override_shapes=[[1, m, live_k], [1, live_k, NB], [1, m, NB]], + override_strides=[[MB * KB, KB, 1], [KB * NB, 1, KB], [MB * NB, NB, 1]], + ) torch.cuda.synchronize() - torch.testing.assert_close(c.float(), torch.einsum("bmk,bnk->bmn", a.float(), b.float()), atol=2e-1, rtol=2e-2) + ref = torch.einsum("bmk,bnk->bmn", a[:, :m, :live_k].float(), b[:, :, :live_k].float()) + torch.testing.assert_close(c[:, :m].float(), ref, atol=2e-1, rtol=2e-2) @requires_sm100 From 992ac5630811631884d04a26e07baeed93ea51e7 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 02:20:14 -0700 Subject: [PATCH 3/9] Name the pack's positions and buffers for what they are "slot" was doing two jobs -- an int POSITION in the pack, and the OBJECT at that position -- so neither read as anything. They are now index_of() / operands() and OperandBuffer, and the C type, the native methods and every caller follow. PortSlots is PortIndices for the same reason. The pair that mattered most is on the compiled object: `run_views` (the interpreter) and `launch` (the emitted straight line) gave no hint they were two implementations of one thing. They are now `launch` and `lowered`, and the engine binds `_launch = lowered or launch` once at build. Also splits the per-call gate along the line the two halves actually fall on: check_shapes asks whether the extents agree with the problem size this call runs, check_alignment asks whether each buffer meets the width its role's accesses were compiled for. Neither name is new vocabulary. test_execute_recipe gains a degenerate-extent sweep. An extent of 1 leaves its axis's stride free, so two majors can look alike; what keeps that from mattering is that the TMA modulus is at least 4 and 1 divides none of them, so a unit contiguous extent never reaches a launch. The test asserts the property rather than the argument: every degenerate shape is refused or matches the backend. All eight K == 1 shapes take the refusal, at plan time. docs/python_graph_and_execution_backends.md gets the measured budget a python execute path has to fit in, and the pattern that fits it. An engine that re-derives its per-call facts lands near 40 us; the same kernel reading them once is 17.5. Co-Authored-By: Claude Opus 5 (1M context) --- docs/python_graph_and_execution_backends.md | 72 ++- python/cudnn/_pygraph.py | 4 +- python/cudnn/engines/base.py | 50 +- python/cudnn/frost/README.md | 6 + python/cudnn/frost/buffers.py | 2 +- python/cudnn/frost/workspace.py | 2 +- python/cudnn/gemm/frost/compiler.py | 50 +- python/cudnn/gemm/frost/engine.py | 30 +- python/cudnn/gemm/frost/recipe.py | 109 ++-- python/cudnn/linear_attention/engine_utils.py | 2 +- python/pygraph/variant_pack.cpp | 486 +++++++++--------- test/python/gemm/frost/test_execute_recipe.py | 79 ++- 12 files changed, 525 insertions(+), 367 deletions(-) diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index 149735ce1..ede9859ac 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -102,10 +102,16 @@ caller does not control. themselves, held in a C type (`pygraph/variant_pack.cpp`) as one `DLTensor` each. That type both consumes `__dlpack_c_exchange_api__` — the C function table a producer publishes on its type, which is how one crossing reads the -whole pack — and implements it, so a kernel reads a slot through the same fast -path it has for a framework tensor. `address` is the pointer array +whole pack — and implements it, so a kernel reads an operand through the same +fast path it has for a framework tensor. `address` is the pointer array `_execute_with_raw_ptrs` takes. +The pack's vocabulary distinguishes the position from the thing at it: +`pack.index_of(tensor_or_uid)` gives an operand's POSITION, and +`pack.operands(indices)` turns positions into `OperandBuffer`s — one caller +buffer described (pointer, shape, stride, dtype), non-owning. Resolve positions +once; ask for buffers per call. + Each operand's OWN dim/stride/data_type is what the pack holds, deliberately not the graph's declaration: the two may differ and one engine relies on it — `frost_gemm` takes its M/N/K from the buffers, so a plan built for one problem @@ -139,6 +145,68 @@ set it still receives the caller's `{uid: buffer}` map and reaches ports through has moved. `execute()` builds the `Tensor`s only for a plan that sets the flag — measured, normalizing for a plan that will not read the result costs more than it saves. + +#### What a per-execute path costs + +An engine owns its internals, and this section does not change that. It exists +because the default outcome is expensive: an engine that re-derives its per-call +facts lands around **40 µs of host time per execute**, and for a single-kernel +op that is most of what the caller pays. The same kernel with those facts read +once is **17.5**. Both numbers are `frost_gemm` at 256×256×128 bf16, host +enqueue, min over 25 reps of a 64-call burst from a drained queue. + +The budget it has to fit in, all measured on SM100: + +| | µs | +|---|---| +| `cuLaunchKernelEx`, untraced | 1.85 | +| one CuTe-DSL entry | ~3.6 | +| `graph.execute()` entry + `_normalize` | ~8 | +| **everything else is the engine's** | | + +Do not read a per-call cost out of an nsys trace: CUPTI adds ~2.2 µs per traced +API call, which is more than the call. + +**Split the facts by when they are decided.** Operand roles and majors, packing +factors, alignment requirements, output shape rules, which outputs need a seed — +all fixed when the kernel compiled. M/N/K, strides and pointers arrive per call. +Read the first set into a table at build (`gemm/frost/recipe.py` is the worked +example) and let the call read the table. That alone is 44 → 35. + +**Then lower the table for the shape you actually run.** Emitting one closure +per plan, with the table's constants inlined, is 35 → 17.5. Two rules make that +safe: + +- **The lowered path never raises.** Anything it is not certain of it hands to + the interpreting path, which serves every flavor and owns every rejection + message. It can then only ever accept a subset, and there is no second set of + error strings to drift. +- **Both read the same table.** A differential test between them catches + divergence, but never a misconception they share — so the table is where a + fact lives exactly once, and the tests that matter are against intended + semantics, at the shapes where two encodings coincide. + +This is a pattern to copy, not a framework to import. Sharing the code across +engines would couple their kernels' ABIs, which is the thing engine autonomy +buys; sharing the shape of the solution costs nothing. + +**Costs that are easy to miss, each measured:** + +- A `from x import y` inside a per-call function: **1.1 µs**. It was 65% of + what `_check_plan_device` cost. +- `torch.Tensor.permute()`: **1.4 µs** per call, per operand. +- Rebuilding a `{id(tensor): buffer}` map to look operands back up, when the + operand order was settled at build and a list index would do. +- Recomputing a pure function of values every call. `tensor_alignment`'s + layout half is **1.5 µs** and memoizes on `(shape, stride, elem_bytes)` — + values, so there is nothing to invalidate; only the pointer half is per call. +- Reading an operand through the exchange vtable is **0.08 µs** against 1.5 for + the python attribute walk. Framework neutrality is not what costs. + +**Measure from a drained queue, and sweep the burst size.** A number that is +flat in the burst size is host-bound; one that climbs with it is the device +rate, and back-to-back timing reads the device rate whenever host and device +are close. - **An engine does not propose its own plans.** Which configs to try, in what order, and where the backend's entries belong is one comparison across every candidate, and no engine can make it from the inside — it sees neither its diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 949e06db0..22fc71c09 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -1896,7 +1896,7 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= # slot that borrowed one is named here. from_graph.append(i) ptr, tensor = self._describe(data, order[i]) - native.set_slot(i, ptr, tuple(tensor.dim), tuple(tensor.stride), *_dlpack_code_bits(tensor.data_type)) + native.set_operand(i, ptr, tuple(tensor.dim), tuple(tensor.stride), *_dlpack_code_bits(tensor.data_type)) if strict: hole = native.first_unfilled() if hole >= 0: @@ -1917,7 +1917,7 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= i = slot_of.get(uid) if i is None: raise ValueError(f"override_uids names tensor uid {uid}, which is not an operand of this graph") - native.override_slot(i, *_in_axis_order_of(tuple(override_shapes[j]), tuple(override_strides[j]), native.stride(i))) + native.override_operand(i, *_in_axis_order_of(tuple(override_shapes[j]), tuple(override_strides[j]), native.stride(i))) # The workspace has no uid, so it is not an operand — but an engine has # to bounds-check its carves, and reading its size here is the same read # every other buffer gets rather than a second probe further down. diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 532fbd8af..e3d5dd46a 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -133,7 +133,7 @@ class VariantPack: pointers — silently, because every pointer in it is individually valid. """ - __slots__ = ("uids", "native", "_slot_of", "workspace", "workspace_bytes", "_device", "graph_described") + __slots__ = ("uids", "native", "_index_of", "workspace", "workspace_bytes", "_device", "graph_described") def __init__(self, uids, native, workspace_ptr: int = 0, workspace_bytes: int = 0, graph_described=()): self.uids = uids @@ -145,25 +145,25 @@ def __init__(self, uids, native, workspace_ptr: int = 0, workspace_bytes: int = # axis position needs this: the graph and the caller order a matmul's B # differently, and the description does not say which one it is. self.graph_described = graph_described - self._slot_of = None # built on first lookup: the backend never does one + self._index_of = None # built on first lookup: the backend never does one self._device = None @property def address(self) -> int: - """The ``void*[]`` in slot order, for ``_execute_with_raw_ptrs``.""" + """The ``void*[]`` in operand order, for ``_execute_with_raw_ptrs``.""" return self.native.address def all_contiguous(self): - """``(ok, slot)`` over every filled operand, decided from the strides + """``(ok, index)`` over every filled operand, decided from the strides the native pack already holds.""" ok, offender = self.native.all_contiguous() return ok, (int(offender) if offender else -1) @property - def slot_of(self): - if self._slot_of is None: - self._slot_of = {u: i for i, u in enumerate(self.uids)} - return self._slot_of + def index_of_uid(self): + if self._index_of is None: + self._index_of = {u: i for i, u in enumerate(self.uids)} + return self._index_of @property def device(self) -> int: @@ -178,20 +178,20 @@ def device(self) -> int: self._device = current_device() return self._device - def slot(self, tensor_or_uid) -> int: + def index_of(self, tensor_or_uid) -> int: """Index of a tensor's operand. KeyError when it is not caller-filled (a virtual intermediate, or a value the graph itself supplies).""" uid = tensor_or_uid if isinstance(tensor_or_uid, int) else tensor_or_uid.uid - return self.slot_of[uid] + return self.index_of_uid[uid] def ptr(self, tensor_or_uid) -> int: - return self.native.pointer(self.slot(tensor_or_uid)) + return self.native.pointer(self.index_of(tensor_or_uid)) - def views(self, slots): - """The DLPack producers for ``slots``, in one crossing.""" - return self.native.views(list(slots), self.device) + def operands(self, indices): + """The buffers for ``indices``, in one crossing.""" + return self.native.operands(list(indices), self.device) - def view(self, slot: int): + def operand(self, index: int): """A DLPack producer over one operand, for a kernel that needs an object rather than an address. @@ -203,15 +203,15 @@ def view(self, slot: int): which is an argument for making the producer a C type, not for keeping the caller's object. """ - return self.native.view(slot, self.device) + return self.native.operand(index, self.device) def __len__(self) -> int: return len(self.uids) @dataclass(frozen=True) -class PortSlots: - """Per-node ``{port_name: slot}``. A port with no caller operand — a +class PortIndices: + """Per-node ``{port_name: operand index}``. A port with no caller operand — a virtual intermediate — is ABSENT, so ``.get(port) is None`` keeps meaning what it meant when these were buffers.""" @@ -219,24 +219,24 @@ class PortSlots: outputs: Dict[str, int] -def bind_ports(graph: "pygraph", variant_pack: VariantPack) -> Dict[Any, PortSlots]: +def bind_ports(graph: "pygraph", variant_pack: VariantPack) -> Dict[Any, PortIndices]: """Join each node's wired ports with the operand layout. Strict: every non-virtual port must have an operand.""" def resolve(node, ports, direction): - slots = {} + indices = {} for port, t in ports.items(): if t is None: continue - slot = variant_pack.slot_of.get(t.uid) - if slot is None: + index = variant_pack.index_of_uid.get(t.uid) + if index is None: if t.is_virtual: continue # engine-internal intermediate raise ValueError(f"node {node.name!r}: no buffer for {direction} port {port!r} (tensor {t.name!r})") - slots[port] = slot - return slots + indices[port] = index + return indices - return {node: PortSlots(resolve(node, node.inputs, "input"), resolve(node, node.outputs, "output")) for node in graph.nodes} + return {node: PortIndices(resolve(node, node.inputs, "input"), resolve(node, node.outputs, "output")) for node in graph.nodes} def _view_over_address(address: int, tensor, node, port: str): diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index 348c8f488..84dfc1a71 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -249,6 +249,12 @@ CompiledPlan.execute(graph, uid_to_data, ctx) (the hot path) engine may hard-code a stream, reach into private graph state, or allocate hidden workspace. `uid_to_data` is the caller's variant pack (tensor uid -> device buffer), exactly as the classic backend receives it. +- **The pack's vocabulary is `index` and `OperandBuffer`,** and the two are not + the same thing. `pack.index_of(tensor_or_uid)` gives an operand's POSITION in + the pack; `pack.operands(indices)` turns positions into `OperandBuffer`s -- + one caller buffer described (pointer, shape, stride, dtype), non-owning, and + itself a DLPack producer. An engine resolves positions once at first execute + and asks for buffers per call. `python/cudnn/gemm/frost/engine.py` is the worked example, deliberately thin: `check_support` delegates to `probe_supported` and diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 280928e71..014d12631 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -180,7 +180,7 @@ def __dlpack__(self, *, stream=None, **_kwargs): is a use-after-free the moment a consumer outlives the view. """ code, bits = DTYPES[self.dtype] - return _pybind_module.make_slot(self._ptr, list(self.shape), code, bits, self._device_id).__dlpack__() + return _pybind_module.make_operand_buffer(self._ptr, list(self.shape), code, bits, self._device_id).__dlpack__() class DeviceBuffer(DeviceView): diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index 1089bfe5f..0322712a3 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -152,7 +152,7 @@ def view(self, offset: int, dtype: str, shape): count *= int(extent) self._check_span(offset, count * buffers.DTYPE_ITEMSIZE[dtype]) code, bits = buffers.DTYPES[dtype] - return _pybind_module.make_slot(self._ptr + offset, list(shape), code, bits, self._device) + return _pybind_module.make_operand_buffer(self._ptr + offset, list(shape), code, bits, self._device) def carve(self, plan): """Every region a :func:`carve_plan` describes, in one crossing.""" diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 820a65932..6c974341b 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -65,9 +65,10 @@ def _as_custream(stream): REDUCTION_INIT_VALUE, _output_rule, build as build_recipe, + check_alignment, + check_shapes, contiguous_modulus, expected_shape, - gate as gate_recipe, ) from .graph_analyzer import ( GemmBinding, @@ -1857,14 +1858,14 @@ class CompiledFusedGemm: # constructed without a binding and so has no call path. recipe: Any = field(default=None, init=False, repr=False, compare=False) # The recipe as one straight line, when this graph is a shape it emits. - launch: Any = field(default=None, init=False, repr=False, compare=False) + lowered: Any = field(default=None, init=False, repr=False, compare=False) bound: Any = field(default=(), init=False, repr=False, compare=False) def __post_init__(self) -> None: if self.binding is not None: self.bound = tuple(self.binding.bound_tensors()) self.recipe = build_recipe(self) - self.launch = self._lower() + self.lowered = self._lower() def __call__(self, variant_pack, stream=None): # The runtime call is a variant-pack dict keyed by cuDNN tensor object @@ -1879,13 +1880,13 @@ def __call__(self, variant_pack, stream=None): def run_resolved(self, resolved, stream=None): """Launch over ``{id(bound_tensor): buffer}``, already resolved.""" - views = [] + operands = [] for i, t in enumerate(self.bound): buf = resolved.get(id(t)) if buf is None: raise KeyError(f"variant pack is missing a buffer for {self.recipe.roles[i]}") - views.append(buf) - return (self.launch or self.run_views)(views, stream=stream) + operands.append(buf) + return (self.lowered or self.launch)(operands, stream=stream) def _lower(self): """The recipe as one straight line, or None for a shape it does not emit. @@ -1898,7 +1899,7 @@ def _lower(self): rebuilds it. The emitted body never raises. Anything it is not sure of it hands to - ``run_views``, which serves every flavor and owns every rejection + ``launch``, which serves every flavor and owns every rejection message -- so this can only ever accept a subset of what the general path accepts, and there is no second set of error strings to drift. @@ -1929,7 +1930,7 @@ def _lower(self): declared_stride = op.declared_layout[1] if op.dc != op.kc and declared_stride and declared_stride[op.kc] == 1: return None - ai, bi, ci = a.view, b.view, out.view + ai, bi, ci = a.index, b.index, out.index # kc is both the axis whose stride must be 1 and the axis whose extent # enters the TMA rule -- they are the same axis by definition of major. a_kc, b_kc = a.kc, b.kc @@ -1942,18 +1943,18 @@ def _lower(self): batch = out_batch if r.has_output_specs else max(a_batch, b_batch) device = self.device launchable = self._launchable - general = self.run_views + general = self.launch - def launch(views, graph_order=None, stream=None): + def lowered(operands, graph_order=None, stream=None): _check_plan_device(device) if graph_order: # An operand described from the graph is in the graph's axis # order; this body reads the caller's. - return general(views, graph_order, stream=stream) - av, bv, cv = views[ai], views[bi], views[ci] + return general(operands, graph_order, stream=stream) + av, bv, cv = operands[ai], operands[bi], operands[ci] a_sh, b_sh, c_sh = av.shape, bv.shape, cv.shape if len(a_sh) != 3 or len(b_sh) != 3 or len(c_sh) != 3: - return general(views, stream=stream) + return general(operands, stream=stream) a_st, b_st, c_st = av.stride(), bv.stride(), cv.stride() m, n, k = a_sh[1], b_sh[1], a_sh[2] * a_kpack if ( @@ -1971,7 +1972,7 @@ def launch(views, graph_order=None, stream=None): or _pow2_floor(bv.data_ptr()) < 16 or tensor_alignment(tuple(c_sh), tuple(c_st), cv.element_size(), ptr=cv.data_ptr()) < out_align ): - return general(views, stream=stream) + return general(operands, stream=stream) problem_size = ( m, n, @@ -1993,9 +1994,9 @@ def launch(views, graph_order=None, stream=None): # coincide, and a shape where they do not is not emitted here. return launchable(problem_size, av.permute(1, 2, 0), bv.permute(1, 2, 0), cv.permute(1, 2, 0), stream=_as_custream(stream)) - return launch + return lowered - def run_views(self, views, graph_order=None, stream=None): + def launch(self, operands, graph_order=None, stream=None): """Launch over the operand buffers, in bound-tensor order. Which bound tensor holds which operand is fixed when the plan compiles, @@ -2008,19 +2009,22 @@ def run_views(self, views, graph_order=None, stream=None): """ _check_plan_device(self.device) recipe = self.recipe - mnk, axes = recipe.problem(views, graph_order) - gate_recipe(recipe, views, mnk, axes) + mnk, axes = recipe.problem(operands, graph_order) + check_shapes(recipe, operands, mnk, axes) + check_alignment(recipe, operands, axes) - out_bufs = [views[o.view] for o in recipe.outputs] - aux_bufs = [views[x.view] for x in recipe.aux] + out_bufs = [operands[o.index] for o in recipe.outputs] + aux_bufs = [operands[x.index] for x in recipe.aux] c_arg = out_bufs if len(out_bufs) > 1 else out_bufs[0] # Everything below indexes an operand's axes by position, so an operand # that arrived in the graph's order is re-expressed in the kernel's here # rather than threaded through every launcher. - operands = [(views[op.view] if ax is KERNEL_AXES else views[op.view].permute(ax[AX_BATCH], ax[AX_MN], ax[AX_K])) for op, ax in zip(recipe.inputs, axes)] - a_bufs, b_bufs = operands[: recipe.b_at], operands[recipe.b_at :] + in_bufs = [ + (operands[op.index] if ax is KERNEL_AXES else operands[op.index].permute(ax[AX_BATCH], ax[AX_MN], ax[AX_K])) for op, ax in zip(recipe.inputs, axes) + ] + a_bufs, b_bufs = in_bufs[: recipe.b_at], in_bufs[recipe.b_at :] - sf_bufs = [views[s.view] for s in recipe.sf] + sf_bufs = [operands[s.index] for s in recipe.sf] if self.chain.is_multi_gemm: sfa, sfb = sf_bufs[: recipe.b_at], sf_bufs[recipe.b_at :] if self.block_scale: diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index 33a438b39..8f1f3c871 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -31,36 +31,36 @@ def __init__(self, compiled): # roles (matmul(A, A)), and resolve_variant_pack treats a repeated uid # as ambiguous. self._tensors = list(compiled.binding.bound_tensors()) - self._slots = None + self._operand_indices = None # Which call path this plan uses is a property of the compiled kernel, - # so it is chosen here and not re-asked per execute. ``launch`` is the + # so it is chosen here and not re-asked per execute. ``lowered`` is the # straight line the recipe lowers to when the kernel is a shape it - # emits; ``run_views`` serves everything else. - self._lowered = getattr(compiled, "launch", None) - self._run_views = self._lowered or getattr(compiled, "run_views", None) + # emits; ``launch`` serves everything else. + self._lowered = getattr(compiled, "lowered", None) + self._launch = self._lowered or getattr(compiled, "launch", None) def get_workspace_size(self) -> int: return int(getattr(self._compiled, "workspace_bytes", 0) or 0) def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: - slots = self._slots - if slots is None: + indices = self._operand_indices + if indices is None: try: - slots = self._slots = [variant_pack.slot(t.get_uid()) for t in self._tensors] + indices = self._operand_indices = [variant_pack.index_of(t.get_uid()) for t in self._tensors] except KeyError as exc: raise ValueError(f"frost_gemm: tensor uid {exc} is bound by the kernel but is not an operand of this graph") from exc # The kernel reads its M/N/K off these, so they must be the pack's -- # which carry the shape this execute runs, override_shapes included. - views = variant_pack.views(slots) + operands = variant_pack.operands(indices) required = self.get_workspace_size() if required: # Scratch is carved from the CALLER's workspace: stable pointers, so # a plan stays safe to capture in a CUDA graph. Only the MoE # launchers need it, and they take the variant-pack dict. - self._compiled(dict(zip(self._tensors, views)), Workspace.over(variant_pack, required, "frost_gemm"), stream=ctx.stream) + self._compiled(dict(zip(self._tensors, operands)), Workspace.over(variant_pack, required, "frost_gemm"), stream=ctx.stream) return - run_views = self._run_views - if run_views is not None: + launch = self._launch + if launch is not None: # Which bound tensor holds which operand was settled at build, so # the buffers arrive in that order and the launcher indexes them. # Which AXIS ORDER each one arrived in is a per-call fact only the @@ -70,11 +70,11 @@ def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: graph_order = None borrowed = variant_pack.graph_described if borrowed: - flags = tuple(s in borrowed for s in slots) + flags = tuple(i in borrowed for i in indices) graph_order = flags if any(flags) else None - run_views(views, graph_order, stream=ctx.stream) + launch(operands, graph_order, stream=ctx.stream) else: - self._compiled(dict(zip(self._tensors, views)), stream=ctx.stream) + self._compiled(dict(zip(self._tensors, operands)), stream=ctx.stream) class FrostGemmEngine(BaseEngine): diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py index 98e1f92de..291fe1d7e 100644 --- a/python/cudnn/gemm/frost/recipe.py +++ b/python/cudnn/gemm/frost/recipe.py @@ -9,11 +9,14 @@ else. This module writes the first set down, so that neither the interpreted path nor the straight line lowered from it re-derives them per call. -The two consumers read the same recipe but do not share a body: :func:`gate` -interprets it, and ``CompiledFusedGemm._lower`` emits a straight line with the -constants inlined. That is a compiler beside its interpreter, kept honest the -way those always are -- ``test_execute_recipe.py`` runs both over the same -accepts and rejects and requires the same answer. +The two consumers read the same recipe but do not share a body: +``CompiledFusedGemm.launch`` interprets it (through :func:`check_shapes` and +:func:`check_alignment`), and ``CompiledFusedGemm.lowered`` is the straight line +``_lower`` emits with the constants inlined. That is a compiler beside its +interpreter, kept honest the way those always are -- ``test_execute_recipe.py`` +runs both over the same accepts and rejects and requires the same answer. What +it CANNOT catch is a misconception they share, which is how the axis-order bug +survived it. """ from __future__ import annotations @@ -114,7 +117,7 @@ def _output_rule(spec, chain: FusionChain) -> tuple: class Operand: """One A or B operand: what a call must satisfy, with the rest baked.""" - view: int # position in the view list the engine hands over + index: int # position in the operand list the engine hands over role: str # names this operand in a rejection message major: str kpack: int # 2 when fp4 stores two elements per slot @@ -145,7 +148,7 @@ def axes(self, shape, stride, graph_order: bool) -> "tuple | None": @dataclass(frozen=True) class Output: - view: int + index: int role: str rule: tuple align: int @@ -156,7 +159,7 @@ class Output: @dataclass(frozen=True) class Aux: - view: int + index: int role: str align: int ref: Any # TensorRef, for the fake-shape reshape @@ -164,7 +167,7 @@ class Aux: @dataclass(frozen=True) class ScaleFactor: - view: int + index: int role: str is_a: bool operand_at: int # the A or B operand this one scales, as an index into inputs @@ -176,7 +179,7 @@ class GemmRecipe: outputs: tuple aux: tuple sf: tuple - roles: tuple # what occupies each view position, for a missing-buffer message + roles: tuple # what occupies each operand position, for a missing-buffer message a_at: int # M and K are read off inputs[a_at] b_at: int # N off inputs[b_at] has_output_specs: bool @@ -192,19 +195,19 @@ def a(self) -> Operand: def b(self) -> Operand: return self.inputs[self.b_at] - def problem(self, views, graph_order=None) -> tuple: + def problem(self, operands, graph_order=None) -> tuple: """``((M, N, K), axes-per-input)``, located rather than assumed. Reading M off axis 1 assumes the caller laid the buffer out the way the kernel reads it, which a bare device address does not: the pack describes that one from the graph's declaration, which orders B the - other way round. ``graph_order`` is the pack's per-view answer to which + other way round. ``graph_order`` is the pack's per-operand answer to which it was, or None when every operand is the caller's own. """ axes, bad = [], [] for op in self.inputs: - v = views[op.view] - borrowed = bool(graph_order and graph_order[op.view]) + v = operands[op.index] + borrowed = bool(graph_order and graph_order[op.index]) ax = op.axes(v.shape, v.stride(), borrowed) if ax is None: want = op.dc if borrowed else op.kc @@ -217,17 +220,17 @@ def problem(self, views, graph_order=None) -> tuple: raise ValueError("runtime operand layout does not match the layout the kernel was compiled for: " + "; ".join(bad)) a, b = self.a, self.b a_ax, b_ax = axes[self.a_at], axes[self.b_at] - a_shape, b_shape = views[a.view].shape, views[b.view].shape + a_shape, b_shape = operands[a.index].shape, operands[b.index].shape return (a_shape[a_ax[AX_MN]], b_shape[b_ax[AX_MN]], a_shape[a_ax[AX_K]] * a.kpack), tuple(axes) -def _tma_reject(recipe: GemmRecipe, views, axes) -> "str | None": +def _tma_reject(recipe: GemmRecipe, operands, axes) -> "str | None": """TMA encodes the contiguous dimension in 16-byte units; a misaligned extent silently mis-strides every row past the first.""" bad = [] for op, ax in zip(recipe.inputs, axes): role = op.contiguous_role - extent = views[op.view].shape[ax[role]] + extent = operands[op.index].shape[ax[role]] if extent % op.modulus: # Report the LOGICAL extent and modulus: fp4 stores two elements per # slot, so "K % 32" is the rule a user wrote their shape against. @@ -238,13 +241,13 @@ def _tma_reject(recipe: GemmRecipe, views, axes) -> "str | None": return None -def _shape_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": +def _shape_reject(recipe: GemmRecipe, operands, axes, mnk) -> "str | None": """Every operand must agree with the M/N/K read off the first A and B -- the kernel walks the inferred K on all of them.""" m, n, k = mnk bad = [] for op, ax in zip(recipe.inputs, axes): - shape = views[op.view].shape + shape = operands[op.index].shape want = (op.batch, n if op.is_b else m, k // op.kpack) got = (shape[ax[AX_BATCH]], shape[ax[AX_MN]], shape[ax[AX_K]]) if got != want: @@ -254,11 +257,11 @@ def _shape_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": return None -def _output_shape_reject(recipe: GemmRecipe, views, mnk) -> "str | None": +def _output_shape_reject(recipe: GemmRecipe, operands, mnk) -> "str | None": m, n, _ = mnk bad = [] for out in recipe.outputs: - shape = tuple(views[out.view].shape) + shape = tuple(operands[out.index].shape) want = expected_shape(out.rule, m, n) if shape != want: bad.append(f"{out.role}: expected {want}, got {shape}") @@ -267,7 +270,7 @@ def _output_shape_reject(recipe: GemmRecipe, views, mnk) -> "str | None": return None -def _align_reject(recipe: GemmRecipe, views) -> "str | None": +def _align_reject(recipe: GemmRecipe, operands) -> "str | None": """Every buffer's alignment must meet the width its role's access uses. Inputs and scale factors are TMA-loaded, so only the base pointer is at @@ -276,12 +279,12 @@ def _align_reject(recipe: GemmRecipe, views) -> "str | None": """ bad = [] for item in recipe.inputs + recipe.sf: - ptr = int(views[item.view].data_ptr()) + ptr = int(operands[item.index].data_ptr()) align = _pow2_floor(ptr) if align < 16: bad.append(f"{item.role}: alignment {align}B < required 16B (ptr=0x{ptr:x})") for item in recipe.outputs + recipe.aux: - v = views[item.view] + v = operands[item.index] ptr = int(v.data_ptr()) align = tensor_alignment(tuple(v.shape), tuple(v.stride()), v.element_size(), ptr=ptr) if align < item.align: @@ -291,7 +294,7 @@ def _align_reject(recipe: GemmRecipe, views) -> "str | None": return None -def _sf_blob_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": +def _sf_blob_reject(recipe: GemmRecipe, operands, axes, mnk) -> "str | None": """A block-scale SF reaches the kernel as a base pointer plus a layout the template re-synthesizes from M/N/K, so a blob that is not one dense byte run of at least the required size is read out of bounds with no fault.""" @@ -299,10 +302,10 @@ def _sf_blob_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": k4 = ((k // recipe.block_size) + 3) // 4 bad = [] for sf in recipe.sf: - v = views[sf.view] + v = operands[sf.index] op = recipe.inputs[sf.operand_at] rows = m if sf.is_a else n - batch = views[op.view].shape[axes[sf.operand_at][AX_BATCH]] + batch = operands[op.index].shape[axes[sf.operand_at][AX_BATCH]] required = 512 * k4 * ((rows + 127) // 128) * int(batch) span = 1 + sum((int(s) - 1) * int(st) for s, st in zip(v.shape, v.stride())) if int(v.numel()) != span: @@ -318,21 +321,39 @@ def _sf_blob_reject(recipe: GemmRecipe, views, axes, mnk) -> "str | None": return None -def gate(recipe: GemmRecipe, views, mnk, axes) -> None: - """Raise on anything about this call the compiled kernel cannot serve.""" - reasons = [ - _tma_reject(recipe, views, axes), - _shape_reject(recipe, views, axes, mnk), - _output_shape_reject(recipe, views, mnk), - _align_reject(recipe, views), - ] - if recipe.block_size: - reasons.append(_sf_blob_reject(recipe, views, axes, mnk)) +def _raise_first(reasons) -> None: for reason in reasons: if reason is not None: raise ValueError(reason) +def check_shapes(recipe: GemmRecipe, operands, mnk, axes) -> None: + """Do the extents agree with the problem size this call will run? + + The kernel's M/N/K are symbolic, so one plan serves many problem sizes and + the call's own extents are what has to hold together: every operand against + the M/N/K read off the first A and B, every output against the shape rule + the build recorded, and a block-scale blob against the size the template + re-synthesizes. + """ + reasons = [_shape_reject(recipe, operands, axes, mnk), _output_shape_reject(recipe, operands, mnk)] + if recipe.block_size: + reasons.append(_sf_blob_reject(recipe, operands, axes, mnk)) + _raise_first(reasons) + + +def check_alignment(recipe: GemmRecipe, operands, axes) -> None: + """Does every buffer meet the width its role's accesses were compiled for? + + Two rules, both about 16 bytes and neither about shape: TMA encodes the + contiguous dimension in 16-byte units, so its extent has a modulus; and each + buffer's base (plus, for a stored output, its stride and contiguous extent) + bounds the widest vector the kernel can issue. Below either, the kernel does + not fault -- it mis-strides or reads past the end. + """ + _raise_first([_tma_reject(recipe, operands, axes), _align_reject(recipe, operands)]) + + def _declared_layout(tensor) -> tuple: """The ``(dim, stride)`` the graph declared, or ``((), ())`` when it cannot say. @@ -345,12 +366,12 @@ def _declared_layout(tensor) -> tuple: return (), () -def _operand(view: int, role: str, tensor, *, major: str, dtype: str, batch: int, is_b: bool) -> Operand: +def _operand(index: int, role: str, tensor, *, major: str, dtype: str, batch: int, is_b: bool) -> Operand: declared = _DECLARED_AXES["b" if is_b else "a"] contiguous = AX_K if major == "k" else AX_MN modulus, pack = contiguous_modulus(dtype, contiguous == AX_K) return Operand( - view=view, + index=index, role=role, major=major, kpack=2 if dtype == "fp4_e2m1" else 1, @@ -392,7 +413,7 @@ def build(compiled) -> GemmRecipe: sqrt = red.mode == "norm2" outputs.append( Output( - view=order[id(t)], + index=order[id(t)], role=spec.source, rule=_output_rule(spec, chain), align=out_reqs[i], @@ -403,18 +424,18 @@ def build(compiled) -> GemmRecipe: ) aux = tuple( - Aux(view=order[id(t)], role=f"aux {compiled.aux_names[i]!r}", align=aux_reqs[compiled.aux_names[i]], ref=ref) + Aux(index=order[id(t)], role=f"aux {compiled.aux_names[i]!r}", align=aux_reqs[compiled.aux_names[i]], ref=ref) for i, (t, ref) in enumerate(zip(binding.aux, chain.aux_tensors)) ) na = len(binding.a_operands) sf = tuple( - [ScaleFactor(view=order[id(t)], role=f"SFA[{i}]", is_a=True, operand_at=i) for i, t in enumerate(binding.sfa_operands)] - + [ScaleFactor(view=order[id(t)], role=f"SFB[{j}]", is_a=False, operand_at=na + j) for j, t in enumerate(binding.sfb_operands)] + [ScaleFactor(index=order[id(t)], role=f"SFA[{i}]", is_a=True, operand_at=i) for i, t in enumerate(binding.sfa_operands)] + + [ScaleFactor(index=order[id(t)], role=f"SFB[{j}]", is_a=False, operand_at=na + j) for j, t in enumerate(binding.sfb_operands)] ) roles = [f"bound tensor {i}" for i in range(len(binding.bound_tensors()))] for item in (*inputs, *outputs, *aux, *sf): - roles[item.view] = item.role + roles[item.index] = item.role return GemmRecipe( inputs=tuple(inputs), diff --git a/python/cudnn/linear_attention/engine_utils.py b/python/cudnn/linear_attention/engine_utils.py index d44f366aa..73be1d59c 100644 --- a/python/cudnn/linear_attention/engine_utils.py +++ b/python/cudnn/linear_attention/engine_utils.py @@ -71,7 +71,7 @@ def execute(self, graph, variant_pack, ctx) -> None: node_buffers = {} for node, slots in ports.items(): names = list(slots.inputs) + list(slots.outputs) - views = variant_pack.views(list(slots.inputs.values()) + list(slots.outputs.values())) + views = variant_pack.operands(list(slots.inputs.values()) + list(slots.outputs.values())) split = len(slots.inputs) node_buffers[node] = NodeBuffers(dict(zip(names[:split], views[:split])), dict(zip(names[split:], views[split:]))) required = self._compiled.workspace_bytes() diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 0c25db2b7..281ad728b 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -6,12 +6,12 @@ // `__dlpack_c_exchange_api__` is a vtable on the buffer's TYPE whose // dltensor_from_py_object_no_sync fills a caller-provided DLTensor in place -- // no capsule, no allocation. This file consumes it to read the caller's -// operands and implements it so the slots it hands a kernel are read the same +// operands and implements it so the operands it hands a kernel are read the same // way, which is why nothing is given up by refusing to pass the caller's // object through. // -// A producer without the vtable is not an error: read_slot returns false and -// python fills that slot from its own reader, so a mixed pack costs the sum of +// A producer without the vtable is not an error: read_operand returns false and +// python fills that operand from its own reader, so a mixed pack costs the sum of // its parts. #include "variant_pack.h" @@ -145,8 +145,8 @@ is_dense(const DLTensor &t) { } // One operand. The shape and stride live here rather than behind the DLTensor's -// pointers so a slot stays valid once the producer's own DLTensor is gone. -struct Slot { +// pointers so a operand stays valid once the producer's own DLTensor is gone. +struct Operand { void *data = nullptr; int32_t ndim = 0; DLDataType dtype = {0, 0, 1}; @@ -157,30 +157,30 @@ struct Slot { } // namespace -// A pack's slot, exposed to a kernel. It implements the same exchange protocol +// A pack's operand, exposed to a kernel. It implements the same exchange protocol // it was read through, so a consumer that has the fast path for a torch tensor // has it for this too. -class VariantPackSlot { +class OperandBuffer { public: - VariantPackSlot(Slot slot, int32_t device_id) : slot_(std::move(slot)) { - tensor_.data = slot_.data; + OperandBuffer(Operand operand, int32_t device_id) : operand_(std::move(operand)) { + tensor_.data = operand_.data; tensor_.device = DLDevice{kDLCUDA, device_id}; - tensor_.ndim = slot_.ndim; - tensor_.dtype = slot_.dtype; - tensor_.shape = slot_.shape.empty() ? nullptr : slot_.shape.data(); - tensor_.strides = slot_.stride.empty() ? nullptr : slot_.stride.data(); + tensor_.ndim = operand_.ndim; + tensor_.dtype = operand_.dtype; + tensor_.shape = operand_.shape.empty() ? nullptr : operand_.shape.data(); + tensor_.strides = operand_.stride.empty() ? nullptr : operand_.stride.data(); tensor_.byte_offset = 0; } - // tensor_ points into slot_'s vectors, so a copy would leave the new - // object's DLTensor describing the old one's storage. Slots are always + // tensor_ points into operand_'s vectors, so a copy would leave the new + // object's DLTensor describing the old one's storage. Operands are always // heap-allocated and handed out by pointer, so nothing needs to copy one. - VariantPackSlot(const VariantPackSlot &) = delete; - VariantPackSlot & - operator=(const VariantPackSlot &) = delete; - VariantPackSlot(VariantPackSlot &&) = delete; - VariantPackSlot & - operator=(VariantPackSlot &&) = delete; + OperandBuffer(const OperandBuffer &) = delete; + OperandBuffer & + operator=(const OperandBuffer &) = delete; + OperandBuffer(OperandBuffer &&) = delete; + OperandBuffer & + operator=(OperandBuffer &&) = delete; const DLTensor & tensor() const { @@ -194,24 +194,24 @@ class VariantPackSlot { std::vector shape() const { - return slot_.shape; + return operand_.shape; } std::vector stride() const { - if (!slot_.stride.empty()) return slot_.stride; - std::vector dense(slot_.ndim, 1); - for (int d = slot_.ndim - 2; d >= 0; d--) dense[d] = dense[d + 1] * slot_.shape[d + 1]; + if (!operand_.stride.empty()) return operand_.stride; + std::vector dense(operand_.ndim, 1); + for (int d = operand_.ndim - 2; d >= 0; d--) dense[d] = dense[d + 1] * operand_.shape[d + 1]; return dense; } // One axis of it, the way a framework tensor is asked (stride(-1)). int64_t stride_at(int64_t dim) const { - int64_t axis = dim < 0 ? dim + slot_.ndim : dim; - if (axis < 0 || axis >= slot_.ndim) + int64_t axis = dim < 0 ? dim + operand_.ndim : dim; + if (axis < 0 || axis >= operand_.ndim) throw py::index_error("stride(): dimension " + std::to_string(dim) + " is out of range for a " + - std::to_string(slot_.ndim) + "-D slot"); + std::to_string(operand_.ndim) + "-D operand"); return stride()[axis]; } @@ -220,18 +220,18 @@ class VariantPackSlot { // a torch tensor's "torch.bfloat16" and this "bfloat16" answer the same. std::string dtype() const { - return dtype_name(slot_.dtype); + return dtype_name(operand_.dtype); } int64_t element_size() const { - return slot_.dtype.bits / 8; + return operand_.dtype.bits / 8; } int64_t numel() const { int64_t n = 1; - for (int64_t extent : slot_.shape) n *= extent; + for (int64_t extent : operand_.shape) n *= extent; return n; } @@ -242,7 +242,7 @@ class VariantPackSlot { int64_t length() const { - return slot_.shape.empty() ? 0 : slot_.shape[0]; + return operand_.shape.empty() ? 0 : operand_.shape[0]; } py::tuple @@ -251,14 +251,14 @@ class VariantPackSlot { } // A differently shaped view of the same memory, with one -1 wildcard. Only - // meaningful for a dense slot, which is why a strided one is refused rather + // meaningful for a dense operand, which is why a strided one is refused rather // than silently reinterpreted. - VariantPackSlot * + OperandBuffer * reshape(std::vector shape) const { - if (!slot_.stride.empty() && !is_dense(tensor_)) - throw py::value_error("cannot reshape a non-contiguous variant-pack slot"); + if (!operand_.stride.empty() && !is_dense(tensor_)) + throw py::value_error("cannot reshape a non-contiguous variant-pack operand"); int64_t numel = 1; - for (int64_t extent : slot_.shape) numel *= extent; + for (int64_t extent : operand_.shape) numel *= extent; int64_t fixed = 1; int wildcard = -1; for (size_t d = 0; d < shape.size(); d++) { @@ -276,37 +276,37 @@ class VariantPackSlot { } else if (fixed != numel) { throw py::value_error("cannot reshape " + std::to_string(numel) + " elements to " + std::to_string(fixed)); } - Slot out = slot_; - out.shape = std::move(shape); - out.ndim = static_cast(out.shape.size()); + Operand out = operand_; + out.shape = std::move(shape); + out.ndim = static_cast(out.shape.size()); out.stride.clear(); // dense by construction, as the reshape required - return new VariantPackSlot(out, tensor_.device.device_id); + return new OperandBuffer(out, tensor_.device.device_id); } - // The same memory with its axes relabelled; exact for a strided slot too. - VariantPackSlot * + // The same memory with its axes relabelled; exact for a strided operand too. + OperandBuffer * permute(const std::vector &axes) const { - if (axes.size() != static_cast(slot_.ndim)) - throw py::value_error("permute needs one axis per dimension: this slot is " + std::to_string(slot_.ndim) + - "-D"); + if (axes.size() != static_cast(operand_.ndim)) + throw py::value_error("permute needs one axis per dimension: this operand is " + + std::to_string(operand_.ndim) + "-D"); std::vector seen(axes.size(), false); - Slot out = slot_; - out.stride.assign(slot_.ndim, 1); - if (slot_.stride.empty()) { - for (int d = slot_.ndim - 2; d >= 0; d--) out.stride[d] = out.stride[d + 1] * slot_.shape[d + 1]; + Operand out = operand_; + out.stride.assign(operand_.ndim, 1); + if (operand_.stride.empty()) { + for (int d = operand_.ndim - 2; d >= 0; d--) out.stride[d] = out.stride[d + 1] * operand_.shape[d + 1]; } else { - out.stride = slot_.stride; + out.stride = operand_.stride; } const std::vector from_stride = out.stride; for (size_t d = 0; d < axes.size(); d++) { - int64_t axis = axes[d] < 0 ? axes[d] + slot_.ndim : axes[d]; - if (axis < 0 || axis >= slot_.ndim || seen[axis]) - throw py::value_error("permute axes must be a permutation of the slot's dimensions"); + int64_t axis = axes[d] < 0 ? axes[d] + operand_.ndim : axes[d]; + if (axis < 0 || axis >= operand_.ndim || seen[axis]) + throw py::value_error("permute axes must be a permutation of the operand's dimensions"); seen[axis] = true; - out.shape[d] = slot_.shape[axis]; + out.shape[d] = operand_.shape[axis]; out.stride[d] = from_stride[axis]; } - return new VariantPackSlot(std::move(out), tensor_.device.device_id); + return new OperandBuffer(std::move(out), tensor_.device.device_id); } // Row-major contiguous by construction, so this is the identity a caller @@ -323,11 +323,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. + // outlives this operand rather than aliasing storage the operand 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 + // from_dlpack; tvm-ffi reads a operand through the exchange vtable and never // gets here. py::capsule dlpack(py::object /*stream*/, py::object /*max_version*/) const { @@ -336,7 +336,7 @@ class VariantPackSlot { std::vector shape; std::vector stride; }; - auto *owned = new Owned{{}, slot_.shape, slot_.stride}; + auto *owned = new Owned{{}, operand_.shape, operand_.stride}; owned->managed.dl_tensor = tensor_; owned->managed.dl_tensor.shape = owned->shape.empty() ? nullptr : owned->shape.data(); owned->managed.dl_tensor.strides = owned->stride.empty() ? nullptr : owned->stride.data(); @@ -353,36 +353,36 @@ class VariantPackSlot { } private: - Slot slot_; // owns the shape/stride storage the DLTensor points into + Operand operand_; // owns the shape/stride storage the DLTensor points into DLTensor tensor_{}; }; namespace { int -slot_dltensor_from_py_object(void *py_object, DLTensor *out) { - auto *slot = py::cast(py::handle(static_cast(py_object))); - *out = slot->tensor(); +buffer_dltensor_from_py_object(void *py_object, DLTensor *out) { + auto *operand = py::cast(py::handle(static_cast(py_object))); + *out = operand->tensor(); return 0; } int -slot_managed_from_py_object(void *py_object, DLManagedTensorVersioned **out) { - auto *slot = py::cast(py::handle(static_cast(py_object))); +buffer_managed_from_py_object(void *py_object, DLManagedTensorVersioned **out) { + auto *operand = py::cast(py::handle(static_cast(py_object))); // A managed tensor is the form a consumer is allowed to OUTLIVE the - // producer with, so it cannot point at the slot's vectors: the shape and + // producer with, so it cannot point at the operand's vectors: the shape and // stride are copied and owned here, and the deleter frees them. struct Managed { DLManagedTensorVersioned versioned; std::vector shape; std::vector stride; }; - auto *owned = - new Managed{{}, slot->shape(), slot->tensor().strides == nullptr ? std::vector() : slot->stride()}; + auto *owned = new Managed{ + {}, operand->shape(), operand->tensor().strides == nullptr ? std::vector() : operand->stride()}; auto &tensor = owned->versioned.dl_tensor; - tensor = slot->tensor(); + tensor = operand->tensor(); tensor.shape = owned->shape.empty() ? nullptr : owned->shape.data(); - // a slot with no stride array is dense, and DLPack spells that as null + // a operand with no stride array is dense, and DLPack spells that as null tensor.strides = owned->stride.empty() ? nullptr : owned->stride.data(); owned->versioned.version.major = DLPACK_MAJOR_VERSION; owned->versioned.version.minor = DLPACK_MINOR_VERSION; @@ -393,17 +393,17 @@ slot_managed_from_py_object(void *py_object, DLManagedTensorVersioned **out) { } int -slot_allocator(DLTensor *, - DLManagedTensorVersioned **, - void *error_ctx, - void (*set_error)(void *, const char *, const char *)) { - set_error(error_ctx, "NotImplementedError", "a variant-pack slot views the caller's memory; it never allocates"); +buffer_allocator(DLTensor *, + DLManagedTensorVersioned **, + void *error_ctx, + void (*set_error)(void *, const char *, const char *)) { + set_error(error_ctx, "NotImplementedError", "a variant-pack operand views the caller's memory; it never allocates"); return -1; } int -slot_to_py_object(DLManagedTensorVersioned *, void **) { - PyErr_SetString(PyExc_NotImplementedError, "a variant-pack slot is not an importer"); +buffer_to_py_object(DLManagedTensorVersioned *, void **) { + PyErr_SetString(PyExc_NotImplementedError, "a variant-pack operand is not an importer"); return -1; } @@ -411,23 +411,23 @@ slot_to_py_object(DLManagedTensorVersioned *, void **) { // explicitly. Reporting no producer stream is what tells a consumer to use the // one it was given rather than going looking for ours. int -slot_current_work_stream(DLDeviceType, int32_t, void **out_stream) { +buffer_current_work_stream(DLDeviceType, int32_t, void **out_stream) { *out_stream = nullptr; return 0; } DLPackExchangeAPI & -slot_exchange_api() { +buffer_exchange_api() { static DLPackExchangeAPI api = [] { DLPackExchangeAPI table{}; table.header.version.major = DLPACK_MAJOR_VERSION; table.header.version.minor = DLPACK_MINOR_VERSION; table.header.prev_api = nullptr; - table.managed_tensor_allocator = slot_allocator; - table.managed_tensor_from_py_object_no_sync = slot_managed_from_py_object; - table.managed_tensor_to_py_object_no_sync = slot_to_py_object; - table.dltensor_from_py_object_no_sync = slot_dltensor_from_py_object; - table.current_work_stream = slot_current_work_stream; + table.managed_tensor_allocator = buffer_allocator; + table.managed_tensor_from_py_object_no_sync = buffer_managed_from_py_object; + table.managed_tensor_to_py_object_no_sync = buffer_to_py_object; + table.dltensor_from_py_object_no_sync = buffer_dltensor_from_py_object; + table.current_work_stream = buffer_current_work_stream; return table; }(); return api; @@ -437,35 +437,35 @@ slot_exchange_api() { class VariantPackNative { public: - explicit VariantPackNative(size_t n) : slots_(n), pointers_(n, nullptr) {} + explicit VariantPackNative(size_t n) : operands_(n), pointers_(n, nullptr) {} - // Fill one slot from the caller's buffer. False means its type does not + // Fill one operand from the caller's buffer. False means its type does not // implement the exchange protocol and python must describe it instead. bool - read_slot(size_t index, py::handle buffer) { - Slot &slot = slots_.at(index); + read_operand(size_t index, py::handle buffer) { + Operand &operand = operands_.at(index); DLPackExchangeAPI *api = exchange_api_for(buffer.ptr()); if (api == nullptr || api->dltensor_from_py_object_no_sync == nullptr) return false; DLTensor t{}; if (api->dltensor_from_py_object_no_sync(buffer.ptr(), &t) != 0) throw py::error_already_set(); - slot.data = static_cast(t.data) + t.byte_offset; - slot.ndim = t.ndim; - slot.dtype = t.dtype; - slot.shape.assign(t.shape, t.shape + t.ndim); + operand.data = static_cast(t.data) + t.byte_offset; + operand.ndim = t.ndim; + operand.dtype = t.dtype; + operand.shape.assign(t.shape, t.shape + t.ndim); if (t.strides != nullptr) { - slot.stride.assign(t.strides, t.strides + t.ndim); + operand.stride.assign(t.strides, t.strides + t.ndim); } else { - slot.stride.clear(); + operand.stride.clear(); } - slot.filled = true; - pointers_[index] = slot.data; + operand.filled = true; + pointers_[index] = operand.data; return true; } - // Every slot in one call. Returns the indices whose producer has no vtable, - // for python to describe and report back through set_slot -- crossing the + // Every operand in one call. Returns the indices whose producer has no vtable, + // for python to describe and report back through set_operand -- crossing the // binding once per pack rather than once per operand is 2.53 us against - // 1.0 for eight. A None entry is a slot the caller did not fill. + // 1.0 for eight. A None entry is a operand the caller did not fill. // A uid the map does not carry is left unfilled rather than refused: whether // that is the caller's mistake or an optional port depends on the graph, // which python knows and this does not. @@ -473,104 +473,104 @@ class VariantPackNative { read_from(const py::dict &uid_to_data, const std::vector &uids) { std::vector unread; const size_t n = uids.size(); - for (size_t i = 0; i < n && i < slots_.size(); i++) { + for (size_t i = 0; i < n && i < operands_.size(); i++) { PyObject *buffer = PyDict_GetItem(uid_to_data.ptr(), py::int_(uids[i]).ptr()); if (buffer == nullptr || buffer == Py_None) { - skip_slot(i); - } else if (!read_slot(i, py::handle(buffer))) { + skip_operand(i); + } else if (!read_operand(i, py::handle(buffer))) { unread.push_back(i); } } return unread; } - // The first slot no one filled, or -1. + // The first operand no one filled, or -1. int64_t first_unfilled() const { - for (size_t i = 0; i < slots_.size(); i++) { - if (!slots_[i].filled) return static_cast(i); + for (size_t i = 0; i < operands_.size(); i++) { + if (!operands_[i].filled) return static_cast(i); } return -1; } // The fallback: python read the buffer its own way and reports the result. void - set_slot(size_t index, - int64_t ptr, - std::vector shape, - std::vector stride, - int dtype_code, - int dtype_bits) { - Slot &slot = slots_.at(index); - slot.data = reinterpret_cast(ptr); - slot.ndim = static_cast(shape.size()); - slot.dtype = DLDataType{static_cast(dtype_code), static_cast(dtype_bits), 1}; - slot.shape = std::move(shape); - slot.stride = std::move(stride); - slot.filled = true; - pointers_[index] = slot.data; - } - - // Re-describe a slot at the shape this execute runs, keeping its buffer. + set_operand(size_t index, + int64_t ptr, + std::vector shape, + std::vector stride, + int dtype_code, + int dtype_bits) { + Operand &operand = operands_.at(index); + operand.data = reinterpret_cast(ptr); + operand.ndim = static_cast(shape.size()); + operand.dtype = DLDataType{static_cast(dtype_code), static_cast(dtype_bits), 1}; + operand.shape = std::move(shape); + operand.stride = std::move(stride); + operand.filled = true; + pointers_[index] = operand.data; + } + + // Re-describe a operand at the shape this execute runs, keeping its buffer. // Applying override_shapes here rather than in an engine is what keeps the // two paths answering the same question: an engine that reads the pack // honours the override without knowing the concept exists. void - override_slot(size_t index, std::vector shape, std::vector stride) { - Slot &slot = slots_.at(index); - if (!slot.filled) { - throw py::value_error("variant-pack slot " + std::to_string(index) + " has no buffer to re-describe"); + override_operand(size_t index, std::vector shape, std::vector stride) { + Operand &operand = operands_.at(index); + if (!operand.filled) { + throw py::value_error("variant-pack operand " + std::to_string(index) + " has no buffer to re-describe"); } // ndim comes from the shape and the stride array is read ndim deep, and // this is the one place the two arrive from different lists. if (shape.size() != stride.size()) { throw py::value_error("override shape and stride must have the same rank; got " + std::to_string(shape.size()) + " and " + std::to_string(stride.size()) + - " for slot " + std::to_string(index)); + " for operand " + std::to_string(index)); } - slot.ndim = static_cast(shape.size()); - slot.shape = std::move(shape); - slot.stride = std::move(stride); + operand.ndim = static_cast(shape.size()); + operand.shape = std::move(shape); + operand.stride = std::move(stride); } void - skip_slot(size_t index) { - slots_.at(index).filled = false; - pointers_[index] = nullptr; + skip_operand(size_t index) { + operands_.at(index).filled = false; + pointers_[index] = nullptr; } bool all_contiguous(std::string &offender) const { - for (size_t i = 0; i < slots_.size(); i++) { - const Slot &slot = slots_[i]; - if (!slot.filled || slot.stride.empty()) continue; + for (size_t i = 0; i < operands_.size(); i++) { + const Operand &operand = operands_[i]; + if (!operand.filled || operand.stride.empty()) continue; int64_t expect = 1; - for (int d = slot.ndim - 1; d >= 0; d--) { - if (slot.shape[d] != 1 && slot.stride[d] != expect) { + for (int d = operand.ndim - 1; d >= 0; d--) { + if (operand.shape[d] != 1 && operand.stride[d] != expect) { offender = std::to_string(i); return false; } - expect *= slot.shape[d]; + expect *= operand.shape[d]; } } return true; } bool - slot_contiguous(size_t index) const { - const Slot &slot = slots_.at(index); - if (!slot.filled || slot.stride.empty()) return true; + operand_contiguous(size_t index) const { + const Operand &operand = operands_.at(index); + if (!operand.filled || operand.stride.empty()) return true; int64_t expect = 1; - for (int d = slot.ndim - 1; d >= 0; d--) { - if (slot.shape[d] != 1 && slot.stride[d] != expect) return false; - expect *= slot.shape[d]; + for (int d = operand.ndim - 1; d >= 0; d--) { + if (operand.shape[d] != 1 && operand.stride[d] != expect) return false; + expect *= operand.shape[d]; } return true; } bool is_filled(size_t index) const { - return slots_.at(index).filled; + return operands_.at(index).filled; } int64_t @@ -580,26 +580,26 @@ class VariantPackNative { std::vector shape(size_t index) const { - return slots_.at(index).shape; + return operands_.at(index).shape; } std::vector stride(size_t index) const { - const Slot &slot = slots_.at(index); - if (!slot.stride.empty()) return slot.stride; - std::vector dense(slot.ndim, 1); - for (int d = slot.ndim - 2; d >= 0; d--) dense[d] = dense[d + 1] * slot.shape[d + 1]; + const Operand &operand = operands_.at(index); + if (!operand.stride.empty()) return operand.stride; + std::vector dense(operand.ndim, 1); + for (int d = operand.ndim - 2; d >= 0; d--) dense[d] = dense[d + 1] * operand.shape[d + 1]; return dense; } py::tuple dtype(size_t index) const { - const Slot &slot = slots_.at(index); - return py::make_tuple(slot.dtype.code, slot.dtype.bits); + const Operand &operand = operands_.at(index); + return py::make_tuple(operand.dtype.code, operand.dtype.bits); } // The address the backend's variant pack reads: a contiguous void*[] in - // slot order, so it goes to _execute_with_raw_ptrs with no copy. + // operand order, so it goes to _execute_with_raw_ptrs with no copy. int64_t pointer_array(void) const { return reinterpret_cast(pointers_.data()); @@ -607,48 +607,48 @@ class VariantPackNative { size_t size(void) const { - return slots_.size(); + return operands_.size(); } - // Every requested slot in one crossing, for an engine binding a whole node. - std::vector - views(const std::vector &indices, int32_t device_id) const { - std::vector out; + // Every requested operand in one crossing, for an engine binding a whole node. + std::vector + operands(const std::vector &indices, int32_t device_id) const { + std::vector out; out.reserve(indices.size()); - for (size_t index : indices) out.push_back(view(index, device_id)); + for (size_t index : indices) out.push_back(operand(index, device_id)); return out; } - VariantPackSlot * - view(size_t index, int32_t device_id) const { - const Slot &slot = slots_.at(index); - if (!slot.filled) - throw py::value_error("variant-pack slot " + std::to_string(index) + " was not filled by the caller"); - return new VariantPackSlot(slot, device_id); + OperandBuffer * + operand(size_t index, int32_t device_id) const { + const Operand &operand = operands_.at(index); + if (!operand.filled) + throw py::value_error("variant-pack operand " + std::to_string(index) + " was not filled by the caller"); + return new OperandBuffer(operand, device_id); } private: - std::vector slots_; + std::vector operands_; std::vector pointers_; }; -// A slot over memory that is not a caller operand: the regions a plan carves +// A operand over memory that is not a caller operand: the regions a plan carves // out of the workspace. Same type, so a kernel is handed one kind of buffer // whether it came from the caller or from the workspace, and both are read // through the exchange vtable rather than a per-call capsule. -VariantPackSlot * -make_slot(int64_t ptr, std::vector shape, int dtype_code, int dtype_bits, int32_t device_id) { - Slot slot; - slot.data = reinterpret_cast(ptr); - slot.ndim = static_cast(shape.size()); - slot.dtype = DLDataType{static_cast(dtype_code), static_cast(dtype_bits), 1}; - slot.shape = std::move(shape); - slot.filled = true; // stride left empty: a carve is dense by construction - return new VariantPackSlot(std::move(slot), device_id); +OperandBuffer * +make_operand_buffer(int64_t ptr, std::vector shape, int dtype_code, int dtype_bits, int32_t device_id) { + Operand operand; + operand.data = reinterpret_cast(ptr); + operand.ndim = static_cast(shape.size()); + operand.dtype = DLDataType{static_cast(dtype_code), static_cast(dtype_bits), 1}; + operand.shape = std::move(shape); + operand.filled = true; // stride left empty: a carve is dense by construction + return new OperandBuffer(std::move(operand), device_id); } // (pointer, bytes) for a buffer that publishes the vtable, else None. The -// workspace has no uid and no slot, but an engine still bounds-checks its +// workspace has no uid and no operand, but an engine still bounds-checks its // carves against it. py::object read_buffer_extent(py::handle buffer) { @@ -678,7 +678,7 @@ class WorkspaceCarve { throw py::value_error("a carve region is (offset, dtype_code, dtype_bits, shape)"); } int64_t offset = region[0].cast(); - Slot proto; + Operand proto; proto.dtype = DLDataType{region[1].cast(), region[2].cast(), 1}; proto.shape = region[3].cast>(); proto.ndim = static_cast(proto.shape.size()); @@ -691,9 +691,9 @@ class WorkspaceCarve { } } - std::vector + std::vector carve(int64_t base, int64_t nbytes, int32_t device_id) const { - std::vector out; + std::vector out; out.reserve(protos_.size()); for (size_t i = 0; i < protos_.size(); i++) { if (nbytes != 0 && ends_[i] > nbytes) { // 0 = size unknown (bare address) @@ -701,9 +701,9 @@ class WorkspaceCarve { std::to_string(ends_[i]) + ") exceeds the " + std::to_string(nbytes) + "-byte buffer (sizing bug)"); } - Slot slot = protos_[i]; - slot.data = reinterpret_cast(base + offsets_[i]); - out.push_back(new VariantPackSlot(std::move(slot), device_id)); + Operand operand = protos_[i]; + operand.data = reinterpret_cast(base + offsets_[i]); + out.push_back(new OperandBuffer(std::move(operand), device_id)); } return out; } @@ -715,76 +715,76 @@ class WorkspaceCarve { private: std::string owner_; - std::vector protos_; + std::vector protos_; std::vector offsets_; std::vector ends_; }; void init_variant_pack(py::module_ &m) { - auto slot_class = py::class_(m, "VariantPackSlot", R"( + auto operand_class = py::class_(m, "OperandBuffer", R"( One operand of a variant pack, as a DLPack producer. Implements ``__dlpack_c_exchange_api__``, so a consumer reads it through the same C function table it uses for a framework tensor rather than through a capsule built in python. )") - .def("data_ptr", &VariantPackSlot::data_ptr) - .def_property_readonly("shape", &VariantPackSlot::shape) - .def_property_readonly("dtype", &VariantPackSlot::dtype) - .def_property_readonly("nbytes", &VariantPackSlot::nbytes) - .def( - "stride", - [](const VariantPackSlot &self, py::object dim) -> py::object { - if (dim.is_none()) return py::cast(self.stride()); - return py::cast(self.stride_at(dim.cast())); - }, - py::arg("dim") = py::none()) - .def("element_size", &VariantPackSlot::element_size) - .def("numel", &VariantPackSlot::numel) - .def("__len__", &VariantPackSlot::length) - .def("reshape", - [](const VariantPackSlot &self, py::args dims) { - std::vector shape; - if (dims.size() == 1 && py::isinstance(dims[0]) && - !py::isinstance(dims[0])) { - shape = dims[0].cast>(); - } else { - for (auto d : dims) shape.push_back(d.cast()); - } - return self.reshape(std::move(shape)); - }) - .def("permute", - [](const VariantPackSlot &self, py::args axes) { - std::vector order; - if (axes.size() == 1 && py::isinstance(axes[0]) && - !py::isinstance(axes[0])) { - order = axes[0].cast>(); - } else { - for (auto a : axes) order.push_back(a.cast()); - } - return self.permute(order); - }) - .def("contiguous", [](py::object self) { return self; }) - .def("__dlpack_device__", &VariantPackSlot::dlpack_device) - .def("__dlpack__", - &VariantPackSlot::dlpack, - py::kw_only(), - py::arg("stream") = py::none(), - py::arg("max_version") = py::none()); + .def("data_ptr", &OperandBuffer::data_ptr) + .def_property_readonly("shape", &OperandBuffer::shape) + .def_property_readonly("dtype", &OperandBuffer::dtype) + .def_property_readonly("nbytes", &OperandBuffer::nbytes) + .def( + "stride", + [](const OperandBuffer &self, py::object dim) -> py::object { + if (dim.is_none()) return py::cast(self.stride()); + return py::cast(self.stride_at(dim.cast())); + }, + py::arg("dim") = py::none()) + .def("element_size", &OperandBuffer::element_size) + .def("numel", &OperandBuffer::numel) + .def("__len__", &OperandBuffer::length) + .def("reshape", + [](const OperandBuffer &self, py::args dims) { + std::vector shape; + if (dims.size() == 1 && py::isinstance(dims[0]) && + !py::isinstance(dims[0])) { + shape = dims[0].cast>(); + } else { + for (auto d : dims) shape.push_back(d.cast()); + } + return self.reshape(std::move(shape)); + }) + .def("permute", + [](const OperandBuffer &self, py::args axes) { + std::vector order; + if (axes.size() == 1 && py::isinstance(axes[0]) && + !py::isinstance(axes[0])) { + order = axes[0].cast>(); + } else { + for (auto a : axes) order.push_back(a.cast()); + } + return self.permute(order); + }) + .def("contiguous", [](py::object self) { return self; }) + .def("__dlpack_device__", &OperandBuffer::dlpack_device) + .def("__dlpack__", + &OperandBuffer::dlpack, + py::kw_only(), + py::arg("stream") = py::none(), + py::arg("max_version") = py::none()); // The protocol looks the attribute up on the TYPE, and a pybind11 class is // a heap type, so it takes a plain setattr. - PyObject *capsule = PyCapsule_New(&slot_exchange_api(), "dlpack_exchange_api", nullptr); + PyObject *capsule = PyCapsule_New(&buffer_exchange_api(), "dlpack_exchange_api", nullptr); if (capsule == nullptr) throw py::error_already_set(); - if (PyObject_SetAttrString(slot_class.ptr(), "__dlpack_c_exchange_api__", capsule) < 0) { + if (PyObject_SetAttrString(operand_class.ptr(), "__dlpack_c_exchange_api__", capsule) < 0) { Py_DECREF(capsule); throw py::error_already_set(); } Py_DECREF(capsule); - m.def("make_slot", - &make_slot, + m.def("make_operand_buffer", + &make_operand_buffer, py::arg("ptr"), py::arg("shape"), py::arg("dtype_code"), @@ -808,31 +808,31 @@ instead of one per region. .def("carve", &WorkspaceCarve::carve, py::arg("base"), py::arg("nbytes"), py::arg("device_id")) .def("__len__", &WorkspaceCarve::size); - slot_class.attr("view") = slot_class.attr("reshape"); + operand_class.attr("view") = operand_class.attr("reshape"); py::class_(m, "VariantPackNative", R"( The caller's operands, held as DLTensors. -``read_slot`` returns False for a buffer whose type does not implement +``read_operand`` returns False for a buffer whose type does not implement ``__dlpack_c_exchange_api__``; the caller describes that one itself and reports -it through ``set_slot``, so a pack mixing producers costs exactly the sum of +it through ``set_operand``, so a pack mixing producers costs exactly the sum of its parts. )") .def(py::init()) - .def("read_slot", &VariantPackNative::read_slot) + .def("read_operand", &VariantPackNative::read_operand) .def("read_from", &VariantPackNative::read_from) .def("first_unfilled", &VariantPackNative::first_unfilled) - .def("set_slot", &VariantPackNative::set_slot) - .def("override_slot", &VariantPackNative::override_slot) - .def("skip_slot", &VariantPackNative::skip_slot) - .def("slot_contiguous", &VariantPackNative::slot_contiguous) + .def("set_operand", &VariantPackNative::set_operand) + .def("override_operand", &VariantPackNative::override_operand) + .def("skip_operand", &VariantPackNative::skip_operand) + .def("operand_contiguous", &VariantPackNative::operand_contiguous) .def("is_filled", &VariantPackNative::is_filled) .def("pointer", &VariantPackNative::pointer) .def("shape", &VariantPackNative::shape) .def("stride", &VariantPackNative::stride) .def("dtype", &VariantPackNative::dtype) - .def("view", &VariantPackNative::view) - .def("views", &VariantPackNative::views) + .def("operand", &VariantPackNative::operand) + .def("operands", &VariantPackNative::operands) .def_property_readonly("address", &VariantPackNative::pointer_array) .def("__len__", &VariantPackNative::size) .def("all_contiguous", [](const VariantPackNative &self) { diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index 86534aced..96ec4aaef 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -3,7 +3,7 @@ """The build-time recipe, and the straight line lowered from it. -``_lower`` emits a call path with the recipe's constants inlined; ``run_views`` +``_lower`` emits a call path with the recipe's constants inlined; ``launch`` interprets the same recipe. That is a compiler beside its interpreter, and the two can drift -- an earlier hand-written version of the emitted path lost the operand batch check and pinned an fp4 output at N instead of N/2, both of which @@ -124,7 +124,7 @@ def test_which_flavors_lower(build, lowered): than growing a branch per flavor. """ compiled = jit_from_cudnn_graph(build()) - assert (compiled.launch is not None) is lowered + assert (compiled.lowered is not None) is lowered @requires_sm100 @@ -146,15 +146,15 @@ def test_public_execute_takes_the_lowered_path(): # --- the differential ------------------------------------------------------- -def _views(compiled, a, b, c): +def _bound_buffers(compiled, a, b, c): resolved = resolve_variant_pack(vp(compiled, a, b, c), compiled.binding) return [resolved[id(t)] for t in compiled.binding.bound_tensors()] -def _verdict(run, views, c): +def _verdict(run, operands, c): c.zero_() try: - run(views, stream=None) + run(operands, stream=None) except ValueError: return "rejected", None torch.cuda.synchronize() @@ -220,12 +220,12 @@ def test_lowered_and_interpreted_agree(case): read. """ compiled = jit_from_cudnn_graph(_plain_graph()) - if compiled.launch is None: + if compiled.lowered is None: pytest.skip("this build does not lower (no tvm-ffi front door)") a, b, c = case() - fast, fast_out = _verdict(compiled.launch, _views(compiled, a, b, c), c) - slow, slow_out = _verdict(compiled.run_views, _views(compiled, a, b, c), c) + fast, fast_out = _verdict(compiled.lowered, _bound_buffers(compiled, a, b, c), c) + slow, slow_out = _verdict(compiled.launch, _bound_buffers(compiled, a, b, c), c) assert fast == slow, f"lowered says {fast}, interpreted says {slow}" if fast == "ran": torch.testing.assert_close(fast_out, slow_out, atol=0, rtol=0) @@ -233,6 +233,65 @@ def test_lowered_and_interpreted_agree(case): torch.testing.assert_close(fast_out.float(), ref, atol=2e-1, rtol=2e-2) +def _matmul_on(batch, m, n, k, want_frost, a, b, c): + """Build the plain graph, pin a backend or a FROST plan, run it.""" + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[batch, m, k], stride=[m * k, k, 1]) + B = g.tensor(name="B", dim=[batch, k, n], stride=[k * n, 1, k]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + hits = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id) == want_frost] + if not hits: + return None + g.select_plan(hits[0]) + g.check_support() + g.build_plans() + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute({A: a, B: b, C: c}, ws) + torch.cuda.synchronize() + return c + + +@requires_sm100 +@pytest.mark.parametrize("batch", (1, 2), ids=("b1", "b2")) +@pytest.mark.parametrize("m,n,k", [(m, n, k) for m in (1, 128) for n in (1, 128) for k in (1, 128)], ids=str) +def test_a_degenerate_extent_is_refused_or_matches_the_backend(monkeypatch, batch, m, n, k): + """An extent of 1 leaves its axis's stride free, so two majors can look alike. + + ``(batch, M, 1)`` k-major carries stride 1 on axis 1 AND axis 2, and nothing + in the description says which one the kernel should read as K. What keeps + that from mattering is the TMA rule: the contiguous extent must divide + ``128 // bits``, whose smallest value is 4, and 1 divides none of them -- so + a unit contiguous extent never reaches a launch. This asserts the property + that argument implies rather than the argument: every degenerate shape is + either refused or agrees with the backend. + """ + monkeypatch.setenv("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "1") + torch.manual_seed(0) + a = torch.randn(batch, m, k, dtype=torch.bfloat16, device="cuda") + b = torch.randn(batch, n, k, dtype=torch.bfloat16, device="cuda") + + got = {} + for want_frost in (False, True): + c = torch.zeros(batch, m, n, dtype=torch.bfloat16, device="cuda") + try: + # None = no plan of that kind. For FROST that IS a refusal, and the + # one every K == 1 shape takes: the graph-time gate applies the same + # TMA rule, so the degenerate contiguous extent never gets a plan. + ran = _matmul_on(batch, m, n, k, want_frost, a, b, c) + got[want_frost] = "refused" if ran is None else ran + except (ValueError, NotImplementedError, cudnn.cudnnGraphNotSupportedError): + got[want_frost] = "refused" + if got[False] == "refused": + pytest.skip("the backend has nothing to compare against for this shape") + if got[True] == "refused": + return # refusing is always allowed; computing something else is not + torch.testing.assert_close(got[True].float(), got[False].float(), atol=2e-1, rtol=2e-2) + + @requires_sm100 def test_an_operand_reporting_the_declaration_agrees_with_the_backend(monkeypatch): """At N == K the two axis orders are the SAME shape and the SAME stride. @@ -325,8 +384,8 @@ def test_operand_batch_is_checked(): compiled = jit_from_cudnn_graph(_plain_graph(batch=2)) a, b, c = _operands(batch=2) one_batch_b = b[:1].contiguous() - for run in (compiled.launch, compiled.run_views): + for run in (compiled.lowered, compiled.launch): if run is None: continue with pytest.raises(ValueError): - run(_views(compiled, a, one_batch_b, c), stream=None) + run(_bound_buffers(compiled, a, one_batch_b, c), stream=None) From ae9342397875fda15c591838b44bad99922ad1f7 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 12:10:55 -0700 Subject: [PATCH 4/9] Lower every gemm flavor off the recipe, not just plain The lowered call path served one graph shape: single-GEMM, one dense output, no aux, no scale factors, no reduction seed. Every other flavor fell back to the interpreter at 39-50 us of host time per execute, against 17.7 for the shape that lowered -- and for a single-kernel op that gap is most of what the caller pays. What actually differed between the flavors was which buffers the launch passes and in what order, so that is a recipe field now: `arg_plan`, one (operand index, aux ref) per positional argument after `problem_size`. Three smaller fields carry the rest -- `stride_ins` (whose permuted strides ride in `problem_size`), `shared_layout` (block-scale multi-GEMM collapses its A operands to one stride triple and requires the others to match), and `seeds` (a reduction's identity, packed as its output dtype). With the launch shape as data the emitted body is one loop over flat tuples, and the hand-unrolled straight line is deleted. Measured against it on plain gemm the loop costs 12% and saves one closure body per operand shape; source codegen off the same table is how to buy that back for every flavor at once rather than only for the one worth hand-writing. 256x256x128 bf16, min over 25 reps of a 64-call burst from a drained queue: flavor before now plain 17.73 19.95 epilogue (relu) 17.69 19.89 aux (bias tensor) 38.70 22.25 2 dense outputs 43.89 22.57 reduction output 49.90 26.29 multi-gemm 42.58 22.38 Reduction sits ~4 us above the rest because it still pays one cuMemsetD32Async per tap; that goes when the kernel seeds itself. Still declined, each by a named guard: no tvm-ffi front door, a plan that wants workspace, a norm2 output (its post-kernel sqrt_ is a device operation the engine does not own), and a multi-GEMM with no dense output -- the two multi-GEMM launchers read the batch off cs[0] where the single-GEMM one branches on output_specs, and declining keeps one clean rule in the table instead of that quirk. `device` moves into the recipe alongside `workspace_bytes`, so the only thing the closure captures that is not a table entry is the kernel itself. Tests: test_execute_recipe.py runs all six flavors through both paths and requires the same numbers. A reduction output is compared at the noise floor instead of bit-exactly, because its taps land through cross-CTA float atomics -- measured, the same path run six times spreads 0.0049 while the two paths differ by 0.00024. test_public_execute_flavors.py gains aux, two-output and multi-GEMM cases, so the newly-lowered flavors are exercised through graph.execute() and not only through the direct call. Co-Authored-By: Claude Opus 5 (1M context) --- docs/python_graph_and_execution_backends.md | 19 +- python/cudnn/frost/README.md | 14 +- python/cudnn/gemm/frost/compiler.py | 184 ++++++++++-------- python/cudnn/gemm/frost/engine.py | 4 +- python/cudnn/gemm/frost/recipe.py | 74 +++++-- test/python/gemm/frost/test_execute_recipe.py | 179 ++++++++++++++--- .../gemm/frost/test_public_execute_flavors.py | 59 ++++++ 7 files changed, 399 insertions(+), 134 deletions(-) diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index ede9859ac..b2fa56632 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -152,7 +152,7 @@ An engine owns its internals, and this section does not change that. It exists because the default outcome is expensive: an engine that re-derives its per-call facts lands around **40 µs of host time per execute**, and for a single-kernel op that is most of what the caller pays. The same kernel with those facts read -once is **17.5**. Both numbers are `frost_gemm` at 256×256×128 bf16, host +once is **20**. Both numbers are `frost_gemm` at 256×256×128 bf16, host enqueue, min over 25 reps of a 64-call burst from a drained queue. The budget it has to fit in, all measured on SM100: @@ -173,9 +173,10 @@ all fixed when the kernel compiled. M/N/K, strides and pointers arrive per call. Read the first set into a table at build (`gemm/frost/recipe.py` is the worked example) and let the call read the table. That alone is 44 → 35. -**Then lower the table for the shape you actually run.** Emitting one closure -per plan, with the table's constants inlined, is 35 → 17.5. Two rules make that -safe: +**Then lower the table into one closure per plan**, with its constants captured +and the operand structure flattened into the loop headers, so the call does no +attribute lookup and takes no branch the build already settled. That is 35 → 20. +Two rules make it safe: - **The lowered path never raises.** Anything it is not certain of it hands to the interpreting path, which serves every flavor and owns every rejection @@ -186,6 +187,16 @@ safe: fact lives exactly once, and the tests that matter are against intended semantics, at the shapes where two encodings coincide. +**A loop over a flat table gets almost all of it, so do not hand-unroll per +flavor.** Measured three ways on the same plan and buffers: interpreting the +table 35.8, looping over it flattened 19.7, a hand-written straight line with the +structure unrolled 17.5. The loop is worth 45%; unrolling adds 12% and costs one +closure body per operand shape — six flavors, six bodies to keep in agreement. +One loop over `arg_plan` (the launch argument order as data) serves aux, extra +outputs, multi-GEMM and block scale at 22 µs each, down from 39–50. Source +codegen off the same table is how to buy the last 12% back later, for every +flavor at once rather than for the one that was worth hand-writing. + This is a pattern to copy, not a framework to import. Sharing the code across engines would couple their kernels' ABIs, which is the thing engine autonomy buys; sharing the shape of the solution costs nothing. diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index 84dfc1a71..059ea95a2 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -239,12 +239,14 @@ CompiledPlan.execute(graph, uid_to_data, ctx) (the hot path) operand's role and major, each output's shape rule and required alignment, which outputs are reductions -- all settled when the kernel compiled, and deciding them again per call is most of what a python execute path costs - (measured: 44 -> 18 us for one gemm). `gemm/frost/recipe.py` is the worked - example: one table, read by an interpreter that serves every flavor and by a - straight line lowered for the common one. Emitting a second call path per - flavor instead is how the two disagree, so the lowered one never raises -- - anything it is unsure of it hands back to the interpreter, which owns every - rejection message. + (measured: 40-50 -> 20-22 us for one gemm, across six flavors). + `gemm/frost/recipe.py` is the worked example: one table, read by an + interpreter that walks the operand structure and by a closure that captures + the table and loops over it flat. Even what the kernel's parameter list looks + like is a table entry (`arg_plan`), which is why one loop serves every flavor + -- a call path per flavor is how two of them disagree. The lowered one never + raises: anything it is unsure of it hands back to the interpreter, which owns + every rejection message. - **`ExecutionContext` carries handle, stream and workspace explicitly.** No engine may hard-code a stream, reach into private graph state, or allocate hidden workspace. `uid_to_data` is the caller's variant pack (tensor uid -> diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 6c974341b..87eab914d 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -60,7 +60,6 @@ def _as_custream(stream): AX_MN, CONST, FROM_M, - FROM_N, KERNEL_AXES, REDUCTION_INIT_VALUE, _output_rule, @@ -1889,61 +1888,63 @@ def run_resolved(self, resolved, stream=None): return (self.lowered or self.launch)(operands, stream=stream) def _lower(self): - """The recipe as one straight line, or None for a shape it does not emit. + """The recipe as one loop over flat tuples, or None for a graph it does not serve. Every branch the general path takes -- multi-GEMM or not, block scale or not, how many outputs, which are reductions, which axis of each operand - carries M/N/K -- was settled when the kernel compiled. Taking them again - per call is what stands between this path and its floor: the - load-bearing validation measures 1.1 us against 21 for the walk that - rebuilds it. + carries M/N/K, what order the kernel's parameters come in -- was settled + when the kernel compiled. Taking them again per call is most of what a + python execute path costs: 44 us for the walk that rebuilt them, 35 + reading a recipe through its objects, 20 here. + + What is left per call is the same loop for every flavor, over tuples + flat enough that the body does no attribute lookup and calls nothing it + does not have to. Measured against a hand-unrolled straight line for the + plain flavor, the loop gives back 12% (19.7 us against 17.5) and costs + one body instead of one per shape; source codegen off this same table is + how to get that 12% back for every flavor at once. The emitted body never raises. Anything it is not sure of it hands to ``launch``, which serves every flavor and owns every rejection message -- so this can only ever accept a subset of what the general path accepts, and there is no second set of error strings to drift. - - Emitting another flavor means widening the recipe, not copying this - body: a second body is the drift the recipe exists to prevent. """ r = self.recipe - if r is None or not _TVM_FFI_OK or r.multi_gemm or r.block_size or r.aux or r.workspace_bytes: - return None - if len(r.inputs) != 2 or len(r.outputs) != 1: + if r is None or not _TVM_FFI_OK or r.workspace_bytes: return None - out = r.outputs[0] - # A raw output is a reduction or a quant scale: both want a pre-kernel - # seed and one wants a post-kernel sqrt, which are device operations the - # engine does not own yet (see _initialize_reduction_outputs). - if out.raw or out.init is not None or out.sqrt: + # norm2 takes a square root through the caller's buffer after the kernel; + # that is a device operation the engine does not own yet, and neither is + # seeding a strided reduction output (checked per call, below). + if any(o.sqrt for o in r.outputs): return None - rule = out.rule - if tuple(s for s, _ in rule) != (CONST, FROM_M, FROM_N): + # Scale factors come with a block size to size their blob against, and a + # multi-GEMM's batch with a dense output to read it off. + if bool(r.sf) != bool(r.block_size) or (r.multi_gemm and not r.has_output_specs): return None - - a, b = r.a, r.b # A buffer reporting the declaration is read in the graph's axis order, # which this body does not serve. The stride guard below tells that # apart -- unless the declaration itself would satisfy the guard, which # only a unit extent can arrange, and which is settled here not per call. - for op in (a, b): + for op in r.inputs: declared_stride = op.declared_layout[1] if op.dc != op.kc and declared_stride and declared_stride[op.kc] == 1: return None - ai, bi, ci = a.index, b.index, out.index - # kc is both the axis whose stride must be 1 and the axis whose extent - # enters the TMA rule -- they are the same axis by definition of major. - a_kc, b_kc = a.kc, b.kc - a_mod, b_mod = a.modulus, b.modulus - a_batch, b_batch = a.batch, b.batch - a_kpack, b_kpack = a.kpack, b.kpack - out_batch, out_ndiv, out_align = rule[0][1], rule[2][1], out.align - # Every batch extent the launch could read is pinned by a check above - # it, so the kernel's batch is settled here rather than re-read. - batch = out_batch if r.has_output_specs else max(a_batch, b_batch) - device = self.device - launchable = self._launchable - general = self.launch + + # The recipe flattened into the loop headers: everything the body reads + # is unpacked by the `for`, so it costs no attribute lookup. `kc` is both + # the axis whose stride must be 1 and the axis whose extent enters the + # TMA rule -- the same axis, by the definition of major. + stride_ins = r.stride_ins + ins = tuple((op.index, op.kc, op.modulus, op.batch, op.kpack, op.is_b, i in stride_ins) for i, op in enumerate(r.inputs)) + outs = tuple((o.index, *(x for axis in o.rule for x in axis), o.align) for o in r.outputs) + auxs = tuple((x.index, x.align) for x in r.aux) + sfs = tuple((s.index, s.is_a, r.inputs[s.operand_at].batch) for s in r.sf) + args = r.arg_plan + shared, seeds = r.shared_layout, r.seeds + ai, bi, a_kpack = r.a.index, r.b.index, r.a.kpack + block_size, batch = r.block_size, r.batch + device, launchable, general = r.device, self._launchable, self.launch + fill_word, is_contiguous = buffers.fill_word_async, buffers.is_contiguous def lowered(operands, graph_order=None, stream=None): _check_plan_device(device) @@ -1951,48 +1952,77 @@ def lowered(operands, graph_order=None, stream=None): # An operand described from the graph is in the graph's axis # order; this body reads the caller's. return general(operands, graph_order, stream=stream) - av, bv, cv = operands[ai], operands[bi], operands[ci] - a_sh, b_sh, c_sh = av.shape, bv.shape, cv.shape - if len(a_sh) != 3 or len(b_sh) != 3 or len(c_sh) != 3: - return general(operands, stream=stream) - a_st, b_st, c_st = av.stride(), bv.stride(), cv.stride() + a_sh, b_sh = operands[ai].shape, operands[bi].shape + if len(a_sh) != 3 or len(b_sh) != 3: + return general(operands, None, stream=stream) m, n, k = a_sh[1], b_sh[1], a_sh[2] * a_kpack - if ( - a_st[a_kc] != 1 - or b_st[b_kc] != 1 - or a_sh[a_kc] % a_mod - or b_sh[b_kc] % b_mod - or a_sh[0] != a_batch - or b_sh[0] != b_batch - or b_sh[2] * b_kpack != k - or c_sh[0] != out_batch - or c_sh[1] != m - or c_sh[2] != n // out_ndiv - or _pow2_floor(av.data_ptr()) < 16 - or _pow2_floor(bv.data_ptr()) < 16 - or tensor_alignment(tuple(c_sh), tuple(c_st), cv.element_size(), ptr=cv.data_ptr()) < out_align - ): - return general(operands, stream=stream) - problem_size = ( - m, - n, - k, - batch, - # permute(1, 2, 0) relabels axes, so the kernel's strides are - # that permutation of the ones just read -- no second read. - a_st[1], - a_st[2], - a_st[0], - b_st[1], - b_st[2], - b_st[0], - c_st[1], - c_st[2], - c_st[0], + # permute(1, 2, 0) relabels axes, so the strides the kernel wants are + # that rotation of the ones each check already read -- no second read. + problem = [m, n, k, batch] + for idx, kc, mod, ebatch, kpack, is_b, takes_stride in ins: + v = operands[idx] + sh, st = v.shape, v.stride() + if ( + len(sh) != 3 + or st[kc] != 1 + or sh[kc] % mod + or sh[0] != ebatch + or sh[1] != (n if is_b else m) + or sh[2] * kpack != k + or _pow2_floor(v.data_ptr()) < 16 + ): + return general(operands, None, stream=stream) + if takes_stride: + problem += (st[1], st[2], st[0]) + # Each output axis is a constant, M, or N over a divisor (fp4 packs + # two along N) -- the rule the build recorded, read back per axis. + for idx, k0, v0, k1, v1, k2, v2, align in outs: + v = operands[idx] + sh, st = v.shape, v.stride() + if ( + len(sh) != 3 + or sh[0] != (v0 if k0 == CONST else m if k0 == FROM_M else n // v0) + or sh[1] != (v1 if k1 == CONST else m if k1 == FROM_M else n // v1) + or sh[2] != (v2 if k2 == CONST else m if k2 == FROM_M else n // v2) + or tensor_alignment(tuple(sh), tuple(st), v.element_size(), ptr=v.data_ptr()) < align + ): + return general(operands, None, stream=stream) + problem += (st[1], st[2], st[0]) + for idx, align in auxs: + v = operands[idx] + if tensor_alignment(tuple(v.shape), tuple(v.stride()), v.element_size(), ptr=v.data_ptr()) < align: + return general(operands, None, stream=stream) + if sfs: + k4 = ((k // block_size) + 3) // 4 + for idx, is_a, sf_batch in sfs: + v = operands[idx] + count = int(v.numel()) + if ( + _pow2_floor(v.data_ptr()) < 16 + or count != 1 + sum((int(s) - 1) * int(st) for s, st in zip(v.shape, v.stride())) + or count * v.element_size() < 512 * k4 * (((m if is_a else n) + 127) // 128) * sf_batch + ): + return general(operands, None, stream=stream) + for lead, followers in shared: + st = tuple(operands[lead].stride()) + for j in followers: + if tuple(operands[j].stride()) != st: + return general(operands, None, stream=stream) + if seeds: + # Seeding is a write, so nothing above may still reject: check + # every reduction output first, then fill. + for idx, _word in seeds: + v = operands[idx] + if not is_contiguous(tuple(v.shape), tuple(v.stride())): + return general(operands, None, stream=stream) + for idx, word in seeds: + v = operands[idx] + fill_word(v.data_ptr(), int(v.numel()), word, stream) + return launchable( + tuple(problem), + *(operands[i].permute(1, 2, 0) if ref is None else _reshape_aux_to_fake(operands[i], ref) for i, ref in args), + stream=_as_custream(stream), ) - # One output and no aux: the TMA-store and plain argument orders - # coincide, and a shape where they do not is not emitted here. - return launchable(problem_size, av.permute(1, 2, 0), bv.permute(1, 2, 0), cv.permute(1, 2, 0), stream=_as_custream(stream)) return lowered @@ -2007,8 +2037,8 @@ def launch(self, operands, graph_order=None, stream=None): ``graph_order`` is the pack's per-view "this operand's layout is the graph's, not the caller's", or None when they are all the caller's. """ - _check_plan_device(self.device) recipe = self.recipe + _check_plan_device(recipe.device) mnk, axes = recipe.problem(operands, graph_order) check_shapes(recipe, operands, mnk, axes) check_alignment(recipe, operands, axes) diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index 8f1f3c871..56ad6417e 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -34,8 +34,8 @@ def __init__(self, compiled): self._operand_indices = None # Which call path this plan uses is a property of the compiled kernel, # so it is chosen here and not re-asked per execute. ``lowered`` is the - # straight line the recipe lowers to when the kernel is a shape it - # emits; ``launch`` serves everything else. + # closure the recipe is captured into when the kernel is a graph it + # serves; ``launch`` interprets the same recipe for everything else. self._lowered = getattr(compiled, "lowered", None) self._launch = self._lowered or getattr(compiled, "launch", None) diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py index 291fe1d7e..cfb06dc39 100644 --- a/python/cudnn/gemm/frost/recipe.py +++ b/python/cudnn/gemm/frost/recipe.py @@ -4,19 +4,23 @@ """What one compiled gemm needs per call, settled when it compiles. Operand roles, majors, packing factors, alignment requirements, output shape -rules and which outputs are reductions are all fixed by the time cute hands back -a launchable. A call carries M, N, K, the strides and the pointers, and nothing -else. This module writes the first set down, so that neither the interpreted -path nor the straight line lowered from it re-derives them per call. - -The two consumers read the same recipe but do not share a body: -``CompiledFusedGemm.launch`` interprets it (through :func:`check_shapes` and -:func:`check_alignment`), and ``CompiledFusedGemm.lowered`` is the straight line -``_lower`` emits with the constants inlined. That is a compiler beside its -interpreter, kept honest the way those always are -- ``test_execute_recipe.py`` -runs both over the same accepts and rejects and requires the same answer. What -it CANNOT catch is a misconception they share, which is how the axis-order bug -survived it. +rules, which outputs are reductions and what order the kernel takes its +parameters in are all fixed by the time cute hands back a launchable. A call +carries M, N, K, the strides and the pointers, and nothing else. This module +writes the first set down, so that neither of the two call paths re-derives them. + +They read the same recipe but do not share a body: ``CompiledFusedGemm.launch`` +interprets it (through :func:`check_shapes` and :func:`check_alignment`) by +walking the operand structure, and ``CompiledFusedGemm.lowered`` is the closure +``_lower`` captures it into, where the same walk is a loop over tuples flat +enough to need no attribute lookup. That is a compiler beside its interpreter, +kept honest the way those always are -- ``test_execute_recipe.py`` runs both +over the same accepts and rejects and requires the same answer. What it CANNOT +catch is a misconception they share, which is how the axis-order bug survived it. + +The field that makes one loop serve six flavors is :attr:`GemmRecipe.arg_plan`: +what differs between plain, aux, multi-output, multi-GEMM and block scale is +only which buffers the launch passes and in what order, so that is data. """ from __future__ import annotations @@ -24,6 +28,8 @@ from dataclasses import dataclass from typing import Any +from cudnn.frost.buffers import init_word + from .dtypes import DTYPE_BYTES, _aux_align_reqs, _output_align_reqs, _pow2_floor, tensor_alignment from .fusion_ir import FusionChain @@ -186,6 +192,21 @@ class GemmRecipe: block_size: "int | None" workspace_bytes: int multi_gemm: bool + device: int # the GPU whose SMEM depth / cluster count / SM the kernel is baked for + batch: int # the kernel's batch extent, pinned by the checks above the launch + # The launch call as data. ``arg_plan`` is one entry per positional argument + # after ``problem_size``: an operand index, plus the aux TensorRef whose fake + # shape it reshapes to (None -- every other role -- means permute(1, 2, 0)). + # ``stride_ins`` gives the positions in ``inputs`` whose permuted strides ride + # in ``problem_size``, ahead of every output's. This is the one field that says + # how the six flavors differ, which is why they differ in a table and not in + # six launchers. + arg_plan: tuple + stride_ins: tuple + # ``(leader, followers)`` groups whose strides the launch collapses to one. + shared_layout: tuple + # ``(output index, identity as a dtype-packed word)`` per reduction output. + seeds: tuple @property def a(self) -> Operand: @@ -404,13 +425,14 @@ def build(compiled) -> GemmRecipe: out_reqs = _output_align_reqs(chain, compiled.use_tma_store, vec_bytes=compiled.vec_bytes_epi) aux_reqs = _aux_align_reqs(chain, vec_bytes=compiled.vec_bytes_epi) - outputs = [] + outputs, seeds = [], [] for i, (spec, t) in enumerate(zip(chain.outputs, binding.outputs)): init, sqrt = None, False if spec.is_reduction: red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] init = REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] sqrt = red.mode == "norm2" + seeds.append((order[id(t)], init_word(red.compute_dtype, init))) outputs.append( Output( index=order[id(t)], @@ -437,6 +459,24 @@ def build(compiled) -> GemmRecipe: for item in (*inputs, *outputs, *aux, *sf): roles[item.index] = item.role + # The kernel's signature, in the order the launchers pass it: every distinct + # A, every distinct B, their scale factors, then the outputs and the aux -- + # except under a TMA-store epilogue, where the single dense output binds the + # template's trailing TMA-only parameter and so goes last. + heads = [(op.index, None) for op in inputs] + [(s.index, None) for s in sf] + outs = [(o.index, None) for o in outputs] + auxs = [(x.index, x.ref) for x in aux] + tma = bool(compiled.use_tma_store) and not chain.is_multi_gemm + arg_plan = tuple(heads + (auxs + outs[:1] if tma else outs + auxs)) + + # Block-scale multi-GEMM sends ONE A and ONE B stride triple and requires the + # rest to match it; every other flavor sends each operand's own. + grouped = bool(chain.is_multi_gemm and compiled.block_scale) + stride_ins = (0, na) if grouped else tuple(range(len(inputs))) + shared_layout = () + if grouped: + shared_layout = tuple((group[0].index, tuple(op.index for op in group[1:])) for group in (inputs[:na], inputs[na:]) if len(group) > 1) + return GemmRecipe( inputs=tuple(inputs), outputs=tuple(outputs), @@ -449,4 +489,10 @@ def build(compiled) -> GemmRecipe: block_size=chain.block_scale.block_size if compiled.block_scale else None, workspace_bytes=int(getattr(compiled, "workspace_bytes", 0) or 0), multi_gemm=bool(chain.is_multi_gemm), + device=int(compiled.device), + batch=int(outputs[0].rule[0][1] if chain.output_specs else max(mm.a_batch, mm.b_batch)), + arg_plan=arg_plan, + stride_ins=stride_ins, + shared_layout=shared_layout, + seeds=tuple(seeds), ) diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index 96ec4aaef..46f5d69b8 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -1,27 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The build-time recipe, and the straight line lowered from it. - -``_lower`` emits a call path with the recipe's constants inlined; ``launch`` -interprets the same recipe. That is a compiler beside its interpreter, and the -two can drift -- an earlier hand-written version of the emitted path lost the -operand batch check and pinned an fp4 output at N instead of N/2, both of which -made it accept or reject calls the general path did not. - -So the differential below is the point of this file: every case runs through -BOTH entry points and the two must return the same verdict and the same numbers. -The rest asserts that the fast path is the one a public ``execute()`` actually -takes, and that the flavors it declines are declined on purpose. +"""The build-time recipe, and the closure lowered from it. + +``_lower`` captures the recipe into a call path that loops over it flat; +``launch`` interprets the same recipe by walking the operand structure. That is +a compiler beside its interpreter, and the two can drift -- an earlier +hand-written version of the emitted path lost the operand batch check and pinned +an fp4 output at N instead of N/2, both of which made it accept or reject calls +the general path did not. + +So the differentials below are the point of this file: every case and every +flavor runs through BOTH entry points, and the two must return the same verdict +and the same numbers. The rest asserts that the fast path is the one a public +``execute()`` actually takes, and that what it declines is declined on purpose. """ from __future__ import annotations +from dataclasses import replace from types import SimpleNamespace import pytest import torch -from gemm_test_utils import requires_sm100, vp +from gemm_test_utils import requires_sm100 import cudnn import cudnn.gemm.frost # noqa: F401 — installs the cudnn.pygraph recorder hook @@ -106,30 +108,80 @@ def _aux_graph(): return g -@requires_sm100 -@pytest.mark.parametrize( - "build,lowered", - ( - (_plain_graph, True), - (_reduction_graph, False), - (_aux_graph, False), - ), - ids=("plain", "reduction", "aux"), +def _epilogue_graph(): + g = _plain_graph() + mm = [t for t in g._nodes[-1].outputs.values()][0] + mm.set_output(False) + g.relu(input=mm, name="relu").set_output(True).set_data_type(BF16) + return g + + +def _two_output_graph(): + g = _plain_graph() + mm = [t for t in g._nodes[-1].outputs.values()][0] + g.relu(input=mm, name="relu").set_output(True).set_data_type(BF16) + return g + + +def _multi_gemm_graph(): + 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]) + B0 = g.tensor(name="B0", dim=[1, K, N], stride=[K * N, 1, K]) + B1 = g.tensor(name="B1", dim=[1, K, N], stride=[K * N, 1, K]) + Y = g.add(a=g.matmul(A=A, B=B0, name="mm0"), b=g.matmul(A=A, B=B1, name="mm1"), name="sum") + Y.set_output(True).set_data_type(BF16) + return g + + +# Each entry is (build, how many distinct B operands, how many outputs, how many aux). +_FLAVORS = ( + (_plain_graph, 1, 1, 0), + (_epilogue_graph, 1, 1, 0), + (_aux_graph, 1, 1, 1), + (_two_output_graph, 1, 2, 0), + (_reduction_graph, 1, 2, 0), + (_multi_gemm_graph, 2, 1, 0), ) -def test_which_flavors_lower(build, lowered): - """A declined flavor is declined by a named recipe field, not by accident. +_FLAVOR_IDS = ("plain", "epilogue", "aux", "two_outputs", "reduction", "multi_gemm") - Reductions want a pre-kernel seed and aux wants a fake-shape reshape; both - are work the emitted line does not carry yet, so it hands them over rather - than growing a branch per flavor. + +@requires_sm100 +@pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) +def test_which_flavors_lower(build, n_b, n_out, n_aux): + """Every flavor lowers, because what differs between them is a table entry. + + Aux, extra outputs, a reduction's seed and a multi-GEMM's operand list all + used to be branches in a launcher, so the emitted path served only the one + shape with none of them. They are ``arg_plan`` and ``seeds`` now, which the + same loop reads -- so the list below is a list of table shapes, not of code + paths, and adding to it is data. """ compiled = jit_from_cudnn_graph(build()) - assert (compiled.lowered is not None) is lowered + assert compiled.lowered is not None + r = compiled.recipe + assert (len(r.inputs) - 1, len(r.outputs), len(r.aux)) == (n_b, n_out, n_aux) + # One launch argument per bound operand: nothing dropped, nothing passed twice. + assert len(r.arg_plan) == len(r.inputs) + len(r.outputs) + len(r.aux) + len(r.sf) + + +@requires_sm100 +def test_a_post_kernel_finalize_is_declined(): + """``norm2`` takes a square root through the caller's buffer after the kernel. + + The backend refuses that reduction while the graph is being lowered, so it + never reaches a plan -- but the recipe records ``sqrt`` and the lowered path + declines on it, because the alternative is a device operation the engine + does not own and would have to borrow off whatever the caller passed. + """ + compiled = jit_from_cudnn_graph(_plain_graph()) + assert compiled.lowered is not None + compiled.recipe = replace(compiled.recipe, outputs=(replace(compiled.recipe.outputs[0], sqrt=True),)) + assert compiled._lower() is None @requires_sm100 def test_public_execute_takes_the_lowered_path(): - """The straight line is what a user's ``execute()`` runs, not a side door.""" + """The lowered path is what a user's ``execute()`` runs, not a side door.""" g = _plain_graph() g.validate() g.build_operation_graph() @@ -146,9 +198,74 @@ def test_public_execute_takes_the_lowered_path(): # --- the differential ------------------------------------------------------- +def _bind(compiled, a_bufs, b_bufs, out_bufs, aux_bufs=()): + """The operand list in bound-tensor order, which is what both paths take.""" + bd = compiled.binding + pack = {} + for role, bufs in ((bd.a_operands, a_bufs), (bd.b_operands, b_bufs), (bd.outputs, out_bufs), (bd.aux, aux_bufs)): + pack.update(zip(role, bufs)) + resolved = resolve_variant_pack(pack, bd) + return [resolved[id(t)] for t in bd.bound_tensors()] + + def _bound_buffers(compiled, a, b, c): - resolved = resolve_variant_pack(vp(compiled, a, b, c), compiled.binding) - return [resolved[id(t)] for t in compiled.binding.bound_tensors()] + return _bind(compiled, [a], [b], [c]) + + +def _buffers_for(compiled): + """One buffer per bound tensor, sized from the graph's own declaration. + + Every flavor's operands follow from the recipe, so the differential below + does not need a hand-written variant pack per flavor -- which is the same + reason the launch path does not need a launcher per flavor. + """ + r = compiled.recipe + bufs = [None] * len(compiled.binding.bound_tensors()) + for op in r.inputs: + rows = N if op.is_b else M + bufs[op.index] = torch.randn(op.batch, rows, K // op.kpack, dtype=torch.bfloat16, device="cuda") + for out in r.outputs: + dtype = torch.float32 if out.raw else torch.bfloat16 + bufs[out.index] = torch.zeros(expected_shape(out.rule, M, N), dtype=dtype, device="cuda") + for x in r.aux: + bufs[x.index] = torch.randn(tuple(int(d) for d in x.ref.dim), dtype=torch.float32, device="cuda") + assert all(b is not None for b in bufs) + return bufs + + +@requires_sm100 +@pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) +def test_every_flavor_agrees_with_the_interpreter(build, n_b, n_out, n_aux): + """One loop serves six flavors, so one of them drifting is the failure mode. + + The interpreter is the reference: it reads the same recipe but assembles the + launch by walking the operand structure, which is the thing the lowered path + replaced with a table. Bit-exact is the bar -- the two issue the same kernel + with the same arguments or they do not agree. + + Except for a reduction output, which is not bit-reproducible against ITSELF: + the taps land through cross-CTA float atomics, so the order varies with + scheduling. Measured on this shape, the same path run six times spreads + 0.0049 while the two paths differ by 0.00024 -- twenty times smaller than + the noise, so the tolerance below is the noise floor and not a slackened bar. + """ + compiled = jit_from_cudnn_graph(build()) + if compiled.lowered is None: + pytest.skip("this build does not lower (no tvm-ffi front door)") + operands = _buffers_for(compiled) + outs = compiled.recipe.outputs + + runs = [] + for run in (compiled.lowered, compiled.launch): + for o in outs: + operands[o.index].zero_() + run(operands, stream=None) + torch.cuda.synchronize() + runs.append([operands[o.index].clone() for o in outs]) + for o, fast, slow in zip(outs, *runs): + exact = o.init is None + torch.testing.assert_close(fast, slow, atol=0 if exact else 1e-2, rtol=0 if exact else 1e-5) + assert runs[0][0].abs().sum() > 0 # a path that wrote nothing would also "agree" def _verdict(run, operands, c): diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index ae994c2b8..ad81ee0d1 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -97,6 +97,65 @@ def test_epilogue_fusion(): torch.testing.assert_close(y, torch.relu(ref).to(torch.bfloat16), atol=1e-1, rtol=1e-2) +@_GPU +def test_aux_tensor(): + """An aux operand reaches the kernel reshaped to its fake's rank, not permuted.""" + 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=F32) + Y = g.add(a=g.matmul(A=A, B=B, name="mm"), b=bias, name="bias_add") + Y.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + bias_buf = torch.randn(1, 1, N, dtype=torch.float32, device="cuda") + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, bias: bias_buf, Y: y}) + torch.testing.assert_close(y, (ref + bias_buf).to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_two_dense_outputs(): + """Two stored outputs: each contributes its own stride triple, in order.""" + 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) + Y = g.relu(input=C, name="relu") + Y.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") + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, C: c, Y: y}) + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + torch.testing.assert_close(y, torch.relu(ref).to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_multi_gemm(): + """Two matmuls sharing A and an epilogue: the operands the kernel takes are + the DISTINCT ones, which is what the pack binds.""" + 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]) + B0 = g.tensor(name="B0", dim=[1, K, N], stride=[K * N, 1, K]) + B1 = g.tensor(name="B1", dim=[1, K, N], stride=[K * N, 1, K]) + Y = g.add(a=g.matmul(A=A, B=B0, name="mm0"), b=g.matmul(A=A, B=B1, name="mm1"), name="sum") + Y.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a = torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda") + b0 = torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda") + b1 = torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda") + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B0: b0, B1: b1, Y: y}) + ref = torch.einsum("bmk,bnk->bmn", a.float(), b0.float()) + torch.einsum("bmk,bnk->bmn", a.float(), b1.float()) + torch.testing.assert_close(y, ref.to(torch.bfloat16), atol=2e-1, rtol=2e-2) + + @_GPU @pytest.mark.parametrize( "mode,reference", From b7fec81210b2977bfe580309826b2b8c292d7c72 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 13:14:03 -0700 Subject: [PATCH 5/9] Read an aux buffer's rank off its shape, not off an attribute it may not have An aux operand whose buffer rank differs from the graph's declaration was refused where the backend accepts it. Reported for a rank-1 bias, but it is not rank-1 specific: `_reshape_aux_to_fake` asked for the rank with `getattr(t, "ndim", )`, and the pack's `OperandBuffer` carries no `ndim` -- `__len__` is the first EXTENT, not the rank. The default therefore meant "already matches", so EVERY aux arriving through `graph.execute()` skipped the reshape. The existing public aux test passed only because a rank-3 bias needs no reshape. Same graph, same data, only the handed-over buffer's shape differs: bias declared [1,1,N], given (1,1,N) backend OK frost OK bias declared [1,1,N], given (1,N) backend OK frost ValueError bias declared [1,1,N], given (N,) backend OK frost ValueError That is the rule this branch already settled for operand axis order, applied to rank: the descriptor defines the tensor and the variant pack supplies only a pointer. So the test is a differential against the backend across all three ranks, and `OperandBuffer` grows an `ndim` so the next `getattr` for one cannot silently take a default instead. Also adds the block-scale evidence the recipe claim needs: an nvfp4 lowered-vs -interpreted differential for one and two GEMMs (the two-GEMM case additionally asserting `shared_layout`, which only it populates), and an nvfp4 matmul through the public `graph.execute()` that asserts the plan really took the lowered path. Block scale is the flavor with the most per-call table in it -- its scale factors ride in the launch argument list but not in `problem_size`, and their blob size is re-synthesized from M/N/K rather than read off the buffer. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/gemm/frost/compiler.py | 10 +- python/pygraph/variant_pack.cpp | 93 ++++++++++--------- test/python/gemm/frost/test_execute_recipe.py | 82 +++++++++++++++- .../gemm/frost/test_public_execute_flavors.py | 93 +++++++++++++++++++ 4 files changed, 229 insertions(+), 49 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 87eab914d..edd2df6ec 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1764,11 +1764,17 @@ def _reshape_aux_to_fake(t: object, ref: TensorRef) -> object: """View a broadcast-aux runtime tensor to the fake's rank/unit-dim layout so the tvm-ffi front door's ndim/shape check accepts it. Free view: the aux is consumed via raw pointer + injected strides, so only the wrapper rank/shape is - at stake. No-op off the front door, so the plain path stays bit-identical.""" + at stake. No-op off the front door, so the plain path stays bit-identical. + + The rank comes from ``len(t.shape)``, which every buffer answers. Reading it + off an ``ndim`` attribute with a default meant a buffer that does not carry + one -- the pack's, so every operand arriving through ``graph.execute()`` -- + was taken to match already, and a bias declared ``[1, 1, N]`` but handed over + as ``[N]`` was refused where the backend accepts it.""" if not _TVM_FFI_OK: return t real = _aux_fake_real_axes(ref) - if getattr(t, "ndim", len(real)) == len(real): + if len(t.shape) == len(real): return t extents = [int(e) for e in t.shape if int(e) != 1] if sum(real) != len(extents): diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 281ad728b..8b9b0b351 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -722,56 +722,61 @@ class WorkspaceCarve { void init_variant_pack(py::module_ &m) { - auto operand_class = py::class_(m, "OperandBuffer", R"( + auto operand_class = + py::class_(m, "OperandBuffer", R"( One operand of a variant pack, as a DLPack producer. Implements ``__dlpack_c_exchange_api__``, so a consumer reads it through the same C function table it uses for a framework tensor rather than through a capsule built in python. )") - .def("data_ptr", &OperandBuffer::data_ptr) - .def_property_readonly("shape", &OperandBuffer::shape) - .def_property_readonly("dtype", &OperandBuffer::dtype) - .def_property_readonly("nbytes", &OperandBuffer::nbytes) - .def( - "stride", - [](const OperandBuffer &self, py::object dim) -> py::object { - if (dim.is_none()) return py::cast(self.stride()); - return py::cast(self.stride_at(dim.cast())); - }, - py::arg("dim") = py::none()) - .def("element_size", &OperandBuffer::element_size) - .def("numel", &OperandBuffer::numel) - .def("__len__", &OperandBuffer::length) - .def("reshape", - [](const OperandBuffer &self, py::args dims) { - std::vector shape; - if (dims.size() == 1 && py::isinstance(dims[0]) && - !py::isinstance(dims[0])) { - shape = dims[0].cast>(); - } else { - for (auto d : dims) shape.push_back(d.cast()); - } - return self.reshape(std::move(shape)); - }) - .def("permute", - [](const OperandBuffer &self, py::args axes) { - std::vector order; - if (axes.size() == 1 && py::isinstance(axes[0]) && - !py::isinstance(axes[0])) { - order = axes[0].cast>(); - } else { - for (auto a : axes) order.push_back(a.cast()); - } - return self.permute(order); - }) - .def("contiguous", [](py::object self) { return self; }) - .def("__dlpack_device__", &OperandBuffer::dlpack_device) - .def("__dlpack__", - &OperandBuffer::dlpack, - py::kw_only(), - py::arg("stream") = py::none(), - py::arg("max_version") = py::none()); + .def("data_ptr", &OperandBuffer::data_ptr) + .def_property_readonly("shape", &OperandBuffer::shape) + // ndim, because `__len__` is the first EXTENT and a caller + // reaching for a rank through getattr(.., "ndim", default) + // silently gets the default instead. + .def_property_readonly("ndim", [](const OperandBuffer &self) { return self.shape().size(); }) + .def_property_readonly("dtype", &OperandBuffer::dtype) + .def_property_readonly("nbytes", &OperandBuffer::nbytes) + .def( + "stride", + [](const OperandBuffer &self, py::object dim) -> py::object { + if (dim.is_none()) return py::cast(self.stride()); + return py::cast(self.stride_at(dim.cast())); + }, + py::arg("dim") = py::none()) + .def("element_size", &OperandBuffer::element_size) + .def("numel", &OperandBuffer::numel) + .def("__len__", &OperandBuffer::length) + .def("reshape", + [](const OperandBuffer &self, py::args dims) { + std::vector shape; + if (dims.size() == 1 && py::isinstance(dims[0]) && + !py::isinstance(dims[0])) { + shape = dims[0].cast>(); + } else { + for (auto d : dims) shape.push_back(d.cast()); + } + return self.reshape(std::move(shape)); + }) + .def("permute", + [](const OperandBuffer &self, py::args axes) { + std::vector order; + if (axes.size() == 1 && py::isinstance(axes[0]) && + !py::isinstance(axes[0])) { + order = axes[0].cast>(); + } else { + for (auto a : axes) order.push_back(a.cast()); + } + return self.permute(order); + }) + .def("contiguous", [](py::object self) { return self; }) + .def("__dlpack_device__", &OperandBuffer::dlpack_device) + .def("__dlpack__", + &OperandBuffer::dlpack, + py::kw_only(), + py::arg("stream") = py::none(), + py::arg("max_version") = py::none()); // The protocol looks the attribute up on the TYPE, and a pybind11 class is // a heap type, so it takes a plain setattr. diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index 46f5d69b8..f1847f7bf 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -23,7 +23,7 @@ import pytest import torch -from gemm_test_utils import requires_sm100 +from gemm_test_utils import kw, requires_sm100, to_blocked import cudnn import cudnn.gemm.frost # noqa: F401 — installs the cudnn.pygraph recorder hook @@ -198,11 +198,18 @@ def test_public_execute_takes_the_lowered_path(): # --- the differential ------------------------------------------------------- -def _bind(compiled, a_bufs, b_bufs, out_bufs, aux_bufs=()): +def _bind(compiled, a_bufs, b_bufs, out_bufs, aux_bufs=(), sfa_bufs=(), sfb_bufs=()): """The operand list in bound-tensor order, which is what both paths take.""" bd = compiled.binding pack = {} - for role, bufs in ((bd.a_operands, a_bufs), (bd.b_operands, b_bufs), (bd.outputs, out_bufs), (bd.aux, aux_bufs)): + for role, bufs in ( + (bd.a_operands, a_bufs), + (bd.b_operands, b_bufs), + (bd.sfa_operands, sfa_bufs), + (bd.sfb_operands, sfb_bufs), + (bd.outputs, out_bufs), + (bd.aux, aux_bufs), + ): pack.update(zip(role, bufs)) resolved = resolve_variant_pack(pack, bd) return [resolved[id(t)] for t in bd.bound_tensors()] @@ -372,6 +379,75 @@ def _matmul_on(batch, m, n, k, want_frost, a, b, c): return c +BS_M = BS_N = 128 +BS_K = 256 +BS_BLOCK = 16 + + +def _nvfp4_graph(gemms=1): + """One or two nvfp4 block-scaled matmuls, sharing A when there are two.""" + sf_k = BS_K // BS_BLOCK + fp4, fp8 = cudnn.data_type.FP4_E2M1, cudnn.data_type.FP8_E4M3 + reorder = dict(reordering_type=cudnn.tensor_reordering.F8_128x4) + g = cudnn.pygraph(io_data_type=cudnn.data_type.HALF, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, BS_M, BS_K], stride=[BS_M * BS_K, BS_K, 1], data_type=fp4) + SFA = g.tensor(name="SFA", dim=[1, BS_M, sf_k], stride=[BS_M * sf_k, sf_k, 1], data_type=fp8, **reorder) + Ad = g.block_scale_dequantize(input=A, descale=SFA, block_size=[1, BS_BLOCK]) + products = [] + for i in range(gemms): + B = g.tensor(name=f"B{i}", dim=[1, BS_K, BS_N], stride=[BS_K * BS_N, 1, BS_K], data_type=fp4) + SFB = g.tensor(name=f"SFB{i}", dim=[1, sf_k, BS_N], stride=[sf_k * BS_N, 1, sf_k], data_type=fp8, **reorder) + Bd = g.block_scale_dequantize(input=B, descale=SFB, block_size=[BS_BLOCK, 1]) + products.append(g.matmul(A=Ad, B=Bd, name=f"mm{i}")) + out = products[0] if gemms == 1 else g.add(a=products[0], b=products[1], name="sum") + out.set_output(True).set_data_type(cudnn.data_type.HALF) + return g + + +def _nvfp4_buffers(compiled): + """Packed fp4 operands and their F8_128x4 scale blobs, in bound order.""" + sf_k = BS_K // BS_BLOCK + n_b = len(compiled.binding.b_operands) + a = torch.randint(0, 256, (1, BS_M, BS_K // 2), dtype=torch.uint8, device="cuda").view(torch.float4_e2m1fn_x2) + sfa = to_blocked(torch.randint(1, 4, (BS_M, sf_k), device="cuda").to(torch.float8_e4m3fn)).view(1, BS_M, sf_k) + bs = [torch.randint(0, 256, (1, BS_N, BS_K // 2), dtype=torch.uint8, device="cuda").view(torch.float4_e2m1fn_x2) for _ in range(n_b)] + sfbs = [to_blocked(torch.randint(1, 4, (BS_N, sf_k), device="cuda").to(torch.float8_e4m3fn)).view(1, BS_N, sf_k) for _ in range(n_b)] + out = torch.zeros(1, BS_M, BS_N, dtype=torch.float16, device="cuda") + return _bind(compiled, [a], bs, [out], sfa_bufs=[sfa], sfb_bufs=sfbs), out + + +@requires_sm100 +@pytest.mark.parametrize("gemms", (1, 2), ids=("single", "multi")) +def test_block_scale_lowers_and_agrees_with_the_interpreter(gemms): + """Block scale is the flavor with the most per-call table in it. + + Its scale factors ride in the launch argument list but NOT in + ``problem_size``, its blob size is re-synthesized from M/N/K rather than read + off the buffer, and the multi-GEMM form sends one A stride triple for every + operand instead of one each. Four recipe fields, so it is the one most worth + running against the interpreter. + """ + # Two GEMMs do not fit the auto-selected cta_n=256 in TMEM, so pin a + # geometry that does; which config the engine picks for one GEMM is the + # public execute test's job, not this one's. + cfg = kw("CONFIG_sm100_128x128x128_128x128x32_cluster1x1_1ctamma") if gemms > 1 else {} + compiled = jit_from_cudnn_graph(_nvfp4_graph(gemms), **cfg) + assert compiled.block_scale + if compiled.lowered is None: + pytest.skip("this build does not lower (no tvm-ffi front door)") + assert bool(compiled.recipe.shared_layout) == (gemms > 1) + + operands, out = _nvfp4_buffers(compiled) + runs = [] + for run in (compiled.lowered, compiled.launch): + out.zero_() + run(operands, stream=None) + torch.cuda.synchronize() + runs.append(out.clone()) + torch.testing.assert_close(runs[0], runs[1], atol=0, rtol=0) + assert runs[0].abs().sum() > 0 + + @requires_sm100 @pytest.mark.parametrize("batch", (1, 2), ids=("b1", "b2")) @pytest.mark.parametrize("m,n,k", [(m, n, k) for m in (1, 128) for n in (1, 128) for k in (1, 128)], ids=str) diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index ad81ee0d1..da3e2085d 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -115,6 +115,46 @@ def test_aux_tensor(): torch.testing.assert_close(y, (ref + bias_buf).to(torch.bfloat16), atol=1e-1, rtol=1e-2) +@_GPU +@pytest.mark.parametrize("buffer_shape", ((1, 1, N), (1, N), (N,)), ids=("rank3", "rank2", "rank1")) +def test_aux_buffer_rank_need_not_match_the_declaration(buffer_shape): + """A bias declared ``[1, 1, N]`` may be handed over as ``[N]``. + + Same rule as an operand's axis order: the descriptor defines the tensor and + the pack supplies only a pointer, so the backend takes any rank here and a + python plan must reach the same answer. The reshape that makes the kernel's + wrapper agree read the rank off an ``ndim`` attribute with a default, and the + pack's buffers carry no ``ndim`` -- so every operand arriving through + ``execute()`` looked like it already matched, and these were refused. + """ + out = {} + for want_frost in (False, True): + 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=F32) + Y = g.add(a=g.matmul(A=A, B=B, name="mm"), b=bias, name="bias_add") + Y.set_output(True).set_data_type(BF16) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + hits = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id) == want_frost] + if not hits: + pytest.skip("no plan of the requested kind for this graph") + g.select_plan(hits[0]) + g.check_support() + g.build_plans() + + torch.manual_seed(0) + a = torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda") + bias_buf = torch.randn(buffer_shape, dtype=torch.float32, device="cuda") + y = torch.zeros(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, bias: bias_buf, Y: y}) + out[want_frost] = y + torch.testing.assert_close(out[True], out[False], atol=1e-1, rtol=1e-2) + + @_GPU def test_two_dense_outputs(): """Two stored outputs: each contributes its own stride triple, in order.""" @@ -225,6 +265,59 @@ def test_int32_reduction_seed_is_packed_as_int32(): assert int(r.item()) == floor +@_GPU +def test_block_scale_matmul(): + """nvfp4 + F8_128x4 scale blobs, through the public entry point. + + The scale factors are the operands the recipe treats least like the others: + they ride in the launch argument list but not in ``problem_size``, and their + size is re-synthesized from M/N/K rather than read off the buffer, so a blob + that is not one dense byte run is read out of bounds with no fault. + """ + from gemm_test_utils import E2M1, to_blocked, unpack_fp4 + + bs_m = bs_n = 128 + bs_k, block = 256, 16 + sf_k = bs_k // block + fp4, fp8 = cudnn.data_type.FP4_E2M1, cudnn.data_type.FP8_E4M3 + reorder = dict(reordering_type=cudnn.tensor_reordering.F8_128x4) + + g = cudnn.pygraph(io_data_type=cudnn.data_type.HALF, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, bs_m, bs_k], stride=[bs_m * bs_k, bs_k, 1], data_type=fp4) + B = g.tensor(name="B", dim=[1, bs_k, bs_n], stride=[bs_k * bs_n, 1, bs_k], data_type=fp4) + SFA = g.tensor(name="SFA", dim=[1, bs_m, sf_k], stride=[bs_m * sf_k, sf_k, 1], data_type=fp8, **reorder) + SFB = g.tensor(name="SFB", dim=[1, sf_k, bs_n], stride=[sf_k * bs_n, 1, sf_k], data_type=fp8, **reorder) + C = g.matmul( + A=g.block_scale_dequantize(input=A, descale=SFA, block_size=[1, block]), + B=g.block_scale_dequantize(input=B, descale=SFB, block_size=[block, 1]), + name="mm", + ) + C.set_output(True).set_data_type(cudnn.data_type.HALF) + _pin_frost(g) + assert g._compiled_plans[g._plan_index]._lowered is not None + + torch.manual_seed(0) + lut = torch.tensor(E2M1, dtype=torch.float32, device="cuda") + a_u8 = torch.randint(0, 256, (1, bs_m, bs_k // 2), dtype=torch.uint8, device="cuda") + b_u8 = torch.randint(0, 256, (1, bs_n, bs_k // 2), dtype=torch.uint8, device="cuda") + sfa = torch.randint(1, 4, (bs_m, sf_k), device="cuda").to(torch.float8_e4m3fn) + sfb = torch.randint(1, 4, (bs_n, sf_k), device="cuda").to(torch.float8_e4m3fn) + c = torch.zeros(1, bs_m, bs_n, dtype=torch.float16, device="cuda") + _run( + g, + { + A: a_u8.view(torch.float4_e2m1fn_x2), + B: b_u8.view(torch.float4_e2m1fn_x2), + SFA: to_blocked(sfa).view(1, bs_m, sf_k), + SFB: to_blocked(sfb).view(1, bs_n, sf_k), + C: c, + }, + ) + a_s = unpack_fp4(a_u8, lut).view(bs_m, bs_k) * sfa.float().repeat_interleave(block, 1) + b_s = unpack_fp4(b_u8, lut).view(bs_n, bs_k) * sfb.float().repeat_interleave(block, 1) + torch.testing.assert_close(c.float(), (a_s @ b_s.t()).reshape(1, bs_m, bs_n), atol=2e-1, rtol=2e-2) + + @_GPU def test_norm2_reduction_is_refused_at_build(): """``norm2`` is the one reduction mode with a post-kernel finalize. From 815ad8bb3bd88cd51c6b1619cb5165a23352124b Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 14:28:45 -0700 Subject: [PATCH 6/9] Count every reason a call leaves the fast path The interpreter is migration scaffolding: the fast path hands it anything it is not certain of. That makes "no legal call needs the interpreter" the claim the whole two-track structure rests on, and until now it was only an assertion. The day the fallback is deleted is the day an unnoticed dependency on it becomes a regression. So both halves of "why not the fast path" are named data now: - `compiled.declined` says which rule denied this graph a fast path at build -- "needs workspace", "post-kernel sqrt", "multi-GEMM without a dense output", "no tvm-ffi front door", "scale factors without a block size", "declared layout is indistinguishable from the kernel's" -- instead of a bare None. - `compiled.deferrals` counts the eight per-call sites, by reason. Incremented only on the path that is already giving up, so the fast path pays nothing. The counts separate two populations that reading the code does not. Six of the eight sites mean the CALL IS ILLEGAL -- wrong major, a TMA-misaligned extent, a misaligned output base, an SF blob that is not a dense run -- and exist only so the interpreter can produce the message. The other two are calls that are perfectly LEGAL and that the fast path cannot serve: an operand described from the graph (a bare device address, or a buffer that reports the declaration), and a strided reduction output, which needs a fill the engine does not own. Only the second population is a reason to widen the plan; the first never will be. Three tests: every flavor's legal call defers for nothing, a decline names its rule, and a bare address is counted under exactly that reason and no other. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/gemm/frost/compiler.py | 44 ++++++++++++++++--- test/python/gemm/frost/test_execute_recipe.py | 42 ++++++++++++++++++ 2 files changed, 80 insertions(+), 6 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index edd2df6ec..34e72a7df 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -11,6 +11,7 @@ import functools import hashlib +from collections import Counter import importlib.util import logging import os @@ -1865,6 +1866,13 @@ class CompiledFusedGemm: # The recipe as one straight line, when this graph is a shape it emits. lowered: Any = field(default=None, init=False, repr=False, compare=False) bound: Any = field(default=(), init=False, repr=False, compare=False) + # Why this graph has no fast path, or None when it has one. Set at build. + declined: "str | None" = field(default=None, init=False, repr=False, compare=False) + # Reason -> count for calls the fast path handed to the interpreter. The + # interpreter is migration scaffolding, so "no production call needs it" has + # to be measurable rather than asserted; a test reads this per flavor. + # Counted only on the path already being given up, so the fast path is free. + deferrals: Any = field(default_factory=Counter, init=False, repr=False, compare=False) def __post_init__(self) -> None: if self.binding is not None: @@ -1915,18 +1923,31 @@ def _lower(self): message -- so this can only ever accept a subset of what the general path accepts, and there is no second set of error strings to drift. """ - r = self.recipe - if r is None or not _TVM_FFI_OK or r.workspace_bytes: + + def decline(reason: str) -> None: + self.declined = reason return None + + r = self.recipe + if r is None: + return decline("no binding") + if not _TVM_FFI_OK: + # An executor capability, not a property of this graph: without the + # front door the arguments need the legacy DLPack wrappers. + return decline("no tvm-ffi front door") + if r.workspace_bytes: + return decline("needs workspace") # norm2 takes a square root through the caller's buffer after the kernel; # that is a device operation the engine does not own yet, and neither is # seeding a strided reduction output (checked per call, below). if any(o.sqrt for o in r.outputs): - return None + return decline("post-kernel sqrt") # Scale factors come with a block size to size their blob against, and a # multi-GEMM's batch with a dense output to read it off. - if bool(r.sf) != bool(r.block_size) or (r.multi_gemm and not r.has_output_specs): - return None + if bool(r.sf) != bool(r.block_size): + return decline("scale factors without a block size") + if r.multi_gemm and not r.has_output_specs: + return decline("multi-GEMM without a dense output") # A buffer reporting the declaration is read in the graph's axis order, # which this body does not serve. The stride guard below tells that # apart -- unless the declaration itself would satisfy the guard, which @@ -1934,7 +1955,7 @@ def _lower(self): for op in r.inputs: declared_stride = op.declared_layout[1] if op.dc != op.kc and declared_stride and declared_stride[op.kc] == 1: - return None + return decline("declared layout is indistinguishable from the kernel's") # The recipe flattened into the loop headers: everything the body reads # is unpacked by the `for`, so it costs no attribute lookup. `kc` is both @@ -1951,15 +1972,20 @@ def _lower(self): block_size, batch = r.block_size, r.batch device, launchable, general = r.device, self._launchable, self.launch fill_word, is_contiguous = buffers.fill_word_async, buffers.is_contiguous + # Named so a test can assert which reason a call deferred for, and that a + # legal call defers for none. Incremented only on the giving-up path. + gave_up = self.deferrals def lowered(operands, graph_order=None, stream=None): _check_plan_device(device) if graph_order: # An operand described from the graph is in the graph's axis # order; this body reads the caller's. + gave_up["graph-described operand"] += 1 return general(operands, graph_order, stream=stream) a_sh, b_sh = operands[ai].shape, operands[bi].shape if len(a_sh) != 3 or len(b_sh) != 3: + gave_up["A or B is not rank 3"] += 1 return general(operands, None, stream=stream) m, n, k = a_sh[1], b_sh[1], a_sh[2] * a_kpack # permute(1, 2, 0) relabels axes, so the strides the kernel wants are @@ -1977,6 +2003,7 @@ def lowered(operands, graph_order=None, stream=None): or sh[2] * kpack != k or _pow2_floor(v.data_ptr()) < 16 ): + gave_up["input layout"] += 1 return general(operands, None, stream=stream) if takes_stride: problem += (st[1], st[2], st[0]) @@ -1992,11 +2019,13 @@ def lowered(operands, graph_order=None, stream=None): or sh[2] != (v2 if k2 == CONST else m if k2 == FROM_M else n // v2) or tensor_alignment(tuple(sh), tuple(st), v.element_size(), ptr=v.data_ptr()) < align ): + gave_up["output layout"] += 1 return general(operands, None, stream=stream) problem += (st[1], st[2], st[0]) for idx, align in auxs: v = operands[idx] if tensor_alignment(tuple(v.shape), tuple(v.stride()), v.element_size(), ptr=v.data_ptr()) < align: + gave_up["aux alignment"] += 1 return general(operands, None, stream=stream) if sfs: k4 = ((k // block_size) + 3) // 4 @@ -2008,11 +2037,13 @@ def lowered(operands, graph_order=None, stream=None): or count != 1 + sum((int(s) - 1) * int(st) for s, st in zip(v.shape, v.stride())) or count * v.element_size() < 512 * k4 * (((m if is_a else n) + 127) // 128) * sf_batch ): + gave_up["scale-factor blob"] += 1 return general(operands, None, stream=stream) for lead, followers in shared: st = tuple(operands[lead].stride()) for j in followers: if tuple(operands[j].stride()) != st: + gave_up["operands do not share a layout"] += 1 return general(operands, None, stream=stream) if seeds: # Seeding is a write, so nothing above may still reject: check @@ -2020,6 +2051,7 @@ def lowered(operands, graph_order=None, stream=None): for idx, _word in seeds: v = operands[idx] if not is_contiguous(tuple(v.shape), tuple(v.stride())): + gave_up["strided reduction seed"] += 1 return general(operands, None, stream=stream) for idx, word in seeds: v = operands[idx] diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index f1847f7bf..bd4aa179c 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -164,6 +164,48 @@ def test_which_flavors_lower(build, n_b, n_out, n_aux): assert len(r.arg_plan) == len(r.inputs) + len(r.outputs) + len(r.aux) + len(r.sf) +@requires_sm100 +@pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) +def test_a_legal_call_defers_to_the_interpreter_for_nothing(build, n_b, n_out, n_aux): + """The interpreter is migration scaffolding, so that has to be measurable. + + Every reason the fast path can hand a call over is counted, and a legal call + of every flavor must trigger none of them -- otherwise "the fallback only + catches what the fast path declines" is a claim with nothing behind it, and + the day it is deleted is the day the regression appears. + """ + compiled = jit_from_cudnn_graph(build()) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") + operands = _buffers_for(compiled) + compiled.lowered(operands, stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {} + + +@requires_sm100 +def test_every_decline_names_its_reason(): + """A plan without a fast path says which rule denied it, not just ``None``.""" + compiled = jit_from_cudnn_graph(_plain_graph()) + assert compiled.lowered is not None and compiled.declined is None + compiled.recipe = replace(compiled.recipe, workspace_bytes=4096) + assert compiled._lower() is None + assert compiled.declined == "needs workspace" + + +@requires_sm100 +def test_a_deferral_is_counted_under_the_rule_that_caused_it(): + """A bare address is the one deferral a legal call still takes.""" + compiled = jit_from_cudnn_graph(_plain_graph()) + if compiled.lowered is None: + pytest.skip("this build does not lower") + a, b, c = _operands() + operands = _bound_buffers(compiled, a, b, c) + compiled.lowered(operands, (True, False, False), stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {"graph-described operand": 1} + + @requires_sm100 def test_a_post_kernel_finalize_is_declined(): """``norm2`` takes a square root through the caller's buffer after the kernel. From 8813dd85c19c4d84d290247fad2301a491788cbb Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 15:33:27 -0700 Subject: [PATCH 7/9] Make the lowered path the only executor, and the slow path a checker A call the fast path was unsure of went to an interpreter that ran it anyway, so "the fast path refused a legal call" was a performance bug at worst and nothing asserted otherwise. The deferral counters from 98578b5b6 split the eight hand-over sites into two populations that reading the code does not: six mean the call is illegal and only need a message, two meant it was legal and the loop could not serve it. Both of the second are gone. An operand in the GRAPH's axis order. cuDNN declares a matmul's B [b, K, N] where this kernel reads (b, N, K), so a bare device address -- and any buffer reporting the declared (dim, stride) -- arrived the other way round and paid a whole interpreted pass for being legal. `lowered` re-labels it with one permute whose shape the recipe already carries, before any check runs, so everything below still reads one order. Measured on the plain flavor: 19.95 -> 19.76 us, i.e. the pre-pass is free on a call that does not use it. That also retires the "declared layout is indistinguishable from the kernel's" build-time decline, which existed because the stride guard alone cannot tell the two orders apart at a unit extent. Comparing the full (dim, stride) can. A strided reduction seed. fill_word_strided_async collapses the layout to the runs a memset can cover and issues cuMemsetD2D32Async, which takes a pitch: one call per remaining outer point (the batch, usually 1) rather than one per row, which is the reading that made this look expensive. It was the last place the engine wrote through a buffer it does not own -- tensor.fill_() works only while the caller passes a torch tensor, and queues on torch's stream rather than the one the kernel will run on. Three build-time declines also kept launch alive as an executor: no tvm-ffi front door, a norm2 output, a multi-GEMM with no dense output. They are _check_executable now, called from probe_supported, so the ENGINE declines the graph and it goes to the backend -- where it would have gone had FROST never been asked. _lower reads that same function instead of restating it. Then launch loses its second half. explain() runs the same recipe.problem / check_shapes / check_alignment and raises; _call_positional, _call_multi_gemm, _call_block_scale_multi_gemm and launch's assembly are deleted, 267 lines. MoE is untouched -- it is >= 2 launches with its own workspace and no recipe, which is why _maybe_wrap_layout, _wrap_raw_tensor, _initialize_reduction_outputs and _finalize_reductions all survive. The invariant tightens from "the fast path may accept a subset" to: the set of calls the launch path refuses must EQUAL the set of illegal calls. `deferrals` enforces it -- a legal call of every flavor must leave it empty, asserted per flavor and for the bare-address form. Refusing a legal call is now a bug. Execution stops being duplicated but the RULES are still written twice, fused in lowered's guards and readable in explain. If they drift, a call is refused and the checker finds nothing wrong, so explain falling off the end raises a distinct RuntimeError naming which guard fired. The lowered-vs-interpreted differential goes with the interpreter. Two readings of the same wrong plan agreeing proves nothing; that is exactly how the axis-order bug survived one. What replaces it is fast-vs-BACKEND, already here, plus the per-flavor torch references in test_public_execute_flavors.py, which run through the public graph.execute() rather than beside it. All six flavors still lower and their host cost is unchanged: plain 19.76, epilogue 19.69, aux 22.26, 2 outputs 22.65, reduction 26.12, multi-gemm 22.54 (min over 25 reps of a 64-call burst, 256x256x128 bf16, SM100). Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/gemm/frost/compiler.py | 469 +++++------------- python/cudnn/gemm/frost/engine.py | 13 +- python/cudnn/gemm/frost/recipe.py | 30 +- test/python/gemm/frost/test_execute_recipe.py | 322 ++++++++---- .../gemm/frost/test_matmul_epilogue_fusion.py | 30 +- .../gemm/frost/test_public_execute_flavors.py | 5 +- 6 files changed, 411 insertions(+), 458 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 34e72a7df..86a921b4a 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -56,9 +56,6 @@ def _as_custream(stream): from .epilogue_codegen import EpilogueSnippets, generate from .fusion_ir import ZERO_PRESERVING_OPS, FusionChain, TensorRef from .recipe import ( - AX_BATCH, - AX_K, - AX_MN, CONST, FROM_M, KERNEL_AXES, @@ -1832,8 +1829,9 @@ class CompiledFusedGemm: `col + vsize <= N`, the TMA-store path relies on HW dropping OOB coords. Construct via :func:`jit_from_cudnn_graph`, then call with a variant-pack - dict. See ``_maybe_wrap_layout`` for the leading-dim policy (A=-1, B=-2, - C=-1, aux=-1). + dict. ``lowered`` is the launch, and the only one: a graph it cannot serve + was declined by :func:`_check_executable` before a plan existed, and a call + it refuses goes to ``explain``, which raises rather than running anything. Do NOT set `oob_fill=nan_request_zero_fma` on sm100: the "NONE" enum name is misleading (bit 0 = zero-fill), and NaN-request is harmful because @@ -1863,15 +1861,15 @@ class CompiledFusedGemm: # once here rather than rebuilt per execute. None when this object was # constructed without a binding and so has no call path. recipe: Any = field(default=None, init=False, repr=False, compare=False) - # The recipe as one straight line, when this graph is a shape it emits. + # The recipe as one loop: this kernel's launch path, and the only one. lowered: Any = field(default=None, init=False, repr=False, compare=False) bound: Any = field(default=(), init=False, repr=False, compare=False) - # Why this graph has no fast path, or None when it has one. Set at build. + # Why this object has no launch path, or None when it has one. Set at build. declined: "str | None" = field(default=None, init=False, repr=False, compare=False) - # Reason -> count for calls the fast path handed to the interpreter. The - # interpreter is migration scaffolding, so "no production call needs it" has - # to be measurable rather than asserted; a test reads this per flavor. - # Counted only on the path already being given up, so the fast path is free. + # Reason -> count for calls the launch path refused. Refusing is the answer + # now, not a slower route, so this is the invariant rather than observability: + # a legal call of any flavor must leave it empty. Counted only on the path + # that is already raising, so a served call pays nothing for it. deferrals: Any = field(default_factory=Counter, init=False, repr=False, compare=False) def __post_init__(self) -> None: @@ -1899,17 +1897,19 @@ def run_resolved(self, resolved, stream=None): if buf is None: raise KeyError(f"variant pack is missing a buffer for {self.recipe.roles[i]}") operands.append(buf) - return (self.lowered or self.launch)(operands, stream=stream) + if self.lowered is None: + raise NotImplementedError(f"cudnn.frost gemm: this kernel has no launch path -- {self.declined}") + return self.lowered(operands, stream=stream) def _lower(self): - """The recipe as one loop over flat tuples, or None for a graph it does not serve. + """The recipe as one loop over flat tuples: this kernel's only call path. - Every branch the general path takes -- multi-GEMM or not, block scale or - not, how many outputs, which are reductions, which axis of each operand - carries M/N/K, what order the kernel's parameters come in -- was settled - when the kernel compiled. Taking them again per call is most of what a - python execute path costs: 44 us for the walk that rebuilt them, 35 - reading a recipe through its objects, 20 here. + Every branch a per-call walk would take -- multi-GEMM or not, block + scale or not, how many outputs, which are reductions, which axis of each + operand carries M/N/K, what order the kernel's parameters come in -- was + settled when the kernel compiled. Taking them again per call is most of + what a python execute path costs: 44 us for the walk that rebuilt them, + 35 reading a recipe through its objects, 20 here. What is left per call is the same loop for every flavor, over tuples flat enough that the body does no attribute lookup and calls nothing it @@ -1918,10 +1918,14 @@ def _lower(self): one body instead of one per shape; source codegen off this same table is how to get that 12% back for every flavor at once. - The emitted body never raises. Anything it is not sure of it hands to - ``launch``, which serves every flavor and owns every rejection - message -- so this can only ever accept a subset of what the general - path accepts, and there is no second set of error strings to drift. + The body never raises. What it refuses it hands to ``explain``, which + owns every rejection message and does not run anything -- so the set of + calls this refuses must EQUAL the set of illegal calls, and a legal one + it will not serve is a bug rather than a slower route. ``deferrals`` + counts each refusal under its rule, which is what makes that testable. + + None means this object has no call path at all, which the engine's + support gate already refused; ``declined`` says why. """ def decline(reason: str) -> None: @@ -1931,31 +1935,28 @@ def decline(reason: str) -> None: r = self.recipe if r is None: return decline("no binding") - if not _TVM_FFI_OK: - # An executor capability, not a property of this graph: without the - # front door the arguments need the legacy DLPack wrappers. - return decline("no tvm-ffi front door") + # The engine's support gate, read again rather than restated: a graph + # this cannot run was declined before a plan existed, so reaching here + # means a caller built the kernel directly. + try: + _check_executable(self.chain) + except NotImplementedError as exc: + return decline(str(exc)) if r.workspace_bytes: return decline("needs workspace") - # norm2 takes a square root through the caller's buffer after the kernel; - # that is a device operation the engine does not own yet, and neither is - # seeding a strided reduction output (checked per call, below). - if any(o.sqrt for o in r.outputs): - return decline("post-kernel sqrt") - # Scale factors come with a block size to size their blob against, and a - # multi-GEMM's batch with a dense output to read it off. + # Scale factors come with a block size to size their blob against. if bool(r.sf) != bool(r.block_size): return decline("scale factors without a block size") - if r.multi_gemm and not r.has_output_specs: - return decline("multi-GEMM without a dense output") - # A buffer reporting the declaration is read in the graph's axis order, - # which this body does not serve. The stride guard below tells that - # apart -- unless the declaration itself would satisfy the guard, which - # only a unit extent can arrange, and which is settled here not per call. - for op in r.inputs: - declared_stride = op.declared_layout[1] - if op.dc != op.kc and declared_stride and declared_stride[op.kc] == 1: - return decline("declared layout is indistinguishable from the kernel's") + + # Operands whose graph axis order is not the kernel's -- B, which cuDNN + # declares [b, K, N] where the kernel reads (b, N, K). Re-labelling one + # into kernel order is a permute, so the checks below read one order and + # the second axis map costs nothing on a call that does not use it. + renorm = tuple( + (op.index, tuple(op.declared), op.declared_layout[0] if len(op.declared_layout[0]) == 3 else None, op.declared_layout[1]) + for op in r.inputs + if op.declared != KERNEL_AXES + ) # The recipe flattened into the loop headers: everything the body reads # is unpacked by the `for`, so it costs no attribute lookup. `kc` is both @@ -1970,29 +1971,46 @@ def decline(reason: str) -> None: shared, seeds = r.shared_layout, r.seeds ai, bi, a_kpack = r.a.index, r.b.index, r.a.kpack block_size, batch = r.block_size, r.batch - device, launchable, general = r.device, self._launchable, self.launch - fill_word, is_contiguous = buffers.fill_word_async, buffers.is_contiguous - # Named so a test can assert which reason a call deferred for, and that a - # legal call defers for none. Incremented only on the giving-up path. + device, launchable, refuse = r.device, self._launchable, self.explain + fill_word, fill_strided, is_contiguous = buffers.fill_word_async, buffers.fill_word_strided_async, buffers.is_contiguous + # Named so a test can assert which rule refused a call, and that a legal + # call trips none. Incremented only on the path that is already raising. gave_up = self.deferrals def lowered(operands, graph_order=None, stream=None): _check_plan_device(device) - if graph_order: - # An operand described from the graph is in the graph's axis - # order; this body reads the caller's. - gave_up["graph-described operand"] += 1 - return general(operands, graph_order, stream=stream) - a_sh, b_sh = operands[ai].shape, operands[bi].shape + # Which axis order an operand arrived in, by the backend's own rule: + # the descriptor defines the tensor and the pack supplies a pointer, + # so a slot the pack described FROM the graph, or a buffer reporting + # exactly the declared (dim, stride), is the declaration. Everything + # else is the caller's own labelling. Settled here, once, so the + # checks below all read the kernel's order. + ops = operands + for idx, perm, dsh, dst in renorm: + v = operands[idx] + sh = v.shape + if len(sh) != 3: + continue # no axis map fits it; the loop below is what refuses it + declared = bool(graph_order) and graph_order[idx] + if not declared and dsh is not None: + declared = sh[0] == dsh[0] and sh[1] == dsh[1] and sh[2] == dsh[2] + if declared: + st = v.stride() + declared = st[0] == dst[0] and st[1] == dst[1] and st[2] == dst[2] + if declared: + if ops is operands: + ops = list(operands) + ops[idx] = v.permute(*perm) + a_sh, b_sh = ops[ai].shape, ops[bi].shape if len(a_sh) != 3 or len(b_sh) != 3: gave_up["A or B is not rank 3"] += 1 - return general(operands, None, stream=stream) + return refuse(operands, graph_order) m, n, k = a_sh[1], b_sh[1], a_sh[2] * a_kpack # permute(1, 2, 0) relabels axes, so the strides the kernel wants are # that rotation of the ones each check already read -- no second read. problem = [m, n, k, batch] for idx, kc, mod, ebatch, kpack, is_b, takes_stride in ins: - v = operands[idx] + v = ops[idx] sh, st = v.shape, v.stride() if ( len(sh) != 3 @@ -2004,13 +2022,13 @@ def lowered(operands, graph_order=None, stream=None): or _pow2_floor(v.data_ptr()) < 16 ): gave_up["input layout"] += 1 - return general(operands, None, stream=stream) + return refuse(operands, graph_order) if takes_stride: problem += (st[1], st[2], st[0]) # Each output axis is a constant, M, or N over a divisor (fp4 packs # two along N) -- the rule the build recorded, read back per axis. for idx, k0, v0, k1, v1, k2, v2, align in outs: - v = operands[idx] + v = ops[idx] sh, st = v.shape, v.stride() if ( len(sh) != 3 @@ -2020,17 +2038,17 @@ def lowered(operands, graph_order=None, stream=None): or tensor_alignment(tuple(sh), tuple(st), v.element_size(), ptr=v.data_ptr()) < align ): gave_up["output layout"] += 1 - return general(operands, None, stream=stream) + return refuse(operands, graph_order) problem += (st[1], st[2], st[0]) for idx, align in auxs: - v = operands[idx] + v = ops[idx] if tensor_alignment(tuple(v.shape), tuple(v.stride()), v.element_size(), ptr=v.data_ptr()) < align: gave_up["aux alignment"] += 1 - return general(operands, None, stream=stream) + return refuse(operands, graph_order) if sfs: k4 = ((k // block_size) + 3) // 4 for idx, is_a, sf_batch in sfs: - v = operands[idx] + v = ops[idx] count = int(v.numel()) if ( _pow2_floor(v.data_ptr()) < 16 @@ -2038,299 +2056,57 @@ def lowered(operands, graph_order=None, stream=None): or count * v.element_size() < 512 * k4 * (((m if is_a else n) + 127) // 128) * sf_batch ): gave_up["scale-factor blob"] += 1 - return general(operands, None, stream=stream) + return refuse(operands, graph_order) for lead, followers in shared: - st = tuple(operands[lead].stride()) + st = tuple(ops[lead].stride()) for j in followers: - if tuple(operands[j].stride()) != st: + if tuple(ops[j].stride()) != st: gave_up["operands do not share a layout"] += 1 - return general(operands, None, stream=stream) - if seeds: - # Seeding is a write, so nothing above may still reject: check - # every reduction output first, then fill. - for idx, _word in seeds: - v = operands[idx] - if not is_contiguous(tuple(v.shape), tuple(v.stride())): - gave_up["strided reduction seed"] += 1 - return general(operands, None, stream=stream) - for idx, word in seeds: - v = operands[idx] + return refuse(operands, graph_order) + # Seeding is a write, so it goes last: every rule above has had its + # say. A padded tap costs one 2D memset per batch instead of one + # dense one, which is a price and not a reason to leave. + for idx, word in seeds: + v = ops[idx] + sh, st = v.shape, v.stride() + if is_contiguous(sh, st): fill_word(v.data_ptr(), int(v.numel()), word, stream) + else: + fill_strided(v.data_ptr(), sh, st, v.element_size(), word, stream) return launchable( tuple(problem), - *(operands[i].permute(1, 2, 0) if ref is None else _reshape_aux_to_fake(operands[i], ref) for i, ref in args), + *(ops[i].permute(1, 2, 0) if ref is None else _reshape_aux_to_fake(ops[i], ref) for i, ref in args), stream=_as_custream(stream), ) return lowered - def launch(self, operands, graph_order=None, stream=None): - """Launch over the operand buffers, in bound-tensor order. + def explain(self, operands, graph_order=None): + """Say what is wrong with a call the launch path refused, and raise. - Which bound tensor holds which operand is fixed when the plan compiles, - so a caller that knows it -- the engine, which binds the graph's slots - once -- indexes rather than resolves, and the gate below reads a table - built at that same moment instead of rebuilding one per call. + The rules are written twice on purpose: once fused into ``lowered``'s + guards, where the cost is per call and the answer is a bool, and once + here, where the cost is nothing -- this only ever runs on a call that + has already failed -- and the answer names the operand and both numbers. - ``graph_order`` is the pack's per-view "this operand's layout is the - graph's, not the caller's", or None when they are all the caller's. + What this is NOT is a second way to run the kernel. Two executors are + two answers to what the graph computes, and the differential that + policed them could never catch a misconception they shared, which is + exactly how the axis-order bug survived one. + + Falling off the end means the two readings have drifted apart, which is + the one failure mode splitting them introduces. Loud beats silent. """ recipe = self.recipe _check_plan_device(recipe.device) mnk, axes = recipe.problem(operands, graph_order) check_shapes(recipe, operands, mnk, axes) check_alignment(recipe, operands, axes) - - out_bufs = [operands[o.index] for o in recipe.outputs] - aux_bufs = [operands[x.index] for x in recipe.aux] - c_arg = out_bufs if len(out_bufs) > 1 else out_bufs[0] - # Everything below indexes an operand's axes by position, so an operand - # that arrived in the graph's order is re-expressed in the kernel's here - # rather than threaded through every launcher. - in_bufs = [ - (operands[op.index] if ax is KERNEL_AXES else operands[op.index].permute(ax[AX_BATCH], ax[AX_MN], ax[AX_K])) for op, ax in zip(recipe.inputs, axes) - ] - a_bufs, b_bufs = in_bufs[: recipe.b_at], in_bufs[recipe.b_at :] - - sf_bufs = [operands[s.index] for s in recipe.sf] - if self.chain.is_multi_gemm: - sfa, sfb = sf_bufs[: recipe.b_at], sf_bufs[recipe.b_at :] - if self.block_scale: - pairs = [((a_bufs[ai], sfa[ai]), (b_bufs[bi], sfb[bi])) for ai, bi in self.chain.gemm_operands] - else: - pairs = [(a_bufs[ai], b_bufs[bi]) for ai, bi in self.chain.gemm_operands] - args = (pairs, c_arg, mnk) - else: - args = (a_bufs[0], b_bufs[0], c_arg, mnk, *sf_bufs) - r = self._call_positional(*args, *aux_bufs, stream=stream) - _finalize_reductions(self.chain, out_bufs) - return r - - def _call_positional(self, *args, stream=None): - # Internal launcher (called by __call__ after variant-pack resolve). - # Single-GEMM: (a, b, c, (M,N,K), *aux). Multi-GEMM: first arg is a list - # of per-GEMM (a, b) pairs deduped into the JIT-fixed distinct slots. - if self.chain.is_multi_gemm: - if self.block_scale: - return self._call_block_scale_multi_gemm(*args, stream=stream) - return self._call_multi_gemm(*args, stream=stream) - a, b, c, mnk, *aux = args - # `c` is a single Tensor or a list/tuple in `self.chain.outputs` order - # (outputs bind in chain.outputs order). - outputs_spec = self.chain.outputs - if isinstance(c, (list, tuple)): - cs = list(c) - else: - cs = [c] - if len(cs) != len(outputs_spec): - raise ValueError( - f"this graph has {len(outputs_spec)} output(s) " - f"({[o.source for o in outputs_spec]}); got {len(cs)} runtime " - f"output tensor(s). Pass a list of tensors in slot order." - ) - - expected_a_batch = self.chain.matmul.a_batch - expected_b_batch = self.chain.matmul.b_batch - bad_shapes = len(a.shape) != 3 or len(b.shape) != 3 or a.shape[0] != expected_a_batch or b.shape[0] != expected_b_batch - for spec, ci in zip(outputs_spec, cs): - if len(ci.shape) != 3 or tuple(ci.shape) != _expected_output_shape(spec, self.chain, mnk): - bad_shapes = True - if bad_shapes: - raise ValueError( - f"runtime tensors must be rank-3 with shapes matching the graph " - f"A batch={expected_a_batch}, B batch={expected_b_batch}, " - f"outputs={[ _expected_output_shape(o, self.chain, mnk) for o in outputs_spec ]}; " - 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, stream) - - if self.chain.output_specs: - base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) - else: - base_problem = (mnk[0], mnk[1], mnk[2], max(a.shape[0], b.shape[0])) - # Block-scale: first two aux are the SFA/SFB scale factors (128x4-blocked - # layout), placed right after a/b; the rest are epilogue-fusion tensors. - sf_args: tuple = () - if self.block_scale: - if len(aux) < 2: - raise ValueError("block-scaled matmul call needs sfa, sfb after c: " "compiled(a, b, c, (M,N,K), sfa, sfb, *epilogue_aux)") - sfa, sfb = aux[0], aux[1] - aux = aux[2:] - sfa = _maybe_wrap_layout(sfa.permute(1, 2, 0), _LEADING_DIM_AUX) - sfb = _maybe_wrap_layout(sfb.permute(1, 2, 0), _LEADING_DIM_AUX) - sf_args = (sfa, sfb) - a = a.permute(1, 2, 0) - b = b.permute(1, 2, 0) - cs = [ci.permute(1, 2, 0) for ci in cs] - output_strides = tuple(stride for _spec, ci in zip(outputs_spec, cs) for stride in ci.stride()) - problem_size = ( - *base_problem, - *tuple(a.stride()), - *tuple(b.stride()), - *output_strides, + raise RuntimeError( + "cudnn.frost gemm: the launch path refused this call and no rule explains why -- its " + "per-call guards and the diagnostics here have drifted apart. Guards that fired: " + f"{dict(self.deferrals)}." ) - a = _maybe_wrap_layout(a, _LEADING_DIM_A) - b = _maybe_wrap_layout(b, _LEADING_DIM_B) - cs = [ - (_wrap_raw_tensor(ci) if (spec.is_reduction or spec.is_quant_scale) else _maybe_wrap_layout(ci, _LEADING_DIM_C)) - for spec, ci in zip(outputs_spec, cs) - ] - aux = tuple(_maybe_wrap_layout(_reshape_aux_to_fake(t, ref), _LEADING_DIM_AUX) for ref, t in zip(self.chain.aux_tensors, aux)) - # cute.compile fixes one param per output at JIT time, so pass them flat: - # plain: (a, b, *outputs, mnk, *aux) - # block-scale: (a, b, sfa, sfb, *outputs, mnk, *aux) - if self.use_tma_store: - # TMA mode: the single dense output binds the trailing TMA-only c param. - return self._launchable(problem_size, a, b, *sf_args, *aux, cs[0], stream=_as_custream(stream)) - return self._launchable(problem_size, a, b, *sf_args, *cs, *aux, stream=_as_custream(stream)) - - def _call_multi_gemm(self, gemm_pairs, c, mnk, *aux, stream=None): - """Multi-GEMM call: ``compiled([(A,B0),(A,B1),...], c, (M,N,K), *aux)``. - - Dedup the (a, b) pairs by tensor identity into the JIT-fixed distinct A/B - slots, verify the sharing pattern matches the chain, then pass - ``(a_0.., b_0.., c, mnk4, *aux)`` in kernel-signature order.""" - chain = self.chain - if not isinstance(gemm_pairs, (list, tuple)) or not all(isinstance(p, (list, tuple)) and len(p) == 2 for p in gemm_pairs): - raise ValueError("multi-GEMM call expects a list of (a, b) tensor pairs as the " f"first argument; got {type(gemm_pairs).__name__}") - if len(gemm_pairs) != chain.num_gemms: - raise ValueError(f"this graph has {chain.num_gemms} GEMM(s); got " f"{len(gemm_pairs)} (a, b) pair(s)") - na, nb = chain.num_a_operands, chain.num_b_operands - a_slots: list = [None] * na - b_slots: list = [None] * nb - for (A_g, B_g), (ai, bi) in zip(gemm_pairs, chain.gemm_operands): - for slots, idx, t, role in ( - (a_slots, ai, A_g, "A"), - (b_slots, bi, B_g, "B"), - ): - if slots[idx] is None: - slots[idx] = t - elif slots[idx].data_ptr() != t.data_ptr(): - raise ValueError( - f"multi-GEMM operand sharing mismatch: distinct {role} slot " - f"{idx} was given two different tensors. The runtime sharing " - "pattern must match the graph the kernel was compiled from." - ) - if any(s is None for s in a_slots) or any(s is None for s in b_slots): - raise ValueError("multi-GEMM: not every distinct A/B operand slot was filled") - - outputs_spec = chain.outputs - cs = list(c) if isinstance(c, (list, tuple)) else [c] - if len(cs) != len(outputs_spec): - raise ValueError( - f"this graph has {len(outputs_spec)} output(s) " - f"({[o.source for o in outputs_spec]}); got {len(cs)}. " - "Pass a list of output tensors in slot order." - ) - for spec, ci in zip(outputs_spec, cs): - if len(ci.shape) != 3 or tuple(ci.shape) != _expected_output_shape(spec, chain, mnk): - raise ValueError(f"multi-GEMM output {spec.source!r} must have shape " f"{_expected_output_shape(spec, chain, mnk)}; " f"got {tuple(ci.shape)}") - for role, slots in (("A", a_slots), ("B", b_slots)): - 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, 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] - b_permuted = [t.permute(1, 2, 0) for t in b_slots] - c_permuted = [ci.permute(1, 2, 0) for ci in cs] - output_strides = tuple(stride for _spec, ci in zip(outputs_spec, c_permuted) for stride in ci.stride()) - if not self.block_scale: - problem_size = ( - *base_problem, - *(x for t in a_permuted for x in t.stride()), - *(x for t in b_permuted for x in t.stride()), - *output_strides, - ) - else: - problem_size = base_problem - a_wrapped = [_maybe_wrap_layout(t, _LEADING_DIM_A) for t in a_permuted] - b_wrapped = [_maybe_wrap_layout(t, _LEADING_DIM_B) for t in b_permuted] - cs_wrapped = [ - (_wrap_raw_tensor(ci) if (spec.is_reduction or spec.is_quant_scale) else _maybe_wrap_layout(ci, _LEADING_DIM_C)) - for spec, ci in zip(outputs_spec, c_permuted) - ] - aux = tuple(_maybe_wrap_layout(_reshape_aux_to_fake(t, ref), _LEADING_DIM_AUX) for ref, t in zip(chain.aux_tensors, aux)) - return self._launchable(problem_size, *a_wrapped, *b_wrapped, *cs_wrapped, *aux, stream=_as_custream(stream)) - - def _call_block_scale_multi_gemm(self, gemm_pairs, c, mnk, *aux, stream=None): - """Block-scale multi-GEMM call: - ``compiled([((A,SFA),(B0,SFB0)), ((A,SFA),(B1,SFB1))], c, (M,N,K), *epi_aux)``. - - Each operand is a (packed_data, SF) pair; dedup by packed-data identity - (SF travels with its data → a shared dequant collapses to one operand). - Grouped by kind in the kernel signature (a.., b.., sfa.., sfb..).""" - chain = self.chain - ok = ( - isinstance(gemm_pairs, (list, tuple)) - and gemm_pairs - and all(isinstance(p, (list, tuple)) and len(p) == 2 and all(isinstance(o, (list, tuple)) and len(o) == 2 for o in p) for p in gemm_pairs) - ) - if not ok: - raise ValueError("block-scale multi-GEMM call expects a list of " "((a,sfa),(b,sfb)) pairs as the first argument") - if len(gemm_pairs) != chain.num_gemms: - raise ValueError(f"this graph has {chain.num_gemms} GEMM(s); got {len(gemm_pairs)} pair(s)") - na, nb = chain.num_a_operands, chain.num_b_operands - a_slots: list = [None] * na # (packed_a, sfa) - b_slots: list = [None] * nb - for ((A_g, SFA_g), (B_g, SFB_g)), (ai, bi) in zip(gemm_pairs, chain.gemm_operands): - for slots, idx, data, sf, role in ( - (a_slots, ai, A_g, SFA_g, "A"), - (b_slots, bi, B_g, SFB_g, "B"), - ): - if slots[idx] is None: - slots[idx] = (data, sf) - elif slots[idx][0].data_ptr() != data.data_ptr(): - raise ValueError(f"block-scale multi-GEMM operand sharing mismatch: distinct " f"{role} slot {idx} got two different packed tensors.") - if any(s is None for s in a_slots) or any(s is None for s in b_slots): - raise ValueError("block-scale multi-GEMM: not every distinct operand slot was filled") - - # chain.outputs order. No-epilogue - # → one output per GEMM in GEMM order; fused → one output. - outputs_spec = chain.outputs - cs = list(c) if isinstance(c, (list, tuple)) else [c] - if len(cs) != len(outputs_spec): - raise ValueError( - f"this graph has {len(outputs_spec)} output(s) " - f"({[o.source for o in outputs_spec]}); got {len(cs)}. " - f"Pass a list of output tensors in slot order." - ) - for spec, ci in zip(outputs_spec, cs): - if len(ci.shape) != 3 or tuple(ci.shape) != _expected_output_shape(spec, chain, mnk): - 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, 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] - b_permuted = [d.permute(1, 2, 0) for d, _ in b_slots] - c_permuted = [ci.permute(1, 2, 0) for ci in cs] - a_stride = tuple(a_permuted[0].stride()) - b_stride = tuple(b_permuted[0].stride()) - if any(tuple(t.stride()) != a_stride for t in a_permuted[1:]): - raise ValueError("block-scale multi-GEMM requires all distinct A operands to share layout") - if any(tuple(t.stride()) != b_stride for t in b_permuted[1:]): - raise ValueError("block-scale multi-GEMM requires all distinct B operands to share layout") - output_strides = tuple(stride for _spec, ci in zip(outputs_spec, c_permuted) for stride in ci.stride()) - problem_size = ( - *base_problem, - *a_stride, - *b_stride, - *output_strides, - ) - a_w = [_maybe_wrap_layout(t, _LEADING_DIM_A) for t in a_permuted] - b_w = [_maybe_wrap_layout(t, _LEADING_DIM_B) for t in b_permuted] - sfa_w = [_maybe_wrap_layout(s.permute(1, 2, 0), _LEADING_DIM_AUX) for _, s in a_slots] - sfb_w = [_maybe_wrap_layout(s.permute(1, 2, 0), _LEADING_DIM_AUX) for _, s in b_slots] - cs_w = [ - (_wrap_raw_tensor(t) if (spec.is_reduction or spec.is_quant_scale) else _maybe_wrap_layout(t, _LEADING_DIM_C)) - for spec, t in zip(outputs_spec, c_permuted) - ] - aux = tuple(_maybe_wrap_layout(_reshape_aux_to_fake(t, ref), _LEADING_DIM_AUX) for ref, t in zip(chain.aux_tensors, aux)) - return self._launchable(problem_size, *a_w, *b_w, *sfa_w, *sfb_w, *cs_w, *aux, stream=_as_custream(stream)) def _mma_a_dtype(chain: FusionChain) -> str: @@ -2736,6 +2512,28 @@ def _check_block_quant_supported( _FORCE_STG_EPI = False +def _check_executable(chain: FusionChain) -> None: + """Can this engine RUN the graph, or only render a kernel for it? + + A dense or block-scale chain has ONE call path -- the one lowered from the + recipe -- so what that path cannot serve is a graph this engine declines, + not a call that takes a slower route. Declining sends it to the backend, + which is where it would have gone had this engine never been asked; keeping + a second executor for it is how the two drift apart. + + MoE is the exception and stays on its own launchers: it is >= 2 launches + with a workspace, and has no recipe to lower from. + """ + if chain.has_moe: + return + if not _TVM_FFI_OK: + raise NotImplementedError("the tvm-ffi front door is not installed, and the launch path this engine has needs it (pip install apache-tvm-ffi)") + if any(red.mode == "norm2" for red in chain.reductions): + raise NotImplementedError("a norm2 reduction takes a square root after the kernel, which is a device operation this engine does not own") + if chain.is_multi_gemm and not chain.output_specs: + raise NotImplementedError("a multi-GEMM with no dense output has no tensor to read the batch extent off") + + def probe_supported( graph: cudnn.pygraph, config: TileConfig = DEFAULT_CONFIG, @@ -2751,6 +2549,7 @@ def probe_supported( Block-scale / MoE gate inside their ``_jit_*`` compile paths; here a successful analysis is treated as eligible (full validation at compile).""" chain, _binding = analyze_with_binding(graph) + _check_executable(chain) if chain.has_moe or chain.has_block_scale: return # specialized paths validate at compile if chain.is_multi_gemm: diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index 56ad6417e..cda0e63c5 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -32,12 +32,12 @@ def __init__(self, compiled): # as ambiguous. self._tensors = list(compiled.binding.bound_tensors()) self._operand_indices = None - # Which call path this plan uses is a property of the compiled kernel, - # so it is chosen here and not re-asked per execute. ``lowered`` is the - # closure the recipe is captured into when the kernel is a graph it - # serves; ``launch`` interprets the same recipe for everything else. + # The one launch path a dense or block-scale kernel has: the closure the + # recipe is captured into. A graph it cannot serve was declined at + # check_support, so there is nothing to fall back TO -- None here means + # MoE, which takes the variant-pack dict and its own workspace below. self._lowered = getattr(compiled, "lowered", None) - self._launch = self._lowered or getattr(compiled, "launch", None) + self._launch = self._lowered def get_workspace_size(self) -> int: return int(getattr(self._compiled, "workspace_bytes", 0) or 0) @@ -65,8 +65,7 @@ def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: # the buffers arrive in that order and the launcher indexes them. # Which AXIS ORDER each one arrived in is a per-call fact only the # pack knows, since a bare address wears the graph's layout. None - # means every operand here is the caller's own, which is what both - # launchers are written for. + # means every operand here is the caller's own. graph_order = None borrowed = variant_pack.graph_described if borrowed: diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py index cfb06dc39..b570e3b61 100644 --- a/python/cudnn/gemm/frost/recipe.py +++ b/python/cudnn/gemm/frost/recipe.py @@ -7,16 +7,21 @@ rules, which outputs are reductions and what order the kernel takes its parameters in are all fixed by the time cute hands back a launchable. A call carries M, N, K, the strides and the pointers, and nothing else. This module -writes the first set down, so that neither of the two call paths re-derives them. - -They read the same recipe but do not share a body: ``CompiledFusedGemm.launch`` -interprets it (through :func:`check_shapes` and :func:`check_alignment`) by -walking the operand structure, and ``CompiledFusedGemm.lowered`` is the closure -``_lower`` captures it into, where the same walk is a loop over tuples flat -enough to need no attribute lookup. That is a compiler beside its interpreter, -kept honest the way those always are -- ``test_execute_recipe.py`` runs both -over the same accepts and rejects and requires the same answer. What it CANNOT -catch is a misconception they share, which is how the axis-order bug survived it. +writes the first set down, so that no call re-derives them. + +One reading of it RUNS: ``CompiledFusedGemm.lowered``, the closure ``_lower`` +captures it into, where every check is a loop over tuples flat enough to need no +attribute lookup. The other only EXPLAINS -- ``CompiledFusedGemm.explain``, which +walks the operand structure through :func:`check_shapes` and +:func:`check_alignment` to name what is wrong with a call the first one refused, +and never launches anything. Two executors would have been two answers to what +the graph computes, and a differential between them cannot catch a misconception +they share, which is how the axis-order bug survived one. + +So the rules here are written twice and the launch is written once: the fast +form pays per call and answers a bool, the readable form runs only on a call +that has already failed. If they ever disagree, ``explain`` finds nothing and +says so rather than returning quietly. The field that makes one loop serve six flavors is :attr:`GemmRecipe.arg_plan`: what differs between plain, aux, multi-output, multi-GEMM and block scale is @@ -160,7 +165,6 @@ class Output: align: int raw: bool # the kernel takes only its pointer init: Any = None # reduction identity, seeded before the kernel runs - sqrt: bool = False # norm2 takes a square root after @dataclass(frozen=True) @@ -427,11 +431,10 @@ def build(compiled) -> GemmRecipe: aux_reqs = _aux_align_reqs(chain, vec_bytes=compiled.vec_bytes_epi) outputs, seeds = [], [] for i, (spec, t) in enumerate(zip(chain.outputs, binding.outputs)): - init, sqrt = None, False + init = None if spec.is_reduction: red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] init = REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] - sqrt = red.mode == "norm2" seeds.append((order[id(t)], init_word(red.compute_dtype, init))) outputs.append( Output( @@ -441,7 +444,6 @@ def build(compiled) -> GemmRecipe: align=out_reqs[i], raw=bool(spec.is_reduction or spec.is_quant_scale), init=init, - sqrt=sqrt, ) ) diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index bd4aa179c..2f7701dc6 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -1,19 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The build-time recipe, and the closure lowered from it. - -``_lower`` captures the recipe into a call path that loops over it flat; -``launch`` interprets the same recipe by walking the operand structure. That is -a compiler beside its interpreter, and the two can drift -- an earlier -hand-written version of the emitted path lost the operand batch check and pinned -an fp4 output at N instead of N/2, both of which made it accept or reject calls -the general path did not. - -So the differentials below are the point of this file: every case and every -flavor runs through BOTH entry points, and the two must return the same verdict -and the same numbers. The rest asserts that the fast path is the one a public -``execute()`` actually takes, and that what it declines is declined on purpose. +"""The build-time recipe, and the one call path lowered from it. + +``_lower`` captures the recipe into the loop that launches; ``explain`` reads the +same recipe to say what is wrong with a call that loop refused, and runs nothing. +There is no second executor to differential against, and that is deliberate: two +readings of one wrong plan agreeing proves nothing, which is exactly how the +axis-order bug survived a differential that ran both. + +What replaces it is the BACKEND -- the two tests below that run a shape on a +cuDNN plan and a FROST plan and require the same numbers -- and the invariant +that makes a rejection meaningful: the set of calls the launch path refuses must +equal the set of illegal calls, asserted per flavor through ``deferrals``. """ from __future__ import annotations @@ -28,7 +27,8 @@ import cudnn import cudnn.gemm.frost # noqa: F401 — installs the cudnn.pygraph recorder hook from cudnn.engines import is_python_engine -from cudnn.gemm.frost.compiler import jit_from_cudnn_graph +from cudnn.frost.buffers import collapse_layout, fill_word_strided_async, init_word +from cudnn.gemm.frost.compiler import jit_from_cudnn_graph, probe_supported from cudnn.gemm.frost.graph_analyzer import resolve_variant_pack from cudnn.gemm.frost.recipe import _output_rule, expected_shape @@ -77,6 +77,66 @@ def test_a_quant_scale_output_is_fixed_at_build(): assert expected_shape(rule, 999, 999) == (1, 128, 4) +# --- seeding a 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 below to a single memset for a + contiguous tap and one per batch for a padded one, rather than one per row. + """ + assert collapse_layout(shape, stride) == want + + +@requires_sm100 +@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_strided_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 kernel's. What it has to get + right is exactly this: every element the view covers, and nothing between + them. + """ + 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")) + + +@pytest.mark.parametrize( + "shape,stride,why", + [ + ((1, 8, 1), (0, 0, 1), "alias"), # stride 0 over a real extent + ((1, 4, 4), (16, 2, 1), "overlap"), # rows closer together than they are wide + ], +) +def test_a_reduction_output_that_writes_a_byte_twice_is_rejected(shape, stride, why): + """Two elements at one address is a write race, not a layout to support. + + Named here rather than left to the driver, which reports the second as a + pitch smaller than the width and says nothing about whose buffer it was. + """ + with pytest.raises(ValueError, match=why): + fill_word_strided_async(0, shape, stride, 4, 0, None) + + # --- which flavors lower, and which decline --------------------------------- @@ -97,6 +157,15 @@ def _reduction_graph(): return g +def _row_reduction_graph(pad=4): + """A per-row tap, declared into a padded buffer when ``pad`` > 1.""" + g = _plain_graph() + Y = [t for t in g._nodes[-1].outputs.values()][0] + R = g.reduction(input=Y, mode=cudnn.reduction_mode.ADD, name="red") + R.set_dim([1, M, 1]).set_stride([M * pad, pad, 1]).set_output(True).set_data_type(F32) + return g + + def _aux_graph(): 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]) @@ -166,23 +235,51 @@ def test_which_flavors_lower(build, n_b, n_out, n_aux): @requires_sm100 @pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) -def test_a_legal_call_defers_to_the_interpreter_for_nothing(build, n_b, n_out, n_aux): - """The interpreter is migration scaffolding, so that has to be measurable. +def test_a_legal_call_of_every_flavor_is_refused_for_nothing(build, n_b, n_out, n_aux): + """The invariant, per flavor: refusing a legal call is a bug, not a detour. - Every reason the fast path can hand a call over is counted, and a legal call - of every flavor must trigger none of them -- otherwise "the fallback only - catches what the fast path declines" is a claim with nothing behind it, and - the day it is deleted is the day the regression appears. + While there was an interpreter behind it, a fast path that gave up on a + legal call cost time and nothing else, so nothing asserted it did not. This + is what that assertion looks like: every reason the launch path can refuse + is counted, and a legal call of every flavor must trigger none of them. """ compiled = jit_from_cudnn_graph(build()) if compiled.lowered is None: pytest.skip(f"this build does not lower: {compiled.declined}") operands = _buffers_for(compiled) + for o in compiled.recipe.outputs: + operands[o.index].zero_() compiled.lowered(operands, stream=None) torch.cuda.synchronize() + assert operands[compiled.recipe.outputs[0].index].abs().sum() > 0 assert dict(compiled.deferrals) == {} +@requires_sm100 +def test_a_padded_reduction_output_stays_on_the_fast_path(): + """A tap declared into a padded buffer is a legal call, not a slow one. + + The fast path could seed exactly one dense run, so it handed a padded tap + to the interpreter -- which seeded it by calling ``fill_()`` on whatever the + caller passed. Both are gone: the seed is the driver's 2D memset, and it + must touch every element the view covers and nothing between them. + """ + compiled = jit_from_cudnn_graph(_row_reduction_graph()) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") + operands = _buffers_for(compiled) + tap = [o for o in compiled.recipe.outputs if o.init is not None][0] + pad = torch.full((1, M, 4), -1.0, dtype=torch.float32, device="cuda") + operands[tap.index] = pad[:, :, :1] + + compiled.lowered(operands, stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {} + a, b = (operands[op.index] for op in compiled.recipe.inputs) + torch.testing.assert_close(pad[:, :, 0], torch.einsum("bmk,bnk->bm", a.float(), b.float()), atol=1.0, rtol=2e-2) + assert torch.equal(pad[:, :, 1:], torch.full((1, M, 3), -1.0, device="cuda")) + + @requires_sm100 def test_every_decline_names_its_reason(): """A plan without a fast path says which rule denied it, not just ``None``.""" @@ -195,30 +292,66 @@ def test_every_decline_names_its_reason(): @requires_sm100 def test_a_deferral_is_counted_under_the_rule_that_caused_it(): - """A bare address is the one deferral a legal call still takes.""" + """Only an ILLEGAL call leaves the fast path, and it says which rule sent it.""" compiled = jit_from_cudnn_graph(_plain_graph()) if compiled.lowered is None: pytest.skip("this build does not lower") - a, b, c = _operands() + a, b, c = _wrong_major() operands = _bound_buffers(compiled, a, b, c) - compiled.lowered(operands, (True, False, False), stream=None) - torch.cuda.synchronize() - assert dict(compiled.deferrals) == {"graph-described operand": 1} + with pytest.raises(ValueError): + compiled.lowered(operands, stream=None) + assert dict(compiled.deferrals) == {"input layout": 1} @requires_sm100 -def test_a_post_kernel_finalize_is_declined(): +@pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) +def test_an_operand_in_the_graph_s_axis_order_stays_on_the_fast_path(build, n_b, n_out, n_aux): + """A bare device address wears the graph's layout, and that is legal. + + cuDNN declares a matmul's B ``[batch, K, N]`` where this kernel reads + ``(batch, N, K)``, so an operand the pack described FROM the graph -- a bare + address has no geometry of its own -- arrives with its axes the other way + round. Re-labelling one is a permute the recipe already knows the shape of, + so it is not a reason to leave the fast path; it used to be, and a legal + call paid a whole interpreted pass for it. + """ + compiled = jit_from_cudnn_graph(build()) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") + operands = _buffers_for(compiled) + outs = compiled.recipe.outputs + borrowed = tuple(True for _ in operands) + + runs = [] + for order in (None, borrowed): + for o in outs: + operands[o.index].zero_() + # Same buffers either way: what changes is which axis order they claim, + # and the graph's is the one the declaration already describes them in. + compiled.lowered(_as_declared(compiled, operands) if order else operands, order, stream=None) + torch.cuda.synchronize() + runs.append([operands[o.index].clone() for o in outs]) + assert dict(compiled.deferrals) == {} + for o, own, graph in zip(outs, *runs): + exact = o.init is None + torch.testing.assert_close(own, graph, atol=0 if exact else 1e-2, rtol=0 if exact else 1e-5) + + +@requires_sm100 +def test_a_post_kernel_finalize_is_refused_before_a_plan_exists(): """``norm2`` takes a square root through the caller's buffer after the kernel. - The backend refuses that reduction while the graph is being lowered, so it - never reaches a plan -- but the recipe records ``sqrt`` and the lowered path - declines on it, because the alternative is a device operation the engine - does not own and would have to borrow off whatever the caller passed. + That is a device operation this engine does not own, and it has one call + path -- so the GRAPH is declined and goes to the backend, rather than being + compiled into a kernel only a second executor could run. Which executor a + graph needs is not a question this engine wants to be able to ask. """ - compiled = jit_from_cudnn_graph(_plain_graph()) - assert compiled.lowered is not None - compiled.recipe = replace(compiled.recipe, outputs=(replace(compiled.recipe.outputs[0], sqrt=True),)) - assert compiled._lower() is None + g = _plain_graph() + Y = [t for t in g._nodes[-1].outputs.values()][0] + R = g.reduction(input=Y, mode=cudnn.reduction_mode.NORM2, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]).set_output(True).set_data_type(F32) + with pytest.raises(NotImplementedError, match="square root"): + probe_supported(g) @requires_sm100 @@ -282,39 +415,19 @@ def _buffers_for(compiled): return bufs -@requires_sm100 -@pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) -def test_every_flavor_agrees_with_the_interpreter(build, n_b, n_out, n_aux): - """One loop serves six flavors, so one of them drifting is the failure mode. - - The interpreter is the reference: it reads the same recipe but assembles the - launch by walking the operand structure, which is the thing the lowered path - replaced with a table. Bit-exact is the bar -- the two issue the same kernel - with the same arguments or they do not agree. - - Except for a reduction output, which is not bit-reproducible against ITSELF: - the taps land through cross-CTA float atomics, so the order varies with - scheduling. Measured on this shape, the same path run six times spreads - 0.0049 while the two paths differ by 0.00024 -- twenty times smaller than - the noise, so the tolerance below is the noise floor and not a slackened bar. - """ - compiled = jit_from_cudnn_graph(build()) - if compiled.lowered is None: - pytest.skip("this build does not lower (no tvm-ffi front door)") - operands = _buffers_for(compiled) - outs = compiled.recipe.outputs +def _as_declared(compiled, operands): + """The same memory, each input re-labelled the way the graph declares it. - runs = [] - for run in (compiled.lowered, compiled.launch): - for o in outs: - operands[o.index].zero_() - run(operands, stream=None) - torch.cuda.synchronize() - runs.append([operands[o.index].clone() for o in outs]) - for o, fast, slow in zip(outs, *runs): - exact = o.init is None - torch.testing.assert_close(fast, slow, atol=0 if exact else 1e-2, rtol=0 if exact else 1e-5) - assert runs[0][0].abs().sum() > 0 # a path that wrote nothing would also "agree" + What the pack hands over for a bare address: with no geometry of its own, + the declaration IS the description, so B arrives ``[batch, K, N]``. + """ + out = list(operands) + for op in compiled.recipe.inputs: + inverse = [0, 0, 0] + for role, axis in enumerate(op.declared): + inverse[axis] = role + out[op.index] = operands[op.index].permute(*inverse) + return out def _verdict(run, operands, c): @@ -374,29 +487,43 @@ def _padded_rows(): @requires_sm100 @pytest.mark.parametrize( - "case", - (_good, _short_k, _wrong_major, _wrong_output_shape, _misaligned_output, _misaligned_a, _padded_rows), - ids=lambda f: f.__name__.strip("_"), + "case,verdict", + ( + (_good, "ran"), + (_padded_rows, "ran"), # the outer stride is free + (_short_k, "rejected"), + (_wrong_major, "rejected"), + (_wrong_output_shape, "rejected"), + (_misaligned_output, "rejected"), + (_misaligned_a, "rejected"), + ), + ids=lambda x: x if isinstance(x, str) else x.__name__.strip("_"), ) -def test_lowered_and_interpreted_agree(case): - """Same operands, both entry points, one verdict. - - An unsound fast path shows up here as an accept where the general path - rejects -- which is exactly how the two regressions this replaced would have - read. +def test_the_launch_path_accepts_exactly_the_legal_calls(case, verdict): + """The invariant the second executor used to make untestable. + + While ``lowered`` could hand anything it was unsure of to an interpreter + that ran it anyway, "the fast path rejects a legal call" was a performance + bug at worst and nothing asserted otherwise. There is one path now, so a + rejection IS the answer -- which makes the two directions both real + failures: refusing a legal call breaks it, and accepting an illegal one runs + a kernel over memory the caller did not describe. + + ``deferrals`` is what pins that down. Empty means every rule passed; a + rejection must name the rule that caused it, and a rejection with nothing + named is drift between the guards and the diagnostics, which ``explain`` + raises on separately. """ compiled = jit_from_cudnn_graph(_plain_graph()) if compiled.lowered is None: - pytest.skip("this build does not lower (no tvm-ffi front door)") + pytest.skip(f"this build does not lower: {compiled.declined}") a, b, c = case() - fast, fast_out = _verdict(compiled.lowered, _bound_buffers(compiled, a, b, c), c) - slow, slow_out = _verdict(compiled.launch, _bound_buffers(compiled, a, b, c), c) - assert fast == slow, f"lowered says {fast}, interpreted says {slow}" - if fast == "ran": - torch.testing.assert_close(fast_out, slow_out, atol=0, rtol=0) - ref = torch.einsum("bmk,bnk->bmn", a.float(), b.float()) - torch.testing.assert_close(fast_out.float(), ref, atol=2e-1, rtol=2e-2) + got, out = _verdict(compiled.lowered, _bound_buffers(compiled, a, b, c), c) + assert got == verdict + assert bool(compiled.deferrals) == (verdict == "rejected") + if verdict == "ran": + torch.testing.assert_close(out.float(), torch.einsum("bmk,bnk->bmn", a.float(), b.float()), atol=2e-1, rtol=2e-2) def _matmul_on(batch, m, n, k, want_frost, a, b, c): @@ -460,14 +587,16 @@ def _nvfp4_buffers(compiled): @requires_sm100 @pytest.mark.parametrize("gemms", (1, 2), ids=("single", "multi")) -def test_block_scale_lowers_and_agrees_with_the_interpreter(gemms): +def test_block_scale_lowers_and_runs(gemms): """Block scale is the flavor with the most per-call table in it. Its scale factors ride in the launch argument list but NOT in ``problem_size``, its blob size is re-synthesized from M/N/K rather than read off the buffer, and the multi-GEMM form sends one A stride triple for every - operand instead of one each. Four recipe fields, so it is the one most worth - running against the interpreter. + operand instead of one each. Four recipe fields, so it is the flavor most + likely to reach the launch with an argument list that does not match the + kernel's signature -- which is what running it here catches. What the + numbers should BE is checked against torch through the public entry point. """ # Two GEMMs do not fit the auto-selected cta_n=256 in TMEM, so pin a # geometry that does; which config the engine picks for one GEMM is the @@ -480,14 +609,10 @@ def test_block_scale_lowers_and_agrees_with_the_interpreter(gemms): assert bool(compiled.recipe.shared_layout) == (gemms > 1) operands, out = _nvfp4_buffers(compiled) - runs = [] - for run in (compiled.lowered, compiled.launch): - out.zero_() - run(operands, stream=None) - torch.cuda.synchronize() - runs.append(out.clone()) - torch.testing.assert_close(runs[0], runs[1], atol=0, rtol=0) - assert runs[0].abs().sum() > 0 + compiled.lowered(operands, stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {} + assert out.abs().sum() > 0 @requires_sm100 @@ -617,10 +742,9 @@ def test_operand_batch_is_checked(): """The graph pins each operand's batch, and a launch that ignored it read one batch of A against three of B.""" compiled = jit_from_cudnn_graph(_plain_graph(batch=2)) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") a, b, c = _operands(batch=2) one_batch_b = b[:1].contiguous() - for run in (compiled.lowered, compiled.launch): - if run is None: - continue - with pytest.raises(ValueError): - run(_bound_buffers(compiled, a, one_batch_b, c), stream=None) + with pytest.raises(ValueError): + compiled.lowered(_bound_buffers(compiled, a, one_batch_b, c), stream=None) diff --git a/test/python/gemm/frost/test_matmul_epilogue_fusion.py b/test/python/gemm/frost/test_matmul_epilogue_fusion.py index bd9768237..e7b63013d 100644 --- a/test/python/gemm/frost/test_matmul_epilogue_fusion.py +++ b/test/python/gemm/frost/test_matmul_epilogue_fusion.py @@ -2531,8 +2531,6 @@ def _pw_aux_order(compiled, aux_bufs): "avg_row": (cudnn.reduction_mode.AVG, (1, _PW_M, 1), lambda s: s.mean(dim=1, keepdim=True).view(1, _PW_M, 1), 0.0), "avg_col": (cudnn.reduction_mode.AVG, (1, 1, _PW_N), lambda s: s.mean(dim=0, keepdim=True).view(1, 1, _PW_N), 0.0), "norm1_row": (cudnn.reduction_mode.NORM1, (1, _PW_M, 1), lambda s: s.abs().sum(dim=1, keepdim=True).view(1, _PW_M, 1), 0.0), - "norm2_row": (cudnn.reduction_mode.NORM2, (1, _PW_M, 1), lambda s: s.pow(2).sum(dim=1, keepdim=True).sqrt().view(1, _PW_M, 1), 0.0), - "norm2_full": (cudnn.reduction_mode.NORM2, (1, 1, 1), lambda s: s.pow(2).sum().sqrt().view(1, 1, 1), 0.0), "mul_row": (cudnn.reduction_mode.MUL, (1, _PW_M, 1), lambda s: s.prod(dim=1, keepdim=True).view(1, _PW_M, 1), 1.0), "mul_no_zeros_row": ( cudnn.reduction_mode.MUL_NO_ZEROS, @@ -2585,3 +2583,31 @@ def test_reduction_mode_coverage(case): s = s.to(torch.bfloat16).float() ref = ref_fn(s.double()).float() torch.testing.assert_close(r, ref, atol=5e-2, rtol=2e-2) + + +@requires_sm100 +def test_norm2_is_the_one_reduction_mode_this_engine_declines(): + """It is the only mode that is not finished when the kernel is. + + Its taps land through cross-CTA atomics, so the square root cannot go in the + epilogue -- it has to run after every CTA has contributed. That is a device + operation this engine does not own, and the version that borrowed the + caller's ``sqrt_()`` worked only while the caller passed a torch tensor. + + Declining costs nothing reachable: the BACKEND refuses a norm2 reduction + descriptor while the graph is still being lowered, so no public + ``execute()`` ever gets a plan for one either (see + ``test_public_execute_flavors.py::test_norm2_reduction_is_refused_at_build``). + Recorded here because the mode is otherwise in every list of the reductions + the epilogue supports. + """ + from cudnn.gemm.frost.compiler import probe_supported + + g, C = _pw_matmul_graph() + src = g.swish(input=C, name="s") + src.set_data_type(cudnn.data_type.BFLOAT16).set_output(True) + red = g.reduction(input=src, mode=cudnn.reduction_mode.NORM2, name="red") + red.set_dim([1, _PW_M, 1]).set_stride([_PW_M, 1, 1]) + red.set_output(True).set_data_type(cudnn.data_type.FLOAT) + with pytest.raises(NotImplementedError, match="square root"): + probe_supported(g) diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py index da3e2085d..084f81eff 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -348,7 +348,9 @@ def test_bare_address_operands(): The operand a bare address describes is the GRAPH's declaration, which orders a matmul's B ``[batch, K, N]`` where a caller allocates it ``(batch, N, K)``. Reading an extent by axis position answers one of those - and not the other, so the recipe records which axis carries M/N/K instead. + and not the other, so the recipe records which axis carries M/N/K instead -- + and re-labelling one into the other is a permute, which is why this stays on + the fast path rather than costing an interpreted pass for being legal. """ 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]) @@ -361,6 +363,7 @@ 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) + assert dict(g._compiled_plans[g._plan_index]._compiled.deferrals) == {} @_GPU From cd962a844f566108ad7e70895219ced42ac0c9f7 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 16:25:04 -0700 Subject: [PATCH 8/9] Fix three defects in the one-executor launch path All three are the same shape: a per-ROLE fact treated as a per-SLOT one, or a write issued before every rule has had its say. Found by review on the pushed head, not by the suite -- each is recorded here with the test that now pins it. matmul(A, A) binds ONE slot as both operands: recipe.build keys the operand map on id(tensor), so the A role and the B role share an index. The graph-order pre-pass re-labelled that slot in place, so permuting it for B also handed A a transposed view of the same memory -- wrong numbers, no guard fired, and the checker then correctly reported that no rule explained the refusal. `lowered` now builds one view per ROLE and never rewrites a slot; `shared_layout` moved to role positions for the same reason. The extra list costs nothing measurable. The reduction seed wrote before the launch could reject the call. A fp32 tap bound to a fp16 buffer took a 32-bit word times numel and ran past the end of the allocation; the tvm-ffi front door does refuse the dtype, but only after the fill. Two taps could also leave the first seeded and the second refused. And the overlap rule was too weak -- `pitch >= width` accepts shape (2, 2) stride (2, 2), whose two axes land on the same element. So the fill is planned before it is issued: strided_fill_plan returns a verified plan or None, the element width is a rule of the call, and every seed is planned before any is written. Two fast guards had no matching rule in the checker: block-scale shared_layout, and a rank-2 operand, which additionally made recipe.problem raise IndexError from an axis index rather than a message. Both now have one, and problem checks the rank before it indexes anything. A multi-GEMM with no dense output is supported again rather than declined. The decline was written when two multi-GEMM launchers read the batch off cs[0]; those launchers are deleted, the recipe already carries a batch for the case, and it is a flavor in the test table now. That makes norm2 the only capability this branch removes, which is what the PR claims. The invariant claim was too strong and is corrected in the code, the tests, the handoff and the PR: `deferrals` empty on a legal call proves nothing legal is REFUSED. It cannot prove nothing illegal is ACCEPTED, because an accepted call never reaches the counter. That direction is covered case by case and by the backend differentials; closing it properly means generating the fused guards and the readable diagnostics from one ordered list of checks, which is still open. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/gemm/frost/compiler.py | 114 +++++++---- python/cudnn/gemm/frost/recipe.py | 69 ++++++- test/python/gemm/frost/test_execute_recipe.py | 182 +++++++++++------- 3 files changed, 244 insertions(+), 121 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 86a921b4a..2969bfc9a 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1919,10 +1919,17 @@ def _lower(self): how to get that 12% back for every flavor at once. The body never raises. What it refuses it hands to ``explain``, which - owns every rejection message and does not run anything -- so the set of - calls this refuses must EQUAL the set of illegal calls, and a legal one - it will not serve is a bug rather than a slower route. ``deferrals`` - counts each refusal under its rule, which is what makes that testable. + owns every rejection message and does not run anything -- so the goal is + that the set of calls this refuses EQUALS the set of illegal calls, and + a legal one it will not serve is a bug rather than a slower route. + + ``deferrals`` gets one direction of that, and only one: a legal call of + every flavor must leave it empty, so nothing legal is refused. It says + nothing about the other direction, because a call this ACCEPTS never + reaches the counter -- an illegal call slipping through is caught, if at + all, by the per-case rejection tests and by the backend differentials. + Closing it properly means generating both the fused guards and the + readable diagnostics from one ordered list of checks. None means this object has no call path at all, which the engine's support gate already refused; ``declined`` says why. @@ -1952,9 +1959,12 @@ def decline(reason: str) -> None: # declares [b, K, N] where the kernel reads (b, N, K). Re-labelling one # into kernel order is a permute, so the checks below read one order and # the second axis map costs nothing on a call that does not use it. + # Keyed by POSITION in `inputs`, never by operand index: matmul(A, A) + # binds ONE slot as both operands, and re-labelling the slot would hand + # the other role a transposed view of the same memory. renorm = tuple( - (op.index, tuple(op.declared), op.declared_layout[0] if len(op.declared_layout[0]) == 3 else None, op.declared_layout[1]) - for op in r.inputs + (j, tuple(op.declared), op.declared_layout[0] if len(op.declared_layout[0]) == 3 else None, op.declared_layout[1]) + for j, op in enumerate(r.inputs) if op.declared != KERNEL_AXES ) @@ -1963,45 +1973,50 @@ def decline(reason: str) -> None: # the axis whose stride must be 1 and the axis whose extent enters the # TMA rule -- the same axis, by the definition of major. stride_ins = r.stride_ins - ins = tuple((op.index, op.kc, op.modulus, op.batch, op.kpack, op.is_b, i in stride_ins) for i, op in enumerate(r.inputs)) + in_slots = tuple(op.index for op in r.inputs) + ins = tuple((op.kc, op.modulus, op.batch, op.kpack, op.is_b, i in stride_ins) for i, op in enumerate(r.inputs)) outs = tuple((o.index, *(x for axis in o.rule for x in axis), o.align) for o in r.outputs) auxs = tuple((x.index, x.align) for x in r.aux) sfs = tuple((s.index, s.is_a, r.inputs[s.operand_at].batch) for s in r.sf) - args = r.arg_plan + # Every input is one of the FIRST arguments, in order, so the plan only + # has to carry the ones after them. + tail = r.arg_plan[len(r.inputs) :] shared, seeds = r.shared_layout, r.seeds - ai, bi, a_kpack = r.a.index, r.b.index, r.a.kpack + a_pos, b_pos, a_kpack = r.a_at, r.b_at, r.a.kpack block_size, batch = r.block_size, r.batch device, launchable, refuse = r.device, self._launchable, self.explain - fill_word, fill_strided, is_contiguous = buffers.fill_word_async, buffers.fill_word_strided_async, buffers.is_contiguous + fill_word, fill_plan, apply_fill = buffers.fill_word_async, buffers.strided_fill_plan, buffers.apply_fill_plan + is_contiguous = buffers.is_contiguous # Named so a test can assert which rule refused a call, and that a legal # call trips none. Incremented only on the path that is already raising. gave_up = self.deferrals def lowered(operands, graph_order=None, stream=None): _check_plan_device(device) - # Which axis order an operand arrived in, by the backend's own rule: + # Which axis order each input arrived in, by the backend's own rule: # the descriptor defines the tensor and the pack supplies a pointer, # so a slot the pack described FROM the graph, or a buffer reporting # exactly the declared (dim, stride), is the declaration. Everything - # else is the caller's own labelling. Settled here, once, so the - # checks below all read the kernel's order. - ops = operands - for idx, perm, dsh, dst in renorm: - v = operands[idx] + # else is the caller's own labelling. + # + # One view per ROLE. Re-labelling the SLOT would be cheaper and + # wrong: matmul(A, A) binds a single buffer as both operands, and + # the two roles read it through different axis maps. + vs = [operands[i] for i in in_slots] + for j, perm, dsh, dst in renorm: + v = vs[j] sh = v.shape if len(sh) != 3: continue # no axis map fits it; the loop below is what refuses it - declared = bool(graph_order) and graph_order[idx] + declared = bool(graph_order) and graph_order[in_slots[j]] if not declared and dsh is not None: declared = sh[0] == dsh[0] and sh[1] == dsh[1] and sh[2] == dsh[2] if declared: st = v.stride() declared = st[0] == dst[0] and st[1] == dst[1] and st[2] == dst[2] if declared: - if ops is operands: - ops = list(operands) - ops[idx] = v.permute(*perm) - a_sh, b_sh = ops[ai].shape, ops[bi].shape + vs[j] = v.permute(*perm) + a_sh, b_sh = vs[a_pos].shape, vs[b_pos].shape if len(a_sh) != 3 or len(b_sh) != 3: gave_up["A or B is not rank 3"] += 1 return refuse(operands, graph_order) @@ -2009,8 +2024,7 @@ def lowered(operands, graph_order=None, stream=None): # permute(1, 2, 0) relabels axes, so the strides the kernel wants are # that rotation of the ones each check already read -- no second read. problem = [m, n, k, batch] - for idx, kc, mod, ebatch, kpack, is_b, takes_stride in ins: - v = ops[idx] + for (kc, mod, ebatch, kpack, is_b, takes_stride), v in zip(ins, vs): sh, st = v.shape, v.stride() if ( len(sh) != 3 @@ -2028,7 +2042,7 @@ def lowered(operands, graph_order=None, stream=None): # Each output axis is a constant, M, or N over a divisor (fp4 packs # two along N) -- the rule the build recorded, read back per axis. for idx, k0, v0, k1, v1, k2, v2, align in outs: - v = ops[idx] + v = operands[idx] sh, st = v.shape, v.stride() if ( len(sh) != 3 @@ -2041,14 +2055,14 @@ def lowered(operands, graph_order=None, stream=None): return refuse(operands, graph_order) problem += (st[1], st[2], st[0]) for idx, align in auxs: - v = ops[idx] + v = operands[idx] if tensor_alignment(tuple(v.shape), tuple(v.stride()), v.element_size(), ptr=v.data_ptr()) < align: gave_up["aux alignment"] += 1 return refuse(operands, graph_order) if sfs: k4 = ((k // block_size) + 3) // 4 for idx, is_a, sf_batch in sfs: - v = ops[idx] + v = operands[idx] count = int(v.numel()) if ( _pow2_floor(v.data_ptr()) < 16 @@ -2058,24 +2072,41 @@ def lowered(operands, graph_order=None, stream=None): gave_up["scale-factor blob"] += 1 return refuse(operands, graph_order) for lead, followers in shared: - st = tuple(ops[lead].stride()) + st = tuple(vs[lead].stride()) for j in followers: - if tuple(ops[j].stride()) != st: + if tuple(vs[j].stride()) != st: gave_up["operands do not share a layout"] += 1 return refuse(operands, graph_order) - # Seeding is a write, so it goes last: every rule above has had its - # say. A padded tap costs one 2D memset per batch instead of one - # dense one, which is a price and not a reason to leave. - for idx, word in seeds: - v = ops[idx] - sh, st = v.shape, v.stride() - if is_contiguous(sh, st): - fill_word(v.data_ptr(), int(v.numel()), word, stream) - else: - fill_strided(v.data_ptr(), sh, st, v.element_size(), word, stream) + # Seeding is the one thing here that WRITES, so it is planned in full + # before any of it is issued: a second reduction output that turns + # out to be unseedable must not find the first already filled, and a + # 32-bit word into a narrower element would run off the end of the + # caller's allocation before the launch could reject the dtype. + if seeds: + fills = [] + for idx, word, elem_bytes in seeds: + v = operands[idx] + if v.element_size() != elem_bytes: + gave_up["reduction seed dtype"] += 1 + return refuse(operands, graph_order) + sh, st = v.shape, v.stride() + if is_contiguous(sh, st): + fills.append((v.data_ptr(), None, int(v.numel()), word)) + continue + plan = fill_plan(sh, st) + if plan is None: + gave_up["reduction output writes an element twice"] += 1 + return refuse(operands, graph_order) + fills.append((v.data_ptr(), plan, 0, word)) + for ptr, plan, count, word in fills: + if plan is None: + fill_word(ptr, count, word, stream) + else: + apply_fill(ptr, plan, word, stream) return launchable( tuple(problem), - *(ops[i].permute(1, 2, 0) if ref is None else _reshape_aux_to_fake(ops[i], ref) for i, ref in args), + *(v.permute(1, 2, 0) for v in vs), + *(operands[i].permute(1, 2, 0) if ref is None else _reshape_aux_to_fake(operands[i], ref) for i, ref in tail), stream=_as_custream(stream), ) @@ -2527,11 +2558,12 @@ def _check_executable(chain: FusionChain) -> None: if chain.has_moe: return if not _TVM_FFI_OK: + # Not a narrowing of the install surface: apache-tvm-ffi ships in the + # same `cutedsl` extra as the DSL these kernels are written in, so a + # build without it has no DSL either and was already declining. raise NotImplementedError("the tvm-ffi front door is not installed, and the launch path this engine has needs it (pip install apache-tvm-ffi)") if any(red.mode == "norm2" for red in chain.reductions): raise NotImplementedError("a norm2 reduction takes a square root after the kernel, which is a device operation this engine does not own") - if chain.is_multi_gemm and not chain.output_specs: - raise NotImplementedError("a multi-GEMM with no dense output has no tensor to read the batch extent off") def probe_supported( diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py index b570e3b61..ab997eca5 100644 --- a/python/cudnn/gemm/frost/recipe.py +++ b/python/cudnn/gemm/frost/recipe.py @@ -33,7 +33,7 @@ from dataclasses import dataclass from typing import Any -from cudnn.frost.buffers import init_word +from cudnn.frost.buffers import init_word, is_contiguous, strided_fill_plan from .dtypes import DTYPE_BYTES, _aux_align_reqs, _output_align_reqs, _pow2_floor, tensor_alignment from .fusion_ir import FusionChain @@ -207,9 +207,15 @@ class GemmRecipe: # six launchers. arg_plan: tuple stride_ins: tuple - # ``(leader, followers)`` groups whose strides the launch collapses to one. + # ``(leader, followers)`` groups whose strides the launch collapses to one, + # as POSITIONS in ``inputs``. Positions and not operand indices, because one + # buffer can occupy two roles -- ``matmul(A, A)`` binds a single slot as both + # operands -- and a role is what carries an axis order. shared_layout: tuple - # ``(output index, identity as a dtype-packed word)`` per reduction output. + # ``(output index, identity as a dtype-packed word, the dtype's byte width)`` + # per reduction output. The width is checked before the seed is written: the + # word is 32 bits and the count is the buffer's numel, so a narrower element + # would put the fill past the end of the caller's allocation. seeds: tuple @property @@ -230,8 +236,16 @@ def problem(self, operands, graph_order=None) -> tuple: it was, or None when every operand is the caller's own. """ axes, bad = [], [] + rank = [] for op in self.inputs: v = operands[op.index] + if len(v.shape) != 3: + # Everything below indexes three named axes, so a buffer with a + # different rank has to be answered here and not by an + # IndexError three frames down. + rank.append(f"{op.role}: expected a rank-3 buffer, got shape={tuple(v.shape)}") + axes.append(KERNEL_AXES) + continue borrowed = bool(graph_order and graph_order[op.index]) ax = op.axes(v.shape, v.stride(), borrowed) if ax is None: @@ -241,6 +255,8 @@ def problem(self, operands, graph_order=None) -> tuple: ) ax = KERNEL_AXES axes.append(ax) + if rank: + raise ValueError("the kernel reads three axes off every operand: " + "; ".join(rank)) if bad: raise ValueError("runtime operand layout does not match the layout the kernel was compiled for: " + "; ".join(bad)) a, b = self.a, self.b @@ -346,6 +362,41 @@ def _sf_blob_reject(recipe: GemmRecipe, operands, axes, mnk) -> "str | None": return None +def _shared_layout_reject(recipe: GemmRecipe, operands) -> "str | None": + """Block-scale multi-GEMM sends ONE A stride triple for every A operand, so + operands that share it must actually be laid out alike.""" + bad = [] + for lead, followers in recipe.shared_layout: + want = tuple(operands[recipe.inputs[lead].index].stride()) + for j in followers: + got = tuple(operands[recipe.inputs[j].index].stride()) + if got != want: + bad.append(f"{recipe.inputs[j].role} has stride {got} where {recipe.inputs[lead].role} has {want}") + if bad: + return "this kernel sends one stride triple per operand pool, so the operands in a pool must share a layout: " + "; ".join(bad) + return None + + +def _seed_reject(recipe: GemmRecipe, operands) -> "str | None": + """A reduction output is seeded with its identity before the kernel runs, so + the seed's own preconditions are the call's -- and they are checked before + the first byte is written, because a seed that fails halfway has already + scribbled on a caller's buffer.""" + bad = [] + for index, _word, elem_bytes in recipe.seeds: + v = operands[index] + got = v.element_size() + if got != elem_bytes: + bad.append(f"{recipe.roles[index]}: the reduction accumulates in {elem_bytes}-byte elements but this buffer stores {got}-byte ones") + continue + shape, stride = tuple(v.shape), tuple(v.stride()) + if not is_contiguous(shape, stride) and strided_fill_plan(shape, stride) is None: + bad.append(f"{recipe.roles[index]}: shape {shape} stride {stride} writes some element twice") + if bad: + return "a reduction output must be seedable: " + "; ".join(bad) + return None + + def _raise_first(reasons) -> None: for reason in reasons: if reason is not None: @@ -361,7 +412,12 @@ def check_shapes(recipe: GemmRecipe, operands, mnk, axes) -> None: the build recorded, and a block-scale blob against the size the template re-synthesizes. """ - reasons = [_shape_reject(recipe, operands, axes, mnk), _output_shape_reject(recipe, operands, mnk)] + reasons = [ + _shape_reject(recipe, operands, axes, mnk), + _output_shape_reject(recipe, operands, mnk), + _shared_layout_reject(recipe, operands), + _seed_reject(recipe, operands), + ] if recipe.block_size: reasons.append(_sf_blob_reject(recipe, operands, axes, mnk)) _raise_first(reasons) @@ -435,7 +491,7 @@ def build(compiled) -> GemmRecipe: if spec.is_reduction: red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] init = REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] - seeds.append((order[id(t)], init_word(red.compute_dtype, init))) + seeds.append((order[id(t)], init_word(red.compute_dtype, init), DTYPE_BYTES[red.compute_dtype])) outputs.append( Output( index=order[id(t)], @@ -477,7 +533,8 @@ def build(compiled) -> GemmRecipe: stride_ins = (0, na) if grouped else tuple(range(len(inputs))) shared_layout = () if grouped: - shared_layout = tuple((group[0].index, tuple(op.index for op in group[1:])) for group in (inputs[:na], inputs[na:]) if len(group) > 1) + pools = (tuple(range(na)), tuple(range(na, len(inputs)))) + shared_layout = tuple((pool[0], pool[1:]) for pool in pools if len(pool) > 1) return GemmRecipe( inputs=tuple(inputs), diff --git a/test/python/gemm/frost/test_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index 2f7701dc6..4d0d4de0e 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -11,8 +11,13 @@ What replaces it is the BACKEND -- the two tests below that run a shape on a cuDNN plan and a FROST plan and require the same numbers -- and the invariant -that makes a rejection meaningful: the set of calls the launch path refuses must -equal the set of illegal calls, asserted per flavor through ``deferrals``. +that makes a rejection meaningful: the set of calls the launch path refuses +should equal the set of illegal calls. + +Only one direction of that is cheap to assert. ``deferrals`` must be empty for a +legal call of every flavor, which says nothing legal is refused; a call the fast +path ACCEPTS never reaches the counter, so the other direction is covered +case by case below and by the backend differentials. """ from __future__ import annotations @@ -27,7 +32,6 @@ import cudnn import cudnn.gemm.frost # noqa: F401 — installs the cudnn.pygraph recorder hook from cudnn.engines import is_python_engine -from cudnn.frost.buffers import collapse_layout, fill_word_strided_async, init_word from cudnn.gemm.frost.compiler import jit_from_cudnn_graph, probe_supported from cudnn.gemm.frost.graph_analyzer import resolve_variant_pack from cudnn.gemm.frost.recipe import _output_rule, expected_shape @@ -77,66 +81,6 @@ def test_a_quant_scale_output_is_fixed_at_build(): assert expected_shape(rule, 999, 999) == (1, 128, 4) -# --- seeding a 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 below to a single memset for a - contiguous tap and one per batch for a padded one, rather than one per row. - """ - assert collapse_layout(shape, stride) == want - - -@requires_sm100 -@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_strided_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 kernel's. What it has to get - right is exactly this: every element the view covers, and nothing between - them. - """ - 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")) - - -@pytest.mark.parametrize( - "shape,stride,why", - [ - ((1, 8, 1), (0, 0, 1), "alias"), # stride 0 over a real extent - ((1, 4, 4), (16, 2, 1), "overlap"), # rows closer together than they are wide - ], -) -def test_a_reduction_output_that_writes_a_byte_twice_is_rejected(shape, stride, why): - """Two elements at one address is a write race, not a layout to support. - - Named here rather than left to the driver, which reports the second as a - pitch smaller than the width and says nothing about whose buffer it was. - """ - with pytest.raises(ValueError, match=why): - fill_word_strided_async(0, shape, stride, 4, 0, None) - - # --- which flavors lower, and which decline --------------------------------- @@ -157,6 +101,20 @@ def _reduction_graph(): return g +def _shared_operand_graph(d=K): + """``matmul(A, A)``: one tensor in two roles, so both bind ONE pack slot. + + At d x d the same buffer is a legal A (k-major) and a legal B (n-major), and + the graph declares it once -- which is exactly the shape where re-labelling + the SLOT into B's axis order also re-labels what A reads. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, d, d], stride=[d * d, d, 1]) + C = g.matmul(A=A, B=A, name="mm") + C.set_output(True).set_data_type(BF16) + return g + + def _row_reduction_graph(pad=4): """A per-row tap, declared into a padded buffer when ``pad`` > 1.""" g = _plain_graph() @@ -192,16 +150,24 @@ def _two_output_graph(): return g -def _multi_gemm_graph(): +def _multi_gemm_graph(dense_output=True): 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]) B0 = g.tensor(name="B0", dim=[1, K, N], stride=[K * N, 1, K]) B1 = g.tensor(name="B1", dim=[1, K, N], stride=[K * N, 1, K]) Y = g.add(a=g.matmul(A=A, B=B0, name="mm0"), b=g.matmul(A=A, B=B1, name="mm1"), name="sum") - Y.set_output(True).set_data_type(BF16) + Y.set_data_type(BF16).set_output(dense_output) + if not dense_output: + R = g.reduction(input=Y, mode=cudnn.reduction_mode.ADD, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]).set_output(True).set_data_type(F32) return g +def _multi_gemm_reduction_only_graph(): + """Two matmuls whose only output is a tap: the batch comes off the operands.""" + return _multi_gemm_graph(dense_output=False) + + # Each entry is (build, how many distinct B operands, how many outputs, how many aux). _FLAVORS = ( (_plain_graph, 1, 1, 0), @@ -210,8 +176,9 @@ def _multi_gemm_graph(): (_two_output_graph, 1, 2, 0), (_reduction_graph, 1, 2, 0), (_multi_gemm_graph, 2, 1, 0), + (_multi_gemm_reduction_only_graph, 2, 1, 0), ) -_FLAVOR_IDS = ("plain", "epilogue", "aux", "two_outputs", "reduction", "multi_gemm") +_FLAVOR_IDS = ("plain", "epilogue", "aux", "two_outputs", "reduction", "multi_gemm", "multi_gemm_tap_only") @requires_sm100 @@ -505,14 +472,14 @@ def test_the_launch_path_accepts_exactly_the_legal_calls(case, verdict): While ``lowered`` could hand anything it was unsure of to an interpreter that ran it anyway, "the fast path rejects a legal call" was a performance bug at worst and nothing asserted otherwise. There is one path now, so a - rejection IS the answer -- which makes the two directions both real - failures: refusing a legal call breaks it, and accepting an illegal one runs - a kernel over memory the caller did not describe. - - ``deferrals`` is what pins that down. Empty means every rule passed; a - rejection must name the rule that caused it, and a rejection with nothing - named is drift between the guards and the diagnostics, which ``explain`` - raises on separately. + rejection IS the answer -- which makes both directions real failures: + refusing a legal call breaks it, and accepting an illegal one runs a kernel + over memory the caller did not describe. + + ``deferrals`` pins down the first direction only, and this table is how the + second is covered: each illegal case is named and must be refused. A + rejection with nothing counted would be drift between the guards and the + diagnostics, which ``explain`` raises on separately. """ compiled = jit_from_cudnn_graph(_plain_graph()) if compiled.lowered is None: @@ -526,6 +493,73 @@ def test_the_launch_path_accepts_exactly_the_legal_calls(case, verdict): torch.testing.assert_close(out.float(), torch.einsum("bmk,bnk->bmn", a.float(), b.float()), atol=2e-1, rtol=2e-2) +@requires_sm100 +def test_one_buffer_in_two_roles_reads_each_role_s_own_axis_order(): + """``matmul(A, A)`` binds ONE slot as both operands. + + The two roles read that memory through different axis maps -- A as + ``[b, m, k]``, B as the graph's ``[b, k, n]`` -- so the launch path carries + a view per ROLE. Re-labelling the slot instead is cheaper and silently hands + the other role a transposed view: the numbers come out as ``A @ A`` where + the graph says ``A @ A.T``, and no rule fires, so the checker then finds + nothing wrong and raises the drift error instead. + """ + compiled = jit_from_cudnn_graph(_shared_operand_graph()) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") + assert len({op.index for op in compiled.recipe.inputs}) == 1 # one slot, two roles + + a = torch.randn(1, K, K, dtype=torch.bfloat16, device="cuda") + c = torch.zeros(1, K, K, dtype=torch.bfloat16, device="cuda") + compiled.lowered(_bind(compiled, [a], [a], [c]), stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {} + # The graph declares B [b, K, N], so B's N axis is the buffer's LAST -- which + # makes the product A @ A, not A @ A.T. Asymmetric by construction. + torch.testing.assert_close(c.float(), (a.float() @ a.float()), atol=2e-1, rtol=2e-2) + + +@requires_sm100 +def test_a_reduction_output_of_the_wrong_width_is_refused_before_anything_is_written(): + """The seed is a 32-bit word and the count is the buffer's numel. + + A tap bound to a narrower buffer therefore writes twice the bytes it owns, + and the launch's own dtype check comes too late to help -- by then the fill + has already run. So the width is a rule of this call, checked with the + caller's memory still untouched. The canaries are the assertion: rejecting + is not enough if it rejects afterwards. + """ + compiled = jit_from_cudnn_graph(_reduction_graph()) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") + operands = _buffers_for(compiled) + tap = [o for o in compiled.recipe.outputs if o.init is not None][0] + block = torch.full((16,), -7.0, dtype=torch.float16, device="cuda") + operands[tap.index] = block[:1].view(1, 1, 1) + + with pytest.raises(ValueError, match="element"): + compiled.lowered(operands, stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {"reduction seed dtype": 1} + assert torch.equal(block, torch.full((16,), -7.0, dtype=torch.float16, device="cuda")) + + +@requires_sm100 +def test_the_checker_is_loud_when_it_cannot_find_the_fault(): + """The one failure mode that writing the rules twice introduces. + + ``lowered``'s guards are fused for speed and ``explain``'s are written for + the message: two spellings of one set of rules. If they drift, a call gets + refused and the checker then finds nothing wrong with it -- so a call the + checker considers legal has to be loud rather than a quiet return, and + distinct from the ValueError a real rejection raises. + """ + compiled = jit_from_cudnn_graph(_plain_graph()) + a, b, c = _good() + with pytest.raises(RuntimeError, match="no rule explains"): + compiled.explain(_bound_buffers(compiled, a, b, c)) + + def _matmul_on(batch, m, n, k, want_frost, a, b, c): """Build the plain graph, pin a backend or a FROST plan, run it.""" g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) From 1667f73494822fd88df2cdcc47c7774d063b8575 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Wed, 12 Aug 2026 18:11:29 -0700 Subject: [PATCH 9/9] Answer a scale factor's rank, and stop the docs promising an interpreter Three things review found on the pushed head. A scale-factor operand had no rank check. Its blob rules -- a 16-byte base, one dense run, at least the size the template re-synthesizes from M/N/K -- all pass for a flat rank-1 blob, and it then reached permute(1, 2, 0) in the launch argument list, because build() puts scale factors in `heads` alongside the inputs. So it raised from inside the body whose whole contract is that it does not raise, and the checker could not name a cause. Both sides answer it now, which makes it an ordinary rejection with a reason. Two documents still described the interpreter this branch deletes. docs/python_graph_and_execution_backends.md and python/cudnn/frost/README.md both said the lowered path hands what it is unsure of to an interpreting path that "serves every flavor and owns every rejection message". They now say what is actually there: one launch, a checker that names a rule and raises without running anything, and a graph the closure cannot serve declined when the engine is asked to support it. Both also say why a reference executor kept for diagnostics is not the cheaper option -- it is a second answer to what the graph computes, and a differential between two readings of one plan cannot catch a misconception they share. The docstring edit in the three linear-attention kernels left a 103-character line. It is a paragraph break now, which is what the rest of those docstrings do. Co-Authored-By: Claude Opus 5 (1M context) --- docs/python_graph_and_execution_backends.md | 20 ++++++++++------- python/cudnn/frost/README.md | 15 ++++++++----- python/cudnn/gemm/frost/compiler.py | 6 ++++- python/cudnn/gemm/frost/recipe.py | 6 +++++ .../frost/kernel/gdn2_prefill_f16.py | 4 +++- .../frost/kernel/gdn_bprop_f16.py | 4 +++- .../frost/kernel/kda_prefill_f16.py | 4 +++- test/python/gemm/frost/test_execute_recipe.py | 22 +++++++++++++++++++ 8 files changed, 63 insertions(+), 18 deletions(-) diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index b2fa56632..9a595cfd1 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -178,14 +178,18 @@ and the operand structure flattened into the loop headers, so the call does no attribute lookup and takes no branch the build already settled. That is 35 → 20. Two rules make it safe: -- **The lowered path never raises.** Anything it is not certain of it hands to - the interpreting path, which serves every flavor and owns every rejection - message. It can then only ever accept a subset, and there is no second set of - error strings to drift. -- **Both read the same table.** A differential test between them catches - divergence, but never a misconception they share — so the table is where a - fact lives exactly once, and the tests that matter are against intended - semantics, at the shapes where two encodings coincide. +- **The lowered path never raises, and it is the only path that runs.** What it + refuses it hands to a checker that reads the same table, names the rule and + raises — it launches nothing. A graph the closure cannot serve at all is + declined when the engine is asked to support it, so it goes to the backend + rather than to a second executor. +- **So a refusal is the answer, not a slower route.** The set of calls the + closure refuses should equal the set of illegal calls; a legal call it will + not serve is a bug. Keeping a reference executor instead would buy a + differential that catches divergence but never a misconception the two share — + which is exactly how an axis-order bug survived one here. The tests that + matter are against intended semantics and against the BACKEND, at the shapes + where two encodings coincide. **A loop over a flat table gets almost all of it, so do not hand-unroll per flavor.** Measured three ways on the same plan and buffers: interpreting the diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index 059ea95a2..eb0215c51 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -240,13 +240,16 @@ CompiledPlan.execute(graph, uid_to_data, ctx) (the hot path) which outputs are reductions -- all settled when the kernel compiled, and deciding them again per call is most of what a python execute path costs (measured: 40-50 -> 20-22 us for one gemm, across six flavors). - `gemm/frost/recipe.py` is the worked example: one table, read by an - interpreter that walks the operand structure and by a closure that captures - the table and loops over it flat. Even what the kernel's parameter list looks + `gemm/frost/recipe.py` is the worked example: one table, captured into a + closure that loops over it flat. Even what the kernel's parameter list looks like is a table entry (`arg_plan`), which is why one loop serves every flavor - -- a call path per flavor is how two of them disagree. The lowered one never - raises: anything it is unsure of it hands back to the interpreter, which owns - every rejection message. + -- a call path per flavor is how two of them disagree. That closure never + raises and it is the ONLY thing that launches: what it refuses goes to a + checker that reads the same table, names the rule and raises without running + anything, and a graph it cannot serve at all is declined at `check_support`. + A second executor kept for diagnostics is still a second answer to what the + graph computes, and a differential between two readings of one plan cannot + catch a misconception they share. - **`ExecutionContext` carries handle, stream and workspace explicitly.** No engine may hard-code a stream, reach into private graph state, or allocate hidden workspace. `uid_to_data` is the caller's variant pack (tensor uid -> diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 2969bfc9a..49e720979 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -2065,7 +2065,11 @@ def lowered(operands, graph_order=None, stream=None): v = operands[idx] count = int(v.numel()) if ( - _pow2_floor(v.data_ptr()) < 16 + # A scale factor rides the launch as a rank-3 permute + # like every other head, so its rank is a rule here and + # not something for the permute to discover. + len(v.shape) != 3 + or _pow2_floor(v.data_ptr()) < 16 or count != 1 + sum((int(s) - 1) * int(st) for s, st in zip(v.shape, v.stride())) or count * v.element_size() < 512 * k4 * (((m if is_a else n) + 127) // 128) * sf_batch ): diff --git a/python/cudnn/gemm/frost/recipe.py b/python/cudnn/gemm/frost/recipe.py index ab997eca5..be960fce9 100644 --- a/python/cudnn/gemm/frost/recipe.py +++ b/python/cudnn/gemm/frost/recipe.py @@ -344,6 +344,12 @@ def _sf_blob_reject(recipe: GemmRecipe, operands, axes, mnk) -> "str | None": bad = [] for sf in recipe.sf: v = operands[sf.index] + if len(v.shape) != 3: + # It reaches the kernel through the same rank-3 relabelling as every + # other head, so a different rank is answered here rather than by + # the permute three frames down. + bad.append(f"{sf.role}: expected a rank-3 buffer, got shape={tuple(v.shape)}") + continue op = recipe.inputs[sf.operand_at] rows = m if sf.is_a else n batch = operands[op.index].shape[axes[sf.operand_at][AX_BATCH]] 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 ff80aad30..4203786d8 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,9 @@ 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 f9fd8c670..0b36946fb 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,9 @@ 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 dfd622955..71263b08b 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,9 @@ 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_execute_recipe.py b/test/python/gemm/frost/test_execute_recipe.py index 4d0d4de0e..d5dc59b24 100644 --- a/test/python/gemm/frost/test_execute_recipe.py +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -649,6 +649,28 @@ def test_block_scale_lowers_and_runs(gemms): assert out.abs().sum() > 0 +@requires_sm100 +def test_a_scale_factor_of_the_wrong_rank_is_refused_rather_than_crashing(): + """A scale factor is relabelled like every other head, so its rank is a rule. + + The blob checks -- alignment, one dense run, the size the template + re-synthesizes -- all pass for a flat rank-1 blob, and it then reached + ``permute(1, 2, 0)`` in the launch argument list and raised from INSIDE the + body that is not allowed to raise. Both the guard and the checker name it + now, so it is a rejection with a reason like any other. + """ + compiled = jit_from_cudnn_graph(_nvfp4_graph(1)) + if compiled.lowered is None: + pytest.skip(f"this build does not lower: {compiled.declined}") + operands, _out = _nvfp4_buffers(compiled) + sf = compiled.recipe.sf[0] + operands[sf.index] = operands[sf.index].reshape(-1) + + with pytest.raises(ValueError, match="rank-3"): + compiled.lowered(operands, stream=None) + assert dict(compiled.deferrals) == {"scale-factor blob": 1} + + @requires_sm100 @pytest.mark.parametrize("batch", (1, 2), ids=("b1", "b2")) @pytest.mark.parametrize("m,n,k", [(m, n, k) for m in (1, 128) for n in (1, 128) for k in (1, 128)], ids=str)