From c3ffaf0437d9d413087e7303f872e4b355076e77 Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Wed, 15 Jul 2026 04:11:36 +0300 Subject: [PATCH 1/2] feat: Backend-specific prefill chunk sizes - The default adaptive_gpu_chunk is too conservative, For example it'd fire up 28 times for 204 tokens causing large delays in prefill with 800 ms~ prefill latency if used with FA. - By using larger chunk sizes that are adjusted for each backend, we can safely decrease this delay (378 ms with FA) to increase performance. --- src/s2_model.cpp | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/s2_model.cpp b/src/s2_model.cpp index 9ced623..55633e9 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -815,12 +815,21 @@ bool SlowARModel::prefill(const std::vector & flat_tokens, int32_t n_to backend_requires_single_token_semantic_prefill(backend_gpu_)) { chunk_tokens = 1; } else { + // Backend-specific prefill chunk sizes for naive attention path. + // Vulkan: Limited by ggml_soft_max shader workgroup/shared memory limits + // CUDA/Metal: Highly optimized soft_max kernels handle large batches + // CPU: Memory bandwidth bound, scale with available parallelism + int32_t adaptive_gpu_chunk = 64; // Safe fallback + +#if defined(GGML_USE_CUDA) || defined(GGML_USE_METAL) + adaptive_gpu_chunk = 512; +#elif defined(GGML_USE_VULKAN) + adaptive_gpu_chunk = 64; +#else + // CPU or unknown backend: scale with thread count, capped at 256 + adaptive_gpu_chunk = std::min(256, std::max(64, n_threads * 16)); +#endif - const int32_t adaptive_gpu_chunk = std::clamp( - 128 / std::max(1, n_gpu_layers_), - 8, - 64 - ); if (n_tokens > adaptive_gpu_chunk) { chunk_tokens = adaptive_gpu_chunk; } From 42e1a93b3d8608f404713df0c0a52b42d99dd81c Mon Sep 17 00:00:00 2001 From: Skyrion9 Date: Wed, 15 Jul 2026 07:17:10 +0300 Subject: [PATCH 2/2] fix(prefill): use runtime backend detection and fix unreachable CPU chunking - The CPU fallback was nested inside (backend_gpu_ != nullptr) guard. We can instead establish a CPU baseline first and adapt it to backend down the line. --- src/s2_model.cpp | 48 +++++++++++++++++++++++++++++++++++------------- 1 file changed, 35 insertions(+), 13 deletions(-) diff --git a/src/s2_model.cpp b/src/s2_model.cpp index 55633e9..f505a0c 100755 --- a/src/s2_model.cpp +++ b/src/s2_model.cpp @@ -819,19 +819,41 @@ bool SlowARModel::prefill(const std::vector & flat_tokens, int32_t n_to // Vulkan: Limited by ggml_soft_max shader workgroup/shared memory limits // CUDA/Metal: Highly optimized soft_max kernels handle large batches // CPU: Memory bandwidth bound, scale with available parallelism - int32_t adaptive_gpu_chunk = 64; // Safe fallback - -#if defined(GGML_USE_CUDA) || defined(GGML_USE_METAL) - adaptive_gpu_chunk = 512; -#elif defined(GGML_USE_VULKAN) - adaptive_gpu_chunk = 64; -#else - // CPU or unknown backend: scale with thread count, capped at 256 - adaptive_gpu_chunk = std::min(256, std::max(64, n_threads * 16)); -#endif - - if (n_tokens > adaptive_gpu_chunk) { - chunk_tokens = adaptive_gpu_chunk; + int32_t adaptive_chunk = std::min(256, std::max(64, resolve_n_threads(n_threads) * 16)); + + if (backend_gpu_ != nullptr && n_gpu_layers_ > 0) { + const char * backend_name = ggml_backend_name(backend_gpu_); + if (backend_name != nullptr) { + std::string name(backend_name); + if (name.find("CUDA") != std::string::npos || + name.find("Metal") != std::string::npos) { + adaptive_chunk = 512; + } else if (name.find("Vulkan") != std::string::npos) { + adaptive_chunk = 64; + } + } + + int32_t semantic_prompt_tokens = 0; + for (int32_t t = 0; t < n_tokens; ++t) { + const int32_t semantic = flat_tokens[static_cast(t) * codebook_dim]; + if (semantic >= hparams_.semantic_begin_id && + semantic <= hparams_.semantic_end_id) { + semantic_prompt_tokens++; + if (semantic_prompt_tokens > 1) { + break; + } + } + } + + if (semantic_prompt_tokens > 1 && + backend_requires_single_token_semantic_prefill(backend_gpu_)) { + adaptive_chunk = 1; + } + } + + int32_t chunk_tokens = n_tokens; + if (n_tokens > adaptive_chunk) { + chunk_tokens = adaptive_chunk; } } }