From 4e49c8151801a1c31b2f39afefcde72e7a3b11e4 Mon Sep 17 00:00:00 2001 From: LeSingh1 Date: Sun, 9 Aug 2026 13:43:51 -0700 Subject: [PATCH] fix(core): count blocks, not clusters, in the cooperative-launch check `_check_cooperative_launch` compares the requested grid against a residency limit expressed in thread blocks: max_grid_size = ( kernel.occupancy.max_active_blocks_per_multiprocessor(...) * num_sm ) if prod(config.grid) > max_grid_size: but `config.grid` counts CLUSTERS, not blocks, whenever `cluster` is set. LaunchConfig says so in its own docstring ("When cluster is specified, the grid parameter represents the number of clusters (not blocks)"), and `_to_native_launch_config` implements exactly that, multiplying grid by cluster before filling in gridDimX/Y/Z. So with e.g. cluster=(2, 2, 1) the guard under-counts by 4x: the driver is asked for prod(grid) * prod(cluster) blocks while the check only looks at prod(grid). A cooperative launch that genuinely over-subscribes the device sails past the guard and fails (or deadlocks) inside the driver instead of getting the clean ValueError this function exists to raise. The message was misleading too -- it printed cluster counts as "grid size" and compared them against a block limit. Route the comparison through a small `_cooperative_block_count(config)` helper that applies the same cluster multiplication as `_to_native_launch_config`, and spell the block count out in the error. --- cuda_core/cuda/core/_launcher.pyx | 23 ++++++++++-- cuda_core/docs/source/release/1.2.0-notes.rst | 7 ++++ cuda_core/tests/test_launcher.py | 35 +++++++++++++++++++ 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/cuda_core/cuda/core/_launcher.pyx b/cuda_core/cuda/core/_launcher.pyx index 036189790e0..8d17867c6fd 100644 --- a/cuda_core/cuda/core/_launcher.pyx +++ b/cuda_core/cuda/core/_launcher.pyx @@ -70,14 +70,33 @@ def launch( HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL)) +def _cooperative_block_count(config: LaunchConfig): + """Number of thread blocks a launch of ``config`` asks the driver for. + + ``config.grid`` counts *clusters*, not blocks, whenever ``cluster`` is set + -- see :class:`LaunchConfig` and ``_to_native_launch_config``, which + multiplies the two together before filling in ``gridDim``. Residency limits + are expressed in blocks, so any comparison against one has to go through + here. + """ + if config.cluster is None: + return prod(config.grid) + return prod(config.grid) * prod(config.cluster) + + cdef _check_cooperative_launch(kernel: Kernel, config: LaunchConfig, stream: Stream): dev = stream.device num_sm = dev.properties.multiprocessor_count max_grid_size = ( kernel.occupancy.max_active_blocks_per_multiprocessor(prod(config.block), config.shmem_size) * num_sm ) - if prod(config.grid) > max_grid_size: + num_blocks = _cooperative_block_count(config) + if num_blocks > max_grid_size: # For now let's try not to be smart and adjust the grid size behind users' back. # We explicitly ask users to adjust. x, y, z = config.grid - raise ValueError(f"The specified grid size ({x} * {y} * {z}) exceeds the limit ({max_grid_size})") + detail = f"{x} * {y} * {z}" + if config.cluster is not None: + cx, cy, cz = config.cluster + detail = f"{detail} clusters of {cx} * {cy} * {cz} blocks = {num_blocks} blocks" + raise ValueError(f"The specified grid size ({detail}) exceeds the limit ({max_grid_size} blocks)") diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 120d2c2a253..8fbe0e172e5 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -73,6 +73,13 @@ Fixes and enhancements Windows, both ``ctypes.CFUNCTYPE`` and ``ctypes.WINFUNCTYPE`` are accepted. (`#2439 `__) +- The cooperative-launch residency check now counts blocks when + :class:`LaunchConfig` specifies a ``cluster``. ``grid`` counts clusters in + that case, so the check compared cluster counts against a per-device block + limit and under-counted by ``prod(cluster)`` -- a cooperative launch that + genuinely over-subscribes the device passed the guard. The error message now + spells out the block count as well. + Deprecation Notices ------------------- diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index e5cf05b435d..65b018ade67 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -157,6 +157,41 @@ class _FakeDev: LaunchConfig(grid=1, block=1, is_cooperative=True) +@pytest.mark.agent_authored(model="claude-opus-5") +def test_cooperative_block_count_counts_blocks_not_clusters(monkeypatch): + """The cooperative residency check compares against a *block* limit. + + ``config.grid`` counts clusters whenever ``cluster`` is set (that is what + ``_to_native_launch_config`` multiplies out before filling in ``gridDim``), + so comparing ``prod(config.grid)`` against + ``max_active_blocks_per_multiprocessor * num_sm`` under-counted by + ``prod(config.cluster)`` and let an over-subscribed cooperative launch + through the guard it exists to trip. + + Device is mocked so this runs on any GPU. + """ + from cuda.core import _launch_config as _lc_mod + from cuda.core._launcher import _cooperative_block_count + + class _FakeProps: + cooperative_launch = True + + class _FakeDev: + compute_capability = (9, 0) + properties = _FakeProps() + + monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev()) + + # No cluster: grid is already a block count. + config = LaunchConfig(grid=(2, 3, 1), block=32, is_cooperative=True) + assert _cooperative_block_count(config) == 6 + + # With a cluster the driver is asked for prod(grid) * prod(cluster) blocks. + config = LaunchConfig(grid=(2, 3, 1), cluster=(2, 2, 1), block=32, is_cooperative=True) + assert config.grid == (2, 3, 1) # grid is still stored in cluster units + assert _cooperative_block_count(config) == 24 + + def test_to_native_launch_config_cooperative(monkeypatch): """Covers the is_cooperative branch of _to_native_launch_config; Device is mocked so it runs on any GPU.""" from cuda.bindings import driver