ROCm: accelerate GLM 5.3 prefill and reduce MTP overhead on gfx1151

ds4.c

@@ -47515,14 +47515,15 @@ static bool glm_graph_mtp_matmul(
  * hidden h[pos] (g->cur for GLM-5.2, g->hc_cur for GLM-5.3) and next_token
  * (= token[pos+1]), writes the nextn KV at slot pos, and returns the
  * drafted token[pos+2] by greedy argmax. Clobbers the decode scratch. */
-static bool glm_graph_mtp_step(
+static bool glm_graph_mtp_step_impl(
         ds4_glm_gpu_graph *g,
         const ds4_model   *model,
         const ds4_weights *weights,
         int                next_token,
         uint32_t           pos,
         uint32_t           min_pos,
-        int               *draft_out) {
+        int               *draft_out,
+        bool               cache_only) {
     if (!g || !model || !weights || !draft_out) return false;
     if (DS4_N_NEXTN_PREDICT == 0) return false;
     const uint32_t cache_cap = glm_graph_mtp_cache_cap(g);
@@ -47725,6 +47726,20 @@ static bool glm_graph_mtp_step(
                                                      DS4_N_KV_LORA,
                                                      DS4_N_ROT,
                                                      glm_graph_compact_cache_is_f16()) != 0;
+    /* The first accepted MTP step only needs to publish its private KV row.
+     * Complete that prefix before the caller uploads the next hidden row. */
+    if (cache_only) {
+        DS4_GLM_MTP_STAGE("cache_end");
+        if (ok) ok = ds4_gpu_end_commands() != 0;
+        else (void)ds4_gpu_synchronize();
+        ds4_gpu_tensor_free(enorm_view);
+        ds4_gpu_tensor_free(hnorm_view);
+        if (!ok) {
+            fprintf(stderr, "ds4: glm mtp cache step failed at stage '%s' (pos %u)\n",
+                    mtp_stage, pos);
+        }
+        return ok;
+    }
     DS4_GLM_MTP_STAGE("qk_low");
     if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_tensor(g->qk_low,
                                                      g->q,
@@ -47930,6 +47945,39 @@ static bool glm_graph_mtp_step(
     return true;
 }
 
+static bool glm_graph_mtp_step(
+        ds4_glm_gpu_graph *g,
+        const ds4_model *model,
+        const ds4_weights *weights,
+        int next_token,
+        uint32_t pos,
+        uint32_t min_pos,
+        int *draft_out) {
+    return glm_graph_mtp_step_impl(g, model, weights, next_token, pos,
+                                    min_pos, draft_out, false);
+}
+
+static bool glm_graph_mtp_accepted_first_step(
+        ds4_glm_gpu_graph *g,
+        const ds4_model *model,
+        const ds4_weights *weights,
+        int next_token,
+        uint32_t pos,
+        uint32_t min_pos,
+        int *draft_out,
+        bool greedy) {
+    bool cache_only = false;
+#ifdef DS4_ROCM_BUILD
+    cache_only = greedy && g && g->glm53 && !g->ssd_streaming &&
+                 !g->placement && g->tp_world <= 1 &&
+                 ds4_gpu_dspark_gfx1151_fast_path() != 0;
+#else
+    (void)greedy;
+#endif
+    return glm_graph_mtp_step_impl(g, model, weights, next_token, pos,
+                                    min_pos, draft_out, cache_only);
+}
+
 static bool glm_graph_forward_token(
         ds4_glm_gpu_graph *g,
         const ds4_model   *model,
@@ -65903,8 +65951,9 @@ static int ds4_session_glm_spec_cycle_impl(
         const bool cu =
             ds4_gpu_tensor_write(target_hidden, 0, s->glm_mtp_hc,
                                  hc_row_bytes) != 0 &&
-            glm_graph_mtp_step(g, &e->model, &e->weights, d, pos,
-                               s->glm_mtp_min_pos, &dummy) &&
+            glm_graph_mtp_accepted_first_step(g, &e->model, &e->weights, d, pos,
+                                             s->glm_mtp_min_pos, &dummy,
+                                             !exact_sampling && temperature <= 0.0f) &&
             ds4_gpu_tensor_write(target_hidden,
                                  0,
                                  s->glm_mtp_hc + hc_row_values,
@@ -67024,6 +67073,12 @@ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, c
             start = s->checkpoint.len;
             resumed_checkpoint = true;
             s->mtp_draft_valid = false;
+            if (prompt->len > start) {
+                /* Prompt rows do not populate the private nextn KV window. */
+                s->glm_mtp_have = 0;
+                s->glm_mtp_rollback_valid = false;
+                s->glm_mtp_min_pos = 0;
+            }
         } else {
             s->checkpoint.len = 0;
             s->checkpoint_valid = false;
@@ -68825,6 +68880,7 @@ static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp,
         if (!s->glm_spec_inside) {
             s->glm_mtp_rollback_valid = false;
             s->glm_mtp_have = 0;
+            s->glm_mtp_min_pos = 0;
         }
         const uint32_t pos = (uint32_t)s->checkpoint.len;
         const bool updates_dense =

ds4_rocm.cu

@@ -54,6 +54,9 @@ extern "C" int ds4_mmq_iq2_xxs_moe_pair(
     const int32_t *ids, float *out_a, float *out_b,
     int M, int K, int n_tokens, int n_experts, int n_expert_used,
     cudaStream_t stream);
+extern "C" int ds4_mmq_q4_K_dense(
+    const void *W, const float *X_f32, float *out_f32,
+    int M, int N, int K, cudaStream_t stream);
 extern "C" int ds4_mmq_q8_0_dense_vec(
     const void *W, const float *X_f32, float *out_f32,
     int M, int N, int K, cudaStream_t stream);

rocm/ds4_rocm_glm.cuh

@@ -457,6 +457,21 @@ extern "C" int ds4_gpu_matmul_q4_K_tensor(
     const char *weight = cuda_model_range_ptr(
         model_map, weight_offset, weight_bytes, "GLM-5.3 Q4_K matrix");
     if (!weight) return 0;
+    // GLM-5.3 KDA Q/K bulk projection. MMQ uses 32-value Q8_1
+    // activation scales instead of this entry's 256-value Q8_K scales.
+    // Keep quality, streaming, small-row and other-model arithmetic intact.
+    if (g_glm_model && !g_quality_mode && !g_ssd_streaming_mode &&
+        ds4_rocm_is_gfx1151() && in_dim == 4096u && out_dim == 8192u &&
+        n_rows >= 128u && n_rows <= 2048u) {
+        const int rc = ds4_mmq_q4_K_dense(
+            weight, (const float *)x->ptr, (float *)out->ptr,
+            (int)out_dim, (int)n_rows, (int)in_dim, (cudaStream_t)0);
+        if (rc != 0) {
+            fprintf(stderr, "ds4: GLM-5.3 Q4_K MMQ failed (%d)\n", rc);
+            return 0;
+        }
+        return 1;
+    }
     cuda_block_q8_K *xq = (cuda_block_q8_K *)cuda_tmp_alloc(
         xq_bytes,
         "GLM-5.3 Q4_K activations");
@@ -1513,6 +1528,56 @@ __global__ static void glm_indexer_scores_batch_kernel(
     if (threadIdx.x == 0u) *dst = score;
 }
 
+/* GLM-5.3 bulk scorer on gfx1151: one wave32 per pooled key/query pair.
+ * Keep FP32 queries, sequential head accumulation, and the scalar tree's
+ * 64/32/16/8/4/2/1 grouping. Four independent waves need no block barriers. */
+__global__ static void glm53_indexer_scores_wave32_kernel(
+        float *scores,
+        const float *q,
+        const float *weights,
+        const __half *indexer_key_cache,
+        uint32_t n_rows,
+        uint32_t n_tokens,
+        uint32_t pos0,
+        float scale) {
+    const uint32_t lane = threadIdx.x & 31u;
+    const uint32_t row = blockIdx.x * 4u + (threadIdx.x >> 5u);
+    const uint32_t token = blockIdx.y;
+    if (row >= n_rows || token >= n_tokens) return;
+    float *dst = scores + (uint64_t)token * n_rows + row;
+    if (row >= (pos0 + token + 1u) / 4u) {
+        if (lane == 0u) *dst = -INFINITY;
+        return;
+    }
+    const __half *kr = indexer_key_cache + (uint64_t)row * 128u + lane;
+    const float k0 = __half2float(kr[0]);
+    const float k1 = __half2float(kr[32]);
+    const float k2 = __half2float(kr[64]);
+    const float k3 = __half2float(kr[96]);
+    float score = 0.0f;
+    for (uint32_t h = 0; h < 32u; ++h) {
+        const float *qh = q + ((uint64_t)token * 32u + h) * 128u + lane;
+        float dot;
+        {
+            // Preserve the explicit pair tree. The gfx1151 backend can still
+            // fuse product pairs, so rounding differs from the scalar kernel.
+#pragma clang fp reassociate(off)
+#pragma clang fp contract(off)
+            const float p0 = qh[0] * k0;
+            const float p1 = qh[32] * k1;
+            const float p2 = qh[64] * k2;
+            const float p3 = qh[96] * k3;
+            dot = (p0 + p2) + (p1 + p3);
+        }
+        dot = warp_sum_f32(dot);
+        if (lane == 0u) {
+            score += fmaxf(dot * scale, 0.0f) *
+                     weights[(uint64_t)token * 32u + h];
+        }
+    }
+    if (lane == 0u) *dst = score;
+}
+
 __global__ static void glm_attention_indexed_lora_kernel(
         float *lora_out,
         const float *q,
@@ -2945,6 +3010,19 @@ extern "C" int ds4_gpu_glm53_indexer_scores_batch_tensor(
         !glm_rocm_tensor_has_cache2(indexer_key_cache, n_rows, head_dim, elem)) {
         return 0;
     }
+    if (g_glm_model && !g_quality_mode && !g_ssd_streaming_mode &&
+        ds4_rocm_is_gfx1151() &&
+        cache_f16 && n_head == 32u && head_dim == 128u &&
+        pool_size == 4u && n_tokens >= 128u) {
+        const dim3 grid(n_rows / 4u + (n_rows % 4u != 0u), n_tokens, 1u);
+        glm53_indexer_scores_wave32_kernel<<<grid, 128u>>>(
+                (float *)scores->ptr, (const float *)q->ptr,
+                (const float *)weights->ptr,
+                (const __half *)indexer_key_cache->ptr,
+                n_rows, n_tokens, pos0, scale);
+        return cuda_ok(cudaGetLastError(),
+                       "GLM-5.3 grouped indexer wave32 scores launch");
+    }
     const dim3 grid(n_rows, n_tokens, 1u);
     glm_indexer_scores_batch_kernel<<<grid, 256u>>>(
         (float *)scores->ptr, (const float *)q->ptr,

rocm/ds4_rocm_matmul.cuh

@@ -1,3 +1,8 @@
+/* Existing engine query reads the active model variant; its engine argument
+ * is unused. Keep GLM 5.2's overlapping 2048 -> 4096 indexer unchanged. */
+struct ds4_engine;
+extern "C" bool ds4_engine_is_glm53(struct ds4_engine *e);
+
 __global__ static void matmul_f16_tiny_batch_wave_kernel(
         float *out,
         const half *w,
@@ -428,7 +433,21 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode
         !cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
         !cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
         x->bytes < x_bytes || out->bytes < out_bytes) return 0;
-    if (n_tok > 1 && !g_quality_mode &&
+    /* Tested GLM 5.3 Q8 projections reuse each weight for two quantized rows.
+     * The small 128 -> 8192 KDA projections retain their existing dispatch. */
+    const bool glm_q8_tok2_preq =
+        g_glm_model && !g_ssd_streaming_mode && !g_quality_mode &&
+        n_tok == 2u &&
+        ((in_dim == 4096u &&
+          (out_dim == 64u || out_dim == 128u || out_dim == 512u ||
+           out_dim == 1536u || out_dim == 2048u || out_dim == 8192u ||
+           out_dim == 12288u)) ||
+         (out_dim == 4096u &&
+          (in_dim == 2048u || in_dim == 8192u || in_dim == 12288u ||
+           in_dim == 16384u)) ||
+         (in_dim == 1536u && out_dim == 16384u)) &&
+        ds4_rocm_is_gfx1151() && ds4_engine_is_glm53(NULL);
+    if (n_tok > 1 && !g_quality_mode && !glm_q8_tok2_preq &&
         cuda_runtime_config()->shared_down_cublas && in_dim == 2048u && out_dim == 4096u &&
         cuda_matmul_q8_0_tensor_f16_gemm(out, model_map, model_size, weight_offset,
                                          in_dim, out_dim, x, n_tok, label ? label : "shared_expert")) {
@@ -493,7 +512,7 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode
                 blocks);
         return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 warp launch");
     }
-    if (n_tok > 1) {
+    if (n_tok > 1 && !glm_q8_tok2_preq) {
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
         if (!g_quality_mode &&
             g_dspark_verify_mode && n_tok <= 6u &&
@@ -611,7 +630,7 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode
                 blocks);
         return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch warp launch");
     }
-    if (g_cublas_ready && n_tok > 1) {
+    if (g_cublas_ready && n_tok > 1 && !glm_q8_tok2_preq) {
         const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
         if (w_f16) {
             const uint64_t xh_count = n_tok * in_dim;
@@ -661,6 +680,14 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode
     dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1);
     quantize_q8_0_f32_kernel<<<qgrid, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
     if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0;
+    if (glm_q8_tok2_preq) {
+        matmul_q8_0_preq_batch_warp8_tok2_kernel<<<
+                ((unsigned)out_dim + 7u) / 8u, 256u>>>(
+                (float *)out->ptr,
+                reinterpret_cast<const unsigned char *>(wptr),
+                xq, xscale, in_dim, out_dim, blocks, use_dp4a);
+        return cuda_ok(cudaGetLastError(), "matmul_q8_0 GLM tok2 prequant launch");
+    }
     if (n_tok == 1) {
         const uint32_t rows_per_block = cfg->q8_decode_rpb;
         matmul_q8_0_preq_rows_w32_kernel<<<

rocm/ds4_rocm_moe_launch.cuh

@@ -813,9 +813,17 @@ static int routed_moe_launch(
             mxfp4_path && use_expert_tiles && !use_mxfp4_tile32 && !use_mxfp4_ldsB &&
             !use_mxfp4_tile4 && n_tokens >= 8u &&
             getenv("DS4_ROCM_ENABLE_MXFP4_ROW64") != NULL;
+        /* GLM-5.3 Flash has 288 experts, beyond the generic admission cap.
+         * Admit only its resident IQ2 gate/up and Q2 down topology on gfx1151;
+         * the existing checks below still exclude quality and small batches. */
+        const uint32_t glm53_iq2_mmq_topology =
+            g_glm_model && !g_ssd_streaming_mode && ds4_rocm_is_gfx1151() &&
+            n_total_expert == DS4_ROCM_GLM53_N_EXPERT &&
+            n_expert == 8u && expert_in_dim == 4096u &&
+            expert_mid_dim == 2048u && out_dim == 4096u;
         const uint32_t use_rocm_mmq_gateup =
             ok && iq2_path && n_tokens >= 128u && !g_quality_mode &&
-            n_total_expert <= 256u &&
+            (n_total_expert <= 256u || glm53_iq2_mmq_topology) &&
             !batch_stream_selected && !batch_stream_split_selected &&
             !split_selected && !compact_selected && gate_w && up_w &&
             (stream_full_layer || full_table_cached) &&

rocm/ds4_rocm_q8.cuh

@@ -286,6 +286,60 @@ __global__ static void matmul_q8_0_preq_batch_warp8_kernel(
     if (lane == 0) out[tok * out_dim + row] = acc;
 }
 
+/* Reuse packed Q8 weights across two activation rows, with the same
+ * per-token integer dots and wave reduction as ordinary prequant decode. */
+__global__ static void matmul_q8_0_preq_batch_warp8_tok2_kernel(
+        float *out,
+        const unsigned char *w,
+        const int8_t *xq,
+        const float *xscale,
+        uint64_t in_dim,
+        uint64_t out_dim,
+        uint64_t blocks,
+        int use_dp4a) {
+    const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
+    const uint32_t lane = threadIdx.x & 31u;
+    if (row >= out_dim) return;
+
+    const unsigned char *wr = w + row * blocks * 34u;
+    const int8_t *xqr0 = xq;
+    const int8_t *xqr1 = xq + blocks * 32u;
+    const float *xsr0 = xscale;
+    const float *xsr1 = xscale + blocks;
+    float acc0 = 0.0f;
+    float acc1 = 0.0f;
+    for (uint64_t b = lane; b < blocks; b += 32u) {
+        const uint64_t i0 = b * 32u;
+        const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u;
+        const __half *scale_h = (const __half *)(wr + b * 34u);
+        const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u);
+        const int8_t *xqb0 = xqr0 + b * 32u;
+        const int8_t *xqb1 = xqr1 + b * 32u;
+        int dot0 = 0;
+        int dot1 = 0;
+        if (use_dp4a && bn == 32u) {
+#pragma unroll
+            for (uint32_t i = 0; i < 32u; i += 4u) {
+                const int32_t w4 = load_i8x4_i32_unaligned(qs + i);
+                dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0);
+                dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1);
+            }
+        } else {
+            dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
+            dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
+        }
+        const float ws = __half2float(*scale_h);
+        acc0 += ws * xsr0[b] * (float)dot0;
+        acc1 += ws * xsr1[b] * (float)dot1;
+    }
+    acc0 = warp_sum_f32(acc0);
+    acc1 = warp_sum_f32(acc1);
+    if (lane == 0) {
+        out[row] = acc0;
+        out[out_dim + row] = acc1;
+    }
+}
+
 __device__ static float q8_0_scale_scalar(const unsigned char *blk) {
     const uint16_t bits = (uint16_t)blk[0] | ((uint16_t)blk[1] << 8);
     return __half2float(__ushort_as_half((unsigned short)bits));