diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index 149735ce1..9a595cfd1 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,83 @@ 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 **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: + +| | µ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 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, 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 +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. + +**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 39a7011d1..22fc71c09 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -1882,12 +1882,21 @@ 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)) + 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: @@ -1908,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. @@ -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..e3d5dd46a 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -133,32 +133,37 @@ class VariantPack: pointers — silently, because every pointer in it is individually valid. """ - __slots__ = ("uids", "native", "_slot_of", "workspace", "workspace_bytes", "_device") + __slots__ = ("uids", "native", "_index_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 - self._slot_of = None # built on first lookup: the backend never does one + # 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._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: @@ -173,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. @@ -198,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.""" @@ -214,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 e76a14c7f..eb0215c51 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -235,10 +235,31 @@ 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: 40-50 -> 20-22 us for one gemm, across six flavors). + `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. 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 -> 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 011322cf8..49e720979 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 @@ -19,11 +20,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 +55,18 @@ def _as_custream(stream): ) from .epilogue_codegen import EpilogueSnippets, generate from .fusion_ir import ZERO_PRESERVING_OPS, FusionChain, TensorRef +from .recipe import ( + CONST, + FROM_M, + KERNEL_AXES, + REDUCTION_INIT_VALUE, + _output_rule, + build as build_recipe, + check_alignment, + check_shapes, + contiguous_modulus, + expected_shape, +) from .graph_analyzer import ( GemmBinding, analyze_with_binding, @@ -711,8 +725,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 +738,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( @@ -1751,11 +1762,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): @@ -1765,42 +1782,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 +1797,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 @@ -1846,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 @@ -1873,6 +1857,26 @@ 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 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 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 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: + if self.binding is not None: + self.bound = tuple(self.binding.bound_tensors()) + self.recipe = build_recipe(self) + self.lowered = self._lower() def __call__(self, variant_pack, stream=None): # The runtime call is a variant-pack dict keyed by cuDNN tensor object @@ -1886,324 +1890,258 @@ 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. - - 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. + """Launch over ``{id(bound_tensor): buffer}``, already resolved.""" + 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]}") + operands.append(buf) + 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: this kernel's only call path. + + 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 + 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 body never raises. What it refuses it hands to ``explain``, which + 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. """ - _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) - _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) - - if self.chain.is_multi_gemm: - 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 + def decline(reason: str) -> None: + self.declined = reason + return None - 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) - _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." - ) + r = self.recipe + if r is None: + return decline("no binding") + # 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") + # 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") + + # 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. + # 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( + (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 + ) - 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]}" + # 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 + 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) + # 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 + 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_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 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. + # + # 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[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: + 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) + 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 (kc, mod, ebatch, kpack, is_b, takes_stride), v in zip(ins, vs): + 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 + ): + gave_up["input layout"] += 1 + 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] + 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 + ): + gave_up["output layout"] += 1 + return refuse(operands, graph_order) + 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 refuse(operands, graph_order) + if sfs: + k4 = ((k // block_size) + 3) // 4 + for idx, is_a, sf_batch in sfs: + v = operands[idx] + count = int(v.numel()) + if ( + # 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 + ): + gave_up["scale-factor blob"] += 1 + return refuse(operands, graph_order) + for lead, followers in shared: + st = tuple(vs[lead].stride()) + for j in followers: + if tuple(vs[j].stride()) != st: + gave_up["operands do not share a layout"] += 1 + return refuse(operands, graph_order) + # 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), + *(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), ) - _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, - ) - 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") + return lowered - 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 explain(self, operands, graph_order=None): + """Say what is wrong with a call the launch path refused, and raise. - 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)``. + 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. - 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") + 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. - # 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, + 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) + 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_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: @@ -2362,76 +2300,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, @@ -2670,6 +2547,29 @@ 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: + # 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") + + def probe_supported( graph: cudnn.pygraph, config: TileConfig = DEFAULT_CONFIG, @@ -2685,6 +2585,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/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..cda0e63c5 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -31,33 +31,49 @@ 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 + # 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 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() - # 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, operands)), Workspace.over(variant_pack, required, "frost_gemm"), stream=ctx.stream) + return + launch = self._launch + if launch 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. + # 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. + graph_order = None + borrowed = variant_pack.graph_described + if borrowed: + flags = tuple(i in borrowed for i in indices) + graph_order = flags if any(flags) else None + launch(operands, graph_order, stream=ctx.stream) else: - self._compiled(dict(zip(self._tensors, views)), *extra, 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 new file mode 100644 index 000000000..be960fce9 --- /dev/null +++ b/python/cudnn/gemm/frost/recipe.py @@ -0,0 +1,563 @@ +# 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, 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 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 +only which buffers the launch passes and in what order, so that is data. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +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 + +# 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]; 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) + +# 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.""" + + 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 + 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_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, 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) +class Output: + index: 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 + + +@dataclass(frozen=True) +class Aux: + index: int + role: str + align: int + ref: Any # TensorRef, for the fake-shape reshape + + +@dataclass(frozen=True) +class ScaleFactor: + index: 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 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 + 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, + # 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, 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 + def a(self) -> Operand: + return self.inputs[self.a_at] + + @property + def b(self) -> Operand: + return self.inputs[self.b_at] + + 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-operand answer to which + 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: + want = op.dc if borrowed else op.kc + bad.append( + 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) + 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 + a_ax, b_ax = axes[self.a_at], axes[self.b_at] + 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, 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 = 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. + 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, 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 = 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: + 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, operands, mnk) -> "str | None": + m, n, _ = mnk + bad = [] + for out in recipe.outputs: + 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}") + if bad: + return "runtime tensors must be rank-3 with shapes matching the graph: " + "; ".join(bad) + return 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 + 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(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 = 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: + 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, 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.""" + m, n, k = mnk + k4 = ((k // recipe.block_size) + 3) // 4 + 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]] + 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 _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: + 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), + _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) + + +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. + + 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()), tuple(int(s) for s in tensor.get_stride()) + except Exception: # noqa: BLE001 -- an analyzer-synthesized ref has no dims + return (), () + + +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( + index=index, + 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_layout=_declared_layout(tensor), + declared=declared, + 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, seeds = [], [] + for i, (spec, t) in enumerate(zip(chain.outputs, binding.outputs)): + init = None + 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), DTYPE_BYTES[red.compute_dtype])) + outputs.append( + Output( + index=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, + ) + ) + + aux = tuple( + 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(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.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: + 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), + 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), + 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/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/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/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 0c25db2b7..8b9b0b351 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,81 @@ 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) + // 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. - 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 +813,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 new file mode 100644 index 000000000..d5dc59b24 --- /dev/null +++ b/test/python/gemm/frost/test_execute_recipe.py @@ -0,0 +1,806 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""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 +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 + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch +from gemm_test_utils import kw, requires_sm100, to_blocked + +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, probe_supported +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 _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() + 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]) + 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 + + +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(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_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), + (_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), + (_multi_gemm_reduction_only_graph, 2, 1, 0), +) +_FLAVOR_IDS = ("plain", "epilogue", "aux", "two_outputs", "reduction", "multi_gemm", "multi_gemm_tap_only") + + +@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 + 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 +@pytest.mark.parametrize("build,n_b,n_out,n_aux", _FLAVORS, ids=_FLAVOR_IDS) +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. + + 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``.""" + 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(): + """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 = _wrong_major() + operands = _bound_buffers(compiled, a, b, c) + with pytest.raises(ValueError): + compiled.lowered(operands, stream=None) + assert dict(compiled.deferrals) == {"input layout": 1} + + +@requires_sm100 +@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. + + 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. + """ + 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 +def test_public_execute_takes_the_lowered_path(): + """The lowered path 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 _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.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()] + + +def _bound_buffers(compiled, a, b, c): + 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 + + +def _as_declared(compiled, operands): + """The same memory, each input re-labelled the way the graph declares it. + + 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): + c.zero_() + try: + run(operands, 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,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_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 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: + pytest.skip(f"this build does not lower: {compiled.declined}") + + a, b, c = case() + 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) + + +@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) + 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 + + +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_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 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 + # 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) + compiled.lowered(operands, stream=None) + torch.cuda.synchronize() + assert dict(compiled.deferrals) == {} + 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) +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. + + 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() + 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 +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() + 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 068e159c4..084f81eff 100644 --- a/test/python/gemm/frost/test_public_execute_flavors.py +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -97,6 +97,105 @@ 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 +@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.""" + 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", @@ -166,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. @@ -190,18 +342,16 @@ 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 -- + 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]) B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) @@ -213,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