Repository navigation
Description
Problem
Follow-up to upstream ml-explore/mlx#4738 (53c18705, CUDA gather-QMM int64 strides). The ROCm default gather paths already use int64 expert offsets (quantized/qmm.hip 4759-4766, 5010-5018, 5206-5224, 5675-5690, 5779-5790, 5877-5891, 5998-6001; quantized/qmv_tiled_kernel.hip 436-443), but 32-bit sites and a possible launch limit remain. Paths are under src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/.
32-bit sites
quantized/qmv_tiled_kernel.hip:342-345:gather_qmv_tiled_kernelcomputesrhs_idx * N * w_strideandlhs_idx * M * Kwithuint32_tindices. It is opt-in throughMLX_ROCM_GATHER_QMV_USE_TILED(qmm.hip:6447-6449).- int output index
batch * M * N + row * N + colatqmv_tiled_kernel.hip:394, 501andqmm.hip:5870, 5929(gather_qmv_kernel).qmm.hip:5608, 5714, 5819already use int64 and are the pattern to copy. gemms/gemv.hip:272:out + indices_idx * rowswithint indices_idxandint rows.quantized/qmv_kernel.hiprepeats the pattern but is dead: not inCMakeLists.txtand included nowhere.
Possible launch limit (inferred, not verified)
The gather kernels put batch B = tokens * top_k in grid z (qmm.hip:5170, 5643, 5749, 5844; fast_grid at 6214; tiled gt_grid at 6500). The header of tests/rocm_hip_grid_clamp.rs:18-26 states y and z both top out at 65535 blocks, but only the y cap was observed failing (:48-52). If z has the cap, B > 65535 fails to launch; a 2048-token prefill chunk (src/lib/mlxcel-core/src/prefill_plan.rs:61) at top_k 32 reaches it.
Acceptance Criteria
Implement all of it in one PR with one verification pass at the end.
- Decisive test first in
tests/rocm_hip_grid_clamp.rs: unsorted gather_qmm, bf16, affine 4-bit, group 64, E=128, M=1, N=K=64, B=65,537, against CPU. Record in the PR whether the z cap is real - If it is, every gather QMV kernel indexing batch by
blockIdx.zfolds the batch into a grid capped at 65535 and loops over batch with a stride (not needed: the z cap was ruled out on gfx1151, see fix(rocm): widen gather QMV output offsets to 64 bits #2624) - Listed offsets widened to int64; an
#[ignore]test (f32 activations, N=65,536, K=64, B=32,769, output past 2^31 elements) checks the last row against CPU - Optional: delete
quantized/qmv_kernel.hipand its mention inqdequant.hpp:2 -
LOCAL_FIXES.mdentry (next free number) andmake verify-rocm-overlaygreen
Verification
bash scripts/rocm_gpu_guard.sh --hold -- cargo test --profile test-fast --features rocm --test rocm_hip_grid_clamp -- --test-threads=1 --include-ignored
make verify-rocm-overlayThis edits the same files as #2617 (native NVFP4) and should land before it.
Activity
test(rocm): run the bf16 gather QMV cases in bf16
fix(rocm): widen gather QMV output offsets to 64 bits (lablup#2624)