Skip to content

feat(rocm): port the BitNet bitlinear_matmul kernel - #1870

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-1862-rocm-bitlinear-kernel
Sep 13, 2026
Merged

inureyes merged 1 commit into
mainfrom
feature/issue-1862-rocm-bitlinear-kernel

Conversation

@inureyes

Copy link
Copy Markdown
Member

bitlinear_matmul is the one mlxcel fused kernel with no graph fallback, so until #1803 a BitNet checkpoint ended the process on any backend without a port. #1803 turned that into a refusal at load. This gives ROCm the kernel instead, so BitNet runs there.

The body is the CUDA port 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 and the width is stated rather than implied. And the fold starts at 16, which is right only for a 32-lane wavefront: on a wave64 target (CDNA, gfx90a and gfx942) it would fold half the lanes and the result would stay finite and plausible while being wrong, so the source refuses to compile there. The check is a preprocessor test on __AMDGCN_WAVEFRONT_SIZE__, not static_assert(warpSize == 32), because HIP's warpSize is an object with an operator int() and that form does not compile.

The backend selection gains its third arm, and bitlinear_kernel_available() joins custom_kernels_available() rather than replacing it. They answered the same question while Metal and CUDA were the only backends with ports; they stop being the same the moment one backend has this kernel and not the rest, which is now.

bitlinear_matmul_folds_every_lane is new because the existing test could not have caught a wrong fold. It runs in_features = 4, and the kernel has lane L start at L * 4 and stride by 128, so only lane 0 enters the loop and the fold adds one real value to 31 zeros. Any fold width passes it. That gap was never ROCm-specific: the CUDA arm's reduction has never been covered either. The new case uses 256 inputs so every lane runs twice, and its expectation is computed from the packed weights rather than written down, because a hand-computed reference over 256 inputs is easier to get wrong than the kernel and a wrong expectation that passes is worse than no test.

Verified on gfx1151: both BitNet tests pass, where before this branch the known-ternary case aborted the whole mlxcel test binary.

Closes #1862

Stacked on #1869 (issue #1803), which adds the backend routing this selects through. Review that one first; the diff here is the kernel, its selection arm, the bitlinear_kernel_available predicate and the new test.

Validated on gfx1151: cargo fmt --all -- --check, scripts/ci/check_kernel_dtype_keys.py (this launcher keys on {"T", T} and stays in scope), and the two BitNet tests.

@inureyes inureyes added priority:high High priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:linux Linux (CUDA / packaging) specific status:review Under review type:enhancement New features, capabilities, or significant additions labels Sep 12, 2026
@inureyes
inureyes force-pushed the feature/issue-1803-kernel-backend-kind branch from dda6a7e to 314ecaa Compare September 13, 2026 13:20
`bitlinear_matmul` is the one mlxcel fused kernel with no graph fallback, so until #1803 a BitNet checkpoint ended the process on any backend without a port. #1803 turned that into a refusal at load. This gives ROCm the kernel instead, so BitNet runs there.

The body is the CUDA port 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 and the width is stated rather than implied. And the fold starts at 16, which is right only for a 32-lane wavefront: on a wave64 target (CDNA, gfx90a and gfx942) it would fold half the lanes and the result would stay finite and plausible while being wrong, so the source refuses to compile there. The check is a preprocessor test on `__AMDGCN_WAVEFRONT_SIZE__`, not `static_assert(warpSize == 32)`, because HIP's `warpSize` is an object with an `operator int()` and that form does not compile.

The backend selection gains its third arm, and `bitlinear_kernel_available()` joins `custom_kernels_available()` rather than replacing it. They answered the same question while Metal and CUDA were the only backends with ports; they stop being the same the moment one backend has this kernel and not the rest, which is now.

`bitlinear_matmul_folds_every_lane` is new because the existing test could not have caught a wrong fold. It runs `in_features = 4`, and the kernel has lane `L` start at `L * 4` and stride by 128, so only lane 0 enters the loop and the fold adds one real value to 31 zeros. Any fold width passes it. That gap was never ROCm-specific: the CUDA arm's reduction has never been covered either. The new case uses 256 inputs so every lane runs twice, and its expectation is computed from the packed weights rather than written down, because a hand-computed reference over 256 inputs is easier to get wrong than the kernel and a wrong expectation that passes is worse than no test.

Verified on gfx1151: both BitNet tests pass, where before this branch the known-ternary case aborted the whole `mlxcel` test binary.

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

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:linux Linux (CUDA / packaging) specific priority:high High priority status:review Under review type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant