Normalize the variant pack once, into a C type that is also the DLPack producer - #547
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughGraph execution now normalizes caller buffers into ordered operands and uses the native raw-pointer API. SDPA and documentation use ChangesGraph execution and runtime contracts
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant GraphExecute
participant VariantPackNative
participant CompiledPlan
participant PyGraph
Caller->>GraphExecute: pass uid_to_tensor and workspace
GraphExecute->>VariantPackNative: normalize buffers and resolve operand order
GraphExecute->>CompiledPlan: pass VariantPack and execution context
CompiledPlan->>PyGraph: submit raw pointers and plan_index
PyGraph-->>CompiledPlan: return execution status
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-547-d6be41a |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (3)
python/cudnn/engines/base.py (2)
205-215: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winRead the dtype through
get_data_type().
_view_over_addressreadstensor.data_typedirectly.Tensor.get_data_type()inpython/cudnn/graph_types.pyexists precisely because a user may set a torch dtype, and it converts that to the cuDNN enum. Every engine in the tree usesget_data_type()(for exampleGdnCuTileEngine.check_support). With a torch dtype stored on the tensor,_cudnn_to_frost_dtype_namereturnsNoneand this function raises "the graph declares no data_type", which is wrong and hard to act on.♻️ Proposed change
- dtype = _cudnn_to_frost_dtype_name(tensor.data_type) + dtype = _cudnn_to_frost_dtype_name(tensor.get_data_type())🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/engines/base.py` around lines 205 - 215, Update _view_over_address to obtain the tensor dtype via tensor.get_data_type() before passing it to _cudnn_to_frost_dtype_name, matching the established engine usage and supporting user-supplied torch dtypes. Preserve the existing validation and DeviceView construction behavior.
127-137: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument that
tensorsisNonewhen the caller skipped description.
python/cudnn/_pygraph.pybuildsOperandswithtuple(records) if describe else None, anddescribeiseng is not None. On the backend pathtensorsis thereforeNone. The docstring states unconditionally thattensors[i]describes what the caller passed. A future engine that readsoperands.tensorswithout checking gets aTypeErroron a path that currently never reaches an engine.State the
Nonecase in the docstring, or set the attribute to an empty tuple.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/engines/base.py` around lines 127 - 137, Update the docstring for the operands/tensors attribute in the relevant class or initializer to state that tensors is None when descriptions are skipped on the backend path, while retaining the existing per-call tensor semantics when descriptions are available. Do not imply tensors[i] is always accessible.python/cudnn/datatypes.py (1)
116-126: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueHoist the cuDNN-to-Frost dtype mapping to a module-level constant.
_cudnn_to_frost_dtype_namerebuilds the dictionary for each bare-address port.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/datatypes.py` around lines 116 - 126, Hoist the cuDNN-to-Frost dtype mapping used by _cudnn_to_frost_dtype_name into a module-level constant, then have the function perform lookups against that shared mapping instead of rebuilding the dictionary on each call.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/adding_torch_custom_ops.md`:
- Around line 82-86: Update the documentation around graph.execute and the
workspace guidance: remove the claim that execute reuses a pointer array, and
revise the caching advice so users cache only the graph and workspace size, not
workspace tensors. Keep the execute usage example and its performance/dispatch
guidance accurate and consistent with the instructions around lines 101-103 and
177-178.
In `@python/cudnn/__init__.py`:
- Around line 297-307: Update the __getattr__ handling for the "experimental"
name to route its import through _load_optional_symbol, ensuring torch-related
ImportError failures produce the optional-dependency diagnostic and installation
hint. Preserve the existing direct importlib path for "ops" and continue caching
successfully loaded modules in globals().
In `@python/cudnn/_pygraph.py`:
- Around line 1725-1727: Update the operand preparation around selected_engine
and the python-engine execution branch so override_uids, override_shapes, or
override_strides do not leave operands as None for Python plans. Always
normalize uid_to_data for the Python-engine path, preserving the backend uid-map
overload behavior for non-Python execution and ensuring plans with
takes_operands=True receive Operands rather than a plain dict.
- Around line 1806-1811: Update the fallback ordering logic in the method
containing _sorted_uids to exclude tensors whose pass_by_value and scalar_type
are both non-None, while retaining non-virtual pass-by-value tensors without
embedded values. Apply this filter when building order so _normalize() does not
require unavailable buffers, preserving the existing empty-order and sorted
result behavior.
In `@test/python/test_import_boundaries.py`:
- Around line 102-103: Add the required `@pytest.mark.L0` marker to
test_lazy_top_level_attribute_resolves, placing it before the existing
parametrization decorator and leaving the test logic unchanged.
In `@test/python/test_variant_pack_normalization.py`:
- Around line 100-122: Update the concurrency test around worker so each thread
allocates its own workspace using the graph’s required size, preserving the
minimum one-byte allocation for zero-workspace plans. Also create and use a
thread-local cuDNN handle inside worker instead of sharing the outer handle
across threads, and pass both per-thread resources to g.execute.
- Around line 41-50: Add a module-level pytest skip in
test_variant_pack_normalization.py that requires torch.cuda.is_available() and a
CUDA compute capability of SM80 or newer, using
torch.cuda.get_device_capability(); ensure all tests in the module are skipped
before creating CUDA bfloat16 tensors on unsupported hosts.
In `@test/python/test_wrapper_graph.py`:
- Around line 29-57: Add the repository’s supported-capability skip checks
before both test_workspace_alloc_default_allocates and
test_workspace_alloc_false_is_honored run, covering CUDA availability, cuDNN
9.12+, and supported GPU architecture via the established support-check helpers,
cudnn.backend_version(), and torch.cuda.get_device_capability(). Keep the
existing test assertions and workspace behavior unchanged.
---
Nitpick comments:
In `@python/cudnn/datatypes.py`:
- Around line 116-126: Hoist the cuDNN-to-Frost dtype mapping used by
_cudnn_to_frost_dtype_name into a module-level constant, then have the function
perform lookups against that shared mapping instead of rebuilding the dictionary
on each call.
In `@python/cudnn/engines/base.py`:
- Around line 205-215: Update _view_over_address to obtain the tensor dtype via
tensor.get_data_type() before passing it to _cudnn_to_frost_dtype_name, matching
the established engine usage and supporting user-supplied torch dtypes. Preserve
the existing validation and DeviceView construction behavior.
- Around line 127-137: Update the docstring for the operands/tensors attribute
in the relevant class or initializer to state that tensors is None when
descriptions are skipped on the backend path, while retaining the existing
per-call tensor semantics when descriptions are available. Do not imply
tensors[i] is always accessible.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 86777b88-34a9-48f0-bfad-23ac1057fe7d
📒 Files selected for processing (25)
docs/adding_torch_custom_ops.mdpython/cudnn/__init__.pypython/cudnn/_pygraph.pypython/cudnn/datatypes.pypython/cudnn/engines/base.pypython/cudnn/experimental/ops/sdpa.pypython/cudnn/frost/device.pypython/cudnn/frost/workspace.pypython/cudnn/gemm/frost/compiler.pypython/cudnn/gemm/frost/tile_config.pypython/cudnn/graph_types.pypython/cudnn/linear_attention/cutile/gdn_engine.pypython/cudnn/linear_attention/cutile/kda_engine.pypython/cudnn/linear_attention/frost/gdn2_engine.pypython/cudnn/linear_attention/frost/gdn_engine.pypython/cudnn/linear_attention/frost/kda_engine.pypython/cudnn/nodes.pypython/cudnn/wrapper.pypython/pygraph/pygraph.cpppython/pygraph/pygraph.htest/python/test_decline_types.pytest/python/test_graph_native.pytest/python/test_import_boundaries.pytest/python/test_variant_pack_normalization.pytest/python/test_wrapper_graph.py
💤 Files with no reviewable changes (2)
- python/cudnn/graph_types.py
- python/cudnn/nodes.py
| `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. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove stale pointer-array and workspace-cache guidance.
graph.execute creates its ctypes pointer array per call. Do not state that it reuses one pointer array. Line 177 also contradicts Lines 101-103 by telling users to cache workspace tensors. Cache only the graph and workspace size.
Proposed documentation fix
- 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,
+ It normalizes the mapping and creates its pointer array for each call. The
+ sorted-pointer path runs underneath, so 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.
- [ ] Cache graph + uid_order + workspace in module-level dict
+ [ ] Cache graph + workspace size in a module-level dictAlso applies to: 177-178
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/adding_torch_custom_ops.md` around lines 82 - 86, Update the
documentation around graph.execute and the workspace guidance: remove the claim
that execute reuses a pointer array, and revise the caching advice so users
cache only the graph and workspace size, not workspace tensors. Keep the execute
usage example and its performance/dispatch guidance accurate and consistent with
the instructions around lines 101-103 and 177-178.
| if name in ("ops", "experimental"): | ||
| # Use importlib rather than "from . import <name>" to avoid infinite | ||
| # recursion. The cycle: | ||
| # 1. cudnn.ops accessed → __getattr__("ops") fires | ||
| # 2. "from . import ops" → _handle_fromlist(cudnn, ["ops"], ...) | ||
| # 3. _handle_fromlist calls hasattr(cudnn, "ops") | ||
| # 4. "ops" not in __dict__ yet → __getattr__("ops") again → goto 1 | ||
| # 1. cudnn.<name> accessed → __getattr__("<name>") fires | ||
| # 2. "from . import <name>" → _handle_fromlist(cudnn, ["<name>"], ...) | ||
| # 3. _handle_fromlist calls hasattr(cudnn, "<name>") | ||
| # 4. not in __dict__ yet → __getattr__("<name>") again → goto 1 | ||
| # importlib.import_module bypasses _handle_fromlist entirely. | ||
| _ops = importlib.import_module(".ops", __name__) | ||
| globals()["ops"] = _ops | ||
| return _ops | ||
|
|
||
| if name == "experimental": | ||
| from . import experimental as _experimental | ||
|
|
||
| globals()["experimental"] = _experimental | ||
| return _experimental | ||
| module = importlib.import_module(f".{name}", __name__) | ||
| globals()[name] = module | ||
| return module |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find eager torch/cutlass imports reachable from cudnn.experimental.
set -euo pipefail
fd -t f '__init__.py' python/cudnn/experimental --exec cat -n {}
rg -nP --type=py '^\s*(import|from)\s+(torch|cutlass)\b' python/cudnn/experimental/ || echo "no module-scope torch/cutlass imports found"
rg -n -C 6 '_load_optional_symbol' python/cudnn/__init__.pyRepository: NVIDIA/cudnn-frontend
Length of output: 2581
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- python/cudnn/__init__.py ---'
sed -n '245,318p' python/cudnn/__init__.py
printf '%s\n' '--- python/cudnn/experimental/__init__.py ---'
cat -n python/cudnn/experimental/__init__.py
printf '%s\n' '--- python/cudnn/experimental/ops/__init__.py ---'
cat -n python/cudnn/experimental/ops/__init__.py
printf '%s\n' '--- python/cudnn/experimental/ops/sdpa.py imports ---'
sed -n '1,70p' python/cudnn/experimental/ops/sdpa.py
printf '%s\n' '--- static import-chain verifier ---'
python3 - <<'PY'
import ast
from pathlib import Path
files = [
Path("python/cudnn/experimental/__init__.py"),
Path("python/cudnn/experimental/ops/__init__.py"),
Path("python/cudnn/experimental/ops/sdpa.py"),
Path("python/cudnn/__init__.py"),
]
for path in files:
tree = ast.parse(path.read_text(), filename=str(path))
imports = []
for node in tree.body:
if isinstance(node, ast.Import):
imports.extend(alias.name for alias in node.names)
elif isinstance(node, ast.ImportFrom):
imports.append(f"from {node.module or ''} import " +
", ".join(alias.name for alias in node.names))
print(f"{path}:")
for item in imports:
print(f" {item}")
root = ast.parse(Path("python/cudnn/__init__.py").read_text())
for node in ast.walk(root):
if isinstance(node, ast.FunctionDef) and node.name == "__getattr__":
body = ast.get_source_segment(Path("python/cudnn/__init__.py").read_text(), node)
print("\n__getattr__ contains diagnostic wrapper call:",
"_load_optional_symbol" in body)
print("experimental branch uses importlib.import_module:",
'importlib.import_module(f".{name}", __name__)' in body)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 8503
Route cudnn.experimental failures through the optional-dependency diagnostic.
cudnn.experimental imports sdpa.py, which imports torch at module scope. This path bypasses _load_optional_symbol, so missing torch raises a bare ImportError without the pip install nvidia-cudnn-frontend[cutedsl] hint.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/__init__.py` around lines 297 - 307, Update the __getattr__
handling for the "experimental" name to route its import through
_load_optional_symbol, ensuring torch-related ImportError failures produce the
optional-dependency diagnostic and installation hint. Preserve the existing
direct importlib path for "ops" and continue caching successfully loaded modules
in globals().
Source: Coding guidelines
| else: | ||
| order = sorted(t.uid for t in self._tensor_by_uid.values() if not t.is_virtual) | ||
| if not order: | ||
| return None | ||
| self._sorted_uids = order | ||
| return order |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find non-virtual pass-by-value tensors on python-only graph paths.
set -euo pipefail
rg -nP --type=py -C 4 'tensor_scalar\(|is_pass_by_value\s*=\s*True' python/cudnn/linear_attention/ python/cudnn/gemm/ || echo "no matches in python-only engine trees"
rg -n -C 6 'def _operand_uids' python/cudnn/_pygraph.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1073
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- _operand_uids and normalization ---'
sed -n '1775,1845p' python/cudnn/_pygraph.py
rg -n -C 8 'missing a buffer for tensor uid|is_pass_by_value|tensor_scalar|def tensor\(' python/cudnn/_pygraph.py python/cudnn
printf '%s\n' '--- Python graph engine files ---'
git ls-files 'python/cudnn' | rg 'linear_attention|gemm|gdn|kda|pygraph' | head -200
printf '%s\n' '--- pass-by-value creation and uses ---'
rg -n -C 5 'tensor_scalar\(|is_pass_by_value\s*=\s*True|is_pass_by_value' python test docsRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- scalar and pass-by-value use in Python graph paths ---'
rg -n -C 3 --glob '*.py' 'tensor_scalar\(|is_pass_by_value\s*=\s*True|set_is_pass_by_value' \
python/cudnn/linear_attention python/cudnn/gemm python/cudnn/_pygraph.py \
| head -300
printf '%s\n' '--- graph execution and variant-pack normalization call sites ---'
rg -n -C 6 --glob '*.py' '_normalize\(|_operand_uids\(|uid_to_data|variant_pack' \
python/cudnn/_pygraph.py python/cudnn/linear_attention python/cudnn/gemm \
| head -500
printf '%s\n' '--- Tensor registration and backend-lowerability paths ---'
rg -n -C 8 --glob '*.py' 'def _register_tensor|def _backend_lowerable|def execute|def _execute|_engine_for' \
python/cudnn/_pygraph.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- C++ variant-pack treatment of pass-by-value tensors ---'
rg -n -C 8 'is_pass_by_value|variant.*pack|variant_pack|pass.by.value' \
include python/pygraph test | head -500
printf '%s\n' '--- Python graph construction for GDN/KDA/GDN2 ---'
rg -n -C 8 'def .*graph|Graph\(|tensor_scalar|tensor_like|_make_tensor|is_pass_by_value' \
python/cudnn/linear_attention/ops python/cudnn/linear_attention/engine_utils.py \
| head -500
printf '%s\n' '--- auto bindings and input classification ---'
sed -n '430,560p' python/cudnn/_pygraph.py
sed -n '1870,1910p' python/cudnn/_pygraph.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 12 'get_variant_pack_uids_sorted|prepare_variant_pack_template' \
include python | head -300
printf '%s\n' '--- pass-by-value tests and execution packs ---'
rg -n -C 8 'is_pass_by_value=True|is_pass_by_value.*True|pass_by_value' \
test/python | head -250Repository: NVIDIA/cudnn-frontend
Length of output: 42529
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '2267,2392p' include/cudnn_frontend/graph_interface.hRepository: NVIDIA/cudnn-frontend
Length of output: 5895
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- C++ pass-by-value cache population ---'
rg -n -C 10 'cached_pass_by_value|deserialized_pass_by_value|variant_pack_uids|extend_tensor_map_with_pass_by_value_tensors_' \
include/cudnn_frontend/graph_interface.h include/cudnn_frontend/graph/graph.h include/cudnn_frontend \
| head -500
printf '%s\n' '--- Python Tensor pass-by-value fields and lowering ---'
sed -n '330,355p' python/cudnn/_pygraph.py
sed -n '1980,2010p' python/cudnn/_pygraph.pyRepository: NVIDIA/cudnn-frontend
Length of output: 46125
Exclude embedded pass-by-value scalars from the IR fallback order. tensor_scalar() creates a non-virtual tensor with an embedded value, so _normalize() requires a buffer that the caller does not provide. Retain tensor(is_pass_by_value=True) tensors without embedded values because callers can supply those buffers. Filter tensors with pass_by_value is not None and scalar_type is not None.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/_pygraph.py` around lines 1806 - 1811, Update the fallback
ordering logic in the method containing _sorted_uids to exclude tensors whose
pass_by_value and scalar_type are both non-None, while retaining non-virtual
pass-by-value tensors without embedded values. Apply this filter when building
order so _normalize() does not require unavailable buffers, preserving the
existing empty-order and sorted result behavior.
| @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") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Gate these tests on CUDA availability and bfloat16 support.
Every test in this file allocates bfloat16 tensors on device="cuda" with no guard. bfloat16 needs SM80 or later, and the module also needs a CUDA device to exist. On an unsupported host these tests error instead of skipping.
Add a module-level skip based on torch.cuda.is_available() and torch.cuda.get_device_capability().
💚 Proposed guard
M = N = K = 64
+
+pytestmark = pytest.mark.skipif(
+ not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8,
+ reason="requires a CUDA device with bfloat16 support (SM80+)",
+)As per coding guidelines: "Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks, cudnn.backend_version(), and torch.cuda.get_device_capability()."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/test_variant_pack_normalization.py` around lines 41 - 50, Add a
module-level pytest skip in test_variant_pack_normalization.py that requires
torch.cuda.is_available() and a CUDA compute capability of SM80 or newer, using
torch.cuda.get_device_capability(); ensure all tests in the module are skipped
before creating CUDA bfloat16 tensors on unsupported hosts.
Source: Coding guidelines
| @pytest.mark.L0 | ||
| def test_workspace_alloc_default_allocates(): | ||
| """The default path allocates a workspace the caller never has to think about.""" | ||
| graph = _matmul_graph() | ||
| assert torch.is_tensor(graph._Graph__workspace) | ||
|
|
||
|
|
||
| @pytest.mark.L0 | ||
| def test_workspace_alloc_false_is_honored(): | ||
| """``workspace_alloc=False`` means the CALLER owns the workspace. | ||
|
|
||
| Regression: the sentinel is written as ``self.__workspace`` (mangled to | ||
| ``_Graph__workspace``) but was read back with ``hasattr(self, "__workspace")`` | ||
| — a plain string, which is NOT name-mangled. That probe was therefore always | ||
| False, the sentinel was overwritten with a fresh allocation on every | ||
| ``__exit__``, and the "Need to specify workspace" guard below was unreachable. | ||
| """ | ||
| graph = _matmul_graph(workspace_alloc=False) | ||
| assert graph._Graph__workspace is False | ||
|
|
||
| x = torch.randn(1, 64, 64, dtype=torch.half, device="cuda") | ||
| w = torch.randn(1, 64, 64, dtype=torch.half, device="cuda") | ||
|
|
||
| with pytest.raises(RuntimeError, match="Need to specify workspace"): | ||
| graph(x, w) | ||
|
|
||
| workspace = torch.empty(max(graph.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") | ||
| out = graph(x, w, workspace=workspace) | ||
| torch.testing.assert_close(out.float(), (x @ w).float(), atol=1e-2, rtol=1e-2) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Add supported-capability gates for these CUDA tests.
These tests unconditionally require CUDA and cuDNN 9.12 or later. Unsupported backend versions or GPU architectures will fail instead of skip. Add the repository support checks before both tests run.
As per coding guidelines, “Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks, cudnn.backend_version(), and torch.cuda.get_device_capability().”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/test_wrapper_graph.py` around lines 29 - 57, Add the repository’s
supported-capability skip checks before both
test_workspace_alloc_default_allocates and test_workspace_alloc_false_is_honored
run, covering CUDA availability, cuDNN 9.12+, and supported GPU architecture via
the established support-check helpers, cudnn.backend_version(), and
torch.cuda.get_device_capability(). Keep the existing test assertions and
workspace behavior unchanged.
Source: Coding guidelines
d6be41a to
db42121
Compare
Update: tried the engine migration, it does not pay yet — and it found a contract bugMigrated A real bug in the layout, now fixed here. A gdn graph marks its own output The migration itself is reverted. Measured on GDN forward, sm100:
So Numbers after all of the above
Neither path regresses; the backend gains.
|
db42121 to
7daf876
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/python_graph_and_execution_backends.md`:
- Around line 83-85: Clarify the `CompiledPlan.execute()` documentation and
related sections to state that native execution and plans with
`takes_operands=True` consume `Operands`, while legacy plans with
`takes_operands=False` continue receiving the raw `{uid: buffer}` map until
migration is complete; ensure the metadata-record description is limited to
migrated plans.
- Around line 122-126: Update the documentation around bind_ports to clarify
that an unfilled slot may be omitted only when the port is virtual or explicitly
optional; missing non-virtual, non-optional ports must raise ValueError.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 88a73912-5d80-4269-abde-b49e109785b7
📒 Files selected for processing (2)
docs/python_graph_and_execution_backends.mdpython/cudnn/_pygraph.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudnn/_pygraph.py
YangXu1990uiuc
changed the title
Third commit: the DLPack struct is filled from a per-layout prototypeChased the engine-migration measurement to its root rather than leaving the PR half-done.
So: fill once per (shape, dtype, device), memmove per call.
No engine was touched for that GDN number — the ten workspace views are all it takes. This is the shape the backend already uses for kernel arguments ( The struct stays fresh per capsule: cute's Incidentally this removes a leak: the old implementation appended every struct it ever built to On the engine migration, which is deliberately NOT hereI tried it (three The migration itself measured +34 us, and the honest reading is that the DLPack producer was the problem, not the contract — which this commit now fixes. Migrating engines is the next PR, on top of a producer that is 2.6x faster. Totals
167 tests pass ( |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/frost/buffers.py`:
- Line 118: Sort the entries in the _DLPackProto.__slots__ tuple alphabetically
to satisfy Ruff RUF023, preserving the same slot names and behavior.
- Around line 236-241: Update DeviceView.__dlpack__ and its capsule creation
flow so each DLManagedTensor and prototype is registered against the DLPack
consumer lifetime rather than only appended to self._live. Install a real
DLManagedTensor deleter that releases the registry entry after consumption, and
add a capsule destructor that handles unconsumed dltensor capsules while
preserving correct lifetime of the ctypes storage.
In `@test/python/test_dlpack_proto.py`:
- Around line 27-28: Remove the redundant per-file CUDA marker from
test_capsule_decodes_to_the_right_buffer, and gate the bfloat16 cases in that
test with an SM80-or-newer capability check. Apply the check specifically to the
cases around lines 50 and 60-65, requiring CUDA compute capability at least (8,
0), while leaving non-bfloat16 coverage unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ed2fbd95-0e94-46f4-a5e6-732f3e2bb2a1
📒 Files selected for processing (6)
python/cudnn/_pygraph.pypython/cudnn/datatypes.pypython/cudnn/engines/base.pypython/cudnn/frost/buffers.pytest/python/test_dlpack_proto.pytest/python/test_variant_pack_normalization.py
🚧 Files skipped from review as they are similar to previous changes (2)
- test/python/test_variant_pack_normalization.py
- python/cudnn/_pygraph.py
| @pytest.mark.L0 | ||
| def test_capsule_decodes_to_the_right_buffer(): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 3 'skipif|cuda\.is_available|get_device_capability|bfloat16|pytestmark' \
test/python/conftest.py test/python/test_dlpack_proto.pyRepository: NVIDIA/cudnn-frontend
Length of output: 2505
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test/python/test_dlpack_proto.py ---'
cat -n test/python/test_dlpack_proto.py
printf '%s\n' '--- test/python/conftest.py setup ---'
sed -n '130,180p' test/python/conftest.py
printf '%s\n' '--- capability and skip patterns in Python tests ---'
rg -n -C 2 'get_device_capability|cuda\.is_available|bfloat16|skipif|support' test/python -g '*.py' | head -n 240Repository: NVIDIA/cudnn-frontend
Length of output: 23457
🌐 Web query:
PyTorch torch.cuda.is_bf16_supported compute capability bfloat16 CUDA support
💡 Result:
In PyTorch, the function torch.cuda.is_bf16_supported is used to check if the current CUDA device supports the bfloat16 data type [1][2]. Native Hardware Support Bfloat16 precision is natively supported on NVIDIA GPU architectures with compute capability 8.0 or higher (e.g., Ampere and newer architectures) [3][4][5][6]. On these devices, bfloat16 operations are executed using specialized hardware, such as Tensor Cores [6]. Software Emulation PyTorch allows for bfloat16 usage even on older hardware that lacks native bfloat16 support [7][8]. In such cases, the system performs software emulation, where bfloat16 operations are typically converted and executed using float32 compute [7][9]. While this enables compatibility, it does not provide the performance benefits of native hardware acceleration and may be slower than standard float32 operations [10][7][9]. Using torch.cuda.is_bf16_supported You can refine your query in PyTorch to distinguish between native hardware support and software emulation using the including_emulation parameter [7]: - torch.cuda.is_bf16_supported(including_emulation=True) (Default): Returns True if the device supports bfloat16, either natively or via emulation [2][7]. - torch.cuda.is_bf16_supported(including_emulation=False): Returns True only if the device has native hardware support for bfloat16 [7]. For optimal performance in deep learning workflows, native support (compute capability 8.0+) is highly recommended [3][11]. Relying on emulation is generally discouraged for training scenarios due to the lack of performance gains [10][7].
Citations:
- 1: https://docs.pytorch.org/docs/stable/generated/torch.cuda.is_bf16_supported.html
- 2: https://docs.pytorch.org/docs/2.11/generated/torch.cuda.is_bf16_supported.html
- 3: https://docs.nvidia.com/cuda/archive/13.2.1/cuda-programming-guide/05-appendices/compute-capabilities.html
- 4: https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html
- 5: https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html.md
- 6: https://docs.nvidia.com/cuda/archive/12.0.0/pdf/Ampere_Tuning_Guide.pdf
- 7: use bfloat16 on nvidia V100 GPU pytorch/pytorch#124996
- 8: https://discuss.pytorch.org/t/bfloat16-on-nvidia-v100-gpu/201629
- 9: https://discuss.pytorch.org/t/bfloat16-native-support/117155
- 10:
torch.cuda.is_bf16_supported()seem to not work properly pytorch/pytorch#75427 - 11: https://www.dgl.ai/dgl_docs/guide/mixed_precision.html
🌐 Web query:
site:github.com/pytorch/pytorch "bfloat16" "V100" torch.ones CUDA
💡 Result:
The NVIDIA V100 GPU (based on the Volta architecture, compute capability 7.0) does not provide native hardware support for bfloat16 [1][2]. While you can create and store torch.bfloat16 tensors on a V100 using PyTorch (e.g., torch.ones(..., device='cuda', dtype=torch.bfloat16)), executing operations on them often relies on software emulation or fallbacks [1][3]. Key points regarding bfloat16 on V100: 1. Hardware Support: Native bfloat16 instructions require an NVIDIA GPU with compute capability 8.0 or higher (e.g., Ampere architecture and newer) [1][4][5]. 2. PyTorch Behavior: - Eager mode may allow basic operations through software emulation [1][6]. - Compiling code with torch.compile (using TorchInductor) often fails on V100 because the generated Triton kernels target hardware features (like .bf16 instructions) that are unavailable on Volta [4][5][7]. - Functions like torch.cuda.is_bf16_supported() may return True depending on the PyTorch version and whether emulation is included, but this refers to the ability to handle the format via software, not native hardware acceleration [1][6]. 3. Performance: Because V100 lacks hardware-accelerated bfloat16, operations performed in this format on Volta GPUs do not provide the performance benefits seen on newer architectures and are typically as slow as, or slower than, float32 computations [1]. If you encounter issues such as "Feature '.bf16' requires.target sm_80 or higher," it is because the backend is attempting to use hardware-native features that your V100 does not support [4][5]. Using torch.float32 or torch.float16 (which is supported natively on V100) is generally recommended for production workflows on this architecture [2].
Citations:
- 1: use bfloat16 on nvidia V100 GPU pytorch/pytorch#124996
- 2: DISABLED test_index_add_correctness (__main__.TestTorch) pytorch/pytorch#91184
- 3: NVFuser takes a fallback with bfloat16 inputs on V100 pytorch/pytorch#75089
- 4: Some V100 inductor unit tests fail with message
Feature '.bf16' requires .target sm_80 or higherpytorch/pytorch#103993 - 5:
test_scatter_bf16_cudafails on V100 pytorch/pytorch#118581 - 6:
torch.cuda.is_bf16_compatible()output inconsistent with with TorchInductor support pytorch/pytorch#118122 - 7: Many test failures when GPU does not support BF16 pytorch/pytorch#172085
Gate the bfloat16 cases by GPU capability.
test/python/conftest.py already asserts CUDA availability, so a per-file CUDA marker is redundant. The bfloat16 cases at lines 50 and 60-65 still need an SM80-or-newer check because native CUDA bfloat16 support starts at compute capability (8, 0).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/test_dlpack_proto.py` around lines 27 - 28, Remove the redundant
per-file CUDA marker from test_capsule_decodes_to_the_right_buffer, and gate the
bfloat16 cases in that test with an SM80-or-newer capability check. Apply the
check specifically to the cases around lines 50 and 60-65, requiring CUDA
compute capability at least (8, 0), while leaving non-bfloat16 coverage
unchanged.
Source: Coding guidelines
_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) <noreply@anthropic.com>
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 NVIDIA#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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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.
0147228 to
3e5da3b
Compare
YangXu1990uiuc
changed the title
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-547-3e5da3b |
…oducer
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.
3e5da3b to
64089d9
Compare
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudnn/graph_types.py (1)
109-112: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftRestore protection for direct
Tensorattribute writes.
_freeze()seals onlydimandstride. Direct writes touidand other metadata remain possible. A directtensor.uid = ...bypasses_reuid_tensor(), leaving_tensor_by_uid,_data_bindings, and_cpp_tensorskeyed by the old UID while lowering and execution use the new UID. Restore guarded or immutable writes and add regression tests for direct writes after freezing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/graph_types.py` around lines 109 - 112, Extend Tensor attribute protection beyond dim and stride so direct writes to uid and other frozen metadata are guarded or rejected after _freeze(). Route allowed UID changes through _reuid_tensor() to keep _tensor_by_uid, _data_bindings, and _cpp_tensors consistent, and add regression tests covering direct writes after freezing.
🧹 Nitpick comments (5)
python/cudnn/engines/base.py (1)
144-144: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__slots__.Ruff reports RUF023 for this tuple. Apply a natural sort to keep the lint clean.
♻️ Proposed change
- __slots__ = ("uids", "native", "_tensors", "_slot_of", "workspace", "workspace_bytes", "_device") + __slots__ = ("_device", "_slot_of", "_tensors", "native", "uids", "workspace", "workspace_bytes")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/engines/base.py` at line 144, Apply natural alphabetical sorting to the __slots__ tuple in the relevant class, preserving every existing slot name and removing the Ruff RUF023 violation.Source: Linters/SAST tools
python/pygraph/variant_pack.cpp (3)
267-272: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
contiguousmember.The binding at Line 586 registers a lambda, not this method.
VariantPackSlot::contiguousis never referenced. Either bind it or delete it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pygraph/variant_pack.cpp` around lines 267 - 272, Remove the unused VariantPackSlot::contiguous member function, including its declaration or definition as applicable; do not modify the separate lambda binding registered near the binding code.
410-423: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
read_alltolerates a length mismatch silently.The loop stops at
min(len(buffers), slots_.size()). If Python passes fewer buffers than there are slots, the trailing slots stay unfilled and no caller learns why. If it passes more, the extra buffers are dropped.
_pygraph._normalizebuilds the list fromorder, so the lengths match today. Raise on a mismatch to keep that invariant enforced.♻️ Proposed change
std::vector<size_t> unread; const size_t n = py::len(buffers); + if (n != slots_.size()) + throw py::value_error("read_all got " + std::to_string(n) + " buffers for " + + std::to_string(slots_.size()) + " slots"); for (size_t i = 0; i < n && i < slots_.size(); i++) {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pygraph/variant_pack.cpp` around lines 410 - 423, Update read_all to validate that py::len(buffers) equals slots_.size() before processing any slots, and raise an appropriate exception when the lengths differ. Preserve the existing skip_slot, read_slot, and unread handling for matching lengths, rather than silently truncating or leaving slots unfilled.
164-174: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDelete the copy and move operations of
VariantPackSlot.
tensor_.shapeandtensor_.stridespoint intoslot_.shapeandslot_.stride. The implicit copy constructor copies both members, so the copy'stensor_still points into the source object's vectors. Any future copy therefore produces a danglingDLTensor.No current path copies a slot. Make the hazard impossible.
🛡️ Proposed change
VariantPackSlot(const Slot &slot, int32_t device_id) : slot_(slot) { + // tensor_ points into slot_, so a copy would alias the source's storageprivate: + VariantPackSlot(const VariantPackSlot &) = delete; + VariantPackSlot &operator=(const VariantPackSlot &) = delete; Slot slot_; // owns the shape/stride storage the DLTensor points into DLTensor tensor_{}; };Also applies to: 305-308
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pygraph/variant_pack.cpp` around lines 164 - 174, Delete the copy and move constructors and assignment operators for VariantPackSlot, leaving the existing constructor and destructor behavior unchanged. Ensure all copy and move operations are explicitly disabled so tensor_.shape and tensor_.strides cannot be copied with stale pointers into another instance’s slot storage.python/CMakeLists.txt (1)
21-26: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe custom failure message is unreachable.
find_package(dlpack 1.3 REQUIRED)aborts the configure step when the package is missing or too old. CMake prints its own error, so theelse()branch at Lines 24-26 never runs and the guidance aboutDLPackExchangeAPIandCUDNN_FRONTEND_USE_SYSTEM_DLPACKnever reaches the user.Drop
REQUIREDso the custom message runs, or remove the dead branch.♻️ Proposed change
- find_package(dlpack 1.3 REQUIRED) + find_package(dlpack 1.3) if(dlpack_FOUND) message(STATUS "Found system dlpack ${dlpack_VERSION}") else() message(FATAL_ERROR "dlpack >= 1.3 not found (needed for DLPackExchangeAPI); unset CUDNN_FRONTEND_USE_SYSTEM_DLPACK to fetch it") endif()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/CMakeLists.txt` around lines 21 - 26, Update the find_package(dlpack 1.3) call and its surrounding conditional so the custom missing-or-outdated-package guidance can execute; either remove REQUIRED while preserving the existing fatal message, or remove the unreachable else branch.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/_pygraph.py`:
- Around line 1877-1884: Update the workspace handling in the VariantPack
construction flow to use the required workspace size when workspace is supplied
as an integer address, rather than deriving it from the geometry-less Tensor
returned by _describe(). Preserve the existing _byte_size(workspace_tensor)
behavior for tensor-backed workspaces, and ensure Workspace.over() receives the
actual available byte count.
In `@python/cudnn/datatypes.py`:
- Around line 333-356: Ensure DLPack dtype tables are initialized before reverse
lookups: expose a shared initializer from datatypes.py, have _dlpack_code_bits
reuse it, and invoke it before VariantPack.tensors reads
_FROST_DTYPE_CODE_TO_CUDNN. Add/use _cudnn_for_dlpack_code_bits for the reverse
mapping so all-torch packs resolve their dtypes correctly.
In `@python/cudnn/engines/base.py`:
- Around line 51-52: Remove the unused module-level imports
_CUDNN_TO_FROST_DTYPE_NAME and DeviceView from base.py, and update
_view_over_address to reference DeviceView through the existing buffers module
if needed, preserving the lazy Frost import boundary.
- Around line 252-269: Update bind_ports.resolve to skip slots where
native.is_filled(slot) is false for optional ports, while preserving the
existing ValueError for missing required non-virtual ports. Adjust Frost GDN
consumers to check bound-port presence before accessing optional final_state and
H outputs, avoiding direct nb.outputs lookups when those ports are unbound.
In `@python/cudnn/graph_types.py`:
- Around line 267-270: Update byte_size() to obtain the data type through
tensor.get_data_type() before looking it up in _CUDNN_TO_FROST_DTYPE_NAME,
preserving the existing empty-dimension and unknown-type handling. Add a
regression test covering a non-empty Tensor with torch.float16 and verifying its
correct byte size.
In `@python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py`:
- Around line 2988-3042: Update the _build_descs docstring to remove the
outdated claim that builder launches do not recur during steady-state replay,
and document that cache["build_descs"](...) is invoked on every launch while the
compiled builder remains cached.
In `@python/pygraph/variant_pack.cpp`:
- Around line 319-334: In python/pygraph/variant_pack.cpp#L319-L334, update
slot_managed_from_py_object so the returned DLManagedTensorVersioned owns
independent shape and stride copies, and make its deleter release those
allocations along with the managed tensor, matching the capsule path. In
python/pygraph/variant_pack.cpp#L164-L174, delete VariantPackSlot’s copy
constructor and copy assignment to prevent tensor_ from retaining pointers into
another slot’s storage.
- Around line 83-105: Update the type cache used by exchange_api_for so each
cached PyTypeObject* is held with a strong Python reference, including existing
entries and newly inserted types. Ensure cache cleanup or replacement decrements
those references appropriately while preserving null API-result caching.
- Around line 282-303: Update the dlpack method to negotiate max_version: when
it requests DLPack 1.x, construct and return a DLManagedTensorVersioned in a
"dltensor_versioned" capsule, including a capsule destructor that recognizes and
transfers ownership for "used_dltensor_versioned". For unsupported or
unversioned-only requests, explicitly document the limitation and reject the
request rather than silently returning the unversioned tensor.
---
Outside diff comments:
In `@python/cudnn/graph_types.py`:
- Around line 109-112: Extend Tensor attribute protection beyond dim and stride
so direct writes to uid and other frozen metadata are guarded or rejected after
_freeze(). Route allowed UID changes through _reuid_tensor() to keep
_tensor_by_uid, _data_bindings, and _cpp_tensors consistent, and add regression
tests covering direct writes after freezing.
---
Nitpick comments:
In `@python/CMakeLists.txt`:
- Around line 21-26: Update the find_package(dlpack 1.3) call and its
surrounding conditional so the custom missing-or-outdated-package guidance can
execute; either remove REQUIRED while preserving the existing fatal message, or
remove the unreachable else branch.
In `@python/cudnn/engines/base.py`:
- Line 144: Apply natural alphabetical sorting to the __slots__ tuple in the
relevant class, preserving every existing slot name and removing the Ruff RUF023
violation.
In `@python/pygraph/variant_pack.cpp`:
- Around line 267-272: Remove the unused VariantPackSlot::contiguous member
function, including its declaration or definition as applicable; do not modify
the separate lambda binding registered near the binding code.
- Around line 410-423: Update read_all to validate that py::len(buffers) equals
slots_.size() before processing any slots, and raise an appropriate exception
when the lengths differ. Preserve the existing skip_slot, read_slot, and unread
handling for matching lengths, rather than silently truncating or leaving slots
unfilled.
- Around line 164-174: Delete the copy and move constructors and assignment
operators for VariantPackSlot, leaving the existing constructor and destructor
behavior unchanged. Ensure all copy and move operations are explicitly disabled
so tensor_.shape and tensor_.strides cannot be copied with stale pointers into
another instance’s slot storage.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 22be3612-b910-47cf-93db-fa5f3ec52ea0
📒 Files selected for processing (24)
dlpack_version.txtdocs/python_graph_and_execution_backends.mdpython/CMakeLists.txtpython/cudnn/__init__.pypython/cudnn/_pygraph.pypython/cudnn/datatypes.pypython/cudnn/engines/base.pypython/cudnn/frost/buffers.pypython/cudnn/frost/device.pypython/cudnn/frost/workspace.pypython/cudnn/gemm/frost/compiler.pypython/cudnn/graph_types.pypython/cudnn/linear_attention/engine_utils.pypython/cudnn/linear_attention/frost/gdn2_engine.pypython/cudnn/linear_attention/frost/gdn_engine.pypython/cudnn/linear_attention/frost/kda_engine.pypython/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.pypython/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.pypython/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.pypython/cudnn/linear_attention/frost/kernel/kda_prefill_f16.pypython/pycudnn.cpppython/pygraph/variant_pack.cpppython/pygraph/variant_pack.htest/python/test_variant_pack_normalization.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudnn/init.py
| py::capsule | ||
| dlpack(py::object /*stream*/, py::object /*max_version*/) const { | ||
| struct Owned { | ||
| DLManagedTensor managed; | ||
| std::vector<int64_t> shape; | ||
| std::vector<int64_t> 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<Owned *>(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<DLManagedTensor *>(PyCapsule_GetPointer(capsule, "dltensor")); | ||
| if (managed != nullptr && managed->deleter != nullptr) managed->deleter(managed); | ||
| } | ||
| }); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
DLPack 1.3 __dlpack__ max_version dltensor_versioned capsule producer requirements
💡 Result:
In DLPack 1.3, the exchange of tensor data through Python's dlpack mechanism involves version negotiation via the max_version parameter, the use of the DLManagedTensorVersioned structure, and specific PyCapsule naming conventions [1][2][3]. Max Version and Negotiation The max_version parameter in dlpack (introduced in the Python array API standard) allows a consumer to specify the maximum DLPack version (as a (major, minor) tuple) it supports [2][3]. Producers are encouraged to use this to determine which struct to export [1][2][3]: - If the producer's version is less than or equal to the consumer's max_version, it should return a capsule matching the requested version [2][3]. - If the producer cannot support the consumer's requested version, it may return a different version (which the consumer must verify) or raise a BufferError if compatibility is impossible [2][3]. Capsule Requirements Producers must wrap the tensor data in a PyCapsule [1]. The naming convention for these capsules depends on the struct version: - When using the DLManagedTensorVersioned struct (the current standard), the capsule name must be "dltensor_versioned" [1][2][3]. - For backward compatibility with older libraries or legacy data structures (DLManagedTensor), the name "dltensor" is used [1][2][3]. - Ownership transfer: When a consumer imports the tensor (via from_dlpack), it must rename the capsule from "dltensor_versioned" to "used_dltensor_versioned" (or "used_dltensor" for legacy) to indicate that ownership has been transferred and to prevent the original PyCapsule_Destructor from prematurely invoking the deleter [1][2]. Struct Requirements The DLManagedTensorVersioned struct is the standard for DLPack 1.0 and later [1]. It includes a DLPackVersion field (major/minor) that must be checked by the consumer for major version compatibility before accessing the structure [4][5][6]. A major version mismatch (e.g., consumer expects 1.x but receives 2.x) indicates an incompatible ABI change, in which case the consumer must call the capsule's deleter and abort access [4][5]. Minor version updates are generally additive and safe to interpret if the consumer is prepared for them [4][5].
Citations:
- 1: https://dmlc.github.io/dlpack/latest/python_spec.html
- 2: https://data-apis.org/array-api/2025.12/API_specification/generated/array_api.array.__dlpack__.html
- 3: https://data-apis.org/array-api/draft/API_specification/generated/array_api.array.__dlpack__.html
- 4: https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h
- 5: https://dmlc.github.io/dlpack/latest/c_api.html
- 6: dmlc/dlpack@ca4d00a
🏁 Script executed:
fd -i 'variant_pack.cpp|dlpack.h|pygraph' . && \
sed -n '240,330p' python/pygraph/variant_pack.cpp && \
rg -n "__dlpack__|max_version|DLManagedTensorVersioned|dltensor_versioned|dltensor" python include CMakeLists.txt pyproject.toml setup.py 2>/dev/null | head -200Repository: NVIDIA/cudnn-frontend
Length of output: 9151
🏁 Script executed:
sed -n '1,45p' python/pygraph/variant_pack.cpp
sed -n '560,610p' python/pygraph/variant_pack.cpp
sed -n '1,180p' python/pygraph/pygraph.cpp
rg -n "DLPACK_(MAJOR|MINOR)_VERSION|DLPACK_VERSION|dlpack.*version|dlpack.h" . -g '!*.pyc' | head -200Repository: NVIDIA/cudnn-frontend
Length of output: 13425
🏁 Script executed:
cat dlpack_version.txt
sed -n '1,80p' python/CMakeLists.txt
sed -n '55,90p' python/pygraph/variant_pack.cpp
sed -n '315,340p' python/pygraph/variant_pack.cpp
git ls-files | rg '(^|/)(dlpack|DLPack)(/|_|\.h)' | head -100Repository: NVIDIA/cudnn-frontend
Length of output: 5008
🌐 Web query:
site:dmlc.github.io/dlpack/latest/python_spec.html __dlpack__ max_version dltensor_versioned producer None
💡 Result:
The max_version argument in the dlpack method is used by a consumer to signal the maximal DLPack version it supports to the producer [1]. When using DLManagedTensorVersioned, the PyCapsule names must include a _versioned suffix, specifically dltensor_versioned and used_dltensor_versioned, rather than the standard dltensor and used_dltensor [1]. While documentation notes that max_version allows producers to choose the appropriate struct (DLManagedTensor or DLManagedTensorVersioned), standard naming conventions for the capsules must be updated accordingly when the versioned struct is in use [1].
Citations:
🌐 Web query:
site:github.com/dmlc/dlpack v1.3 python_spec.h __dlpack__ max_version versioned capsule
💡 Result:
In the context of DLPack and its Python integration, the concept of a "versioned capsule" and the max_version parameter (often discussed in the context of dlpack and ABI evolution) relate to ensuring compatibility between producers and consumers of tensor data [1][2]. 1. DLPack Versioning and Capsules: As DLPack has evolved, it moved from a singular DLManagedTensor struct to a versioned approach [3][4]. The introduction of DLManagedTensorVersioned allows for ABI stability and forward compatibility [4]. When dlpack is called, a producer may return a PyCapsule containing either the legacy DLManagedTensor or the newer DLManagedTensorVersioned [1][2]. The specific capsule name (e.g., "dltensor" vs. "versioned_dltensor") or the version information encoded within the struct helps the consumer determine how to interpret the data [4][2]. 2. The role of max_version: The max_version parameter was proposed as a mechanism for consumers to signal the maximum DLPack version they support to the producer [3][1][5]. This allows the producer to adapt and return a version of the tensor data that the consumer is guaranteed to understand, preventing ABI mismatches [5][2]. While early RFCs discussed including max_version directly in the dlpack call, later developments in the DLPack exchange API have shifted towards more robust C-level exchange conventions [5][6]. 3. v1.3 and C Exchange API: The release of DLPack v1.3 (January 2026) notably updated the fast C exchange API convention to use PyCapsules [7]. This API, accessed via dlpack_c_exchange_api on the Python type, provides a structured, versioned way for libraries to exchange tensors and stream information without the overhead of the Python interpreter [6][7]. This system relies on a DLPackExchangeAPI struct that includes its own version header (major/minor) and allows for safe, backward-compatible interaction between frameworks [6]. In summary, these components (versioned capsules, max_version, and the modern C exchange API) are part of DLPack's strategy to provide a stable, efficient, and cross-framework ABI for tensor memory exchange [4][6][7].
Citations:
- 1: [DISCUSS][RFC] DLPack Versioning and ABI Update dmlc/dlpack#104
- 2: Future ABI compatibility dmlc/dlpack#34
- 3: https://github.com/dmlc/dlpack/releases
- 4: dmlc/dlpack@ca4d00a
- 5: [RFC] C Functions for Speed Exchange and Stream Handling dmlc/dlpack#175
- 6: https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h
- 7: https://github.com/dmlc/dlpack/releases/tag/v1.3
🏁 Script executed:
python3 - <<'PY'
from urllib.request import urlopen
urls = [
"https://raw.githubusercontent.com/dmlc/dlpack/v1.3/docs/source/python_spec.rst",
"https://raw.githubusercontent.com/dmlc/dlpack/v1.3/include/dlpack/dlpack.h",
]
for url in urls:
print(f"\n=== {url} ===")
try:
text = urlopen(url, timeout=10).read().decode()
except Exception as exc:
print(exc)
continue
lines = text.splitlines()
for i, line in enumerate(lines):
if any(term in line for term in ("max_version", "dltensor_versioned", "DLManagedTensorVersioned")):
start = max(0, i - 8)
end = min(len(lines), i + 15)
print("\n".join(f"{j + 1}: {lines[j]}" for j in range(start, end)))
PYRepository: NVIDIA/cudnn-frontend
Length of output: 601
🏁 Script executed:
for url in \
https://raw.githubusercontent.com/dmlc/dlpack/v1.3/docs/source/python_spec.rst \
https://raw.githubusercontent.com/dmlc/dlpack/v1.3/include/dlpack/dlpack.h; do
echo "=== $url ==="
curl -k -L --fail --silent "$url" |
nl -ba |
awk '/max_version|dltensor_versioned|DLManagedTensorVersioned/ {
for (i = NR - 8; i <= NR + 14; ++i) keep[i] = 1
}
{ lines[NR] = $0 }
END {
for (i = 1; i <= NR; ++i)
if (keep[i]) print lines[i]
}'
doneRepository: NVIDIA/cudnn-frontend
Length of output: 401
🏁 Script executed:
python3 - <<'PY'
import ssl
from urllib.request import urlopen
ctx = ssl._create_unverified_context()
for url in (
"https://raw.githubusercontent.com/dmlc/dlpack/v1.3/docs/source/python_spec.rst",
"https://raw.githubusercontent.com/dmlc/dlpack/v1.3/include/dlpack/dlpack.h",
):
print(f"=== {url} ===")
text = urlopen(url, context=ctx, timeout=10).read().decode().splitlines()
hits = [i for i, line in enumerate(text)
if any(x in line for x in ("max_version", "dltensor_versioned", "DLManagedTensorVersioned"))]
shown = set()
for i in hits:
for j in range(max(0, i - 8), min(len(text), i + 15)):
if j not in shown:
print(f"{j + 1}: {text[j]}")
shown.add(j)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 11960
Handle DLPack version negotiation in __dlpack__.
When max_version requests DLPack 1.x, return DLManagedTensorVersioned in a "dltensor_versioned" capsule and support "used_dltensor_versioned" ownership transfer. Otherwise, document and reject the unversioned-only limitation.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/pygraph/variant_pack.cpp` around lines 282 - 303, Update the dlpack
method to negotiate max_version: when it requests DLPack 1.x, construct and
return a DLManagedTensorVersioned in a "dltensor_versioned" capsule, including a
capsule destructor that recognizes and transfers ownership for
"used_dltensor_versioned". For unsupported or unversioned-only requests,
explicitly document the limitation and reject the request rather than silently
returning the unversioned tensor.
There was a problem hiding this comment.
Real gap, and left open deliberately for now. __dlpack__ ignores max_version and always returns an unversioned "dltensor" capsule. The consumer that matters here never takes this path -- tvm-ffi reads the slot through __dlpack_c_exchange_api__, and the capsule exists for cute's compile-time from_dlpack -- so the versioned protocol is unexercised rather than broken. I would rather implement it against a consumer that asks for it than guess at the ownership-transfer semantics; the limitation is now stated at the method.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline python/pygraph/variant_pack.cpp --items all --type method,function --match 'dlpack|VariantPackSlot' --view expanded
rg -n -A45 -B15 '^\s*dlpack\s*\(' python/pygraph/variant_pack.cpp
rg -n -A5 -B5 '__dlpack__|__dlpack_c_exchange_api__|dltensor_versioned|max_version' python/pygraph/variant_pack.cppLength of output: 7756
@YangXu1990uiuc, verification shows that __dlpack__ still ignores max_version and returns the legacy "dltensor" capsule at Line 328 through Line 340.
The deferred implementation is reasonable because tvm-ffi uses __dlpack_c_exchange_api__. The source also provides a versioned managed tensor for that exchange API at Line 365.
However, the current __dlpack__ comment does not state the unversioned-only limitation. It describes capsule ownership only. The finding remains open for a minimal documentation update or a consumer-driven implementation of version negotiation.
You are interacting with an AI system.
_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 NVIDIA#547.
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.
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.
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.
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.
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.
|
Verification on SM100 (B200), against the pushed head: Two things worth stating about what is not covered, so a reviewer does not
clang-format v21.1.6 and black both clean. |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-547-16bd949 |
NVIDIA#547 normalized the variant pack, which changed what an engine is handed: buffers arrive as slots that describe memory rather than as tensors. Two places in the gemm engine were writing THROUGH the caller's buffer with torch methods, which worked only while the buffer happened to be a torch tensor: - a reduction output is seeded with its identity before the kernel runs, via tensor.fill_() - a norm2 reduction is finalized with tensor.sqrt_() The engine owns the first one now. Every seed a reduction uses (0, 1, +-inf) is a 32-bit pattern, so buffers.fill_f32_async drives cuMemsetD32Async and no kernel is needed. The second is unreachable: a norm2 reduction is refused while the graph is lowered, so no plan exists to execute -- recorded as a test rather than left as a live-looking path. A bare device address as the WORKSPACE measured 0 bytes, and Workspace.over read that as "empty" and refused any engine that needs scratch. It means "the pack could not measure it": a raw pointer carries no size and the backend takes one without checking, so refusing here made the same call depend on which plan ran. Zero now skips the size check, here and in the C carve's bounds check. The reason none of this was caught: no test drove a non-trivial gemm flavor through graph.execute(). The direct-call tests construct a fusion chain and invoke the compiled object with torch tensors, which skips operand binding, the pack, and the conversion execute() performs -- exactly the part that changed. test_public_execute_flavors.py covers plain matmul, epilogue fusion, the four reduction modes, norm2's refusal, and the bare-address operand form through the entry point a caller actually has. That last one is xfail: frost reads its extents by axis position, and a bare address describes the operand the way the GRAPH declares it (a matmul's B is [batch, K, N]) rather than the way a caller's buffer reports it. It was broken before this too -- the geometry-less Tensor made it an IndexError instead. The fix is the engine recording which axis is M/N/K at build. Also: two unused imports in engines/base.py that broke the lazy frost boundary, and three kernel docstrings claiming the descriptor builders do not recur.
* Make the public gemm path work for the flavors it claims #547 normalized the variant pack, which changed what an engine is handed: buffers arrive as slots that describe memory rather than as tensors. Two places in the gemm engine were writing THROUGH the caller's buffer with torch methods, which worked only while the buffer happened to be a torch tensor: - a reduction output is seeded with its identity before the kernel runs, via tensor.fill_() - a norm2 reduction is finalized with tensor.sqrt_() The engine owns the first one now. Every seed a reduction uses (0, 1, +-inf) is a 32-bit pattern, so buffers.fill_f32_async drives cuMemsetD32Async and no kernel is needed. The second is unreachable: a norm2 reduction is refused while the graph is lowered, so no plan exists to execute -- recorded as a test rather than left as a live-looking path. A bare device address as the WORKSPACE measured 0 bytes, and Workspace.over read that as "empty" and refused any engine that needs scratch. It means "the pack could not measure it": a raw pointer carries no size and the backend takes one without checking, so refusing here made the same call depend on which plan ran. Zero now skips the size check, here and in the C carve's bounds check. The reason none of this was caught: no test drove a non-trivial gemm flavor through graph.execute(). The direct-call tests construct a fusion chain and invoke the compiled object with torch tensors, which skips operand binding, the pack, and the conversion execute() performs -- exactly the part that changed. test_public_execute_flavors.py covers plain matmul, epilogue fusion, the four reduction modes, norm2's refusal, and the bare-address operand form through the entry point a caller actually has. That last one is xfail: frost reads its extents by axis position, and a bare address describes the operand the way the GRAPH declares it (a matmul's B is [batch, K, N]) rather than the way a caller's buffer reports it. It was broken before this too -- the geometry-less Tensor made it an IndexError instead. The fix is the engine recording which axis is M/N/K at build. Also: two unused imports in engines/base.py that broke the lazy frost boundary, and three kernel docstrings claiming the descriptor builders do not recur. * Seed a reduction output through the driver, on the kernel's stream Three passes over the same fix, each caught by measuring rather than by the suite going green: 1. tensor.fill_() does not exist on a slot -> use cuMemsetD32Async. 2. That lost the STRIDE. torch's fill_ sets every element; a memset writes a contiguous byte range, and the two agree only for a dense buffer. Six strided reduction tests went NaN -- the elements past the first run were never written. 3. One memset per contiguous run is correct and 572 us for a per-row scalar output, against 3.5 for the single kernel torch launches. Correct is not mergeable; nothing in the suite would have flagged it. So: the driver where it is right, which is every contiguous buffer -- and that also puts the seed on the stream the kernel will run on, where tensor.fill_() queues on torch's current stream and is the same stream only by luck. A padded output keeps the torch path it already had, and a padded output arriving as a slot -- the public path, which is where this was broken and where nothing could seed it -- is refused with the measurement in a TODO. The fix is a fill kernel; this is the last place the engine writes through the caller's buffer. Also from review: remaining() refuses rather than returning a negative extent when the workspace size is unknown, and three cases cover the unknown-capacity path end to end (bare-address workspace, undersized-but-known, and the tail refusal). * Pack a reduction identity as the output's dtype, and seed on the MoE stream A memset moves bits, not numbers, so the identity has to be packed as the dtype the kernel reads it back as. int32's identities are the ends of its range and are exactly where that bites: -2**31 packed as float is 0xcf000000, so an int32 MAX reduction returned -822083584 for every input below it. fill_f32_async becomes fill_word_async over a 32-bit pattern, with init_word turning a value into one. The four MoE launchers seeded on the null stream rather than the execute one -- they were the call sites that had no stream to pass before this path moved to the driver, and passing None was not the same thing afterwards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Seed a padded reduction output through the driver too The contiguous case already went through cuMemsetD32Async on the execute-time stream. The padded one fell back to tensor.fill_(), which queues on torch's CURRENT stream -- the same stream only by luck -- and exists at all only while the caller happened to pass a torch tensor. That contradicted what this branch claims to do, so it is gone. strided_fill_plan collapses the layout to the runs a memset can cover and returns the 2D memsets that cover it exactly once. cuMemsetD2D32Async takes a pitch, so a per-row scalar tap is ONE call rather than one per row -- that reading, 572 us at one memset per row, is why the fallback was there. What remains is one call per point of whatever axis is left outside the 2D region, which for a rank-3 output is the batch and is usually 1. The plan is returned before anything is written, and is None for a layout that would write an element twice: a stride of 0 over a real extent, or an outer stride that does not clear the axis below it. Checking only the innermost pair (pitch >= width) is not enough -- shape (2, 2) stride (2, 2) has width 1 and lands both axes on element 2. Also in this commit: a missing space in three linear-attention kernel docstrings (review caught two of the three), and a raw-string pytest.raises pattern. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
3 participants
Normalize the variant pack once, into a C type that is also the DLPack producer
graph.execute()converts whatever the caller passed into one object at thetop. Below that point the cuDNN backend and every python engine see only that,
so neither can behave differently on account of what the caller happened to
hold.
Why (the bugs this closes)
Three cases where the same public call answered two ways, decided by which plan
the heuristics happened to pick — which the caller does not control:
engine got the caller's object untouched and
frost.buffers.proberefused itwith "buffer of type int exposes neither
__cuda_array_interface__nor__dlpack__".override_shapeson a FROST plan.override_shapesspeaks the graph'sdeclaration (a matmul's B is
[batch, K, N]); the buffer is allocated(batch, N, K). Same memory, two axis orders. FROST read its extents off thebuffer, so a max-allocation with a smaller live shape — which is what
override shape is for — ran the whole allocation while the backend ran the
named one. Now the override is applied to the pack, re-expressed in the axis
order the operand already uses, and an engine honours it without knowing the
concept exists.
test_override_shape_inside_a_max_allocation_matches_the_backendruns both paths over the same buffers and compares each to a reference; the
existing coverage only checked FROST against itself.
api_dsl.pycompared caller buffers against torchdtypes by identity, so a JAX or CuPy fp32 buffer failed with
"must be float32; got float32".
Two memory-safety fixes, both found by review:
override_slottookndimfrom the shape but stored whatever stride it wasgiven; a shorter stride array was then read
ndimdeep by any consumer.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 owns its copies now.
DeviceView.__dlpack__had thesame shape and is fixed the same way.
Plus: a non-dense workspace was measured by element count and carved as a byte
range; a short
override_shapessilently kept the original metadata where thebackend rejects the request; and the exchange-vtable cache kept null answers,
so a graph normalized before tvm-ffi's import was pinned to the python fallback
for the life of the process.
How
VariantPackholds the operands in a C type (python/pygraph/variant_pack.cpp)as one
DLTensoreach. That type both consumes__dlpack_c_exchange_api__— the C function table a producer publishes on its type — and implements
it, so a kernel reads a slot through the same fast path it has for a framework
tensor. A producer without the vtable is not an error: python describes those
slots with the existing reader and a mixed pack costs the sum of its parts.
The workspace is carved the same way. Which regions a plan cuts, at what
offsets and shapes, is fixed when the engine builds; only the base pointer
arrives per execute, so one crossing serves them all.
dlpack_version.txtmoves 1.1 → 1.3, whereDLPackExchangeAPIis declared.The wire structs are byte-identical between the two —
sizeofand everyoffsetofofDLTensorandDLManagedTensormatch — so this is acompile-time requirement only; capsules exchanged with a consumer built against
1.1 are unaffected.
Cost
Host time for one
graph.execute(), GDN forward, 4096 tokens / 4 heads / D=128on SM100, measured from a drained queue and swept over burst size:
Of what remains, 8
cudaLaunchKernelExare ~15 (1.85 each, measured untraced —nsys reports 4.06 because CUPTI adds ~2.2 per traced call). Reducing that
launch count is the GDN kernel's own item, not this PR's.
Reading an operand through the vtable is 0.08 us against 1.5 in python, and a
slot converts to a kernel argument slightly cheaper than the torch tensor it
replaces — so framework neutrality is not paid for here.
Where to look hardest
python/pygraph/variant_pack.cpplifetimes. A slot'sDLTensorpointsinto its own vectors, which is why the copy and move constructors are deleted
and the managed export owns copies.
_in_axis_order_ofin_pygraph.py. The permutation between the graph'sdeclared axis order and the buffer's. Both orders rank their axes the same
way by stride — that is what makes them the same memory — and matching the
two rankings is the whole derivation.
_normalize. The alternative was for eachengine to learn about overrides; putting them in the pack is what makes the
two paths answer the same question.
Not in this PR
backend_graph.execute's monkey-patch removal is in here and deserves anexplicit nod. The gemm execute-path rewrite (a closure emitted once per plan,
frost gemm 44 → 18 us) is on a separate branch and will come after this lands.
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