diff --git a/dlpack_version.txt b/dlpack_version.txt index 9459d4ba2..7e32cd569 100644 --- a/dlpack_version.txt +++ b/dlpack_version.txt @@ -1 +1 @@ -1.1 +1.3 diff --git a/docs/adding_torch_custom_ops.md b/docs/adding_torch_custom_ops.md index d3f90f2d5..207ea059f 100644 --- a/docs/adding_torch_custom_ops.md +++ b/docs/adding_torch_custom_ops.md @@ -77,25 +77,25 @@ graph.check_support() graph.build_plans() # also prepares variant pack template ``` -## Execute: use sorted_ptrs path +## Execute -Cache `uid_order` alongside the graph. Use `_execute_with_ptrs` instead of dict-based execute: +`graph.execute(uid_to_tensor, workspace, handle=handle)` is the only form you need. +It already caches the backend's operand order and reuses one pointer array, so the +sorted-pointer path is what runs underneath — there is nothing faster to reach for, +and hand-rolling it costs you the dynamic-shape overrides and the python-engine +dispatch that `execute()` handles. ```python if cache_key not in _cache: - graph, ws_size = _build_graph(...) - uid_order = graph._get_variant_pack_uids_sorted() - _cache[cache_key] = (graph, ws_size, uid_order) + _cache[cache_key] = _build_graph(...) # (graph, ws_size) -graph, ws_size, uid_order = _cache[cache_key] +graph, ws_size = _cache[cache_key] # Allocate workspace per-call (PyTorch's caching allocator recycles efficiently) workspace = torch.empty(max(ws_size, 1), device=x.device, dtype=torch.uint8) -# Build uid→tensor map, extract sorted ptrs uid_to_tensor = {X.get_uid(): x, W.get_uid(): w, Y.get_uid(): y_out} -ptrs = [uid_to_tensor[uid].data_ptr() for uid in uid_order] -graph._execute_with_ptrs(ptrs, workspace.data_ptr(), int(handle)) +graph.execute(uid_to_tensor, workspace, handle=handle) ``` **Do NOT cache workspace tensors** — they can race on different CUDA streams. @@ -175,7 +175,7 @@ def my_op(x, w, eps=1e-5, bias=None): - [ ] Use `torch.Library.define/impl`, not `@torch.library.custom_op` - [ ] Cache graph + uid_order + workspace in module-level dict -- [ ] Use `graph._execute_with_ptrs(sorted_ptrs)` not `graph.execute(dict)` +- [ ] Use `graph.execute(uid_to_tensor, workspace, handle=handle)` — it takes the sorted-pointer path internally - [ ] Use explicit UIDs (IntEnum) for stable cache keys - [ ] Cache cuDNN handle per device - [ ] Allocate workspace per-call (do NOT cache — stream safety) @@ -192,7 +192,7 @@ def my_op(x, w, eps=1e-5, bias=None): | `torch.empty` per output tensor | ~2 each | consider `out=` or pre-alloc | | cache key build + lookup | ~1.5 | tuple construction + dict hash | | uid→tensor dict + list comp | ~1 | Python overhead | -| `graph._execute_with_ptrs` | ~19 | 1.7 us FE + 0.8 us varpack + 5.6 us backend | +| `graph.execute` | ~19 | 1.7 us FE + 0.8 us varpack + 5.6 us backend | | **Total (well-optimized)** | **~52** | vs ~18 us for native ATen ops | The ~34 us gap vs native ATen is torch.ops dispatcher + autograd overhead. diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index 2b3f6f17c..149735ce1 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -80,15 +80,65 @@ create_execution_plans([heur_mode.A, ...]) _pygraph.py - `BaseEngine`: `check_support(graph)` (accept, or decline by raising), `build_plan(graph, plan, ctx) → CompiledPlan` (the expensive JIT step, once per graph/plan, cached on the graph), - `CompiledPlan.execute(graph, uid_to_data, ExecutionContext)` with explicit - handle/stream/workspace/overrides. `uid_to_data` is the caller's variant - pack, exactly as the classic backend receives it; engines that address - buffers by port name call `resolve_node_buffers(graph, uid_to_data)` - (`engines/base.py`), which joins the pack with each node's wired ports — - strict missing-buffer validation, torch tensors detached once (DLPack/CAI - refuse `requires_grad` export) — into per-node `NodeBuffers` - (`{port_name: caller buffer}`). Simple eager engines implement - `execute()` only. + `CompiledPlan.execute(graph, operands, ExecutionContext)` with explicit + handle/stream/workspace. Dynamic-shape overrides are a backend-path feature: + a python plan is compiled for the shapes the graph declared, so `execute()` + refuses them rather than silently running a different problem. Simple eager + engines implement `execute()` only. + +#### The variant pack is normalized once + +`graph.execute()` converts whatever the caller passed — a torch tensor, a +`DeviceView`, any `__dlpack__` / `__cuda_array_interface__` producer, or a bare +device address — into a `VariantPack` at the top, and everything below reads that. +**Do not add a branch on what the caller's object is.** This exists because +there used to be two such branches: the backend path accepted a bare address +(`_native_var_pack`'s `if type(d) is int: return d`) while an engine got the +object untouched and `frost.buffers.probe` refused it. One public call, two +answers, decided by which plan the heuristics happened to pick — which the +caller does not control. + +`VariantPack` carries the caller-filled uids ascending and the operands +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 +`_execute_with_raw_ptrs` takes. + +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 +size runs another bit-exactly. **Read the IR port for the shape the plan was +built for; read the pack for the shape about to run.** + +Two rules that are easy to break by accident: + +- **The pointer array is per call.** Two threads may execute one graph + concurrently with different buffers. A shared array hands each thread the + other's pointers, and each pointer in it is individually valid, so the failure + is a wrong number rather than a raise. +- **The operand order has exactly one source, never a union.** The lowered + graph's variant-pack template when the graph has one — only C++ sees every + user slot, since a tensor's `ragged_offset` is an operand but hangs off the + `Tensor` rather than off a node port, and the slots the graph fills itself + (pass-by-value scalars, slice replacement destinations, workspace + modifications) must be excluded. The python IR only for the python-only ops + that never lower. The two sides do not have to agree: each indexes the layout + it was handed. + +`is_virtual` on a python-only graph does not mean "the caller supplies +nothing" — it is a statement about the backend's lowering, and a gdn graph marks +its own `O` virtual while the caller passes a buffer for it. So the layout there +is every wired port, and an unfilled slot is an optional port the caller did not +request. + +`CompiledPlan.takes_variant_pack` is the migration flag. An engine that has not +set it still receives the caller's `{uid: buffer}` map and reaches ports through +`resolve_node_buffers`; the flag and that function both go once the last engine +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. - **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/CMakeLists.txt b/python/CMakeLists.txt index 36c0c2e48..997dec23c 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -13,11 +13,16 @@ set(BUILD_MOCK OFF) option(CUDNN_FRONTEND_USE_SYSTEM_DLPACK "Whether dlpack should use the system version or fetched." OFF) IF(CUDNN_FRONTEND_USE_SYSTEM_DLPACK) - find_package(dlpack REQUIRED) + # 1.3 is where DLPackExchangeAPI is declared, which pygraph/variant_pack.cpp + # needs. The wire structs are byte-identical to 1.1 -- sizeof and every + # offsetof of DLTensor and DLManagedTensor match -- so this is a + # compile-time requirement only; capsules exchanged with a consumer built + # against 1.1 are unaffected. + find_package(dlpack 1.3 REQUIRED) if(dlpack_FOUND) - message(STATUS "Found system dlpack") + message(STATUS "Found system dlpack ${dlpack_VERSION}") else() - message(FATAL_ERROR "dlpack not found") + message(FATAL_ERROR "dlpack >= 1.3 not found (needed for DLPackExchangeAPI); unset CUDNN_FRONTEND_USE_SYSTEM_DLPACK to fetch it") endif() else() @@ -68,6 +73,7 @@ python_add_library( pygraph/norm.cpp pygraph/sdpa.cpp pygraph/pointwise.cpp + pygraph/variant_pack.cpp WITH_SOABI ) diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 9b4fee01d..391cb4aea 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -120,93 +120,6 @@ def _set_data_type( _pybind_module.backend_graph.tensor = _tensor -def _library_device_pointer(input_tensor): - # either pass in pointers directly - if type(input_tensor) is int: - return input_tensor - # directly extract data pointer for torch tensors - elif _is_torch_tensor(input_tensor): - return input_tensor.data_ptr() - # fall back to dlpack support by library - else: - return _pybind_module._get_data_ptr(input_tensor) - - -def _execute( - self, - tensor_to_device_buffer, - workspace, - handle=None, - override_uids=None, - override_shapes=None, - override_strides=None, -): - """ - Execute a cudnn graph. - - Args: - tensor_to_device_buffer (dict(cudnn_tensor, Union[torch.Tensor, int, __dlpack__])): The dimensions of the tensor. - workspace (Union[torch.Tensor, int, __dlpack__]): The name of the tensor. - handle: cudnn_handle created with cudnn.create_handle() - Returns: - None - """ - uid_to_tensor_pointer = { - x if type(x) is int else x.get_uid(): _library_device_pointer(pointer) for x, pointer in tensor_to_device_buffer.items() if x is not None - } - - workspace_pointer = _library_device_pointer(workspace) - self._execute( - uid_to_tensor_pointer, - workspace_pointer, - handle, - override_uids, - override_shapes, - override_strides, - ) - - -def _execute_plan_at_index( - self, - tensor_to_device_buffer, - workspace, - index, - handle=None, - override_uids=None, - override_shapes=None, - override_strides=None, -): - """ - Execute a cudnn graph. - - Args: - tensor_to_device_buffer (dict(cudnn_tensor, Union[torch.Tensor, int, __dlpack__])): The dimensions of the tensor. - workspace (Union[torch.Tensor, int, __dlpack__]): The name of the tensor. - index(int): Location of execution plan to use. - handle: cudnn_handle created with cudnn.create_handle() - Returns: - None - """ - uid_to_tensor_pointer = { - x if type(x) is int else x.get_uid(): _library_device_pointer(pointer) for x, pointer in tensor_to_device_buffer.items() if x is not None - } - - workspace_pointer = _library_device_pointer(workspace) - self._execute_plan_at_index( - uid_to_tensor_pointer, - workspace_pointer, - index, - handle, - override_uids, - override_shapes, - override_strides, - ) - - -_pybind_module.backend_graph.execute = _execute -_pybind_module.backend_graph.execute_plan_at_index = _execute_plan_at_index - - def load_cudnn(): # First look at python site packages lib_path = glob.glob(os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn/bin/cudnn64_9.dll")) diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 088874bcc..25d901b20 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -19,17 +19,52 @@ >>> graph.execute({C: c_tensor}) # routes to a supporting engine, else cuDNN """ +import ctypes from dataclasses import dataclass import logging import weakref from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -from .graph_types import NodeType, Tensor +from cudnn import _pybind_module + +from .datatypes import _buffer_dtype_to_cudnn, _dlpack_code_bits, _torch_to_cudnn_data_type +from .engines.base import ExecutionContext, VariantPack +from .engines.engine_ids import is_python_engine +from .graph_types import NodeType, Tensor, byte_size as _byte_size, describing_tensor from .nodes import Node, _row_major_stride _LOG = logging.getLogger("cudnn.pygraph") +def _is_dense(dim, stride) -> bool: + """Row-major compact.""" + expect = 1 + for extent, step in zip(reversed(tuple(dim)), reversed(tuple(stride))): + if extent != 1 and step != expect: + return False + expect *= extent + return True + + +def _in_axis_order_of(shape, stride, reference_stride): + """``(shape, stride)`` re-expressed in the axis order ``reference_stride`` uses. + + ``override_shapes`` speaks the GRAPH's declaration (a matmul's B is + ``[batch, K, N]``); the slot holds what the caller's buffer reports (B is + allocated ``(batch, N, K)``). Same memory, two orders — so an engine + indexing an extent by position would read the wrong one. Both orders rank + their axes the same way by stride, which gives the permutation. + """ + if len(shape) != len(stride) or len(stride) != len(reference_stride): + return tuple(shape), tuple(stride) + by_stride = sorted(range(len(stride)), key=lambda i: -stride[i]) + reference = sorted(range(len(reference_stride)), key=lambda i: -reference_stride[i]) + permutation = [0] * len(stride) + for rank, axis in enumerate(by_stride): + permutation[reference[rank]] = axis + return tuple(shape[a] for a in permutation), tuple(stride[a] for a in permutation) + + def cudnn_graph_not_supported(message: str) -> Exception: """The classic unsupported-graph error (built lazily: importing cudnn at module scope here would be circular).""" @@ -50,11 +85,6 @@ class GraphContext: intermediate_data_type: Any = None compute_data_type: Any = None - def __setattr__(self, name, value): - if getattr(self, "_frozen", False) and name != "_frozen": - raise RuntimeError("the graph is frozen after lowering/planning — build a new graph to change its configuration") - object.__setattr__(self, name, value) - class pygraph: """Pure Python graph representation. @@ -150,6 +180,10 @@ def __init__( self._workspace_limit: Optional[int] = None # deselect_workspace_greater_than() self._note_filters: List[Any] = [] # (kind, note, keep) from the classic note filters self._plan_pinned: bool = False # select_plan() => the walk is strict + # Backend operand order + the reusable pointer array handed to execute. + # Both are properties of the frozen graph, so they outlive any one call. + self._sorted_uids: Optional[List[int]] = None + self._selected_engine_cache = None # (plan config, engine); see selected_engine # ========================================================================= # Routing @@ -162,8 +196,6 @@ def plans(self) -> List[Any]: @property def _selected_plan_config(self) -> Optional[Any]: - from .engines.engine_ids import is_python_engine - if not self._plans or not 0 <= self._plan_index < len(self._plans): return None cfg = self._plans[self._plan_index] @@ -171,8 +203,6 @@ def _selected_plan_config(self) -> Optional[Any]: def _engine_for(self, cfg) -> Optional["BaseEngine"]: """The python engine that owns ``cfg``'s id, or None for a backend entry.""" - from .engines.engine_ids import is_python_engine - if cfg is None or not is_python_engine(cfg.engine_id): return None owners = self._owners_for_id(cfg.engine_id) @@ -223,8 +253,21 @@ def _barred_indices(self) -> set: @property def selected_engine(self) -> Optional["BaseEngine"]: """The python engine for the currently selected plan entry, or None for - the backend path. Populated after create_execution_plans().""" - return self._engine_for(self._selected_plan_config) + the backend path. Populated after create_execution_plans(). + + Cached, because ``execute()`` asks on every call and answering means + walking every registered engine for the one declaring this id — 2.75 us + to re-derive something that only ``select_plan`` can change. Keyed on + the config OBJECT, so replanning invalidates it without needing a hook + on every writer of ``_plan_index``. + """ + cfg = self._selected_plan_config + cached = self._selected_engine_cache + if cached is not None and cached[0] is cfg: + return cached[1] + engine = self._engine_for(cfg) + self._selected_engine_cache = (cfg, engine) + return engine # ========================================================================= # Tensor Creation @@ -360,14 +403,17 @@ def _check_mutable(self, what: str) -> None: self._is_validated = False def _freeze(self) -> None: - """Freeze the ENTIRE public graph surface (not just the fluent API). - - Called at lowering and at planning, whichever happens first. After - this, every mutation path raises: fluent setters and op builders (via - _check_mutable), attribute writes on Tensor/Node/GraphContext (their - __setattr__ guards), dict writes on node.inputs/outputs/params - (MappingProxy), and in-place list mutation of dim/stride (tuples). - The inspection surface stays fully readable for engines.""" + """Freeze the ENTIRE public graph surface. + + Called at lowering and at planning, whichever happens first. + + Frozen-ness is ONE flag, on the graph. Every mutation route the public + API offers goes through _check_mutable (the chained setters via + Tensor._guard / Node._guard, and the op builders), so the flag alone is + the guard. The structures the caller could otherwise mutate behind the + API's back are made immutable in their own right rather than watched: + node.inputs/outputs/params become MappingProxy views and dim/stride + become tuples. The inspection surface stays fully readable for engines.""" if self._frozen: return from types import MappingProxyType @@ -376,12 +422,9 @@ def _freeze(self) -> None: node.inputs = MappingProxyType(dict(node.inputs)) node.outputs = MappingProxyType(dict(node.outputs)) node.params = MappingProxyType(dict(node.params)) - node._frozen = True for t in self._tensor_by_uid.values(): t.dim = tuple(t.dim) if t.dim else t.dim t.stride = tuple(t.stride) if t.stride else t.stride - t._frozen = True - self._context._frozen = True self._frozen = True def _rename_tensor(self, t: Tensor, name: str) -> None: @@ -1222,8 +1265,6 @@ def _resolve_stream(self, handle: Any) -> Any: return cudnn.get_stream(handle) def _build_context(self, handle: Any = None) -> Any: - from .engines.base import ExecutionContext - h = handle if handle is not None else self._handle return ExecutionContext(handle=h, stream=self._resolve_stream(h)) @@ -1718,37 +1759,244 @@ def execute( self.build(ctx=caller_ctx) uid_to_data = self._uid_to_data(tensor_dict) + # The dynamic-shape overrides live only on the backend's uid-map + # overload, so a call carrying them takes that path and normalizes + # nothing. + overriding = override_uids is not None or override_shapes is not None or override_strides is not None eng = self.selected_engine - if eng is not None: # python engine (plan id in the reserved region) - from .engines.base import ExecutionContext + if eng is not None: # python engine (plan id in the reserved region) h = handle if handle is not None else self._handle - ctx = ExecutionContext( - handle=h, - stream=self._resolve_stream(h), - workspace=workspace, - override_uids=override_uids, - override_shapes=override_shapes, - override_strides=override_strides, - ) + ctx = ExecutionContext(handle=h, stream=self._resolve_stream(h), workspace=workspace) if self._plan_index not in self._compiled_plans: # compile with the CALLER's context (execute-supplied handle # and its stream reach the JIT build) self._compiled_plans[self._plan_index] = eng.build_plan(self, self._selected_plan_config, ctx) self._is_built = True - self._compiled_plans[self._plan_index].execute(self, uid_to_data, ctx) + plan = self._compiled_plans[self._plan_index] + # Overrides go INTO the pack rather than around it: they describe + # what this execute runs, so an engine reading the pack agrees with + # the backend without knowing they exist. + if plan.takes_variant_pack: + pack = self._normalize(uid_to_data, workspace, override_uids, override_shapes, override_strides) + plan.execute(self, pack, ctx) + else: + plan.execute(self, uid_to_data, ctx) return - # Backend path. Variant-pack keys are IR uids == the C++ uids by - # construction. Address the plan the WALK built, not the backend's own + variant_pack = None if overriding else self._normalize(uid_to_data, workspace) + + # Backend path. Address the plan the WALK built, not the backend's own # selection: they differ once the walk has skipped an entry. - var_pack, ws_ptr = self._native_var_pack(uid_to_data, workspace) cfg = self._materialize_backend_plan(self._plan_index) if self._plans else None - if cfg is not None and cfg.cpp_index is not None: - self._lowered_graph._execute_plan_at_index(var_pack, ws_ptr, cfg.cpp_index, handle, override_uids, override_shapes, override_strides) + cpp_index = cfg.cpp_index if cfg is not None else None + + # C++ turns a uid map into sorted pointers anyway (graph_interface.h, + # "uid map -> extract sorted ptrs, delegate to the sorted_ptrs + # implementation"), so handing it the sorted array directly skips one + # dict build here, one map copy in pybind, and one hash lookup per + # operand there. + if variant_pack is not None: + self._lowered_graph._execute_with_raw_ptrs( + variant_pack.address, + len(variant_pack), + variant_pack.workspace, + handle or 0, + -1 if cpp_index is None else cpp_index, + ) + return + + var_pack, ws_ptr = self._native_var_pack(uid_to_data, workspace) + if cpp_index is not None: + self._lowered_graph._execute_plan_at_index(var_pack, ws_ptr, cpp_index, handle, override_uids, override_shapes, override_strides) return self._lowered_graph._execute(var_pack, ws_ptr, handle, override_uids, override_shapes, override_strides) + def _variant_pack_uids(self) -> Optional[List[int]]: + """The graph's caller-filled variant_pack, ASCENDING by uid. + + Taken from the lowered graph whenever there is one: C++ is the only + side that can see every user slot, including the ones a walk over node + ports cannot name (a tensor's ragged_offset hangs off the Tensor, not + off a port) and correctly excluding the slots the graph fills itself + (pass-by-value scalars it already knows, slice replacement + destinations, cached workspace modifications). + + A python-only graph — gdn / kda / gdn2, which cannot lower by + construction — has no C++ side, so its variant_pack come from the IR. The + two never have to agree: each side indexes the layout it was given. + """ + order = self._sorted_uids + if order is not None: + return order + lowered = self._lowered_graph + if lowered is not None: + try: + # The order lives in the variant-pack template, which C++ builds + # lazily inside execute; the query itself does not trigger it, so + # ask explicitly or it answers with an empty list. + lowered._prepare_variant_pack_template() + order = list(lowered._get_variant_pack_uids_sorted()) + except Exception: # noqa: BLE001 — no template available yet + order = [] + else: + # Every tensor wired to a port, virtual or not. is_virtual is a + # statement about the BACKEND's lowering — an intermediate it fuses + # away — and a python-only op never lowers, so it does not mean + # "the caller supplies nothing": a gdn graph marks its own O virtual + # and the caller passes a buffer for it regardless. A slot nobody + # fills stays empty; which ports are optional is the engine's own + # business, and it already reads them with .get(). + order = sorted({t.uid for node in self._nodes for t in list(node.inputs.values()) + list(node.outputs.values()) if t is not None}) + if not order: + return None + self._sorted_uids = order + return order + + def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids=None, override_shapes=None, override_strides=None): + """Turn the caller's variant pack into :class:`VariantPack`, once. + + This is the ONLY place a caller's object is inspected. Everything below + — the backend and every python engine — reads the pointers and Tensors + built here, so the two paths cannot disagree about what the caller + passed. Returns None when the operand layout is not known yet, which + puts the caller back on the uid-map path. + + Overrides are applied here, to the slot, because they are part of the + same answer: ``override_shapes`` says the caller allocated at a cache + shape and is running a smaller one this call, so the pack must describe + the shape about to run rather than the allocation. An engine that reads + the pack then honours them without knowing the concept exists — which + is the difference between one answer and two, since the backend + re-describes the tensor from the overrides either way. + """ + order = self._variant_pack_uids() + if order is None: + return None + native = _pybind_module.VariantPackNative(len(order)) + # One crossing for the whole pack, uid lookups included. What comes back + # is the slots whose producer publishes no exchange vtable; those are + # described here without taking the rest down with them. + unread = native.read_from(uid_to_data, order) + # The backend's layout is exactly the slots it REQUIRES, so a hole is the + # 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 + for i in unread: + data = uid_to_data.get(order[i]) + if data is None: + continue # named below if this graph requires it + ptr, tensor = self._describe(data, order[i]) + native.set_slot(i, ptr, tuple(tensor.dim), tuple(tensor.stride), *_dlpack_code_bits(tensor.data_type)) + if strict: + hole = native.first_unfilled() + if hole >= 0: + uid = order[hole] + declared = self._tensor_by_uid.get(uid) + name = f" ({declared.name!r})" if declared is not None and declared.name else "" + raise ValueError(f"the variant pack is missing a buffer for tensor uid {uid}{name}") + if override_uids: + # The backend refuses a partial override; a short list must not + # quietly mean "keep the rest" here. + if len(override_shapes or ()) != len(override_uids) or len(override_strides or ()) != len(override_uids): + raise ValueError( + f"override_uids, override_shapes and override_strides must name the same tensors: got " + f"{len(override_uids)}, {len(override_shapes or ())} and {len(override_strides or ())} entries" + ) + slot_of = {uid: i for i, uid in enumerate(order)} + for j, uid in enumerate(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))) + # 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. + workspace_ptr, workspace_bytes = 0, 0 + if workspace is not None: + extent = _pybind_module.read_buffer_extent(workspace) + if extent is None: # a bare address, a non-dense buffer, or no vtable + workspace_ptr, workspace_tensor = self._describe(workspace, -1) + # An engine carves the workspace by byte offset, so a byte + # COUNT is only a byte RANGE when the buffer is dense. + if not _is_dense(workspace_tensor.dim, workspace_tensor.stride): + raise ValueError(f"the workspace buffer must be contiguous; got dim {tuple(workspace_tensor.dim)} stride {tuple(workspace_tensor.stride)}") + workspace_bytes = _byte_size(workspace_tensor) + else: + workspace_ptr, workspace_bytes = extent + return VariantPack(tuple(order), native, workspace_ptr, workspace_bytes) + + def _describe(self, data: Any, uid: int): + """``(pointer, Tensor)`` for one caller buffer. + + The Tensor carries the buffer's OWN dim/stride/data_type, which need + not match what the graph declared — frost_gemm takes its problem size + from here. + + Every framework publishes the same four facts under a different + spelling, so this asks for each spelling in turn. Two differences are + the only ones that matter, and both are handled below: strides are in + BYTES for the array-interface family and in ELEMENTS for torch/DLPack, + and an absent stride means dense row-major rather than unknown. + + Ordered by what it costs to ask (measured, 4096x4x128): reading + attributes is 0.52 us for torch and 0.72 for cupy, while the DLPack + capsule round trip is 1.5-8.6. The expensive case is specifically torch + bfloat16, whose ``__cuda_array_interface__`` raises because the + protocol cannot spell bf16 — which is why the last branch is last and + not the only one. + + A bare address carries no geometry, so it borrows the graph's: the + backend has always accepted a raw pointer, and an engine that reads the + pack for its extents would otherwise fail on an operand form the + backend takes -- one call, two answers, decided by plan selection. + """ + if type(data) is int: # bare device address + declared = self._tensor_by_uid.get(uid) + if declared is None or not declared.dim: + return data, Tensor(uid=uid) + return data, describing_tensor(uid, tuple(declared.dim), tuple(declared.stride), declared.data_type) + dim = getattr(data, "shape", None) + + # torch: pointer from data_ptr(), strides from stride() in ELEMENTS + if dim is not None and hasattr(data, "data_ptr") and callable(getattr(data, "stride", None)): + return data.data_ptr(), describing_tensor(uid, tuple(dim), tuple(data.stride()), _buffer_dtype_to_cudnn(data.dtype)) + + # cupy / numba: pointer from .data.ptr, strides from .strides in BYTES + ptr = getattr(getattr(data, "data", None), "ptr", None) + if dim is not None and ptr is not None: + itemsize = data.dtype.itemsize + strides = getattr(data, "strides", None) + stride = tuple(s // itemsize for s in strides) if strides else _row_major_stride(dim) + return int(ptr), describing_tensor(uid, tuple(dim), stride, _buffer_dtype_to_cudnn(data.dtype)) + + # jax and bare DLPack producers: one capsule read, which is also the + # only reader that gets cupy's byte strides right without knowing it is + # cupy. It declines two different ways and they are NOT the same + # answer, so they are told apart: "no protocol" means a pointer is all + # this buffer will ever yield, while "dtype I cannot name" (fp8, fp4, + # anything sub-byte) still has a real dim and stride worth keeping — + # the torch branch above records those dtypes, and a buffer should not + # be described differently for having come from jax. + from .frost.buffers import _dlpack_geometry + + geometry = _dlpack_geometry(data) + if geometry is None: # neither __dlpack__ nor __cuda_array_interface__ + return self._device_pointer(data), Tensor(uid=uid, dim=tuple(dim) if dim is not None else []) + ptr, dims, strides, name, _device = geometry + # data_type is None for a dtype with no cuDNN enum (fp4 and friends) + return ptr, describing_tensor(uid, tuple(dims), tuple(strides) if strides else _row_major_stride(dims), _buffer_dtype_to_cudnn(name)) + + @staticmethod + def _device_pointer(data: Any) -> int: + if type(data) is int: + return data + if hasattr(data, "data_ptr"): + return data.data_ptr() + import cudnn + + return cudnn._pybind_module._get_data_ptr(data) # dlpack fallback + def _uid_to_data(self, tensor_dict) -> Dict[int, Any]: """Normalize a variant pack keyed by Tensor / name / uid to uid -> data, starting from the graph's auto-bound inputs (user keys win).""" @@ -1833,6 +2081,9 @@ def deserialize(self, *args, **kwargs) -> None: self._lowered_graph = cudnn._pybind_module.backend_graph() self._lowered_graph.deserialize(*args, **kwargs) self._is_built = True + # The loaded graph carries its own variant_pack, so an order cached while + # this container held a different graph no longer describes it. + self._sorted_uids = None def _lower_to_cpp(self) -> Any: """Lower Python graph to C++ (the internal ``_pybind_module.backend_graph``).""" diff --git a/python/cudnn/datatypes.py b/python/cudnn/datatypes.py index 59b6ce963..b407a1227 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -167,6 +167,69 @@ def _torch_to_cudnn_data_type(torch_data_type) -> cudnn_data_type: return None +# cuDNN enum -> the dtype NAME frost.buffers speaks (its DTYPES table is keyed +# by name because a DLPack view needs no tensor library). Built once: this is +# read per operand per execute. +_CUDNN_TO_FROST_DTYPE_NAME = { + cudnn_data_type.FLOAT: "float32", + cudnn_data_type.HALF: "float16", + cudnn_data_type.BFLOAT16: "bfloat16", + cudnn_data_type.DOUBLE: "float64", + cudnn_data_type.INT64: "int64", + cudnn_data_type.INT32: "int32", + cudnn_data_type.INT8: "int8", + cudnn_data_type.UINT8: "uint8", + cudnn_data_type.BOOLEAN: "bool", + cudnn_data_type.FP8_E4M3: "float8_e4m3fn", + cudnn_data_type.FP8_E5M2: "float8_e5m2", + cudnn_data_type.FP8_E8M0: "float8_e8m0fnu", +} + + +def _cudnn_to_frost_dtype_name(data_type): + """Name for a cuDNN dtype in the vocabulary ``frost.buffers.DTYPES`` uses, + or None when the type has no DLPack-expressible name (the sub-byte ones — + fp4 has a DLPack code but an itemsize of 0 bytes, so a caller must pass it + as a typed buffer rather than as a bare address). + + Lives here so the mapping has one home; frost imports it rather than + keeping a second table.""" + return _CUDNN_TO_FROST_DTYPE_NAME.get(data_type) + + +_buffer_dtype_to_cudnn_dict = None + + +def _buffer_dtype_to_cudnn(dtype) -> cudnn_data_type: + """cuDNN enum for however a caller's buffer spells its dtype, or None. + + ONE table for every framework rather than one per framework: a torch dtype, + a numpy/cupy dtype and a bare name string are all hashable and mutually + unequal, so they coexist as keys and the caller needs no branch. numpy has + no bfloat16, which is why the name keys exist at all — that is the dtype a + DLPack read hands back for the case torch's ``__cuda_array_interface__`` + cannot express. + """ + global _buffer_dtype_to_cudnn_dict + if _buffer_dtype_to_cudnn_dict is None: + table = {name: enum for enum, name in _CUDNN_TO_FROST_DTYPE_NAME.items()} + if is_torch_available(): + table.update(_torch_to_cudnn_data_type_dict) + try: + import numpy + except ImportError: + pass + else: + for name, enum in list(table.items()): + if isinstance(name, str): + try: + table[numpy.dtype(name)] = enum + except TypeError: + pass # no numpy spelling (bfloat16); the name key serves it + _buffer_dtype_to_cudnn_dict = table + return _buffer_dtype_to_cudnn_dict.get(dtype) + + def _torch_to_cutlass_data_type(data_type, interpret_uint8_as_fp4x2: bool = False): # A torch dtype can only be passed in if torch is already imported, so probing # sys.modules avoids importing torch on behalf of other frameworks' dtypes. @@ -265,3 +328,27 @@ def _is_jax_array(input_tensor) -> bool: if jax is not None and isinstance(input_tensor, getattr(jax, "Array", ())): return True return type(input_tensor).__module__.startswith(("jax", "jaxlib")) + + +# The DLPack (code, bits) a cuDNN dtype travels as. The native variant pack +# speaks DLPack, so this is the one translation between it and the graph's +# vocabulary. +_CUDNN_TO_DLPACK_CODE_BITS = {} + + +def _init_dlpack_dtype_tables(): + from .frost.buffers import DTYPES + + for enum, name in _CUDNN_TO_FROST_DTYPE_NAME.items(): + code_bits = DTYPES.get(name) + if code_bits is None: + continue + _CUDNN_TO_DLPACK_CODE_BITS[enum] = code_bits + + +def _dlpack_code_bits(data_type): + """``(code, bits)`` for a cuDNN dtype, or ``(0, 0)`` when it has no DLPack + spelling — a slot with no dtype still carries its pointer and shape.""" + if not _CUDNN_TO_DLPACK_CODE_BITS: + _init_dlpack_dtype_tables() + return _CUDNN_TO_DLPACK_CODE_BITS.get(data_type, (0, 0)) diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 2d5d4da01..46c523a19 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -20,9 +20,9 @@ 3. ``CompiledPlan.execute(graph, uid_to_data, ctx)`` — hot path. ``uid_to_data`` is the caller's variant pack (tensor uid -> device buffer, exactly as the classic backend receives it); the - ``ExecutionContext`` carries the caller's handle / stream / workspace / - dynamic-shape overrides explicitly; engines must not hard-code a stream - or silently allocate hidden workspace. Engines that address buffers by + ``ExecutionContext`` carries the caller's handle / stream / workspace + explicitly; engines must not hard-code a stream or silently allocate + hidden workspace. Engines that address buffers by port name call ``resolve_node_buffers(graph, uid_to_data)`` (see below). Simple eager engines only implement ``execute()`` — the default ``build_plan`` @@ -48,6 +48,8 @@ def execute(self, graph, uid_to_data, ctx): from typing import TYPE_CHECKING, Any, Dict, List from .engine_ids import PYTHON_ENGINE_ID_BASE # noqa: F401 — re-exported for engine authors +from ..datatypes import _CUDNN_TO_FROST_DTYPE_NAME +from ..frost.buffers import DeviceView if TYPE_CHECKING: from ..pygraph import pygraph @@ -103,22 +105,163 @@ class ExecutionContext: handle: Any = None stream: Any = None workspace: Any = None - override_uids: Any = None - override_shapes: Any = None - override_strides: Any = None + + +class VariantPack: + """The caller's variant pack, normalized ONCE at the top of execute(). + + Whatever the caller passed — a torch tensor, a DeviceView, any + ``__dlpack__`` / ``__cuda_array_interface__`` producer, or a bare device + address — is converted here and then dropped. Below this point the cuDNN + backend and every python engine see the same two index-aligned sequences + and nothing else, so neither can behave differently on account of what the + caller happened to hold. + + The operands live in ``native``, a C container holding one ``DLTensor`` + each, read through the producer's ``__dlpack_c_exchange_api__`` vtable and + handed to kernels through the same one. + + ``uids`` is ASCENDING, matching the backend's own operand order + (``get_variant_pack_uids_sorted()``), so ``address`` goes straight to + ``_execute_with_raw_ptrs`` with no copy and no per-operand hash lookup. + + What the caller ACTUALLY passed is what is recorded, which need not be what + the graph declared: an engine reads the IR port for the shape the plan was + built for and this pack for the shape about to run. frost_gemm takes its + M/N/K from here; the backend takes only the pointer. + + Allocated per call. Two threads may execute one graph concurrently with + different buffers, and a shared pack would hand each thread the other's + pointers — silently, because every pointer in it is individually valid. + """ + + __slots__ = ("uids", "native", "_slot_of", "workspace", "workspace_bytes", "_device") + + def __init__(self, uids, native, workspace_ptr: int = 0, workspace_bytes: int = 0): + 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 + self._device = None + + @property + def address(self) -> int: + """The ``void*[]`` in slot order, for ``_execute_with_raw_ptrs``.""" + return self.native.address + + def all_contiguous(self): + """``(ok, slot)`` 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 + + @property + def device(self) -> int: + """The GPU this execute is going to, for the views handed to kernels. + + One per pack, not one per operand: cuDNN's own variant pack carries no + device at all. Read on demand; the backend path never asks. + """ + if self._device is None: + from ..frost.device import current_device + + self._device = current_device() + return self._device + + def slot(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] + + def ptr(self, tensor_or_uid) -> int: + return self.native.pointer(self.slot(tensor_or_uid)) + + def views(self, slots): + """The DLPack producers for ``slots``, in one crossing.""" + return self.native.views(list(slots), self.device) + + def view(self, slot: int): + """A DLPack producer over one operand, for a kernel that needs an + object rather than an address. + + This is the whole reason an engine never sees the caller's buffer: the + kernel gets ours, built from the pointer and the geometry recorded at + normalization. It costs more than handing the torch tensor straight + through (measured +1.6 us per operand, because tvm-ffi reads a torch + tensor through a C vtable and any python producer through a capsule) — + 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) + + 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 + virtual intermediate — is ABSENT, so ``.get(port) is None`` keeps meaning + what it meant when these were buffers.""" + + inputs: Dict[str, int] + outputs: Dict[str, int] + + +def bind_ports(graph: "pygraph", variant_pack: VariantPack) -> Dict[Any, PortSlots]: + """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 = {} + for port, t in ports.items(): + if t is None: + continue + slot = variant_pack.slot_of.get(t.uid) + if slot 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 + + return {node: PortSlots(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): + """A DLPack view over a caller-supplied bare address, shaped by the IR. + + Only reachable when the caller passed an int for this port. Requires the + graph to declare a dim and a dtype for it — an address carries neither, so + if the graph does not say, nobody can. + """ + from ..datatypes import _cudnn_to_frost_dtype_name + from ..frost import buffers + + dtype = _cudnn_to_frost_dtype_name(tensor.data_type) + if not tensor.dim or dtype is None: + raise ValueError( + f"node {node.name!r}: port {port!r} was given a bare device address, " + f"but the graph declares no {'dim' if not tensor.dim else 'data_type'} for tensor " + f"{tensor.name!r} — pass a buffer that carries its own shape and dtype, or declare them" + ) + return buffers.DeviceView(address, tuple(tensor.dim), dtype, buffers.current_device_id()) @dataclass(frozen=True) class NodeBuffers: - """Per-node ``{port_name: caller buffer}`` maps, the result of - ``resolve_node_buffers``. Only WIRED, NON-VIRTUAL ports - appear, and every one is guaranteed a buffer (a missing buffer raises at - resolution). Torch tensors arrive detached — both DLPack and - ``__cuda_array_interface__`` refuse to export ``requires_grad`` tensors, - and graph-level gradients are the backward nodes' contract, never - autograd tracing through an engine. Virtual intermediates carry no - caller buffers; engines that chain them across nodes key their own - scratch by ``node.inputs[port].uid``.""" + """DEPRECATED, kept until every engine takes ``VariantPack``. + + Per-node ``{port_name: caller buffer}`` maps, the result of + ``resolve_node_buffers``.""" inputs: Dict[str, Any] outputs: Dict[str, Any] @@ -142,6 +285,13 @@ def resolve(node, ports, direction): 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})") + if type(b) is int: + # A bare device address. The backend has always taken one + # (_pygraph._describe), so a python engine must too, or the + # same graph.execute() call succeeds or fails depending on + # which plan the heuristics happened to pick. The geometry is + # the one the graph declares for this port. + b = _view_over_address(b, t, node, port) bufs[port] = b.detach() if hasattr(b, "detach") else b return bufs @@ -151,11 +301,16 @@ def resolve(node, ports, direction): class CompiledPlan: """A compiled (graph, plan) artifact. Subclass for real JIT engines.""" + # Set True once execute() takes VariantPack. Until then execute() is handed the + # caller's raw {uid: buffer} map, as before. Migration flag; it goes away + # with the last engine that does not set it. + takes_variant_pack: bool = False + def get_workspace_size(self) -> int: """Workspace bytes this plan needs at execute time (default 0).""" return 0 - def execute(self, graph: "pygraph", uid_to_data: Dict[int, Any], ctx: ExecutionContext) -> None: + def execute(self, graph: "pygraph", variant_pack: "VariantPack", ctx: ExecutionContext) -> None: raise NotImplementedError diff --git a/python/cudnn/experimental/ops/sdpa.py b/python/cudnn/experimental/ops/sdpa.py index 2b34a64c0..52617f70b 100644 --- a/python/cudnn/experimental/ops/sdpa.py +++ b/python/cudnn/experimental/ops/sdpa.py @@ -142,12 +142,6 @@ def _packed_bhsd_stride(shape: Tuple[int, int, int, int]) -> Tuple[int, int, int return (s * h * d, d, h * d, 1) -def _get_variant_pack_uid_order(graph, present_uids): - if hasattr(graph, "_get_variant_pack_uids_sorted"): - return graph._get_variant_pack_uids_sorted() - return sorted(present_uids) - - def _validate_d256_oss_args( is_causal: bool, diagonal_alignment: int, @@ -714,23 +708,9 @@ def _sdpa_impl( cumulative_seq_len_q, cumulative_seq_len_kv, ) - uid_order = _get_variant_pack_uid_order( - graph, - [ - int(_UIDs.Q), - int(_UIDs.K), - int(_UIDs.V), - int(_UIDs.O), - int(_UIDs.STATS), - *([int(_UIDs.SEQ_LEN_Q)] if seq_len_q is not None else []), - *([int(_UIDs.SEQ_LEN_KV)] if seq_len_kv is not None else []), - *([int(_UIDs.CUM_SEQ_LEN_Q)] if cumulative_seq_len_q is not None else []), - *([int(_UIDs.CUM_SEQ_LEN_KV)] if cumulative_seq_len_kv is not None else []), - ], - ) - _fprop_cache[cache_key] = (graph, workspace_size, uid_order) + _fprop_cache[cache_key] = (graph, workspace_size) - graph, workspace_size, uid_order = _fprop_cache[cache_key] + graph, workspace_size = _fprop_cache[cache_key] # Allocate outputs and workspace (BHSD layout) # Workspace is per-call — PyTorch's caching allocator recycles the allocation. @@ -757,11 +737,7 @@ def _sdpa_impl( if cumulative_seq_len_kv is not None: uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_KV)] = cumulative_seq_len_kv - if hasattr(graph, "_execute_with_ptrs"): - ptrs = [uid_to_tensor[uid].data_ptr() for uid in uid_order] - graph._execute_with_ptrs(ptrs, workspace.data_ptr(), int(handle)) - else: - graph.execute(uid_to_tensor, workspace, handle=handle) + graph.execute(uid_to_tensor, workspace, handle=handle) return o_gpu, stats_gpu @@ -1007,27 +983,9 @@ def _sdpa_bwd_impl( cumulative_seq_len_kv, is_deterministic, ) - uid_order = _get_variant_pack_uid_order( - graph, - [ - int(_UIDs.Q), - int(_UIDs.K), - int(_UIDs.V), - int(_UIDs.O), - int(_UIDs.DO), - int(_UIDs.STATS), - int(_UIDs.DQ), - int(_UIDs.DK), - int(_UIDs.DV), - *([int(_UIDs.SEQ_LEN_Q)] if seq_len_q is not None else []), - *([int(_UIDs.SEQ_LEN_KV)] if seq_len_kv is not None else []), - *([int(_UIDs.CUM_SEQ_LEN_Q)] if cumulative_seq_len_q is not None else []), - *([int(_UIDs.CUM_SEQ_LEN_KV)] if cumulative_seq_len_kv is not None else []), - ], - ) - _bprop_cache[cache_key] = (graph, workspace_size, uid_order) + _bprop_cache[cache_key] = (graph, workspace_size) - graph, workspace_size, uid_order = _bprop_cache[cache_key] + graph, workspace_size = _bprop_cache[cache_key] # Allocate gradient outputs and workspace (same shapes as Q, K, V) dQ_gpu = torch.empty_like(q) @@ -1056,11 +1014,7 @@ def _sdpa_bwd_impl( if cumulative_seq_len_kv is not None: uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_KV)] = cumulative_seq_len_kv - if hasattr(graph, "_execute_with_ptrs"): - ptrs = [uid_to_tensor[uid].data_ptr() for uid in uid_order] - graph._execute_with_ptrs(ptrs, workspace.data_ptr(), int(handle)) - else: - graph.execute(uid_to_tensor, workspace, handle=handle) + graph.execute(uid_to_tensor, workspace, handle=handle) return dQ_gpu, dK_gpu, dV_gpu diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 1c4364c26..003b761d3 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -19,6 +19,8 @@ import ctypes +from cudnn import _pybind_module + # --------------------------------------------------------------------------- # DLPack ABI (dlpack.h v0.8 layout; the unversioned "dltensor" capsule) # --------------------------------------------------------------------------- @@ -58,15 +60,6 @@ class _DLManagedTensor(ctypes.Structure): ] -@_DELETER_T -def _noop_deleter(_ptr): - # views own no memory: the workspace (or caller buffer) outlives them - pass - - -_PyCapsule_New = ctypes.pythonapi.PyCapsule_New -_PyCapsule_New.restype = ctypes.py_object -_PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] _PyCapsule_GetPointer = ctypes.pythonapi.PyCapsule_GetPointer _PyCapsule_GetPointer.restype = ctypes.c_void_p _PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] @@ -75,6 +68,10 @@ def _noop_deleter(_ptr): _PyCapsule_SetName.argtypes = [ctypes.py_object, ctypes.c_char_p] # name -> (DLPack type code, bits); typestr -> name for the CAI path +# name -> (DLPack type code, bits). Names match how torch spells them, so one +# table serves both a dtype read off a buffer and one named in a graph. +# Sub-byte types are deliberately absent: DTYPE_ITEMSIZE below is bits // 8, so +# fp4 would land on 0 and take every byte/element conversion with it. DTYPES = { "float32": (2, 32), "float16": (2, 16), @@ -85,6 +82,9 @@ def _noop_deleter(_ptr): "int8": (0, 8), "uint8": (1, 8), "bool": (6, 8), + "float8_e4m3fn": (10, 8), + "float8_e5m2": (12, 8), + "float8_e8m0fnu": (14, 8), } _TYPESTR = {" int: return self._ptr @@ -166,23 +170,16 @@ def __dlpack_device__(self): return (_KDL_CUDA, self._device_id) def __dlpack__(self, *, stream=None, **_kwargs): - # a fresh managed struct per call; shape array + struct stay alive on - # the view (the no-op deleter frees nothing) - ndim = len(self.shape) - shape_arr = (ctypes.c_int64 * max(ndim, 1))(*self.shape) + """Delegate to a slot, which owns the struct it hands out. + + A view has nowhere to put a struct that must outlive the capsule: it + cannot know when the consumer is done with it, and cute's from_dlpack + keeps the pointer rather than copying the DLTensor. Holding the struct + on the view and shipping a no-op deleter -- which is what this did -- + is a use-after-free the moment a consumer outlives the view. + """ code, bits = DTYPES[self.dtype] - mt = _DLManagedTensor() - mt.dl_tensor.data = ctypes.c_void_p(self._ptr) - mt.dl_tensor.device = _DLDevice(_KDL_CUDA, self._device_id) - mt.dl_tensor.ndim = ndim - mt.dl_tensor.dtype = _DLDataType(code, bits, 1) - mt.dl_tensor.shape = shape_arr - mt.dl_tensor.strides = None # None = compact row-major - mt.dl_tensor.byte_offset = 0 - mt.manager_ctx = None - mt.deleter = _noop_deleter - self._keepalive.append((mt, shape_arr)) - return _PyCapsule_New(ctypes.addressof(mt), b"dltensor", None) + return _pybind_module.make_slot(self._ptr, list(self.shape), code, bits, self._device_id).__dlpack__() class DeviceBuffer(DeviceView): @@ -212,7 +209,28 @@ def __del__(self): def probe(buf): """(ptr, shape, strides_in_elements_or_None, dtype_name, device_id) of a device buffer, via ``__cuda_array_interface__`` when available (torch, - CuPy, numba) else the ``__dlpack__`` protocol.""" + CuPy, numba) else the ``__dlpack__`` protocol. + + Raises for a buffer it cannot read. Callers that would rather have the + geometry of a buffer whose DTYPE has no name here — fp8, fp4, anything + sub-byte — want :func:`_dlpack_geometry`, which separates the two + failures.""" + geometry = _dlpack_geometry(buf) + if geometry is None: + raise TypeError(f"buffer of type {type(buf).__name__} exposes neither __cuda_array_interface__ nor __dlpack__") + if geometry[3] is None: + raise TypeError(f"unsupported buffer dtype for {type(buf).__name__}") + return geometry + + +def _dlpack_geometry(buf): + """``probe``'s reading, with its two declines made distinguishable. + + Returns None when the buffer exposes neither protocol — nothing but a + pointer will ever come out of it. Returns the 5-tuple with ``dtype_name`` + set to None when the buffer IS readable but its dtype has no name in + ``DTYPES``; dim and stride are real in that case and worth keeping. + """ try: # torch's property RAISES for dtypes CAI can't express (bf16) instead # of being absent — treat any failure as "no CAI" and use DLPack @@ -234,11 +252,11 @@ def probe(buf): dl = getattr(buf, "__dlpack__", None) if dl is None: - raise TypeError(f"buffer of type {type(buf).__name__} exposes neither __cuda_array_interface__ nor __dlpack__") + return None # stream=-1 is DLPack's "the caller handles synchronisation; do no - # bookkeeping". probe() only reads metadata, so it never needed any -- - # and the default makes torch call record_stream, which is illegal inside - # a CUDA graph capture. + # bookkeeping". This only reads metadata, so it never needed any -- and the + # default makes torch call record_stream, which is illegal inside a CUDA + # graph capture. try: capsule = dl(stream=-1) except TypeError: # a producer whose __dlpack__ predates the stream kwarg @@ -248,9 +266,7 @@ def probe(buf): t = mt.dl_tensor shape = tuple(t.shape[i] for i in range(t.ndim)) strides = tuple(t.strides[i] for i in range(t.ndim)) if t.strides else None - dtype = _CODE_BITS.get((t.dtype.code, t.dtype.bits)) - if dtype is None or t.dtype.lanes != 1: - raise TypeError(f"unsupported buffer dtype (code={t.dtype.code}, bits={t.dtype.bits}, lanes={t.dtype.lanes})") + dtype = _CODE_BITS.get((t.dtype.code, t.dtype.bits)) if t.dtype.lanes == 1 else None ptr = (t.data or 0) + t.byte_offset device_id = t.device.device_id # release: mark the capsule consumed and run its deleter diff --git a/python/cudnn/frost/device.py b/python/cudnn/frost/device.py index f6b9c7e58..c61771ddc 100644 --- a/python/cudnn/frost/device.py +++ b/python/cudnn/frost/device.py @@ -3,18 +3,19 @@ """Which GPU a FROST plan is built for — shared by every FROST engine. -The device is then recorded on the compiled plan and :func:`check_buffer_device` -re-checks it at execute time, so running a plan against another GPU's buffers -fails loudly instead of launching a kernel whose baked constants describe the -wrong hardware. +The device is recorded on the compiled plan and compared against +:func:`current_device` at execute time, so a plan whose baked constants +describe one GPU fails loudly instead of launching on another. + +That comparison is about the LAUNCH, not the buffers. cuDNN's own variant pack +carries pointers, uids and a workspace and no device at all, so an operand's +device is not something the front end has an opinion about. """ from __future__ import annotations import functools -_DLPACK_CUDA_KINDS = (2, 13) # kDLCUDA, kDLCUDAManaged - @functools.lru_cache(maxsize=1) def _driver(): @@ -131,35 +132,6 @@ def device_name(device: int) -> str: return _ck(*drv.cuDeviceGetName(256, _device_handle(device))).split(b"\x00")[0].decode() -def buffer_device(buf): - """CUDA ordinal a runtime buffer lives on, or ``None`` when it carries no - CUDA device (host array, raw pointer, int).""" - describe = getattr(buf, "__dlpack_device__", None) - if describe is None: - return None - try: - kind, index = describe() - except Exception: # noqa: BLE001 — a buffer that cannot describe itself is not ours to check - return None - return int(index) if int(kind) in _DLPACK_CUDA_KINDS else None - - -def check_buffer_device(buffers, plan_device: int, *, what: str = "plan") -> None: - """Raise if any runtime buffer lives on a GPU other than the one the plan was - built for — its baked SMEM / cluster / arch constants describe ``plan_device`` - only. Non-tensor entries (raw pointers, ints) carry no device and are skipped.""" - for buf in buffers: - index = buffer_device(buf) - if index is None or index == plan_device: - continue - raise ValueError( - f"cudnn.frost: this {what} was built for cuda:{plan_device} but a buffer is on " - f"cuda:{index}. The kernel's SMEM pipeline depth, cluster count and target SM are " - f"baked at build time, so a plan cannot move between GPUs — rebuild it with " - f"cuda:{index} current." - ) - - class device_context: """Bind ``device``'s primary context for the enclosing block, then restore whatever was bound before. The retain/release pair is refcounted, so this diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index ca9dcaa70..f198cfa4e 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -26,6 +26,8 @@ from __future__ import annotations +from cudnn import _pybind_module + from . import buffers # TMA tensormap patches need their 128-byte slot alignment; every other consumer @@ -66,6 +68,19 @@ def base_align(self) -> int: return self._base_align +def carve_plan(owner: str, regions) -> "_pybind_module.WorkspaceCarve": + """Compile a build-time carve: ``[(offset, dtype, shape), ...]``. + + The regions are fixed once :class:`WorkspaceLayout` has run; only the base + pointer arrives per execute, so one crossing serves them all. + """ + spec = [] + for offset, dtype, shape in regions: + code, bits = buffers.DTYPES[dtype] + spec.append((int(offset), code, bits, [int(extent) for extent in shape])) + return _pybind_module.WorkspaceCarve(owner, spec) + + class Workspace: """Execute-time view onto the caller's workspace buffer. @@ -90,6 +105,9 @@ def __init__(self, buffer, required_bytes: int, owner: str, *, align: int = DEFA raise ValueError(f"{owner}: needs a {required_bytes}-byte workspace, got {nbytes} bytes " "(size it with graph.get_workspace_size())") if ptr % align != 0: raise ValueError(f"{owner}: the workspace buffer must be {align}-byte aligned; got 0x{ptr:x}") + self._init(ptr, nbytes, device, owner, align) + + def _init(self, ptr, nbytes, device, owner, align): self._ptr = ptr self._device = device self._nbytes = nbytes @@ -97,17 +115,44 @@ def __init__(self, buffer, required_bytes: int, owner: str, *, align: int = DEFA self._align = int(align) self._offset = 0 + @classmethod + def over(cls, variant_pack, required_bytes: int, owner: str, *, align: int = DEFAULT_ALIGN) -> "Workspace": + """The same validated carver, over a workspace the pack already read.""" + required_bytes = int(required_bytes) + ptr, nbytes = variant_pack.workspace, variant_pack.workspace_bytes + if not ptr: + raise ValueError( + f"{owner} requires a {required_bytes}-byte workspace but execute() received " + f"none; allocate graph.get_workspace_size() bytes and pass the buffer to execute()" + ) + if nbytes < required_bytes: + raise ValueError(f"{owner}: needs a {required_bytes}-byte workspace, got {nbytes} bytes (size it with graph.get_workspace_size())") + if ptr % align != 0: + raise ValueError(f"{owner}: the workspace buffer must be {align}-byte aligned; got 0x{ptr:x}") + self = cls.__new__(cls) + self._init(ptr, nbytes, variant_pack.device, owner, align) + return self + @property def nbytes(self) -> int: return self._nbytes - def view(self, offset: int, dtype: str, shape) -> buffers.DeviceView: - """The region a :class:`WorkspaceLayout` reserved at ``offset``.""" + def view(self, offset: int, dtype: str, shape): + """The region a :class:`WorkspaceLayout` reserved at ``offset``. + + A carve is the same kind of buffer a caller operand is, so a graph + hands its kernels one buffer type rather than two. + """ count = 1 for extent in shape: count *= int(extent) self._check_span(offset, count * buffers.DTYPE_ITEMSIZE[dtype]) - return buffers.DeviceView(self._ptr + offset, shape, dtype, self._device) + code, bits = buffers.DTYPES[dtype] + return _pybind_module.make_slot(self._ptr + offset, list(shape), code, bits, self._device) + + def carve(self, plan): + """Every region a :func:`carve_plan` describes, in one crossing.""" + return plan.carve(self._ptr, self._nbytes, self._device) def take(self, numel: int, dtype: str) -> buffers.DeviceView: """The next region dealt sequentially: a 1-D ``numel``-element view.""" diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 4665ce6d3..c3fca98a3 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -694,12 +694,27 @@ def _plan_device() -> int: return current_device() -def _check_plan_device(variant_pack, plan_device: int) -> None: +def _check_plan_device(plan_device: int) -> None: """A plan's SMEM depth / cluster count / target SM are baked for ONE GPU; - refuse buffers from another rather than launching a mismatched kernel.""" - from cudnn.frost.device import check_buffer_device + refuse to launch it anywhere else. + + Asks where the launch is going, not where each operand lives. The operands + were never the question — a kernel built for one arch produces garbage on + another whoever owns the memory — and the backend does not look at operand + 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. + """ + from cudnn.frost.device import current_device - check_buffer_device(variant_pack.values(), plan_device, what="FROST plan") + device = current_device() + if device != plan_device: + raise ValueError( + f"cudnn.frost: this FROST plan was built for cuda:{plan_device} but cuda:{device} is " + f"current. The kernel's SMEM pipeline depth, cluster count and target SM are baked at " + f"build time, so a plan cannot move between GPUs — rebuild it with cuda:{device} current." + ) def _grid_num_clusters(cfg: TileConfig, device=None) -> int: @@ -1806,11 +1821,20 @@ def __call__(self, variant_pack, stream=None): raise TypeError( "compiled kernels are called with a variant-pack dict " "{cuDNN tensor | uid | name: buffer}; got " f"{type(variant_pack).__name__}" ) - _check_plan_device(variant_pack, self.device) if self.binding is None: raise NotImplementedError("variant-pack call is not wired up for this graph type") + 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. + """ + _check_plan_device(self.device) b = self.binding - resolved = resolve_variant_pack(variant_pack, b) def pull(t, role): if t is None or id(t) not in resolved: @@ -2801,7 +2825,9 @@ def _moe_carve_workspace(caller, n_slots: int, plan: str): executor converts through the one-time adapter registered above.""" _register_legacy_device_view_adapter() need = n_slots * _MOE_DESC_SLOT_BYTES - ws = Workspace(caller, need, plan, align=_MOE_DESC_SLOT_BYTES) + # already carved when the engine handed one down; a raw buffer only when + # the plan allocated its own (the direct jit_from_cudnn_graph path) + ws = caller if isinstance(caller, Workspace) else Workspace(caller, need, plan, align=_MOE_DESC_SLOT_BYTES) return ws.view(0, "int64", (need // 8,)) @@ -2856,7 +2882,7 @@ def __call__(self, variant_pack, workspace=None, stream=None): raise TypeError( "compiled kernels are called with a variant-pack dict " "{cuDNN tensor | uid | name: buffer}; got " f"{type(variant_pack).__name__}" ) - _check_plan_device(variant_pack, self.device) + _check_plan_device(self.device) return self._call_variant_pack(variant_pack, workspace, stream) def _launch_single(self, token, weight, first_token_offset, output, snke, workspace=None, stream=None): @@ -3232,7 +3258,7 @@ def __call__(self, variant_pack, workspace=None, stream=None): raise TypeError( "compiled kernels are called with a variant-pack dict " "{cuDNN tensor | uid | name: buffer}; got " f"{type(variant_pack).__name__}" ) - _check_plan_device(variant_pack, self.device) + _check_plan_device(self.device) return self._call_variant_pack(variant_pack, workspace, stream) def _launch_single(self, token, weight, sfa, sfb, first_token_offset, output, snke, workspace=None, stream=None): diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index eee374a8c..a30382dfe 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -14,6 +14,7 @@ from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan, ExecutionContext, PlanConfig +from cudnn.frost.workspace import Workspace if TYPE_CHECKING: from cudnn._pygraph import pygraph @@ -22,42 +23,41 @@ class _FrostGemmPlan(CompiledPlan): """A compiled fused-GEMM kernel plus the graph binding it was compiled for.""" + takes_variant_pack = True + def __init__(self, compiled): self._compiled = compiled # Keyed by tensor OBJECT, not uid: one tensor can occupy two operand # roles (matmul(A, A)), and resolve_variant_pack treats a repeated uid # as ambiguous. self._tensors = list(compiled.binding.bound_tensors()) + self._slots = None def get_workspace_size(self) -> int: return int(getattr(self._compiled, "workspace_bytes", 0) or 0) - def execute(self, graph, uid_to_data, ctx: ExecutionContext) -> None: - pack, missing = {}, [] - for t in self._tensors: - buf = uid_to_data.get(t.get_uid()) - if buf is None: - missing.append(t.get_name() or t.get_uid()) - else: - pack[t] = buf - if missing: - raise ValueError(f"frost_gemm: the variant pack is missing buffers for {missing}") + def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: + slots = self._slots + if slots is None: + try: + slots = self._slots = [variant_pack.slot(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) required = self.get_workspace_size() - if required: - _check_workspace(ctx.workspace, required) - self._compiled(pack, ctx.workspace, stream=ctx.stream) + # 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: + # 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) else: - self._compiled(pack, stream=ctx.stream) - - -def _check_workspace(workspace, required: int) -> None: - """A FROST executor carves its scratch out of the CALLER's workspace: no - hidden per-execute allocation, stable pointers, CUDA-graph friendly.""" - if workspace is None: - raise ValueError(f"frost_gemm needs a {required}-byte workspace; execute() got none — allocate graph.get_workspace_size() bytes and pass it") - available = workspace.numel() * workspace.element_size() if hasattr(workspace, "numel") else len(workspace) - if available < required: - raise ValueError(f"frost_gemm needs a {required}-byte workspace; the buffer provides {available}") + self._compiled(dict(zip(self._tensors, views)), *extra, stream=ctx.stream) class FrostGemmEngine(BaseEngine): diff --git a/python/cudnn/graph_types.py b/python/cudnn/graph_types.py index 35283f286..aec72dd72 100644 --- a/python/cudnn/graph_types.py +++ b/python/cudnn/graph_types.py @@ -106,13 +106,6 @@ class Tensor: # (set_name / set_uid) delegate to the graph so its indexes stay coherent. owner: Any = field(default=None, repr=False) - def __setattr__(self, name, value): - # direct attribute writes freeze with the owning graph (the fluent - # setters are guarded separately and give a richer error) - if getattr(self, "_frozen", False) and name != "_frozen": - raise RuntimeError(f"cannot set Tensor.{name}: the owning graph is frozen after lowering/planning") - object.__setattr__(self, name, value) - def _guard(self, what: str = "mutate a tensor attribute") -> None: g = self.owner() if self.owner is not None else None if g is not None: @@ -241,3 +234,40 @@ def validate(self) -> None: # NOTE: hash/eq are object identity (dataclass eq=False). uid and name are # mutable, so value-based hashing would violate the dict-key invariant. + + +def describing_tensor(uid: int, dim, stride, data_type) -> Tensor: + """A Tensor describing a caller's buffer, built without the dataclass + ``__init__``. + + ``execute()`` builds one of these per operand per call, and the generated + ``__init__`` sets seventeen attributes and runs two default factories to do + it: 0.71 us against 0.29 for assigning the four that are known. Every field + left unset resolves to the class attribute the dataclass already installed + for its default, so the result is indistinguishable from ``Tensor(...)`` -- + ``test_describing_tensor_matches_the_dataclass`` compares them field by + field, and fails loudly if a new field arrives with a ``default_factory`` + (those get no class attribute, so reading one would raise). + """ + tensor = object.__new__(Tensor) + attributes = tensor.__dict__ + attributes["uid"] = uid + attributes["dim"] = dim + attributes["stride"] = stride + attributes["data_type"] = data_type + return tensor + + +def byte_size(tensor: Tensor) -> int: + """Bytes a dense tensor of this dim and dtype occupies, or 0 when the dtype + has no known width (a bare address describes neither).""" + from .datatypes import _CUDNN_TO_FROST_DTYPE_NAME + from .frost.buffers import DTYPE_ITEMSIZE + + name = _CUDNN_TO_FROST_DTYPE_NAME.get(tensor.data_type) + if name is None or not tensor.dim: + return 0 + total = DTYPE_ITEMSIZE[name] + for extent in tensor.dim: + total *= int(extent) + return total diff --git a/python/cudnn/linear_attention/engine_utils.py b/python/cudnn/linear_attention/engine_utils.py index cf977689d..d44f366aa 100644 --- a/python/cudnn/linear_attention/engine_utils.py +++ b/python/cudnn/linear_attention/engine_utils.py @@ -6,9 +6,10 @@ from __future__ import annotations -from cudnn.engines.base import CompiledPlan, resolve_node_buffers +from cudnn.engines.base import CompiledPlan, NodeBuffers, bind_ports from cudnn.frost import buffers +from cudnn.frost.workspace import Workspace def _dtype_name(dt) -> str: @@ -44,25 +45,62 @@ def _require_state_pair(engine: str, node) -> None: class _FrostPlan(CompiledPlan): + """A compiled linear-attention kernel, driven from the normalized pack. + + The port-to-slot join is a property of the graph, so it happens once and is + kept; only the addresses change between executes. What the kernel receives + is built from the pack, never the caller's object — the geometry it is + checked against and the geometry it runs on are then the same reading. + """ + + takes_variant_pack = True + def __init__(self, compiled): self._compiled = compiled + self._ports = None + self._name = type(compiled).__name__ def get_workspace_size(self) -> int: return self._compiled.workspace_bytes() - def execute(self, graph, uid_to_data, ctx) -> None: - node_buffers = resolve_node_buffers(graph, uid_to_data) - self._compiled(node_buffers, workspace=getattr(ctx, "workspace", None), stream=getattr(ctx, "stream", None)) - - -def _check_contiguous(plan_name: str, **bufs) -> None: - """Contiguity gate over the caller's buffers (pass-through, no staging).""" - for name, b in bufs.items(): - if b is None: - continue - _ptr, shape, strides, _dtype, _dev = buffers.probe(b) - if not buffers.is_contiguous(shape, strides): - raise ValueError(f"{plan_name}: buffer for {name!r} must be contiguous (buffers pass straight to the kernel)") + def execute(self, graph, variant_pack, ctx) -> None: + ports = self._ports + if ports is None: + ports = self._ports = bind_ports(graph, variant_pack) + _check_contiguous(variant_pack, ports) + 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())) + 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() + workspace = Workspace.over(variant_pack, required, self._name) if required else None + self._compiled(node_buffers, workspace=workspace, stream=ctx.stream) + + +def _check_contiguous(variant_pack, ports) -> None: + """Contiguity gate over the whole pack, decided from the strides it holds. + + The dim and stride were taken from the caller's object once, at + normalization; probing each buffer again cost 8.6 us apiece — nine per GDN + forward — to learn what the pack already knows. The scan itself is in the + native pack, 0.24 us for eight operands, so only naming the offender costs + anything and that happens once, on the way to raising. + + One gate for every kernel rather than a call per compiled callable naming + its own ports: the rule was the same list every time, and a port added to a + node but forgotten there would have gone unchecked. + """ + ok, offender = variant_pack.all_contiguous() + if ok: + return + for node, slots in ports.items(): + for direction in (slots.inputs, slots.outputs): + for port, slot in direction.items(): + if slot == offender: + raise ValueError(f"cudnn.frost {node.name!r}: buffer for {port!r} must be contiguous (buffers pass straight to the kernel)") + raise ValueError(f"cudnn.frost: the buffer at variant-pack slot {offender} must be contiguous") _pinned_engines = None # e.g. ("gdn_cutile",) -- set by a suite, None => the manifest decides diff --git a/python/cudnn/linear_attention/frost/gdn2_engine.py b/python/cudnn/linear_attention/frost/gdn2_engine.py index 47db66ab6..3bcd0fa8a 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -15,8 +15,8 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout -from ..engine_utils import _FrostPlan, _check_contiguous, _require_dtype, _require_state_pair +from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan +from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair def _the_gdn2_node(graph): @@ -144,6 +144,18 @@ def __init__(self, node, kernel_mod): self._off_tensormaps = layout.add(self._tensormap_bytes, align=128) self._ws_bytes = layout.size + self._carve = carve_plan( + "gdn2", + [ + (self._off_sched, "int32", (2,)), + (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_work_count, "int32", (1,)), + (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), + (self._off_tensormaps, "int64", (self._tensormap_bytes // 8,)), + ], + ) + def workspace_bytes(self) -> int: return self._ws_bytes @@ -159,18 +171,11 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: s0 = nb.inputs.get("initial_state") o = nb.outputs["O"] fs = nb.outputs["final_state"] if self._has_fs else None - _check_contiguous("Gdn2FrostEngine (GDN2)", q=q, k=k, v=v, g=g, beta=beta, w=w, cu_seqlens=cu, initial_state=s0, O=o, final_state=fs) stream = stream if stream is not None else 0 - ws = Workspace(workspace, self._ws_bytes, "Gdn2FrostEngine (GDN2)") - sched_ctr = ws.view(self._off_sched, "int32", (2,)) - from .common.split_k import WORK_ITEM_FIELDS - - work_items = ws.view(self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - item_scratch = ws.view(self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - work_count = ws.view(self._off_work_count, "int32", (1,)) - chunk_scratch = ws.view(self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, self._n_heads_out)) + ws = workspace + sched_ctr, work_items, item_scratch, work_count, chunk_scratch, tensormaps = ws.carve(self._carve) from .common.split_k import build_split_table build_split_table( @@ -206,7 +211,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: work_items=work_items, work_count=work_count, sched_ctr=sched_ctr, - tensormap_workspace=ws.view(self._off_tensormaps, "int64", (self._tensormap_bytes // 8,)), + tensormap_workspace=tensormaps, stream=stream, ) return None diff --git a/python/cudnn/linear_attention/frost/gdn_engine.py b/python/cudnn/linear_attention/frost/gdn_engine.py index 4c43a5a7b..b5fd88d3b 100644 --- a/python/cudnn/linear_attention/frost/gdn_engine.py +++ b/python/cudnn/linear_attention/frost/gdn_engine.py @@ -14,8 +14,8 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout -from ..engine_utils import _FrostPlan, _check_contiguous, _require_dtype, _require_state_pair +from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan +from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair def _the_gdn_node(graph): @@ -198,6 +198,19 @@ def __init__(self, node, kernel_mod): self._off_chunk_scratch = layout.add(self._chunk_scratch_rows * HO * 4) self._ws_bytes = layout.size + regions = [ + (self._off_tensormaps, "int64", (self._tensormap_words,)), + (self._off_sched, "int32", (2,)), + ] + if self._split: + regions += [ + (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_work_count, "int32", (1,)), + (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), + ] + self._carve = carve_plan("gdn", regions) + def workspace_bytes(self) -> int: return self._ws_bytes @@ -213,19 +226,12 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: o = nb.outputs["O"] fs = nb.outputs["final_state"] if self._has_fs else None h = nb.outputs["H"] if self._has_h else None - _check_contiguous("GdnFrostEngine (GDN)", q=q, k=k, v=v, g=g, beta=beta, cu_seqlens=cu, initial_state=s0, O=o, final_state=fs) - ws = Workspace(workspace, self._ws_bytes, "GdnFrostEngine (GDN)") + ws = workspace stream = stream if stream is not None else 0 - sched_ctr = ws.view(self._off_sched, "int32", (2,)) work_items = work_count = None if self._split: - from .common.split_k import WORK_ITEM_FIELDS - - work_items = ws.view(self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - item_scratch = ws.view(self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - work_count = ws.view(self._off_work_count, "int32", (1,)) - chunk_scratch = ws.view(self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, self._n_heads_out)) + tensormaps, sched_ctr, work_items, item_scratch, work_count, chunk_scratch = ws.carve(self._carve) _splits_for( g, cu, @@ -242,6 +248,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: sched_ctr=sched_ctr, ) else: + tensormaps, sched_ctr = ws.carve(self._carve) buffers.memset_zero_async(sched_ctr.data_ptr(), sched_ctr.nbytes, stream) self._kernel.chunk_gdn_sm100( @@ -261,7 +268,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: checkpoint_every_n_tokens=self._ckpt, output_h=h, log_gate=True, - workspace=ws.view(self._off_tensormaps, "int64", (self._tensormap_words,)), + workspace=tensormaps, stream=stream, ) return None @@ -333,6 +340,26 @@ def __init__(self, node, kernel_mod): self._shapes = (total, HQ, HV, HO, K, V, B) self._ws_bytes = layout.size + # The regions every backward carves. The rest (io state, regenerated H, + # head-group scratch) hang off build-time branches that are usually + # off, and stay on view() rather than turning this into index + # bookkeeping. The three scheduler views overlap on purpose: one ring + # each for the regen and bwd kernels, and both at once for the split + # pipeline that zeroes them. + self._carve = carve_plan( + "gdn_bwd", + [ + (self._off_sched, "int32", (2,)), + (self._off_sched + 8, "int32", (2,)), + (self._off_sched, "int32", (4,)), + (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_work_count, "int32", (1,)), + (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), + (self._off_tensormaps, "int64", (self._tensormap_words,)), + ], + ) + def workspace_bytes(self) -> int: return self._ws_bytes @@ -354,38 +381,11 @@ def __call__(self, node_buffers, *, workspace=None, stream=None): dv = nb.outputs["dV"] dg = nb.outputs["dG"] dbeta = nb.outputs["dBeta"] - _check_contiguous( - "GdnFrostEngine (GDN_BWD)", - q=q, - k=k, - v=v, - g=g, - beta=beta, - cu_seqlens=cu, - dO=do, - h=h_in, - initial_state=s0, - d_final_state=dht, - d_initial_state=ds0, - dQ=dq, - dK=dk, - dV=dv, - dG=dg, - dBeta=dbeta, - ) - - ws = Workspace(workspace, self._ws_bytes, "GdnFrostEngine (GDN_BWD)") + ws = workspace total, HQ, HV, HO, K, V, B = self._shapes stream = stream if stream is not None else 0 - sched_fwd = ws.view(self._off_sched, "int32", (2,)) - sched_bwd = ws.view(self._off_sched + 8, "int32", (2,)) - from .common.split_k import WORK_ITEM_FIELDS - - work_items = ws.view(self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - item_scratch = ws.view(self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - work_count = ws.view(self._off_work_count, "int32", (1,)) - chunk_scratch = ws.view(self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)) + sched_fwd, sched_bwd, sched_all, work_items, item_scratch, work_count, chunk_scratch, tensormaps = ws.carve(self._carve) _splits_for( g, cu, @@ -399,7 +399,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None): self._b_t, stream, log_gate=True, - sched_ctr=ws.view(self._off_sched, "int32", (4,)), + sched_ctr=sched_all, ) s0_io = None @@ -459,7 +459,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None): work_count=work_count, sched_ctr=sched_bwd if self._bwd_dyn_sched else None, log_gate=True, - workspace=ws.view(self._off_tensormaps, "int64", (self._tensormap_words,)), + workspace=tensormaps, stream=stream, ) if self._is_gva: diff --git a/python/cudnn/linear_attention/frost/kda_engine.py b/python/cudnn/linear_attention/frost/kda_engine.py index 11a9d9246..85f58be11 100644 --- a/python/cudnn/linear_attention/frost/kda_engine.py +++ b/python/cudnn/linear_attention/frost/kda_engine.py @@ -14,8 +14,8 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout -from ..engine_utils import _FrostPlan, _check_contiguous, _require_dtype, _require_state_pair +from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan +from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair def _the_kda_node(graph): @@ -156,6 +156,18 @@ def __init__(self, node, kernel_mod): self._off_tensormaps = layout.add(self._tensormap_bytes, align=128) self._ws_bytes = layout.size + self._carve = carve_plan( + "kda", + [ + (self._off_sched, "int32", (2,)), + (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), + (self._off_work_count, "int32", (1,)), + (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), + (self._off_tensormaps, "int64", (self._tensormap_bytes // 8,)), + ], + ) + def workspace_bytes(self) -> int: return self._ws_bytes @@ -170,18 +182,11 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: s0 = nb.inputs.get("initial_state") o = nb.outputs["O"] fs = nb.outputs["final_state"] if self._has_fs else None - _check_contiguous("KdaFrostEngine (KDA)", q=q, k=k, v=v, g=g, beta=beta, cu_seqlens=cu, initial_state=s0, O=o, final_state=fs) stream = stream if stream is not None else 0 - ws = Workspace(workspace, self._ws_bytes, "KdaFrostEngine (KDA)") - sched_ctr = ws.view(self._off_sched, "int32", (2,)) - from .common.split_k import WORK_ITEM_FIELDS - - work_items = ws.view(self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - item_scratch = ws.view(self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)) - work_count = ws.view(self._off_work_count, "int32", (1,)) - chunk_scratch = ws.view(self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, self._n_heads_out)) + ws = workspace + sched_ctr, work_items, item_scratch, work_count, chunk_scratch, tensormaps = ws.carve(self._carve) from .common.split_k import build_split_table build_split_table( @@ -224,7 +229,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: work_items=work_items, work_count=work_count, sched_ctr=sched_ctr, - tensormap_workspace=ws.view(self._off_tensormaps, "int64", (self._tensormap_bytes // 8,)), + tensormap_workspace=tensormaps, stream=stream, ) return None diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index 42abd2eab..83bf5b409 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -2732,14 +2732,6 @@ def _data_ptr(t) -> int: return t.__cuda_array_interface__["data"][0] -def _stream_capturing(stream) -> bool: - """True when ``stream`` is inside CUDA-graph capture.""" - from cuda.bindings import runtime as _rt - - err, status = _rt.cudaStreamIsCapturing(int(stream)) - return int(err) == 0 and int(status) != 0 - - def _cutlass_io_dtype(dtype): name = str(dtype) if "bfloat16" in name: @@ -3081,89 +3073,62 @@ def chunk_gdn2_sm100( compiled = cache["compiled"] - # desc key: buffer identity + cu _version so address reuse forces a rebuild + # The descriptors encode cu_seqlens' CONTENTS, which no key built from the + # buffers can track. The skip this replaces asked torch's _version counter, + # so it was sound for a torch caller and silently stale for every other + # producer. Rebuilding unconditionally measures free: 131 vs 135 us of host + # time, and 157 either way once the launches are waited on. h_for_descs = output_checkpoints if enable_checkpoints else None - desc_key = ( - _data_ptr(q), - _data_ptr(k), - _data_ptr(v), - _data_ptr(gate), - _data_ptr(beta), - _data_ptr(w), - _data_ptr(output), - _data_ptr(h_for_descs) if h_for_descs is not None else 0, - checkpoint_every_n_tokens, - tuple(h_for_descs.shape) if h_for_descs is not None else (), - tuple(q.shape), - tuple(k.shape), - tuple(v.shape), - tuple(gate.shape), - tuple(beta.shape), - tuple(w.shape), - tuple(output.shape), - ) - cu_versions = (getattr(cu_seqlens, "_version", 0),) - if ( - cache.get("desc_key") != desc_key - or cache.get("desc_cu") is not cu_seqlens - or cache.get("desc_cu_versions") != cu_versions - or cache.get("desc_workspace_ptr") != _data_ptr(tensormap_workspace) - or _stream_capturing(stream) - ): - if cache.get("build_descs_has_h") != (h_for_descs is not None): - cache.pop("build_descs", None) - cache["build_descs_has_h"] = h_for_descs is not None - if "build_descs" not in cache: - io_dtype = _cutlass_io_dtype(q.dtype) - - def _bd3(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - return c - - def _bd4(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) - return c - - cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() - cache["build_descs"] = cute.compile( - _build_descs, - io_dtype, - CFG.B_T, - _bd3(q), - _bd3(k), - _bd3(v), - _bd3(gate), - _bd3(beta), - _bd3(w), - _bd3(output), - None if h_for_descs is None else _bd4(h_for_descs), - cu_bd, - ws_bd, - cutlass.Int32(checkpoint_every_n_tokens), - cu_stream, - options="--enable-tvm-ffi", - ) - cache["build_descs"]( - q, - k, - v, - gate, - beta, - w, - output, - h_for_descs, - cu_seqlens, - tensormap_workspace, - checkpoint_every_n_tokens, + if cache.get("build_descs_has_h") != (h_for_descs is not None): + cache.pop("build_descs", None) + cache["build_descs_has_h"] = h_for_descs is not None + if "build_descs" not in cache: + io_dtype = _cutlass_io_dtype(q.dtype) + + def _bd3(t): + c = from_dlpack(t, assumed_align=16) + c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + return c + + def _bd4(t): + c = from_dlpack(t, assumed_align=16) + c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + return c + + cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + _build_descs, + io_dtype, + CFG.B_T, + _bd3(q), + _bd3(k), + _bd3(v), + _bd3(gate), + _bd3(beta), + _bd3(w), + _bd3(output), + None if h_for_descs is None else _bd4(h_for_descs), + cu_bd, + ws_bd, + cutlass.Int32(checkpoint_every_n_tokens), cu_stream, + options="--enable-tvm-ffi", ) - cache["desc_key"] = desc_key - cache["desc_cu"] = cu_seqlens - cache["desc_cu_versions"] = cu_versions - cache["desc_workspace_ptr"] = _data_ptr(tensormap_workspace) + cache["build_descs"]( + q, + k, + v, + gate, + beta, + w, + output, + h_for_descs, + cu_seqlens, + tensormap_workspace, + checkpoint_every_n_tokens, + cu_stream, + ) compiled( q, 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 faa4107f1..b78ebf950 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -5251,14 +5251,6 @@ def compile( ) -def _stream_capturing(stream) -> bool: - """True when ``stream`` is inside CUDA-graph capture.""" - from cuda.bindings import runtime as _rt - - err, status = _rt.cudaStreamIsCapturing(int(stream)) - return int(err) == 0 and int(status) != 0 - - def chunk_gdn_bwd_sm100( q, k, @@ -5421,71 +5413,46 @@ def _tok2(t): compiled = cache["compiled"] - # desc key: cu object identity + _version so address reuse forces a rebuild - desc_key = ( - _data_ptr(q), - _data_ptr(k), - _data_ptr(v), - _data_ptr(do), - _data_ptr(h), - _data_ptr(dq), - _data_ptr(dk), - _data_ptr(dv), - tuple(q.shape), - tuple(k.shape), - tuple(v.shape), - tuple(do.shape), - tuple(h.shape), - _data_ptr(initial_state) if initial_state is not None else None, - int(B), - ) - cu_versions = (getattr(cu_seqlens, "_version", 0),) - if ( - cache.get("desc_key") != desc_key - or cache.get("desc_cu") is not cu_seqlens - or cache.get("desc_cu_versions") != cu_versions - or cache.get("desc_workspace") is not workspace - or _stream_capturing(stream) - ): - if "build_descs" not in cache: - - def _tok3_bc(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - return c - - h_bc = from_dlpack(h, assumed_align=16) - h_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) - cu_bc = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() - s0_bc = None - if initial_state is not None: - s0_bc = from_dlpack(initial_state, assumed_align=16) - s0_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) - - ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() - cache["build_descs"] = cute.compile( - _build_descs, - io_dtype, - CFG.B_T, - _tok3_bc(q), - _tok3_bc(k), - _tok3_bc(v), - _tok3_bc(do), - _tok3_bc(dq), - _tok3_bc(dk), - _tok3_bc(dv), - h_bc, - cu_bc, - s0_bc, - ws_bc, - cu_stream, - options="--enable-tvm-ffi", - ) - cache["build_descs"](q, k, v, do, dq, dk, dv, h, cu_seqlens, initial_state, workspace, cu_stream) - cache["desc_key"] = desc_key - cache["desc_cu"] = cu_seqlens - cache["desc_cu_versions"] = cu_versions - cache["desc_workspace"] = workspace + # The descriptors encode cu_seqlens' CONTENTS, which no key built from the + # buffers can track. The skip this replaces asked torch's _version counter, + # so it was sound for a torch caller and silently stale for every other + # producer. Rebuilding unconditionally measures free: 131 vs 135 us of host + # time, and 157 either way once the launches are waited on. + if "build_descs" not in cache: + + def _tok3_bc(t): + c = from_dlpack(t, assumed_align=16) + c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + return c + + h_bc = from_dlpack(h, assumed_align=16) + h_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + cu_bc = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() + s0_bc = None + if initial_state is not None: + s0_bc = from_dlpack(initial_state, assumed_align=16) + s0_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + + ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + _build_descs, + io_dtype, + CFG.B_T, + _tok3_bc(q), + _tok3_bc(k), + _tok3_bc(v), + _tok3_bc(do), + _tok3_bc(dq), + _tok3_bc(dk), + _tok3_bc(dv), + h_bc, + cu_bc, + s0_bc, + ws_bc, + cu_stream, + options="--enable-tvm-ffi", + ) + cache["build_descs"](q, k, v, do, dq, dk, dv, h, cu_seqlens, initial_state, workspace, cu_stream) compiled( q, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py index 26de36c29..a19bc6864 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py @@ -2181,9 +2181,10 @@ def _build_descs( stream: cuda.CUstream, ): """Build the 5 per-(b,h) TMA-descriptor arrays (Q, K, V, O, S) into - ``tensormap_workspace``. Compiled + launched separately from the main - kernel and cached by input identity in the host bridge, so the builder - launches do not recur in steady-state replay. + ``tensormap_workspace``. Compiled and launched separately from the main + kernel, and 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 that CUDA-graph capture forbids. The H descriptor is 3-D ``(dv, dk, h)`` over the packed ``[total_h, HO, DK, DV]`` H tensor; ``build_h_descs_kernel`` derives the @@ -3199,14 +3200,6 @@ def compile( ) -def _stream_capturing(stream) -> bool: - """True when ``stream`` is inside CUDA-graph capture.""" - from cuda.bindings import runtime as _rt - - err, status = _rt.cudaStreamIsCapturing(int(stream)) - return int(err) == 0 and int(status) != 0 - - def chunk_gdn_sm100( q, k, @@ -3404,77 +3397,54 @@ def chunk_gdn_sm100( compiled = cache["compiled"] - # desc key: cu object identity + _version so address reuse forces a rebuild - desc_key = ( - _data_ptr(q), - _data_ptr(k), - _data_ptr(v), - _data_ptr(output) if enable_o else 0, - _data_ptr(output_h) if enable_h else 0, - tuple(q.shape), - tuple(k.shape), - tuple(v.shape), - tuple(output.shape) if enable_o else (), - tuple(output_h.shape) if enable_h else (), - int(B), - int(checkpoint_every_n_tokens) if enable_h else (), - ) - cu_versions = (getattr(cu_seqlens, "_version", 0),) - if ( - cache.get("desc_key") != desc_key - or cache.get("desc_cu") is not cu_seqlens - or cache.get("desc_cu_versions") != cu_versions - or cache.get("desc_workspace") is not workspace - or _stream_capturing(stream) - ): - if "build_descs" not in cache: - q_bc = from_dlpack(q, assumed_align=16) - q_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_bc = from_dlpack(k, assumed_align=16) - k_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_bc = from_dlpack(v, assumed_align=16) - v_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - o_bc = None - if enable_o: - o_bc = from_dlpack(output, assumed_align=16) - o_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - cu_bc = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() - s_bc = None - cu_ckpt_bc = None - if enable_h: - s_bc = from_dlpack(output_h, assumed_align=16) - s_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) - ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() - cache["build_descs"] = cute.compile( - _build_descs, - io_dtype, - CFG.B_T, - q_bc, - k_bc, - v_bc, - o_bc, - cu_bc, - s_bc, - cutlass.Int32(checkpoint_every_n_tokens if enable_h else 1), - ws_bc, - cu_stream, - options="--enable-tvm-ffi", - ) - cache["build_descs"]( - q, - k, - v, - output, - cu_seqlens, - output_h, - checkpoint_every_n_tokens if enable_h else 1, - workspace, + # The descriptors encode cu_seqlens' CONTENTS, which no key built from the + # buffers can track. The skip this replaces asked torch's _version counter, + # so it was sound for torch callers and silently stale for every other + # producer. Rebuilding unconditionally measured free: 131 vs 135 us of host + # time, 157 either way once the launches are waited on. + if "build_descs" not in cache: + q_bc = from_dlpack(q, assumed_align=16) + q_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + k_bc = from_dlpack(k, assumed_align=16) + k_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + v_bc = from_dlpack(v, assumed_align=16) + v_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + o_bc = None + if enable_o: + o_bc = from_dlpack(output, assumed_align=16) + o_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + cu_bc = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() + s_bc = None + if enable_h: + s_bc = from_dlpack(output_h, assumed_align=16) + s_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + _build_descs, + io_dtype, + CFG.B_T, + q_bc, + k_bc, + v_bc, + o_bc, + cu_bc, + s_bc, + cutlass.Int32(checkpoint_every_n_tokens if enable_h else 1), + ws_bc, cu_stream, + options="--enable-tvm-ffi", ) - cache["desc_key"] = desc_key - cache["desc_cu"] = cu_seqlens - cache["desc_cu_versions"] = cu_versions - cache["desc_workspace"] = workspace + cache["build_descs"]( + q, + k, + v, + output, + cu_seqlens, + output_h, + checkpoint_every_n_tokens if enable_h else 1, + workspace, + cu_stream, + ) compiled( q, 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 1d41a3d24..235d10e38 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -2646,14 +2646,6 @@ def _data_ptr(t) -> int: return t.__cuda_array_interface__["data"][0] -def _stream_capturing(stream) -> bool: - """True when ``stream`` is inside CUDA-graph capture.""" - from cuda.bindings import runtime as _rt - - err, status = _rt.cudaStreamIsCapturing(int(stream)) - return int(err) == 0 and int(status) != 0 - - def _cutlass_io_dtype(dtype): name = str(dtype) if "bfloat16" in name: @@ -2993,84 +2985,61 @@ def chunk_kda_sm100( compiled = cache["compiled"] - # desc key: buffer identity + cu _version so address reuse forces a rebuild + # The descriptors encode cu_seqlens' CONTENTS, which no key built from the + # buffers can track. The skip this replaces asked torch's _version counter, + # so it was sound for a torch caller and silently stale for every other + # producer. Rebuilding unconditionally measures free: 131 vs 135 us of host + # time, and 157 either way once the launches are waited on. h_for_descs = output_checkpoints if enable_checkpoints else None - desc_key = ( - _data_ptr(q), - _data_ptr(k), - _data_ptr(v), - _data_ptr(gate), - _data_ptr(output), - _data_ptr(h_for_descs) if h_for_descs is not None else 0, - checkpoint_every_n_tokens, - tuple(q.shape), - tuple(k.shape), - tuple(v.shape), - tuple(gate.shape), - tuple(output.shape), - tuple(h_for_descs.shape) if h_for_descs is not None else (), - ) - cu_versions = (getattr(cu_seqlens, "_version", 0),) - if ( - cache.get("desc_key") != desc_key - or cache.get("desc_cu") is not cu_seqlens - or cache.get("desc_cu_versions") != cu_versions - or cache.get("desc_workspace_ptr") != _data_ptr(tensormap_workspace) - or _stream_capturing(stream) - ): - if cache.get("build_descs_has_h") != (h_for_descs is not None): - cache.pop("build_descs", None) - cache["build_descs_has_h"] = h_for_descs is not None - if "build_descs" not in cache: - io_dtype = _cutlass_io_dtype(q.dtype) - q_bd = from_dlpack(q, assumed_align=16) - q_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_bd = from_dlpack(k, assumed_align=16) - k_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_bd = from_dlpack(v, assumed_align=16) - v_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_bd = from_dlpack(gate, assumed_align=16) - gate_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - o_bd = from_dlpack(output, assumed_align=16) - o_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() - h_bd = None - if h_for_descs is not None: - h_bd = from_dlpack(h_for_descs, assumed_align=16) - h_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) - cache["build_descs"] = cute.compile( - _build_descs, - io_dtype, - CFG.B_T, - q_bd, - k_bd, - v_bd, - gate_bd, - o_bd, - h_bd, - cu_bd, - ws_bd, - cutlass.Int32(checkpoint_every_n_tokens), - cu_stream, - options="--enable-tvm-ffi", - ) - cache["build_descs"]( - q, - k, - v, - gate, - output, - h_for_descs, - cu_seqlens, - tensormap_workspace, - checkpoint_every_n_tokens, + if cache.get("build_descs_has_h") != (h_for_descs is not None): + cache.pop("build_descs", None) + cache["build_descs_has_h"] = h_for_descs is not None + if "build_descs" not in cache: + io_dtype = _cutlass_io_dtype(q.dtype) + q_bd = from_dlpack(q, assumed_align=16) + q_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + k_bd = from_dlpack(k, assumed_align=16) + k_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + v_bd = from_dlpack(v, assumed_align=16) + v_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + gate_bd = from_dlpack(gate, assumed_align=16) + gate_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + o_bd = from_dlpack(output, assumed_align=16) + o_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + h_bd = None + if h_for_descs is not None: + h_bd = from_dlpack(h_for_descs, assumed_align=16) + h_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + cache["build_descs"] = cute.compile( + _build_descs, + io_dtype, + CFG.B_T, + q_bd, + k_bd, + v_bd, + gate_bd, + o_bd, + h_bd, + cu_bd, + ws_bd, + cutlass.Int32(checkpoint_every_n_tokens), cu_stream, + options="--enable-tvm-ffi", ) - cache["desc_key"] = desc_key - cache["desc_cu"] = cu_seqlens - cache["desc_cu_versions"] = cu_versions - cache["desc_workspace_ptr"] = _data_ptr(tensormap_workspace) + cache["build_descs"]( + q, + k, + v, + gate, + output, + h_for_descs, + cu_seqlens, + tensormap_workspace, + checkpoint_every_n_tokens, + cu_stream, + ) compiled( q, diff --git a/python/cudnn/nodes.py b/python/cudnn/nodes.py index c7cfea510..9b407604b 100644 --- a/python/cudnn/nodes.py +++ b/python/cudnn/nodes.py @@ -42,13 +42,6 @@ def __init__( self.outputs: Dict[str, Tensor] = {} self.params: Dict[str, Any] = {} - def __setattr__(self, name, value): - # attribute writes freeze with the owning graph; the port/param dicts - # themselves become MappingProxy views at freeze time - if getattr(self, "_frozen", False) and name != "_frozen": - raise RuntimeError(f"cannot set Node.{name}: the owning graph is frozen after lowering/planning") - object.__setattr__(self, name, value) - def validate(self) -> None: """Validate node configuration.""" for port_name, tensor in self.inputs.items(): diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 7712dcb6b..347fe2faa 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -38,6 +38,19 @@ ) +def dtype_name(buffer) -> str: + """The buffer's dtype as a bare name, whoever produced it. + + A caller buffer reaches these checks as whatever the graph normalized it + into, which is a variant-pack slot rather than a torch tensor. Comparing + ``buffer.dtype is torch.float32`` therefore rejects a perfectly good fp32 + buffer with "must be float32; got float32". Names are the one spelling + every producer agrees on -- torch prints ``torch.float32``, numpy and the + slot print ``float32``. + """ + return str(buffer.dtype).rsplit(".", 1)[-1] + + def _require_reciprocal_s_scales(descale_s: float, scale_s: float) -> None: """Guard for a kernel that converts P to e4m3 UNSCALED (the SM100 FP8 row). @@ -449,7 +462,7 @@ def _amax_slot(self, tensor, name: str, device: torch.device) -> torch.Tensor: # Both callers re-view the slot as Int32 for the kernel ABI, so a wider # element would yield two int32s and the kernel would write only the low # word -- the caller then reads a corrupted value. - if tensor.dtype is not torch.float32: + if dtype_name(tensor) != "float32": raise ValueError(f"{name} must be float32; got {tensor.dtype}") try: return tensor.view(-1)[:1] @@ -474,7 +487,7 @@ def _checked_lse_view(self, lse_tensor: torch.Tensor) -> torch.Tensor: receive the output and be dropped, leaving the caller's LSE unwritten. """ self._value_error_if( - lse_tensor.dtype != torch.float32, + dtype_name(lse_tensor) != "float32", f"lse_tensor must be float32; got {lse_tensor.dtype}", ) expected = self.batch_size * self.h_q * self.s_q_max @@ -496,7 +509,7 @@ def _checked_sinks_1d(self, sinks: torch.Tensor) -> torch.Tensor: on the execute hot path (and break CUDA-graph pointer stability). """ self._value_error_if( - sinks.dtype != torch.float32, + dtype_name(sinks) != "float32", f"sinks must be float32; got {sinks.dtype}", ) self._value_error_if( @@ -517,7 +530,7 @@ def _checked_seq_lens(self, seq_lens: torch.Tensor, name: str) -> torch.Tensor: pointer stability). """ self._value_error_if( - seq_lens.dtype != torch.int32, + dtype_name(seq_lens) != "int32", f"{name} must be int32; got {seq_lens.dtype}", ) self._value_error_if( diff --git a/python/pycudnn.cpp b/python/pycudnn.cpp index b334dc7d2..95c85e10d 100644 --- a/python/pycudnn.cpp +++ b/python/pycudnn.cpp @@ -143,6 +143,10 @@ create_kernel_cache_submodule(py::module_ &); void init_properties(py::module_ &); +// pybinds for the native variant pack +void +init_variant_pack(py::module_ &); + void set_dlhandle_cudnn(std::intptr_t dlhandle) { #ifdef _WIN32 @@ -158,6 +162,7 @@ PYBIND11_MODULE(_compiled_module, m) { init_properties(m); init_pygraph_submodule(m); + init_variant_pack(m); m.def("_set_dlhandle_cudnn", &set_dlhandle_cudnn); diff --git a/python/pygraph/pygraph.cpp b/python/pygraph/pygraph.cpp index f56181eb2..b0e08c2bd 100644 --- a/python/pygraph/pygraph.cpp +++ b/python/pygraph/pygraph.cpp @@ -736,31 +736,22 @@ PyGraph::prepare_variant_pack_template() { throw_if(status.is_bad(), status.get_code(), status.get_message()); } -void -PyGraph::execute_with_ptrs(std::vector const& user_ptrs, - std::intptr_t workspace, - std::intptr_t exec_handle) { - std::vector ptrs(user_ptrs.size()); - for (size_t i = 0; i < user_ptrs.size(); i++) { - ptrs[i] = (void*)user_ptrs[i]; - } - cudnnHandle_t h = exec_handle ? static_cast((void*)exec_handle) : handle; - auto status = graph->execute(h, ptrs.data(), (int)ptrs.size(), (void*)workspace); - throw_if(status.is_bad(), status.get_code(), status.get_message()); -} - void PyGraph::execute_with_raw_ptrs(std::intptr_t user_ptrs_array, int64_t n_user, std::intptr_t workspace, - std::intptr_t exec_handle) { + std::intptr_t exec_handle, + int64_t plan_index) { static_assert(sizeof(std::intptr_t) == sizeof(void*), "intptr_t and void* must be the same size"); throw_if(n_user < 0, error_code_t::INVALID_VALUE, "n_user must be non-negative"); throw_if(n_user > 0 && user_ptrs_array == 0, error_code_t::INVALID_VALUE, "user_ptrs_array is null"); // user_ptrs_array points to a contiguous intptr_t[] of device pointers — zero copy void** ptrs = reinterpret_cast(user_ptrs_array); cudnnHandle_t h = exec_handle ? static_cast((void*)exec_handle) : handle; - auto status = graph->execute(h, ptrs, (int)n_user, (void*)workspace); + // Address the plan the python walk built, which stops being the graph's own + // candidate once the walk has skipped an entry. -1 defers to the candidate. + auto status = plan_index < 0 ? graph->execute(h, ptrs, (int)n_user, (void*)workspace) + : graph->execute_plan_at_index(h, ptrs, (int)n_user, (void*)workspace, plan_index); throw_if(status.is_bad(), status.get_code(), status.get_message()); } @@ -1368,17 +1359,13 @@ init_pygraph_submodule(py::module_& m) { py::arg("override_strides") = py::none()) .def("_prepare_variant_pack_template", &PyGraph::prepare_variant_pack_template) .def("_get_variant_pack_uids_sorted", &PyGraph::get_variant_pack_uids_sorted) - .def("_execute_with_ptrs", - &PyGraph::execute_with_ptrs, - py::arg("user_ptrs"), - py::arg("workspace"), - py::arg("handle")) .def("_execute_with_raw_ptrs", &PyGraph::execute_with_raw_ptrs, py::arg("user_ptrs_array"), py::arg("n_user"), py::arg("workspace"), - py::arg("handle")) + py::arg("handle"), + py::arg("plan_index") = -1) .def("populate_cuda_graph", &PyGraph::populate_cuda_graph) .def("update_cuda_graph", &PyGraph::update_cuda_graph) .def("serialize", &PyGraph::serialize) diff --git a/python/pygraph/pygraph.h b/python/pygraph/pygraph.h index 63f557aa6..99951832b 100644 --- a/python/pygraph/pygraph.h +++ b/python/pygraph/pygraph.h @@ -730,15 +730,16 @@ class PyGraph { return graph->get_variant_pack_uids_sorted(); } - void - execute_with_ptrs(std::vector const& user_ptrs, std::intptr_t workspace, std::intptr_t exec_handle); - - // Raw pointer version: takes a pointer to an array of device pointers (no pybind11 copy) + // Takes a pointer to a contiguous array of device pointers, ordered as + // get_variant_pack_uids_sorted() reports — no pybind11 container copy and no + // per-operand uid->pointer hash lookup. ``plan_index`` selects the plan; + // -1 means the graph's own candidate. void execute_with_raw_ptrs(std::intptr_t user_ptrs_array, int64_t n_user, std::intptr_t workspace, - std::intptr_t exec_handle); + std::intptr_t exec_handle, + int64_t plan_index); std::vector get_behavior_notes(); diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp new file mode 100644 index 000000000..8ef85175c --- /dev/null +++ b/python/pygraph/variant_pack.cpp @@ -0,0 +1,841 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// The variant pack, held as DLTensors rather than python objects. +// +// `__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 +// 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 +// its parts. +#include "variant_pack.h" + +#include +#include +#include +#include + +#include +#include + +#include "dlpack/dlpack.h" + +namespace py = pybind11; + +namespace cudnn_frontend { +namespace python_bindings { + +namespace { + +// The vtable is a property of the type, and the DLPack docs tell consumers to +// cache it per type. A handful of buffer types occur in one process. +constexpr int kTypeCacheSlots = 8; + +struct TypeCache { + PyTypeObject *types[kTypeCacheSlots] = {}; + DLPackExchangeAPI *apis[kTypeCacheSlots] = {}; + int count = 0; +}; + +TypeCache & +type_cache() { + static TypeCache cache; + return cache; +} + +// The newest table in the producer's chain whose layout matches the one this +// was compiled against, or null. +// +// The protocol requires this walk: a table is only safe to call when its major +// version is ours, and a producer that has moved on keeps older tables reachable +// through prev_api for exactly this reason. Skipping the check would mean +// calling function pointers at offsets that a future major version is free to +// move -- a crash years later, in code that had been correct all along. +DLPackExchangeAPI * +compatible_api(DLPackExchangeAPI *api) { + for (int hops = 0; api != nullptr && hops < 8; hops++) { + if (api->header.version.major == DLPACK_MAJOR_VERSION) return api; + api = reinterpret_cast(api->header.prev_api); + } + return nullptr; +} + +// The producer's exchange vtable, or null when its type does not implement the +// protocol at a version we speak. A missing attribute is the common case for +// older frameworks, not an error, so the python exception it raises is +// swallowed and the caller falls back to reading the buffer from python. +DLPackExchangeAPI * +exchange_api_for(PyObject *obj) { + PyTypeObject *type = Py_TYPE(obj); + TypeCache &cache = type_cache(); + for (int i = 0; i < cache.count; i++) { + if (cache.types[i] == type) return cache.apis[i]; + } + PyObject *capsule = PyObject_GetAttrString(reinterpret_cast(type), "__dlpack_c_exchange_api__"); + DLPackExchangeAPI *api = nullptr; + if (capsule == nullptr) { + PyErr_Clear(); + } else { + api = compatible_api(static_cast(PyCapsule_GetPointer(capsule, "dlpack_exchange_api"))); + Py_DECREF(capsule); + if (api == nullptr) PyErr_Clear(); + } + // Only a hit is cached. A type can acquire the vtable AFTER we first look: + // on torch builds without it natively, tvm-ffi installs one when it is + // imported, and a graph normalized before that import would otherwise be + // pinned to the python fallback for the life of the process. + // + // Keyed on the type's ADDRESS, so the entry must own a reference: a heap + // type that got collected could be replaced by a different type allocated + // at the same address, and this would hand out its vtable. The cache never + // evicts, so this pins at most kTypeCacheSlots types. + if (api != nullptr && cache.count < kTypeCacheSlots) { + Py_INCREF(type); + cache.types[cache.count] = type; + cache.apis[cache.count] = api; + cache.count++; + } + return api; +} + +// The name each DLPack (code, bits) travels under in the kernels' vocabulary, +// which is torch's spelling minus the "torch." prefix. +std::string +dtype_name(DLDataType dtype) { + const int code = dtype.code; + const int bits = dtype.bits; + if (code == kDLFloat) { + if (bits == 16) return "float16"; + if (bits == 32) return "float32"; + if (bits == 64) return "float64"; + } else if (code == kDLBfloat && bits == 16) { + return "bfloat16"; + } else if (code == kDLInt) { + if (bits == 8) return "int8"; + if (bits == 32) return "int32"; + if (bits == 64) return "int64"; + } else if (code == kDLUInt && bits == 8) { + return "uint8"; + } else if (code == kDLBool) { + return "bool"; + } else if (code == kDLFloat8_e4m3fn) { + return "float8_e4m3fn"; + } else if (code == kDLFloat8_e5m2) { + return "float8_e5m2"; + } else if (code == kDLFloat8_e8m0fnu) { + return "float8_e8m0fnu"; + } + return "code" + std::to_string(code) + "_" + std::to_string(bits); +} + +bool +is_dense(const DLTensor &t) { + if (t.strides == nullptr) return true; // compact by definition + int64_t expect = 1; + for (int d = t.ndim - 1; d >= 0; d--) { + if (t.shape[d] != 1 && t.strides[d] != expect) return false; + expect *= t.shape[d]; + } + return true; +} + +// 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 { + void *data = nullptr; + int32_t ndim = 0; + DLDataType dtype = {0, 0, 1}; + std::vector shape; + std::vector stride; // empty means compact row-major + bool filled = false; +}; + +} // namespace + +// A pack's slot, 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 { + public: + VariantPackSlot(Slot slot, int32_t device_id) : slot_(std::move(slot)) { + tensor_.data = slot_.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_.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 + // 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; + + const DLTensor & + tensor() const { + return tensor_; + } + + int64_t + data_ptr() const { + return reinterpret_cast(tensor_.data); + } + + std::vector + shape() const { + return slot_.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]; + 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) + throw py::index_error("stride(): dimension " + std::to_string(dim) + " is out of range for a " + + std::to_string(slot_.ndim) + "-D slot"); + return stride()[axis]; + } + + // The bare dtype NAME, which is what a kernel means when it asks a buffer + // for its dtype: they all reach it through str(x.dtype).split(".")[-1], so + // a torch tensor's "torch.bfloat16" and this "bfloat16" answer the same. + std::string + dtype() const { + return dtype_name(slot_.dtype); + } + + int64_t + element_size() const { + return slot_.dtype.bits / 8; + } + + int64_t + numel() const { + int64_t n = 1; + for (int64_t extent : slot_.shape) n *= extent; + return n; + } + + int64_t + nbytes() const { + return numel() * element_size(); + } + + int64_t + length() const { + return slot_.shape.empty() ? 0 : slot_.shape[0]; + } + + py::tuple + dlpack_device() const { + return py::make_tuple(static_cast(kDLCUDA), tensor_.device.device_id); + } + + // 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 + // than silently reinterpreted. + VariantPackSlot * + reshape(std::vector shape) const { + if (!slot_.stride.empty() && !is_dense(tensor_)) + throw py::value_error("cannot reshape a non-contiguous variant-pack slot"); + int64_t numel = 1; + for (int64_t extent : slot_.shape) numel *= extent; + int64_t fixed = 1; + int wildcard = -1; + for (size_t d = 0; d < shape.size(); d++) { + if (shape[d] == -1) { + if (wildcard >= 0) throw py::value_error("reshape accepts at most one -1"); + wildcard = static_cast(d); + } else { + fixed *= shape[d]; + } + } + if (wildcard >= 0) { + if (fixed == 0 || numel % fixed != 0) + throw py::value_error("cannot reshape " + std::to_string(numel) + " elements to the requested shape"); + shape[wildcard] = numel / fixed; + } 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()); + out.stride.clear(); // dense by construction, as the reshape required + return new VariantPackSlot(out, tensor_.device.device_id); + } + + // The same memory with its axes relabelled; exact for a strided slot too. + VariantPackSlot * + 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"); + 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]; + } else { + out.stride = slot_.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"); + seen[axis] = true; + out.shape[d] = slot_.shape[axis]; + out.stride[d] = from_stride[axis]; + } + return new VariantPackSlot(std::move(out), tensor_.device.device_id); + } + + // Row-major contiguous by construction, so this is the identity a caller + // written against a framework tensor expects to be able to call. + py::object + contiguous(py::object self) const { + return self; + } + + // The capsule form of the same tensor, for a consumer that does not read + // the exchange vtable -- cute's from_dlpack at compile time is the one that + // matters here. It costs an allocation where the vtable costs none, which + // is why it is not what the per-launch path uses. + // + // 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. + py::capsule + dlpack(py::object /*stream*/, py::object /*max_version*/) const { + struct Owned { + DLManagedTensor managed; + std::vector shape; + std::vector stride; + }; + auto *owned = new Owned{{}, slot_.shape, slot_.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(); + owned->managed.manager_ctx = owned; + owned->managed.deleter = [](DLManagedTensor *self) { delete static_cast(self->manager_ctx); }; + return py::capsule(&owned->managed, "dltensor", [](PyObject *capsule) { + // only reached when nobody consumed it: a consumer renames the + // capsule to "used_dltensor" and takes the deleter over + if (PyCapsule_IsValid(capsule, "dltensor")) { + auto *managed = static_cast(PyCapsule_GetPointer(capsule, "dltensor")); + if (managed != nullptr && managed->deleter != nullptr) managed->deleter(managed); + } + }); + } + + private: + Slot slot_; // 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(); + return 0; +} + +int +slot_managed_from_py_object(void *py_object, DLManagedTensorVersioned **out) { + auto *slot = 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 + // 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 &tensor = owned->versioned.dl_tensor; + tensor = slot->tensor(); + tensor.shape = owned->shape.empty() ? nullptr : owned->shape.data(); + // a slot 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; + owned->versioned.manager_ctx = owned; + owned->versioned.deleter = [](DLManagedTensorVersioned *self) { delete static_cast(self->manager_ctx); }; + *out = &owned->versioned; + return 0; +} + +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"); + return -1; +} + +int +slot_to_py_object(DLManagedTensorVersioned *, void **) { + PyErr_SetString(PyExc_NotImplementedError, "a variant-pack slot is not an importer"); + return -1; +} + +// The graph launches on the stream its handle carries, which execute() passes +// 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) { + *out_stream = nullptr; + return 0; +} + +DLPackExchangeAPI & +slot_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; + return table; + }(); + return api; +} + +} // namespace + +class VariantPackNative { + public: + explicit VariantPackNative(size_t n) : slots_(n), pointers_(n, nullptr) {} + + // Fill one slot 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); + 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); + if (t.strides != nullptr) { + slot.stride.assign(t.strides, t.strides + t.ndim); + } else { + slot.stride.clear(); + } + slot.filled = true; + pointers_[index] = slot.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 + // 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. + // 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. + std::vector + 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++) { + 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))) { + unread.push_back(i); + } + } + return unread; + } + + // The first slot 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); + } + 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. + // 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"); + } + // 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)); + } + slot.ndim = static_cast(shape.size()); + slot.shape = std::move(shape); + slot.stride = std::move(stride); + } + + void + skip_slot(size_t index) { + slots_.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; + int64_t expect = 1; + for (int d = slot.ndim - 1; d >= 0; d--) { + if (slot.shape[d] != 1 && slot.stride[d] != expect) { + offender = std::to_string(i); + return false; + } + expect *= slot.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; + 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]; + } + return true; + } + + bool + is_filled(size_t index) const { + return slots_.at(index).filled; + } + + int64_t + pointer(size_t index) const { + return reinterpret_cast(pointers_.at(index)); + } + + std::vector + shape(size_t index) const { + return slots_.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]; + 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); + } + + // 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. + int64_t + pointer_array(void) const { + return reinterpret_cast(pointers_.data()); + } + + size_t + size(void) const { + return slots_.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; + out.reserve(indices.size()); + for (size_t index : indices) out.push_back(view(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); + } + + private: + std::vector slots_; + std::vector pointers_; +}; + +// A slot 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); +} + +// (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 +// carves against it. +py::object +read_buffer_extent(py::handle buffer) { + DLPackExchangeAPI *api = exchange_api_for(buffer.ptr()); + if (api == nullptr || api->dltensor_from_py_object_no_sync == nullptr) return py::none(); + DLTensor t{}; + if (api->dltensor_from_py_object_no_sync(buffer.ptr(), &t) != 0) throw py::error_already_set(); + // A byte count is only a byte RANGE when the buffer is dense; a carve + // bounds-checked against a strided one would write outside the allocation. + if (!is_dense(t)) return py::none(); + int64_t numel = 1; + for (int d = 0; d < t.ndim; d++) numel *= t.shape[d]; + const int64_t itemsize = (static_cast(t.dtype.bits) * t.dtype.lanes + 7) / 8; + return py::make_tuple(reinterpret_cast(static_cast(t.data) + t.byte_offset), numel * itemsize); +} + +// A workspace carve, planned once: the regions are fixed when the engine +// builds and only the base pointer arrives per execute. +class WorkspaceCarve { + public: + WorkspaceCarve(std::string owner, const std::vector ®ions) : owner_(std::move(owner)) { + protos_.reserve(regions.size()); + offsets_.reserve(regions.size()); + ends_.reserve(regions.size()); + for (const py::tuple ®ion : regions) { + if (region.size() != 4) { + throw py::value_error("a carve region is (offset, dtype_code, dtype_bits, shape)"); + } + int64_t offset = region[0].cast(); + Slot proto; + proto.dtype = DLDataType{region[1].cast(), region[2].cast(), 1}; + proto.shape = region[3].cast>(); + proto.ndim = static_cast(proto.shape.size()); + proto.filled = true; // stride left empty: a carve is dense by construction + int64_t numel = 1; + for (int64_t extent : proto.shape) numel *= extent; + offsets_.push_back(offset); + ends_.push_back(offset + numel * ((proto.dtype.bits + 7) / 8)); + protos_.push_back(std::move(proto)); + } + } + + std::vector + carve(int64_t base, int64_t nbytes, int32_t device_id) const { + std::vector out; + out.reserve(protos_.size()); + for (size_t i = 0; i < protos_.size(); i++) { + if (ends_[i] > nbytes) { + throw py::value_error(owner_ + ": workspace overrun -- region [" + std::to_string(offsets_[i]) + ", " + + std::to_string(ends_[i]) + ") exceeds the " + std::to_string(nbytes) + + "-byte buffer (sizing bug)"); + } + Slot slot = protos_[i]; + slot.data = reinterpret_cast(base + offsets_[i]); + out.push_back(new VariantPackSlot(std::move(slot), device_id)); + } + return out; + } + + size_t + size() const { + return protos_.size(); + } + + private: + std::string owner_; + std::vector protos_; + std::vector offsets_; + std::vector ends_; +}; + +void +init_variant_pack(py::module_ &m) { + auto slot_class = py::class_(m, "VariantPackSlot", 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()); + + // 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); + if (capsule == nullptr) throw py::error_already_set(); + if (PyObject_SetAttrString(slot_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, + py::arg("ptr"), + py::arg("shape"), + py::arg("dtype_code"), + py::arg("dtype_bits"), + py::arg("device_id"), + "A DLPack producer over memory the caller did not supply -- a workspace carve."); + + m.def("read_buffer_extent", + &read_buffer_extent, + py::arg("buffer"), + "(pointer, bytes) through the exchange vtable, or None when the type does not publish one."); + + py::class_(m, "WorkspaceCarve", R"( +A workspace carve compiled once, at build. + +Regions are ``(offset, dtype_code, dtype_bits, shape)``. Only the base pointer +arrives per execute, so ``carve`` hands back every region in one crossing +instead of one per region. +)") + .def(py::init>(), py::arg("owner"), py::arg("regions")) + .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"); + + py::class_(m, "VariantPackNative", R"( +The caller's operands, held as DLTensors. + +``read_slot`` 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 +its parts. +)") + .def(py::init()) + .def("read_slot", &VariantPackNative::read_slot) + .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("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_property_readonly("address", &VariantPackNative::pointer_array) + .def("__len__", &VariantPackNative::size) + .def("all_contiguous", [](const VariantPackNative &self) { + std::string offender; + bool ok = self.all_contiguous(offender); + return py::make_tuple(ok, offender); + }); +} + +} // namespace python_bindings +} // namespace cudnn_frontend diff --git a/python/pygraph/variant_pack.h b/python/pygraph/variant_pack.h new file mode 100644 index 000000000..5083b6a6e --- /dev/null +++ b/python/pygraph/variant_pack.h @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +namespace cudnn_frontend { +namespace python_bindings { + +void +init_variant_pack(pybind11::module_ &); + +} // namespace python_bindings +} // namespace cudnn_frontend diff --git a/test/python/gemm/frost/test_frontend_integration.py b/test/python/gemm/frost/test_frontend_integration.py index c76aad88d..343ea351a 100644 --- a/test/python/gemm/frost/test_frontend_integration.py +++ b/test/python/gemm/frost/test_frontend_integration.py @@ -509,3 +509,57 @@ def test_wrapper_graph_path(): g({A: a, B: b, bias: bias_t, Y: y}) torch.cuda.synchronize() torch.testing.assert_close(y, ref, atol=1e-1, rtol=1e-2) + + +@_GPU +def test_override_shape_inside_a_max_allocation_matches_the_backend(): + """The case override shape is FOR: allocate once at a cache shape, name the + live shape per call. + + A caller does not pick the plan, so the two paths have to answer the same + question. FROST reads its M/N/K off the operands, so this only works if + ``execute`` applies the overrides to what it hands the engine -- and in the + axis order the operand already uses, since ``override_shapes`` speaks the + graph's declaration (B is ``[batch, K, N]``) while the buffer is allocated + ``(batch, N, K)``. Getting that wrong reads N as K and the kernel rejects + the launch, or worse runs the whole allocation. + """ + mb, nb, kb = 256, 256, 128 + m, n, k = 128, 192, 64 + a = torch.empty(1, mb, kb, dtype=torch.int32).random_(-2, 2).to(torch.bfloat16).cuda() + b = torch.empty(1, nb, kb, dtype=torch.int32).random_(-2, 2).to(torch.bfloat16).cuda() + ref = torch.einsum("bmk,bnk->bmn", a[:, :m, :k].float(), b[:, :n, :k].float()).to(torch.bfloat16) + + uids = [1, 2, 3] + shapes = [[1, m, k], [1, k, n], [1, m, n]] + strides = [[mb * kb, kb, 1], [nb * kb, 1, kb], [mb * nb, nb, 1]] + + results = {} + for want_frost in (False, True): + g = cudnn.pygraph( + io_data_type=cudnn.data_type.BFLOAT16, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + is_dynamic_shape_enabled=True, + is_override_shape_enabled=True, + ) + A = g.tensor(name="A", uid=1, dim=[1, mb, kb], stride=[mb * kb, kb, 1], data_type=cudnn.data_type.BFLOAT16) + B = g.tensor(name="B", uid=2, dim=[1, kb, nb], stride=[kb * nb, 1, kb], data_type=cudnn.data_type.BFLOAT16) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(cudnn.data_type.BFLOAT16).set_uid(3) + _plan(g) + index = _index_of(g, _FROST) if want_frost else _first_backend_index(g) + g.select_plan(index) + g.check_support() + g.build_plans() + + h = cudnn.create_handle() + wsz = g.get_workspace_size_plan_at_index(index, h, uids, shapes, strides) + ws = torch.empty(max(wsz, 1), device="cuda", dtype=torch.uint8) + c = torch.zeros(1, mb, nb, dtype=torch.bfloat16, device="cuda") + g.execute({1: a, 2: b, 3: c}, ws, handle=h, override_uids=uids, override_shapes=shapes, override_strides=strides) + torch.cuda.synchronize() + results[g.get_plan_name_at_index(index)] = c[:, :m, :n].clone() + + for name, got in results.items(): + torch.testing.assert_close(got, ref, atol=0, rtol=0, msg=lambda s, name=name: f"{name} ran the wrong shape\n{s}") diff --git a/test/python/test_dlpack_export.py b/test/python/test_dlpack_export.py new file mode 100644 index 000000000..db607ad7d --- /dev/null +++ b/test/python/test_dlpack_export.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A DeviceView exports DLPack through a slot, which owns the struct it hands out. + +The view itself has nowhere to put a ``DLManagedTensor`` that must outlive the +capsule: cute's ``from_dlpack`` aliases the struct rather than copying it, and +the view cannot know when the consumer is done. Holding the struct on the view +behind a no-op deleter — which is what this did — reads freed memory as soon as +a consumer outlives the view. +""" + +import ctypes +import gc + +import pytest +import torch + +from cudnn.frost import buffers + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="DLPack export is over device pointers") + + +def _capsule_address(capsule): + fn = ctypes.pythonapi.PyCapsule_GetPointer + fn.restype = ctypes.c_void_p + fn.argtypes = [ctypes.py_object, ctypes.c_char_p] + return fn(capsule, b"dltensor") + + +@pytest.mark.L0 +def test_capsule_decodes_to_the_right_buffer(): + t = torch.arange(24, dtype=torch.float32, device="cuda").reshape(2, 3, 4) + view = buffers.DeviceView(t.data_ptr(), (2, 3, 4), "float32", t.device.index or 0) + back = torch.from_dlpack(view) + torch.testing.assert_close(back, t) + + +@pytest.mark.L0 +def test_two_capsules_from_one_view_do_not_share_a_struct(): + """cute aliases the struct, so a shared one would be rewritten under it.""" + t = torch.zeros(8, dtype=torch.float32, device="cuda") + view = buffers.DeviceView(t.data_ptr(), (8,), "float32", t.device.index or 0) + a, b = view.__dlpack__(), view.__dlpack__() + assert _capsule_address(a) != _capsule_address(b) + + +@pytest.mark.L0 +def test_capsule_outlives_the_view_it_came_from(): + """The regression: the struct belongs to the capsule, not to the view. + + An unconsumed capsule used to point at a struct the view held in a list, + so dropping the view left the consumer decoding freed memory. + """ + t = torch.arange(16, dtype=torch.float32, device="cuda") + view = buffers.DeviceView(t.data_ptr(), (16,), "float32", t.device.index or 0) + capsule = view.__dlpack__() + del view + gc.collect() + torch.testing.assert_close(torch.from_dlpack(capsule), t) + + +@pytest.mark.L0 +@pytest.mark.parametrize("dtype,torch_dtype", [("float32", torch.float32), ("bfloat16", torch.bfloat16), ("int32", torch.int32), ("uint8", torch.uint8)]) +def test_dtypes_round_trip(dtype, torch_dtype): + t = torch.ones(6, dtype=torch_dtype, device="cuda") + view = buffers.DeviceView(t.data_ptr(), (6,), dtype, t.device.index or 0) + back = torch.from_dlpack(view) + assert back.dtype is torch_dtype and back.shape == (6,) + torch.testing.assert_close(back, t) diff --git a/test/python/test_graph_native.py b/test/python/test_graph_native.py index 34524670c..9255ea243 100644 --- a/test/python/test_graph_native.py +++ b/test/python/test_graph_native.py @@ -574,9 +574,20 @@ def execute(self, graph, tensor_data, ctx=None): mutate() def test_freeze_covers_public_surface(self, monkeypatch): - """Review round 5: the freeze must close EVERY public mutation path, - not only the fluent API — attribute writes, live containers, in-place - list edits, node params, and graph context.""" + """Freezing is ONE flag, on the graph. + + Every mutation route the public API offers runs through _check_mutable, + so the flag alone guards them. What the caller could otherwise change + behind the API's back is made immutable in its own right — dim/stride + become tuples, the port and param dicts become MappingProxy views — + rather than watched by a per-object __setattr__ guard. + + Those guards used to exist on Tensor, Node and GraphContext. They cost + a python-level call per field on EVERY construction (measured 2.1 us of + a 2.8 us Tensor, and a graph builds one per tensor and per node) to + catch a write that bypasses the setters anyway. Assigning + `t.dim = [...]` directly on a frozen graph is therefore no longer an + error; it is also not something the API asks anyone to do.""" from cudnn.engines import BaseEngine from test_dispatch import _FAKE, _offer @@ -595,18 +606,18 @@ def execute(self, graph, tensor_data, ctx=None): g.create_execution_plans() node = g.nodes[0] + # the setters and op builders — every mutation route the API offers — are closed with pytest.raises(RuntimeError, match="frozen"): - A.dim = [9, 9] # direct attribute write + A.set_dim([9, 9]) + with pytest.raises(RuntimeError, match="frozen"): + g.matmul(A, C) + # and the containers are immutable in their own right, at no per-call cost with pytest.raises(TypeError): A.dim[:] = [9] # sealed to a tuple: no in-place edits with pytest.raises(TypeError): node.params["padding"] = 123 # MappingProxy with pytest.raises(TypeError): node.inputs["A"] = C # MappingProxy - with pytest.raises(RuntimeError, match="frozen"): - node.inputs = {} # attribute write on the node - with pytest.raises(RuntimeError, match="frozen"): - g.context.compute_data_type = "HALF" # graph context # live-container laundering: the public views are copies g.nodes.clear() g.tensors.clear() diff --git a/test/python/test_import_boundaries.py b/test/python/test_import_boundaries.py index ac7d23dc4..3336d5b97 100644 --- a/test/python/test_import_boundaries.py +++ b/test/python/test_import_boundaries.py @@ -55,10 +55,12 @@ def _assert_absent(mods: set, stage: str) -> None: assert not present, f"{stage} imported {present}; it must not" +@pytest.mark.L0 def test_importing_cudnn_pulls_no_framework(): _assert_absent(_imported_by("import cudnn"), "import cudnn") +@pytest.mark.L0 def test_describing_a_graph_pulls_no_framework(): """Build and validate an SDPA graph through the graph API alone.""" _assert_absent( @@ -79,6 +81,7 @@ def test_describing_a_graph_pulls_no_framework(): ) +@pytest.mark.L0 def test_classification_and_facts_pull_no_framework(): """The manifest classifies and the analyzer describes; neither lowers.""" _assert_absent( @@ -87,6 +90,7 @@ def test_classification_and_facts_pull_no_framework(): ) +@pytest.mark.L0 @pytest.mark.parametrize("module", ["cudnn.sdpa.fwd.engines", "cudnn.sdpa.bwd.engines"]) def test_support_check_pulls_no_framework(module): """Capabilities and mismatch() are pure data and comparisons. diff --git a/test/python/test_variant_pack_normalization.py b/test/python/test_variant_pack_normalization.py new file mode 100644 index 000000000..7b03df5f5 --- /dev/null +++ b/test/python/test_variant_pack_normalization.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""graph.execute() normalizes the variant pack once, and both paths see the same thing. + +The property under test is that a caller never has to know whether the +heuristics landed the graph on the cuDNN backend or on a python engine. Before +normalization the backend accepted a bare device address (its `_ptr` had +`if type(d) is int: return d`) while a python engine did not — `resolve_node_buffers` +handed the engine the caller's object untouched and `frost.buffers.probe` then +raised "buffer of type int exposes neither __cuda_array_interface__ nor +__dlpack__". Same call, two answers, and the caller does not pick the plan. +""" + +import threading + +import pytest +import torch + +import cudnn + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="execute() needs a device to run a plan on") + +M = N = K = 64 + + +def _matmul_graph(): + """A graph the cuDNN backend serves.""" + a = torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, K, N, dtype=torch.bfloat16, device="cuda") + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + g = cudnn.pygraph(io_data_type=cudnn.data_type.BFLOAT16, compute_data_type=cudnn.data_type.FLOAT) + A, B = g.tensor_like(a), g.tensor_like(b) + C = g.matmul(A=A, B=B) + C.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + g.build_plans() + return g, {A: a, B: b, C: c}, (a, b, c) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "form", + ["tensor_keys", "uid_keys", "int_values", "int_values_and_workspace"], +) +def test_every_variant_pack_form_still_works(form): + """The four shapes a variant pack has always been allowed to take.""" + g, vp, (a, b, c) = _matmul_graph() + handle = cudnn.create_handle() + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + ws_arg = ws + + if form == "uid_keys": + vp = {t.get_uid(): v for t, v in vp.items()} + elif form == "int_values": + vp = {t: v.data_ptr() for t, v in vp.items()} + elif form == "int_values_and_workspace": + vp = {t.get_uid(): v.data_ptr() for t, v in vp.items()} + ws_arg = ws.data_ptr() + + g.execute(vp, ws_arg, handle=handle) + torch.cuda.synchronize() + ref = (a.float() @ b.float()).to(torch.bfloat16) + torch.testing.assert_close(c, ref, atol=0.2, rtol=0.05) + + +@pytest.mark.L0 +def test_missing_operand_names_the_tensor(): + g, vp, _ = _matmul_graph() + handle = cudnn.create_handle() + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + partial = dict(list(vp.items())[:-1]) + with pytest.raises(ValueError, match="missing a buffer for tensor uid"): + g.execute(partial, ws, handle=handle) + + +@pytest.mark.L0 +def test_operand_order_is_the_backend_order(): + """The layout is the backend's own, ascending by uid — not a python guess. + + A walk over node ports cannot produce it: a tensor's ragged_offset is a + user operand but hangs off the Tensor rather than off a port, and the slots + the graph fills itself (pass-by-value scalars, slice replacement + destinations, workspace modifications) must be excluded. + """ + g, _, _ = _matmul_graph() + order = g._variant_pack_uids() + assert order == sorted(order), f"not ascending: {order}" + assert order == list(g._lowered_graph._get_variant_pack_uids_sorted()) + + +@pytest.mark.L0 +def test_execute_is_reentrant(): + """One built graph, many threads, each with its own buffers. + + The pointer array is per call for this reason. Sharing one across calls + hands each thread the other's pointers — silently, because every pointer in + it is individually valid, so the failure is a wrong number and not a raise. + """ + g, _, _ = _matmul_graph() + handle = cudnn.create_handle() + uids = g._variant_pack_uids() + wrong = [0] * 8 + + def worker(i): + a = torch.full((1, M, K), float(i + 1), dtype=torch.bfloat16, device="cuda") + b = torch.eye(K, N, dtype=torch.bfloat16, device="cuda").unsqueeze(0) + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + # one workspace per thread: it is scratch the plan writes, so sharing + # it would be the very crossing this test is looking for + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + want = float(i + 1) + for _ in range(200): + g.execute({uids[0]: a, uids[1]: b, uids[2]: c}, ws, handle=handle) + torch.cuda.synchronize() + if c[0, 0, 0].item() != want: + wrong[i] += 1 + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + assert sum(wrong) == 0, f"crossed buffers between threads: {wrong}" + + +@pytest.mark.L0 +def test_describing_tensor_matches_the_dataclass(): + """``describing_tensor`` skips ``Tensor.__init__``, so every field it does + not set has to resolve to the same default the dataclass would have given + it. A new field with a ``default_factory`` gets no class attribute and + would raise here rather than reach an engine as a missing attribute.""" + import dataclasses + + from cudnn.graph_types import Tensor, describing_tensor + + fast = describing_tensor(7, (4, 3), (6, 1), cudnn.data_type.FLOAT) + slow = Tensor(uid=7, dim=(4, 3), stride=(6, 1), data_type=cudnn.data_type.FLOAT) + for f in dataclasses.fields(Tensor): + assert getattr(fast, f.name) == getattr(slow, f.name), f.name + + +@pytest.mark.L0 +def test_shape_overrides_reach_a_migrated_plan_as_a_pack(): + """Overrides do not change what a python plan is handed. + + They exist so the backend can re-describe a tensor it lowered at another + shape. A python engine reads the shape off the buffer, which is what the + pack already carries. Branching on them used to send a migrated plan the + raw uid map, which it cannot read. + """ + total, h, d, nseq = 256, 4, 128, 2 + dt = cudnn.data_type.BFLOAT16 + g = cudnn.pygraph() + q = g.tensor([total, h, d], data_type=dt, name="q") + k = g.tensor([total, h, d], data_type=dt, name="k") + v = g.tensor([total, h, d], data_type=dt, name="v") + gate = g.tensor([total, h], data_type=cudnn.data_type.FLOAT, name="g") + beta = g.tensor([total, h], data_type=cudnn.data_type.FLOAT, name="beta") + cu = g.tensor([nseq + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + out, _fs, _h = g.gdn(q=q, k=k, v=v, g=gate, beta=beta, cu_seqlens=cu, scale=1.0 / d**0.5, name="gdn") + out.set_output(True).set_data_type(dt) + try: + g.build() + except Exception as exc: # no python engine on this arch -- nothing to assert + pytest.skip(f"no GDN engine here: {exc}") + if not g._compiled_plans[g._plan_index].takes_variant_pack: + pytest.skip("the selected plan has not migrated to the variant pack") + + per = total // nseq + data = { + q: torch.randn(total, h, d, dtype=torch.bfloat16, device="cuda"), + k: torch.randn(total, h, d, dtype=torch.bfloat16, device="cuda"), + v: torch.randn(total, h, d, dtype=torch.bfloat16, device="cuda"), + gate: torch.rand(total, h, device="cuda").log(), + beta: torch.rand(total, h, device="cuda"), + cu: torch.tensor([0, per, 2 * per], dtype=torch.int32, device="cuda"), + out: torch.empty(total, h, d, dtype=torch.bfloat16, device="cuda"), + } + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute(data, ws) + torch.cuda.synchronize() + plain = data[out].clone() + + data[out].zero_() + g.execute(data, ws, override_uids=[q.get_uid()], override_shapes=[[total, h, d]], override_strides=[[h * d, d, 1]]) + torch.cuda.synchronize() + torch.testing.assert_close(data[out], plain)