diff --git a/Makefile b/Makefile index 089bcd76de..5162b2b5fe 100644 --- a/Makefile +++ b/Makefile @@ -55,6 +55,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) @@ -62,7 +63,7 @@ DS4_LINK_LIBS ?= $(CUDA_LDLIBS) METAL_LDLIBS := $(LDLIBS) endif -.PHONY: all help clean test test-metal-session-batch test-mxfp4-cuda 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-rocm test-metal-session-batch test-mxfp4-cuda 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) .PHONY: metal-decode-schedule-bench metal-prefill-variant-bench check-mxfp4-half-lut @@ -156,6 +157,8 @@ 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 test-rocm Core regression suite on ROCm-only hosts" @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" @@ -179,12 +182,32 @@ 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)" rocm: strix-halo +# Core regression suite for ROCm-only hosts: the CUDA-specific binaries +# (tests/test_sampling, the CUDA session/mixed-batch oracles) are not part +# of this target; run them through `make test` / `make cuda-regression` on +# CUDA hosts. Everything else mirrors `make test`. +test-rocm: + $(MAKE) -B ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \ + tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \ + ds4 ds4-server ds4-bench 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) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD" \ + DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \ + DS4_LINK_LIBS="$(ROCM_LDLIBS)" + ./ds4-eval --self-test-extractors + ./ds4_agent_test + ./ds4_test + ./tests/test_layer_pack + ./tests/test_engine_mgpu_placement + ./tests/test_gpu_args + ./tests/test_gpu_args_cli.sh + ds4: ds4_cli.o ds4_help.o linenoise.o ds4_gpu_args.o $(CORE_OBJS) $(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS) @@ -201,7 +224,7 @@ ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_ar $(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS) gguf-tools/quality-testing/score_official.o: gguf-tools/quality-testing/score_official.c ds4.h - $(CC) $(filter-out -ffast-math,$(QUALITY_CFLAGS)) -I. -c -o $@ $< + $(CC) $(filter-out -ffast-math,$(QUALITY_CFLAGS)) $(ROCM_HOST_CFLAGS) -I. -c -o $@ $< gguf-tools/quality-testing/score_official: gguf-tools/quality-testing/score_official.o $(CORE_OBJS) rax.o ds4_gpu_args.o $(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS) @@ -331,6 +354,21 @@ cuda/mmq/ds4_repack.o: cuda/mmq/ds4_repack.cu cuda/mmq/ds4_repack.h 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) + +tests/bench_mxfp4_rocm.o: tests/bench_mxfp4_rocm.c ds4_gpu.h + $(CC) $(filter-out -ffast-math,$(CFLAGS)) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD -I. -c -o $@ $< + +tests/bench_mxfp4_rocm: tests/bench_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 @@ -430,14 +468,14 @@ ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) ifeq ($(UNAME_S),Darwin) $(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS) else - $(NVCC) $(NVCCFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(CUDA_LDLIBS) + $(DS4_LINK) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(DS4_LINK_LIBS) endif ds4_agent_test: ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) ifeq ($(UNAME_S),Darwin) $(CC) $(CFLAGS) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(METAL_LDLIBS) else - $(NVCC) $(NVCCFLAGS) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(CUDA_LDLIBS) + $(DS4_LINK) -o $@ ds4_agent_test.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o $(CORE_OBJS) $(DS4_LINK_LIBS) endif test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \ @@ -488,4 +526,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 gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_cuda 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 gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda 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 a74d50b5ad..64622beec6 100644 --- a/QA_BEFORE_RELEASES.md +++ b/QA_BEFORE_RELEASES.md @@ -463,6 +463,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-0731.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 0f5324dec5..db5b4688b9 100644 --- a/README.md +++ b/README.md @@ -420,6 +420,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 481232bdfc..229a769c42 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 449140b523..2d12da18c9 100644 --- a/ds4.c +++ b/ds4.c @@ -50946,7 +50946,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 46ba223df6..f1f7547357 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 684f867294..12d8998f21 100644 --- a/rocm/ds4_rocm_moe.cuh +++ b/rocm/ds4_rocm_moe.cuh @@ -288,6 +288,273 @@ __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; +} + +/* Aligned-activation variant of the 8-pair chunk dot. The staged copies + * hold each pair's 256 quant bytes 16-byte aligned with the Q8_K scale in + * a separate array, so one b128 load replaces four dword loads per half + * sub-block: this kernel family is issue-bound, not bandwidth-bound, and + * the load count per dp4a is what gates it. The low/high nibble sums use + * separate integer accumulators; integer reassociation is exact, so the + * result is bit-identical to the struct-layout helper. */ +__device__ static void dev_dot_mxfp4_q8a_block8( + const cuda_block_mxfp4 *x8, + const int8_t *q0, const int8_t *q1, const int8_t *q2, const int8_t *q3, + const int8_t *q4, const int8_t *q5, const int8_t *q6, const int8_t *q7, + float d0, float d1, float d2, float d3, + float d4, float d5, float d6, float d7, + uint32_t n, + float acc[8]) { + const int8_t *qs[8] = { q0, q1, q2, q3, q4, q5, q6, q7 }; + const float ds[8] = { d0, d1, d2, d3, d4, d5, d6, d7 }; + 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; + /* One misaligned b128 covers the block's 16 packed bytes; gfx11 + * global loads support arbitrary alignment. */ + const uint4 w4 = *(const uint4 *)(x->qs); + int32_t wlo[4], whi[4]; + wlo[0] = (int32_t)dev_mxfp4_unpack4(w4.x); + whi[0] = (int32_t)dev_mxfp4_unpack4(w4.x >> 4u); + wlo[1] = (int32_t)dev_mxfp4_unpack4(w4.y); + whi[1] = (int32_t)dev_mxfp4_unpack4(w4.y >> 4u); + wlo[2] = (int32_t)dev_mxfp4_unpack4(w4.z); + whi[2] = (int32_t)dev_mxfp4_unpack4(w4.z >> 4u); + wlo[3] = (int32_t)dev_mxfp4_unpack4(w4.w); + whi[3] = (int32_t)dev_mxfp4_unpack4(w4.w >> 4u); + const float d = dev_e8m0_to_f32(x->e); + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < n) { + const int4 lo = *(const int4 *)(qs[p] + sb * 32u); + const int4 hi = *(const int4 *)(qs[p] + sb * 32u + 16u); + int32_t s = 0; + int32_t t = 0; + s = __dp4a(wlo[0], lo.x, s); + t = __dp4a(whi[0], hi.x, t); + s = __dp4a(wlo[1], lo.y, s); + t = __dp4a(whi[1], hi.y, t); + s = __dp4a(wlo[2], lo.z, s); + t = __dp4a(whi[2], hi.z, t); + s = __dp4a(wlo[3], lo.w, s); + t = __dp4a(whi[3], hi.w, t); + chunk[p] += d * (float)(s + t); + } + } + } + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < n) acc[p] += 0.5f * ds[p] * chunk[p]; + } +} + +/* Fused gate+up flavor: one activation load pair feeds both weight + * matrices, halving the dominant b128 traffic of the gate/up tile + * kernel. Per-matrix accumulation is unchanged. */ +__device__ static void dev_dot_mxfp4_q8a_pair_block8( + const cuda_block_mxfp4 *g8, + const cuda_block_mxfp4 *u8, + const int8_t *q0, const int8_t *q1, const int8_t *q2, const int8_t *q3, + const int8_t *q4, const int8_t *q5, const int8_t *q6, const int8_t *q7, + float d0, float d1, float d2, float d3, + float d4, float d5, float d6, float d7, + uint32_t n, + float gate[8], + float up[8]) { + const int8_t *qs[8] = { q0, q1, q2, q3, q4, q5, q6, q7 }; + const float ds[8] = { d0, d1, d2, d3, d4, d5, d6, d7 }; + float gchunk[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float uchunk[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 uint4 gw = *(const uint4 *)(g8[sb].qs); + const uint4 uw = *(const uint4 *)(u8[sb].qs); + int32_t glo[4], ghi[4], ulo[4], uhi[4]; + glo[0] = (int32_t)dev_mxfp4_unpack4(gw.x); + ghi[0] = (int32_t)dev_mxfp4_unpack4(gw.x >> 4u); + glo[1] = (int32_t)dev_mxfp4_unpack4(gw.y); + ghi[1] = (int32_t)dev_mxfp4_unpack4(gw.y >> 4u); + glo[2] = (int32_t)dev_mxfp4_unpack4(gw.z); + ghi[2] = (int32_t)dev_mxfp4_unpack4(gw.z >> 4u); + glo[3] = (int32_t)dev_mxfp4_unpack4(gw.w); + ghi[3] = (int32_t)dev_mxfp4_unpack4(gw.w >> 4u); + ulo[0] = (int32_t)dev_mxfp4_unpack4(uw.x); + uhi[0] = (int32_t)dev_mxfp4_unpack4(uw.x >> 4u); + ulo[1] = (int32_t)dev_mxfp4_unpack4(uw.y); + uhi[1] = (int32_t)dev_mxfp4_unpack4(uw.y >> 4u); + ulo[2] = (int32_t)dev_mxfp4_unpack4(uw.z); + uhi[2] = (int32_t)dev_mxfp4_unpack4(uw.z >> 4u); + ulo[3] = (int32_t)dev_mxfp4_unpack4(uw.w); + uhi[3] = (int32_t)dev_mxfp4_unpack4(uw.w >> 4u); + const float gd = dev_e8m0_to_f32(g8[sb].e); + const float ud = dev_e8m0_to_f32(u8[sb].e); + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < n) { + const int4 lo = *(const int4 *)(qs[p] + sb * 32u); + const int4 hi = *(const int4 *)(qs[p] + sb * 32u + 16u); + int32_t gs = 0; + int32_t gt = 0; + int32_t us = 0; + int32_t ut = 0; + gs = __dp4a(glo[0], lo.x, gs); + gt = __dp4a(ghi[0], hi.x, gt); + us = __dp4a(ulo[0], lo.x, us); + ut = __dp4a(uhi[0], hi.x, ut); + gs = __dp4a(glo[1], lo.y, gs); + gt = __dp4a(ghi[1], hi.y, gt); + us = __dp4a(ulo[1], lo.y, us); + ut = __dp4a(uhi[1], hi.y, ut); + gs = __dp4a(glo[2], lo.z, gs); + gt = __dp4a(ghi[2], hi.z, gt); + us = __dp4a(ulo[2], lo.z, us); + ut = __dp4a(uhi[2], hi.z, ut); + gs = __dp4a(glo[3], lo.w, gs); + gt = __dp4a(ghi[3], hi.w, gt); + us = __dp4a(ulo[3], lo.w, us); + ut = __dp4a(uhi[3], hi.w, ut); + gchunk[p] += gd * (float)(gs + gt); + uchunk[p] += ud * (float)(us + ut); + } + } + } + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < n) { + gate[p] += 0.5f * ds[p] * gchunk[p]; + up[p] += 0.5f * ds[p] * uchunk[p]; + } + } +} + +__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 +2105,501 @@ __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) { + /* Grid is (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 and the expert's weights stream from DRAM ~once + * instead of once per 8-token tile. Pure scheduling change; every + * output element's arithmetic is untouched. */ + uint32_t tile = blockIdx.x; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.y * 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]; + /* Staged activations live as 16-byte-aligned quant slices with the + * Q8_K scales in a separate array: the dot helper is issue-bound, and + * the aligned layout lets it load b128 instead of four dwords. The + * layout also drops the unused 32-byte bsums from every staged chunk. + * Sized by the launch to 8 * xq_blocks * (256 + 4) bytes. */ + extern __shared__ int4 sxq4[]; + int8_t *sqs = (int8_t *)sxq4; + float *sds = (float *)(sqs + (uint64_t)8u * xq_blocks * 256u); + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + /* Fixed-count predicated setup keeps pair/xqb in registers instead of + * dynamically indexed scratch. The tail xqb entries stay at valid + * dummy rows and are never read past np. */ + const cuda_block_q8_K *xqb[8] = { xq, xq, xq, xq, xq, xq, xq, xq }; + uint32_t np = count - local_start; + if (np > 8u) np = 8u; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < np) { + pair[p] = sorted_pairs[offsets[expert] + local_start + p]; + xqb[p] = xq + (uint64_t)(pair[p] / n_expert) * xq_blocks; + } + } + const int staged = xq_blocks <= 28u; + if (staged) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks * 16u; i += blockDim.x) { + const uint32_t u = i / (xq_blocks * 16u); + const uint32_t rem = i - u * (xq_blocks * 16u); + const uint32_t b = rem >> 4u; + const uint32_t k = rem & 15u; + const uint32_t sp = sorted_pairs[offsets[expert] + local_start + u]; + const int32_t *src = (const int32_t *) + xq[(uint64_t)(sp / n_expert) * xq_blocks + b].qs + k * 4u; + ((int4 *)sqs)[(u * xq_blocks + b) * 16u + k] = + make_int4(src[0], src[1], src[2], src[3]); + } + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + const uint32_t u = i / xq_blocks; + const uint32_t b = i - u * xq_blocks; + const uint32_t sp = sorted_pairs[offsets[expert] + local_start + u]; + sds[i] = xq[(uint64_t)(sp / n_expert) * xq_blocks + b].d; + } + __syncthreads(); + } + 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}; + const uint64_t sq = (uint64_t)xq_blocks << 8u; + 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); + if (staged) { + const int8_t *qb = sqs + ((uint64_t)b << 8u); + dev_dot_mxfp4_q8a_pair_block8(gb, ub, + qb, qb + sq, qb + 2u * sq, qb + 3u * sq, + qb + 4u * sq, qb + 5u * sq, qb + 6u * sq, qb + 7u * sq, + sds[b], sds[xq_blocks + b], sds[2u * xq_blocks + b], sds[3u * xq_blocks + b], + sds[4u * xq_blocks + b], sds[5u * xq_blocks + b], sds[6u * xq_blocks + b], sds[7u * xq_blocks + b], + np, gate, up); + } else { + dev_dot_mxfp4_q8_K_block8(gb, xqb[0] + b, xqb[1] + b, xqb[2] + b, xqb[3] + b, + xqb[4] + b, xqb[5] + b, xqb[6] + b, xqb[7] + b, np, gate); + dev_dot_mxfp4_q8_K_block8(ub, xqb[0] + b, xqb[1] + b, xqb[2] + b, xqb[3] + b, + xqb[4] + b, xqb[5] + b, xqb[6] + b, xqb[7] + b, np, up); + } + } + /* pair == tok * n_expert + slot, so it indexes weights directly. */ + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p >= np) continue; + 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[pair[p]]; + } + } +} + +/* Generic-tile MXFP4 gate/up kernel: TILE_TOKENS tokens per expert tile, + * QWARPS 32-row slots per block (blockDim = QWARPS*8). Lets the launch + * trade staging footprint against resident-warp count: tile4 stages + * TILE_TOKENS*xq_blocks*260 B (29.1 KiB at 7168 dims -> 4 blocks/WGP), + * row64 keeps 8-token staging but runs 512 threads (2 blocks/WGP -> 32 + * resident warps). Both target the tile8 kernel's latency-bound profile + * (1-2 blocks/WGP today) without touching arithmetic. + * + * Exactness: per (token,row) output the accumulation order is identical + * to moe_gate_up_mid_mxfp4_expert_tile8_row32_kernel - same 8-lane split + * of the q8_K block loop, same dev_dot helpers, same quarter-warp + * reduction, same epilogue. Only grid/tile geometry changes. */ +template +__global__ static void moe_gate_up_mid_mxfp4_expert_tileN_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.x; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.y * QWARPS + (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]; + extern __shared__ int4 sxq4[]; + int8_t *sqs = (int8_t *)sxq4; + float *sds = (float *)(sqs + (uint64_t)TILE_TOKENS * xq_blocks * 256u); + uint32_t pair[TILE_TOKENS] = {0}; + const cuda_block_q8_K *xqb[TILE_TOKENS]; + #pragma unroll + for (uint32_t p = 0; p < TILE_TOKENS; p++) xqb[p] = xq; + uint32_t np = count - local_start; + if (np > TILE_TOKENS) np = TILE_TOKENS; + #pragma unroll + for (uint32_t p = 0; p < TILE_TOKENS; p++) { + if (p < np) { + pair[p] = sorted_pairs[offsets[expert] + local_start + p]; + xqb[p] = xq + (uint64_t)(pair[p] / n_expert) * xq_blocks; + } + } + const int staged = xq_blocks <= 28u; + if (staged) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks * 16u; i += blockDim.x) { + const uint32_t u = i / (xq_blocks * 16u); + const uint32_t rem = i - u * (xq_blocks * 16u); + const uint32_t b = rem >> 4u; + const uint32_t k = rem & 15u; + const uint32_t sp = sorted_pairs[offsets[expert] + local_start + u]; + const int32_t *src = (const int32_t *) + xq[(uint64_t)(sp / n_expert) * xq_blocks + b].qs + k * 4u; + ((int4 *)sqs)[(u * xq_blocks + b) * 16u + k] = + make_int4(src[0], src[1], src[2], src[3]); + } + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + const uint32_t u = i / xq_blocks; + const uint32_t b = i - u * xq_blocks; + const uint32_t sp = sorted_pairs[offsets[expert] + local_start + u]; + sds[i] = xq[(uint64_t)(sp / n_expert) * xq_blocks + b].d; + } + __syncthreads(); + } + 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}; + const uint64_t sq = (uint64_t)xq_blocks << 8u; + 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); + if (staged) { + const int8_t *qb = sqs + ((uint64_t)b << 8u); + const int8_t *qp[8]; + float sd[8]; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + qp[p] = qb + (p < TILE_TOKENS ? p : 0u) * sq; + sd[p] = sds[(p < TILE_TOKENS ? p : 0u) * xq_blocks + b]; + } + dev_dot_mxfp4_q8a_pair_block8(gb, ub, + qp[0], qp[1], qp[2], qp[3], qp[4], qp[5], qp[6], qp[7], + sd[0], sd[1], sd[2], sd[3], sd[4], sd[5], sd[6], sd[7], + np, gate, up); + } else { + const cuda_block_q8_K *xq8[8]; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) xq8[p] = xqb[p < TILE_TOKENS ? p : 0u] + b; + dev_dot_mxfp4_q8_K_block8(gb, xq8[0], xq8[1], xq8[2], xq8[3], + xq8[4], xq8[5], xq8[6], xq8[7], np, gate); + dev_dot_mxfp4_q8_K_block8(ub, xq8[0], xq8[1], xq8[2], xq8[3], + xq8[4], xq8[5], xq8[6], xq8[7], np, up); + } + } + /* pair == tok * n_expert + slot, so it indexes weights directly. */ + #pragma unroll + for (uint32_t p = 0; p < TILE_TOKENS; p++) { + if (p >= np) continue; + 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[pair[p]]; + } + } +} + +/* Prefill-oriented MXFP4 gate/up kernel: one block covers 8 output rows + * for up to 128 tokens of a single expert, with the 8 gate+up weight + * rows staged in LDS (16 * gate_row_bytes; 60.9 KiB at row_bytes=3808). + * Each expert's weights are then read from DRAM ~once per 128-token + * tile instead of once per 8-token tile (16x re-read reduction at + * balanced routing), while the quantized activations stream from global + * (the expert's ~1 MiB xq working set stays L2-resident across its row + * blocks). + * + * Exactness: per (token,row) output the accumulation order is identical + * to moe_gate_up_mid_mxfp4_expert_tile8_row32_kernel - same 8-lane split + * of the q8_K block loop, same dev_dot_mxfp4_q8_K_block8 helper, same + * quarter-warp reduction, same epilogue. Only data movement changes. */ +__global__ static void moe_gate_up_mid_mxfp4_expert_row8_ldsB_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; + const uint32_t expert = tile_experts[tile]; + const uint32_t count = counts[expert]; + if (max_count != 0u && count >= max_count) return; + const uint32_t local_start = tile_starts[tile]; + const uint32_t row0 = blockIdx.x * 8u; + if (row0 >= expert_mid_dim) return; + + /* LDS layout: 16 segments (gate rows 0..7 then up rows 0..7), each + * gate_row_bytes long, laid out exactly like the global weight row so + * the dot helper can point into it directly. */ + extern __shared__ char sB[]; + { + const uint32_t nseg4 = (uint32_t)(gate_row_bytes >> 4u); /* 16 B units per row */ + const int4 *g4 = (const int4 *)(gate_base + (uint64_t)expert * gate_expert_bytes) + (uint64_t)row0 * nseg4; + const int4 *u4 = (const int4 *)(up_base + (uint64_t)expert * gate_expert_bytes) + (uint64_t)row0 * nseg4; + int4 *dst = (int4 *)sB; + for (uint32_t m = 0; m < 2u; m++) { + const int4 *src = m == 0u ? g4 : u4; + for (uint32_t r = 0; r < 8u; r++) { + for (uint32_t i = threadIdx.x; i < nseg4; i += blockDim.x) { + dst[(m * 8u + r) * nseg4 + i] = src[r * nseg4 + i]; + } + } + } + __syncthreads(); + } + + const uint32_t lane = threadIdx.x & 7u; + const uint32_t qwarp = threadIdx.x >> 3u; /* 0..31 */ + const uint32_t row_in = qwarp & 7u; /* row within the block's 8 */ + const uint32_t tok_grp = qwarp >> 3u; /* 4 token groups of 8 per round */ + const uint32_t row = row0 + row_in; + const uint32_t avail = count - local_start; + + for (uint32_t round = 0u; round * 32u < avail; round++) { + const uint32_t t0 = local_start + round * 32u + tok_grp * 8u; + uint32_t np = (t0 < count) ? count - t0 : 0u; + if (np > 8u) np = 8u; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = { xq, xq, xq, xq, xq, xq, xq, xq }; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < np) { + pair[p] = sorted_pairs[offsets[expert] + t0 + p]; + xqb[p] = xq + (uint64_t)(pair[p] / n_expert) * xq_blocks; + } + } + const char *sB_gate = sB + (uint64_t)row_in * gate_row_bytes; + const char *sB_up = sB + (uint64_t)(8u + row_in) * gate_row_bytes; + const uint64_t 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 *)(sB_gate + (uint64_t)b * chunk_bytes); + const cuda_block_mxfp4 *ub = (const cuda_block_mxfp4 *)(sB_up + (uint64_t)b * chunk_bytes); + dev_dot_mxfp4_q8_K_block8(gb, xqb[0] + b, xqb[1] + b, xqb[2] + b, xqb[3] + b, + xqb[4] + b, xqb[5] + b, xqb[6] + b, xqb[7] + b, np, gate); + dev_dot_mxfp4_q8_K_block8(ub, xqb[0] + b, xqb[1] + b, xqb[2] + b, xqb[3] + b, + xqb[4] + b, xqb[5] + b, xqb[6] + b, xqb[7] + b, np, up); + } + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p >= np) continue; + 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[pair[p]]; + } + } + } +} + +/* MXFP4 prefill gate/up tile32 kernel: each quarter-warp owns one output + * row and FOUR 8-token groups (32 tokens), so the 17-byte weight chunks it + * loads per q8_K block are reused across 32 tokens instead of 8. At + * balanced routing (128 tokens/expert) this cuts the expert-weight DRAM + * re-read factor from 16x to 4x; the production tile8 path is bound by + * exactly that traffic (~64 GB per 10.5k-token prefill). + * + * Exactness: identical to the tile8 kernel per (token,row) output - same + * 8-lane split of the q8_K block loop, same dev_dot_mxfp4_q8_K_block8 + * helper per 8-token group, same quarter-warp reduction and epilogue. + * Only the weight-load reuse changes. */ +__global__ static void moe_gate_up_mid_mxfp4_expert_tile32_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) { + const uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t expert = tile_experts[tile]; + const uint32_t count = counts[expert]; + if (max_count != 0u && count >= max_count) return; + const uint32_t local_start = tile_starts[tile]; + 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 chunk_bytes = gate_row_bytes / xq_blocks; + + uint32_t pair[32]; + uint32_t np4[4] = {0u, 0u, 0u, 0u}; + #pragma unroll + for (uint32_t g = 0; g < 4u; g++) { + const uint32_t t0 = local_start + g * 8u; + uint32_t np = (t0 < count) ? count - t0 : 0u; + if (np > 8u) np = 8u; + np4[g] = np; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + pair[g * 8u + p] = (p < np) + ? sorted_pairs[offsets[expert] + t0 + p] + : 0u; + } + } + float gate[32]; + float up[32]; + #pragma unroll + for (uint32_t i = 0; i < 32u; i++) { gate[i] = 0.0f; up[i] = 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 * chunk_bytes); + const cuda_block_mxfp4 *ub = (const cuda_block_mxfp4 *)(up_row + (uint64_t)b * chunk_bytes); + #pragma unroll + for (uint32_t g = 0; g < 4u; g++) { + if (np4[g] == 0u) continue; + const uint32_t tok0 = pair[g * 8u + 0] / n_expert; + const cuda_block_q8_K *xg = xq + (uint64_t)tok0 * xq_blocks; + /* Groups are 8 consecutive sorted pairs of one expert; pairs + * share the expert, but the tokens are arbitrary, so the xq + * base must be per-token. Build the 8 pointers for this + * group (registers reused across the b loop). */ + const cuda_block_q8_K *xp[8]; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + xp[p] = (p < np4[g]) + ? xq + (uint64_t)(pair[g * 8u + p] / n_expert) * xq_blocks + b + : xg + b; + } + dev_dot_mxfp4_q8_K_block8(gb, xp[0], xp[1], xp[2], xp[3], + xp[4], xp[5], xp[6], xp[7], np4[g], gate + g * 8u); + dev_dot_mxfp4_q8_K_block8(ub, xp[0], xp[1], xp[2], xp[3], + xp[4], xp[5], xp[6], xp[7], np4[g], up + g * 8u); + } + } + #pragma unroll + for (uint32_t g = 0; g < 4u; g++) { + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p >= np4[g]) continue; + const uint32_t i = g * 8u + p; + gate[i] = quarter_warp_sum_f32(gate[i], lane); + up[i] = quarter_warp_sum_f32(up[i], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[i] > clamp) gate[i] = clamp; + if (up[i] > clamp) up[i] = clamp; + if (up[i] < -clamp) up[i] = -clamp; + } + const uint64_t off = (uint64_t)pair[i] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[i]; + up_out[off] = up[i]; + } + mid_out[off] = (gate[i] / (1.0f + expf(-gate[i]))) * up[i] * weights[pair[i]]; + } + } + } +} + __global__ static DS4_ROCM_UNUSED void moe_down_kernel( float *down_out, const char *down_base, @@ -2002,6 +2764,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 +3101,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 +3336,107 @@ __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 n_expert, + uint32_t atomic_out, + uint32_t row_groups) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + /* Aligned quant slices plus separate scales, as in the gate/up tile + * kernel, so the dot helper loads b128 from LDS. */ + __shared__ int4 sxq4[8u * 8u * 16u]; + __shared__ float sds[8u * 8u]; + int8_t *sqs = (int8_t *)sxq4; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + /* Fixed-count predicated setup keeps pair/xqb in registers; the tail + * xqb entries stay at valid dummy rows and are never read past np. */ + const cuda_block_q8_K *xqb[8] = { midq, midq, midq, midq, midq, midq, midq, midq }; + uint32_t count = counts[expert]; + uint32_t np = count - local_start; + if (np > 8u) np = 8u; + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p < np) { + pair[p] = sorted_pairs[offsets[expert] + local_start + p]; + xqb[p] = midq + (uint64_t)pair[p] * midq_blocks; + } + } + const int staged = midq_blocks <= 8u; + if (staged) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks * 16u; i += blockDim.x) { + const uint32_t u = i / (midq_blocks * 16u); + const uint32_t rem = i - u * (midq_blocks * 16u); + const uint32_t b = rem >> 4u; + const uint32_t k = rem & 15u; + const uint32_t sp = sorted_pairs[offsets[expert] + local_start + u]; + const int32_t *src = (const int32_t *) + midq[(uint64_t)sp * midq_blocks + b].qs + k * 4u; + ((int4 *)sqs)[(u * midq_blocks + b) * 16u + k] = + make_int4(src[0], src[1], src[2], src[3]); + } + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + const uint32_t u = i / midq_blocks; + const uint32_t b = i - u * midq_blocks; + const uint32_t sp = sorted_pairs[offsets[expert] + local_start + u]; + sds[i] = midq[(uint64_t)sp * midq_blocks + b].d; + } + __syncthreads(); + } + /* Each block covers row_groups consecutive 32-row groups against one + * staged copy of the tile's activations: activation staging traffic + * and block-launch count both divide by row_groups. Per (token,row) + * the arithmetic is untouched (same helper, same lane-strided block + * loop, same reduction), so results are bit-identical to + * row_groups=1. */ + for (uint32_t rg = 0u; rg < row_groups; rg++) { + uint32_t row = (blockIdx.x * row_groups + rg) * 32u + (threadIdx.x >> 3u); + if (row >= out_dim) continue; + 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}; + const uint64_t sq = (uint64_t)midq_blocks << 8u; + 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); + if (staged) { + const int8_t *qb = sqs + ((uint64_t)b << 8u); + dev_dot_mxfp4_q8a_block8(wb, + qb, qb + sq, qb + 2u * sq, qb + 3u * sq, + qb + 4u * sq, qb + 5u * sq, qb + 6u * sq, qb + 7u * sq, + sds[b], sds[midq_blocks + b], sds[2u * midq_blocks + b], sds[3u * midq_blocks + b], + sds[4u * midq_blocks + b], sds[5u * midq_blocks + b], sds[6u * midq_blocks + b], sds[7u * midq_blocks + b], + np, acc); + } else { + dev_dot_mxfp4_q8_K_block8(wb, xqb[0] + b, xqb[1] + b, xqb[2] + b, xqb[3] + b, + xqb[4] + b, xqb[5] + b, xqb[6] + b, xqb[7] + b, np, acc); + } + } + #pragma unroll + for (uint32_t p = 0; p < 8u; p++) { + if (p >= np) continue; + 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 51fafe3d38..317181ecdf 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,61 @@ 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; + const uint32_t use_mxfp4_ldsB = + mxfp4_path && use_expert_tiles && n_tokens >= 128u && + 16u * gate_row_bytes <= 60u * 1024u && + (expert_mid_dim % 8u) == 0u && + getenv("DS4_ROCM_ENABLE_MXFP4_LDSB") != NULL; + const uint32_t use_mxfp4_tile32 = + mxfp4_path && use_expert_tiles && n_tokens >= 32u && + (expert_mid_dim % 32u) == 0u && + getenv("DS4_ROCM_ENABLE_MXFP4_TILE32") != NULL; + /* Occupancy variants (bit-exact; see kernel comments): tile4 cuts + * the staged-activation LDS to ~29 KiB (4 blocks/WGP), row64 runs + * 512 threads per block (2 blocks/WGP = 32 resident warps). Both + * target the tile8 kernel's latency-bound profile. Default off. */ + const uint32_t use_mxfp4_tile4 = + mxfp4_path && use_expert_tiles && !use_mxfp4_tile32 && !use_mxfp4_ldsB && + n_tokens >= 4u && + getenv("DS4_ROCM_ENABLE_MXFP4_TILE4") != NULL; + const uint32_t use_mxfp4_row64 = + mxfp4_path && use_expert_tiles && !use_mxfp4_tile32 && !use_mxfp4_ldsB && + !use_mxfp4_tile4 && n_tokens >= 8u && + getenv("DS4_ROCM_ENABLE_MXFP4_ROW64") != NULL; + uint32_t down_row_groups = 1u; + { + const char *rge = getenv("DS4_ROCM_MXFP4_DOWN_RGROUP"); + if (rge && rge[0]) { + const long rv = strtol(rge, NULL, 10); + if (rv >= 1 && rv <= 8) down_row_groups = (uint32_t)rv; + } + } uint32_t *sorted_pairs = NULL; uint32_t *sorted_offsets = NULL; uint32_t *sorted_counts = NULL; @@ -772,9 +813,29 @@ static int routed_moe_launch( uint32_t *tile16_total = NULL; uint32_t *tile16_experts = NULL; uint32_t *tile16_starts = NULL; + uint32_t *tile128_total = NULL; + uint32_t *tile128_experts = NULL; + uint32_t *tile128_starts = NULL; + uint32_t *tile32_total = NULL; + uint32_t *tile32_experts = NULL; + uint32_t *tile32_starts = NULL; + uint32_t *tile4_total = NULL; + uint32_t *tile4_experts = NULL; + uint32_t *tile4_starts = NULL; uint32_t *iq2_gate_hot_dev = NULL; uint32_t tile_capacity = 0; uint32_t tile16_capacity = 0; + uint32_t tile128_capacity = 0; + uint32_t tile32_capacity = 0; + uint32_t tile4_capacity = 0; + if (getenv("DS4_ROCM_MOE_PATH_DEBUG") != NULL) { + fprintf(stderr, + "ds4: moe path n=%u mxfp4=%d stream_full=%d full_cached=%d " + "batch_stream=%d split=%d compact=%d sorted=%u tiles=%u\n", + n_tokens, mxfp4_path, stream_full_layer, full_table_cached, + batch_stream_selected, batch_stream_split_selected, + compact_selected, use_sorted_pairs, use_expert_tiles); + } dim3 xq_grid(xq_blocks, n_tokens, 1); q8_K_quantize_kernel<<>>(xq, (const float *)x->ptr, expert_in_dim, n_tokens); ok = cuda_ok(cudaGetLastError(), "routed_moe x quantize launch"); @@ -897,6 +958,9 @@ static int routed_moe_launch( const uint64_t sorted_bytes = (uint64_t)pair_count * sizeof(uint32_t); tile_capacity = (pair_count + expert_tile_m - 1u) / expert_tile_m + bucket_count; tile16_capacity = use_down_tile16 ? ((pair_count + 15u) / 16u + bucket_count) : 0u; + tile128_capacity = use_mxfp4_ldsB ? ((pair_count + 127u) / 128u + bucket_count) : 0u; + tile32_capacity = use_mxfp4_tile32 ? ((pair_count + 31u) / 32u + bucket_count) : 0u; + tile4_capacity = use_mxfp4_tile4 ? ((pair_count + 3u) / 4u + bucket_count) : 0u; const uint64_t tile_offsets_bytes = (uint64_t)(bucket_count + 1u) * sizeof(uint32_t); const uint64_t tile_total_bytes = sizeof(uint32_t); const uint64_t tile_experts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t); @@ -915,7 +979,31 @@ static int routed_moe_launch( const uint64_t tile16_starts_off = tile16_experts_off + tile16_experts_bytes; const uint64_t iq2_gate_hot_off = tile16_starts_off + tile16_starts_bytes; const uint64_t iq2_gate_hot_bytes = (uint64_t)bucket_count * sizeof(uint32_t); - const uint64_t scratch_bytes = iq2_gate_hot_off + iq2_gate_hot_bytes; + const uint64_t tile128_offsets_off = iq2_gate_hot_off + iq2_gate_hot_bytes; + const uint64_t tile128_offsets_bytes = use_mxfp4_ldsB ? (uint64_t)(bucket_count + 1u) * sizeof(uint32_t) : 0u; + const uint64_t tile128_total_off = tile128_offsets_off + tile128_offsets_bytes; + const uint64_t tile128_total_bytes = use_mxfp4_ldsB ? sizeof(uint32_t) : 0u; + const uint64_t tile128_experts_off = tile128_total_off + tile128_total_bytes; + const uint64_t tile128_experts_bytes = (uint64_t)tile128_capacity * sizeof(uint32_t); + const uint64_t tile128_starts_off = tile128_experts_off + tile128_experts_bytes; + const uint64_t tile128_starts_bytes = (uint64_t)tile128_capacity * sizeof(uint32_t); + const uint64_t tile32_offsets_off = tile128_starts_off + tile128_starts_bytes; + const uint64_t tile32_offsets_bytes = use_mxfp4_tile32 ? (uint64_t)(bucket_count + 1u) * sizeof(uint32_t) : 0u; + const uint64_t tile32_total_off = tile32_offsets_off + tile32_offsets_bytes; + const uint64_t tile32_total_bytes = use_mxfp4_tile32 ? sizeof(uint32_t) : 0u; + const uint64_t tile32_experts_off = tile32_total_off + tile32_total_bytes; + const uint64_t tile32_experts_bytes = (uint64_t)tile32_capacity * sizeof(uint32_t); + const uint64_t tile32_starts_off = tile32_experts_off + tile32_experts_bytes; + const uint64_t tile32_starts_bytes = (uint64_t)tile32_capacity * sizeof(uint32_t); + const uint64_t tile4_offsets_off = tile32_starts_off + tile32_starts_bytes; + const uint64_t tile4_offsets_bytes = use_mxfp4_tile4 ? (uint64_t)(bucket_count + 1u) * sizeof(uint32_t) : 0u; + const uint64_t tile4_total_off = tile4_offsets_off + tile4_offsets_bytes; + const uint64_t tile4_total_bytes = use_mxfp4_tile4 ? sizeof(uint32_t) : 0u; + const uint64_t tile4_experts_off = tile4_total_off + tile4_total_bytes; + const uint64_t tile4_experts_bytes = (uint64_t)tile4_capacity * sizeof(uint32_t); + const uint64_t tile4_starts_off = tile4_experts_off + tile4_experts_bytes; + const uint64_t tile4_starts_bytes = (uint64_t)tile4_capacity * sizeof(uint32_t); + const uint64_t scratch_bytes = tile4_starts_off + tile4_starts_bytes; uint8_t *scratch = (uint8_t *)cuda_tmp_alloc(scratch_bytes, "routed_moe sorted pairs"); if (!scratch) { @@ -936,6 +1024,18 @@ static int routed_moe_launch( tile16_experts = use_down_tile16 ? (uint32_t *)(scratch + tile16_experts_off) : NULL; tile16_starts = use_down_tile16 ? (uint32_t *)(scratch + tile16_starts_off) : NULL; iq2_gate_hot_dev = (uint32_t *)(scratch + iq2_gate_hot_off); + uint32_t *tile128_offsets = use_mxfp4_ldsB ? (uint32_t *)(scratch + tile128_offsets_off) : NULL; + tile128_total = use_mxfp4_ldsB ? (uint32_t *)(scratch + tile128_total_off) : NULL; + tile128_experts = use_mxfp4_ldsB ? (uint32_t *)(scratch + tile128_experts_off) : NULL; + tile128_starts = use_mxfp4_ldsB ? (uint32_t *)(scratch + tile128_starts_off) : NULL; + uint32_t *tile32_offsets = use_mxfp4_tile32 ? (uint32_t *)(scratch + tile32_offsets_off) : NULL; + tile32_total = use_mxfp4_tile32 ? (uint32_t *)(scratch + tile32_total_off) : NULL; + tile32_experts = use_mxfp4_tile32 ? (uint32_t *)(scratch + tile32_experts_off) : NULL; + tile32_starts = use_mxfp4_tile32 ? (uint32_t *)(scratch + tile32_starts_off) : NULL; + uint32_t *tile4_offsets = use_mxfp4_tile4 ? (uint32_t *)(scratch + tile4_offsets_off) : NULL; + tile4_total = use_mxfp4_tile4 ? (uint32_t *)(scratch + tile4_total_off) : NULL; + tile4_experts = use_mxfp4_tile4 ? (uint32_t *)(scratch + tile4_experts_off) : NULL; + tile4_starts = use_mxfp4_tile4 ? (uint32_t *)(scratch + tile4_starts_off) : NULL; ok = cuda_ok(cudaMemset(counts, 0, counts_bytes), "routed_moe sorted counts clear"); if (ok) { moe_count_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( @@ -962,6 +1062,33 @@ static int routed_moe_launch( moe_build_expert_tile_offsets_kernel<<<1, 1>>>(tile_offsets, tile_total, counts, expert_tile_m, bucket_count); ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile offsets launch"); } + if (ok && use_mxfp4_ldsB) { + moe_build_expert_tile_offsets_kernel<<<1, 1, 0, 0 /* default stream */>>>(tile128_offsets, tile128_total, counts, 128u, bucket_count); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile128 offsets launch"); + } + if (ok && use_mxfp4_ldsB) { + moe_build_expert_tiles_kernel<<<(bucket_count + 255u) / 256u, 256, 0, 0 /* default stream */>>>( + tile128_experts, tile128_starts, tile128_offsets, counts, 128u, bucket_count); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile128 build launch"); + } + if (ok && use_mxfp4_tile32) { + moe_build_expert_tile_offsets_kernel<<<1, 1, 0, 0 /* default stream */>>>(tile32_offsets, tile32_total, counts, 32u, bucket_count); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile32 offsets launch"); + } + if (ok && use_mxfp4_tile32) { + moe_build_expert_tiles_kernel<<<(bucket_count + 255u) / 256u, 256, 0, 0 /* default stream */>>>( + tile32_experts, tile32_starts, tile32_offsets, counts, 32u, bucket_count); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile32 build launch"); + } + if (ok && use_mxfp4_tile4) { + moe_build_expert_tile_offsets_kernel<<<1, 1, 0, 0 /* default stream */>>>(tile4_offsets, tile4_total, counts, 4u, bucket_count); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile4 offsets launch"); + } + if (ok && use_mxfp4_tile4) { + moe_build_expert_tiles_kernel<<<(bucket_count + 255u) / 256u, 256, 0, 0 /* default stream */>>>( + tile4_experts, tile4_starts, tile4_offsets, counts, 4u, bucket_count); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile4 build launch"); + } if (ok && use_expert_tiles) { moe_build_expert_tiles_kernel<<<(bucket_count + 255u) / 256u, 256>>>( tile_experts, tile_starts, tile_offsets, counts, expert_tile_m, bucket_count); @@ -1116,6 +1243,14 @@ static int routed_moe_launch( } } if (ok && !split_gateup_done) { + if (getenv("DS4_ROCM_MOE_PATH_DEBUG") != NULL) { + fprintf(stderr, + "ds4: moe gate/up launch check: sorted_pairs=%p offsets=%p counts=%p " + "tile_total=%p tile_experts=%p tile_starts=%p\n", + (void *)sorted_pairs, (void *)sorted_offsets, + (void *)sorted_counts, (void *)tile_total, + (void *)tile_experts, (void *)tile_starts); + } dim3 mgrid((expert_mid_dim + 31u) / 32u, pair_count, 1); if (ok && sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && tile_total && tile_experts && tile_starts) { if (q4k_path) { @@ -1135,6 +1270,112 @@ 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) { + if (use_mxfp4_tile32 && tile32_total && tile32_experts && tile32_starts) { + dim3 t32grid((expert_mid_dim + 31u) / 32u, tile32_capacity, 1); + moe_gate_up_mid_mxfp4_expert_tile32_row32_kernel<<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile32_total, tile32_experts, tile32_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_mxfp4_ldsB && tile128_total && tile128_experts && tile128_starts) { + /* B-staged prefill path: 8 output rows x up to 128 + * tokens per block; expert weights read ~once. */ + const uint32_t ldsB_shmem = 16u * (uint32_t)gate_row_bytes; + if (ldsB_shmem > 48u * 1024u) { + static int ldsB_shmem_attr_set = 0; + if (!ldsB_shmem_attr_set) { + const cudaError_t attr_err = cudaFuncSetAttribute( + (const void *)moe_gate_up_mid_mxfp4_expert_row8_ldsB_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 16u * 3808u); + if (attr_err != cudaSuccess) { + (void)cudaGetLastError(); + return 0; + } + ldsB_shmem_attr_set = 1; + } + } + dim3 bgrid(expert_mid_dim / 8u, tile128_capacity, 1); + moe_gate_up_mid_mxfp4_expert_row8_ldsB_kernel<<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile128_total, tile128_experts, tile128_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_mxfp4_tile4) { + const uint32_t tile4_shmem = xq_blocks <= 28u ? + 4u * xq_blocks * (256u + (uint32_t)sizeof(float)) : 0u; + dim3 tgrid(tile4_capacity, (expert_mid_dim + 31u) / 32u, 1); + moe_gate_up_mid_mxfp4_expert_tileN_kernel<4u, 32u><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile4_total, tile4_experts, tile4_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_mxfp4_row64) { + const uint32_t row64_shmem = xq_blocks <= 28u ? + 8u * xq_blocks * (256u + (uint32_t)sizeof(float)) : 0u; + if (row64_shmem > 48u * 1024u) { + static int row64_shmem_attr_set = 0; + if (!row64_shmem_attr_set) { + /* Hoisted: the compat macro cannot see a + * comma inside a template-id argument. */ + static const void *row64_fn = + (const void *)&moe_gate_up_mid_mxfp4_expert_tileN_kernel<8u, 64u>; + const cudaError_t attr_err = cudaFuncSetAttribute( + row64_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 8u * 28u * 260u); + if (attr_err != cudaSuccess) { + (void)cudaGetLastError(); + return 0; + } + row64_shmem_attr_set = 1; + } + } + dim3 tgrid(tile_capacity, (expert_mid_dim + 63u) / 64u, 1); + moe_gate_up_mid_mxfp4_expert_tileN_kernel<8u, 64u><<>>( + (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 { + dim3 tgrid(tile_capacity, (expert_mid_dim + 31u) / 32u, 1); + if (getenv("DS4_ROCM_MOE_PATH_DEBUG") != NULL) { + fprintf(stderr, + "ds4: moe mxfp4 gate/up tile8 launch grid=(%u,%u) mid_dim=%u tile_capacity=%u xq_blocks=%u\n", + tgrid.x, tgrid.y, expert_mid_dim, tile_capacity, xq_blocks); + } + /* LDS staging of 8 activation rows as aligned quant + * slices plus scales, sized to the actual xq_blocks. + * Raised to 28 blocks (7168-dim x at Q8_K=256): 56.9 KiB + * needs the >48 KiB dynamic-smem opt-in, set once. */ + const uint32_t tile8_shmem = xq_blocks <= 28u ? + 8u * xq_blocks * (256u + (uint32_t)sizeof(float)) : 0u; + if (tile8_shmem > 48u * 1024u) { + static int tile8_shmem_attr_set = 0; + if (!tile8_shmem_attr_set) { + const cudaError_t attr_err = cudaFuncSetAttribute( + (const void *)moe_gate_up_mid_mxfp4_expert_tile8_row32_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 8u * 28u * 260u); + if (attr_err != cudaSuccess) { + (void)cudaGetLastError(); + return 0; + } + tile8_shmem_attr_set = 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 +1493,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 +1712,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 +1773,15 @@ 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) { + const uint32_t row_blocks = (out_dim + 31u) / 32u; + dim3 tgrid((row_blocks + down_row_groups - 1u) / down_row_groups, 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, n_expert, 0u, down_row_groups); } else if (use_down_row2048) { if (down_row_span == 512u) { dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); diff --git a/tests/bench_mxfp4_rocm.c b/tests/bench_mxfp4_rocm.c new file mode 100644 index 0000000000..7090869428 --- /dev/null +++ b/tests/bench_mxfp4_rocm.c @@ -0,0 +1,221 @@ +/* MXFP4 routed-MoE prefill benchmark (ROCm). + * + * Times ds4_gpu_routed_moe_batch_tensor at production prefill shapes: + * n_tokens=4096, 256 experts, top-8, model_dim=7168, ffn_dim=2048. + * Weights are synthetic (kernel cost is data-independent); selection is + * uniform-random with a fixed seed, matching the production ~128 + * tokens/expert distribution. + * + * Not a correctness harness — tests/test_mxfp4_rocm.c owns exactness. + * This exists so kernel retiling work has a seconds-scale iteration loop + * instead of 75-second distributed probes. + * + * Usage: tests/bench_mxfp4_rocm [n_tokens] [iters] + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ds4_gpu.h" + +#define MXFP4_TYPE 39u +#define QK_MXFP4 32u + +#ifndef MODEL_DIM +#define MODEL_DIM 7168u +#endif +#ifndef FFN_DIM +#define FFN_DIM 2048u +#endif +#ifndef N_TOTAL_EXPERT +#define N_TOTAL_EXPERT 256u +#endif +#ifndef N_EXPERT +#define N_EXPERT 8u +#endif +#define CLAMP 7.0f + +typedef __attribute__((__may_alias__)) struct { + uint8_t e; + uint8_t qs[16]; +} block_mxfp4; + +static uint64_t align_up_u64(uint64_t v, uint64_t a) { + return (v + a - 1u) / a * a; +} + +static double now_sec(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9; +} + +static void fill_matrix(block_mxfp4 *m, uint32_t rows, uint32_t cols, uint32_t seed) { + /* Cheap deterministic fill: any nibble pattern exercises the same + * kernel cost; keep the e8m0 scale nonzero to avoid denormal edges. */ + uint64_t state = seed; + const uint64_t blocks = (uint64_t)rows * (cols / QK_MXFP4); + for (uint64_t i = 0; i < blocks; i++) { + state = state * 6364136223846793005ull + 1442695040888963407ull; + m[i].e = (uint8_t)(126u + (state >> 59u)); + for (uint32_t k = 0; k < 16u; k++) { + state = state * 6364136223846793005ull + 1442695040888963407ull; + m[i].qs[k] = (uint8_t)(state >> 56u); + } + } +} + +int main(int argc, char **argv) { + const uint32_t n_tokens = argc > 1 ? (uint32_t)strtoul(argv[1], NULL, 10) : 4096u; + const int iters = argc > 2 ? (int)strtol(argv[2], NULL, 10) : 6; + + 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_u64(gate_tensor_bytes, 4096u); + const uint64_t down_offset = align_up_u64(up_offset + gate_tensor_bytes, 4096u); + const uint64_t model_size = align_up_u64(down_offset + down_tensor_bytes, 4096u); + + fprintf(stderr, + "bench: tokens=%u iters=%d model=%.2f GiB (gate/up %.2f GiB each, down %.2f GiB)\n", + n_tokens, iters, (double)model_size / 1073741824.0, + (double)gate_tensor_bytes / 1073741824.0, + (double)down_tensor_bytes / 1073741824.0); + + FILE *model_file = tmpfile(); + void *model = MAP_FAILED; + 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); + } + if (!model_file || model == MAP_FAILED) { + fprintf(stderr, "bench: model image allocation failed\n"); + return 1; + } + 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); + + if (!ds4_gpu_init()) { + fprintf(stderr, "bench: ds4_gpu_init failed\n"); + return 1; + } + 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; + if (!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, + 3u, max_tensor_bytes)) { + fprintf(stderr, "bench: model cache setup failed\n"); + return 1; + } + + 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; + + 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)); + if (!x || !selected || !weights) { + fprintf(stderr, "bench: host allocation failed\n"); + return 1; + } + uint64_t rng = 0x243f6a8885a308d3ull; + for (uint64_t i = 0; i < token_x_count; i++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + x[i] = ((double)(rng >> 40u) / (double)(1ull << 24u) - 0.5) * 0.03125; + } + /* Uniform top-8 without replacement, fixed seed: ~n_tokens*8/256 + * tokens per expert, matching production routing balance. */ + for (uint32_t t = 0; t < n_tokens; t++) { + uint32_t pool[N_TOTAL_EXPERT]; + for (uint32_t e = 0; e < N_TOTAL_EXPERT; e++) pool[e] = e; + for (uint32_t s = 0; s < N_EXPERT; s++) { + rng = rng * 6364136223846793005ull + 1442695040888963407ull; + const uint32_t j = s + (uint32_t)((rng >> 33u) % (N_TOTAL_EXPERT - s)); + const uint32_t tmp = pool[s]; pool[s] = pool[j]; pool[j] = tmp; + selected[(uint64_t)t * N_EXPERT + s] = (int32_t)pool[s]; + weights[(uint64_t)t * N_EXPERT + s] = 1.0f / (float)N_EXPERT; + } + } + + ds4_gpu_tensor *x_t = ds4_gpu_tensor_alloc(token_x_count * sizeof(float)); + ds4_gpu_tensor *sel_t = ds4_gpu_tensor_alloc(route_count * sizeof(int32_t)); + ds4_gpu_tensor *w_t = ds4_gpu_tensor_alloc(route_count * sizeof(float)); + ds4_gpu_tensor *gate_t = ds4_gpu_tensor_alloc(mid_count * sizeof(float)); + ds4_gpu_tensor *up_t = ds4_gpu_tensor_alloc(mid_count * sizeof(float)); + ds4_gpu_tensor *mid_t = ds4_gpu_tensor_alloc(mid_count * sizeof(float)); + ds4_gpu_tensor *down_t = ds4_gpu_tensor_alloc(down_count * sizeof(float)); + ds4_gpu_tensor *out_t = ds4_gpu_tensor_alloc(token_x_count * sizeof(float)); + if (!x_t || !sel_t || !w_t || !gate_t || !up_t || !mid_t || !down_t || !out_t) { + fprintf(stderr, "bench: tensor allocation failed\n"); + return 1; + } + if (!ds4_gpu_tensor_write(x_t, 0u, x, token_x_count * sizeof(float)) || + !ds4_gpu_tensor_write(sel_t, 0u, selected, route_count * sizeof(int32_t)) || + !ds4_gpu_tensor_write(w_t, 0u, weights, route_count * sizeof(float))) { + fprintf(stderr, "bench: tensor upload failed\n"); + return 1; + } + + double best = 1e30; + double total = 0.0; + int done = 0; + for (int it = 0; it < iters; it++) { + bool mid_is_f16 = false; + const double t0 = now_sec(); + const int ok = ds4_gpu_routed_moe_batch_tensor( + out_t, gate_t, up_t, mid_t, down_t, + 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, + sel_t, w_t, N_TOTAL_EXPERT, N_EXPERT, + CLAMP, x_t, 0u, n_tokens, &mid_is_f16, true); + if (!ok) { + fprintf(stderr, "bench: iteration %d launch failed\n", it); + return 1; + } + if (!ds4_gpu_synchronize()) { + fprintf(stderr, "bench: iteration %d sync failed\n", it); + return 1; + } + const double dt = now_sec() - t0; + if (it > 0) { /* first iteration absorbs plan/scratch warmup */ + if (dt < best) best = dt; + total += dt; + done++; + } + fprintf(stderr, "bench: iter %d %.1f ms\n", it, dt * 1000.0); + } + const double unique_bytes = + (double)gate_tensor_bytes * 2.0 + (double)down_tensor_bytes; + const double avg = total / (double)(done > 0 ? done : 1); + fprintf(stderr, + "bench: best %.1f ms avg %.1f ms | unique weights %.2f GiB -> %.1f GB/s (best)\n", + best * 1000.0, avg * 1000.0, + unique_bytes / 1073741824.0, + unique_bytes / best / 1e9); + + ds4_gpu_cleanup(); + return 0; +} diff --git a/tests/test_mxfp4_rocm.c b/tests/test_mxfp4_rocm.c new file mode 100644 index 0000000000..db3727108a --- /dev/null +++ b/tests/test_mxfp4_rocm.c @@ -0,0 +1,655 @@ +/* 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 +/* Default to the DS4 Flash routed-expert shape: model_dim 4096 means 16 + * Q8_K activation chunks per token, which is what the resident tile and + * decode kernels see in production. Smaller synthetic dims miss the + * no-staging and multi-chunk code paths entirely. */ +#ifndef N_TOTAL_EXPERT +#define N_TOTAL_EXPERT 256u +#endif +#ifndef N_EXPERT +#define N_EXPERT 6u +#endif +#ifndef MODEL_DIM +#define MODEL_DIM 4096u +#endif +#ifndef FFN_DIM +#define FFN_DIM 2048u +#endif +#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]; +} 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) { + 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); + } +} + +static int compare_repeated(const char *name, + const float *actual, + uint32_t n_tokens, + uint32_t token_elems, + const reference_pattern patterns[N_PATTERN], + 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 = patterns[token % N_PATTERN].mid; + 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; +} + +/* The GPU quantizes its own gate/up mid before the down projection, and a + * one-LSB rounding flip against the analytic reference mid is legitimate: + * both mids sit within float tolerance of each other, yet one flipped + * activation moves a down dot by up to d * 2^(e-127) * 6, which real-shape + * scales push past any fixed tolerance. Judge the down stage on the GPU's + * own mid instead: quantize it with the identical CPU algorithm and require + * the down kernels to match that reference tightly. Expected rows are + * cached per pattern; repeated tokens reuse them after a bitwise mid check. */ +static int check_out_from_gpu_mid(const float *out_actual, + const float *mid_actual, + uint32_t n_tokens, + const reference_pattern patterns[N_PATTERN], + const block_mxfp4 *down_matrix, + float abs_tolerance, + float rel_tolerance) { + static ref_block_q8_K midq[N_EXPERT][FFN_DIM / QK_K]; + static float expected[N_PATTERN][MODEL_DIM]; + int32_t cached_token[N_PATTERN] = { -1, -1, -1, -1 }; + 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 * MODEL_DIM; + + for (uint32_t token = 0; token < n_tokens; token++) { + const uint32_t p = token % N_PATTERN; + const reference_pattern *pattern = &patterns[p]; + const float *mid = mid_actual + (uint64_t)token * N_EXPERT * FFN_DIM; + const uint64_t mid_bytes = (uint64_t)N_EXPERT * FFN_DIM * sizeof(float); + if (cached_token[p] < 0 || + memcmp(mid_actual + (uint64_t)cached_token[p] * N_EXPERT * FFN_DIM, + mid, mid_bytes) != 0) { + for (uint32_t slot = 0; slot < N_EXPERT; slot++) { + quantize_q8_K(midq[slot], mid + (uint64_t)slot * FFN_DIM, FFN_DIM); + } + for (uint32_t row = 0; row < MODEL_DIM; row++) { + float want = 0.0f; + for (uint32_t slot = 0; slot < N_EXPERT; slot++) { + const uint32_t expert = (uint32_t)pattern->selected[slot]; + want += dot_mxfp4_q8_K( + matrix_row(down_matrix, expert, row, MODEL_DIM, FFN_DIM), + midq[slot], FFN_DIM); + } + expected[p][row] = want; + } + cached_token[p] = (int32_t)token; + } + for (uint32_t i = 0; i < MODEL_DIM; i++) { + const uint64_t index = (uint64_t)token * MODEL_DIM + i; + const float got = out_actual[index]; + const float want = expected[p][i]; + if (!isfinite(got) || !isfinite(want)) { + fprintf(stderr, + "MXFP4 ROCm tokens=%u out non-finite at token=%u element=%u " + "got=%g expected=%g\n", + n_tokens, 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 out max_abs=%-10g at=%llu " + "max_tol_ratio=%g at=%llu failures=%llu/%llu\n", + n_tokens, 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], + float *snapshot_out, + const float *expect_out) { + 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, + 1.0e-4f, 1.0e-4f); + const int out_ok = check_out_from_gpu_mid( + out_actual, mid_actual, n_tokens, patterns, + (const block_mxfp4 *)((const char *)model + down_offset), + 2.0e-4f, 1.0e-4f); + ok = mid_ok && out_ok; + } + if (ok && snapshot_out) { + memcpy(snapshot_out, out_actual, out_count * sizeof(float)); + } + if (ok && expect_out && + memcmp(expect_out, out_actual, out_count * sizeof(float)) != 0) { + fprintf(stderr, + "MXFP4 ROCm tokens=%u bitwise mismatch vs default tile path\n", + n_tokens); + ok = 0; + } + + 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); + for (uint32_t p = 0; p < N_PATTERN; p++) { + build_reference(&patterns[p], gate_matrix, up_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; + } + + float *snap128 = (float *)malloc(128u * MODEL_DIM * sizeof(float)); + float *snap512 = (float *)malloc(512u * MODEL_DIM * sizeof(float)); + ok = ok && snap128 && snap512; + for (uint32_t i = 0; ok && i < sizeof(token_cases) / sizeof(token_cases[0]); i++) { + float *snap = token_cases[i] == 128u ? snap128 : + token_cases[i] == 512u ? snap512 : NULL; + 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, + snap, NULL)) { + ok = 0; + } + } + /* Occupancy/path variants must reproduce the default tile path bit + * for bit: same accumulation order by construction, so the out + * tensors compare with memcmp, not a tolerance. */ + if (ok) { + const struct { + const char *name; + const char *value; + } variant_envs[] = { + { "DS4_ROCM_ENABLE_MXFP4_TILE4", "1" }, + { "DS4_ROCM_ENABLE_MXFP4_ROW64", "1" }, + { "DS4_ROCM_MXFP4_DOWN_RGROUP", "4" }, + }; + for (uint32_t v = 0; v < sizeof(variant_envs) / sizeof(variant_envs[0]); v++) { + setenv(variant_envs[v].name, variant_envs[v].value, 1); + const int vok = + run_case(128u, model, model_size, + gate_offset, up_offset, down_offset, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, patterns, + NULL, snap128) && + run_case(512u, model, model_size, + gate_offset, up_offset, down_offset, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, patterns, + NULL, snap512); + unsetenv(variant_envs[v].name); + fprintf(stderr, "MXFP4 ROCm variant %s=%s: %s\n", + variant_envs[v].name, variant_envs[v].value, + vok ? "bitwise OK" : "MISMATCH"); + if (!vok) ok = 0; + } + } + free(snap128); + free(snap512); + +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; +}