Apply ROCm grid-overflow cap hygiene to _float_to_fusednbitrowwise_cuda_kernel

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2931

Tier-1 follow-up to D104903707 / D104937969 / D105282095 / D105286434
— apply the canonical ROCm grid-overflow cap to a `quantize_ops/`
launch site identified by the audit at:
`/home/bensonma415/.llms/plans/permute_metric_layout_etc_rocm_grid_overflow_audit.plan.md`

`_float_to_fusednbitrowwise_cuda_kernel` is launched with grid
`cuda_calc_xblock_count(nrows, 256)` and block `256`. The host fn
reads `const int nrows = input.size(0)`, so total threads are
naturally bounded below 2**31 by valid inputs (the int truncation
silently caps at INT_MAX rows). Thus this kernel cannot strictly
trigger HIP's 2**32 cap with valid input — **the cap added here is
defensive**, applied for hygienic consistency with the broader audit
and to match the canonical pattern from D104903707 / D104937969.

The kernel grid-strides over `row` (line 31), so the cap is
correctness-preserving regardless. NVIDIA codegen sees the host-side
`#ifdef USE_ROCM` block as a plain alias and remains a no-op delta
in PTX/SASS.

Reviewed By: henrylhtsang

Differential Revision: D105286713

fbshipit-source-id: 95f99caa203861b75b45a283f8aa83813ad1bc4e

Author: q10

fbgemm_gpu/src/quantize_ops/quantize_fused_nbit_rowwise.cu

@@ -7,6 +7,7 @@
  */
 
 #include "common.cuh"
+#include "fbgemm_gpu/utils/cuda_utilities.cuh"
 
 using Tensor = at::Tensor;
 
@@ -144,7 +145,12 @@ Tensor _float_to_fusednbitrowwise_gpu_t(
   }
 
   constexpr auto threads_per_block = 256;
-  const auto num_blocks = cuda_calc_xblock_count(nrows, threads_per_block);
+  // HIP enforces a hard limit of 2^32 total threads per launch.
+  // _float_to_fusednbitrowwise_cuda_kernel grid-strides over `row`, so
+  // capping is correctness-preserving.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto num_blocks = utils::cuda::cap_grid_dim_x_from_workload(
+      nrows, threads_per_block, at::cuda::getCurrentCUDAStream());
   // think unsigned as we use 0, 255
 
   FBGEMM_DISPATCH_FLOATING_TYPES(

fbgemm_gpu/test/quantize/fused_nbit_rowwise_test.py

@@ -28,9 +28,14 @@
 # pyre-fixme[16]: Module `common` has no attribute `open_source`.
 if open_source:
     # pyre-ignore[21]
-    from test_utils import gpu_available, optests
+    from test_utils import gpu_available, gpu_memory_lt_gb, gpu_unavailable, optests
 else:
-    from fbgemm_gpu.test.test_utils import gpu_available, optests
+    from fbgemm_gpu.test.test_utils import (
+        gpu_available,
+        gpu_memory_lt_gb,
+        gpu_unavailable,
+        optests,
+    )
 
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
 
@@ -412,6 +417,61 @@ def test_quantize_and_dequantize_op_cpu_and_cuda(
                 )
             torch.testing.assert_close(dequantized_data.cpu(), reference)
 
+    @optests.dontGenerateOpCheckTests("Large grid HIP regression — uses ~4 GB HBM")
+    @unittest.skipIf(*gpu_unavailable)
+    # pyre-ignore[56]: gpu_memory_lt_gb is typed (bool, str), but Pyre cannot
+    # infer the unpacked arg type through the open-source import shim above.
+    @unittest.skipIf(*gpu_memory_lt_gb(8))
+    def test_float_to_fusednbitrowwise_large_grid(self) -> None:
+        """
+        Regression test for the HIP grid-cap on
+        `_float_to_fusednbitrowwise_cuda_kernel`
+        (quantize_ops/quantize_fused_nbit_rowwise.cu line ~152).
+
+        Block: 256. Grid: cuda_calc_xblock_count(nrows, 256). Total threads
+        ~= nrows. The host fn reads `nrows = input.size(0)` into an `int`,
+        so total threads are naturally bounded below 2**31 by valid inputs;
+        thus this kernel cannot strictly trigger HIP's 2**32 cap. The cap
+        added here is defensive — consistent with the canonical pattern
+        applied across the audit and harmless for kernels that grid-stride.
+
+        Two-part test:
+        1. Launch-survival at large grid (nrows ~ 2**28) — verifies the
+           cap-path compiles and the launch succeeds.
+        2. Small-scale Tier-A CPU-oracle correctness — verifies the
+           kernel produces the same quantized bytes as the CPU dispatch
+           for non-trivial input.
+        """
+        device = torch.device(torch.accelerator.current_accelerator() or "cuda")
+
+        # ---- Step 1: launch-survival at large nrows. ----
+        nrows = 1 << 28  # 256M rows
+        ncols = 4  # smallest valid: ncols % (2 * num_elem_per_byte) == 0 for bit_rate=4
+        inp = torch.zeros((nrows, ncols), dtype=torch.float32, device=device)
+        output = torch.ops.fbgemm.FloatToFusedNBitRowwiseQuantizedSBHalf(inp, 4)
+        # Output shape: (nrows, ncols / num_elem_per_byte + 2 * sizeof(at::Half))
+        #             = (nrows, 4/2 + 4) = (nrows, 6).
+        self.assertEqual(output.shape[0], nrows)
+        del inp, output
+        torch.cuda.empty_cache()
+
+        # ---- Step 2: Tier-A CPU-oracle correctness at small scale. ----
+        small_nrows = 2 * 256 + 7  # multi-block on the small kernel grid
+        small_ncols = 8  # multi-element per row, valid for bit_rate=4
+        # Sentinel non-zero values at start / middle / end of each row.
+        small_inp_cpu = torch.zeros((small_nrows, small_ncols), dtype=torch.float32)
+        small_inp_cpu[0, 0] = 1.0
+        small_inp_cpu[small_nrows // 2, small_ncols // 2] = 2.0
+        small_inp_cpu[-1, -1] = 3.0
+
+        out_cpu = torch.ops.fbgemm.FloatToFusedNBitRowwiseQuantizedSBHalf(
+            small_inp_cpu, 4
+        )
+        out_gpu = torch.ops.fbgemm.FloatToFusedNBitRowwiseQuantizedSBHalf(
+            small_inp_cpu.to(device), 4
+        )
+        torch.testing.assert_close(out_gpu.cpu(), out_cpu)
+
 
 if __name__ == "__main__":
     unittest.main()