diff --git a/Makefile b/Makefile index 3cc7ccef4..028d03457 100644 --- a/Makefile +++ b/Makefile @@ -40,6 +40,7 @@ CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o CUDA_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas HIPCC ?= $(shell command -v hipcc 2>/dev/null || echo /opt/rocm/bin/hipcc) ROCM_ARCH ?= gfx1151 +ROCM_HOST_CFLAGS ?= -fPIC ROCM_CFLAGS ?= -O3 -ffast-math -g -fno-finite-math-only -pthread -D__HIP_PLATFORM_AMD__ -Wno-unused-command-line-argument --offload-arch=$(ROCM_ARCH) ROCM_LDLIBS ?= -lm -pthread -lhipblas -lhipblaslt DS4_LINK ?= $(NVCC) $(NVCCFLAGS) @@ -47,7 +48,7 @@ DS4_LINK_LIBS ?= $(CUDA_LDLIBS) METAL_LDLIBS := $(LDLIBS) endif -.PHONY: all help clean test test-metal-session-batch test-cuda-session-batch test-cuda-mixed-batch dspark-acceptance dspark-verify-depth mtp-verify-depth cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm +.PHONY: all help clean test test-metal-session-batch test-mxfp4-rocm test-cuda-session-batch test-cuda-mixed-batch dspark-acceptance dspark-verify-depth mtp-verify-depth cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm ifeq ($(UNAME_S),Darwin) all: ds4 ds4-server ds4-bench ds4-eval ds4-agent @@ -116,6 +117,7 @@ help: @echo " make cuda CUDA_ARCH=sm_N Build CUDA with an explicit nvcc -arch value" @echo " make strix-halo Build ROCm for Strix Halo / gfx1151" @echo " make rocm Alias for make strix-halo" + @echo " make test-mxfp4-rocm Build and run the synthetic ROCm MXFP4 MoE test" @echo " make cpu Build CPU-only ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent" @echo " make test Build and run tests" @echo " make dspark-verify-depth Run DSpark speculative verification smoke if support GGUF is present" @@ -139,7 +141,7 @@ cuda: strix-halo: $(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent \ CORE_OBJS="ds4.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_rocm.o ds4_rocm_compat.o ds4_rocm_unavailable.o ds4_layer_pack.o" \ - CFLAGS="$(CFLAGS) -DDS4_ROCM_BUILD" \ + CFLAGS="$(CFLAGS) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD" \ DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \ DS4_LINK_LIBS="$(ROCM_LDLIBS)" @@ -258,6 +260,15 @@ ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_iq2_tables_cuda.inc ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS) $(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm.cu +tests/test_mxfp4_rocm.o: tests/test_mxfp4_rocm.c ds4_gpu.h + $(CC) $(filter-out -ffast-math,$(CFLAGS)) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD -I. -c -o $@ $< + +tests/test_mxfp4_rocm: tests/test_mxfp4_rocm.o ds4_rocm.o + $(HIPCC) $(ROCM_CFLAGS) -o $@ $^ $(ROCM_LDLIBS) + +test-mxfp4-rocm: tests/test_mxfp4_rocm + ./tests/test_mxfp4_rocm + ds4_rocm_compat.o: ds4_rocm_compat.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_gpu_args.h $(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm_compat.cu @@ -415,4 +426,4 @@ mxfp4-dot-test: tests/test_mxfp4_dot.c ./tests/test_mxfp4_dot clean: - rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_metal_session_batch tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o + rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_metal_session_batch tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o diff --git a/QA_BEFORE_RELEASES.md b/QA_BEFORE_RELEASES.md index b2375ba2b..49bf7adee 100644 --- a/QA_BEFORE_RELEASES.md +++ b/QA_BEFORE_RELEASES.md @@ -349,6 +349,9 @@ a substitute for CUDA or Metal release testing. - Build: `make clean && make strix-halo`. - Require the ROCm build to complete without compiler warnings. +- After MXFP4 or ROCm routed-MoE changes, run `make test-mxfp4-rocm`. Require + zero `failures` for both `mid` and `out` at 1, 3, 32, 128, and 512 tokens, + followed by `MXFP4 ROCm routed MoE: PASS`. - Use the q2 Flash imatrix GGUF for release smoke tests: `DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf`. - Do not use the mixed q2-q4 or Q4 Flash GGUFs for routine Strix Halo QA yet. diff --git a/README.md b/README.md index 032224bae..371ada019 100644 --- a/README.md +++ b/README.md @@ -393,6 +393,12 @@ To build an initial mental model, here are the high level concepts: 4. Each worker keeps its slice of the KV cache. 5. Communication is worker-to-worker, there is no need to use the coordinator as relay, so if your coordinator is `A`, and you make a request, activations will flow in `A -> B -> C -> back to A`. +The resident ROCm MXFP4 routed-expert path supports the same pipeline mode. A +tested two-host Strix Halo split uses `--layers 0:21` on the coordinator and +`--layers 22:output` on the worker. This is a capacity configuration for a +model that does not fit on one 128 GB system; it does not add ROCm SSD +streaming support for Flash. + ### How it works and how to configure it The prefill path is pipelined (this is why it can go faster than in a single machine). diff --git a/STRIXHALO.md b/STRIXHALO.md index 481232bdf..229a769c4 100644 --- a/STRIXHALO.md +++ b/STRIXHALO.md @@ -113,6 +113,16 @@ make strix-halo -j"$(nproc)" `make rocm` is an alias for `make strix-halo`. +After changes to the ROCm MXFP4 or routed-MoE kernels, run the standalone +CPU-oracle regression: + +```sh +make test-mxfp4-rocm +``` + +This test does not require a full model GGUF. It covers resident decode and +batched routed-MoE execution at 1, 3, 32, 128, and 512 tokens. + ## 5. Use the right GGUF Use the standard IQ2XXS/Q2K/Q8 imatrix GGUF: diff --git a/ds4.c b/ds4.c index 66621ff89..e61498640 100644 --- a/ds4.c +++ b/ds4.c @@ -49606,7 +49606,8 @@ int ds4_engine_routed_quant_bits(ds4_engine *e) { for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_tensor *gate = e->weights.layer[il].ffn_gate_exps; if (!gate) continue; - return gate->type == DS4_TENSOR_Q4_K ? 4 : 2; + return (gate->type == DS4_TENSOR_Q4_K || + gate->type == DS4_TENSOR_MXFP4) ? 4 : 2; } return 0; } diff --git a/ds4_rocm.cu b/ds4_rocm.cu index 46ba223df..f1f754735 100644 --- a/ds4_rocm.cu +++ b/ds4_rocm.cu @@ -87,6 +87,20 @@ typedef struct { uint16_t qs[CUDA_QK_K / 8]; } cuda_block_iq2_xxs; +typedef struct { + uint8_t e; + uint8_t qs[16]; +} cuda_block_mxfp4; + +static_assert(sizeof(cuda_block_mxfp4) == 17, "cuda_block_mxfp4 must match the GGUF MXFP4 block layout"); + +/* Twice the MXFP4 values so each 32-value sub-block can use signed-int8 + * dp4a; the factor of 1/2 is folded into the sub-block scale. */ +__device__ __constant__ static const int8_t cuda_mxfp4_values_x2[16] = { + 0, 1, 2, 3, 4, 6, 8, 12, + 0, -1, -2, -3, -4, -6, -8, -12, +}; + #include "ds4_iq2_tables_cuda.inc" #include "rocm/ds4_rocm_runtime.cuh" diff --git a/rocm/ds4_rocm_moe.cuh b/rocm/ds4_rocm_moe.cuh index 684f86729..ebd56077e 100644 --- a/rocm/ds4_rocm_moe.cuh +++ b/rocm/ds4_rocm_moe.cuh @@ -288,6 +288,135 @@ __device__ static float dev_dot_q4_K_q8_K_block(const cuda_block_q4_K *x, const return y->d * xd * (float)isum - y->d * xmin * (float)summs; } +__device__ __forceinline__ static float dev_e8m0_to_f32(uint8_t e) { + const uint32_t bits = e == 0u ? 0x00400000u : (uint32_t)e << 23u; + return __uint_as_float(bits); +} + +/* Expand four MXFP4 codes to signed, doubled int8 values in one register. + * On AMDGPU, V_PERM_B32 treats selector bytes 0..3 as bytes from src1, + * 4..7 as bytes from src0, 12 as zero, and 13 as 0xff. Keeping both + * halves of the tiny value table in SGPR-immediate operands avoids four + * divergent global LUT reads for every packed word. */ +__device__ __forceinline__ static uint32_t dev_mxfp4_unpack4(uint32_t codes) { +#if defined(__AMDGCN__) + const uint32_t selectors = codes & 0x07070707u; + const uint32_t positive = __builtin_amdgcn_perm( + 0x0c080604u, 0x03020100u, selectors); + const uint32_t negative = __builtin_amdgcn_perm( + 0xf4f8fafcu, 0xfdfeff00u, selectors); + const uint32_t sign_selectors = + 0x0c0c0c0cu + ((codes & 0x08080808u) >> 3u); + const uint32_t sign_mask = __builtin_amdgcn_perm( + 0u, 0u, sign_selectors); + return (positive & ~sign_mask) | (negative & sign_mask); +#else + uint32_t result = 0u; + #pragma unroll + for (uint32_t i = 0; i < 4u; i++) { + const uint32_t code = (codes >> (8u * i)) & 0x0fu; + const uint32_t base = code & 7u; + int32_t value = (int32_t)(base + (base > 4u ? base - 4u : 0u) + + (base == 7u ? 2u : 0u)); + if ((code & 8u) != 0u) value = -value; + result |= (uint32_t)(uint8_t)value << (8u * i); + } + return result; +#endif +} + +__device__ __forceinline__ static void dev_mxfp4_unpack2x4( + const uint8_t *q, + int32_t *low, + int32_t *high) { + const uint32_t packed = *(const uint32_t *)q; + *low = (int32_t)dev_mxfp4_unpack4(packed); + *high = (int32_t)dev_mxfp4_unpack4(packed >> 4u); +} + +/* One q8_K chunk covers eight consecutive 32-value MXFP4 blocks. MXFP4 + * stores the first 16 values in the low nibbles and the second 16 in the + * high nibbles, rather than interleaving them. */ +__device__ static float dev_dot_mxfp4_q8_K_block( + const cuda_block_mxfp4 *x8, + const cuda_block_q8_K *y) { + float chunk = 0.0f; + #pragma unroll + for (uint32_t sb = 0; sb < 8u; sb++) { + const cuda_block_mxfp4 *x = x8 + sb; + const int8_t *q8 = y->qs + sb * 32u; + int32_t bsum = 0; + #pragma unroll + for (uint32_t j = 0; j < 16u; j += 4u) { + int32_t wlo, whi; + dev_mxfp4_unpack2x4(x->qs + j, &wlo, &whi); + bsum = __dp4a(wlo, *(const int32_t *)(q8 + j), bsum); + bsum = __dp4a(whi, *(const int32_t *)(q8 + 16u + j), bsum); + } + chunk += dev_e8m0_to_f32(x->e) * (float)bsum; + } + return 0.5f * y->d * chunk; +} + +/* Split one 32-value MXFP4 block across a pair of lanes. A wave therefore + * reads 16 consecutive 17-byte blocks instead of having each quarter-wave + * lane jump by a full 136-byte Q8_K chunk. This mirrors the coalesced Metal + * decode layout while retaining the faster Q8_K activation path on gfx1151. */ +__device__ __forceinline__ static float dev_dot_mxfp4_q8_K_half_block( + const cuda_block_mxfp4 *x, + const cuda_block_q8_K *y, + uint32_t subblock, + uint32_t half) { + const uint32_t weight_offset = half * 8u; + const uint32_t activation_offset = subblock * 32u + weight_offset; + const int8_t *q8_lo = y->qs + activation_offset; + const int8_t *q8_hi = y->qs + activation_offset + 16u; + const uint8_t *q = x->qs + weight_offset; + int32_t bsum = 0; + #pragma unroll + for (uint32_t j = 0; j < 8u; j += 4u) { + int32_t wlo, whi; + dev_mxfp4_unpack2x4(q + j, &wlo, &whi); + bsum = __dp4a(wlo, *(const int32_t *)(q8_lo + j), bsum); + bsum = __dp4a(whi, *(const int32_t *)(q8_hi + j), bsum); + } + return 0.5f * y->d * dev_e8m0_to_f32(x->e) * (float)bsum; +} + +__device__ static void dev_dot_mxfp4_q8_K_block8( + const cuda_block_mxfp4 *x8, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + const cuda_block_q8_K *y4, + const cuda_block_q8_K *y5, + const cuda_block_q8_K *y6, + const cuda_block_q8_K *y7, + uint32_t n, + float acc[8]) { + const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; + float chunk[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + #pragma unroll + for (uint32_t sb = 0; sb < 8u; sb++) { + const cuda_block_mxfp4 *x = x8 + sb; + int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + #pragma unroll + for (uint32_t j = 0; j < 16u; j += 4u) { + int32_t wlo, whi; + dev_mxfp4_unpack2x4(x->qs + j, &wlo, &whi); + for (uint32_t p = 0; p < n; p++) { + const int8_t *q8 = ys[p]->qs + sb * 32u; + bsum[p] = __dp4a(wlo, *(const int32_t *)(q8 + j), bsum[p]); + bsum[p] = __dp4a(whi, *(const int32_t *)(q8 + 16u + j), bsum[p]); + } + } + const float d = dev_e8m0_to_f32(x->e); + for (uint32_t p = 0; p < n; p++) chunk[p] += d * (float)bsum[p]; + } + for (uint32_t p = 0; p < n; p++) acc[p] += 0.5f * ys[p]->d * chunk[p]; +} + __device__ static void dev_dot_q4_K_q8_K_block4( const cuda_block_q4_K *x, const cuda_block_q8_K *y0, @@ -1838,6 +1967,96 @@ __global__ static void moe_gate_up_mid_q4K_expert_tile8_row32_kernel( } } +__global__ static void moe_gate_up_mid_mxfp4_expert_tile8_row32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t max_count, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + uint32_t count = counts[expert]; + if (max_count != 0u && count >= max_count) return; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][16]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= count) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + tok[np] = pair[np] / n_expert; + slot[np] = pair[np] - tok[np] * n_expert; + xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; + } + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + uint32_t p = i / xq_blocks; + uint32_t b = i - p * xq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= expert_mid_dim) return; + const char *gate_row = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes; + const char *up_row = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes; + const uint64_t gate_chunk_bytes = gate_row_bytes / xq_blocks; + float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + const cuda_block_mxfp4 *gb = (const cuda_block_mxfp4 *)(gate_row + (uint64_t)b * gate_chunk_bytes); + const cuda_block_mxfp4 *ub = (const cuda_block_mxfp4 *)(up_row + (uint64_t)b * gate_chunk_bytes); + dev_dot_mxfp4_q8_K_block8(gb, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate); + dev_dot_mxfp4_q8_K_block8(ub, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up); + } + for (uint32_t p = 0; p < np; p++) { + gate[p] = quarter_warp_sum_f32(gate[p], lane); + up[p] = quarter_warp_sum_f32(up[p], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[p] > clamp) gate[p] = clamp; + if (up[p] > clamp) up[p] = clamp; + if (up[p] < -clamp) up[p] = -clamp; + } + const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[p]; + up_out[off] = up[p]; + } + mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; + } + } +} + __global__ static DS4_ROCM_UNUSED void moe_down_kernel( float *down_out, const char *down_base, @@ -2002,6 +2221,88 @@ __global__ static void moe_gate_up_mid_decode_q4K_qwarp32_kernel( } } +__global__ static void moe_gate_up_mid_decode_mxfp4_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + constexpr uint32_t rows_per_wave = 1u; + constexpr uint32_t waves_per_block = 8u; + constexpr uint32_t rows_per_block = rows_per_wave * waves_per_block; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t wave = threadIdx.x >> 5u; + const uint32_t block_lane = lane >> 1u; + const uint32_t half = lane & 1u; + const uint32_t first_row = blockIdx.x * rows_per_block + wave * rows_per_wave; + const uint32_t pair = blockIdx.y; + const uint32_t tok = pair / n_expert; + const uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + float gate[rows_per_wave] = {0.0f}; + float up[rows_per_wave] = {0.0f}; + const uint32_t mxfp4_blocks = xq_blocks * 8u; + + for (uint32_t mb = block_lane; mb < mxfp4_blocks; mb += 16u) { + const cuda_block_q8_K *yb = xqb + (mb >> 3u); + #pragma unroll + for (uint32_t rr = 0; rr < rows_per_wave; rr++) { + const uint32_t row = first_row + rr; + if (row >= expert_mid_dim) continue; + const cuda_block_mxfp4 *gate_blocks = + (const cuda_block_mxfp4 *)(gate_base + + (uint64_t)expert * gate_expert_bytes + + (uint64_t)row * gate_row_bytes); + const cuda_block_mxfp4 *up_blocks = + (const cuda_block_mxfp4 *)(up_base + + (uint64_t)expert * gate_expert_bytes + + (uint64_t)row * gate_row_bytes); + const uint32_t subblock = mb & 7u; + gate[rr] += dev_dot_mxfp4_q8_K_half_block( + gate_blocks + mb, yb, subblock, half); + up[rr] += dev_dot_mxfp4_q8_K_half_block( + up_blocks + mb, yb, subblock, half); + } + } + + #pragma unroll + for (uint32_t rr = 0; rr < rows_per_wave; rr++) { + const uint32_t row = first_row + rr; + if (row >= expert_mid_dim) continue; + gate[rr] = warp_sum_f32(gate[rr]); + up[rr] = warp_sum_f32(up[rr]); + if (lane == 0u) { + float g = gate[rr]; + float u = up[rr]; + if (clamp > 1.0e-6f) { + if (g > clamp) g = clamp; + if (u > clamp) u = clamp; + if (u < -clamp) u = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = g; + up_out[off] = u; + } + mid_out[off] = (g / (1.0f + expf(-g))) * u * + weights[(uint64_t)tok * n_expert + slot]; + } + } +} + __global__ static void moe_gate_up_mid_q2K_decode_q8_qwarp32_kernel( float *gate_out, float *up_out, @@ -2257,6 +2558,64 @@ __global__ static void moe_down_q4K_sum6_qwarp32_kernel( if (lane == 0) out[row] = total; } +template +__global__ static void moe_down_mxfp4_sum6_qwarp32_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + constexpr uint32_t rows_per_wave = 1u; + constexpr uint32_t waves_per_block = 8u; + constexpr uint32_t rows_per_block = rows_per_wave * waves_per_block; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t wave = threadIdx.x >> 5u; + const uint32_t block_lane = lane >> 1u; + const uint32_t half = lane & 1u; + const uint32_t first_row = blockIdx.x * rows_per_block + wave * rows_per_wave; + const uint32_t tok = Batch ? blockIdx.y : 0u; + const int32_t *token_selected = selected + (uint64_t)tok * n_expert; + const cuda_block_q8_K *token_midq = + midq + (uint64_t)tok * n_expert * midq_blocks; + float *token_out = out + (uint64_t)tok * out_dim; + float total[rows_per_wave] = {0.0f}; + const uint32_t mxfp4_blocks = midq_blocks * 8u; + + #pragma unroll + for (uint32_t slot = 0; slot < DS4_ROCM_N_EXPERT_USED; slot++) { + if (slot >= n_expert) continue; + int32_t expert_i = token_selected[slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q8_K *xq = token_midq + (uint64_t)slot * midq_blocks; + for (uint32_t mb = block_lane; mb < mxfp4_blocks; mb += 16u) { + const cuda_block_q8_K *yb = xq + (mb >> 3u); + #pragma unroll + for (uint32_t rr = 0; rr < rows_per_wave; rr++) { + const uint32_t row = first_row + rr; + if (row >= out_dim) continue; + const cuda_block_mxfp4 *down_blocks = + (const cuda_block_mxfp4 *)(down_base + + (uint64_t)(uint32_t)expert_i * down_expert_bytes + + (uint64_t)row * down_row_bytes); + total[rr] += dev_dot_mxfp4_q8_K_half_block( + down_blocks + mb, yb, mb & 7u, half); + } + } + } + + #pragma unroll + for (uint32_t rr = 0; rr < rows_per_wave; rr++) { + const uint32_t row = first_row + rr; + if (row >= out_dim) continue; + total[rr] = warp_sum_f32(total[rr]); + if (lane == 0u) token_out[row] = total[rr]; + } +} + __global__ static void moe_down_q4K_qwarp32_kernel( float *down_out, const char *down_base, @@ -2434,6 +2793,64 @@ __global__ static void moe_down_q4K_expert_tile8_row32_kernel( } } +__global__ static void moe_down_mxfp4_expert_tile8_row32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][8]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= out_dim) return; + const char *down_row = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes; + const uint64_t down_chunk_bytes = down_row_bytes / midq_blocks; + float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + const cuda_block_mxfp4 *wb = (const cuda_block_mxfp4 *)(down_row + (uint64_t)b * down_chunk_bytes); + dev_dot_mxfp4_q8_K_block8(wb, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc); + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) { + down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } +} + __global__ static void moe_down_sorted_qwarp32_kernel( float *down_out, const char *down_base, diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index 51fafe3d3..e7c525b73 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -443,6 +443,7 @@ typedef struct { int iq2_path; int iq2_iq2_path; int q2k_path; + int mxfp4_path; uint64_t gate_bytes; uint64_t down_bytes; } routed_moe_launch_plan; @@ -493,8 +494,10 @@ static int routed_moe_build_plan( plan->iq2_path = (gate_type == 16u && down_type == 10u); plan->iq2_iq2_path = (gate_type == 16u && down_type == 16u); plan->q2k_path = (gate_type == 10u && down_type == 10u); + plan->mxfp4_path = (gate_type == 39u && down_type == 39u); if (!plan->q4k_path && !plan->iq2_path && - !plan->iq2_iq2_path && !plan->q2k_path) return 0; + !plan->iq2_iq2_path && !plan->q2k_path && + !plan->mxfp4_path) return 0; if (!cuda_u64_mul_checked(n_total_expert, gate_expert_bytes, &plan->gate_bytes) || !cuda_u64_mul_checked(n_total_expert, down_expert_bytes, &plan->down_bytes) || !cuda_model_range_fits(model_size, gate_offset, plan->gate_bytes) || @@ -559,6 +562,7 @@ static int routed_moe_launch( const int iq2_iq2_path = plan.iq2_iq2_path; const int iq2_gate_path = iq2_path || iq2_iq2_path; const int q2k_path = plan.q2k_path; + const int mxfp4_path = plan.mxfp4_path; const uint64_t gate_bytes = plan.gate_bytes; const uint64_t down_bytes = plan.down_bytes; uint64_t pair_count64 = 0; @@ -745,24 +749,32 @@ static int routed_moe_launch( /* Correctness rollback for the optimized resident IQ2 prefill path. */ const uint32_t disable_resident_iq2_sorted = iq2_gate_path && getenv("DS4_ROCM_DISABLE_RESIDENT_IQ2_SORTED") != NULL; + const uint32_t use_mxfp4_tiny_batch = + mxfp4_path && n_tokens <= 4u; const uint32_t use_sorted_pairs = n_tokens > 1u && + !use_mxfp4_tiny_batch && (!q4k_path || n_tokens >= 32u) && !disable_resident_iq2_sorted; const uint32_t use_expert_tiles = use_sorted_pairs; const uint32_t expert_tile_m = 8u; const uint32_t write_gate_up = 0u; const uint32_t use_p2_sorted = 0u; - const uint32_t use_atomic_down = use_expert_tiles && n_tokens >= 128u; - const uint32_t use_gate_row2048 = !q4k_path && use_expert_tiles && n_tokens >= 128u; - const uint32_t use_down_tile16 = !q4k_path && use_atomic_down && n_tokens >= 128u; + const uint32_t use_atomic_down = + !mxfp4_path && use_expert_tiles && n_tokens >= 128u; + const uint32_t use_gate_row2048 = + !q4k_path && !mxfp4_path && use_expert_tiles && n_tokens >= 128u; + const uint32_t use_down_tile16 = + !q4k_path && !mxfp4_path && use_atomic_down && n_tokens >= 128u; const uint32_t use_decode_lut_gate = - n_tokens == 1u && xq_blocks <= 16u; + !mxfp4_path && n_tokens == 1u && xq_blocks <= 16u; const uint32_t gate_row_span = 1024u; const uint32_t down_row_span = 1024u; - const uint32_t use_down_row2048 = !q4k_path && use_atomic_down && use_down_tile16; + const uint32_t use_down_row2048 = + !q4k_path && !mxfp4_path && use_atomic_down && use_down_tile16; const uint32_t use_direct_down_sum6 = - n_tokens == 1u && n_expert <= DS4_ROCM_N_EXPERT_USED; + (n_tokens == 1u || use_mxfp4_tiny_batch) && + n_expert <= DS4_ROCM_N_EXPERT_USED; uint32_t *sorted_pairs = NULL; uint32_t *sorted_offsets = NULL; uint32_t *sorted_counts = NULL; @@ -1135,6 +1147,14 @@ static int routed_moe_launch( gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, 0u, write_gate_up, clamp); } + } else if (mxfp4_path) { + dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); + moe_gate_up_mid_mxfp4_expert_tile8_row32_kernel<<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + 0u, write_gate_up, clamp); } else if (use_gate_row2048) { if (gate_row_span == 512u) { dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); @@ -1252,6 +1272,24 @@ static int routed_moe_launch( n_expert, write_gate_up, clamp); + } else if (mxfp4_path) { + dim3 mxgrid((expert_mid_dim + 7u) / 8u, pair_count, 1); + moe_gate_up_mid_decode_mxfp4_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected_exec->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); } else if (use_decode_lut_gate) { moe_gate_up_mid_decode_lut_qwarp32_kernel<<>>( (float *)gate->ptr, @@ -1453,6 +1491,31 @@ static int routed_moe_launch( midq_blocks, out_dim, n_expert); + } else if (mxfp4_path) { + dim3 mxgrid((out_dim + 7u) / 8u, n_tokens, 1); + if (n_tokens == 1u) { + moe_down_mxfp4_sum6_qwarp32_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected_exec->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } else { + moe_down_mxfp4_sum6_qwarp32_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected_exec->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } } else { moe_down_sum6_qwarp32_kernel<<>>( (float *)out->ptr, @@ -1489,6 +1552,14 @@ static int routed_moe_launch( down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, midq_blocks, out_dim, n_expert, use_atomic_down); } + } else if (mxfp4_path) { + dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); + moe_down_mxfp4_expert_tile8_row32_kernel<<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, + down_expert_bytes, down_row_bytes, + midq_blocks, out_dim); } else if (use_down_row2048) { if (down_row_span == 512u) { dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); diff --git a/tests/test_mxfp4_rocm.c b/tests/test_mxfp4_rocm.c new file mode 100644 index 000000000..0065aa078 --- /dev/null +++ b/tests/test_mxfp4_rocm.c @@ -0,0 +1,525 @@ +/* Synthetic end-to-end test for the ROCm MXFP4 routed-MoE paths. + * + * The ROCm kernels quantize both the input and the fused FP32 SwiGLU mid + * activation to Q8_K. The CPU oracle below independently mirrors that + * quantization before applying the MXFP4 weights. Four repeated routing + * patterns keep the 512-token reference inexpensive while still exercising + * token indexing, expert bucketing, rectangular expert matrices, multiple + * Q8_K chunks per row, and expert IDs at both ends of a 256-expert table. + */ + +#include "ds4_gpu.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#define MXFP4_TYPE 39u +#define QK_MXFP4 32u +#define QK_K 256u +#define N_TOTAL_EXPERT 256u +#define N_EXPERT 6u +#define MODEL_DIM 512u +#define FFN_DIM 256u +#define N_PATTERN 4u +#define CLAMP 7.0f + +typedef struct { + uint8_t e; + uint8_t qs[QK_MXFP4 / 2u]; +} block_mxfp4; + +typedef struct { + float d; + int8_t qs[QK_K]; +} ref_block_q8_K; + +typedef struct { + float x[MODEL_DIM]; + int32_t selected[N_EXPERT]; + float weights[N_EXPERT]; + float mid[N_EXPERT * FFN_DIM]; + float out[MODEL_DIM]; +} reference_pattern; + +static const float mxfp4_values[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, +}; + +static uint64_t align_up(uint64_t value, uint64_t alignment) { + return (value + alignment - 1u) / alignment * alignment; +} + +static uint32_t mix32(uint32_t x) { + x ^= x >> 16u; + x *= 0x7feb352du; + x ^= x >> 15u; + x *= 0x846ca68bu; + x ^= x >> 16u; + return x; +} + +static float e8m0_to_f32(uint8_t e) { + uint32_t bits = e == 0u ? 0x00400000u : (uint32_t)e << 23u; + float value; + memcpy(&value, &bits, sizeof(value)); + return value; +} + +static void quantize_q8_K_block(ref_block_q8_K *out, const float *x) { + float amax = 0.0f; + float maxv = 0.0f; + for (uint32_t i = 0; i < QK_K; i++) { + const float ax = fabsf(x[i]); + /* The device reduction also keeps the lower index on an exact tie. */ + if (ax > amax) { + amax = ax; + maxv = x[i]; + } + } + if (amax == 0.0f) { + out->d = 0.0f; + memset(out->qs, 0, sizeof(out->qs)); + return; + } + + const float iscale = -127.0f / maxv; + out->d = 1.0f / iscale; + for (uint32_t i = 0; i < QK_K; i++) { + long q = lrintf(iscale * x[i]); + if (q > 127) q = 127; + if (q < -128) q = -128; + out->qs[i] = (int8_t)q; + } +} + +static void quantize_q8_K(ref_block_q8_K *out, + const float *x, + uint32_t n) { + for (uint32_t block = 0; block < n / QK_K; block++) { + quantize_q8_K_block(out + block, x + (uint64_t)block * QK_K); + } +} + +static float dot_mxfp4_q8_K(const block_mxfp4 *row, + const ref_block_q8_K *x, + uint32_t input_dim) { + float sum = 0.0f; + const uint32_t mxfp4_per_q8 = QK_K / QK_MXFP4; + for (uint32_t block = 0; block < input_dim / QK_MXFP4; block++) { + const block_mxfp4 *b = row + block; + const ref_block_q8_K *xb = x + block / mxfp4_per_q8; + const uint32_t q8_offset = (block % mxfp4_per_q8) * QK_MXFP4; + const float scale = e8m0_to_f32(b->e) * xb->d; + for (uint32_t i = 0; i < QK_MXFP4 / 2u; i++) { + const uint8_t q = b->qs[i]; + sum += scale * mxfp4_values[q & 15u] * + (float)xb->qs[q8_offset + i]; + sum += scale * mxfp4_values[q >> 4u] * + (float)xb->qs[q8_offset + i + QK_MXFP4 / 2u]; + } + } + return sum; +} + +static void fill_matrix(block_mxfp4 *matrix, + uint32_t rows, + uint32_t input_dim, + uint32_t salt) { + const uint32_t blocks_per_row = input_dim / QK_MXFP4; + for (uint32_t expert = 0; expert < N_TOTAL_EXPERT; expert++) { + for (uint32_t row = 0; row < rows; row++) { + block_mxfp4 *blocks = matrix + + ((uint64_t)expert * rows + row) * blocks_per_row; + for (uint32_t block = 0; block < blocks_per_row; block++) { + block_mxfp4 *b = blocks + block; + const uint32_t key = salt ^ (expert * 0x9e3779b9u) ^ + (row * 0x85ebca6bu) ^ + (block * 0xc2b2ae35u); + b->e = (uint8_t)(120u + mix32(key) % 5u); + for (uint32_t i = 0; i < QK_MXFP4 / 2u; i++) { + const uint32_t h = mix32(key + i * 0x27d4eb2du); + b->qs[i] = (uint8_t)((h & 15u) | (((h >> 9u) & 15u) << 4u)); + } + } + } + } +} + +static const block_mxfp4 *matrix_row(const block_mxfp4 *matrix, + uint32_t expert, + uint32_t row, + uint32_t rows, + uint32_t input_dim) { + return matrix + ((uint64_t)expert * rows + row) * + (input_dim / QK_MXFP4); +} + +static void init_patterns(reference_pattern patterns[N_PATTERN]) { + static const int32_t selected[N_PATTERN][N_EXPERT] = { + { 0, 1, 2, 3, 4, 255 }, + { 17, 63, 127, 128, 200, 254 }, + { 255, 0, 129, 42, 11, 201 }, + { 5, 85, 170, 250, 13, 199 }, + }; + static const float base_weights[N_EXPERT] = { + 0.25f, 0.20f, 0.18f, 0.15f, 0.12f, 0.10f, + }; + + memset(patterns, 0, sizeof(reference_pattern) * N_PATTERN); + for (uint32_t p = 0; p < N_PATTERN; p++) { + memcpy(patterns[p].selected, selected[p], sizeof(selected[p])); + for (uint32_t slot = 0; slot < N_EXPERT; slot++) { + patterns[p].weights[slot] = + base_weights[(slot + p) % N_EXPERT]; + } + for (uint32_t i = 0; i < MODEL_DIM; i++) { + const uint32_t h = mix32(i + 1u + p * 0x9e3779b9u); + patterns[p].x[i] = + (float)((int32_t)(h % 255u) - 127) / 256.0f; + } + /* Give each vector a unique signed maximum so Q8_K scale selection + * is deterministic on both the CPU and GPU. */ + patterns[p].x[19u + p * 47u] = (p & 1u) ? 0.875f : -0.875f; + } +} + +static void build_reference(reference_pattern *pattern, + const block_mxfp4 *gate_matrix, + const block_mxfp4 *up_matrix, + const block_mxfp4 *down_matrix) { + ref_block_q8_K xq[MODEL_DIM / QK_K]; + ref_block_q8_K midq[N_EXPERT][FFN_DIM / QK_K]; + quantize_q8_K(xq, pattern->x, MODEL_DIM); + + for (uint32_t slot = 0; slot < N_EXPERT; slot++) { + const uint32_t expert = (uint32_t)pattern->selected[slot]; + float *mid = pattern->mid + (uint64_t)slot * FFN_DIM; + for (uint32_t row = 0; row < FFN_DIM; row++) { + float gate = dot_mxfp4_q8_K( + matrix_row(gate_matrix, expert, row, FFN_DIM, MODEL_DIM), + xq, MODEL_DIM); + float up = dot_mxfp4_q8_K( + matrix_row(up_matrix, expert, row, FFN_DIM, MODEL_DIM), + xq, MODEL_DIM); + if (gate > CLAMP) gate = CLAMP; + if (up > CLAMP) up = CLAMP; + if (up < -CLAMP) up = -CLAMP; + mid[row] = (gate / (1.0f + expf(-gate))) * up * + pattern->weights[slot]; + } + quantize_q8_K(midq[slot], mid, FFN_DIM); + } + + for (uint32_t row = 0; row < MODEL_DIM; row++) { + float sum = 0.0f; + for (uint32_t slot = 0; slot < N_EXPERT; slot++) { + const uint32_t expert = (uint32_t)pattern->selected[slot]; + sum += dot_mxfp4_q8_K( + matrix_row(down_matrix, expert, row, MODEL_DIM, FFN_DIM), + midq[slot], FFN_DIM); + } + pattern->out[row] = sum; + } +} + +static int compare_repeated(const char *name, + const float *actual, + uint32_t n_tokens, + uint32_t token_elems, + const reference_pattern patterns[N_PATTERN], + bool compare_mid, + float abs_tolerance, + float rel_tolerance) { + float max_abs = 0.0f; + float max_ratio = 0.0f; + uint64_t max_abs_index = 0u; + uint64_t max_ratio_index = 0u; + uint64_t failures = 0u; + const uint64_t count = (uint64_t)n_tokens * token_elems; + + for (uint32_t token = 0; token < n_tokens; token++) { + const float *expected = compare_mid ? + patterns[token % N_PATTERN].mid : + patterns[token % N_PATTERN].out; + for (uint32_t i = 0; i < token_elems; i++) { + const uint64_t index = (uint64_t)token * token_elems + i; + const float got = actual[index]; + const float want = expected[i]; + if (!isfinite(got) || !isfinite(want)) { + fprintf(stderr, + "MXFP4 ROCm tokens=%u %s non-finite at token=%u element=%u " + "got=%g expected=%g\n", + n_tokens, name, token, i, got, want); + return 0; + } + const float error = fabsf(got - want); + const float allowed = abs_tolerance + rel_tolerance * fabsf(want); + const float ratio = allowed > 0.0f ? error / allowed : error; + if (error > max_abs) { + max_abs = error; + max_abs_index = index; + } + if (ratio > max_ratio) { + max_ratio = ratio; + max_ratio_index = index; + } + if (error > allowed) failures++; + } + } + + fprintf(stderr, + "MXFP4 ROCm tokens=%-3u %-3s max_abs=%-10g at=%llu " + "max_tol_ratio=%g at=%llu failures=%llu/%llu\n", + n_tokens, name, max_abs, (unsigned long long)max_abs_index, + max_ratio, (unsigned long long)max_ratio_index, + (unsigned long long)failures, (unsigned long long)count); + return failures == 0u; +} + +static int run_case(uint32_t n_tokens, + const void *model, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + const reference_pattern patterns[N_PATTERN]) { + const uint64_t token_x_count = (uint64_t)n_tokens * MODEL_DIM; + const uint64_t route_count = (uint64_t)n_tokens * N_EXPERT; + const uint64_t mid_count = route_count * FFN_DIM; + const uint64_t down_count = route_count * MODEL_DIM; + const uint64_t out_count = token_x_count; + float *x = (float *)calloc((size_t)token_x_count, sizeof(float)); + int32_t *selected = (int32_t *)calloc((size_t)route_count, sizeof(int32_t)); + float *weights = (float *)calloc((size_t)route_count, sizeof(float)); + float *mid_actual = (float *)calloc((size_t)mid_count, sizeof(float)); + float *out_actual = (float *)calloc((size_t)out_count, sizeof(float)); + ds4_gpu_tensor *x_tensor = NULL; + ds4_gpu_tensor *selected_tensor = NULL; + ds4_gpu_tensor *weights_tensor = NULL; + ds4_gpu_tensor *gate_tensor = NULL; + ds4_gpu_tensor *up_tensor = NULL; + ds4_gpu_tensor *mid_tensor = NULL; + ds4_gpu_tensor *experts_tensor = NULL; + ds4_gpu_tensor *out_tensor = NULL; + int ok = x && selected && weights && mid_actual && out_actual; + + for (uint32_t token = 0; ok && token < n_tokens; token++) { + const reference_pattern *pattern = &patterns[token % N_PATTERN]; + memcpy(x + (uint64_t)token * MODEL_DIM, + pattern->x, sizeof(pattern->x)); + memcpy(selected + (uint64_t)token * N_EXPERT, + pattern->selected, sizeof(pattern->selected)); + memcpy(weights + (uint64_t)token * N_EXPERT, + pattern->weights, sizeof(pattern->weights)); + } + + if (ok) x_tensor = ds4_gpu_tensor_alloc(token_x_count * sizeof(float)); + if (ok) selected_tensor = ds4_gpu_tensor_alloc(route_count * sizeof(int32_t)); + if (ok) weights_tensor = ds4_gpu_tensor_alloc(route_count * sizeof(float)); + if (ok) gate_tensor = ds4_gpu_tensor_alloc(mid_count * sizeof(float)); + if (ok) up_tensor = ds4_gpu_tensor_alloc(mid_count * sizeof(float)); + if (ok) mid_tensor = ds4_gpu_tensor_alloc(mid_count * sizeof(float)); + if (ok) experts_tensor = ds4_gpu_tensor_alloc(down_count * sizeof(float)); + if (ok) out_tensor = ds4_gpu_tensor_alloc(out_count * sizeof(float)); + ok = ok && x_tensor && selected_tensor && weights_tensor && gate_tensor && + up_tensor && mid_tensor && experts_tensor && out_tensor; + + ok = ok && ds4_gpu_tensor_write( + x_tensor, 0u, x, token_x_count * sizeof(float)); + ok = ok && ds4_gpu_tensor_write( + selected_tensor, 0u, selected, route_count * sizeof(int32_t)); + ok = ok && ds4_gpu_tensor_write( + weights_tensor, 0u, weights, route_count * sizeof(float)); + + if (ok && n_tokens == 1u) { + ok = ds4_gpu_routed_moe_one_tensor( + out_tensor, gate_tensor, up_tensor, mid_tensor, experts_tensor, + model, model_size, gate_offset, up_offset, down_offset, + MXFP4_TYPE, MXFP4_TYPE, gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + MODEL_DIM, FFN_DIM, MODEL_DIM, + selected_tensor, weights_tensor, N_TOTAL_EXPERT, N_EXPERT, + CLAMP, x_tensor, NULL, 0u, true); + } else if (ok) { + bool mid_is_f16 = true; + ok = ds4_gpu_routed_moe_batch_tensor( + out_tensor, gate_tensor, up_tensor, mid_tensor, experts_tensor, + model, model_size, gate_offset, up_offset, down_offset, + MXFP4_TYPE, MXFP4_TYPE, gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + MODEL_DIM, FFN_DIM, MODEL_DIM, + selected_tensor, weights_tensor, N_TOTAL_EXPERT, N_EXPERT, + CLAMP, x_tensor, 0u, n_tokens, &mid_is_f16, true); + if (ok && mid_is_f16) { + fprintf(stderr, + "MXFP4 ROCm tokens=%u unexpectedly reported FP16 mid storage\n", + n_tokens); + ok = 0; + } + } + + if (!ok) { + fprintf(stderr, "MXFP4 ROCm tokens=%u launch failed\n", n_tokens); + } + ok = ok && ds4_gpu_tensor_read( + mid_tensor, 0u, mid_actual, mid_count * sizeof(float)); + ok = ok && ds4_gpu_tensor_read( + out_tensor, 0u, out_actual, out_count * sizeof(float)); + + if (ok) { + /* Gate/up are optional scratch outputs in the optimized ROCm paths; + * mid is the public, stable result of that fused stage. */ + const int mid_ok = compare_repeated( + "mid", mid_actual, n_tokens, N_EXPERT * FFN_DIM, patterns, true, + 1.0e-4f, 1.0e-4f); + const int out_ok = compare_repeated( + "out", out_actual, n_tokens, MODEL_DIM, patterns, false, + 2.0e-4f, 1.0e-4f); + ok = mid_ok && out_ok; + } + + ds4_gpu_tensor_free(out_tensor); + ds4_gpu_tensor_free(experts_tensor); + ds4_gpu_tensor_free(mid_tensor); + ds4_gpu_tensor_free(up_tensor); + ds4_gpu_tensor_free(gate_tensor); + ds4_gpu_tensor_free(weights_tensor); + ds4_gpu_tensor_free(selected_tensor); + ds4_gpu_tensor_free(x_tensor); + free(out_actual); + free(mid_actual); + free(weights); + free(selected); + free(x); + return ok; +} + +int main(void) { + /* Two through four tokens exercise the direct tiny-batch path; five + * tokens returns to expert-sorted tiles. Larger cases cover prefill. */ + static const uint32_t token_cases[] = { + 1u, 2u, 3u, 4u, 5u, 32u, 128u, 512u, + }; + const uint64_t gate_row_bytes = + (MODEL_DIM / QK_MXFP4) * sizeof(block_mxfp4); + const uint64_t gate_expert_bytes = FFN_DIM * gate_row_bytes; + const uint64_t gate_tensor_bytes = + N_TOTAL_EXPERT * gate_expert_bytes; + const uint64_t down_row_bytes = + (FFN_DIM / QK_MXFP4) * sizeof(block_mxfp4); + const uint64_t down_expert_bytes = MODEL_DIM * down_row_bytes; + const uint64_t down_tensor_bytes = + N_TOTAL_EXPERT * down_expert_bytes; + const uint64_t gate_offset = 0u; + const uint64_t up_offset = align_up(gate_tensor_bytes, 4096u); + const uint64_t down_offset = + align_up(up_offset + gate_tensor_bytes, 4096u); + const uint64_t model_size = + align_up(down_offset + down_tensor_bytes, 4096u); + FILE *model_file = NULL; + void *model = MAP_FAILED; + reference_pattern *patterns = NULL; + int initialized = 0; + int ok = sizeof(block_mxfp4) == 17u; + + if (!ok) { + fprintf(stderr, "MXFP4 ROCm unexpected block size %zu (expected 17)\n", + sizeof(block_mxfp4)); + return 1; + } + + model_file = tmpfile(); + if (model_file && + ftruncate(fileno(model_file), (off_t)model_size) == 0) { + model = mmap(NULL, (size_t)model_size, PROT_READ | PROT_WRITE, + MAP_SHARED, fileno(model_file), 0); + } + patterns = (reference_pattern *)calloc(N_PATTERN, sizeof(*patterns)); + if (!model_file || model == MAP_FAILED || !patterns) { + fprintf(stderr, "MXFP4 ROCm host allocation failed\n"); + ok = 0; + goto cleanup; + } + memset(model, 0, (size_t)model_size); + fill_matrix((block_mxfp4 *)((uint8_t *)model + gate_offset), + FFN_DIM, MODEL_DIM, 0x12345678u); + fill_matrix((block_mxfp4 *)((uint8_t *)model + up_offset), + FFN_DIM, MODEL_DIM, 0x9abcdef0u); + fill_matrix((block_mxfp4 *)((uint8_t *)model + down_offset), + MODEL_DIM, FFN_DIM, 0x0f1e2d3cu); + init_patterns(patterns); + + const block_mxfp4 *gate_matrix = + (const block_mxfp4 *)((const uint8_t *)model + gate_offset); + const block_mxfp4 *up_matrix = + (const block_mxfp4 *)((const uint8_t *)model + up_offset); + const block_mxfp4 *down_matrix = + (const block_mxfp4 *)((const uint8_t *)model + down_offset); + for (uint32_t p = 0; p < N_PATTERN; p++) { + build_reference(&patterns[p], gate_matrix, up_matrix, down_matrix); + } + + fprintf(stderr, + "MXFP4 ROCm synthetic model: %.2f MiB, experts=%u, " + "model_dim=%u, ffn_dim=%u, selected=%u\n", + (double)model_size / 1048576.0, N_TOTAL_EXPERT, + MODEL_DIM, FFN_DIM, N_EXPERT); + ok = ds4_gpu_init(); + initialized = ok; + if (!ok) { + fprintf(stderr, "MXFP4 ROCm ds4_gpu_init failed\n"); + goto cleanup; + } + ds4_gpu_set_quality(false); + ds4_gpu_set_ssd_streaming(false); + const uint64_t model_offsets[] = { + gate_offset, up_offset, down_offset, + }; + const uint64_t model_sizes[] = { + gate_tensor_bytes, gate_tensor_bytes, down_tensor_bytes, + }; + const uint64_t max_tensor_bytes = + gate_tensor_bytes > down_tensor_bytes ? + gate_tensor_bytes : down_tensor_bytes; + ok = ds4_gpu_set_model_map(model, model_size) && + ds4_gpu_set_model_fd(fileno(model_file)) && + ds4_gpu_set_model_map_spans( + model, model_size, model_offsets, model_sizes, + sizeof(model_offsets) / sizeof(model_offsets[0]), + max_tensor_bytes); + if (!ok) { + fprintf(stderr, "MXFP4 ROCm model cache setup failed\n"); + goto cleanup; + } + + for (uint32_t i = 0; i < sizeof(token_cases) / sizeof(token_cases[0]); i++) { + if (!run_case(token_cases[i], model, model_size, + gate_offset, up_offset, down_offset, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, patterns)) { + ok = 0; + } + } + +cleanup: + if (initialized) { + ds4_gpu_set_model_fd(-1); + ds4_gpu_cleanup(); + } + free(patterns); + if (model != MAP_FAILED) munmap(model, (size_t)model_size); + if (model_file) fclose(model_file); + fprintf(stderr, "MXFP4 ROCm routed MoE: %s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +}