Skip to content

fix(wgpu): clamp compute dispatch cube count to u16::MAX (NVIDIA Vulkan)#1388

Open
thewtex wants to merge 6 commits into
tracel-ai:mainfrom
thewtex:wgpu-nvidia-param
Open

fix(wgpu): clamp compute dispatch cube count to u16::MAX (NVIDIA Vulkan)#1388
thewtex wants to merge 6 commits into
tracel-ai:mainfrom
thewtex:wgpu-nvidia-param

Conversation

@thewtex

@thewtex thewtex commented Jun 19, 2026

Copy link
Copy Markdown

Summary

A small batch of independent correctness fixes in cubecl-wgpu, plus the same
rank-0 stride bug fixed in two sibling crates that carry a copy of the helper.
Each fix is its own commit so they can be reviewed.

Fixes

1. Clamp dispatch cube count to the enforced per-dimension limit (cubecl-wgpu)

Some adapters — notably the Vulkan backend on NVIDIA hardware — report a
max_compute_workgroups_per_dimension larger than the limit the driver
actually validates against at dispatch time. create_server fed that
adapter-reported value straight into HardwareProperties::max_cube_count, so
the launch logic treated an oversized [N, 1, 1] cube count as legal and let
it reach wgpu, which then rejected the dispatch with a validation error once
N exceeded the real cap of u16::MAX (65_535).

This bites any single large elementwise dispatch — e.g. a 52,428,800-element
launch needs 52_428_800 / 256 = 204_800 workgroups along X, well past 65_535.

The fix clamps the reported value to u16::MAX via a new clamp_max_cube_count
helper, keeping it consistent with the cap already advertised by
<WgpuRuntime as Runtime>::max_cube_count.

2. Detect 8-bit storage via the 8-bit feature struct (cubecl-wgpu)

In register_types, storage8 was computed from ext_feat.buf_16 and its
uniform_and_storage_buffer16_bit_access field — a copy-paste of the
storage16 line above it. As a result 8-bit storage capability tracked the
16-bit device feature, so I8/U8 types could be marked storable based on the
wrong feature. Now reads ext_feat.buf_8.uniform_and_storage_buffer8_bit_access.

3. Name AUTO_GRAPHICS_BACKEND in its invalid-value error (cubecl-wgpu)

The error printed for an unrecognised backend reads the AUTO_GRAPHICS_BACKEND
env var but told the user to check GRAPHICS_BACKEND, sending them to the wrong
variable. Corrected the message.

4. Guard contiguous_strides against rank-0 shapes (cubecl-wgpu, cubecl-cpu, cubecl-runtime)

contiguous_strides iterated (0..rank - 1), which underflows and panics when
rank == 0 (a scalar / rank-0 shape). Changed to (0..rank.saturating_sub(1))
so a rank-0 shape yields empty strides. The same helper is duplicated in
cubecl-cpu and cubecl-runtime, fixed in all three.

Notes:

  • cubecl-std's contiguous_strides uses a different, subtraction-free
    implementation and is not affected.
  • The pitched-allocation path in ContiguousMemoryLayoutPolicy
    (cubecl-runtime) already guards its rank - 2 loop behind rank > 2, so
    only the standalone helper needed the fix.

Testing

  • New unit tests for clamp_max_cube_count (caps inflated limits, preserves
    honest ones, agrees with WgpuRuntime::max_cube_count) and for
    contiguous_strides (rank-0 returns empty, row-major correctness) in each of
    the three crates.
  • cargo test, cargo clippy --all-targets, and cargo fmt --check pass for
    the touched crates; cubecl-wgpu additionally checked under --features spirv
    to compile the Vulkan path.

thewtex added 6 commits June 19, 2026 16:27
Some adapters — notably the Vulkan backend on NVIDIA hardware — report a `max_compute_workgroups_per_dimension` larger than the limit the driver actually validates against at dispatch time. `create_server` fed that adapter-reported value straight into `HardwareProperties::max_cube_count`, so the launch logic treated an oversized `[N, 1, 1]` cube count as legal and let it reach wgpu, which rejected the dispatch with a validation error once `N` exceeded the real cap of `u16::MAX` (65_535).

This bites any single large elementwise dispatch: e.g. a 52,428,800-element launch needs 204,800 workgroups along X (52_428_800 / 256), well past 65_535.

Clamp the reported limit to `u16::MAX`, keeping it consistent with the cap already advertised by `<WgpuRuntime as Runtime>::max_cube_count`. The clamp is factored into a small `clamp_max_cube_count` helper with unit tests.
`register_vulkan_features` computed `storage8` from `ext_feat.buf_16` and its `uniform_and_storage_buffer16_bit_access` field — a copy-paste of the `storage16` line above it. As a result 8-bit storage capability tracked the 16-bit feature rather than the actual `PhysicalDevice8BitStorageFeatures`, so I8/U8 types could be marked storable based on the wrong device feature.

Read `ext_feat.buf_8.uniform_and_storage_buffer8_bit_access` instead.

Not unit-testable without a live Vulkan device; verified to compile under `--features spirv`.
The error printed for an unrecognised backend reads the `AUTO_GRAPHICS_BACKEND` env var but told the user to check `GRAPHICS_BACKEND`, sending them to the wrong variable. Use the correct name in the message.
`contiguous_strides` iterated `(0..rank - 1)`, which underflows and panics when `rank == 0` (a scalar/rank-0 shape). Use `rank.saturating_sub(1)` so a rank-0 shape yields empty strides instead of panicking. Adds unit tests for the rank-0 case and row-major correctness.
Same rank-0 underflow as the wgpu `contiguous_strides` fix: the `(0..rank - 1)` loop bound underflows and panics when `rank == 0` (a scalar/rank-0 shape). Use `rank.saturating_sub(1)` so a rank-0 shape yields empty strides. Adds unit tests for the rank-0 case and row-major correctness.
Same rank-0 underflow as the wgpu `contiguous_strides` fix: the `(0..rank - 1)` loop bound underflows and panics when `rank == 0` (a scalar/rank-0 shape). Use `rank.saturating_sub(1)` so a rank-0 shape yields empty strides. The pitched-allocation path in `ContiguousMemoryLayoutPolicy` already guards its `rank - 2` loop behind `rank > 2`, so only this helper needed the fix. Adds unit tests for the rank-0 case and row-major correctness.
@thewtex thewtex changed the title wgpu nvidia param fix(wgpu): clamp compute dispatch cube count to u16::MAX (NVIDIA Vulkan) Jun 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant