diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index c295aed46..05612d91f 100644 --- a/server/CMakeLists.txt +++ b/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 diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq-streamk-schedule.h b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq-streamk-schedule.h new file mode 100644 index 000000000..7a0f4cfc3 --- /dev/null +++ b/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 + +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(ntiles_dst) + nsm - 1) / nsm; + const int64_t tiles_efficiency_percent = + 100 * static_cast(ntiles_dst) / (static_cast(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(ntiles_dst) >= tiles_needed_to_fill) { + return nsm; + } + const int64_t total_iters = static_cast(ntiles_dst) * iters_per_tile; + return total_iters > 0 ? static_cast(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; +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 8fc24e91d..8afacb239 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/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 #include @@ -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 tmp_fixup(pool); diff --git a/server/test/test_mmq_streamk_schedule.cpp b/server/test/test_mmq_streamk_schedule.cpp new file mode 100644 index 000000000..392e8d49c --- /dev/null +++ b/server/test/test_mmq_streamk_schedule.cpp @@ -0,0 +1,63 @@ +#include "mmq-streamk-schedule.h" + +#include +#include +#include + +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; +}