python: add the ensure_current_context #612 imports but never defined by vedaanta · Pull Request #638 · NVIDIA/cudnn-frontend

vedaanta · 2026-08-17T22:56:15Z

Fixes #634.

What

#612 added from ._device import ensure_current_context to _pygraph.py (used on the python-engine execute path at _pygraph.py:1783) without adding the function to _device.py, so import cudnn fails with ImportError on any package built from current develop, and every python test fails at collection.

Hoist the established pattern from linear_attention/cutile/kernels/common.py (ensure_cuda_context) into _device.py using this module's own driver shim: bind a driver context to the calling thread when none is bound. A JIT engine talks to the driver directly, which reads the calling THREAD's context stack - and an autograd backward runs on a worker thread where cudaSetDevice has only moved the runtime's thread-local slot. Prefer the stream's context (cuStreamGetCtx), else retain the runtime-current device's primary one; best-effort like the cutile original (a context this cannot establish fails at the launch, with the launch's own diagnostics).

A follow-up could de-duplicate cutile's ensure_cuda_context onto this; kept out of this PR to stay minimal.

Verification (H100, cuDNN dev build, package built from this branch)

import cudnn succeeds; create_handle() returns a Handle.
A bare worker thread (no bound context) calling ensure_current_context(None) after the process established a primary context ends up with a bound context.
test_mhas_v2 -k "random_bwd_ragged or fp8_fwd_ragged": 258 passed / 7 skipped / 23 failed - the 23 are the pre-existing no-implementation typing issue fixed separately in sdpa: throw cudnnGraphNotSupportedException when no implementation supports the attributes #633, unrelated to this import fix.

Summary by CodeRabbit

Bug Fixes
- Improved CUDA context handling to help ensure operations run on the correct device and stream.
- Added more resilient behavior when CUDA is unavailable, device information cannot be determined, or a context is already active.
- Prevented context-related errors from interrupting supported workflows.

…efined NVIDIA#612 added 'from ._device import ensure_current_context' to _pygraph.py (used on the python-engine execute path) without adding the function to _device.py, so 'import cudnn' fails with ImportError on any package built from develop. Hoist the established pattern from linear_attention/cutile/kernels/common.py (ensure_cuda_context) into _device.py using this module's driver shim: bind a driver context to the calling thread when none is bound - a JIT engine talks to the driver directly, which reads the calling THREAD's context stack, and an autograd backward runs on a worker thread where cudaSetDevice has only moved the runtime's thread-local slot. Prefer the stream's context, else retain the runtime-current device's primary one; best-effort like the cutile original. Fixes NVIDIA#634. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

coderabbitai · 2026-08-17T22:56:32Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f4c6875e-35dc-4910-a173-3a2e4e60a4b2

📥 Commits

Reviewing files that changed from the base of the PR and between 189966f and e69273b.

📒 Files selected for processing (1)

python/cudnn/_device.py

Included review availability: Your plan includes up to 12 reviews per rolling hour; 8 remain after this review.

📝 Walkthrough

Walkthrough

The change adds ensure_current_context(stream=None) to bind a suitable CUDA driver context to the calling thread when needed. It leaves execution unchanged when CUDA or context lookup is unavailable.

Changes

CUDA context initialization

Layer / File(s)	Summary
Best-effort context binding `python/cudnn/_device.py`	Adds `ensure_current_context`, which prefers the supplied stream context and otherwise uses the runtime-current device’s retained primary context. It suppresses failures and does nothing when a context is already current.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to e6927

The PR adds the missing context-binding helper needed for package imports and execution paths; no actionable merge-blocking risk remains beyond normal checks and review.

Possibly related PRs

NVIDIA/cudnn-frontend#635: Adds the same ensure_current_context(stream=None) function in python/cudnn/_device.py.

Suggested reviewers: yangxu1990uiuc

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title identifies the missing ensure_current_context function and the import defect addressed by the pull request.
Description check	✅ Passed	The description explains the defect, implementation, linked issue, compatibility context, and verification results, with sufficient detail despite missing template headings.
Linked Issues check	✅ Passed	The change adds the missing function required by issue `#634` and resolves the ImportError that prevented cudnn imports and test collection.
Out of Scope Changes check	✅ Passed	The changes are limited to implementing ensure_current_context and directly support the linked issue objective.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

vedaanta · 2026-08-17T23:42:45Z

@cudnn-ci-bot run frost

cudnn-ci-bot · 2026-08-17T23:43:09Z

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-638-e69273b
Pipeline: 63199028
Targets: frost

…g cannot name it (#635) ### What `DeviceInfo.oversized_shared_memory_per_block` gated the query on BOTH the driver version and whether cuda-python's `CUdevice_attribute` carries `CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK`. Those are independent axes: the DRIVER decides whether the mode exists, the BINDING only decides how to ask. Refusing on the second gives up the carveout on a live combination — driver 13.5 with a cuda-python 13.3.1 binding reports 0 although the device really offers 327 KiB. On SM 10.7 block-scale GEMM that costs 3 AB pipeline stages (8 → 5) for nothing. ### Change Keep the driver gate, drop the binding one: name the enum member when the binding has it, else pass its ordinal, which that binding forwards fine. Only bindings old enough to reject a bare int (they read `attrib.value`) genuinely cannot make the query, and those return 0 through a narrow `AttributeError` arm. A real driver failure still raises rather than being masked. ### Test SM 10.7, driver 13050, cuda-python 13.3.1: - the query goes 0 → 334848, back to the ceiling the part reports - `test_device_info` + `test_public_execute_flavors` + `test_build_device` + `test_stream_respect` + `test_tile_select_analytic`: 46 passed / 2 skipped - sm107 block-scale kernels (`-k "sm107 and (128x128 or mixed_cga)"`): 42 passed ### Notes - The `ensure_current_context` fix this branch originally carried landed upstream as #638; that commit was dropped from this PR. - Supersedes #615, which fixed the same class of bug in `frost/device.py` before #612 moved the query into `cudnn/_device.py` — #615 can be closed.

…efined (NVIDIA#638) NVIDIA#612 added 'from ._device import ensure_current_context' to _pygraph.py (used on the python-engine execute path) without adding the function to _device.py, so 'import cudnn' fails with ImportError on any package built from develop. Hoist the established pattern from linear_attention/cutile/kernels/common.py (ensure_cuda_context) into _device.py using this module's driver shim: bind a driver context to the calling thread when none is bound - a JIT engine talks to the driver directly, which reads the calling THREAD's context stack, and an autograd backward runs on a worker thread where cudaSetDevice has only moved the runtime's thread-local slot. Prefer the stream's context, else retain the runtime-current device's primary one; best-effort like the cutile original. Fixes NVIDIA#634. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…g cannot name it (NVIDIA#635) ### What `DeviceInfo.oversized_shared_memory_per_block` gated the query on BOTH the driver version and whether cuda-python's `CUdevice_attribute` carries `CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK`. Those are independent axes: the DRIVER decides whether the mode exists, the BINDING only decides how to ask. Refusing on the second gives up the carveout on a live combination — driver 13.5 with a cuda-python 13.3.1 binding reports 0 although the device really offers 327 KiB. On SM 10.7 block-scale GEMM that costs 3 AB pipeline stages (8 → 5) for nothing. ### Change Keep the driver gate, drop the binding one: name the enum member when the binding has it, else pass its ordinal, which that binding forwards fine. Only bindings old enough to reject a bare int (they read `attrib.value`) genuinely cannot make the query, and those return 0 through a narrow `AttributeError` arm. A real driver failure still raises rather than being masked. ### Test SM 10.7, driver 13050, cuda-python 13.3.1: - the query goes 0 → 334848, back to the ceiling the part reports - `test_device_info` + `test_public_execute_flavors` + `test_build_device` + `test_stream_respect` + `test_tile_select_analytic`: 46 passed / 2 skipped - sm107 block-scale kernels (`-k "sm107 and (128x128 or mixed_cga)"`): 42 passed ### Notes - The `ensure_current_context` fix this branch originally carried landed upstream as NVIDIA#638; that commit was dropped from this PR. - Supersedes NVIDIA#615, which fixed the same class of bug in `frost/device.py` before NVIDIA#612 moved the query into `cudnn/_device.py` — NVIDIA#615 can be closed.

brandonfzhang approved these changes Aug 17, 2026

View reviewed changes

vedaanta merged commit 0e708e3 into NVIDIA:develop Aug 18, 2026
1 check passed

brandonfzhang mentioned this pull request Aug 19, 2026

Restructure attention benchmarks and update benchmarking artifacts (backend 9.26.0.17, frontend b21cd3f) #657

Merged

YangXu1990uiuc mentioned this pull request Aug 20, 2026

fix: bind a CUDA context on the calling thread — the right one, and on both sides of the boundary #626

Merged

3 tasks

Navigation Menu

Uh oh!

python: add the ensure_current_context #612 imports but never defined - #638

python: add the ensure_current_context #612 imports but never defined#638
vedaanta merged 1 commit into
NVIDIA:developfrom
vedaanta:vagarwalla/fix-ensure-current-context

vedaanta commented Aug 17, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Aug 17, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

vedaanta commented Aug 17, 2026

Uh oh!

cudnn-ci-bot commented Aug 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

vedaanta commented Aug 17, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Verification (H100, cuDNN dev build, package built from this branch)

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

vedaanta commented Aug 17, 2026

Uh oh!

cudnn-ci-bot commented Aug 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vedaanta commented Aug 17, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 17, 2026 •

edited

Loading