perf(cuda): cap SM86 Stream-K grid to useful chunks

server/CMakeLists.txt

@@ -977,6 +977,12 @@ if(DFLASH27B_TESTS)
             ggml-base)
         list(APPEND _raw_unit_test_targets test_mmq_streamk_iq4_xs)
     endif()
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_mmq_streamk_schedule.cpp")
+        add_executable(test_mmq_streamk_schedule test/test_mmq_streamk_schedule.cpp)
+        target_include_directories(test_mmq_streamk_schedule PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src/ggml-cuda)
+        list(APPEND _raw_unit_test_targets test_mmq_streamk_schedule)
+    endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_drafter_tail_capture_guard.cpp")
         # RED phase binary: same source WITHOUT the fix flag — documents the bug.
         add_executable(test_drafter_tail_capture_guard_red

server/deps/llama.cpp/ggml/src/ggml-cuda/mmq-streamk-schedule.h

@@ -0,0 +1,51 @@
+// Pure host-side MMQ stream-k grid selection (no CUDA types).
+// Shared by mmq.cuh and host unit tests.
+#pragma once
+
+#include <cstdint>
+
+static inline int mmq_stream_k_nblocks(
+        const int ntiles_dst,
+        const int nsm,
+        const int64_t ncols_x,
+        const int iter_k,
+        const bool is_nvidia,
+        const bool enable_useful_chunk_cap) {
+    if (ntiles_dst <= 0 || nsm <= 0 || ncols_x <= 0 || iter_k <= 0) {
+        return 1;
+    }
+
+    const int64_t tiles_nwaves = (static_cast<int64_t>(ntiles_dst) + nsm - 1) / nsm;
+    const int64_t tiles_efficiency_percent =
+        100 * static_cast<int64_t>(ntiles_dst) / (static_cast<int64_t>(nsm) * tiles_nwaves);
+
+    // Preserve the pre-existing NVIDIA >=90% pure-tiling escape.
+    if (is_nvidia && tiles_efficiency_percent >= 90) {
+        return ntiles_dst;
+    }
+
+    // Fail closed outside the specifically validated hardware path.
+    if (!enable_useful_chunk_cap) {
+        return nsm;
+    }
+
+    // The device aligns CTA boundaries down to complete iter_k/qk chunks.
+    // Keep a partial K tail with the tile's final CTA instead of counting it
+    // as another CTA; otherwise the aligned partition creates empty CTAs.
+    const int64_t iters_per_tile = ncols_x >= iter_k ? ncols_x / iter_k : 1;
+    // We only need the exact product when it is below nsm. Saturate first so
+    // neither division nor multiplication can overflow.
+    if (iters_per_tile >= nsm) {
+        return nsm;
+    }
+    const int64_t tiles_needed_to_fill = (nsm + iters_per_tile - 1) / iters_per_tile;
+    if (static_cast<int64_t>(ntiles_dst) >= tiles_needed_to_fill) {
+        return nsm;
+    }
+    const int64_t total_iters = static_cast<int64_t>(ntiles_dst) * iters_per_tile;
+    return total_iters > 0 ? static_cast<int>(total_iters) : 1;
+}
+
+static inline bool mmq_stream_k_fixup_needed(const int ntiles_dst, const int stream_k_blocks) {
+    return stream_k_blocks > 0 && ntiles_dst % stream_k_blocks != 0;
+}

server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh

@@ -3,6 +3,7 @@
 #include "common.cuh"
 #include "vecdotq.cuh"
 #include "mma.cuh"
+#include "mmq-streamk-schedule.h"
 
 #include <climits>
 #include <cstdint>
@@ -4312,14 +4313,19 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
 
     // For the stream-k kernel it is possible to run it with tiling by setting the number of CUDA blocks equal to the number of tiles.
     // This is worthwhile if the efficiency of tiling is high and skipping the fixup kernel is more important.
+    // On the validated SM86 path, also avoid launching more CTAs than useful MMQ iteration chunks.
     const int ntiles_dst = ntx * nty * ntzw;
-    const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm;
-    const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves);
-    const dim3 block_nums_stream_k(GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90 ? ntiles_dst : nsm, 1, 1);
+    // SM86 always uses the regular MMQ iteration width. MXFP4 only switches
+    // to MMQ_ITER_K_MXFP4_FP4 in the Blackwell device path.
+    const int iter_k = MMQ_ITER_K;
+    const bool enable_useful_chunk_cap = cc == 860; // NVIDIA SM86 only; fail closed elsewhere.
+    const int stream_k_blocks = mmq_stream_k_nblocks(
+        ntiles_dst, nsm, args.ncols_x, iter_k, GGML_CUDA_CC_IS_NVIDIA(cc), enable_useful_chunk_cap);
+    const dim3 block_nums_stream_k(stream_k_blocks, 1, 1);
 
     GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow.
 
-    const bool fixup_needed = ntiles_dst % block_nums_stream_k.x != 0;
+    const bool fixup_needed = mmq_stream_k_fixup_needed(ntiles_dst, stream_k_blocks);
 
     ggml_cuda_pool & pool = ctx.pool(id);
     ggml_cuda_pool_alloc<float> tmp_fixup(pool);

server/test/test_mmq_streamk_schedule.cpp

@@ -0,0 +1,63 @@
+#include "mmq-streamk-schedule.h"
+
+#include <climits>
+#include <cstdint>
+#include <cstdio>
+
+namespace {
+constexpr int kIterK = 256;
+constexpr int kNsm86 = 82;
+int failures = 0;
+
+void expect(const char * name, int actual, int wanted) {
+    const bool pass = actual == wanted;
+    std::printf("[%s] %-34s actual=%d expected=%d\n", pass ? "PASS" : "FAIL", name, actual, wanted);
+    failures += pass ? 0 : 1;
+}
+} // namespace
+
+int main() {
+    // Validated SM86 path: avoid empty CTAs for shallow K.
+    expect("sm86 shallow pure tiling", mmq_stream_k_nblocks(20, 82, 256, kIterK, true, true), 20);
+    expect("sm86 partial fill",        mmq_stream_k_nblocks(20, 82, 512, kIterK, true, true), 40);
+    expect("sm86 deep K unchanged",    mmq_stream_k_nblocks(20, 82, 5120, kIterK, true, true), 82);
+    expect("sm86 single shallow",      mmq_stream_k_nblocks(1, 82, 256, kIterK, true, true), 1);
+    expect("sm86 single deep capped",  mmq_stream_k_nblocks(1, 82, 5120, kIterK, true, true), 20);
+
+    // qk-aligned K tails stay with the tile's final CTA. Counting ceil(K/256)
+    // would produce alternating empty CTAs and a spurious fixup launch.
+    const int k320_stream_blocks = mmq_stream_k_nblocks(20, 82, 320, kIterK, true, true);
+    expect("sm86 K320 qk32 tail",       k320_stream_blocks, 20);
+    expect("sm86 K320 skips fixup",     mmq_stream_k_fixup_needed(20, k320_stream_blocks), 0);
+    expect("sm86 K384 qk32 tail",       mmq_stream_k_nblocks(20, 82, 384, kIterK, true, true), 20);
+    const int k576_stream_blocks = mmq_stream_k_nblocks(20, 82, 576, kIterK, true, true);
+    expect("sm86 K576 split remains",   k576_stream_blocks, 40);
+    expect("sm86 K576 needs fixup",     mmq_stream_k_fixup_needed(20, k576_stream_blocks), 1);
+
+    // Existing >=90% NVIDIA tiling behavior remains unchanged.
+    expect("nvidia 90pct tiling",       mmq_stream_k_nblocks(74, 82, 5120, kIterK, true, false), 74);
+    expect("nvidia full SM tiling",     mmq_stream_k_nblocks(82, 82, 5120, kIterK, true, false), 82);
+
+    // Fail closed: non-SM86 and non-NVIDIA paths retain the old nsm schedule.
+    expect("other nvidia unchanged",    mmq_stream_k_nblocks(20, 82, 256, kIterK, true, false), 82);
+    expect("non-nvidia unchanged",      mmq_stream_k_nblocks(20, 82, 256, kIterK, false, false), 82);
+
+    // Invalid inputs remain bounded, and 64-bit arithmetic avoids overflow.
+    expect("invalid ntiles",            mmq_stream_k_nblocks(0, 82, 256, kIterK, true, true), 1);
+    expect("invalid nsm",               mmq_stream_k_nblocks(20, 0, 256, kIterK, true, true), 1);
+    expect("invalid K",                 mmq_stream_k_nblocks(20, 82, 0, kIterK, true, true), 1);
+    expect("invalid iter",              mmq_stream_k_nblocks(20, 82, 256, 0, true, true), 1);
+    expect("large shape no overflow",   mmq_stream_k_nblocks(INT_MAX / 2, 82, INT64_MAX / 4, kIterK, true, true), INT_MAX / 2);
+
+    // Fixup predicate itself remains unchanged.
+    expect("tiling skips fixup", mmq_stream_k_fixup_needed(20, 20), 0);
+    expect("partial needs fixup", mmq_stream_k_fixup_needed(20, 40), 1);
+    expect("deep needs fixup",    mmq_stream_k_fixup_needed(20, 82), 1);
+
+    if (failures) {
+        std::fprintf(stderr, "FAILED: %d schedule checks\n", failures);
+        return 1;
+    }
+    std::puts("ALL PASS: SM86 Stream-K schedule checks");
+    return 0;
+}