Use independent version-matched hipBLASLt 2379 plans for batched attention-A at verifier widths 2 through 6. Preserve handle-stream ordering, shape/model guards, failure isolation and the original fallback.
Production-shape full-output, input and canary checks pass with documented rounding differences. The 5/6-row operator is about 2.55x faster; native and objective coding checks retain the prefill gain.
Use zero-workspace F32-output hipBLASLt plans 2537/2539 for the qualified DeepSeek projection shapes. Require matching gfx1151, model, library versions and supported descriptors; retain the original fallback.
Warm 16K native prefill improves 262.64 to 301.30 tokens/s. Operator safety and official scoring pass; long-prompt NLL changes +0.65% at default 4K chunks while the four coding tasks pass 24/24 checks.
Delete the two obsolete kernels that quantized activations separately for every output row. Neither has callers or a testing role; the active prequantized and reference paths remain unchanged. CUDA sm121a and ROCm gfx1151 builds pass without warnings.
Use the model-defined 2051-token dense boundary, handle padded selections, and synchronize reused reduction storage in Metal, CUDA and ROCm, including vision normalization.
Enable bounded ROCm selected-attention GEMM for GLM 5.3 without RoPE. Add independent attention references, boundary tests and finite, repeatable vision checks. Focused Strix continued prefill improves from about 21 to 68 tokens per second; stored official scores are unchanged.
Exclude CMA pages from host availability, guard pinned and integrated-GPU allocations, and reserve space for the OS and required runtime buffers. Defer optional Q8 expansion and avoid wasting large model arenas.
Add memory parsing and allocation-admission tests. Focused Strix checks cover matched Flash/DSpark, resident GLM and automatic, oversized and one-expert SSD cache budgets.
The r9/r11 campaign trees pass ds4_rocm_stream() to these launches; that
helper belongs to the HIP graph-capture work, which is intentionally not
part of this PR. It returns the default stream whenever capture is not
armed, so passing 0 here is behaviorally identical for this branch.
The tile8 gate/up kernel is latency-bound, not format-bound: 57 KiB of
staged-activation LDS pins residency to 1-2 blocks/WGP and each block
has only 8 warps to hide the serial block-loop latency. The down
kernel has the opposite shape - ~917k tiny blocks per layer-chunk at
4096 tokens, each re-staging the same 16.6 KiB of activations.
Three variants, all preserving the exact per-(token,row) accumulation
order (same 8-lane q8_K block split, same dp4a helpers, same
quarter-warp reduction, same epilogue - bit-exact by construction):
- DS4_ROCM_ENABLE_MXFP4_TILE4: 4-token expert tiles for gate/up
(templated tileN<TILE_TOKENS, QWARPS> kernel) - 29 KiB LDS at 7168
dims, 4 blocks/WGP.
- DS4_ROCM_ENABLE_MXFP4_ROW64: 512-thread gate/up blocks (64 rows x
8 tokens against one staged copy) - 2 blocks/WGP = 32 resident warps.
- DS4_ROCM_MXFP4_DOWN_RGROUP=1..8: down kernel loops over consecutive
32-row groups against one staged activation copy - staging traffic
and block count divide by R.
tests/test_mxfp4_rocm now snapshots the default path's out tensor at
n=128/512 and memcmp-asserts each variant against it (cross-path
bitwise equality, stronger than the tolerance checks).
Compile-verified on gfx1151; hardware A/B and the acceptance gate run
in the next test window. The WMMA rewrite is deferred until these
say whether occupancy or format is the binding constraint.
Production kernel trace of the 10.5k-token prefill (rocprofv3 on the
coordinator) showed the MXFP4 routed-MoE tile kernels dominate prefill:
gate/up 247.7 ms and down 174.9 ms per layer-chunk, bound not by
compute (~0.6% of fp32 peak) but by expert-weight DRAM re-reads - every
8-token tile re-reads its experts' full gate+up rows.
Changes (all bit-exact; per-element accumulation order untouched):
- gate/up tile8 kernel grid swapped to (tile, row-block): consecutive
blocks re-walk the same 32-row weight slice across an expert's token
tiles, so the slice stays L2-resident. Production: 247.7 -> 222.3 ms
per layer-chunk (-10.3%), 10.5k prefill 61.5 s -> 59.9 s.
- A-activation LDS staging threshold raised 16 -> 28 q8_K blocks (the
7168-dim case; 56.9 KiB dynamic smem with the >48 KiB opt-in).
Neutral in production (L2 was already absorbing the re-reads), large
in the isolated bench; staged and unstaged dot helpers are bit-equal
(integer bsum order is exact; verified in-tree and at 7168 dims).
- down-kernel grid swap measured NEGATIVE (174.9 -> 184.8 ms) and was
reverted.
- Two experimental kernels parked default-off behind env gates with
honest numbers: tile32 B-register reuse (DS4_ROCM_ENABLE_MXFP4_TILE32;
register spills, 2155 ms bench) and row8-ldsB B-staging
(DS4_ROCM_ENABLE_MXFP4_LDSB; 154 s production - scattered A reads).
- tests/bench_mxfp4_rocm: production-shape MoE bench harness. NOTE:
it is a poor proxy for production kernel time (isolated process
memory/L2 behavior differs ~5x); production traces are the arbiter.
- DS4_ROCM_MOE_PATH_DEBUG: launch-path diagnostics.
Validation: tests/test_mxfp4_rocm PASS (default dims + 7168/top-8
variant), acceptance gate df07199e5a292872 holds, decode 15.38 t/s,
10.5k probe 59.9 s (176 t/s avg).
The MXFP4 tile kernels ran at roughly three percent of dp4a peak and
were insensitive to occupancy, scratch, DRAM weight traffic, and chain
depth: the limiter was instruction issue, with four dword activation
loads and loop control paired against every eight dp4a.
Stage activations as 16-byte-aligned quant slices with the Q8_K scales
in a separate array, so the dot helper loads one b128 per half
sub-block instead of four dwords; the layout also drops the unused
32-byte bsums from every staged chunk. Weight bytes load as one
misaligned b128 per block, the pair loop is a fixed predicated eight,
and the gate/up kernel computes both matrices from a single activation
load through a fused helper, halving its dominant LDS traffic. Low and
high nibble sums use separate integer accumulators; all integer
reassociation is exact and the outputs stay bit-identical, with the
struct-layout helper retained for the unstaged wide-activation path.
gfx1151, DS4 Flash shape (4096/2048, 256 experts, 6 used), rocprofv3
min-of-3 on the synthetic sweep: gate/up tile kernel 106.0ms -> 50.1ms
(2.11x), down tile kernel 69.6ms -> 34.0ms (2.05x). Distributed pair
(coordinator 0:21 + worker 22:output, dist chunk 4096): prefill 70.4 ->
120.6 tok/s serial single-chunk (1.71x) and 103.9 -> 176.5 tok/s on a
three-chunk pipelined frontier (1.70x); decode unchanged at 7.7 tok/s.
The MXFP4 expert-tile kernels built their per-tile pair state with
runtime-count loops, so pair/tok/slot/xqb were dynamically indexed and
spilled: the compiler reported 48 bytes/lane of scratch on the gate/up
kernel, pairing private-memory traffic with the inner dp4a loop. Setup
and finalize now run a fixed, fully unrolled 8 iterations with a uniform
p < np guard, which lets every array scalarize. tok/slot are gone
entirely since pair == tok * n_expert + slot already indexes the route
weights.
The gate/up kernel also staged activations into a static [8][16] Q8_K
tile even when xq_blocks is smaller; the buffer is now dynamic shared
memory sized by the launch.
gfx1151 compiler-reported resources for the gate/up tile kernel go from
99 VGPRs / 48 B scratch / 12 waves to 83 VGPRs / no scratch / 16 waves.
On the DS4 Flash shape (xq_blocks=16) the LDS footprint is unchanged and
the kernel stays weight-stream bound, so end-to-end prefill is neutral
there; the change is codegen hygiene plus wins for smaller shapes.
Numerics are identical since the summation order is preserved.
Coalesce MXFP4 weight reads across wave32 lanes and unpack nibbles with register permutations. Keep the gate/up and down reductions register-resident, route two-to-four-token batches through the direct kernels, and extend the ROCm CPU-oracle coverage across that dispatch boundary.
Implement resident MXFP4 gate/up and down paths for decode and batched routed MoE on ROCm. Add a synthetic CPU-oracle regression covering full, partial, and large expert tiles, report MXFP4 as a 4-bit routed quant, and make Strix host objects position independent for TheRock linking.
Integrate the MXFP4 Metal decode/prefill optimizations (rounds 7-9) as one
commit. Keep both the MXFP4 half_lut/tail_cull down-pipeline selection and
main's mpp_mask override for IQ2_XXS gate/up and Q2_K/IQ2_XXS down; they
target disjoint quant types and coexist.
Donato Capitella
Salvatore Sanfilippo
Donato Capitella
Jonathan Yates
ivanfioravanti