You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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.
#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
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.
Follow mlxcel's JIT kernel rule that every input dtype appears in the template arguments, which make verify-kernel-dtype-keys checks.
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.
Part of #1801
Problem / Background
bitlinear_matmulis the only mlxcel fused kernel with no graph fallback.src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cppholds a Metal port (get_bitlinear_kernel, around line 76) and a CUDA port (get_bitlinear_kernel_cuda, around line 141) and picks between them atmlx_cxx_kernels.cpp:167, with nothing to fall back to. On ROCm neither exists, so a BitNet checkpoint used to reachfast::cuda_kernel, whose no-CUDA stub throws, and the throw crosses the cxx bridge into anoexceptextern and ends the process. The ROCm gate showed this as an abort inmodels::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_matmulis 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_kernelport ofbitlinear_matmulinsidesrc/lib/mlx-cpp/patches-rocm/, or as an mlxcel-side kernel next to the Metal and CUDA ports insrc/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp, and remove the load-time rejection that #1803 adds once it lands.fast::hip_kernelis already available in the overlay (src/lib/mlx-cpp/patches-rocm/mlx/fast.h:96), the counterpart offast::metal_kernelandfast::cuda_kernel.Implementation plan
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.make verify-kernel-dtype-keyschecks.mlxcel::gpu_kernel_backend(), planned forsrc/lib/mlx-cpp/turbo/gpu_backend.h) rather than reintroducing a Metal-or-CUDA boolean like the currentuse_cudaatmlx_cxx_kernels.cpp:166.src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.mdif the kernel lands in the overlay.Acceptance Criteria
bitlinear_matmulon 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_casepasses ongfx1151instead of aborting.Verification
Notes
No BitNet checkpoint is present on the
gfx1151development 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 indocs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md.