Repository navigation
feat(rocm): port the BitNet bitlinear_matmul kernel - #1870
Merged
Merged
Conversation
9 of 10 tasks
inureyes
force-pushed
the
feature/issue-1803-kernel-backend-kind
branch
from
September 13, 2026 13:20
dda6a7e to
314ecaa
Compare
`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
inureyes
force-pushed
the
feature/issue-1862-rocm-bitlinear-kernel
branch
from
September 13, 2026 13:23
bb4a34b to
7d03956
Compare
inureyes
changed the base branch from
feature/issue-1803-kernel-backend-kind
to
main
September 13, 2026 13:23
This was referenced Sep 29, 2026
5 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
bitlinear_matmulis 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_syncexists 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__, notstatic_assert(warpSize == 32), because HIP'swarpSizeis an object with anoperator int()and that form does not compile.The backend selection gains its third arm, and
bitlinear_kernel_available()joinscustom_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_laneis new because the existing test could not have caught a wrong fold. It runsin_features = 4, and the kernel has laneLstart atL * 4and 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
mlxceltest 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_availablepredicate 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.