Skip to content

feat(rocm): port the BitNet bitlinear_matmul kernel to mlxcelverse #1862

Description

@inureyes

Part of #1801

Problem / Background

bitlinear_matmul is the only mlxcel fused kernel with no graph fallback. src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp holds a Metal port (get_bitlinear_kernel, around line 76) and a CUDA port (get_bitlinear_kernel_cuda, around line 141) and picks between them at mlx_cxx_kernels.cpp:167, with nothing to fall back to. On ROCm neither exists, so a BitNet checkpoint used to reach fast::cuda_kernel, whose no-CUDA stub throws, and the throw crosses the cxx bridge into a noexcept extern and ends the process. The ROCm gate showed this as an abort in models::bitnet::tests::bitlinear_matmul_known_ternary_case (src/models/bitnet.rs:510).

#1803 does not fix that by adding a kernel. It rejects BitNet checkpoints at load on backends without the kernel, with a message naming the backend, so the failure is a clear refusal instead of a process abort. This issue is the follow-up that makes BitNet actually run on ROCm.

Why this is separate from #1814

#1814 ports the fused kernels that matter for decode throughput (sampling, rejection sampling, fused add and RMSNorm, RoPE and append, paged attention, fused MoE, SSM, gated delta), each of which already has a working graph fallback, so it is a performance issue gated by profiling. bitlinear_matmul is a capability gap: without the kernel the model family does not run at all, and there is nothing to profile against. Ordering it by decode impact alongside the others would rank it last, which is the wrong call for the only path that is refused rather than slow.

Scope

Add a fast::hip_kernel port of bitlinear_matmul inside src/lib/mlx-cpp/patches-rocm/, or as an mlxcel-side kernel next to the Metal and CUDA ports in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp, and remove the load-time rejection that #1803 adds once it lands. fast::hip_kernel is already available in the overlay (src/lib/mlx-cpp/patches-rocm/mlx/fast.h:96), the counterpart of fast::metal_kernel and fast::cuda_kernel.

Implementation plan

  1. Port the CUDA kernel body (BITLINEAR_CUDA_SOURCE), whose ternary weight unpacking and scale handling are the closest starting point, keeping the same template arguments (T, in_features, out_features, invert_weight_scales) so the JIT cache keys the same way.
  2. Follow mlxcel's JIT kernel rule that every input dtype appears in the template arguments, which make verify-kernel-dtype-keys checks.
  3. Wire the ROCm branch into the backend selection added by refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 (mlxcel::gpu_kernel_backend(), planned for src/lib/mlx-cpp/turbo/gpu_backend.h) rather than reintroducing a Metal-or-CUDA boolean like the current use_cuda at mlx_cxx_kernels.cpp:166.
  4. Drop the load-time rejection and its message.
  5. Record the change in src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md if the kernel lands in the overlay.

Acceptance Criteria

  • bitlinear_matmul on ROCm matches the CPU stream within float tolerance for the dtypes and shapes the Metal and CUDA ports cover.
  • models::bitnet::tests::bitlinear_matmul_known_ternary_case passes on gfx1151 instead of aborting.
  • A BitNet checkpoint generates coherent text on ROCm with greedy decoding.
  • The load-time rejection added by refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 is removed, and no code path reports BitNet as unsupported on ROCm any more.
  • Metal and CUDA select exactly the same kernel as before.

Verification

make verify-kernel-dtype-keys
cargo test -p mlxcel models::bitnet

Notes

No BitNet checkpoint is present on the gfx1151 development host, so whoever takes this needs to fetch one. The numeric reference should come from the CPU stream rather than from a Metal host, since the two backends are not expected to be byte-identical; see the cross-backend criteria in docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:modelsModel architectures, weights, loading, metadataplatform:linuxLinux (CUDA / packaging) specificpriority:mediumMedium prioritystatus:readyReady to be worked ontype:enhancementNew features, capabilities, or significant additions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions