From 9631a25a39c22b19f8cd85c45ebe5d2d8cfe3d9e Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 10 Aug 2026 19:23:57 -0700 Subject: [PATCH 01/11] Keep frozen-ness on the graph, where it is one flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _freeze() stored a _frozen flag on the graph AND on every Tensor, every Node and the GraphContext, and gave the latter three a __setattr__ guard so a direct attribute write would raise. Freezing is a property of the graph; four copies of the state, and three guards to read them, is not what enforcing it needs. The guards were also expensive in a way nothing measured. A dataclass __init__ assigns field by field, so overriding __setattr__ turns construction into one python-level call plus one failed `getattr(self, "_frozen", False)` lookup PER FIELD. Tensor has fifteen. Measured: 2.53 us to construct a Tensor, of which 2.14 us was the guard, on an object that is by definition not yet frozen. A graph pays it once per tensor and once per node, every build. What actually closes the mutation routes is unchanged: _check_mutable guards every setter and op builder, node.inputs/outputs/params become MappingProxy views, and dim/stride become tuples. Those are structural — they cost nothing per call and they cannot be bypassed. What is no longer an error is assigning `t.dim = [...]` directly on a frozen graph, which was never a route the API offered; the test now pins the routes it does offer. Tensor construction: 2.53 -> 0.56 us. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/_pygraph.py | 27 +++++++++++---------------- python/cudnn/graph_types.py | 7 ------- python/cudnn/nodes.py | 7 ------- test/python/test_graph_native.py | 27 +++++++++++++++++++-------- 4 files changed, 30 insertions(+), 38 deletions(-) diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 088874bcc..bec820c70 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -50,11 +50,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. @@ -360,14 +355,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. What the caller could otherwise change behind the API's back + is made immutable in its 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 +374,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: diff --git a/python/cudnn/graph_types.py b/python/cudnn/graph_types.py index 35283f286..ecabe83e0 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: 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/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() From 47997bddb0c55fca38051e2bf5c43ce4dfd7d941 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 10 Aug 2026 19:36:45 -0700 Subject: [PATCH 02/11] Normalize the variant pack once; both paths read the same operands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit graph.execute() inspected the caller's buffers twice and differently. The backend path built a {uid: pointer} dict (_native_var_pack), whose _ptr accepted a bare device address. A python engine got the caller's objects untouched via resolve_node_buffers and reached them through frost.buffers.probe, which raised "buffer of type int exposes neither __cuda_array_interface__ nor __dlpack__" for that same address. One public call, two answers, and the caller does not choose which plan the heuristics land on. Normalization now happens once, at the top of execute(), into Operands: the caller-filled uids ascending, a ctypes pointer array, and — when a python engine will read them — a Tensor record per operand carrying the buffer's own dim/stride/data_type. Below that line the backend takes ctypes.addressof(ptrs) and every engine takes pointers plus records. A bare address that the backend took now reaches an engine too, shaped by the geometry the graph declares for that port. The order comes from exactly one source, never a union: the lowered graph's variant-pack template when there is one (only C++ can see every user slot — 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 must be excluded), and the IR only for the python-only ops that cannot lower at all. The two sides never have to agree: each indexes the layout it was handed. C++ already turned a uid map into sorted pointers internally ("uid map -> extract sorted ptrs, delegate to the sorted_ptrs implementation", graph_interface.h), so passing the array directly drops one dict build here, one map copy in pybind and one hash lookup per operand there. execute_with_raw_ptrs gains a plan_index because it only ever ran plans.candidate, which stops being the plan the python walk built once the walk has skipped an entry; the vector overload it duplicated had no callers and goes. The pointer array is allocated PER CALL. Two threads may execute one graph concurrently with different buffers, and a shared array hands each thread the other's pointers — silently, since each pointer in it is individually valid. The new test fails with [0,2,7,0,1,2,2,14] crossed results when the array is shared. Also deleted: the 87-line execute/execute_plan_at_index pair monkey-patched onto backend_graph in __init__.py, unreachable since #336 made cudnn.pygraph a python class that defines both names itself; the two always-false `hasattr(graph, "_execute_with_ptrs")` fast paths in experimental/ops/sdpa.py and the uid_order cache feeding them; and the five places docs/adding_torch_custom_ops.md told authors to hand-roll that path, which raises AttributeError as written. Backend execute on a 128^3 bf16 matmul: 16.17 -> 14.76 us. Co-Authored-By: Claude Opus 5 (1M context) --- docs/adding_torch_custom_ops.md | 22 +-- docs/python_graph_and_execution_backends.md | 58 +++++- python/cudnn/__init__.py | 87 --------- python/cudnn/_pygraph.py | 165 +++++++++++++++++- python/cudnn/datatypes.py | 21 +++ python/cudnn/engines/base.py | 134 ++++++++++++-- python/cudnn/experimental/ops/sdpa.py | 58 +----- python/pygraph/pygraph.cpp | 29 +-- python/pygraph/pygraph.h | 11 +- .../python/test_variant_pack_normalization.py | 122 +++++++++++++ 10 files changed, 505 insertions(+), 202 deletions(-) create mode 100644 test/python/test_variant_pack_normalization.py 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..29e1e1f7e 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -80,15 +80,57 @@ 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 + `CompiledPlan.execute(graph, operands, ExecutionContext)` with explicit + handle/stream/workspace/overrides. 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 `Operands` 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. + +`Operands` carries the caller-filled uids ascending, a ctypes pointer array +(`address` goes straight to `_execute_with_raw_ptrs`), and a `Tensor` record per +operand holding the buffer's OWN dim/stride/data_type. That record is +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 record 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_operands` 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 records 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/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 bec820c70..01103bc1a 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -19,11 +19,14 @@ >>> 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 .datatypes import _torch_to_cudnn_data_type +from .engines.base import Operands from .graph_types import NodeType, Tensor from .nodes import Node, _row_major_stride @@ -145,6 +148,9 @@ 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 # ========================================================================= # Routing @@ -362,8 +368,8 @@ def _freeze(self) -> None: 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. What the caller could otherwise change behind the API's back - is made immutable in its own right rather than watched: + 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: @@ -1713,7 +1719,12 @@ 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 @@ -1731,19 +1742,154 @@ def execute( # 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] + # Normalize only for a plan that reads the result. Building records + # an engine will not look at is pure cost, and the ones that have + # not migrated still take the caller's objects. + if plan.takes_operands and not overriding: + plan.execute(self, self._normalize(uid_to_data, workspace, describe=True), 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 + operands = None if overriding else self._normalize(uid_to_data, workspace, describe=False) + + # 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 operands is not None: + self._lowered_graph._execute_with_raw_ptrs( + operands.address, + len(operands), + operands.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 _operand_uids(self) -> Optional[List[int]]: + """The graph's caller-filled operands, 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 operands 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, describe: bool): + """Turn the caller's variant pack into :class:`Operands`, once. + + This is the ONLY place a caller's object is inspected. Everything below + — the backend and every python engine — reads pointers and the records + 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. + """ + order = self._operand_uids() + if order is None: + return None + n = len(order) + ptrs = (ctypes.c_void_p * n)() + records = [None] * n if describe else None + # The backend's layout is exactly the slots it REQUIRES, so a hole there + # is the caller's mistake and is named. A python-only graph's layout is + # every wired port, which includes the optional ones (gdn's final_state, + # H); a hole is simply "not requested", and the engine reads it back as + # a missing port. + strict = self._lowered_graph is not None + for i, uid in enumerate(order): + data = uid_to_data.get(uid) + if data is None: + if strict: + 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}") + continue # leaves ptrs[i] NULL and records[i] None + if describe: + ptrs[i], records[i] = self._describe(data, uid) + else: + # The backend reads geometry from its own descriptors; building + # a record it will not look at is pure cost. + ptrs[i] = self._device_pointer(data) + return Operands(tuple(order), tuple(records) if describe else None, ptrs, self._device_pointer(workspace) if workspace is not None else 0) + + def _describe(self, data: Any, uid: int): + """``(pointer, Tensor)`` for one caller buffer. + + The record 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. Reading these straight off a torch tensor costs ~0.5 us; + going through DLPack (``frost.buffers.probe``) costs 2-9 us, which is + why this does not. + + A caller who passed a bare address gets a record with no geometry: a + pointer carries none, and the backend has always accepted that. + """ + if type(data) is int: + return data, Tensor(uid=uid) + dim = getattr(data, "shape", None) + if dim is not None and hasattr(data, "stride") and hasattr(data, "data_ptr"): + return data.data_ptr(), Tensor(uid=uid, dim=tuple(dim), stride=tuple(data.stride()), data_type=_torch_to_cudnn_data_type(data.dtype)) + ptr = self._device_pointer(data) + if dim is not None: + return ptr, Tensor(uid=uid, dim=tuple(dim)) + return ptr, Tensor(uid=uid) + + @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).""" @@ -1828,6 +1974,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 operands, 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..bf2383757 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -167,6 +167,27 @@ def _torch_to_cudnn_data_type(torch_data_type) -> cudnn_data_type: return None +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 and + block-scaled ones — a caller must pass those as a typed buffer, not as a + bare address). + + Lives here so the mapping has one home; frost imports it rather than + keeping a second table.""" + return { + 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", + }.get(data_type) + + 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. diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 2d5d4da01..9bd2c8f29 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -43,6 +43,7 @@ def execute(self, graph, uid_to_data, ctx): ... # write results into caller-provided output buffers """ +import ctypes from abc import ABC from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List @@ -108,17 +109,118 @@ class ExecutionContext: override_strides: Any = None +class Operands: + """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. + + ``uids`` is ASCENDING, matching the backend's own operand order + (``get_variant_pack_uids_sorted()``), so ``ctypes.addressof(ptrs)`` goes + straight to ``_execute_with_raw_ptrs`` with no copy and no per-operand + hash lookup. + + ``tensors[i]`` describes what the caller ACTUALLY passed for ``uids[i]``: + its ``dim`` / ``stride`` / ``data_type`` are the buffer's, which is not + necessarily what the graph declared. An engine reads whichever it means — + the IR port for the shape the plan was built for, this record 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 array would hand each thread the other's + pointers — silently, because every pointer in it is individually valid. + """ + + __slots__ = ("uids", "tensors", "ptrs", "address", "_slot_of", "workspace") + + def __init__(self, uids, tensors, ptrs, workspace_ptr: int): + self.uids = uids + self.tensors = tensors + self.ptrs = ptrs + self.address = ctypes.addressof(ptrs) + self.workspace = workspace_ptr + self._slot_of = None # built on first lookup: the backend never does one + + @property + def slot_of(self): + if self._slot_of is None: + self._slot_of = {u: i for i, u in enumerate(self.uids)} + return self._slot_of + + def 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.ptrs[self.slot(tensor_or_uid)] or 0 + + 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", operands: Operands) -> 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 = operands.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 ``Operands``. + + Per-node ``{port_name: caller buffer}`` maps, the result of + ``resolve_node_buffers``.""" inputs: Dict[str, Any] outputs: Dict[str, Any] @@ -142,6 +244,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 +260,16 @@ def resolve(node, ports, direction): class CompiledPlan: """A compiled (graph, plan) artifact. Subclass for real JIT engines.""" + # Set True once execute() takes Operands. 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_operands: 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", operands: "Operands", 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/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/test/python/test_variant_pack_normalization.py b/test/python/test_variant_pack_normalization.py new file mode 100644 index 000000000..54cf40da7 --- /dev/null +++ b/test/python/test_variant_pack_normalization.py @@ -0,0 +1,122 @@ +# 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 + +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._operand_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() + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + uids = g._operand_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") + 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}" From c71239023c6cfa3766733160068878a0cee87372 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 10 Aug 2026 20:41:31 -0700 Subject: [PATCH 03/11] Fill a DLPack struct from a per-layout prototype, not field by field MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every field of a `DLManagedTensor` except `data` is a property of the layout, yet `DeviceView.__dlpack__` built a fresh shape array and assigned nine ctypes fields on every call — and a graph makes ten of these per execute, one per workspace region carved for the kernel. Fill the struct once per (shape, dtype, device) and copy it: 1.68 us of field assignment becomes a 0.45 us memmove of 72 bytes. Measured 3.62 -> 1.39 us per `__dlpack__`, and GDN forward 154.5 -> 137.6 us end to end with no engine touched, because the ten workspace views are all it takes. This is the shape the backend already uses for kernel arguments (src/common/include/runtimeKernel.h): a prefilled blob plus, per mutable field, an (offset, uid, UpdateMethod) saying what execute writes where. Here there is exactly one mutable field, `data`, at a fixed offset, with update method POINTER, so the bookkeeping collapses to a memmove and one assignment. The struct stays FRESH per capsule. cute's from_dlpack aliases it rather than copying the DLTensor, so a struct shared between two capsules — or between two threads executing one graph — is read after someone else re-pointed it. Only the prototype is shared, and it is immutable. Also renames Operands to VariantPack: it IS the variant pack, normalized, and the python-side one being slightly wider than the C++ template's is not worth a second word. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/_pygraph.py | 33 ++++--- python/cudnn/datatypes.py | 28 +++--- python/cudnn/engines/base.py | 14 +-- python/cudnn/frost/buffers.py | 93 +++++++++++++++---- test/python/test_dlpack_proto.py | 66 +++++++++++++ .../python/test_variant_pack_normalization.py | 4 +- 6 files changed, 186 insertions(+), 52 deletions(-) create mode 100644 test/python/test_dlpack_proto.py diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 01103bc1a..8b44d82cf 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from .datatypes import _torch_to_cudnn_data_type -from .engines.base import Operands +from .engines.base import VariantPack from .graph_types import NodeType, Tensor from .nodes import Node, _row_major_stride @@ -1746,13 +1746,13 @@ def execute( # Normalize only for a plan that reads the result. Building records # an engine will not look at is pure cost, and the ones that have # not migrated still take the caller's objects. - if plan.takes_operands and not overriding: + if plan.takes_variant_pack and not overriding: plan.execute(self, self._normalize(uid_to_data, workspace, describe=True), ctx) else: plan.execute(self, uid_to_data, ctx) return - operands = None if overriding else self._normalize(uid_to_data, workspace, describe=False) + variant_pack = None if overriding else self._normalize(uid_to_data, workspace, describe=False) # Backend path. Address the plan the WALK built, not the backend's own # selection: they differ once the walk has skipped an entry. @@ -1764,11 +1764,11 @@ def execute( # 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 operands is not None: + if variant_pack is not None: self._lowered_graph._execute_with_raw_ptrs( - operands.address, - len(operands), - operands.workspace, + variant_pack.address, + len(variant_pack), + variant_pack.workspace, handle or 0, -1 if cpp_index is None else cpp_index, ) @@ -1780,8 +1780,8 @@ def execute( return self._lowered_graph._execute(var_pack, ws_ptr, handle, override_uids, override_shapes, override_strides) - def _operand_uids(self) -> Optional[List[int]]: - """The graph's caller-filled operands, ASCENDING by uid. + 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 @@ -1791,7 +1791,7 @@ def _operand_uids(self) -> Optional[List[int]]: destinations, cached workspace modifications). A python-only graph — gdn / kda / gdn2, which cannot lower by - construction — has no C++ side, so its operands come from the IR. The + 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 @@ -1822,7 +1822,7 @@ def _operand_uids(self) -> Optional[List[int]]: return order def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool): - """Turn the caller's variant pack into :class:`Operands`, once. + """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 pointers and the records @@ -1830,7 +1830,7 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool passed. Returns None when the operand layout is not known yet, which puts the caller back on the uid-map path. """ - order = self._operand_uids() + order = self._variant_pack_uids() if order is None: return None n = len(order) @@ -1856,7 +1856,12 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool # The backend reads geometry from its own descriptors; building # a record it will not look at is pure cost. ptrs[i] = self._device_pointer(data) - return Operands(tuple(order), tuple(records) if describe else None, ptrs, self._device_pointer(workspace) if workspace is not None else 0) + return VariantPack( + tuple(order), + tuple(records) if describe else None, + ptrs, + self._device_pointer(workspace) if workspace is not None else 0, + ) def _describe(self, data: Any, uid: int): """``(pointer, Tensor)`` for one caller buffer. @@ -1974,7 +1979,7 @@ 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 operands, so an order cached while + # 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 diff --git a/python/cudnn/datatypes.py b/python/cudnn/datatypes.py index bf2383757..e665bc277 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -167,6 +167,22 @@ 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", +} + + 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 and @@ -175,17 +191,7 @@ def _cudnn_to_frost_dtype_name(data_type): Lives here so the mapping has one home; frost imports it rather than keeping a second table.""" - return { - 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", - }.get(data_type) + return _CUDNN_TO_FROST_DTYPE_NAME.get(data_type) def _torch_to_cutlass_data_type(data_type, interpret_uint8_as_fp4x2: bool = False): diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 9bd2c8f29..59df5e5be 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -109,7 +109,7 @@ class ExecutionContext: override_strides: Any = None -class Operands: +class VariantPack: """The caller's variant pack, normalized ONCE at the top of execute(). Whatever the caller passed — a torch tensor, a DeviceView, any @@ -175,7 +175,7 @@ class PortSlots: outputs: Dict[str, int] -def bind_ports(graph: "pygraph", operands: Operands) -> Dict[Any, PortSlots]: +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.""" @@ -184,7 +184,7 @@ def resolve(node, ports, direction): for port, t in ports.items(): if t is None: continue - slot = operands.slot_of.get(t.uid) + slot = variant_pack.slot_of.get(t.uid) if slot is None: if t.is_virtual: continue # engine-internal intermediate @@ -217,7 +217,7 @@ def _view_over_address(address: int, tensor, node, port: str): @dataclass(frozen=True) class NodeBuffers: - """DEPRECATED, kept until every engine takes ``Operands``. + """DEPRECATED, kept until every engine takes ``VariantPack``. Per-node ``{port_name: caller buffer}`` maps, the result of ``resolve_node_buffers``.""" @@ -260,16 +260,16 @@ def resolve(node, ports, direction): class CompiledPlan: """A compiled (graph, plan) artifact. Subclass for real JIT engines.""" - # Set True once execute() takes Operands. Until then execute() is handed the + # 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_operands: bool = False + 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", operands: "Operands", ctx: ExecutionContext) -> None: + def execute(self, graph: "pygraph", variant_pack: "VariantPack", ctx: ExecutionContext) -> None: raise NotImplementedError diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 1c4364c26..aa25ba3df 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -98,18 +98,80 @@ def dtype_name(buf) -> str: return str(buf.dtype).split(".")[-1] +class _DLPackProto: + """A filled-in ``DLManagedTensor`` for one (shape, dtype, device), minus the + address. + + Every field but ``data`` is a property of the layout, so building them once + and copying the struct beats assigning nine ctypes fields per call: measured + 1.68 us to fill vs 0.45 us to memmove the 72 bytes. Cached per (shape, + dtype, device) — a graph reuses a handful of layouts across its whole run. + + This is the degenerate case of what the backend already does for kernel + arguments (``src/common/include/runtimeKernel.h``): a prefilled blob plus, + per mutable field, an (offset, uid, UpdateMethod) telling execute what to + write where. Here there is exactly one mutable field, ``data``, at a fixed + offset, and its update method is always POINTER — so the bookkeeping + collapses to a memmove and one assignment. + """ + + __slots__ = ("_template", "_shape_arr", "_nbytes") + + def __init__(self, shape, dtype: str, device_id: int): + ndim = len(shape) + self._shape_arr = (ctypes.c_int64 * max(ndim, 1))(*shape) + code, bits = DTYPES[dtype] + mt = _DLManagedTensor() + mt.dl_tensor.device = _DLDevice(_KDL_CUDA, device_id) + mt.dl_tensor.ndim = ndim + mt.dl_tensor.dtype = _DLDataType(code, bits, 1) + mt.dl_tensor.shape = self._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._template = mt + self._nbytes = ctypes.sizeof(_DLManagedTensor) + + def instantiate(self, ptr: int): + """A fresh struct at ``ptr``. Fresh, not shared: a capsule handed to + CuTe ALIASES the struct rather than copying it, and two threads + executing one graph must not be writing the same one.""" + mt = _DLManagedTensor() + ctypes.memmove(ctypes.byref(mt), ctypes.byref(self._template), self._nbytes) + mt.dl_tensor.data = ctypes.c_void_p(ptr) + return mt + + +_PROTO_CACHE = {} + + +def dlpack_proto(shape, dtype: str, device_id: int) -> _DLPackProto: + key = (tuple(shape), dtype, device_id) + proto = _PROTO_CACHE.get(key) + if proto is None: + proto = _PROTO_CACHE[key] = _DLPackProto(key[0], dtype, device_id) + return proto + + class DeviceView: """Zero-copy DLPack view over a raw CUDA pointer. The view owns no memory — the underlying allocation (workspace or caller - buffer) must outlive it. Row-major contiguous.""" + buffer) must outlive it. Row-major contiguous. + + Not a concept an engine has to learn: it is what a variant-pack slot or a + workspace region turns into on the way to a kernel, because CuTe needs an + object exposing ``__dlpack__`` and neither a pointer nor a Tensor record + is one.""" def __init__(self, ptr: int, shape, dtype: str, device_id: int): self._ptr = int(ptr) self.shape = tuple(int(s) for s in shape) self.dtype = dtype self._device_id = int(device_id) - self._keepalive = [] + self._proto = None + self._live = [] def data_ptr(self) -> int: return self._ptr @@ -166,22 +228,17 @@ 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) - 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)) + # A fresh struct per call, copied from the layout's prototype and + # re-pointed. Fresh because the consumer ALIASES it — cute's + # from_dlpack keeps the pointer rather than copying the DLTensor — so a + # struct shared between two capsules, or between two threads executing + # one graph, would be read after someone else rewrote it. + if self._proto is None: + self._proto = dlpack_proto(self.shape, self.dtype, self._device_id) + mt = self._proto.instantiate(self._ptr) + # The struct and the prototype's shape array must outlive the capsule: + # the deleter is a no-op, so nothing else keeps them alive. + self._live.append(mt) return _PyCapsule_New(ctypes.addressof(mt), b"dltensor", None) diff --git a/test/python/test_dlpack_proto.py b/test/python/test_dlpack_proto.py new file mode 100644 index 000000000..964083ecf --- /dev/null +++ b/test/python/test_dlpack_proto.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A DeviceView's DLPack struct is copied from a per-layout prototype. + +Everything in a ``DLManagedTensor`` except ``data`` is a property of the layout, +so the struct is filled once per (shape, dtype, device) and memmove'd per call: +measured 1.68 us to assign nine ctypes fields against 0.45 us to copy the 72 +bytes. Same shape as the backend's kernel-argument handling +(``src/common/include/runtimeKernel.h``), where a prefilled blob carries an +(offset, uid, UpdateMethod) per mutable field; here there is exactly one +mutable field. + +The struct must be FRESH per capsule even so: cute's ``from_dlpack`` aliases it +rather than copying, so a struct shared between two capsules is read after +someone else re-pointed it. +""" + +import ctypes + +import pytest +import torch + +from cudnn.frost import buffers + + +@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__() + addr = ctypes.pythonapi.PyCapsule_GetPointer + addr.restype = ctypes.c_void_p + addr.argtypes = [ctypes.py_object, ctypes.c_char_p] + assert addr(a, b"dltensor") != addr(b, b"dltensor") + + +@pytest.mark.L0 +def test_prototypes_are_shared_across_views_of_one_layout(): + """The prototype is the cache; the struct is not.""" + t = torch.zeros(4, 5, dtype=torch.bfloat16, device="cuda") + dev = t.device.index or 0 + v1 = buffers.DeviceView(t.data_ptr(), (4, 5), "bfloat16", dev) + v2 = buffers.DeviceView(t.data_ptr() + 64, (4, 5), "bfloat16", dev) + v1.__dlpack__() + v2.__dlpack__() + assert v1._proto is v2._proto + + +@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_variant_pack_normalization.py b/test/python/test_variant_pack_normalization.py index 54cf40da7..44321614e 100644 --- a/test/python/test_variant_pack_normalization.py +++ b/test/python/test_variant_pack_normalization.py @@ -84,7 +84,7 @@ def test_operand_order_is_the_backend_order(): destinations, workspace modifications) must be excluded. """ g, _, _ = _matmul_graph() - order = g._operand_uids() + order = g._variant_pack_uids() assert order == sorted(order), f"not ascending: {order}" assert order == list(g._lowered_graph._get_variant_pack_uids_sorted()) @@ -100,7 +100,7 @@ def test_execute_is_reentrant(): g, _, _ = _matmul_graph() handle = cudnn.create_handle() ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") - uids = g._operand_uids() + uids = g._variant_pack_uids() wrong = [0] * 8 def worker(i): From 1c11effaa1d358bcd37d0590161f314ca9eb3362 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 01:02:54 -0700 Subject: [PATCH 04/11] Drive the linear-attention engines from the normalized variant pack execute() already normalized the caller's operands once; the engines were still reading the caller's objects a second time. Every port went through buffers.probe(), which for bfloat16 falls out of __cuda_array_interface__ and into torch's __dlpack__ at 8.6 us apiece -- nine per GDN forward, 47.7 us, to learn dim and stride that the pack was already holding. _FrostPlan now takes the pack. The port-to-slot join is a property of the graph, so it is computed once and kept; between executes only the addresses move. What reaches the kernel is built from the pack rather than passed through, so the geometry a buffer is checked against and the geometry it runs on are the same reading. Contiguity moves with it, and becomes one gate instead of one call per compiled callable naming its own ports. That list was the same every time and had to be maintained by hand: a port added to a node but forgotten there went unchecked. Workspace joins too -- it needs a pointer, a device and a size, and the pack now carries all three, so Workspace.over() replaces a tenth probe. Two costs are added on purpose. Building eight DeviceViews is 10.1 us, and handing them to the kernel instead of the caller's tensors is another 17.6, because tvm-ffi reads a torch tensor through a C vtable (__dlpack_c_exchange_api__) and any python producer through a capsule. Both are the same fact -- python cannot build a fast DLPack producer -- and both go when the producer becomes a C type. Keeping the caller's tensor to avoid them would mean torch is a hard dependency of the engine path, which is the thing this removes. GDN forward, SM100, total=4096 H=4 D=128 4 seqs: before after contiguity gate 47.7 4.2 resolve_node_buffers 8.3 0 (bound once, kept) workspace probe 3.5 0.2 normalize 0 13.9 (now reads the workspace too) building the views 0 10.1 execute() 127 117 Also here, found while measuring: - selected_engine is a property execute() calls every time, and answering it walked every registered engine for the one declaring this id: 2.75 -> 0.48 us, cached against the plan config's identity so replanning invalidates it without a hook on every writer of _plan_index. - Two in-function imports of things the module already imports at the top. The one in VariantPack.view() ran per operand and cost 19 us of the GDN forward on its own. - _describe asked a torch tensor for its facts and gave everyone else a half-filled Tensor: no stride, no data_type. It now asks each producer in its own spelling -- torch's element-unit stride() and data_ptr(), cupy's byte-unit .strides and .data.ptr, one DLPack read for the rest -- so the same buffer is described the same way whoever produced it. fp8 is in the dtype table for the same reason; fp4 is deliberately not, since DTYPE_ITEMSIZE would make it zero bytes wide. - probe() declined two different ways through one exception, so an operand whose dtype has no name here lost its dim and stride as well. The two are told apart now. - The descriptor-skip cache in four kernels had a 0% hit rate: one of its guards compared against ws.view(...), a fresh object every call. It asked torch's _version counter whether cu_seqlens had changed, which was sound for torch callers and silently stale for everyone else. Deleting it is 5 us faster than keeping it. - check_buffer_device walked every operand asking which GPU it was on. cuDNN's own variant pack carries no device at all, and the question that matters is where the launch is going, not where the memory is: one current_device() read, 0.74 us against 1.45 per operand. 419 linear-attention tests pass, 1769 skipped. --- docs/python_graph_and_execution_backends.md | 23 +-- python/cudnn/_pygraph.py | 120 ++++++++++----- python/cudnn/datatypes.py | 42 +++++- python/cudnn/engines/base.py | 48 +++++- python/cudnn/frost/buffers.py | 42 +++++- python/cudnn/frost/device.py | 42 +----- python/cudnn/frost/workspace.py | 26 ++++ python/cudnn/gemm/frost/compiler.py | 29 +++- python/cudnn/graph_types.py | 37 +++++ python/cudnn/linear_attention/engine_utils.py | 60 ++++++-- .../linear_attention/frost/gdn2_engine.py | 5 +- .../linear_attention/frost/gdn_engine.py | 27 +--- .../linear_attention/frost/kda_engine.py | 5 +- .../frost/kernel/gdn2_prefill_f16.py | 141 +++++++----------- .../frost/kernel/gdn_bprop_f16.py | 113 +++++--------- .../frost/kernel/gdn_prefill_f16.py | 123 ++++++--------- .../frost/kernel/kda_prefill_f16.py | 135 +++++++---------- .../python/test_variant_pack_normalization.py | 16 ++ 18 files changed, 566 insertions(+), 468 deletions(-) diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index 29e1e1f7e..0c6182aad 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -88,7 +88,7 @@ create_execution_plans([heur_mode.A, ...]) _pygraph.py `graph.execute()` converts whatever the caller passed — a torch tensor, a `DeviceView`, any `__dlpack__` / `__cuda_array_interface__` producer, or a bare -device address — into `Operands` at the top, and everything below reads that. +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 @@ -96,13 +96,14 @@ 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. -`Operands` carries the caller-filled uids ascending, a ctypes pointer array -(`address` goes straight to `_execute_with_raw_ptrs`), and a `Tensor` record per -operand holding the buffer's OWN dim/stride/data_type. That record is -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 record for the shape about to run.** +`VariantPack` carries the caller-filled uids ascending, a ctypes pointer array +(`address` goes straight to `_execute_with_raw_ptrs`), and one `Tensor` per +operand — the same class `graph.tensor()` returns — holding the buffer's OWN +dim/stride/data_type. It is 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 variant pack's +`Tensor` for the shape about to run.** Two rules that are easy to break by accident: @@ -125,10 +126,10 @@ 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_operands` is the migration flag. An engine that has not set -it still receives the caller's `{uid: buffer}` map and reaches ports through +`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 records only for a plan that sets the flag — +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 diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 8b44d82cf..cb6c2c575 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -25,9 +25,10 @@ import weakref from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -from .datatypes import _torch_to_cudnn_data_type -from .engines.base import VariantPack -from .graph_types import NodeType, Tensor +from .datatypes import _buffer_dtype_to_cudnn, _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") @@ -151,6 +152,7 @@ def __init__( # 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 @@ -163,8 +165,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] @@ -172,8 +172,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) @@ -224,8 +222,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 @@ -1223,8 +1234,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)) @@ -1726,8 +1735,6 @@ def execute( eng = self.selected_engine if eng is not None: # python engine (plan id in the reserved region) - from .engines.base import ExecutionContext - h = handle if handle is not None else self._handle ctx = ExecutionContext( handle=h, @@ -1743,7 +1750,7 @@ def execute( self._compiled_plans[self._plan_index] = eng.build_plan(self, self._selected_plan_config, ctx) self._is_built = True plan = self._compiled_plans[self._plan_index] - # Normalize only for a plan that reads the result. Building records + # Normalize only for a plan that reads the result. Building Tensors # an engine will not look at is pure cost, and the ones that have # not migrated still take the caller's objects. if plan.takes_variant_pack and not overriding: @@ -1825,7 +1832,7 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool """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 pointers and the records + — 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. @@ -1835,7 +1842,7 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool return None n = len(order) ptrs = (ctypes.c_void_p * n)() - records = [None] * n if describe else None + tensors = [None] * n if describe else None # The backend's layout is exactly the slots it REQUIRES, so a hole there # is the caller's mistake and is named. A python-only graph's layout is # every wired port, which includes the optional ones (gdn's final_state, @@ -1849,41 +1856,86 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool 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}") - continue # leaves ptrs[i] NULL and records[i] None + continue # leaves ptrs[i] NULL and tensors[i] None if describe: - ptrs[i], records[i] = self._describe(data, uid) + ptrs[i], tensors[i] = self._describe(data, uid) else: # The backend reads geometry from its own descriptors; building - # a record it will not look at is pure cost. + # a Tensor it will not look at is pure cost. ptrs[i] = self._device_pointer(data) + # 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: + if describe: + workspace_ptr, workspace_tensor = self._describe(workspace, -1) + workspace_bytes = _byte_size(workspace_tensor) + else: + workspace_ptr = self._device_pointer(workspace) return VariantPack( tuple(order), - tuple(records) if describe else None, + tuple(tensors) if describe else None, ptrs, - self._device_pointer(workspace) if workspace is not None else 0, + workspace_ptr, + workspace_bytes, ) def _describe(self, data: Any, uid: int): """``(pointer, Tensor)`` for one caller buffer. - The record carries the buffer's OWN dim/stride/data_type, which need + 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. Reading these straight off a torch tensor costs ~0.5 us; - going through DLPack (``frost.buffers.probe``) costs 2-9 us, which is - why this does not. - - A caller who passed a bare address gets a record with no geometry: a + 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 caller who passed a bare address gets a Tensor with no geometry: a pointer carries none, and the backend has always accepted that. """ if type(data) is int: - return data, Tensor(uid=uid) + return data, Tensor(uid=uid) # bare device address dim = getattr(data, "shape", None) - if dim is not None and hasattr(data, "stride") and hasattr(data, "data_ptr"): - return data.data_ptr(), Tensor(uid=uid, dim=tuple(dim), stride=tuple(data.stride()), data_type=_torch_to_cudnn_data_type(data.dtype)) - ptr = self._device_pointer(data) - if dim is not None: - return ptr, Tensor(uid=uid, dim=tuple(dim)) - return ptr, Tensor(uid=uid) + + # 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: diff --git a/python/cudnn/datatypes.py b/python/cudnn/datatypes.py index e665bc277..1e738263a 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -180,20 +180,56 @@ def _torch_to_cudnn_data_type(torch_data_type) -> cudnn_data_type: 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 and - block-scaled ones — a caller must pass those as a typed buffer, not as a - bare address). + 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. diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 59df5e5be..f7c8c853b 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -49,6 +49,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 @@ -124,10 +126,11 @@ class VariantPack: straight to ``_execute_with_raw_ptrs`` with no copy and no per-operand hash lookup. - ``tensors[i]`` describes what the caller ACTUALLY passed for ``uids[i]``: - its ``dim`` / ``stride`` / ``data_type`` are the buffer's, which is not + ``tensors[i]`` is a ``graph_types.Tensor`` — the same class ``graph.tensor()`` + returns — describing what the caller ACTUALLY passed for ``uids[i]``: its + ``dim`` / ``stride`` / ``data_type`` are the buffer's, which is not necessarily what the graph declared. An engine reads whichever it means — - the IR port for the shape the plan was built for, this record for the shape + the IR port for the shape the plan was built for, this one for the shape about to run. frost_gemm takes its M/N/K from here; the backend takes only the pointer. @@ -136,15 +139,17 @@ class VariantPack: pointers — silently, because every pointer in it is individually valid. """ - __slots__ = ("uids", "tensors", "ptrs", "address", "_slot_of", "workspace") + __slots__ = ("uids", "tensors", "ptrs", "address", "_slot_of", "workspace", "workspace_bytes", "_device") - def __init__(self, uids, tensors, ptrs, workspace_ptr: int): + def __init__(self, uids, tensors, ptrs, workspace_ptr: int, workspace_bytes: int = 0): self.uids = uids self.tensors = tensors self.ptrs = ptrs self.address = ctypes.addressof(ptrs) 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 slot_of(self): @@ -152,6 +157,21 @@ def slot_of(self): 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: an execute launches on the current + device and cuDNN's own variant pack carries no device at all + (``create_variant_pack`` sets pointers, uids and the workspace). Read + on demand — 0.74 us, and 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).""" @@ -161,6 +181,24 @@ def slot(self, tensor_or_uid) -> int: def ptr(self, tensor_or_uid) -> int: return self.ptrs[self.slot(tensor_or_uid)] or 0 + 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. + """ + tensor = self.tensors[slot] + name = _CUDNN_TO_FROST_DTYPE_NAME.get(tensor.data_type) + if name is None: + raise ValueError(f"operand uid {self.uids[slot]} has no DLPack-expressible dtype ({tensor.data_type})") + return DeviceView(self.ptrs[slot] or 0, tensor.dim, name, self.device) + def __len__(self) -> int: return len(self.uids) diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index aa25ba3df..0fccaf5b8 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -75,6 +75,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 +89,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 = {" 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..3fb3a46f7 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -90,6 +90,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,6 +100,29 @@ 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. + + ``execute()`` measures the caller's workspace with the same reader it + gives every other buffer, so re-probing it here cost 3.5 us to learn + what the pack is holding. + """ + 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 diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 4665ce6d3..cb3d2a3e4 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,7 +1821,7 @@ 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) + _check_plan_device(self.device) if self.binding is None: raise NotImplementedError("variant-pack call is not wired up for this graph type") b = self.binding @@ -2856,7 +2871,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 +3247,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/graph_types.py b/python/cudnn/graph_types.py index ecabe83e0..aec72dd72 100644 --- a/python/cudnn/graph_types.py +++ b/python/cudnn/graph_types.py @@ -234,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..02790ad72 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,56 @@ 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) + node_buffers = {} + for node, slots in ports.items(): + _check_contiguous(node.name, variant_pack, slots) + node_buffers[node] = NodeBuffers( + {port: variant_pack.view(slot) for port, slot in slots.inputs.items()}, + {port: variant_pack.view(slot) for port, slot in slots.outputs.items()}, + ) + 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(node_name: str, variant_pack, slots) -> None: + """Contiguity gate over every port this node binds, read off the pack. + + 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. + + 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 here would have gone unchecked. + """ + for direction in (slots.inputs, slots.outputs): + for port, slot in direction.items(): + tensor = variant_pack.tensors[slot] + if not buffers.is_contiguous(tensor.dim, tensor.stride): + raise ValueError(f"cudnn.frost {node_name!r}: buffer for {port!r} must be contiguous (buffers pass straight to the kernel)") _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..8531a0652 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -16,7 +16,7 @@ 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 ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair def _the_gdn2_node(graph): @@ -159,11 +159,10 @@ 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)") + ws = workspace sched_ctr = ws.view(self._off_sched, "int32", (2,)) from .common.split_k import WORK_ITEM_FIELDS diff --git a/python/cudnn/linear_attention/frost/gdn_engine.py b/python/cudnn/linear_attention/frost/gdn_engine.py index 4c43a5a7b..396afb239 100644 --- a/python/cudnn/linear_attention/frost/gdn_engine.py +++ b/python/cudnn/linear_attention/frost/gdn_engine.py @@ -15,7 +15,7 @@ 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 ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair def _the_gdn_node(graph): @@ -213,9 +213,8 @@ 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 @@ -354,27 +353,7 @@ 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 diff --git a/python/cudnn/linear_attention/frost/kda_engine.py b/python/cudnn/linear_attention/frost/kda_engine.py index 11a9d9246..b7428afe0 100644 --- a/python/cudnn/linear_attention/frost/kda_engine.py +++ b/python/cudnn/linear_attention/frost/kda_engine.py @@ -15,7 +15,7 @@ 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 ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair def _the_kda_node(graph): @@ -170,11 +170,10 @@ 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)") + ws = workspace sched_ctr = ws.view(self._off_sched, "int32", (2,)) from .common.split_k import WORK_ITEM_FIELDS 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..bc01c368a 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py @@ -3199,14 +3199,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 +3396,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/test/python/test_variant_pack_normalization.py b/test/python/test_variant_pack_normalization.py index 44321614e..3bc6c2499 100644 --- a/test/python/test_variant_pack_normalization.py +++ b/test/python/test_variant_pack_normalization.py @@ -120,3 +120,19 @@ def worker(i): 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 From 64089d9091180a6f4e11a3c710200e026345f0fb Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 02:03:05 -0700 Subject: [PATCH 05/11] Hold the variant pack as DLTensors, in one C type that is also the producer The pack was python objects: a Tensor per operand, and a DeviceView per operand and per workspace carve to hand the kernel. Both halves cost more than the work they describe. Reading a buffer meant asking a python object four questions one method call at a time and building a Tensor to hold the answers, 1.5 us each. Handing one back meant building a DLPack capsule, which tvm-ffi reads at 1.86 us where it reads a torch tensor at 0.35 -- through a C function table, __dlpack_c_exchange_api__, that no python producer can offer. Both halves are that one protocol, so this consumes it and implements it. VariantPackNative reads each operand through the caller's vtable into a DLTensor it keeps; the slots it hands out carry the vtable themselves, so a kernel reads ours through the same path it reads a framework tensor -- at 0.30 us, cheaper than the tensor it replaces. Refusing to pass the caller's object through therefore costs nothing, where insisting on it used to cost 17.6 us of kernel-argument conversion. A producer without the vtable is not an error and not a cliff: read_all returns the slots it could not take, python describes those with the reader it already had, and a mixed pack costs the sum of its parts. The workspace carves are the same type as the operands now, so a graph hands its kernels one kind of buffer rather than two, and DeviceView is off the hot path entirely. GDN forward, SM100, total=4096 H=4 D=128 4 seqs: before after normalize 13.9 6.3 contiguity gate 4.2 0.4 building the views 10.1 2.3 kernel-argument penalty 17.6 0 execute() 117 58 The backend path picks this up without a line changed: describe= had stopped selecting anything once reading was a single C call, so both paths take it and _execute_with_raw_ptrs reads the native pointer array directly. The parameter is gone. tensors[] is materialized on first access rather than built eagerly -- 16.9 us for eight operands, more than twice the whole normalize. Nothing on the per-execute path asks for it; frost_gemm will, for its M/N/K, and should read the native shapes instead when it migrates. No flag decides this: the laziness is the gate, and it needs no engine to declare anything. dlpack_version.txt moves 1.1 -> 1.3 for the DLPackExchangeAPI declarations. FetchContent keeps its checkout, so an incremental build needs _deps/dlpack-* cleared to actually pick the new tag up. Two things the migration surfaced, both in kernel code the forward path never reaches: - cute's from_dlpack at compile time does not read the vtable, so a slot needs __dlpack__ as well. It transfers ownership properly -- its own copy of the shape and stride plus a real deleter -- rather than aliasing storage the slot owns, which is how DeviceView's no-op deleter became a use-after-free whenever a consumer outlived the view. - the bprop state downcast reshapes its operand, so slots reshape too. A non-contiguous one is refused rather than silently reinterpreted: DeviceView could skip that check because it was row-major by construction, and a slot is whatever the caller passed. 430 linear-attention and dispatch tests pass, 1769 skipped. Two notes for anyone reading the numbers: - The fast path needs the producer's type to carry the vtable. torch 2.13 has it natively; on older torch tvm-ffi installs it by JIT-building a small extension, which is why flashinfer gets the same path there. So what decides it is whether tvm-ffi has been imported, not the torch version -- a backend-only process on old torch takes the python fallback, correctly and slowly. - The vtable is only called after walking prev_api for a table whose major version matches the header this was built against. The protocol requires that walk and keeps older tables reachable for it; without it a producer that moved to a new major version would have us calling function pointers at offsets it was free to move. --- dlpack_version.txt | 2 +- python/CMakeLists.txt | 12 +- python/cudnn/_pygraph.py | 54 +- python/cudnn/datatypes.py | 26 + python/cudnn/engines/base.py | 84 ++- python/cudnn/frost/workspace.py | 15 +- python/cudnn/linear_attention/engine_utils.py | 34 +- python/pycudnn.cpp | 5 + python/pygraph/variant_pack.cpp | 646 ++++++++++++++++++ python/pygraph/variant_pack.h | 14 + 10 files changed, 821 insertions(+), 71 deletions(-) create mode 100644 python/pygraph/variant_pack.cpp create mode 100644 python/pygraph/variant_pack.h 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/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/_pygraph.py b/python/cudnn/_pygraph.py index cb6c2c575..a009adee7 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -25,7 +25,7 @@ import weakref from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -from .datatypes import _buffer_dtype_to_cudnn, _torch_to_cudnn_data_type +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 @@ -1754,12 +1754,12 @@ def execute( # an engine will not look at is pure cost, and the ones that have # not migrated still take the caller's objects. if plan.takes_variant_pack and not overriding: - plan.execute(self, self._normalize(uid_to_data, workspace, describe=True), ctx) + plan.execute(self, self._normalize(uid_to_data, workspace), ctx) else: plan.execute(self, uid_to_data, ctx) return - variant_pack = None if overriding else self._normalize(uid_to_data, workspace, describe=False) + 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. @@ -1828,7 +1828,7 @@ def _variant_pack_uids(self) -> Optional[List[int]]: self._sorted_uids = order return order - def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool): + def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any): """Turn the caller's variant pack into :class:`VariantPack`, once. This is the ONLY place a caller's object is inspected. Everything below @@ -1841,45 +1841,47 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, describe: bool if order is None: return None n = len(order) - ptrs = (ctypes.c_void_p * n)() - tensors = [None] * n if describe else None + import cudnn + + native = cudnn._pybind_module.VariantPackNative(n) + # One crossing for the whole pack: the reader is a C function table on + # the buffer's type (__dlpack_c_exchange_api__), so an operand costs + # 0.08 us there against 1.5 to ask a python object the same four + # questions and build a Tensor to hold the answers. What comes back is + # the slots whose producer does not implement it — described here, at + # the price they always cost, without taking the rest down with them. + unread = native.read_all([uid_to_data.get(uid) for uid in order]) # The backend's layout is exactly the slots it REQUIRES, so a hole there # is the caller's mistake and is named. A python-only graph's layout is # every wired port, which includes the optional ones (gdn's final_state, # H); a hole is simply "not requested", and the engine reads it back as # a missing port. strict = self._lowered_graph is not None - for i, uid in enumerate(order): + for i in unread: + uid = order[i] data = uid_to_data.get(uid) if data is None: if strict: 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}") - continue # leaves ptrs[i] NULL and tensors[i] None - if describe: - ptrs[i], tensors[i] = self._describe(data, uid) - else: - # The backend reads geometry from its own descriptors; building - # a Tensor it will not look at is pure cost. - ptrs[i] = self._device_pointer(data) + continue # an optional port the caller did not request + ptr, tensor = self._describe(data, uid) + native.set_slot(i, ptr, tuple(tensor.dim), tuple(tensor.stride), *_dlpack_code_bits(tensor.data_type)) + if strict: + for i, uid in enumerate(order): + if not native.is_filled(i): + 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}") # 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: - if describe: - workspace_ptr, workspace_tensor = self._describe(workspace, -1) - workspace_bytes = _byte_size(workspace_tensor) - else: - workspace_ptr = self._device_pointer(workspace) - return VariantPack( - tuple(order), - tuple(tensors) if describe else None, - ptrs, - workspace_ptr, - workspace_bytes, - ) + workspace_ptr, workspace_tensor = self._describe(workspace, -1) + workspace_bytes = _byte_size(workspace_tensor) + return VariantPack(tuple(order), native, workspace_ptr, workspace_bytes) def _describe(self, data: Any, uid: int): """``(pointer, Tensor)`` for one caller buffer. diff --git a/python/cudnn/datatypes.py b/python/cudnn/datatypes.py index 1e738263a..7007f4e3f 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -328,3 +328,29 @@ 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, and back. The native +# variant pack speaks DLPack, so this is the one translation between it and the +# graph's vocabulary. +_CUDNN_TO_DLPACK_CODE_BITS = {} +_FROST_DTYPE_CODE_TO_CUDNN = {} + + +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 + _FROST_DTYPE_CODE_TO_CUDNN[code_bits] = enum + + +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 f7c8c853b..979656248 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -43,7 +43,6 @@ def execute(self, graph, uid_to_data, ctx): ... # write results into caller-provided output buffers """ -import ctypes from abc import ABC from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List @@ -121,36 +120,73 @@ class VariantPack: 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: reading a buffer through its type's ``__dlpack_c_exchange_api__`` + vtable is 0.08 us against 1.5 to ask a python object the same four + questions, and the slots it hands a kernel are read back through the same + vtable — 0.30 us, cheaper than the caller's torch tensor at 0.35, so + refusing to pass the caller's object through costs nothing. + ``uids`` is ASCENDING, matching the backend's own operand order - (``get_variant_pack_uids_sorted()``), so ``ctypes.addressof(ptrs)`` goes - straight to ``_execute_with_raw_ptrs`` with no copy and no per-operand - hash lookup. - - ``tensors[i]`` is a ``graph_types.Tensor`` — the same class ``graph.tensor()`` - returns — describing what the caller ACTUALLY passed for ``uids[i]``: its - ``dim`` / ``stride`` / ``data_type`` are the buffer's, which is not - necessarily what the graph declared. An engine reads whichever it means — - the IR port for the shape the plan was built for, this one for the shape - about to run. frost_gemm takes its M/N/K from here; the backend takes only - the pointer. + (``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 array would hand each thread the other's + 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", "tensors", "ptrs", "address", "_slot_of", "workspace", "workspace_bytes", "_device") + __slots__ = ("uids", "native", "_tensors", "_slot_of", "workspace", "workspace_bytes", "_device") - def __init__(self, uids, tensors, ptrs, workspace_ptr: int, workspace_bytes: int = 0): + def __init__(self, uids, native, workspace_ptr: int = 0, workspace_bytes: int = 0): self.uids = uids - self.tensors = tensors - self.ptrs = ptrs - self.address = ctypes.addressof(ptrs) + self.native = native self.workspace = workspace_ptr self.workspace_bytes = workspace_bytes + self._tensors = None # built on demand: the hot paths read the native slots 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 + + @property + def tensors(self): + """One ``Tensor`` per operand, materialized on first access. + + The paths that run every execute — contiguity, the views a kernel gets, + the pointer array — read the native slots and never come here. This + exists for an engine that wants the geometry as python objects, and + costs 0.29 us per operand to build when it does. + """ + if self._tensors is None: + from ..graph_types import describing_tensor + from ..datatypes import _FROST_DTYPE_CODE_TO_CUDNN + + native = self.native + self._tensors = tuple( + ( + describing_tensor(uid, tuple(native.shape(i)), tuple(native.stride(i)), _FROST_DTYPE_CODE_TO_CUDNN.get(native.dtype(i))) + if native.is_filled(i) + else describing_tensor(uid, (), (), None) + ) + for i, uid in enumerate(self.uids) + ) + return self._tensors + + 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: @@ -179,7 +215,11 @@ def slot(self, tensor_or_uid) -> int: return self.slot_of[uid] def ptr(self, tensor_or_uid) -> int: - return self.ptrs[self.slot(tensor_or_uid)] or 0 + 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 @@ -193,11 +233,7 @@ def view(self, slot: int): which is an argument for making the producer a C type, not for keeping the caller's object. """ - tensor = self.tensors[slot] - name = _CUDNN_TO_FROST_DTYPE_NAME.get(tensor.data_type) - if name is None: - raise ValueError(f"operand uid {self.uids[slot]} has no DLPack-expressible dtype ({tensor.data_type})") - return DeviceView(self.ptrs[slot] or 0, tensor.dim, name, self.device) + return self.native.view(slot, self.device) def __len__(self) -> int: return len(self.uids) diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index 3fb3a46f7..04ea17f1b 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -127,13 +127,22 @@ def over(cls, variant_pack, required_bytes: int, owner: str, *, align: int = DEF 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 it is the + same type: a kernel reads both through the DLPack C exchange vtable + rather than a capsule built per call (0.30 us against 1.86), and a + graph hands its kernels one buffer type rather than two. + """ + import cudnn + 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 cudnn._pybind_module.make_slot(self._ptr + offset, list(shape), code, bits, 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/linear_attention/engine_utils.py b/python/cudnn/linear_attention/engine_utils.py index 02790ad72..d44f366aa 100644 --- a/python/cudnn/linear_attention/engine_utils.py +++ b/python/cudnn/linear_attention/engine_utils.py @@ -67,34 +67,40 @@ 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(): - _check_contiguous(node.name, variant_pack, slots) - node_buffers[node] = NodeBuffers( - {port: variant_pack.view(slot) for port, slot in slots.inputs.items()}, - {port: variant_pack.view(slot) for port, slot in slots.outputs.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(node_name: str, variant_pack, slots) -> None: - """Contiguity gate over every port this node binds, read off the pack. +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. + 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 here would have gone unchecked. + node but forgotten there would have gone unchecked. """ - for direction in (slots.inputs, slots.outputs): - for port, slot in direction.items(): - tensor = variant_pack.tensors[slot] - if not buffers.is_contiguous(tensor.dim, tensor.stride): - raise ValueError(f"cudnn.frost {node_name!r}: buffer for {port!r} must be contiguous (buffers pass straight to the kernel)") + 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/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/variant_pack.cpp b/python/pygraph/variant_pack.cpp new file mode 100644 index 000000000..d76291b87 --- /dev/null +++ b/python/pygraph/variant_pack.cpp @@ -0,0 +1,646 @@ +// 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. +// +// execute() reads the caller's operands once and everything below works from +// the result. Doing that in python cost 1.5 us per operand -- 0.6 to ask the +// buffer for its facts one method call at a time, 0.9 to build the Tensor that +// carries them -- and each operand then had to be turned back into a DLPack +// producer for the kernel, which python cannot do quickly: tvm-ffi reads a +// torch tensor through a C function table and any python producer through a +// freshly built capsule, 0.32 us against 1.91. +// +// Both halves are the same protocol. `__dlpack_c_exchange_api__` is a vtable +// on the TYPE whose dltensor_from_py_object_no_sync fills a caller-provided +// DLTensor in place, with no capsule and no allocation. This file consumes it +// to read the caller's buffers and implements it so the slots it hands out are +// read the same way. Measured on SM100, eight operands: reading 0.64 us +// against 10.4, contiguity 0.13 against 4.2, and a slot converts in 0.27 -- +// slightly cheaper than the torch tensor it replaces, so 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. The degradation is per operand, +// so a pack mixing torch with something else is exactly as fast as 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(); + } + if (cache.count < kTypeCacheSlots) { + cache.types[cache.count] = type; + cache.apis[cache.count] = api; // a null answer is worth caching too + 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(const Slot &slot, int32_t device_id) : slot_(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; + } + + 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; + } + + // 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); + } + + // 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))); + auto *managed = static_cast(std::calloc(1, sizeof(DLManagedTensorVersioned))); + if (managed == nullptr) { + PyErr_NoMemory(); + return -1; + } + managed->version.major = DLPACK_MAJOR_VERSION; + managed->version.minor = DLPACK_MINOR_VERSION; + managed->dl_tensor = slot->tensor(); + managed->manager_ctx = nullptr; + managed->deleter = [](DLManagedTensorVersioned *self) { std::free(self); }; + *out = managed; + 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. + std::vector + read_all(py::sequence buffers) { + std::vector unread; + const size_t n = py::len(buffers); + for (size_t i = 0; i < n && i < slots_.size(); i++) { + py::handle buffer = buffers[i]; + if (buffer.is_none()) { + skip_slot(i); + } else if (!read_slot(i, buffer)) { + unread.push_back(i); + } + } + return unread; + } + + // 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; + } + + // A slot the caller did not fill: an optional port it did not request. + 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(slot, device_id); +} + +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", &VariantPackSlot::stride) + .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("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."); + + 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_all", &VariantPackNative::read_all) + .def("set_slot", &VariantPackNative::set_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 From 26cc45868cc008c463e9a3db16efb5752185e84c Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 03:03:36 -0700 Subject: [PATCH 06/11] Initialize the dlpack dtype tables from either direction _FROST_DTYPE_CODE_TO_CUDNN was only ever filled by _dlpack_code_bits, which runs on the python fallback. An operand read through the exchange vtable never takes that path, so nothing populated the tables before something came looking for the reverse mapping and every VariantPack.tensors[i].data_type read back None -- silently, and only for the fast path. Verified before: BFLOAT16/FLOAT/INT32 operands all reported data_type=None. After: each reports its own. Nothing on the per-execute path reads this yet, which is why the suites stayed green; frost_gemm will when it migrates. Reported by coderabbit on #547. --- python/cudnn/datatypes.py | 13 +++++++++++++ python/cudnn/engines/base.py | 4 ++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/python/cudnn/datatypes.py b/python/cudnn/datatypes.py index 7007f4e3f..aedf6176d 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -354,3 +354,16 @@ def _dlpack_code_bits(data_type): if not _CUDNN_TO_DLPACK_CODE_BITS: _init_dlpack_dtype_tables() return _CUDNN_TO_DLPACK_CODE_BITS.get(data_type, (0, 0)) + + +def _cudnn_dtype_for_dlpack(code_bits): + """The cuDNN dtype a DLPack ``(code, bits)`` names, or None. + + Both directions initialize the pair, because either can be the first one + asked: an operand read through the exchange vtable never takes the python + fallback, so nothing would have populated the tables before something came + looking for the reverse mapping — and every dtype would have read back None. + """ + if not _FROST_DTYPE_CODE_TO_CUDNN: + _init_dlpack_dtype_tables() + return _FROST_DTYPE_CODE_TO_CUDNN.get(code_bits) diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 979656248..28e914c2d 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -168,12 +168,12 @@ def tensors(self): """ if self._tensors is None: from ..graph_types import describing_tensor - from ..datatypes import _FROST_DTYPE_CODE_TO_CUDNN + from ..datatypes import _cudnn_dtype_for_dlpack native = self.native self._tensors = tuple( ( - describing_tensor(uid, tuple(native.shape(i)), tuple(native.stride(i)), _FROST_DTYPE_CODE_TO_CUDNN.get(native.dtype(i))) + describing_tensor(uid, tuple(native.shape(i)), tuple(native.stride(i)), _cudnn_dtype_for_dlpack(native.dtype(i))) if native.is_filled(i) else describing_tensor(uid, (), (), None) ) From 6d2826cee356a9e4e1f7f1047b8458d676d79f66 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 04:16:05 -0700 Subject: [PATCH 07/11] Carve the workspace in one crossing, and settle the review comments The six regions a GDN forward carves are plan-time constants -- offsets fixed by WorkspaceLayout at build, dtypes and shapes fixed with them -- but every execute rebuilt them one at a time, at 0.9 us each: a bounds check, two dtype table lookups, an int() walk over the shape and a pybind crossing, per region. A carve compiled once at build hands back all six in one crossing (5.5 us to 0.8). GDN forward host time 58.5 to 51.0; the four frost engines all use it. Alongside, the review comments on the PR: - A python plan reached execute() with dynamic-shape overrides used to be handed the raw uid map, which a migrated plan cannot read. It cannot honour the overrides either -- a frost engine bakes the declared extents into the kernel it compiles -- so execute() refuses rather than answer a different problem than the caller asked. ExecutionContext's three override fields go with it: nothing could set them. - VariantPackSlot's DLTensor points into its own vectors, so its copy and move constructors are deleted rather than left to alias. - DeviceView.__dlpack__ handed out a struct it owned behind a no-op deleter, which a consumer outliving the view read after free. It delegates to a slot, whose capsule owns its struct and has a real deleter -- which retires the ctypes prototype machinery the view needed. - The exchange-vtable cache is keyed on a type's address, so it now holds a reference to it. - The reentrancy test gave eight threads one workspace to write. - L0 markers, CUDA gates, and two docs that described deleted code. --- docs/python_graph_and_execution_backends.md | 30 ++++-- python/cudnn/_pygraph.py | 20 ++-- python/cudnn/engines/base.py | 9 +- python/cudnn/frost/buffers.py | 91 +++---------------- python/cudnn/frost/workspace.py | 25 ++++- .../linear_attention/frost/gdn2_engine.py | 24 +++-- .../linear_attention/frost/gdn_engine.py | 59 ++++++++---- .../linear_attention/frost/kda_engine.py | 24 +++-- .../frost/kernel/gdn_prefill_f16.py | 7 +- python/pygraph/variant_pack.cpp | 87 +++++++++++++++++- ..._dlpack_proto.py => test_dlpack_export.py} | 56 ++++++------ test/python/test_import_boundaries.py | 4 + .../python/test_variant_pack_normalization.py | 46 +++++++++- 13 files changed, 307 insertions(+), 175 deletions(-) rename test/python/{test_dlpack_proto.py => test_dlpack_export.py} (50%) diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index 0c6182aad..0e5982c61 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -81,8 +81,10 @@ create_execution_plans([heur_mode.A, ...]) _pygraph.py `build_plan(graph, plan, ctx) → CompiledPlan` (the expensive JIT step, once per graph/plan, cached on the graph), `CompiledPlan.execute(graph, operands, ExecutionContext)` with explicit - handle/stream/workspace/overrides. Simple eager engines implement - `execute()` only. + 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 @@ -96,14 +98,22 @@ 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, a ctypes pointer array -(`address` goes straight to `_execute_with_raw_ptrs`), and one `Tensor` per -operand — the same class `graph.tensor()` returns — holding the buffer's OWN -dim/stride/data_type. It is 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 variant pack's -`Tensor` for the shape about to run.** +`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.** `pack.tensors` +materializes those as `Tensor` records on demand — 17 us for eight operands, so +an engine that only needs pointers and extents should ask the pack directly and +never touch it. Two rules that are easy to break by accident: diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index a009adee7..78e493e2c 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -1735,15 +1735,19 @@ def execute( eng = self.selected_engine if eng is not None: # python engine (plan id in the reserved region) + if overriding: + # A python plan is compiled against the shapes the graph + # declared -- the frost engines read them in their __init__ -- + # so it cannot honour a shape it is told at execute. Refusing + # is the only honest answer: silently running the compiled + # shapes would return numbers for the wrong problem. + raise ValueError( + f"dynamic-shape overrides are a backend-path feature, and this graph selected " + f"the python engine {eng.name!r}, whose plan is compiled for the shapes the " + f"graph declared; rebuild the graph at the shapes you want to run" + ) 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) diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 28e914c2d..f59ee5c71 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`` @@ -105,9 +105,6 @@ 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: diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 0fccaf5b8..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] @@ -105,62 +98,6 @@ def dtype_name(buf) -> str: return str(buf.dtype).split(".")[-1] -class _DLPackProto: - """A filled-in ``DLManagedTensor`` for one (shape, dtype, device), minus the - address. - - Every field but ``data`` is a property of the layout, so building them once - and copying the struct beats assigning nine ctypes fields per call: measured - 1.68 us to fill vs 0.45 us to memmove the 72 bytes. Cached per (shape, - dtype, device) — a graph reuses a handful of layouts across its whole run. - - This is the degenerate case of what the backend already does for kernel - arguments (``src/common/include/runtimeKernel.h``): a prefilled blob plus, - per mutable field, an (offset, uid, UpdateMethod) telling execute what to - write where. Here there is exactly one mutable field, ``data``, at a fixed - offset, and its update method is always POINTER — so the bookkeeping - collapses to a memmove and one assignment. - """ - - __slots__ = ("_template", "_shape_arr", "_nbytes") - - def __init__(self, shape, dtype: str, device_id: int): - ndim = len(shape) - self._shape_arr = (ctypes.c_int64 * max(ndim, 1))(*shape) - code, bits = DTYPES[dtype] - mt = _DLManagedTensor() - mt.dl_tensor.device = _DLDevice(_KDL_CUDA, device_id) - mt.dl_tensor.ndim = ndim - mt.dl_tensor.dtype = _DLDataType(code, bits, 1) - mt.dl_tensor.shape = self._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._template = mt - self._nbytes = ctypes.sizeof(_DLManagedTensor) - - def instantiate(self, ptr: int): - """A fresh struct at ``ptr``. Fresh, not shared: a capsule handed to - CuTe ALIASES the struct rather than copying it, and two threads - executing one graph must not be writing the same one.""" - mt = _DLManagedTensor() - ctypes.memmove(ctypes.byref(mt), ctypes.byref(self._template), self._nbytes) - mt.dl_tensor.data = ctypes.c_void_p(ptr) - return mt - - -_PROTO_CACHE = {} - - -def dlpack_proto(shape, dtype: str, device_id: int) -> _DLPackProto: - key = (tuple(shape), dtype, device_id) - proto = _PROTO_CACHE.get(key) - if proto is None: - proto = _PROTO_CACHE[key] = _DLPackProto(key[0], dtype, device_id) - return proto - - class DeviceView: """Zero-copy DLPack view over a raw CUDA pointer. @@ -177,8 +114,6 @@ def __init__(self, ptr: int, shape, dtype: str, device_id: int): self.shape = tuple(int(s) for s in shape) self.dtype = dtype self._device_id = int(device_id) - self._proto = None - self._live = [] def data_ptr(self) -> int: return self._ptr @@ -235,18 +170,16 @@ def __dlpack_device__(self): return (_KDL_CUDA, self._device_id) def __dlpack__(self, *, stream=None, **_kwargs): - # A fresh struct per call, copied from the layout's prototype and - # re-pointed. Fresh because the consumer ALIASES it — cute's - # from_dlpack keeps the pointer rather than copying the DLTensor — so a - # struct shared between two capsules, or between two threads executing - # one graph, would be read after someone else rewrote it. - if self._proto is None: - self._proto = dlpack_proto(self.shape, self.dtype, self._device_id) - mt = self._proto.instantiate(self._ptr) - # The struct and the prototype's shape array must outlive the capsule: - # the deleter is a no-op, so nothing else keeps them alive. - self._live.append(mt) - return _PyCapsule_New(ctypes.addressof(mt), b"dltensor", None) + """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] + return _pybind_module.make_slot(self._ptr, list(self.shape), code, bits, self._device_id).__dlpack__() class DeviceBuffer(DeviceView): diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index 04ea17f1b..939ecf94a 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,21 @@ 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), ...]``. + + An engine's regions are fixed once :class:`WorkspaceLayout` has run; only + the caller's base pointer arrives per execute. Describing them here instead + of at each :meth:`Workspace.view` is what lets one execute cross into C + once rather than once per region (5.5 us against 0.8 for six). + """ + 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. @@ -135,14 +152,16 @@ def view(self, offset: int, dtype: str, shape): rather than a capsule built per call (0.30 us against 1.86), and a graph hands its kernels one buffer type rather than two. """ - import cudnn - count = 1 for extent in shape: count *= int(extent) self._check_span(offset, count * buffers.DTYPE_ITEMSIZE[dtype]) code, bits = buffers.DTYPES[dtype] - return cudnn._pybind_module.make_slot(self._ptr + offset, list(shape), code, bits, self._device) + 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/linear_attention/frost/gdn2_engine.py b/python/cudnn/linear_attention/frost/gdn2_engine.py index 8531a0652..3bcd0fa8a 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -15,7 +15,7 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout +from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair @@ -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 @@ -163,13 +175,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: stream = stream if stream is not None else 0 ws = workspace - 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)) + 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( @@ -205,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 396afb239..b5fd88d3b 100644 --- a/python/cudnn/linear_attention/frost/gdn_engine.py +++ b/python/cudnn/linear_attention/frost/gdn_engine.py @@ -14,7 +14,7 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout +from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair @@ -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 @@ -216,15 +229,9 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: 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, @@ -241,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( @@ -260,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 @@ -332,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 @@ -357,14 +385,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None): 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, @@ -378,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 @@ -438,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 b7428afe0..85f58be11 100644 --- a/python/cudnn/linear_attention/frost/kda_engine.py +++ b/python/cudnn/linear_attention/frost/kda_engine.py @@ -14,7 +14,7 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout +from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair @@ -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 @@ -174,13 +186,7 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: stream = stream if stream is not None else 0 ws = workspace - 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)) + 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( @@ -223,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/gdn_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py index bc01c368a..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 diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index d76291b87..e124c098a 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -97,6 +97,11 @@ exchange_api_for(PyObject *obj) { if (api == nullptr) PyErr_Clear(); } if (cache.count < kTypeCacheSlots) { + // 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. + Py_INCREF(type); cache.types[cache.count] = type; cache.apis[cache.count] = api; // a null answer is worth caching too cache.count++; @@ -163,7 +168,7 @@ struct Slot { // has it for this too. class VariantPackSlot { public: - VariantPackSlot(const Slot &slot, int32_t device_id) : slot_(slot) { + 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; @@ -173,6 +178,16 @@ class VariantPackSlot { 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_; @@ -552,9 +567,66 @@ make_slot(int64_t ptr, std::vector shape, int dtype_code, int dtype_bit 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(slot, device_id); + return new VariantPackSlot(std::move(slot), device_id); } +// A workspace carve, planned once. Which regions a plan cuts, at what offsets, +// with what dtypes and shapes, is fixed when the engine builds; only the +// caller's base pointer arrives per execute. Describing them here rather than +// at each view() turns one crossing per region into one crossing 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"( @@ -610,6 +682,17 @@ capsule built in python. py::arg("device_id"), "A DLPack producer over memory the caller did not supply -- a workspace carve."); + 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"( diff --git a/test/python/test_dlpack_proto.py b/test/python/test_dlpack_export.py similarity index 50% rename from test/python/test_dlpack_proto.py rename to test/python/test_dlpack_export.py index 964083ecf..db607ad7d 100644 --- a/test/python/test_dlpack_proto.py +++ b/test/python/test_dlpack_export.py @@ -1,28 +1,32 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""A DeviceView's DLPack struct is copied from a per-layout prototype. - -Everything in a ``DLManagedTensor`` except ``data`` is a property of the layout, -so the struct is filled once per (shape, dtype, device) and memmove'd per call: -measured 1.68 us to assign nine ctypes fields against 0.45 us to copy the 72 -bytes. Same shape as the backend's kernel-argument handling -(``src/common/include/runtimeKernel.h``), where a prefilled blob carries an -(offset, uid, UpdateMethod) per mutable field; here there is exactly one -mutable field. - -The struct must be FRESH per capsule even so: cute's ``from_dlpack`` aliases it -rather than copying, so a struct shared between two capsules is read after -someone else re-pointed it. +"""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(): @@ -38,22 +42,22 @@ def test_two_capsules_from_one_view_do_not_share_a_struct(): 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__() - addr = ctypes.pythonapi.PyCapsule_GetPointer - addr.restype = ctypes.c_void_p - addr.argtypes = [ctypes.py_object, ctypes.c_char_p] - assert addr(a, b"dltensor") != addr(b, b"dltensor") + assert _capsule_address(a) != _capsule_address(b) @pytest.mark.L0 -def test_prototypes_are_shared_across_views_of_one_layout(): - """The prototype is the cache; the struct is not.""" - t = torch.zeros(4, 5, dtype=torch.bfloat16, device="cuda") - dev = t.device.index or 0 - v1 = buffers.DeviceView(t.data_ptr(), (4, 5), "bfloat16", dev) - v2 = buffers.DeviceView(t.data_ptr() + 64, (4, 5), "bfloat16", dev) - v1.__dlpack__() - v2.__dlpack__() - assert v1._proto is v2._proto +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 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 index 3bc6c2499..4d0661cdf 100644 --- a/test/python/test_variant_pack_normalization.py +++ b/test/python/test_variant_pack_normalization.py @@ -19,6 +19,8 @@ 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 @@ -99,7 +101,6 @@ def test_execute_is_reentrant(): """ g, _, _ = _matmul_graph() handle = cudnn.create_handle() - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") uids = g._variant_pack_uids() wrong = [0] * 8 @@ -107,6 +108,9 @@ 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) @@ -136,3 +140,43 @@ def test_describing_tensor_matches_the_dataclass(): 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_are_refused_by_a_python_engine(): + """A python plan is compiled for the shapes the graph declared. + + The overrides re-describe a tensor at execute, which only the backend can + act on: a frost engine reads its extents in __init__ and bakes them into + the kernel it compiles. Running the compiled shapes anyway would answer a + different problem than the caller asked, so execute refuses. + """ + 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}") + + 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") + with pytest.raises(ValueError, match="dynamic-shape overrides"): + g.execute(data, ws, override_uids=[q.get_uid()], override_shapes=[[total, h, d]]) From f532986551a319be9411692a695679097c3c94cd Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 04:47:24 -0700 Subject: [PATCH 08/11] Normalize for a migrated plan whether or not overrides are passed The branch that skipped normalization when the caller passed override_uids / shapes / strides handed a migrated plan the raw uid map, which it cannot read. Refusing the overrides instead was wrong: frost_gemm compiles M/N/K symbolically and test_override_shape_frost runs other sizes through this exact call. They are accepted and change nothing here, so the branch goes. Also adds bench_sdpa_gemm_host.py, the counterpart of bench_gdn_host.py for the two engines that have not migrated: frost sdpa fwd 34.5 us, frost gemm 42.9, against GDN's 49.3 -- of which 14.8 is GDN's eight launches, so gemm carries the most host work of the three. --- bench_sdpa_gemm_host.py | 113 ++++++++++++++++++ python/cudnn/_pygraph.py | 26 ++-- .../python/test_variant_pack_normalization.py | 24 ++-- 3 files changed, 140 insertions(+), 23 deletions(-) create mode 100644 bench_sdpa_gemm_host.py diff --git a/bench_sdpa_gemm_host.py b/bench_sdpa_gemm_host.py new file mode 100644 index 000000000..1709063aa --- /dev/null +++ b/bench_sdpa_gemm_host.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host cost of one execute() for the FROST sdpa and gemm engines. + +The counterpart of bench_gdn_host.py for the two engines that have NOT migrated +to the normalized variant pack, so the difference between them and GDN is the +cost the migration is expected to move. Every number is a burst from a drained +queue, swept over burst size: one that climbs with n is the device rate, not +host cost. +""" + +from __future__ import annotations + +import os +import sys +import time + +# the FROST manifest rows are opt-in; set before cudnn reads the manifest +os.environ["CUDNN_FRONTEND_ENABLE_FROST_ENGINES"] = "1" + +import torch # noqa: E402 + +import cudnn # noqa: E402 +from cudnn.engines import is_python_engine # noqa: E402 + +BURSTS = (1, 16, 64) +HALF, F32, BF16 = cudnn.data_type.HALF, cudnn.data_type.FLOAT, cudnn.data_type.BFLOAT16 + + +def burst(fn, n, reps=25): + for _ in range(30): + fn() + torch.cuda.synchronize() + out = [] + for _ in range(reps): + torch.cuda.synchronize() + t0 = time.perf_counter_ns() + for _ in range(n): + fn() + out.append((time.perf_counter_ns() - t0) / n / 1000.0) + torch.cuda.synchronize() + return min(out) + + +def _pin_python(g): + """Select the python engine's plan, or None when none claimed the graph.""" + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + python = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)] + if not python: + return None + g.select_plan(python[0]) + g.check_support() + g.build_plans() + return g + + +def sdpa_case(b=2, h=8, s=256, d=256): + dims, strides = (b, h, s, d), (s * h * d, d, h * d, 1) + g = cudnn.pygraph(io_data_type=HALF, intermediate_data_type=F32, compute_data_type=F32) + q = g.tensor(dim=dims, stride=strides, data_type=HALF, name="q") + k = g.tensor(dim=dims, stride=strides, data_type=HALF, name="k") + v = g.tensor(dim=dims, stride=strides, data_type=HALF, name="v") + o, _ = g.sdpa(name="sdpa", q=q, k=k, v=v, attn_scale=1.0 / d**0.5, is_inference=True, use_causal_mask=True) + o.set_output(True).set_dim(dims).set_stride(strides) + if _pin_python(g) is None: + return None + mk = lambda: torch.randn(b, s, h, d, device="cuda", dtype=torch.float16).transpose(1, 2) # noqa: E731 + data = {q: mk(), k: mk(), v: mk(), o: torch.empty(b, s, h, d, device="cuda", dtype=torch.float16).transpose(1, 2)} + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + return g, data, ws + + +def gemm_case(m=256, n=256, k=128): + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, m, k], stride=[m * k, k, 1], data_type=BF16) + B = g.tensor(name="B", uid=2, dim=[1, k, n], stride=[k * n, 1, k], data_type=BF16) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + if _pin_python(g) is None: + return None + a = torch.randn(1, m, k, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, n, k, dtype=torch.bfloat16, device="cuda") + c = torch.empty(1, m, n, dtype=torch.bfloat16, device="cuda") + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + return g, {1: a, 2: b, 3: c}, ws + + +def report(label, built): + if built is None: + print(f"{label:34s} no python engine claimed this graph") + return + g, data, ws = built + g.execute(data, ws) + torch.cuda.synchronize() + plan = g._compiled_plans[g._plan_index] + row = [burst(lambda: g.execute(data, ws), n) for n in BURSTS] + print(f"{label:34s}" + "".join(f"{v:10.2f}" for v in row) + f" migrated={plan.takes_variant_pack}") + uid_to_data = g._uid_to_data(data) + print(f"{' _uid_to_data':34s}{burst(lambda: g._uid_to_data(data), 64):10.2f}") + print(f"{' _normalize (not on this path yet)':34s}{burst(lambda: g._normalize(uid_to_data, ws), 64):10.2f}") + + +def main(): + print(f"{'':34s}" + "".join(f"{'n=' + str(n):>10s}" for n in BURSTS)) + report("frost sdpa fwd (2,8,256,256)", sdpa_case()) + report("frost gemm (256x256x128)", gemm_case()) + print("\nmin us/call over 25 reps") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 78e493e2c..bfca20d37 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -1735,17 +1735,6 @@ def execute( eng = self.selected_engine if eng is not None: # python engine (plan id in the reserved region) - if overriding: - # A python plan is compiled against the shapes the graph - # declared -- the frost engines read them in their __init__ -- - # so it cannot honour a shape it is told at execute. Refusing - # is the only honest answer: silently running the compiled - # shapes would return numbers for the wrong problem. - raise ValueError( - f"dynamic-shape overrides are a backend-path feature, and this graph selected " - f"the python engine {eng.name!r}, whose plan is compiled for the shapes the " - f"graph declared; rebuild the graph at the shapes you want to run" - ) h = handle if handle is not None else self._handle ctx = ExecutionContext(handle=h, stream=self._resolve_stream(h), workspace=workspace) if self._plan_index not in self._compiled_plans: @@ -1754,10 +1743,17 @@ def execute( self._compiled_plans[self._plan_index] = eng.build_plan(self, self._selected_plan_config, ctx) self._is_built = True plan = self._compiled_plans[self._plan_index] - # Normalize only for a plan that reads the result. Building Tensors - # an engine will not look at is pure cost, and the ones that have - # not migrated still take the caller's objects. - if plan.takes_variant_pack and not overriding: + # Normalize for a plan that reads the result; the ones that have not + # migrated still take the caller's objects. + # + # Overrides do not change this. 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 -- frost_gemm compiles M/N/K symbolically and runs the + # new size from the operands alone. So they are accepted for API + # parity and change nothing here. Branching on them was worse than + # useless: it sent a migrated plan the raw uid map. + if plan.takes_variant_pack: plan.execute(self, self._normalize(uid_to_data, workspace), ctx) else: plan.execute(self, uid_to_data, ctx) diff --git a/test/python/test_variant_pack_normalization.py b/test/python/test_variant_pack_normalization.py index 4d0661cdf..7b03df5f5 100644 --- a/test/python/test_variant_pack_normalization.py +++ b/test/python/test_variant_pack_normalization.py @@ -143,13 +143,13 @@ def test_describing_tensor_matches_the_dataclass(): @pytest.mark.L0 -def test_shape_overrides_are_refused_by_a_python_engine(): - """A python plan is compiled for the shapes the graph declared. +def test_shape_overrides_reach_a_migrated_plan_as_a_pack(): + """Overrides do not change what a python plan is handed. - The overrides re-describe a tensor at execute, which only the backend can - act on: a frost engine reads its extents in __init__ and bakes them into - the kernel it compiles. Running the compiled shapes anyway would answer a - different problem than the caller asked, so execute refuses. + 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 @@ -166,6 +166,8 @@ def test_shape_overrides_are_refused_by_a_python_engine(): 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 = { @@ -178,5 +180,11 @@ def test_shape_overrides_are_refused_by_a_python_engine(): 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") - with pytest.raises(ValueError, match="dynamic-shape overrides"): - g.execute(data, ws, override_uids=[q.get_uid()], override_shapes=[[total, h, d]]) + 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) From 45084030407156de1a53c9d5584377245897b3ba Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 10:53:01 -0700 Subject: [PATCH 09/11] Migrate frost_gemm to the variant pack, and normalize in one crossing frost_gemm read the caller's buffer objects directly, which tied it to whatever framework produced them and -- because the graph declares B as [batch, K, N] while a caller allocates (batch, N, K) -- made it answer a different question than the backend under override_shapes. The two are one fix: the engine reads the pack, and execute() puts the overrides INTO the pack, so an engine honours them without knowing the concept exists. Overrides are re-expressed in the axis order the operand already uses. override_shapes speaks the graph's declaration; the slot holds what the caller's buffer reports. They are the same memory, so they rank their axes the same way by stride, and matching the two rankings gives the permutation. Applying the override verbatim left the pack describing the same bytes in a second language, and reading N off a fixed axis then read K. Test: same graph, same buffers, same override, backend and FROST both against the reference -- the case override shape is FOR, a max allocation with the live shape named per call. The existing coverage only checked FROST against itself, which is why this could diverge unnoticed. Normalization moved into the C pack while the engines were being pointed at it, since every path pays it: - read_from(uid_to_data, uids) does the lookups and the reads in one crossing, retiring the ordered list python built to hand to read_all - read_buffer_extent() reads the workspace through the same vtable; asking python for its size cost as much as reading all eight operands - first_unfilled() replaces a per-operand is_filled loop _normalize 6.0 -> 2.0 us. GDN forward 51.3 -> 45.3, frost gemm 42.9 -> 43.1 (the migration itself is free; what is left is normalization, which every engine now shares). run_resolved lets the engine skip rebuilding the by-object / by-uid / by-name tables on every execute. VariantPackSlot gains permute() and stride(dim) -- the kernel layer calls both on a caller buffer, and neither is visible from the engine directory. --- python/cudnn/_pygraph.py | 108 +++++++++----- python/cudnn/gemm/frost/compiler.py | 17 ++- python/cudnn/gemm/frost/engine.py | 52 +++---- python/cudnn/sdpa/fwd/api_dsl.py | 21 ++- python/pygraph/variant_pack.cpp | 133 +++++++++++++++++- .../gemm/frost/test_frontend_integration.py | 54 +++++++ 6 files changed, 319 insertions(+), 66 deletions(-) diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index bfca20d37..6ddb2b0f5 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -25,6 +25,8 @@ import weakref from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +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 @@ -34,6 +36,30 @@ _LOG = logging.getLogger("cudnn.pygraph") +def _in_axis_order_of(shape, stride, reference_stride): + """``(shape, stride)`` re-expressed in the axis order ``reference_stride`` uses. + + A tensor and its transpose describe the same memory, and the two sides of + an override speak different ones: ``override_shapes`` is in the order the + GRAPH declared (a matmul's B is ``[batch, K, N]``), while the slot holds + the order the caller's buffer reports (B is allocated ``(batch, N, K)``). + Applying the override verbatim would leave the pack describing the same + bytes in a second language, and an engine indexing an extent by position + would read the wrong one. + + Both orders rank their axes the same way by stride — that is what makes + them the same memory — so matching the two rankings 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).""" @@ -1744,17 +1770,13 @@ def execute( self._is_built = True plan = self._compiled_plans[self._plan_index] # Normalize for a plan that reads the result; the ones that have not - # migrated still take the caller's objects. - # - # Overrides do not change this. 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 -- frost_gemm compiles M/N/K symbolically and runs the - # new size from the operands alone. So they are accepted for API - # parity and change nothing here. Branching on them was worse than - # useless: it sent a migrated plan the raw uid map. + # migrated still take the caller's objects. Overrides go INTO the + # pack rather than around it: they are part of describing what this + # execute runs, and an engine reading the pack then agrees with the + # backend without knowing they exist. if plan.takes_variant_pack: - plan.execute(self, self._normalize(uid_to_data, workspace), ctx) + 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 @@ -1828,7 +1850,7 @@ def _variant_pack_uids(self) -> Optional[List[int]]: self._sorted_uids = order return order - def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any): + 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 @@ -1836,21 +1858,28 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any): 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 - n = len(order) - import cudnn - - native = cudnn._pybind_module.VariantPackNative(n) + native = _pybind_module.VariantPackNative(len(order)) # One crossing for the whole pack: the reader is a C function table on # the buffer's type (__dlpack_c_exchange_api__), so an operand costs # 0.08 us there against 1.5 to ask a python object the same four - # questions and build a Tensor to hold the answers. What comes back is - # the slots whose producer does not implement it — described here, at - # the price they always cost, without taking the rest down with them. - unread = native.read_all([uid_to_data.get(uid) for uid in order]) + # questions and build a Tensor to hold the answers. The uid lookups go + # with it — pairing the map with the layout in python cost more than + # the reads did. What comes back is the slots whose producer does not + # implement the protocol; those are described here, at the price they + # always cost, 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 there # is the caller's mistake and is named. A python-only graph's layout is # every wired port, which includes the optional ones (gdn's final_state, @@ -1858,29 +1887,40 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any): # a missing port. strict = self._lowered_graph is not None for i in unread: - uid = order[i] - data = uid_to_data.get(uid) + data = uid_to_data.get(order[i]) if data is None: - if strict: - 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}") - continue # an optional port the caller did not request - ptr, tensor = self._describe(data, uid) + 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: - for i, uid in enumerate(order): - if not native.is_filled(i): - 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}") + 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: + slot_of = {uid: i for i, uid in enumerate(order)} + shapes = override_shapes or () + strides = override_strides or () + 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") + shape = tuple(shapes[j]) if j < len(shapes) else tuple(native.shape(i)) + stride = tuple(strides[j]) if j < len(strides) else tuple(native.stride(i)) + native.override_slot(i, *_in_axis_order_of(shape, stride, 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: - workspace_ptr, workspace_tensor = self._describe(workspace, -1) - workspace_bytes = _byte_size(workspace_tensor) + extent = _pybind_module.read_buffer_extent(workspace) + if extent is None: # a bare address, or a producer without the vtable + workspace_ptr, workspace_tensor = self._describe(workspace, -1) + 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): diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index cb3d2a3e4..c3fca98a3 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1821,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(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: @@ -2816,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,)) diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index eee374a8c..56748da87 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,45 @@ 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 -- + # and not the caller's own objects, which would also tie the engine to + # whatever framework produced them. + 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) + # A FROST executor carves its scratch out of the CALLER's workspace: no + # hidden per-execute allocation, stable pointers, CUDA-graph friendly. + # Workspace.over validates it against what the pack already read. + 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; going + # back through a dict keyed by tensor object only to have the + # compiled plan rebuild its by-object / by-uid / by-name tables is + # work with no answer in it. + 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/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/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index e124c098a..d884b3610 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -211,6 +211,17 @@ class VariantPackSlot { return dense; } + // One axis of it, the way a caller written against a framework tensor asks + // (``stride(-1)`` for the innermost). + 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. @@ -279,6 +290,34 @@ class VariantPackSlot { return new VariantPackSlot(out, tensor_.device.device_id); } + // The same memory with its axes reordered. A kernel layer written against + // framework tensors reaches for this, and unlike reshape it is exact for a + // strided slot too -- it only relabels axes. + 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 @@ -422,6 +461,39 @@ class VariantPackNative { // 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. + // The whole pack from the caller's uid map, in one crossing: look each uid + // up, read what publishes the vtable, and report the rest. Same result as + // building the ordered buffer list in python and handing it to read_all, + // without the list, the comprehension, or the frame around them. + // + // A uid the map does not carry is left unfilled rather than refused here: + // 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. The strict check asks this once + // instead of calling is_filled per operand. + int64_t + first_unfilled() const { + for (size_t i = 0; i < slots_.size(); i++) { + if (!slots_[i].filled) return static_cast(i); + } + return -1; + } + std::vector read_all(py::sequence buffers) { std::vector unread; @@ -456,6 +528,23 @@ class VariantPackNative { } // A slot the caller did not fill: an optional port it did not request. + // Re-describe a slot at the shape this execute is actually running, keeping + // the buffer it was read from. This is what override_shapes means: the + // caller allocated once at a cache shape and names the live shape per call. + // Applying it here rather than in an engine is what keeps the two paths + // answering the same question -- an engine that reads the pack gets 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"); + } + 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; @@ -570,6 +659,23 @@ make_slot(int64_t ptr, std::vector shape, int dtype_code, int dtype_bit return new VariantPackSlot(std::move(slot), device_id); } +// ``(pointer, bytes)`` for a buffer that publishes the vtable, else None. +// +// The workspace is not an operand -- it has no uid and no slot -- but an engine +// still has to bounds-check its carves against it, and asking python for the +// size cost as much as reading all eight operands here. +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(); + 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. Which regions a plan cuts, at what offsets, // with what dtypes and shapes, is fixed when the engine builds; only the // caller's base pointer arrives per execute. Describing them here rather than @@ -640,7 +746,13 @@ capsule built in python. .def_property_readonly("shape", &VariantPackSlot::shape) .def_property_readonly("dtype", &VariantPackSlot::dtype) .def_property_readonly("nbytes", &VariantPackSlot::nbytes) - .def("stride", &VariantPackSlot::stride) + .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) @@ -655,6 +767,17 @@ capsule built in python. } 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__", @@ -682,6 +805,11 @@ capsule built in python. 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. @@ -706,7 +834,10 @@ its parts. .def(py::init()) .def("read_slot", &VariantPackNative::read_slot) .def("read_all", &VariantPackNative::read_all) + .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) 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}") From cc69a2eafe340bad6e6c2442ce4f7b7bac3c1d8e Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 12:37:13 -0700 Subject: [PATCH 10/11] Settle codex's findings, and prototype the gemm gate as one baked table Six findings from a codex review of the variant-pack work, five of them real and two memory-safety: - override_slot took ndim from the shape but stored whatever stride it was given, so a shorter stride array was read ndim deep by any consumer. It is the one place a shape and a stride arrive from two different lists; equal ranks are now required. - slot_managed_from_py_object shallow-copied the slot's DLTensor, whose shape and stride point into the slot's own vectors, with a deleter that freed only the wrapper. A managed tensor is the form a consumer may outlive the producer with, so it now owns copies. Same class of bug as the DeviceView no-op deleter fixed earlier -- the python half was repaired and the C vtable half was not. - read_buffer_extent computed a byte COUNT and Workspace.over read it as a byte RANGE, which is only the same for a dense buffer. A non-dense workspace is now refused rather than carved past its end. - A short override_shapes / override_strides silently kept the original metadata for the entries it did not name, where the backend rejects the request -- the same call, two geometries, decided by plan selection. - A bare device address normalized to a rank-zero unknown-dtype tensor, so an engine reading the pack for its extents failed on an operand form the backend has always accepted. It borrows the graph's declaration now. - The exchange-vtable cache kept null answers, but a type can acquire the vtable later (tvm-ffi installs one on import for torch builds without it), and a graph normalized before that import was pinned to the python fallback for the life of the process. Only hits are cached. The prototype, behind CUDNN_FRONTEND_FROST_GEMM_GATE_TABLE=1: which operand needs what alignment, what major, and how its extents follow M/N/K are settled when the plan compiles, so they are computed once into a table and walked once, instead of rebuilding five lists and walking them four times per execute. Measured on a 256x256x128 matmul: the gate alone 20.0 -> 13.3 us, graph.execute 43.7 -> 37.0. 4122 gemm tests pass with it on. The operands' pointer alignment stays with _alignment_reject rather than being re-checked inline: a gate emits a contract as well as a verdict, and a test matches its wording. One table evaluated twice keeps the message single too. Sizing it first was the point -- frost_gemm_execute_design.md makes the gate's share the acceptance condition, and it is 46% of execute. The same measurement names the next item: _call_positional is 14.1 us against a ~4 us floor of one launch plus one DSL crossing. watch_run.sh returns from a detached run on finished, KILLED or STALLED. The bare `until grep EXIT=` waiter only returns on the first, and a killed job then looks exactly like a running one. --- frost_gemm_execute_design.md | 179 ++++++++++++++++++++++++++ proto_exchange_api/gemm_gate_cost.py | 129 +++++++++++++++++++ proto_exchange_api/tma_stride_gate.py | 84 ++++++++++++ python/cudnn/_pygraph.py | 43 +++++-- python/cudnn/gemm/frost/compiler.py | 129 +++++++++++++++++++ python/pygraph/variant_pack.cpp | 64 ++++++--- watch_run.sh | 46 +++++++ 7 files changed, 646 insertions(+), 28 deletions(-) create mode 100644 frost_gemm_execute_design.md create mode 100644 proto_exchange_api/gemm_gate_cost.py create mode 100644 proto_exchange_api/tma_stride_gate.py create mode 100755 watch_run.sh diff --git a/frost_gemm_execute_design.md b/frost_gemm_execute_design.md new file mode 100644 index 000000000..089ac53bb --- /dev/null +++ b/frost_gemm_execute_design.md @@ -0,0 +1,179 @@ +# frost_gemm's execute path: one gate table, evaluated twice + +What the per-execute path should look like if it is designed from what actually +varies per call, and what that deletes from what is there now. + +Nothing here is implemented. This is the brief for doing it. + +## The question this answers + +`check_support` and the per-execute gate are the same predicate evaluated at two +times — once against the shapes the graph declared, once against the shapes the +call is running. Today they are two bodies of code computing the same facts, so +they can drift, and the execute side rebuilds facts that were settled when the +kernel compiled. + +Dynamic shape is why the checks must exist at execute at all. It is NOT why they +are written twice. `_tma_alignment_reject` is the proof: it is the one predicate +factored as a pure function of `(dtypes, majors, M, N, K)`, and it is called from +both sites — `compiler.py:2427` with the graph's `mm.M/N/K` and `compiler.py:1861` +with the runtime ones. Its docstring says so: *"One rule ... for both the +graph-time and runtime dims."* The other three gates were never factored that +way, so each fact lives twice. + +## What varies per call + +**Baked when the kernel compiles — a runtime value cannot change it** + +- `a_major` / `b_major`: which dim is contiguous, compiled into the TMA + descriptor and the MMA operand descriptor +- every operand's dtype, and the fp4 two-elements-per-byte packing factor +- the epilogue store / aux load vector width, hence each output's and aux's + required alignment (`_output_align_reqs`, `_aux_align_reqs`) +- the modulus in the TMA 16-byte rule (`128 // bits`) +- rank 3, the operand roles and their count, which are SFA/SFB, the block-scale + block size, `is_multi_gemm`, which outputs are `norm2` reductions + +**Free per call** + +- M, N, K, batch +- the pointers +- **the outer strides.** `_contiguous_dim` only asks which dim has stride 1; + nothing constrains the others, and the max-allocation override case depends on + that freedom (a `[m, k]` corner of an `[mb, kb]` buffer has row stride `kb`). + +**Therefore re-checked per call — the intersection** + +| gate | why it re-runs | the baked half | +|---|---|---| +| TMA 16-byte alignment | the contiguous extent IS M/N/K | the modulus, and which of M/N/K is contiguous | +| operand shape agreement | every operand's extents move with M/N/K | A is `(b, M, K/kpack)`, B is `(b, N, K/kpack)` | +| pointer alignment | new call, new buffers | required bytes per role | +| output `tensor_alignment` | it is `min(ptr, stride, shape)`, and stride/shape move | the required vector width | +| layout / major | the buffer is new; its contiguous dim could differ | the wanted major | + +**Recomputed per call although nothing in it can change:** `_output_align_reqs`, +`_aux_align_reqs`, `k_factor`, the `_MAJOR_CONTIGUOUS_DIM[major]` lookup, and +`_finalize_reductions`' `startswith("reduction_")` / `rsplit` walk over +`chain.outputs`. + +## Found while classifying: the TMA gate checks the wrong quantity + +`_tma_alignment_reject` checks `extent * bits % 128 == 0`, where `extent` is K +(k-major) or M/N. Its own docstring says TMA encodes *"the contiguous input +dimension's stride in 16-byte units"* — stride, not extent. The two are the same +number only when rows are dense. + +Since outer strides are free (above), a caller can hand in a `[m, k]` corner of +an `[mb, kb]` allocation whose row stride `kb` is not 16-byte aligned while `k` +is. `k=64, kb=72` at bf16: extent `64*16=1024` passes, row stride +`72*2=144` bytes is not a multiple of 16. The gate accepts it and TMA +mis-strides every row past the first — silently wrong numbers, the failure mode +the gate exists to prevent. + +The existing override test does not discriminate: it uses `kb=128`, where both +quantities are aligned. + +**Not yet reproduced on hardware.** Confirm before fixing: build the override +case with a deliberately unaligned outer stride and compare against the backend. +If confirmed, the gate should take the row stride, which it has (the buffer's +`stride()`), rather than inferring it from the extent. + +## The shape to build + +One table, built where the analyzer already knows these things, evaluated by +both sites: + +```python +# built once, when the plan compiles +gate = GateTable( + operands=[OperandGate(role, axis_of_m_or_n, axis_of_k, major, kpack, + ptr_align, mode, vector_bytes), ...], + tma_modulus=[(role, 128 // bits, which_extent), ...], +) + +# check_support, at the declared shapes +reason = gate.reject(declared_mnk, declared_strides=..., pointers=None) + +# execute, at the runtime shapes +reason = gate.reject(runtime_mnk, slots) +``` + +`reject` is one pass over `operands`. There is no second formulation to drift, +and the baked half is computed once. + +The per-execute path then reads: + +```python +slots = pack.views(self._indices) # one crossing, already there +M, N, K = self._extents(slots) # build-known axes, 3 index reads +if (reason := self._gate.reject((M, N, K), slots)) is not None: + raise ValueError(reason) +self._launch(slots, (M, N, K), stream) +``` + +`self._extents` reading build-recorded axes also closes the class of bug the +override work hit: `shape[1]` / `shape[2]` hard-codes the caller's axis +convention, which is not the graph's. + +## What that deletes + +1. `resolve_variant_pack` and `run_resolved`'s `{id(t): buf}` indirection — the + engine knows the operand order at build, so one `views()` result sliced by + build-time ranges replaces `pull()` and its two dict lookups per operand. +2. The five intermediate lists (`_operands`, the layout comprehension, `_named`, + the SF comprehension, `pairs`) and the four walks over them → one pass. +3. `_output_align_reqs` / `_aux_align_reqs` per execute. +4. `k_factor` and the fp4 branches, evaluated in three places per call. +5. `_finalize_reductions`' string parsing → a build-time index list. +6. The `shape[1]` / `shape[2]` convention assumption. + +## Is pure python good enough? + +For gemm, plausibly yes, and this is the case worth trying it on: **gemm is one +launch**, not GDN's eight. + +Today 43.1 us. The floor is 1 launch (1.85) + one DSL crossing (~2) ≈ 4. +Above that: `graph.execute` entry ~6, `_normalize` 2.0, views ~1, and ~29 of +engine + compiler python. The table design targets that 29; the rest of the +python is already thin. + +If the gate collapses to one pass and the entry shrinks, mid-teens is the +plausible landing zone — the same order as flashinfer's 14.5, which also has no +graph API to pay for. **This is an estimate from subtraction, not a measurement.** +Confirm first: time `run_resolved` minus `_call_positional` to size the gate +machinery alone. If the gate is not most of the 29, this design is aimed at the +wrong thing and the brief should be rewritten before any code moves. + +## Traps + +- **Measure from a drained queue and sweep the burst size.** GDN's device time is + ~52 us against a ~50 us host, so back-to-back timing reads back the device + rate. `proto_exchange_api/fe_floor.py` reported 53 us for a 20 us stage this + way; `fe_floor2.py` is the corrected form. +- **Do not read launch cost out of an nsys trace.** CUPTI adds ~2.2 us per traced + API call: `cudaLaunchKernelEx` is 1.85 us untraced and 4.06 traced. +- **Never overwrite the venv's `.so` while a test run is in flight.** It + segfaults workers, twice observed, and looks exactly like a concurrency bug in + the code under test. +- **Do not trust grep for what a kernel layer calls on a buffer.** `reshape`, + `permute` and `stride(dim)` were each found only by running the whole suite. +- **`cd` to the worktree explicitly.** The shell's cwd resets between commands, + and building from the wrong one produces a `.so` that imports but is missing + symbols. + +## Where to do it + +On this branch. The design leans on the variant pack, `VariantPackSlot` and +`pack.views()`, which are all here; starting elsewhere would mean rebuilding +them or re-deriving the numbers above. Keep it as its own commit so it can be +dropped without touching the migration. + +A fresh session is fine to execute it — that is what this file is for. Read it +and `HANDOFF_variant_pack.md` first; between them nothing above needs +re-deriving. + +``` +note to self: claude::774e8e99-23ad-4a94-be0d-53ed5ee4def9 — "cuDNN FE variant-pack normalization" +cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_gpu/fe_pr1 +``` diff --git a/proto_exchange_api/gemm_gate_cost.py b/proto_exchange_api/gemm_gate_cost.py new file mode 100644 index 000000000..ce1d811e8 --- /dev/null +++ b/proto_exchange_api/gemm_gate_cost.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""How much of frost_gemm's execute is the per-call gate? + +frost_gemm_execute_design.md proposes replacing five intermediate lists and four +walks with one pass over a table built when the plan compiles. That is only +worth doing if the gate is most of what the compiler layer costs, so this sizes +it before any code moves: + + graph.execute - the whole thing + plan.execute - minus graph.execute's entry and _normalize + run_resolved - the gate plus the launch + _call_positional - the launch alone + +gate = run_resolved - _call_positional. +""" + +import os +import sys +import time + +os.environ["CUDNN_FRONTEND_ENABLE_FROST_ENGINES"] = "1" + +import torch # noqa: E402 + +import cudnn # noqa: E402 +from cudnn.engines import is_python_engine # noqa: E402 + +sys.path.insert(0, "/home/scratch.yanxu_gpu/fe_pr1") + +BF16, F32 = cudnn.data_type.BFLOAT16, cudnn.data_type.FLOAT +M = N = 256 +K = 128 + + +def burst(fn, n, reps=25): + for _ in range(40): + fn() + torch.cuda.synchronize() + out = [] + for _ in range(reps): + torch.cuda.synchronize() + t0 = time.perf_counter_ns() + for _ in range(n): + fn() + out.append((time.perf_counter_ns() - t0) / n / 1000.0) + torch.cuda.synchronize() + return min(out) + + +def build(): + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + a = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1], data_type=BF16) + b = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K], data_type=BF16) + c = g.matmul(A=a, B=b, name="mm") + c.set_output(True).set_data_type(BF16).set_uid(3) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + index = next(i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)) + g.select_plan(index) + g.check_support() + g.build_plans() + data = { + 1: torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda"), + 2: torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda"), + 3: torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda"), + } + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + return g, data, ws + + +def main(): + g, data, ws = build() + g.execute(data, ws) + torch.cuda.synchronize() + + from cudnn.engines.base import ExecutionContext, bind_ports # noqa: F401 + + plan = g._compiled_plans[g._plan_index] + compiled = plan._compiled + handle = g._handle + ctx = ExecutionContext(handle=handle, stream=g._resolve_stream(handle), workspace=ws) + uid_to_data = g._uid_to_data(data) + pack = g._normalize(uid_to_data, ws) + + # capture what run_resolved hands _call_positional, to time the launch alone + captured = {} + original = type(compiled)._call_positional + + def spy(self, *args, **kwargs): + captured["args"] = (args, kwargs) + return original(self, *args, **kwargs) + + type(compiled)._call_positional = spy + g.execute(data, ws) + torch.cuda.synchronize() + type(compiled)._call_positional = original + args, kwargs = captured["args"] + + slots = [pack.slot(t.get_uid()) for t in plan._tensors] + resolved = {id(t): v for t, v in zip(plan._tensors, pack.views(slots))} + + rows = [ + ("graph.execute(data, ws)", lambda: g.execute(data, ws)), + (" _uid_to_data", lambda: g._uid_to_data(data)), + (" _normalize", lambda: g._normalize(uid_to_data, ws)), + (" plan.execute(pack)", lambda: plan.execute(g, pack, ctx)), + (" pack.views(slots)", lambda: pack.views(slots)), + (" run_resolved (gate + launch)", lambda: compiled.run_resolved(resolved, stream=ctx.stream)), + (" _call_positional (launch)", lambda: original(compiled, *args, **kwargs)), + ] + + print(f"{'':38s}{'n=1':>9s}{'n=16':>9s}{'n=64':>9s}") + values = {} + for label, fn in rows: + row = [burst(fn, n) for n in (1, 16, 64)] + values[label.strip()] = min(row) + print(f"{label:38s}" + "".join(f"{v:9.2f}" for v in row)) + + gate = values["run_resolved (gate + launch)"] - values["_call_positional (launch)"] + total = values["graph.execute(data, ws)"] + print() + print(f"the gate alone {gate:7.2f} us ({100 * gate / total:.0f}% of execute)") + print(f"everything above it {total - values['run_resolved (gate + launch)']:7.2f} us") + + +if __name__ == "__main__": + main() diff --git a/proto_exchange_api/tma_stride_gate.py b/proto_exchange_api/tma_stride_gate.py new file mode 100644 index 000000000..25142125d --- /dev/null +++ b/proto_exchange_api/tma_stride_gate.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Does the TMA alignment gate check the extent where TMA cares about the stride? + +_tma_alignment_reject checks ``extent * bits % 128 == 0`` -- K for a k-major +operand. Its docstring says TMA encodes "the contiguous input dimension's +STRIDE in 16-byte units". Those are the same number only when rows are dense. + +Outer strides are free at runtime: _contiguous_dim only asks which dim has +stride 1, and the max-allocation override case depends on that freedom. So a +caller can name a [m, k] corner of an [mb, kb] allocation whose ROW STRIDE is +not 16-byte aligned while k is. + + k = 64 at bf16 -> 64 * 2 = 128 bytes, aligned + kb = 72 at bf16 -> 72 * 2 = 144 bytes, NOT a multiple of 16 + +If the gate accepts that and the result disagrees with the backend, the gate is +checking the wrong quantity. Run twice, with CUDNN_FRONTEND_ENABLE_FROST_ENGINES +unset and set to 1. +""" + +import os + +import torch + +import cudnn + +BF16, F32 = cudnn.data_type.BFLOAT16, cudnn.data_type.FLOAT +MB, NB, KB = 256, 256, 72 # KB * 2 bytes = 144, not 16-byte aligned +M, N, K = 128, 192, 64 # K * 2 bytes = 128, aligned + + +def small(*shape): + return torch.empty(*shape, dtype=torch.int32).random_(-2, 2).to(torch.bfloat16).cuda() + + +def run(a, b, c): + g = cudnn.pygraph( + io_data_type=BF16, + intermediate_data_type=F32, + compute_data_type=F32, + 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=BF16) + B = g.tensor(name="B", uid=2, dim=[1, KB, NB], stride=[KB * NB, 1, KB], data_type=BF16) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + g.check_support() + g.build_plans() + name = g.get_plan_name_at_index(g._plan_index) + + handle = cudnn.create_handle() + 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]] + wsz = g.get_workspace_size_plan_at_index(g._plan_index, handle, uids, shapes, strides) + ws = torch.empty(max(wsz, 1), device="cuda", dtype=torch.uint8) + c.zero_() + g.execute({1: a, 2: b, 3: c}, ws, handle=handle, override_uids=uids, override_shapes=shapes, override_strides=strides) + torch.cuda.synchronize() + return name + + +def main(): + tag = "FROST on " if os.environ.get("CUDNN_FRONTEND_ENABLE_FROST_ENGINES") == "1" else "FROST off" + a, b = small(1, MB, KB), small(1, NB, KB) + c = torch.zeros(1, MB, NB, dtype=torch.bfloat16, device="cuda") + ref = torch.einsum("bmk,bnk->bmn", a[:, :M, :K].float(), b[:, :N, :K].float()).to(torch.bfloat16) + try: + name = run(a, b, c) + except Exception as exc: + print(f"{tag} | refused: {type(exc).__name__}: {str(exc)[:150]}") + return + got = c[:, :M, :N] + err = (got.float() - ref.float()).abs().max().item() + print(f"{tag} | {name:14s} accepted the launch; correct={torch.equal(got, ref)} max|d|={err:.1f}") + + +if __name__ == "__main__": + main() diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 6ddb2b0f5..0219b0f7c 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -36,6 +36,16 @@ _LOG = logging.getLogger("cudnn.pygraph") +def _is_dense(dim, stride) -> bool: + """Row-major compact, the way `frost.buffers.is_contiguous` reads it.""" + 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. @@ -1900,24 +1910,32 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= 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, so a short list must not + # quietly mean "keep what you had for the rest" here: that is the + # same call answering two ways depending on which plan ran. + 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)} - shapes = override_shapes or () - strides = override_strides or () 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") - shape = tuple(shapes[j]) if j < len(shapes) else tuple(native.shape(i)) - stride = tuple(strides[j]) if j < len(strides) else tuple(native.stride(i)) - native.override_slot(i, *_in_axis_order_of(shape, stride, native.stride(i))) + 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, or a producer without the vtable + 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 @@ -1943,11 +1961,16 @@ def _describe(self, data: Any, uid: int): protocol cannot spell bf16 — which is why the last branch is last and not the only one. - A caller who passed a bare address gets a Tensor with no geometry: a - pointer carries none, and the backend has always accepted that. + 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: - return data, Tensor(uid=uid) # bare device address + 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 diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index c3fca98a3..2e3a27b57 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -70,6 +70,10 @@ def _as_custream(stream): _TVM_FFI_OK = importlib.util.find_spec("tvm_ffi") is not None _FROST_COMPILE_OPTIONS = "--enable-tvm-ffi" if _TVM_FFI_OK else "" +# prototype: evaluate the per-execute gate from a table built once at compile +# time instead of rebuilding five lists and walking them four times. +_GATE_TABLE = os.environ.get("CUDNN_FRONTEND_FROST_GEMM_GATE_TABLE") == "1" + # --------------------------------------------------------------------------- # Symbolic-shape helpers for aux fake tensors @@ -1841,6 +1845,9 @@ def pull(t, role): raise KeyError(f"variant pack is missing a buffer for {role}") return resolved[id(t)] + if _GATE_TABLE: + return self._run_gated(resolved, pull, stream) + a_bufs = [pull(t, "A operand") for t in b.a_operands] b_bufs = [pull(t, "B operand") for t in b.b_operands] out_bufs = [pull(t, "output") for t in b.outputs] @@ -1919,6 +1926,128 @@ def pull(t, role): _finalize_reductions(self.chain, out_bufs) return r + # ---- prototype: one gate table, built once, walked once -------------- + # + # The gate above rebuilds five lists and walks them four times, and + # recomputes two alignment tables that only depend on facts the kernel + # baked. Which operand needs what alignment, what major, and how its + # extents follow M/N/K are all settled when the plan compiles; only the + # extents and the pointers arrive per call. This builds that table on the + # first execute and evaluates it in one pass. + # + # Guarded by CUDNN_FRONTEND_FROST_GEMM_GATE_TABLE=1 while it is measured + # against the path it replaces. + + def _gate_table(self): + table = getattr(self, "_gate_cache", None) + if table is not None: + return table + b = self.binding + mm = self.chain.matmul + out_reqs = _output_align_reqs(self.chain, self.use_tma_store, vec_bytes=self.vec_bytes_epi) + aux_reqs = _aux_align_reqs(self.chain, vec_bytes=self.vec_bytes_epi) + a_pack = 2 if mm.a_dtype == "fp4_e2m1" else 1 + b_pack = 2 if mm.b_dtype == "fp4_e2m1" else 1 + + # (role, tensor, extent_is_n, kpack, want_contiguous_dim, major) + # extent_is_n picks M or N for dim 1. + operands = [] + for i, t in enumerate(b.a_operands): + operands.append((f"A operand[{i}]", t, False, a_pack, _MAJOR_CONTIGUOUS_DIM[mm.a_major], mm.a_major)) + for j, t in enumerate(b.b_operands): + operands.append((f"B operand[{j}]", t, True, b_pack, _MAJOR_CONTIGUOUS_DIM[mm.b_major], mm.b_major)) + # every buffer whose alignment _alignment_reject checks, in one table -- + # including the A/B operands, so their rejection message stays that + # function's rather than a second one written here + extras = [] + for i, t in enumerate(b.a_operands): + extras.append(("A operand", t, 16, "ptr")) + for j, t in enumerate(b.b_operands): + extras.append(("B operand", t, 16, "ptr")) + for k, t in enumerate(b.outputs): + extras.append((f"output[{k}]", t, out_reqs[k], "full")) + for k, t in enumerate(b.aux): + extras.append((f"aux {self.aux_names[k]!r}", t, aux_reqs[self.aux_names[k]], "full")) + if self.block_scale: + for t in b.sfa_operands: + extras.append(("SFA", t, 16, "ptr")) + for t in b.sfb_operands: + extras.append(("SFB", t, 16, "ptr")) + table = self._gate_cache = (operands, extras, mm, a_pack) + return table + + def _run_gated(self, resolved, pull, stream): + operands, extras, mm, a_pack = self._gate_table() + b = self.binding + a_bufs = [pull(t, "A operand") for t in b.a_operands] + b_bufs = [pull(t, "B operand") for t in b.b_operands] + out_bufs = [pull(t, "output") for t in b.outputs] + aux_bufs = [pull(t, f"aux {self.aux_names[i]!r}") for i, t in enumerate(b.aux)] + + M = a_bufs[0].shape[1] + K = a_bufs[0].shape[2] * a_pack + N = b_bufs[0].shape[1] + + reason = _tma_alignment_reject(mm.a_dtype, mm.b_dtype, mm.a_major, mm.b_major, M, N, K) + if reason is not None: + raise ValueError(reason) + + # one pass: shape agreement and layout together + shape_bad, layout_bad = [], [] + for role, tensor, extent_is_n, kpack, want_dim, major in operands: + buf = resolved[id(tensor)] + shape = tuple(buf.shape) + if len(shape) != 3: + shape_bad.append(f"{role}: expected a rank-3 buffer, got shape {shape}") + continue + want = ((N if extent_is_n else M), K // kpack) + if (shape[1], shape[2]) != want: + shape_bad.append(f"{role}: expected (batch, {want[0]}, {want[1]}), got {shape}") + strides = tuple(buf.stride()) + unit = [i for i, s in enumerate(strides) if s == 1 and shape[i] > 1] + got = unit[0] if len(unit) == 1 else None + if got is not None and got != want_dim: + names = {0: "batch", 1: "M/N", 2: "K"} + layout_bad.append( + f"{role}: graph declares {major}-major (dim {want_dim} contiguous) but the buffer has dim {got} ({names[got]}) contiguous, stride={strides}" + ) + if shape_bad: + raise ValueError(f"runtime operand shapes disagree with the inferred problem size (M={M}, N={N}, K={K}): " + "; ".join(shape_bad)) + if layout_bad: + raise ValueError("runtime operand layout does not match the layout the kernel was compiled for: " + "; ".join(layout_bad)) + + reason = _alignment_reject([(role, resolved[id(t)], req, mode) for role, t, req, mode in extras]) + if reason is not None: + raise ValueError(reason) + + if self.block_scale: + sf_k4 = ((K // self.chain.block_scale.block_size) + 3) // 4 + reason = _sf_blob_reject( + [(f"SFA[{i}]", pull(t, "SFA"), 512 * sf_k4 * ((M + 127) // 128) * int(a_bufs[i].shape[0])) for i, t in enumerate(b.sfa_operands)] + + [(f"SFB[{j}]", pull(t, "SFB"), 512 * sf_k4 * ((N + 127) // 128) * int(b_bufs[j].shape[0])) for j, t in enumerate(b.sfb_operands)] + ) + if reason is not None: + raise ValueError(reason) + + mnk = (M, N, K) + c_arg = out_bufs if len(out_bufs) > 1 else out_bufs[0] + if self.chain.is_multi_gemm: + if self.block_scale: + sfa = [pull(t, "SFA") for t in b.sfa_operands] + sfb = [pull(t, "SFB") for t in b.sfb_operands] + pairs = [((a_bufs[ai], sfa[ai]), (b_bufs[bi], sfb[bi])) for ai, bi in self.chain.gemm_operands] + else: + pairs = [(a_bufs[ai], b_bufs[bi]) for ai, bi in self.chain.gemm_operands] + r = self._call_positional(pairs, c_arg, mnk, *aux_bufs, stream=stream) + elif self.block_scale: + r = self._call_positional( + a_bufs[0], b_bufs[0], c_arg, mnk, pull(b.sfa_operands[0], "SFA"), pull(b.sfb_operands[0], "SFB"), *aux_bufs, stream=stream + ) + else: + r = self._call_positional(a_bufs[0], b_bufs[0], c_arg, mnk, *aux_bufs, stream=stream) + _finalize_reductions(self.chain, out_bufs) + return r + def _call_positional(self, *args, stream=None): # Internal launcher (called by __call__ after variant-pack resolve). # Single-GEMM: (a, b, c, (M,N,K), *aux). Multi-GEMM: first arg is a list diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index d884b3610..774215221 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -96,14 +96,19 @@ exchange_api_for(PyObject *obj) { Py_DECREF(capsule); if (api == nullptr) PyErr_Clear(); } - if (cache.count < kTypeCacheSlots) { - // 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. + // 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; // a null answer is worth caching too + cache.apis[cache.count] = api; cache.count++; } return api; @@ -372,18 +377,27 @@ slot_dltensor_from_py_object(void *py_object, DLTensor *out) { int slot_managed_from_py_object(void *py_object, DLManagedTensorVersioned **out) { - auto *slot = py::cast(py::handle(static_cast(py_object))); - auto *managed = static_cast(std::calloc(1, sizeof(DLManagedTensorVersioned))); - if (managed == nullptr) { - PyErr_NoMemory(); - return -1; - } - managed->version.major = DLPACK_MAJOR_VERSION; - managed->version.minor = DLPACK_MINOR_VERSION; - managed->dl_tensor = slot->tensor(); - managed->manager_ctx = nullptr; - managed->deleter = [](DLManagedTensorVersioned *self) { std::free(self); }; - *out = managed; + 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; } @@ -540,6 +554,15 @@ class VariantPackNative { 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 DLTensor's stride array is read + // ndim deep, so a shorter stride would be read past its end by any + // consumer -- and this is the one place a shape and a stride arrive + // from two 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); @@ -670,6 +693,11 @@ read_buffer_extent(py::handle buffer) { 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 for a dense buffer; a strided or + // broadcast one covers more (or, at stride 0, far less) than its element + // count says, and a carve bounds-checked against that would write outside + // the allocation. Hand it back to python to describe and refuse. + 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; diff --git a/watch_run.sh b/watch_run.sh new file mode 100755 index 000000000..610f5a5d6 --- /dev/null +++ b/watch_run.sh @@ -0,0 +1,46 @@ +#!/bin/bash +# Watch a detached run and return on ANY of: finished, died, stalled. +# +# watch_run.sh [poll-seconds] +# +# The plain `until grep -q "^EXIT=" log; do sleep; done` waiter only returns when +# the job writes EXIT=. If the job is killed -- OOM, a SIGHUP from the launching +# shell, a segfaulting worker taking the session down -- that line never appears +# and the wait never ends, which is indistinguishable from "still running" and +# has stalled several turns. This returns in that case too, and says which. +# +# Exit codes: 0 finished, 1 died without EXIT=, 2 log stopped growing. +set -u +log=$1 +pattern=$2 +poll=${3:-300} +stall_limit=3 # consecutive polls with no new output before calling it stalled + +previous="" +stalls=0 +while true; do + if grep -q "^EXIT=" "$log" 2>/dev/null; then + echo "=== FINISHED after $(( SECONDS / 60 ))m ===" + tail -25 "$log" + exit 0 + fi + if ! pgrep -f "$pattern" > /dev/null 2>&1; then + echo "=== DIED after $(( SECONDS / 60 ))m: no process matching '$pattern', and the log has no EXIT= ===" + echo "the job was killed rather than finishing; the log's last lines:" + tail -25 "$log" 2>/dev/null + exit 1 + fi + size=$(stat -c %s "$log" 2>/dev/null || echo 0) + if [ "$size" = "$previous" ]; then + stalls=$(( stalls + 1 )) + else + stalls=0 + fi + if [ "$stalls" -ge "$stall_limit" ]; then + echo "=== STALLED: $log has not grown in $(( stall_limit * poll / 60 )) minutes, process still alive ===" + tail -25 "$log" 2>/dev/null + exit 2 + fi + previous=$size + sleep "$poll" +done From 16bd94969f9d49e2d2df1224e50e1e50fa959604 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 11 Aug 2026 15:22:16 -0700 Subject: [PATCH 11/11] Drop what nothing consumes, and stop explaining the obvious VariantPack.tensors materialized a Tensor record per operand for an engine that wanted the geometry as python objects. frost_gemm was that engine -- and when it was migrated it read the slots directly, which is cheaper and does not name a dtype the graph's vocabulary has to translate. So the property had no caller, and neither did the reverse dtype table built for it. A bug fixed in it earlier in this branch was a bug in code nobody ran, which is why it survived. read_all goes the same way: read_from does the lookups and the reads together, and nothing calls the older entry. The gemm gate-table prototype moves to its own branch. It is guarded by an env var and off by default, so in this PR it is a diff a reviewer has to read and cannot benefit from. The benchmark and watchdog scripts, and a design note for work that is not in this PR, come out of the tree entirely. Comments trimmed throughout: measurements and history belong in this description, not at a call site. What is left is the non-obvious and load-bearing -- why the override has to be re-expressed in the operand's axis order, why a byte count is not a byte range, why a managed tensor may not point at the slot's vectors. --- bench_sdpa_gemm_host.py | 113 ------------ docs/python_graph_and_execution_backends.md | 5 +- frost_gemm_execute_design.md | 179 -------------------- proto_exchange_api/gemm_gate_cost.py | 129 -------------- proto_exchange_api/tma_stride_gate.py | 84 --------- python/cudnn/_pygraph.py | 49 ++---- python/cudnn/datatypes.py | 21 +-- python/cudnn/engines/base.py | 40 +---- python/cudnn/frost/workspace.py | 19 +-- python/cudnn/gemm/frost/compiler.py | 129 -------------- python/cudnn/gemm/frost/engine.py | 16 +- python/pygraph/variant_pack.cpp | 103 +++-------- watch_run.sh | 46 ----- 13 files changed, 65 insertions(+), 868 deletions(-) delete mode 100644 bench_sdpa_gemm_host.py delete mode 100644 frost_gemm_execute_design.md delete mode 100644 proto_exchange_api/gemm_gate_cost.py delete mode 100644 proto_exchange_api/tma_stride_gate.py delete mode 100755 watch_run.sh diff --git a/bench_sdpa_gemm_host.py b/bench_sdpa_gemm_host.py deleted file mode 100644 index 1709063aa..000000000 --- a/bench_sdpa_gemm_host.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Host cost of one execute() for the FROST sdpa and gemm engines. - -The counterpart of bench_gdn_host.py for the two engines that have NOT migrated -to the normalized variant pack, so the difference between them and GDN is the -cost the migration is expected to move. Every number is a burst from a drained -queue, swept over burst size: one that climbs with n is the device rate, not -host cost. -""" - -from __future__ import annotations - -import os -import sys -import time - -# the FROST manifest rows are opt-in; set before cudnn reads the manifest -os.environ["CUDNN_FRONTEND_ENABLE_FROST_ENGINES"] = "1" - -import torch # noqa: E402 - -import cudnn # noqa: E402 -from cudnn.engines import is_python_engine # noqa: E402 - -BURSTS = (1, 16, 64) -HALF, F32, BF16 = cudnn.data_type.HALF, cudnn.data_type.FLOAT, cudnn.data_type.BFLOAT16 - - -def burst(fn, n, reps=25): - for _ in range(30): - fn() - torch.cuda.synchronize() - out = [] - for _ in range(reps): - torch.cuda.synchronize() - t0 = time.perf_counter_ns() - for _ in range(n): - fn() - out.append((time.perf_counter_ns() - t0) / n / 1000.0) - torch.cuda.synchronize() - return min(out) - - -def _pin_python(g): - """Select the python engine's plan, or None when none claimed the graph.""" - g.validate() - g.build_operation_graph() - g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - python = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)] - if not python: - return None - g.select_plan(python[0]) - g.check_support() - g.build_plans() - return g - - -def sdpa_case(b=2, h=8, s=256, d=256): - dims, strides = (b, h, s, d), (s * h * d, d, h * d, 1) - g = cudnn.pygraph(io_data_type=HALF, intermediate_data_type=F32, compute_data_type=F32) - q = g.tensor(dim=dims, stride=strides, data_type=HALF, name="q") - k = g.tensor(dim=dims, stride=strides, data_type=HALF, name="k") - v = g.tensor(dim=dims, stride=strides, data_type=HALF, name="v") - o, _ = g.sdpa(name="sdpa", q=q, k=k, v=v, attn_scale=1.0 / d**0.5, is_inference=True, use_causal_mask=True) - o.set_output(True).set_dim(dims).set_stride(strides) - if _pin_python(g) is None: - return None - mk = lambda: torch.randn(b, s, h, d, device="cuda", dtype=torch.float16).transpose(1, 2) # noqa: E731 - data = {q: mk(), k: mk(), v: mk(), o: torch.empty(b, s, h, d, device="cuda", dtype=torch.float16).transpose(1, 2)} - ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) - return g, data, ws - - -def gemm_case(m=256, n=256, k=128): - g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) - A = g.tensor(name="A", uid=1, dim=[1, m, k], stride=[m * k, k, 1], data_type=BF16) - B = g.tensor(name="B", uid=2, dim=[1, k, n], stride=[k * n, 1, k], data_type=BF16) - C = g.matmul(A=A, B=B, name="mm") - C.set_output(True).set_data_type(BF16).set_uid(3) - if _pin_python(g) is None: - return None - a = torch.randn(1, m, k, dtype=torch.bfloat16, device="cuda") - b = torch.randn(1, n, k, dtype=torch.bfloat16, device="cuda") - c = torch.empty(1, m, n, dtype=torch.bfloat16, device="cuda") - ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) - return g, {1: a, 2: b, 3: c}, ws - - -def report(label, built): - if built is None: - print(f"{label:34s} no python engine claimed this graph") - return - g, data, ws = built - g.execute(data, ws) - torch.cuda.synchronize() - plan = g._compiled_plans[g._plan_index] - row = [burst(lambda: g.execute(data, ws), n) for n in BURSTS] - print(f"{label:34s}" + "".join(f"{v:10.2f}" for v in row) + f" migrated={plan.takes_variant_pack}") - uid_to_data = g._uid_to_data(data) - print(f"{' _uid_to_data':34s}{burst(lambda: g._uid_to_data(data), 64):10.2f}") - print(f"{' _normalize (not on this path yet)':34s}{burst(lambda: g._normalize(uid_to_data, ws), 64):10.2f}") - - -def main(): - print(f"{'':34s}" + "".join(f"{'n=' + str(n):>10s}" for n in BURSTS)) - report("frost sdpa fwd (2,8,256,256)", sdpa_case()) - report("frost gemm (256x256x128)", gemm_case()) - print("\nmin us/call over 25 reps") - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/docs/python_graph_and_execution_backends.md b/docs/python_graph_and_execution_backends.md index 0e5982c61..149735ce1 100644 --- a/docs/python_graph_and_execution_backends.md +++ b/docs/python_graph_and_execution_backends.md @@ -110,10 +110,7 @@ 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.** `pack.tensors` -materializes those as `Tensor` records on demand — 17 us for eight operands, so -an engine that only needs pointers and extents should ask the pack directly and -never touch it. +built for; read the pack for the shape about to run.** Two rules that are easy to break by accident: diff --git a/frost_gemm_execute_design.md b/frost_gemm_execute_design.md deleted file mode 100644 index 089ac53bb..000000000 --- a/frost_gemm_execute_design.md +++ /dev/null @@ -1,179 +0,0 @@ -# frost_gemm's execute path: one gate table, evaluated twice - -What the per-execute path should look like if it is designed from what actually -varies per call, and what that deletes from what is there now. - -Nothing here is implemented. This is the brief for doing it. - -## The question this answers - -`check_support` and the per-execute gate are the same predicate evaluated at two -times — once against the shapes the graph declared, once against the shapes the -call is running. Today they are two bodies of code computing the same facts, so -they can drift, and the execute side rebuilds facts that were settled when the -kernel compiled. - -Dynamic shape is why the checks must exist at execute at all. It is NOT why they -are written twice. `_tma_alignment_reject` is the proof: it is the one predicate -factored as a pure function of `(dtypes, majors, M, N, K)`, and it is called from -both sites — `compiler.py:2427` with the graph's `mm.M/N/K` and `compiler.py:1861` -with the runtime ones. Its docstring says so: *"One rule ... for both the -graph-time and runtime dims."* The other three gates were never factored that -way, so each fact lives twice. - -## What varies per call - -**Baked when the kernel compiles — a runtime value cannot change it** - -- `a_major` / `b_major`: which dim is contiguous, compiled into the TMA - descriptor and the MMA operand descriptor -- every operand's dtype, and the fp4 two-elements-per-byte packing factor -- the epilogue store / aux load vector width, hence each output's and aux's - required alignment (`_output_align_reqs`, `_aux_align_reqs`) -- the modulus in the TMA 16-byte rule (`128 // bits`) -- rank 3, the operand roles and their count, which are SFA/SFB, the block-scale - block size, `is_multi_gemm`, which outputs are `norm2` reductions - -**Free per call** - -- M, N, K, batch -- the pointers -- **the outer strides.** `_contiguous_dim` only asks which dim has stride 1; - nothing constrains the others, and the max-allocation override case depends on - that freedom (a `[m, k]` corner of an `[mb, kb]` buffer has row stride `kb`). - -**Therefore re-checked per call — the intersection** - -| gate | why it re-runs | the baked half | -|---|---|---| -| TMA 16-byte alignment | the contiguous extent IS M/N/K | the modulus, and which of M/N/K is contiguous | -| operand shape agreement | every operand's extents move with M/N/K | A is `(b, M, K/kpack)`, B is `(b, N, K/kpack)` | -| pointer alignment | new call, new buffers | required bytes per role | -| output `tensor_alignment` | it is `min(ptr, stride, shape)`, and stride/shape move | the required vector width | -| layout / major | the buffer is new; its contiguous dim could differ | the wanted major | - -**Recomputed per call although nothing in it can change:** `_output_align_reqs`, -`_aux_align_reqs`, `k_factor`, the `_MAJOR_CONTIGUOUS_DIM[major]` lookup, and -`_finalize_reductions`' `startswith("reduction_")` / `rsplit` walk over -`chain.outputs`. - -## Found while classifying: the TMA gate checks the wrong quantity - -`_tma_alignment_reject` checks `extent * bits % 128 == 0`, where `extent` is K -(k-major) or M/N. Its own docstring says TMA encodes *"the contiguous input -dimension's stride in 16-byte units"* — stride, not extent. The two are the same -number only when rows are dense. - -Since outer strides are free (above), a caller can hand in a `[m, k]` corner of -an `[mb, kb]` allocation whose row stride `kb` is not 16-byte aligned while `k` -is. `k=64, kb=72` at bf16: extent `64*16=1024` passes, row stride -`72*2=144` bytes is not a multiple of 16. The gate accepts it and TMA -mis-strides every row past the first — silently wrong numbers, the failure mode -the gate exists to prevent. - -The existing override test does not discriminate: it uses `kb=128`, where both -quantities are aligned. - -**Not yet reproduced on hardware.** Confirm before fixing: build the override -case with a deliberately unaligned outer stride and compare against the backend. -If confirmed, the gate should take the row stride, which it has (the buffer's -`stride()`), rather than inferring it from the extent. - -## The shape to build - -One table, built where the analyzer already knows these things, evaluated by -both sites: - -```python -# built once, when the plan compiles -gate = GateTable( - operands=[OperandGate(role, axis_of_m_or_n, axis_of_k, major, kpack, - ptr_align, mode, vector_bytes), ...], - tma_modulus=[(role, 128 // bits, which_extent), ...], -) - -# check_support, at the declared shapes -reason = gate.reject(declared_mnk, declared_strides=..., pointers=None) - -# execute, at the runtime shapes -reason = gate.reject(runtime_mnk, slots) -``` - -`reject` is one pass over `operands`. There is no second formulation to drift, -and the baked half is computed once. - -The per-execute path then reads: - -```python -slots = pack.views(self._indices) # one crossing, already there -M, N, K = self._extents(slots) # build-known axes, 3 index reads -if (reason := self._gate.reject((M, N, K), slots)) is not None: - raise ValueError(reason) -self._launch(slots, (M, N, K), stream) -``` - -`self._extents` reading build-recorded axes also closes the class of bug the -override work hit: `shape[1]` / `shape[2]` hard-codes the caller's axis -convention, which is not the graph's. - -## What that deletes - -1. `resolve_variant_pack` and `run_resolved`'s `{id(t): buf}` indirection — the - engine knows the operand order at build, so one `views()` result sliced by - build-time ranges replaces `pull()` and its two dict lookups per operand. -2. The five intermediate lists (`_operands`, the layout comprehension, `_named`, - the SF comprehension, `pairs`) and the four walks over them → one pass. -3. `_output_align_reqs` / `_aux_align_reqs` per execute. -4. `k_factor` and the fp4 branches, evaluated in three places per call. -5. `_finalize_reductions`' string parsing → a build-time index list. -6. The `shape[1]` / `shape[2]` convention assumption. - -## Is pure python good enough? - -For gemm, plausibly yes, and this is the case worth trying it on: **gemm is one -launch**, not GDN's eight. - -Today 43.1 us. The floor is 1 launch (1.85) + one DSL crossing (~2) ≈ 4. -Above that: `graph.execute` entry ~6, `_normalize` 2.0, views ~1, and ~29 of -engine + compiler python. The table design targets that 29; the rest of the -python is already thin. - -If the gate collapses to one pass and the entry shrinks, mid-teens is the -plausible landing zone — the same order as flashinfer's 14.5, which also has no -graph API to pay for. **This is an estimate from subtraction, not a measurement.** -Confirm first: time `run_resolved` minus `_call_positional` to size the gate -machinery alone. If the gate is not most of the 29, this design is aimed at the -wrong thing and the brief should be rewritten before any code moves. - -## Traps - -- **Measure from a drained queue and sweep the burst size.** GDN's device time is - ~52 us against a ~50 us host, so back-to-back timing reads back the device - rate. `proto_exchange_api/fe_floor.py` reported 53 us for a 20 us stage this - way; `fe_floor2.py` is the corrected form. -- **Do not read launch cost out of an nsys trace.** CUPTI adds ~2.2 us per traced - API call: `cudaLaunchKernelEx` is 1.85 us untraced and 4.06 traced. -- **Never overwrite the venv's `.so` while a test run is in flight.** It - segfaults workers, twice observed, and looks exactly like a concurrency bug in - the code under test. -- **Do not trust grep for what a kernel layer calls on a buffer.** `reshape`, - `permute` and `stride(dim)` were each found only by running the whole suite. -- **`cd` to the worktree explicitly.** The shell's cwd resets between commands, - and building from the wrong one produces a `.so` that imports but is missing - symbols. - -## Where to do it - -On this branch. The design leans on the variant pack, `VariantPackSlot` and -`pack.views()`, which are all here; starting elsewhere would mean rebuilding -them or re-deriving the numbers above. Keep it as its own commit so it can be -dropped without touching the migration. - -A fresh session is fine to execute it — that is what this file is for. Read it -and `HANDOFF_variant_pack.md` first; between them nothing above needs -re-deriving. - -``` -note to self: claude::774e8e99-23ad-4a94-be0d-53ed5ee4def9 — "cuDNN FE variant-pack normalization" -cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_gpu/fe_pr1 -``` diff --git a/proto_exchange_api/gemm_gate_cost.py b/proto_exchange_api/gemm_gate_cost.py deleted file mode 100644 index ce1d811e8..000000000 --- a/proto_exchange_api/gemm_gate_cost.py +++ /dev/null @@ -1,129 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""How much of frost_gemm's execute is the per-call gate? - -frost_gemm_execute_design.md proposes replacing five intermediate lists and four -walks with one pass over a table built when the plan compiles. That is only -worth doing if the gate is most of what the compiler layer costs, so this sizes -it before any code moves: - - graph.execute - the whole thing - plan.execute - minus graph.execute's entry and _normalize - run_resolved - the gate plus the launch - _call_positional - the launch alone - -gate = run_resolved - _call_positional. -""" - -import os -import sys -import time - -os.environ["CUDNN_FRONTEND_ENABLE_FROST_ENGINES"] = "1" - -import torch # noqa: E402 - -import cudnn # noqa: E402 -from cudnn.engines import is_python_engine # noqa: E402 - -sys.path.insert(0, "/home/scratch.yanxu_gpu/fe_pr1") - -BF16, F32 = cudnn.data_type.BFLOAT16, cudnn.data_type.FLOAT -M = N = 256 -K = 128 - - -def burst(fn, n, reps=25): - for _ in range(40): - fn() - torch.cuda.synchronize() - out = [] - for _ in range(reps): - torch.cuda.synchronize() - t0 = time.perf_counter_ns() - for _ in range(n): - fn() - out.append((time.perf_counter_ns() - t0) / n / 1000.0) - torch.cuda.synchronize() - return min(out) - - -def build(): - g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) - a = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1], data_type=BF16) - b = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K], data_type=BF16) - c = g.matmul(A=a, B=b, name="mm") - c.set_output(True).set_data_type(BF16).set_uid(3) - g.validate() - g.build_operation_graph() - g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - index = next(i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)) - g.select_plan(index) - g.check_support() - g.build_plans() - data = { - 1: torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda"), - 2: torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda"), - 3: torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda"), - } - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") - return g, data, ws - - -def main(): - g, data, ws = build() - g.execute(data, ws) - torch.cuda.synchronize() - - from cudnn.engines.base import ExecutionContext, bind_ports # noqa: F401 - - plan = g._compiled_plans[g._plan_index] - compiled = plan._compiled - handle = g._handle - ctx = ExecutionContext(handle=handle, stream=g._resolve_stream(handle), workspace=ws) - uid_to_data = g._uid_to_data(data) - pack = g._normalize(uid_to_data, ws) - - # capture what run_resolved hands _call_positional, to time the launch alone - captured = {} - original = type(compiled)._call_positional - - def spy(self, *args, **kwargs): - captured["args"] = (args, kwargs) - return original(self, *args, **kwargs) - - type(compiled)._call_positional = spy - g.execute(data, ws) - torch.cuda.synchronize() - type(compiled)._call_positional = original - args, kwargs = captured["args"] - - slots = [pack.slot(t.get_uid()) for t in plan._tensors] - resolved = {id(t): v for t, v in zip(plan._tensors, pack.views(slots))} - - rows = [ - ("graph.execute(data, ws)", lambda: g.execute(data, ws)), - (" _uid_to_data", lambda: g._uid_to_data(data)), - (" _normalize", lambda: g._normalize(uid_to_data, ws)), - (" plan.execute(pack)", lambda: plan.execute(g, pack, ctx)), - (" pack.views(slots)", lambda: pack.views(slots)), - (" run_resolved (gate + launch)", lambda: compiled.run_resolved(resolved, stream=ctx.stream)), - (" _call_positional (launch)", lambda: original(compiled, *args, **kwargs)), - ] - - print(f"{'':38s}{'n=1':>9s}{'n=16':>9s}{'n=64':>9s}") - values = {} - for label, fn in rows: - row = [burst(fn, n) for n in (1, 16, 64)] - values[label.strip()] = min(row) - print(f"{label:38s}" + "".join(f"{v:9.2f}" for v in row)) - - gate = values["run_resolved (gate + launch)"] - values["_call_positional (launch)"] - total = values["graph.execute(data, ws)"] - print() - print(f"the gate alone {gate:7.2f} us ({100 * gate / total:.0f}% of execute)") - print(f"everything above it {total - values['run_resolved (gate + launch)']:7.2f} us") - - -if __name__ == "__main__": - main() diff --git a/proto_exchange_api/tma_stride_gate.py b/proto_exchange_api/tma_stride_gate.py deleted file mode 100644 index 25142125d..000000000 --- a/proto_exchange_api/tma_stride_gate.py +++ /dev/null @@ -1,84 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Does the TMA alignment gate check the extent where TMA cares about the stride? - -_tma_alignment_reject checks ``extent * bits % 128 == 0`` -- K for a k-major -operand. Its docstring says TMA encodes "the contiguous input dimension's -STRIDE in 16-byte units". Those are the same number only when rows are dense. - -Outer strides are free at runtime: _contiguous_dim only asks which dim has -stride 1, and the max-allocation override case depends on that freedom. So a -caller can name a [m, k] corner of an [mb, kb] allocation whose ROW STRIDE is -not 16-byte aligned while k is. - - k = 64 at bf16 -> 64 * 2 = 128 bytes, aligned - kb = 72 at bf16 -> 72 * 2 = 144 bytes, NOT a multiple of 16 - -If the gate accepts that and the result disagrees with the backend, the gate is -checking the wrong quantity. Run twice, with CUDNN_FRONTEND_ENABLE_FROST_ENGINES -unset and set to 1. -""" - -import os - -import torch - -import cudnn - -BF16, F32 = cudnn.data_type.BFLOAT16, cudnn.data_type.FLOAT -MB, NB, KB = 256, 256, 72 # KB * 2 bytes = 144, not 16-byte aligned -M, N, K = 128, 192, 64 # K * 2 bytes = 128, aligned - - -def small(*shape): - return torch.empty(*shape, dtype=torch.int32).random_(-2, 2).to(torch.bfloat16).cuda() - - -def run(a, b, c): - g = cudnn.pygraph( - io_data_type=BF16, - intermediate_data_type=F32, - compute_data_type=F32, - 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=BF16) - B = g.tensor(name="B", uid=2, dim=[1, KB, NB], stride=[KB * NB, 1, KB], data_type=BF16) - C = g.matmul(A=A, B=B, name="mm") - C.set_output(True).set_data_type(BF16).set_uid(3) - g.validate() - g.build_operation_graph() - g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - g.check_support() - g.build_plans() - name = g.get_plan_name_at_index(g._plan_index) - - handle = cudnn.create_handle() - 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]] - wsz = g.get_workspace_size_plan_at_index(g._plan_index, handle, uids, shapes, strides) - ws = torch.empty(max(wsz, 1), device="cuda", dtype=torch.uint8) - c.zero_() - g.execute({1: a, 2: b, 3: c}, ws, handle=handle, override_uids=uids, override_shapes=shapes, override_strides=strides) - torch.cuda.synchronize() - return name - - -def main(): - tag = "FROST on " if os.environ.get("CUDNN_FRONTEND_ENABLE_FROST_ENGINES") == "1" else "FROST off" - a, b = small(1, MB, KB), small(1, NB, KB) - c = torch.zeros(1, MB, NB, dtype=torch.bfloat16, device="cuda") - ref = torch.einsum("bmk,bnk->bmn", a[:, :M, :K].float(), b[:, :N, :K].float()).to(torch.bfloat16) - try: - name = run(a, b, c) - except Exception as exc: - print(f"{tag} | refused: {type(exc).__name__}: {str(exc)[:150]}") - return - got = c[:, :M, :N] - err = (got.float() - ref.float()).abs().max().item() - print(f"{tag} | {name:14s} accepted the launch; correct={torch.equal(got, ref)} max|d|={err:.1f}") - - -if __name__ == "__main__": - main() diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 0219b0f7c..25d901b20 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -37,7 +37,7 @@ def _is_dense(dim, stride) -> bool: - """Row-major compact, the way `frost.buffers.is_contiguous` reads it.""" + """Row-major compact.""" expect = 1 for extent, step in zip(reversed(tuple(dim)), reversed(tuple(stride))): if extent != 1 and step != expect: @@ -49,16 +49,11 @@ def _is_dense(dim, stride) -> bool: def _in_axis_order_of(shape, stride, reference_stride): """``(shape, stride)`` re-expressed in the axis order ``reference_stride`` uses. - A tensor and its transpose describe the same memory, and the two sides of - an override speak different ones: ``override_shapes`` is in the order the - GRAPH declared (a matmul's B is ``[batch, K, N]``), while the slot holds - the order the caller's buffer reports (B is allocated ``(batch, N, K)``). - Applying the override verbatim would leave the pack describing the same - bytes in a second language, and an engine indexing an extent by position - would read the wrong one. - - Both orders rank their axes the same way by stride — that is what makes - them the same memory — so matching the two rankings gives the permutation. + ``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) @@ -1779,11 +1774,9 @@ def execute( self._compiled_plans[self._plan_index] = eng.build_plan(self, self._selected_plan_config, ctx) self._is_built = True plan = self._compiled_plans[self._plan_index] - # Normalize for a plan that reads the result; the ones that have not - # migrated still take the caller's objects. Overrides go INTO the - # pack rather than around it: they are part of describing what this - # execute runs, and an engine reading the pack then agrees with the - # backend without knowing they exist. + # 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) @@ -1881,20 +1874,13 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= if order is None: return None native = _pybind_module.VariantPackNative(len(order)) - # One crossing for the whole pack: the reader is a C function table on - # the buffer's type (__dlpack_c_exchange_api__), so an operand costs - # 0.08 us there against 1.5 to ask a python object the same four - # questions and build a Tensor to hold the answers. The uid lookups go - # with it — pairing the map with the layout in python cost more than - # the reads did. What comes back is the slots whose producer does not - # implement the protocol; those are described here, at the price they - # always cost, without taking the rest down with them. + # 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 there - # is the caller's mistake and is named. A python-only graph's layout is - # every wired port, which includes the optional ones (gdn's final_state, - # H); a hole is simply "not requested", and the engine reads it back as - # a missing port. + # 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]) @@ -1910,9 +1896,8 @@ def _normalize(self, uid_to_data: Dict[int, Any], workspace: Any, override_uids= 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, so a short list must not - # quietly mean "keep what you had for the rest" here: that is the - # same call answering two ways depending on which plan ran. + # 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 " diff --git a/python/cudnn/datatypes.py b/python/cudnn/datatypes.py index aedf6176d..b407a1227 100644 --- a/python/cudnn/datatypes.py +++ b/python/cudnn/datatypes.py @@ -330,11 +330,10 @@ def _is_jax_array(input_tensor) -> bool: return type(input_tensor).__module__.startswith(("jax", "jaxlib")) -# The DLPack (code, bits) a cuDNN dtype travels as, and back. The native -# variant pack speaks DLPack, so this is the one translation between it and the -# graph's vocabulary. +# 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 = {} -_FROST_DTYPE_CODE_TO_CUDNN = {} def _init_dlpack_dtype_tables(): @@ -345,7 +344,6 @@ def _init_dlpack_dtype_tables(): if code_bits is None: continue _CUDNN_TO_DLPACK_CODE_BITS[enum] = code_bits - _FROST_DTYPE_CODE_TO_CUDNN[code_bits] = enum def _dlpack_code_bits(data_type): @@ -354,16 +352,3 @@ def _dlpack_code_bits(data_type): if not _CUDNN_TO_DLPACK_CODE_BITS: _init_dlpack_dtype_tables() return _CUDNN_TO_DLPACK_CODE_BITS.get(data_type, (0, 0)) - - -def _cudnn_dtype_for_dlpack(code_bits): - """The cuDNN dtype a DLPack ``(code, bits)`` names, or None. - - Both directions initialize the pair, because either can be the first one - asked: an operand read through the exchange vtable never takes the python - fallback, so nothing would have populated the tables before something came - looking for the reverse mapping — and every dtype would have read back None. - """ - if not _FROST_DTYPE_CODE_TO_CUDNN: - _init_dlpack_dtype_tables() - return _FROST_DTYPE_CODE_TO_CUDNN.get(code_bits) diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index f59ee5c71..46c523a19 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -118,11 +118,8 @@ class VariantPack: caller happened to hold. The operands live in ``native``, a C container holding one ``DLTensor`` - each: reading a buffer through its type's ``__dlpack_c_exchange_api__`` - vtable is 0.08 us against 1.5 to ask a python object the same four - questions, and the slots it hands a kernel are read back through the same - vtable — 0.30 us, cheaper than the caller's torch tensor at 0.35, so - refusing to pass the caller's object through costs nothing. + 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 @@ -138,14 +135,13 @@ class VariantPack: pointers — silently, because every pointer in it is individually valid. """ - __slots__ = ("uids", "native", "_tensors", "_slot_of", "workspace", "workspace_bytes", "_device") + __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._tensors = None # built on demand: the hot paths read the native slots self._slot_of = None # built on first lookup: the backend never does one self._device = None @@ -154,30 +150,6 @@ def address(self) -> int: """The ``void*[]`` in slot order, for ``_execute_with_raw_ptrs``.""" return self.native.address - @property - def tensors(self): - """One ``Tensor`` per operand, materialized on first access. - - The paths that run every execute — contiguity, the views a kernel gets, - the pointer array — read the native slots and never come here. This - exists for an engine that wants the geometry as python objects, and - costs 0.29 us per operand to build when it does. - """ - if self._tensors is None: - from ..graph_types import describing_tensor - from ..datatypes import _cudnn_dtype_for_dlpack - - native = self.native - self._tensors = tuple( - ( - describing_tensor(uid, tuple(native.shape(i)), tuple(native.stride(i)), _cudnn_dtype_for_dlpack(native.dtype(i))) - if native.is_filled(i) - else describing_tensor(uid, (), (), None) - ) - for i, uid in enumerate(self.uids) - ) - return self._tensors - def all_contiguous(self): """``(ok, slot)`` over every filled operand, decided from the strides the native pack already holds.""" @@ -194,10 +166,8 @@ def slot_of(self): def device(self) -> int: """The GPU this execute is going to, for the views handed to kernels. - One per pack, not one per operand: an execute launches on the current - device and cuDNN's own variant pack carries no device at all - (``create_variant_pack`` sets pointers, uids and the workspace). Read - on demand — 0.74 us, and the backend path never asks. + 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 diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index 939ecf94a..f198cfa4e 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -71,10 +71,8 @@ def base_align(self) -> int: def carve_plan(owner: str, regions) -> "_pybind_module.WorkspaceCarve": """Compile a build-time carve: ``[(offset, dtype, shape), ...]``. - An engine's regions are fixed once :class:`WorkspaceLayout` has run; only - the caller's base pointer arrives per execute. Describing them here instead - of at each :meth:`Workspace.view` is what lets one execute cross into C - once rather than once per region (5.5 us against 0.8 for six). + 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: @@ -119,12 +117,7 @@ def _init(self, ptr, nbytes, device, owner, align): @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. - - ``execute()`` measures the caller's workspace with the same reader it - gives every other buffer, so re-probing it here cost 3.5 us to learn - what the pack is holding. - """ + """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: @@ -147,10 +140,8 @@ def nbytes(self) -> int: 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 it is the - same type: a kernel reads both through the DLPack C exchange vtable - rather than a capsule built per call (0.30 us against 1.86), and a - graph hands its kernels one buffer type rather than two. + 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: diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 2e3a27b57..c3fca98a3 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -70,10 +70,6 @@ def _as_custream(stream): _TVM_FFI_OK = importlib.util.find_spec("tvm_ffi") is not None _FROST_COMPILE_OPTIONS = "--enable-tvm-ffi" if _TVM_FFI_OK else "" -# prototype: evaluate the per-execute gate from a table built once at compile -# time instead of rebuilding five lists and walking them four times. -_GATE_TABLE = os.environ.get("CUDNN_FRONTEND_FROST_GEMM_GATE_TABLE") == "1" - # --------------------------------------------------------------------------- # Symbolic-shape helpers for aux fake tensors @@ -1845,9 +1841,6 @@ def pull(t, role): raise KeyError(f"variant pack is missing a buffer for {role}") return resolved[id(t)] - if _GATE_TABLE: - return self._run_gated(resolved, pull, stream) - a_bufs = [pull(t, "A operand") for t in b.a_operands] b_bufs = [pull(t, "B operand") for t in b.b_operands] out_bufs = [pull(t, "output") for t in b.outputs] @@ -1926,128 +1919,6 @@ def pull(t, role): _finalize_reductions(self.chain, out_bufs) return r - # ---- prototype: one gate table, built once, walked once -------------- - # - # The gate above rebuilds five lists and walks them four times, and - # recomputes two alignment tables that only depend on facts the kernel - # baked. Which operand needs what alignment, what major, and how its - # extents follow M/N/K are all settled when the plan compiles; only the - # extents and the pointers arrive per call. This builds that table on the - # first execute and evaluates it in one pass. - # - # Guarded by CUDNN_FRONTEND_FROST_GEMM_GATE_TABLE=1 while it is measured - # against the path it replaces. - - def _gate_table(self): - table = getattr(self, "_gate_cache", None) - if table is not None: - return table - b = self.binding - mm = self.chain.matmul - out_reqs = _output_align_reqs(self.chain, self.use_tma_store, vec_bytes=self.vec_bytes_epi) - aux_reqs = _aux_align_reqs(self.chain, vec_bytes=self.vec_bytes_epi) - a_pack = 2 if mm.a_dtype == "fp4_e2m1" else 1 - b_pack = 2 if mm.b_dtype == "fp4_e2m1" else 1 - - # (role, tensor, extent_is_n, kpack, want_contiguous_dim, major) - # extent_is_n picks M or N for dim 1. - operands = [] - for i, t in enumerate(b.a_operands): - operands.append((f"A operand[{i}]", t, False, a_pack, _MAJOR_CONTIGUOUS_DIM[mm.a_major], mm.a_major)) - for j, t in enumerate(b.b_operands): - operands.append((f"B operand[{j}]", t, True, b_pack, _MAJOR_CONTIGUOUS_DIM[mm.b_major], mm.b_major)) - # every buffer whose alignment _alignment_reject checks, in one table -- - # including the A/B operands, so their rejection message stays that - # function's rather than a second one written here - extras = [] - for i, t in enumerate(b.a_operands): - extras.append(("A operand", t, 16, "ptr")) - for j, t in enumerate(b.b_operands): - extras.append(("B operand", t, 16, "ptr")) - for k, t in enumerate(b.outputs): - extras.append((f"output[{k}]", t, out_reqs[k], "full")) - for k, t in enumerate(b.aux): - extras.append((f"aux {self.aux_names[k]!r}", t, aux_reqs[self.aux_names[k]], "full")) - if self.block_scale: - for t in b.sfa_operands: - extras.append(("SFA", t, 16, "ptr")) - for t in b.sfb_operands: - extras.append(("SFB", t, 16, "ptr")) - table = self._gate_cache = (operands, extras, mm, a_pack) - return table - - def _run_gated(self, resolved, pull, stream): - operands, extras, mm, a_pack = self._gate_table() - b = self.binding - a_bufs = [pull(t, "A operand") for t in b.a_operands] - b_bufs = [pull(t, "B operand") for t in b.b_operands] - out_bufs = [pull(t, "output") for t in b.outputs] - aux_bufs = [pull(t, f"aux {self.aux_names[i]!r}") for i, t in enumerate(b.aux)] - - M = a_bufs[0].shape[1] - K = a_bufs[0].shape[2] * a_pack - N = b_bufs[0].shape[1] - - reason = _tma_alignment_reject(mm.a_dtype, mm.b_dtype, mm.a_major, mm.b_major, M, N, K) - if reason is not None: - raise ValueError(reason) - - # one pass: shape agreement and layout together - shape_bad, layout_bad = [], [] - for role, tensor, extent_is_n, kpack, want_dim, major in operands: - buf = resolved[id(tensor)] - shape = tuple(buf.shape) - if len(shape) != 3: - shape_bad.append(f"{role}: expected a rank-3 buffer, got shape {shape}") - continue - want = ((N if extent_is_n else M), K // kpack) - if (shape[1], shape[2]) != want: - shape_bad.append(f"{role}: expected (batch, {want[0]}, {want[1]}), got {shape}") - strides = tuple(buf.stride()) - unit = [i for i, s in enumerate(strides) if s == 1 and shape[i] > 1] - got = unit[0] if len(unit) == 1 else None - if got is not None and got != want_dim: - names = {0: "batch", 1: "M/N", 2: "K"} - layout_bad.append( - f"{role}: graph declares {major}-major (dim {want_dim} contiguous) but the buffer has dim {got} ({names[got]}) contiguous, stride={strides}" - ) - if shape_bad: - raise ValueError(f"runtime operand shapes disagree with the inferred problem size (M={M}, N={N}, K={K}): " + "; ".join(shape_bad)) - if layout_bad: - raise ValueError("runtime operand layout does not match the layout the kernel was compiled for: " + "; ".join(layout_bad)) - - reason = _alignment_reject([(role, resolved[id(t)], req, mode) for role, t, req, mode in extras]) - if reason is not None: - raise ValueError(reason) - - if self.block_scale: - sf_k4 = ((K // self.chain.block_scale.block_size) + 3) // 4 - reason = _sf_blob_reject( - [(f"SFA[{i}]", pull(t, "SFA"), 512 * sf_k4 * ((M + 127) // 128) * int(a_bufs[i].shape[0])) for i, t in enumerate(b.sfa_operands)] - + [(f"SFB[{j}]", pull(t, "SFB"), 512 * sf_k4 * ((N + 127) // 128) * int(b_bufs[j].shape[0])) for j, t in enumerate(b.sfb_operands)] - ) - if reason is not None: - raise ValueError(reason) - - mnk = (M, N, K) - c_arg = out_bufs if len(out_bufs) > 1 else out_bufs[0] - if self.chain.is_multi_gemm: - if self.block_scale: - sfa = [pull(t, "SFA") for t in b.sfa_operands] - sfb = [pull(t, "SFB") for t in b.sfb_operands] - pairs = [((a_bufs[ai], sfa[ai]), (b_bufs[bi], sfb[bi])) for ai, bi in self.chain.gemm_operands] - else: - pairs = [(a_bufs[ai], b_bufs[bi]) for ai, bi in self.chain.gemm_operands] - r = self._call_positional(pairs, c_arg, mnk, *aux_bufs, stream=stream) - elif self.block_scale: - r = self._call_positional( - a_bufs[0], b_bufs[0], c_arg, mnk, pull(b.sfa_operands[0], "SFA"), pull(b.sfb_operands[0], "SFB"), *aux_bufs, stream=stream - ) - else: - r = self._call_positional(a_bufs[0], b_bufs[0], c_arg, mnk, *aux_bufs, stream=stream) - _finalize_reductions(self.chain, out_bufs) - return r - def _call_positional(self, *args, stream=None): # Internal launcher (called by __call__ after variant-pack resolve). # Single-GEMM: (a, b, c, (M,N,K), *aux). Multi-GEMM: first arg is a list diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index 56748da87..a30382dfe 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -44,21 +44,17 @@ def execute(self, graph, variant_pack, ctx: ExecutionContext) -> None: 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 -- - # and not the caller's own objects, which would also tie the engine to - # whatever framework produced them. + # which carry the shape this execute runs, override_shapes included. views = variant_pack.views(slots) required = self.get_workspace_size() - # A FROST executor carves its scratch out of the CALLER's workspace: no - # hidden per-execute allocation, stable pointers, CUDA-graph friendly. - # Workspace.over validates it against what the pack already read. + # 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; going - # back through a dict keyed by tensor object only to have the - # compiled plan rebuild its by-object / by-uid / by-name tables is - # work with no answer in it. + # 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(dict(zip(self._tensors, views)), *extra, stream=ctx.stream) diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 774215221..8ef85175c 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -3,27 +3,16 @@ // // The variant pack, held as DLTensors rather than python objects. // -// execute() reads the caller's operands once and everything below works from -// the result. Doing that in python cost 1.5 us per operand -- 0.6 to ask the -// buffer for its facts one method call at a time, 0.9 to build the Tensor that -// carries them -- and each operand then had to be turned back into a DLPack -// producer for the kernel, which python cannot do quickly: tvm-ffi reads a -// torch tensor through a C function table and any python producer through a -// freshly built capsule, 0.32 us against 1.91. +// `__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. // -// Both halves are the same protocol. `__dlpack_c_exchange_api__` is a vtable -// on the TYPE whose dltensor_from_py_object_no_sync fills a caller-provided -// DLTensor in place, with no capsule and no allocation. This file consumes it -// to read the caller's buffers and implements it so the slots it hands out are -// read the same way. Measured on SM100, eight operands: reading 0.64 us -// against 10.4, contiguity 0.13 against 4.2, and a slot converts in 0.27 -- -// slightly cheaper than the torch tensor it replaces, so 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. The degradation is per operand, -// so a pack mixing torch with something else is exactly as fast as its parts. - +// 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 @@ -216,8 +205,7 @@ class VariantPackSlot { return dense; } - // One axis of it, the way a caller written against a framework tensor asks - // (``stride(-1)`` for the innermost). + // 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; @@ -295,9 +283,7 @@ class VariantPackSlot { return new VariantPackSlot(out, tensor_.device.device_id); } - // The same memory with its axes reordered. A kernel layer written against - // framework tensors reaches for this, and unlike reshape it is exact for a - // strided slot too -- it only relabels axes. + // 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)) @@ -475,14 +461,9 @@ class VariantPackNative { // 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. - // The whole pack from the caller's uid map, in one crossing: look each uid - // up, read what publishes the vtable, and report the rest. Same result as - // building the ordered buffer list in python and handing it to read_all, - // without the list, the comprehension, or the frame around them. - // - // A uid the map does not carry is left unfilled rather than refused here: - // whether that is the caller's mistake or an optional port depends on the - // graph, which python knows and this does not. + // 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; @@ -498,8 +479,7 @@ class VariantPackNative { return unread; } - // The first slot no one filled, or -1. The strict check asks this once - // instead of calling is_filled per operand. + // The first slot no one filled, or -1. int64_t first_unfilled() const { for (size_t i = 0; i < slots_.size(); i++) { @@ -508,21 +488,6 @@ class VariantPackNative { return -1; } - std::vector - read_all(py::sequence buffers) { - std::vector unread; - const size_t n = py::len(buffers); - for (size_t i = 0; i < n && i < slots_.size(); i++) { - py::handle buffer = buffers[i]; - if (buffer.is_none()) { - skip_slot(i); - } else if (!read_slot(i, buffer)) { - unread.push_back(i); - } - } - return unread; - } - // The fallback: python read the buffer its own way and reports the result. void set_slot(size_t index, @@ -541,23 +506,18 @@ class VariantPackNative { pointers_[index] = slot.data; } - // A slot the caller did not fill: an optional port it did not request. - // Re-describe a slot at the shape this execute is actually running, keeping - // the buffer it was read from. This is what override_shapes means: the - // caller allocated once at a cache shape and names the live shape per call. - // Applying it here rather than in an engine is what keeps the two paths - // answering the same question -- an engine that reads the pack gets the - // override without knowing the concept exists. + // 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 DLTensor's stride array is read - // ndim deep, so a shorter stride would be read past its end by any - // consumer -- and this is the one place a shape and a stride arrive - // from two different lists. + // 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()) + @@ -682,21 +642,17 @@ make_slot(int64_t ptr, std::vector shape, int dtype_code, int dtype_bit return new VariantPackSlot(std::move(slot), device_id); } -// ``(pointer, bytes)`` for a buffer that publishes the vtable, else None. -// -// The workspace is not an operand -- it has no uid and no slot -- but an engine -// still has to bounds-check its carves against it, and asking python for the -// size cost as much as reading all eight operands here. +// (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 for a dense buffer; a strided or - // broadcast one covers more (or, at stride 0, far less) than its element - // count says, and a carve bounds-checked against that would write outside - // the allocation. Hand it back to python to describe and refuse. + // 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]; @@ -704,10 +660,8 @@ read_buffer_extent(py::handle buffer) { return py::make_tuple(reinterpret_cast(static_cast(t.data) + t.byte_offset), numel * itemsize); } -// A workspace carve, planned once. Which regions a plan cuts, at what offsets, -// with what dtypes and shapes, is fixed when the engine builds; only the -// caller's base pointer arrives per execute. Describing them here rather than -// at each view() turns one crossing per region into one crossing per execute. +// 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)) { @@ -861,7 +815,6 @@ its parts. )") .def(py::init()) .def("read_slot", &VariantPackNative::read_slot) - .def("read_all", &VariantPackNative::read_all) .def("read_from", &VariantPackNative::read_from) .def("first_unfilled", &VariantPackNative::first_unfilled) .def("set_slot", &VariantPackNative::set_slot) diff --git a/watch_run.sh b/watch_run.sh deleted file mode 100755 index 610f5a5d6..000000000 --- a/watch_run.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash -# Watch a detached run and return on ANY of: finished, died, stalled. -# -# watch_run.sh [poll-seconds] -# -# The plain `until grep -q "^EXIT=" log; do sleep; done` waiter only returns when -# the job writes EXIT=. If the job is killed -- OOM, a SIGHUP from the launching -# shell, a segfaulting worker taking the session down -- that line never appears -# and the wait never ends, which is indistinguishable from "still running" and -# has stalled several turns. This returns in that case too, and says which. -# -# Exit codes: 0 finished, 1 died without EXIT=, 2 log stopped growing. -set -u -log=$1 -pattern=$2 -poll=${3:-300} -stall_limit=3 # consecutive polls with no new output before calling it stalled - -previous="" -stalls=0 -while true; do - if grep -q "^EXIT=" "$log" 2>/dev/null; then - echo "=== FINISHED after $(( SECONDS / 60 ))m ===" - tail -25 "$log" - exit 0 - fi - if ! pgrep -f "$pattern" > /dev/null 2>&1; then - echo "=== DIED after $(( SECONDS / 60 ))m: no process matching '$pattern', and the log has no EXIT= ===" - echo "the job was killed rather than finishing; the log's last lines:" - tail -25 "$log" 2>/dev/null - exit 1 - fi - size=$(stat -c %s "$log" 2>/dev/null || echo 0) - if [ "$size" = "$previous" ]; then - stalls=$(( stalls + 1 )) - else - stalls=0 - fi - if [ "$stalls" -ge "$stall_limit" ]; then - echo "=== STALLED: $log has not grown in $(( stall_limit * poll / 60 )) minutes, process still alive ===" - tail -25 "$log" 2>/dev/null - exit 2 - fi - previous=$size - sleep "$poll" -done