[Gajesh2007]

GPT-OSS 20B gathered MXFP4 decode uses width 2880, which misses the existing fast vector kernel's 512-element alignment. Add a masked-tail variant: five full blocks followed by 320 values, read only by the 20 participating lanes. Preserve FP32 accumulation and the existing gathered indices/strides.

Enable this route automatically only on physical M4 Max architecture applegpu_g16s, with exact MXFP4/group32/4-bit/E32/K2880/N2880-or5760 and FP32/BF16 gates. MLX_GPTOSS_MXFP4_DECODE_FAST_TAIL=0 restores the original route. Unsupported calls retain the general path. A separately measured 32x32x32 sorted prefill tile remains opt-in through MLX_GPTOSS_MXFP4_PREFILL_TILE=m32n32k32.

flowchart LR
  subgraph Before
    A[GPT-OSS decode: K2880] --> B[gather_qmv alignment check]
    B --> C[General gathered MXFP4 kernel]
    C --> D[Next-token logits]
  end
  subgraph After
    E[GPT-OSS decode: K2880] --> F[gather_qmv shape and physical-device gates]
    F -->|Measured shape/device| G[gather_qmv_fast_tail: 2560 + 320]
    F -->|Fallback or rollback| H[Existing kernel]
    G --> I[Next-token logits]
    H --> I
  end

Loading

Validation: actual Swift GPU tests cover ragged expert boundaries, partial rows, varying scale exponents, FP32/BF16, B=1/2/4/8, sparse tail-edge values, and a bit-exact unsupported-K fallback. The C++ doctest additions were not separately built; the production C++ bridge and generated Metal source were built and exercised by the Swift suite. Two ABBA cycles measured roughly 2–3% more short-context aggregate decode throughput at B=1/2/4, with exact generated sequences. Floating-point reduction order changes, so bit identity is not promised for every input. The prefill option is not a universal default.

The dependent MLX-Swift PR includes the matching generated fp_quantized.cpp and source pin. No production deployment.