From aa9dbb3a39104fa4f753590409011063a03cd83b Mon Sep 17 00:00:00 2001 From: hy-user Date: Fri, 24 Jul 2026 03:54:12 -0700 Subject: [PATCH] fix(cuda): support GLM IQ2_XXS routed down projection --- Makefile | 11 +- ds4.c | 10 ++ ds4.h | 3 + ds4_cuda.cu | 87 ++++++++++- tests/test_cuda_iq2_routed.c | 288 +++++++++++++++++++++++++++++++++++ 5 files changed, 397 insertions(+), 2 deletions(-) create mode 100644 tests/test_cuda_iq2_routed.c diff --git a/Makefile b/Makefile index 81d4881b2..53ba78499 100644 --- a/Makefile +++ b/Makefile @@ -324,6 +324,15 @@ tests/test_sampling.o: tests/test_sampling.c ds4.h tests/test_sampling: tests/test_sampling.o ds4_cuda_test_hooks.o ds4_gpu_args.o ds4_kvstore.o rax.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) +tests/test_cuda_iq2_routed.o: tests/test_cuda_iq2_routed.c ds4.h ds4_gpu.h + $(CC) $(CFLAGS) -DDS4_TEST_HOOKS -I. -I$(CUDA_HOME)/include -c -o $@ $< + +tests/test_cuda_iq2_routed: tests/test_cuda_iq2_routed.o ds4_cuda_test_hooks.o ds4_gpu_args.o ds4_kvstore.o rax.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o + $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) + +test-cuda-iq2-routed: tests/test_cuda_iq2_routed + ./tests/test_cuda_iq2_routed + tests/test_cuda_session_batch.o: tests/test_cuda_session_batch.c ds4.h ds4_gpu_args.h ds4_gpu_mgpu.h $(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $< @@ -402,4 +411,4 @@ q4k-dot-test: tests/test_q4k_dot.c ./tests/test_q4k_dot clean: - rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official tests/test_q4k_dot tests/test_metal_session_batch tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o + rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official tests/test_q4k_dot tests/test_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_iq2_routed 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/ds4.c b/ds4.c index 1a79a01e4..17649d7e5 100644 --- a/ds4.c +++ b/ds4.c @@ -3808,6 +3808,16 @@ static float ds4_vec_dot_iq2_xxs_f32(int n, const block_iq2_xxs *x, const float return sumf; } +#ifdef DS4_TEST_HOOKS +float ds4_test_iq2_xxs_dot_f32(uint32_t n, const void *blocks, const float *values) { + if (!blocks || !values || n == 0 || n % QK_K != 0 || n > INT_MAX) { + return NAN; + } + return ds4_vec_dot_iq2_xxs_f32((int)n, + (const block_iq2_xxs *)blocks, + values); +} +#endif static void ds4_vec_dot_q8_K_q8_K(int n, float *s, const block_q8_K *x, diff --git a/ds4.h b/ds4.h index 42bc11660..272c313b2 100644 --- a/ds4.h +++ b/ds4.h @@ -371,6 +371,9 @@ int ds4_test_sample_logits(const float *logits, uint32_t n_vocab, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, float *prob_scratch); +float ds4_test_iq2_xxs_dot_f32(uint32_t n, + const void *blocks, + const float *values); uint64_t ds4_test_mixed_native_count(void); #endif int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k); diff --git a/ds4_cuda.cu b/ds4_cuda.cu index aaa4df113..ba5d7a266 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -20788,6 +20788,45 @@ __global__ static void moe_down_f32_kernel( } if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0]; } +__global__ static void moe_down_iq2_xxs_f32_kernel( + float *down_out, + const char *down_base, + const float *mid, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_mid_dim, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t row = blockIdx.x; + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + 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 nb = expert_mid_dim / CUDA_QK_K; + const cuda_block_iq2_xxs *wr = + (const cuda_block_iq2_xxs *)(down_base + + (uint64_t)(uint32_t)expert_i * down_expert_bytes + + (uint64_t)row * down_row_bytes); + const float *xr = mid + (uint64_t)pair * expert_mid_dim; + float acc = 0.0f; + for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) { + acc += dev_iq2_xxs_dot_f32( + wr + b, xr + (uint64_t)b * CUDA_QK_K, 1); + } + __shared__ float partial[256]; + partial[threadIdx.x] = acc; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) { + down_out[(uint64_t)pair * out_dim + row] = partial[0]; + } +} static int routed_moe_launch( ds4_gpu_tensor *out, @@ -20834,7 +20873,9 @@ static int routed_moe_launch( return 0; } const int q4k_path = (gate_type == 12u && down_type == 12u); - if (!q4k_path && (gate_type != 16u || down_type != 10u)) return 0; + const int iq2_xxs_path = (gate_type == 16u && down_type == 16u); + if (!q4k_path && !iq2_xxs_path && + (gate_type != 16u || down_type != 10u)) return 0; /* Q4_K routed-MoE dispatch: * n_tokens == 1 and n_expert == 6: * use_direct_down_sum + moe_gate_up_mid_decode_q4K_qwarp32 @@ -20902,6 +20943,50 @@ static int routed_moe_launch( cuda_resolve_weight_ptr(model_map, down_offset, down_bytes, logical_tier, "moe_down"); if (!gate_w || !up_w || !down_w) return 0; + if (iq2_xxs_path) { + dim3 mgrid(expert_mid_dim, n_tokens * n_expert, 1); + moe_gate_up_mid_f32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + (const float *)x->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + expert_in_dim, + expert_mid_dim, + n_expert, + clamp); + int ok = cuda_ok(cudaGetLastError(), "routed_moe IQ2_XXS gate/up launch"); + if (ok) { + dim3 dgrid(out_dim, n_tokens * n_expert, 1); + moe_down_iq2_xxs_f32_kernel<<>>( + (float *)down->ptr, + down_w, + (const float *)mid->ptr, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + expert_mid_dim, + out_dim, + n_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe IQ2_XXS down launch"); + } + if (ok) { + const uint64_t n = (uint64_t)n_tokens * out_dim; + moe_sum_kernel<<<(n + 255) / 256, 256>>>( + (float *)out->ptr, + (const float *)down->ptr, + out_dim, + n_expert, + n_tokens); + ok = cuda_ok(cudaGetLastError(), "routed_moe IQ2_XXS sum launch"); + } + return ok; + } int ok = 1; const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; diff --git a/tests/test_cuda_iq2_routed.c b/tests/test_cuda_iq2_routed.c new file mode 100644 index 000000000..dec03463d --- /dev/null +++ b/tests/test_cuda_iq2_routed.c @@ -0,0 +1,288 @@ +/* Numerical oracle for the CUDA all-IQ2_XXS routed-MoE path. + * + * Builds a small deterministic two-token, three-expert problem, evaluates it + * through the production CUDA dispatcher, and compares every intermediate and + * final row with the scalar CPU IQ2_XXS implementation from ds4.c. + */ + +#include "ds4.h" +#include "ds4_gpu.h" + +#include +#include +#include +#include +#include + +#define DIM 256u +#define N_EXPERT 3u +#define N_USED 2u +#define N_TOKEN 2u +#define IQ2_BLOCK_BYTES 66u +#define IQ2_TYPE 16u + +typedef struct { + uint16_t d; + uint16_t qs[DIM / 8u]; +} test_iq2_block; + +typedef char test_iq2_block_size[(sizeof(test_iq2_block) == IQ2_BLOCK_BYTES) ? 1 : -1]; + +static void fail(const char *what) { + fprintf(stderr, "FAIL: %s\n", what); + exit(1); +} + +static void require(int condition, const char *what) { + if (!condition) fail(what); +} + +static uint32_t next_u32(uint32_t *state) { + *state = *state * 1664525u + 1013904223u; + return *state; +} + +static void fill_matrix(test_iq2_block *matrix, uint32_t salt) { + uint32_t state = 0x9e3779b9u ^ salt; + for (uint32_t expert = 0; expert < N_EXPERT; expert++) { + for (uint32_t row = 0; row < DIM; row++) { + test_iq2_block *block = &matrix[(uint64_t)expert * DIM + row]; + block->d = 0x2800u; /* IEEE binary16 0.03125. */ + for (uint32_t i = 0; i < DIM / 8u; i++) { + block->qs[i] = (uint16_t)(next_u32(&state) >> 8); + } + } + } +} + +static float dot_row(const test_iq2_block *matrix, + uint32_t expert, + uint32_t row, + const float *values) { + return ds4_test_iq2_xxs_dot_f32( + DIM, &matrix[(uint64_t)expert * DIM + row], values); +} + +static void compare_array(const char *name, + const float *got, + const float *want, + uint64_t count, + float abs_tol, + float rel_tol, + float *worst_abs_out, + float *worst_rel_out) { + float worst_abs = 0.0f; + float worst_rel = 0.0f; + uint64_t worst_i = 0; + for (uint64_t i = 0; i < count; i++) { + const float delta = fabsf(got[i] - want[i]); + const float scale = fmaxf(fabsf(got[i]), fabsf(want[i])); + const float rel = scale > 1.0e-7f ? delta / scale : delta; + if (!isfinite(got[i]) || !isfinite(want[i]) || + (delta > abs_tol && rel > rel_tol)) { + fprintf(stderr, + "FAIL: %s[%llu] got=%g want=%g abs=%g rel=%g " + "tolerance=(%g,%g)\n", + name, + (unsigned long long)i, + got[i], + want[i], + delta, + rel, + abs_tol, + rel_tol); + exit(1); + } + if (delta > worst_abs) { + worst_abs = delta; + worst_i = i; + } + if (rel > worst_rel) worst_rel = rel; + } + fprintf(stderr, + "%s: count=%llu worst_abs=%g at=%llu worst_rel=%g\n", + name, + (unsigned long long)count, + worst_abs, + (unsigned long long)worst_i, + worst_rel); + if (worst_abs_out) *worst_abs_out = worst_abs; + if (worst_rel_out) *worst_rel_out = worst_rel; +} + +int main(void) { + const uint64_t rows_per_matrix = (uint64_t)N_EXPERT * DIM; + const uint64_t matrix_bytes = rows_per_matrix * sizeof(test_iq2_block); + const uint64_t model_bytes = 3u * matrix_bytes; + void *model_raw = NULL; + require(posix_memalign(&model_raw, 4096u, model_bytes) == 0, + "allocate aligned model map"); + test_iq2_block *gate_model = model_raw; + test_iq2_block *up_model = (test_iq2_block *)((char *)model_raw + matrix_bytes); + test_iq2_block *down_model = (test_iq2_block *)((char *)model_raw + 2u * matrix_bytes); + fill_matrix(gate_model, 1u); + fill_matrix(up_model, 2u); + fill_matrix(down_model, 3u); + + float x[N_TOKEN * DIM]; + for (uint32_t token = 0; token < N_TOKEN; token++) { + for (uint32_t i = 0; i < DIM; i++) { + x[(uint64_t)token * DIM + i] = + 0.075f * sinf((float)(i + 1u) * (0.031f + 0.007f * token)) + + 0.025f * cosf((float)(i + 3u) * (0.017f + 0.003f * token)); + } + } + const int32_t selected[N_TOKEN * N_USED] = {0, 2, 1, 0}; + const float route_weights[N_TOKEN * N_USED] = {0.65f, 0.35f, 0.55f, 0.45f}; + const uint64_t slot_count = (uint64_t)N_TOKEN * N_USED; + const uint64_t routed_count = slot_count * DIM; + const uint64_t out_count = (uint64_t)N_TOKEN * DIM; + + float *want_gate = calloc(routed_count, sizeof(float)); + float *want_up = calloc(routed_count, sizeof(float)); + float *want_mid = calloc(routed_count, sizeof(float)); + float *want_down = calloc(routed_count, sizeof(float)); + float *want_out = calloc(out_count, sizeof(float)); + float *got_gate = calloc(routed_count, sizeof(float)); + float *got_up = calloc(routed_count, sizeof(float)); + float *got_mid = calloc(routed_count, sizeof(float)); + float *got_down = calloc(routed_count, sizeof(float)); + float *got_out = calloc(out_count, sizeof(float)); + require(want_gate && want_up && want_mid && want_down && want_out && + got_gate && got_up && got_mid && got_down && got_out, + "allocate oracle buffers"); + + for (uint32_t token = 0; token < N_TOKEN; token++) { + const float *token_x = &x[(uint64_t)token * DIM]; + for (uint32_t slot = 0; slot < N_USED; slot++) { + const uint64_t pair = (uint64_t)token * N_USED + slot; + const uint32_t expert = (uint32_t)selected[pair]; + for (uint32_t row = 0; row < DIM; row++) { + const uint64_t index = pair * DIM + row; + const float gate = dot_row(gate_model, expert, row, token_x); + const float up = dot_row(up_model, expert, row, token_x); + want_gate[index] = gate; + want_up[index] = up; + want_mid[index] = + (gate / (1.0f + expf(-gate))) * up * route_weights[pair]; + } + for (uint32_t row = 0; row < DIM; row++) { + const uint64_t index = pair * DIM + row; + want_down[index] = dot_row( + down_model, expert, row, &want_mid[pair * DIM]); + want_out[(uint64_t)token * DIM + row] += want_down[index]; + } + } + } + + require(ds4_gpu_init() != 0, "initialize CUDA backend"); + require(ds4_gpu_set_model_map(model_raw, model_bytes) != 0, + "register test model map"); + + ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(out_count * sizeof(float)); + ds4_gpu_tensor *gate = ds4_gpu_tensor_alloc(routed_count * sizeof(float)); + ds4_gpu_tensor *up = ds4_gpu_tensor_alloc(routed_count * sizeof(float)); + ds4_gpu_tensor *mid = ds4_gpu_tensor_alloc(routed_count * sizeof(float)); + ds4_gpu_tensor *down = ds4_gpu_tensor_alloc(routed_count * sizeof(float)); + ds4_gpu_tensor *selected_gpu = ds4_gpu_tensor_alloc(sizeof(selected)); + ds4_gpu_tensor *weights_gpu = ds4_gpu_tensor_alloc(sizeof(route_weights)); + ds4_gpu_tensor *x_gpu = ds4_gpu_tensor_alloc(sizeof(x)); + require(out && gate && up && mid && down && selected_gpu && weights_gpu && x_gpu, + "allocate CUDA tensors"); + require(ds4_gpu_tensor_write(selected_gpu, 0, selected, sizeof(selected)) != 0, + "upload selected experts"); + require(ds4_gpu_tensor_write(weights_gpu, 0, route_weights, sizeof(route_weights)) != 0, + "upload route weights"); + require(ds4_gpu_tensor_write(x_gpu, 0, x, sizeof(x)) != 0, + "upload activations"); + + bool mid_is_f16 = true; + require(ds4_gpu_routed_moe_batch_tensor( + out, + gate, + up, + mid, + down, + model_raw, + model_bytes, + 0, + matrix_bytes, + 2u * matrix_bytes, + IQ2_TYPE, + IQ2_TYPE, + (uint64_t)DIM * IQ2_BLOCK_BYTES, + IQ2_BLOCK_BYTES, + (uint64_t)DIM * IQ2_BLOCK_BYTES, + IQ2_BLOCK_BYTES, + DIM, + DIM, + DIM, + selected_gpu, + weights_gpu, + N_EXPERT, + N_USED, + 0.0f, + x_gpu, + 0, + N_TOKEN, + &mid_is_f16, + true) != 0, + "launch CUDA IQ2 routed MoE"); + require(!mid_is_f16, "IQ2 routed midpoint remains f32"); + require(ds4_gpu_synchronize() != 0, "synchronize CUDA routed MoE"); + require(ds4_gpu_tensor_read(gate, 0, got_gate, routed_count * sizeof(float)) != 0, + "read gate output"); + require(ds4_gpu_tensor_read(up, 0, got_up, routed_count * sizeof(float)) != 0, + "read up output"); + require(ds4_gpu_tensor_read(mid, 0, got_mid, routed_count * sizeof(float)) != 0, + "read midpoint output"); + require(ds4_gpu_tensor_read(down, 0, got_down, routed_count * sizeof(float)) != 0, + "read down output"); + require(ds4_gpu_tensor_read(out, 0, got_out, out_count * sizeof(float)) != 0, + "read routed output"); + + float worst_abs = 0.0f; + float worst_rel = 0.0f; + compare_array("gate", got_gate, want_gate, routed_count, + 2.0e-5f, 2.0e-5f, NULL, NULL); + compare_array("up", got_up, want_up, routed_count, + 2.0e-5f, 2.0e-5f, NULL, NULL); + compare_array("mid", got_mid, want_mid, routed_count, + 3.0e-5f, 3.0e-5f, NULL, NULL); + compare_array("down", got_down, want_down, routed_count, + 2.0e-4f, 2.0e-4f, NULL, NULL); + compare_array("out", got_out, want_out, out_count, + 3.0e-4f, 3.0e-4f, &worst_abs, &worst_rel); + + ds4_gpu_tensor_free(x_gpu); + ds4_gpu_tensor_free(weights_gpu); + ds4_gpu_tensor_free(selected_gpu); + ds4_gpu_tensor_free(down); + ds4_gpu_tensor_free(mid); + ds4_gpu_tensor_free(up); + ds4_gpu_tensor_free(gate); + ds4_gpu_tensor_free(out); + ds4_gpu_cleanup(); + free(got_out); + free(got_down); + free(got_mid); + free(got_up); + free(got_gate); + free(want_out); + free(want_down); + free(want_mid); + free(want_up); + free(want_gate); + free(model_raw); + + fprintf(stderr, + "test_cuda_iq2_routed PASS tokens=%u experts=%u/%u dim=%u " + "output_worst_abs=%g output_worst_rel=%g\n", + N_TOKEN, + N_USED, + N_EXPERT, + DIM, + worst_abs, + worst_rel); + return 0; +}