feat(rocm): port the BitNet bitlinear_matmul kernel

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp

@@ -4839,6 +4839,13 @@ bool custom_kernels_available() {
     return mlxcel::custom_kernels_available();
 }
 
+// True when this backend has a BitLinear kernel port. Metal, CUDA and, since
+// issue #1862, ROCm. Kept apart from `custom_kernels_available` because kernels
+// are ported one at a time: ROCm has this one and not the rest.
+bool bitlinear_kernel_available() {
+    return mlxcel::gpu_kernel_backend() != mlxcel::GpuKernelBackend::None;
+}
+
 // Top-p (nucleus) filtering.
 // Not compiled: MLX v0.31.x Scan primitive (cumsum) lacks output_shapes,
 // which causes "CumSum cannot infer output shapes" when used inside

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h

@@ -1318,6 +1318,10 @@ bool gpu_backend_available();
 // only says a GPU exists.
 bool custom_kernels_available();
 
+// True when this backend has a BitLinear kernel port: Metal, CUDA or ROCm
+// (issues #1803, #1862).
+bool bitlinear_kernel_available();
+
 // Fused sampling: top-k + top-p + min-p on the untempered distribution, then
 // one temperature scaling, then the categorical draw, in a single function
 // call to minimize FFI round-trips (chain order per issue #1379).

src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp

@@ -144,6 +144,97 @@ namespace {
         return holder;
     }
 
+    // ROCm port (issue #1862). The body is the CUDA one with two changes.
+    //
+    // `__shfl_down_sync` exists in HIP only as a compatibility shim whose mask
+    // argument is ignored, so the native `__shfl_down(var, delta, width)` is
+    // used instead and the width is stated rather than implied.
+    //
+    // The width is 32 because that is the wavefront this backend is built for:
+    // `WARP_SIZE` in mlx/backend/rocm/kernel_utils.hpp is 32 and gfx1151 reports
+    // warp=32 at device bind. A wave64 target (gfx9) would need the reduction to
+    // start at 32, not 16, so the launch below pins the block to 32 lanes and
+    // this constant travels with it.
+    static const char* BITLINEAR_HIP_SOURCE = R"(
+        // The fold below starts at 16, which is correct only for a 32-lane
+        // wavefront. On a wave64 target (CDNA: gfx90a, gfx942) it would fold
+        // half the lanes and the missing half would not be visible: the result
+        // stays finite and plausible, just wrong. Fail to compile instead.
+        //
+        // Checked with the preprocessor, not `static_assert(warpSize == 32)`:
+        // HIP's `warpSize` is an object with an `operator int()`, so that form
+        // does not compile ("non-constexpr function cannot be used in a
+        // constant expression"). Both macro spellings are accepted because the
+        // trailing-underscore one is the older name.
+        #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32
+        #error "bitlinear_matmul_hip assumes a 32-lane wavefront"
+        #endif
+        #if defined(__AMDGCN_WAVEFRONT_SIZE) && __AMDGCN_WAVEFRONT_SIZE != 32
+        #error "bitlinear_matmul_hip assumes a 32-lane wavefront"
+        #endif
+        constexpr int Mc = 4;
+        uint32_t tid       = blockIdx.y;     // batch * out/4
+        uint32_t in_offset = threadIdx.x;    // lane 0..31
+
+        uint32_t batch_idx = tid / (out_features / 4);
+        uint32_t row_idx   = tid % (out_features / 4);
+
+        float sum[4] = {0.0f, 0.0f, 0.0f, 0.0f};
+        for (uint32_t i = in_offset * Mc; i < (uint32_t)in_features; i += 32u * Mc) {
+            float v[Mc];
+            for (int j = 0; j < Mc; j++) {
+                v[j] = (float)x[batch_idx * in_features + i + j];
+            }
+            for (int j = 0; j < Mc; j++) {
+                uint32_t w = packed_weights[row_idx * in_features + i + j];
+                sum[0] += v[j] * ((float)(w & 3u) - 1.0f);
+                sum[1] += v[j] * ((float)((w >> 2) & 3u) - 1.0f);
+                sum[2] += v[j] * ((float)((w >> 4) & 3u) - 1.0f);
+                sum[3] += v[j] * ((float)((w >> 6) & 3u) - 1.0f);
+            }
+        }
+        #pragma unroll
+        for (int o = 16; o > 0; o >>= 1) {
+            sum[0] += __shfl_down(sum[0], o, 32);
+            sum[1] += __shfl_down(sum[1], o, 32);
+            sum[2] += __shfl_down(sum[2], o, 32);
+            sum[3] += __shfl_down(sum[3], o, 32);
+        }
+        if (in_offset == 0u) {
+            float scale = invert_weight_scales ? 1.0f / (float)weight_scale[0]
+                                               : (float)weight_scale[0];
+            for (int i = 0; i < 4; i++) {
+                out[batch_idx * out_features + row_idx + i * (out_features / 4)] =
+                    (T)(sum[i] * scale);
+            }
+        }
+    )";
+
+    struct BitlinearKernelHolderHip {
+        std::optional<mlx::core::fast::CustomKernelFunction> kernel;
+        bool initialized = false;
+        mlx::core::fast::CustomKernelFunction& get() {
+            if (!initialized) {
+#ifdef MLXCEL_BRIDGE_ROCM_BACKEND
+                kernel = mlx::core::fast::hip_kernel(
+                    "bitlinear_matmul_hip",
+                    {"x", "packed_weights", "weight_scale"},
+                    {"out"},
+                    BITLINEAR_HIP_SOURCE);
+#else
+                throw std::runtime_error(
+                    "[bitlinear_matmul] this build has no ROCm backend");
+#endif
+                initialized = true;
+            }
+            return *kernel;
+        }
+    };
+    static BitlinearKernelHolderHip& get_bitlinear_kernel_hip() {
+        static BitlinearKernelHolderHip holder;
+        return holder;
+    }
+
 }
 
 std::unique_ptr<MlxArray> bitlinear_matmul(
@@ -165,24 +256,24 @@ std::unique_ptr<MlxArray> bitlinear_matmul(
 
     // mx.fast.metal_kernel throws on CUDA, so dispatch the cuda_kernel port
     // there. metal::is_available() is false on a CUDA-only build.
-    // Named backends rather than a negation of Metal (issue #1803). Backends
-    // with no port throw here, and the bridge declares this function `Result`,
-    // so that reaches the caller as an error instead of ending the process:
-    // this op is the one fused kernel with no graph fallback, so before this a
-    // CPU-only build and a ROCm build both terminated on a BitNet checkpoint.
+    // Three real ports now (issues #1803, #1862), so the choice names the
+    // backend rather than negating Metal. `None` is the only value left with no
+    // kernel, and the bridge declares this function `Result`, so its throw
+    // reaches the caller as an error instead of ending the process.
     auto& kernel = [&]() -> mlx::core::fast::CustomKernelFunction& {
         switch (mlxcel::gpu_kernel_backend()) {
             case mlxcel::GpuKernelBackend::Metal:
                 return get_bitlinear_kernel().get();
             case mlxcel::GpuKernelBackend::Cuda:
                 return get_bitlinear_kernel_cuda().get();
             case mlxcel::GpuKernelBackend::Rocm:
+                return get_bitlinear_kernel_hip().get();
             case mlxcel::GpuKernelBackend::None:
                 break;
         }
         throw std::runtime_error(
             "[bitlinear_matmul] no BitLinear kernel port for this GPU backend; "
-            "BitNet needs Metal or CUDA (ROCm port: lablup/mlxcel#1862)");
+            "BitNet needs Metal, CUDA or ROCm");
     }();
     std::vector<std::pair<std::string, mlx::core::fast::TemplateArg>> ta = {
         {"T", T},

src/lib/mlxcel-core/src/lib.rs

@@ -1794,13 +1794,13 @@ mod ffi {
         /// [out_features/4, in_features] uint8 (2-bit ternary, 4 rows/byte),
         /// scaled by `weight_scale[0]` (inverted unless linear_class is
         /// autobitlinear).
-        /// BitLinear ternary matmul, with Metal and CUDA kernel ports.
+        /// BitLinear ternary matmul, with Metal, CUDA and ROCm kernel ports.
         ///
-        /// Returns `Err` on a backend with neither, which today means ROCm and
-        /// CPU-only builds, instead of ending the process: the op has no graph
+        /// Returns `Err` on a backend with none of them, which today means a
+        /// CPU-only build, instead of ending the process: the op has no graph
         /// fallback, so before issue #1803 the `fast::*_kernel` throw crossed a
         /// `noexcept` extern and terminated. Callers that want to refuse early
-        /// should ask [`custom_kernels_available`].
+        /// should ask [`bitlinear_kernel_available`].
         fn bitlinear_matmul(
             x: &MlxArray,
             packed_weights: &MlxArray,
@@ -2107,6 +2107,14 @@ mod ffi {
         /// process.
         fn custom_kernels_available() -> bool;
 
+        /// True when this backend has a BitLinear kernel port, that is Metal,
+        /// CUDA or ROCm (issues #1803, #1862). Separate from
+        /// `custom_kernels_available` on purpose: kernels are ported one at a
+        /// time, so "this backend has fused kernels" and "this backend has
+        /// *this* kernel" stopped being the same question the moment ROCm got
+        /// its first port.
+        fn bitlinear_kernel_available() -> bool;
+
         /// True when the MLX Metal backend is available at runtime (macOS
         /// Apple Silicon). False on CUDA-only and CPU-only builds. Mirrors the
         /// `mlx::core::metal::is_available()` gate that picks the Metal vs CUDA

src/models/bitnet.rs

@@ -42,13 +42,12 @@ use std::path::Path;
 /// this refuses only where no kernel exists at all, and it refuses at load
 /// rather than mid-request.
 fn reject_without_bitlinear_kernel() -> Result<(), String> {
-    if mlxcel_core::custom_kernels_available() {
+    if mlxcel_core::bitlinear_kernel_available() {
         return Ok(());
     }
     Err(
-        "BitNet checkpoints need the bitlinear_matmul kernel, which has Metal and CUDA ports only. \
-         This backend has neither and the op has no graph fallback, so the model cannot run here. \
-         The ROCm port is tracked as lablup/mlxcel#1862."
+        "BitNet checkpoints need the bitlinear_matmul kernel, which has Metal, CUDA and ROCm ports. \
+         This build has no GPU backend, and the op has no graph fallback, so the model cannot run here."
             .to_string(),
     )
 }
@@ -534,13 +533,90 @@ mod tests {
     /// packed[0, col] bits {row0,row1,row2,row3} -> bytes [134, 137, 100, 97].
     /// x = [1, 2, 3, 4], scale = 2.0 (not inverted) =>
     ///   y = (x @ W^T) * 2 = [-4, -4, 8, 6].
+    /// Exercises the cross-lane reduction, which the known-case test above does
+    /// not reach.
+    ///
+    /// The kernel walks `in_features` with lane `L` starting at `L * 4` and
+    /// striding by `32 * 4 = 128`, then folds the 32 lanes with `simd_sum` on
+    /// Metal and `__shfl_down` on CUDA and ROCm. At `in_features = 4` only lane
+    /// 0 ever enters that loop, so the fold adds one real value to 31 zeros and
+    /// a wrong fold width produces the right answer anyway. 256 gives every lane
+    /// two iterations, so the width, the stride and the 2-bit unpacking all have
+    /// to be right at once.
+    ///
+    /// The reference is computed here rather than written down: a hand-computed
+    /// expectation over 256 inputs is easier to get wrong than the kernel, and a
+    /// wrong expectation that passes is worse than no test.
+    #[test]
+    fn bitlinear_matmul_folds_every_lane() {
+        if !mlxcel_core::bitlinear_kernel_available() {
+            eprintln!(
+                "skipped bitlinear_matmul_folds_every_lane: this backend has no BitLinear kernel port"
+            );
+            return;
+        }
+        const IN: usize = 256;
+        const OUT: usize = 8;
+        const ROWS: usize = OUT / 4;
+        const BATCH: usize = 2;
+        let scale_value = 0.5f32;
+
+        let mut state = 0x1234_5678u32;
+        let mut next = move || {
+            state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
+            state >> 8
+        };
+
+        let xs: Vec<f32> = (0..BATCH * IN).map(|_| (next() % 9) as f32 - 4.0).collect();
+        let packed: Vec<u8> = (0..ROWS * IN).map(|_| (next() % 256) as u8).collect();
+
+        // out[b][row + q * ROWS] = scale * sum_k x[b][k] * (((packed[row][k] >> 2q) & 3) - 1)
+        let mut expected = vec![0.0f32; BATCH * OUT];
+        for b in 0..BATCH {
+            for row in 0..ROWS {
+                for q in 0..4 {
+                    let mut acc = 0.0f32;
+                    for k in 0..IN {
+                        let w = packed[row * IN + k];
+                        let t = ((w >> (2 * q)) & 3) as f32 - 1.0;
+                        acc += xs[b * IN + k] * t;
+                    }
+                    expected[b * OUT + row + q * ROWS] = acc * scale_value;
+                }
+            }
+        }
+
+        let x = mlxcel_core::from_slice_f32(&xs, &[BATCH as i32, IN as i32]);
+        let packed_arr = mlxcel_core::from_bytes(
+            &packed,
+            &[ROWS as i32, IN as i32],
+            mlxcel_core::dtype::UINT8,
+        );
+        let scale = mlxcel_core::from_slice_f32(&[scale_value], &[1]);
+        let y =
+            mlxcel_core::bitlinear_matmul(&x, &packed_arr, &scale, IN as i32, OUT as i32, false)
+                .expect("bitlinear_matmul");
+        let want = mlxcel_core::from_slice_f32(&expected, &[BATCH as i32, OUT as i32]);
+        let diff = mlxcel_core::abs(&mlxcel_core::subtract(&y, &want));
+        let max_abs = mlxcel_core::item_f32(&mlxcel_core::max_axis(
+            &mlxcel_core::reshape(&diff, &[-1]),
+            -1,
+            false,
+        ));
+        assert!(
+            max_abs < 1e-3,
+            "bitlinear_matmul disagrees with the host reference by {max_abs}; \
+             a fold narrower than 32 lanes shows up here and not in the 4-input case"
+        );
+    }
+
     #[test]
     fn bitlinear_matmul_known_ternary_case() {
         let x = mlxcel_core::from_slice_f32(&[1.0, 2.0, 3.0, 4.0], &[1, 4]);
         let packed =
             mlxcel_core::from_bytes(&[134u8, 137, 100, 97], &[1, 4], mlxcel_core::dtype::UINT8);
         let scale = mlxcel_core::from_slice_f32(&[2.0], &[1]);
-        if !mlxcel_core::custom_kernels_available() {
+        if !mlxcel_core::bitlinear_kernel_available() {
             // No BitLinear port on this backend, so there is nothing to check.
             // Printed rather than silent: a gate run that skips this should say
             // so, not look like it passed (issue #1803).