Skip to content

fix(rocm): gather QMV large-batch launch and 64-bit output offsets #2620

Closed

Description

@inureyes

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_kernel computes rhs_idx * N * w_stride and lhs_idx * M * K with uint32_t indices. It is opt-in through MLX_ROCM_GATHER_QMV_USE_TILED (qmm.hip:6447-6449).
  • int output index batch * M * N + row * N + col at qmv_tiled_kernel.hip:394, 501 and qmm.hip:5870, 5929 (gather_qmv_kernel). qmm.hip:5608, 5714, 5819 already use int64 and are the pattern to copy.
  • gemms/gemv.hip:272: out + indices_idx * rows with int indices_idx and int rows.
  • quantized/qmv_kernel.hip repeats the pattern but is dead: not in CMakeLists.txt and 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.z folds 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.hip and its mention in qdequant.hpp:2
  • LOCAL_FIXES.md entry (next free number) and make verify-rocm-overlay green

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-overlay

This edits the same files as #2617 (native NVFP4) and should land before it.

Activity

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    platform:linuxLinux (CUDA / packaging) specific
    and removed
    platform:linuxLinux (CUDA / packaging) specific
    on Oct 26, 2026
  2. added a commit that references this issue on Nov 3, 2026
    2a1f1e8
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

    priority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions