fix(cuda): support GLM IQ2_XXS routed down projection

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

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,

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);

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<<<mgrid, 256>>>(
+            (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<<<dgrid, 256>>>(
+                (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;