From b9ee0735a5df83ac39875a418374e9358f951eef Mon Sep 17 00:00:00 2001 From: IxMxAMAR Date: Sun, 18 Oct 2026 05:13:34 +0530 Subject: [PATCH 1/2] fix(rocm): cache warp size per device, not for device 0 bnb_host_warp_size() queried device 0 once and reused the result for every device. Its own comment admitted this returns the wrong value on mixed RDNA+CDNA systems, and #2059 did not touch it. That matters because the launch geometry in ops.cu comes from the runtime value while kQuantizeBlockwiseSmall derives NUM_QB from the compile-time BNB_WARP_SIZE. On a wave64 device reached from a wave32 device (or the reverse) the launched block size disagrees with the kernel's NUM_QB: qb_id runs past warp_reduce[]/smem_absmax_value[], absmax writes land in the neighbouring block's slots, and the grid no longer covers all input. Query hipDeviceAttributeWarpSize for the *current* device and fall back to 32 when the query fails, so a failed query degrades to a conservative launch instead of an out-of-range one. --- csrc/ops.cu | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/csrc/ops.cu b/csrc/ops.cu index 16eed4e81..650a8e0d2 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -14,15 +14,31 @@ #include #include -// NOTE: This queries device 0 once and caches the result. On mixed RDNA+CDNA -// systems (warp size 32 vs 64) this will return the wrong value for whichever -// device doesn't match device 0. +// Queried and cached PER DEVICE. The quantization launch configuration below has +// to match the wavefront width of the device the kernel will actually run on: +// kQuantizeBlockwiseSmall derives its NUM_QB from BNB_WARP_SIZE while this derives +// the grid and block size from the runtime value, so a mismatch makes the launched +// block size disagree with the kernel's NUM_QB and indexes smem_absmax_value[] / +// warp_reduce[] out of bounds. Caching a single value taken from device 0 returns +// the wrong width on mixed RDNA (32) + CDNA (64) systems. +// +// A cached value of 0 means "not queried yet"; any failure falls back to the +// narrowest width so a failed query degrades to a conservative launch rather than +// an out-of-range one. +static constexpr int BNB_MAX_CACHED_DEVICES = 16; + static int bnb_host_warp_size() { - static std::atomic warp_size{0}; - int ws = warp_size.load(std::memory_order_relaxed); + static std::atomic warp_sizes[BNB_MAX_CACHED_DEVICES] = {}; + int device = 0; + if (hipGetDevice(&device) != hipSuccess || device < 0 || device >= BNB_MAX_CACHED_DEVICES) { + return 32; + } + int ws = warp_sizes[device].load(std::memory_order_relaxed); if (ws == 0) { - (void)hipDeviceGetAttribute(&ws, hipDeviceAttributeWarpSize, 0); - warp_size.store(ws, std::memory_order_relaxed); + if (hipDeviceGetAttribute(&ws, hipDeviceAttributeWarpSize, device) != hipSuccess || ws <= 0) { + ws = 32; + } + warp_sizes[device].store(ws, std::memory_order_relaxed); } return ws; } From 9404fb1292e2094d62f53c6e5f05a5004b120dc7 Mon Sep 17 00:00:00 2001 From: IxMxAMAR Date: Sun, 18 Oct 2026 06:26:51 +0530 Subject: [PATCH 2/2] fix(rocm): fall back to BNB_WARP_SIZE instead of 32 The kernel derives NUM_QB from the compile-time BNB_WARP_SIZE, so on gfx9 the host side falling back to 32 halves num_qb and sizes the grid for half the work. That is the same out-of-range access this change set is meant to prevent, just moved into the error path. Use the constant on both sides. This also covers device indices past the cache cap, where the previous code returned 32 without querying at all. --- csrc/ops.cu | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/csrc/ops.cu b/csrc/ops.cu index 650a8e0d2..96c58685d 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -22,21 +22,23 @@ // warp_reduce[] out of bounds. Caching a single value taken from device 0 returns // the wrong width on mixed RDNA (32) + CDNA (64) systems. // -// A cached value of 0 means "not queried yet"; any failure falls back to the -// narrowest width so a failed query degrades to a conservative launch rather than -// an out-of-range one. +// A cached value of 0 means "not queried yet". Any failure falls back to the +// compile-time width rather than to 32: the kernel's NUM_QB is derived from that +// same constant, so it is the only value guaranteed to agree with the launch. On +// gfx9 that constant is 64, and falling back to 32 there would reproduce exactly +// the mismatch described above. static constexpr int BNB_MAX_CACHED_DEVICES = 16; static int bnb_host_warp_size() { static std::atomic warp_sizes[BNB_MAX_CACHED_DEVICES] = {}; int device = 0; if (hipGetDevice(&device) != hipSuccess || device < 0 || device >= BNB_MAX_CACHED_DEVICES) { - return 32; + return BNB_WARP_SIZE; } int ws = warp_sizes[device].load(std::memory_order_relaxed); if (ws == 0) { if (hipDeviceGetAttribute(&ws, hipDeviceAttributeWarpSize, device) != hipSuccess || ws <= 0) { - ws = 32; + ws = BNB_WARP_SIZE; } warp_sizes[device].store(ws, std::memory_order_relaxed); }