diff --git a/.github/workflows/unsloth-pin-preflight.yml b/.github/workflows/unsloth-pin-preflight.yml index f390f381e6f..494bca228c8 100644 --- a/.github/workflows/unsloth-pin-preflight.yml +++ b/.github/workflows/unsloth-pin-preflight.yml @@ -28,16 +28,10 @@ permissions: contents: write issues: write -# Two runs of the same ref probe the same pins against the same base, so the -# second adds nothing and just competes for runners. On 08-04 a dispatch and the -# schedule sat queued together for an hour. Newest wins: it sees the newest -# pr-set.json. -# -# Per ref, though, not globally. This file also runs on any push that touches -# pr-set.json, so with one shared group a push to a second branch cancelled the -# first branch's run: observed on 09-03, where the run that would have said -# whether a repin fixed the nightly was cancelled by an unrelated branch, and -# the PR was left showing the failure from before the fix. +# Two runs of the same ref probe the same pins against the same base, so the second only +# competes for runners; newest wins, since it sees the newest pr-set.json. Per ref, not +# globally: this also runs on any push touching pr-set.json, and with one shared group a push +# to a second branch cancelled the first branch's run (09-03). concurrency: group: unsloth-pin-preflight-${{ github.ref }} cancel-in-progress: true @@ -56,10 +50,8 @@ jobs: id: p run: | set -uo pipefail - # Everything below reports through `status`/`details`, so a death - # anywhere else leaves both empty and the alert blank: a red X on a - # scheduled run nobody opens. Report the abort through the same - # channel as a finding, so the repin bot sees a failure either way. + # everything below reports through `status`/`details`, so a death anywhere else + # leaves the alert blank; report the abort through the same channel as a finding trap 'rc=$?; if [ "$rc" != 0 ]; then { echo "status=failure" echo "details</dev/null PROBLEMS="${PROBLEMS}- \`${SRC}#${NUM}\` (\`${SHA:0:10}\`) does not merge onto \`${BASE}\` + the pins before it.\n\n Conflicting files:\n\n\`\`\`\n${FILES}\n\`\`\`\n\n
conflict hunks\n\n\`\`\`diff\n${HUNKS}\n\`\`\`\n\n
\n" @@ -193,11 +180,9 @@ jobs: PROBLEMS="${PROBLEMS}- the merged tree builds, but \`scripts/unsloth/merge_checks.py\` found a resolution that is silently wrong. See the run log for file and line.\n" fi - # The other half of that question. merge_checks.py asks whether the - # tree contains something wrong; this asks whether it still contains - # what each pin carries. A pin that has rotted into a no-op, or an - # arch registration a resolution quietly dropped, is invisible to - # every other check here and to the compiler. + # the other half: merge_checks.py asks whether the tree contains something + # wrong, this asks whether it still contains what each pin carries. A pin rotted + # into a no-op is invisible to every other check here and to the compiler. if ! python3 ../scripts/unsloth/pin_contract.py --root . --base "$BASE" \ --pr-set ../scripts/unsloth/pr-set.json --report "${RUNNER_TEMP}/pin_contract.json" ; then PROBLEMS="${PROBLEMS}- the merged tree is missing code a pin carries. See the run log for the pin and file.\n" @@ -207,12 +192,10 @@ jobs: PROBLEMS="${PROBLEMS}- pins upstream has taken over, safe to delete from \`pr-set.json\`:\n\n\`\`\`\n${NOTES}\n\`\`\`\n" fi - # A clean merge is not a compiling tree. On 09-03 ggml-org#27754 - # merged with no conflicts at all and did not compile: upstream had - # added a parameter to build_attn_mha and the pin's new - # build_attn_sparse still called the old signature. Nothing above - # can see that. CPU only and the `llama` target only, which is where - # that translation unit lives; 59s cold at -j4 with no ccache. + # a clean merge is not a compiling tree: on 09-03 ggml-org#27754 merged with no + # conflicts and did not compile, upstream having added a parameter to + # build_attn_mha that the pin's build_attn_sparse still called without. CPU only, + # 59s cold at -j4 with no ccache. GATE_OK=1 if ! cmake -B "${RUNNER_TEMP}/gate" -DCMAKE_BUILD_TYPE=Release \ -DGGML_CUDA=OFF -DLLAMA_BUILD_TESTS=ON -DLLAMA_BUILD_SERVER=OFF \ @@ -223,10 +206,8 @@ jobs: PROBLEMS="${PROBLEMS}- the pins merge cleanly and the merged tree does not compile. See the run log for the file and line; this is the failure that only shows up in the CUDA leg once the nightly has fanned out.\n" fi - # The last question, and the only one that needs a binary: does each - # feature we ship still work. Everything above is about the source. - # CPU only, because no runner in this pipeline has a GPU -- see the - # note in feature_matrix.py about what that does and does not prove. + # the only question that needs a binary: does each feature we ship still work. + # CPU only, since no runner here has a GPU; see the note in feature_matrix.py. if [ -n "$GATE_OK" ]; then if ! python3 ../scripts/unsloth/feature_matrix.py \ --build-dir "${RUNNER_TEMP}/gate" \ diff --git a/.github/workflows/unsloth-pr-set-lint.yml b/.github/workflows/unsloth-pr-set-lint.yml index 8a898539f02..99b06b74f34 100644 --- a/.github/workflows/unsloth-pr-set-lint.yml +++ b/.github/workflows/unsloth-pr-set-lint.yml @@ -128,9 +128,9 @@ jobs: done exit "$fail" - # A pin nobody decided about is the failure this whole file exists to stop. - # Being in `unchecked` with a reason is a fine answer; being in neither map - # is how DiffusionGemma went five weeks with no coverage and no record of it. + # a pin nobody decided about is the failure this file exists to stop: being in + # `unchecked` with a reason is fine, being in neither map is how DiffusionGemma went + # five weeks with no coverage and no record of it - name: Every pin is either checked or knowingly unchecked run: | set -euo pipefail diff --git a/.github/workflows/unsloth-prebuilt.yml b/.github/workflows/unsloth-prebuilt.yml index 834ff544909..353d88417ca 100644 --- a/.github/workflows/unsloth-prebuilt.yml +++ b/.github/workflows/unsloth-prebuilt.yml @@ -282,8 +282,8 @@ jobs: # .github/workflows, which upstream history routinely does). if [ "$EXISTS" != "true" ] || [ "${{ github.event_name }}" = "workflow_dispatch" ]; then git remote add upstream https://github.com/ggml-org/llama.cpp.git - # The upstream checkout below takes scripts/unsloth/ away. Copy the - # whole dir out, not file by file: see the note above the step. + # the upstream checkout below takes scripts/unsloth/ away; copy the whole dir + # out, not file by file, see the note above the step cp -r scripts/unsloth "${RUNNER_TEMP}/us" ADDITIVE_MERGE="${RUNNER_TEMP}/us/additive_merge.py" if [ "$(jq length <<<"$PRS")" != 0 ]; then @@ -446,8 +446,7 @@ jobs: # A bad pin resolution can still build fine, so it must be caught before the source artifact ships. See merge_checks.py. # Its own step, not more script in `resolve`: GitHub caps one workflow string at 21000 chars and that step is near it. See check_workflow_scalars.py. - # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: every check added - # here would otherwise cost another line inside the capped block, and going over silently disables the whole workflow. + # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: another line inside the capped block per check would eventually go over, which silently disables the whole workflow. - name: Check the merged tree for silently wrong resolutions if: ${{ env.MERGED_PINS == '1' }} run: | @@ -457,13 +456,11 @@ jobs: exit 1 fi - # merge_checks.py asserts the ABSENCE of two known-bad shapes. This asserts the PRESENCE of what each pin carries, which is a different question and - # the one that goes unanswered when a pin rots into a no-op or a resolution quietly drops an arch registration. Free, so it runs before the compile gate. + # merge_checks.py asserts the ABSENCE of two known-bad shapes; this asserts the PRESENCE of what each pin carries, the question that goes unanswered when a pin rots into a no-op. Free, so it runs before the compile gate. - name: Check every pin still contributes what it carries if: ${{ env.MERGED_PINS == '1' }} - # Through env, never interpolated into the script: `prs` carries PR - # titles, which are third-party text, and `${{ }}` pastes them into the - # shell source before bash ever sees it. + # through env, never interpolated: `prs` carries PR titles, which are third-party + # text, and `${{ }}` pastes them into the shell source before bash sees it env: PRS: ${{ steps.r.outputs.prs }} BASE: ${{ steps.r.outputs.base }} @@ -475,12 +472,8 @@ jobs: exit 1 fi - # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, because upstream had added a parameter to - # build_attn_mha and the pin's new build_attn_sparse still called the old signature. Nothing before this point can see that, and without it the release - # dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4 with no ccache, against 20-60 minutes for a CUDA build. - # mtmd is in the gate because `llama` alone is not enough: observed 09-04, ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, - # upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const, so it overrode nothing and the - # vision and audio towers were abstract. Every vision pin lands in mtmd, so a gate that skips it cannot see the whole class. + # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, upstream having added a parameter to build_attn_mha that the pin's build_attn_sparse still called without. Without this the release dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4, against 20-60 minutes for CUDA. + # mtmd is in the gate because `llama` alone is not enough: on 09-04 ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const. Every vision pin lands in mtmd. - name: Compile gate (CPU, llama and mtmd targets) if: ${{ env.MERGED_PINS == '1' }} run: | diff --git a/common/arg.cpp b/common/arg.cpp index 5bfa4adcdf0..e4f2e8f2517 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1304,6 +1304,13 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e exit(0); } params.lr.init(); + + // [TAG_EXACT_CONCURRENCY] refuse a column bound that cannot cover a decode step before + // anything is loaded, rather than running with the guarantee quietly switched off + if (!common_exact_concurrency_init(ctx_arg.params)) { + ctx_arg.params = params_org; + return false; + } } catch (const std::invalid_argument & ex) { fprintf(stderr, "%s\n", ex.what()); ctx_arg.params = params_org; diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002d..b74bee734ae 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1,4 +1,5 @@ #include "ggml.h" +#include "ggml-backend.h" #include "gguf.h" #include "build-info.h" @@ -1289,6 +1290,14 @@ struct common_init_result::impl { common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { + // [TAG_EXACT_CONCURRENCY] before any context exists, so one is never created under a figure + // the explicit bound does not cover; this also covers a caller that skipped + // common_params_parse(). On failure nothing is loaded. + if (!model_only && !common_exact_concurrency_init(params)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n"); + return; + } + auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); @@ -1433,6 +1442,78 @@ std::vector & common_init_result::lora() { return pimpl->lora; } +// [TAG_EXACT_CONCURRENCY] +bool common_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + +// [TAG_EXACT_CONCURRENCY] +int common_exact_decode_width(const common_params & params) { + const int64_t n_slots = std::max(1, params.n_parallel); + + // draft tokens a slot carries into the verify ubatch, from the same place the speculation + // code takes its own width + const int64_t n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); + + // the product is handed to a backend as an int; one that overflows is reported, not wrapped + const int64_t n_cols = n_slots*(1 + n_draft); + + return n_cols > INT32_MAX ? -1 : (int) n_cols; +} + +// [TAG_EXACT_CONCURRENCY] +bool common_exact_concurrency_init(const common_params & params) { + if (!common_exact_concurrency()) { + return true; + } + + // DFlash drafting turns causal attention off on its draft context, which the paged attention + // needs; say so instead of asserting in the graph. DSpark is the same implementation. + for (const auto type : params.speculative.types) { + if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n"); + return false; + } + } + + const int n_cols = common_exact_decode_width(params); + + if (n_cols < 0) { + COM_ERR("LLAMA_EXACT_CONCURRENCY: a decode step of %d slots with %d draft tokens each is too wide to report\n", + std::max(1, params.n_parallel), std::max(0, (int) common_speculative_n_max(¶ms.speculative))); + return false; + } + + const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + if (bound) { + const int max_cols = atoi(bound); + if (max_cols > 0 && max_cols < n_cols) { + COM_ERR("GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at " + "least %d to cover a decode step of %d slots, above which a matmul is left " + "batched and its rows depend on the other rows in the ubatch. Raise it to %d, " + "set it to 0 for no bound, or unset it to let it default to %d.\n", + max_cols, n_cols, std::max(1, params.n_parallel), n_cols, n_cols); + return false; + } + } + + // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step + // is. A context created later reports n_seq_max times that, which is n_cols again; reporting + // n_cols here too covers a caller that decodes first, or contexts it created earlier. + if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) || + !llama_set_exact_decode_width((uint32_t) n_cols)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n"); + return false; + } + + return true; +} + common_init_result_ptr common_init_from_params(common_params & params, bool model_only) { common_init_result_ptr res(new common_init_result(params, model_only)); diff --git a/common/common.h b/common/common.h index c99269f9a96..4616b73153e 100644 --- a/common/common.h +++ b/common/common.h @@ -931,6 +931,17 @@ using common_init_result_ptr = std::unique_ptr; common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false); +// [TAG_EXACT_CONCURRENCY] true when LLAMA_EXACT_CONCURRENCY is set for this process +bool common_exact_concurrency(); + +// the widest ubatch a decode step can build here: one column per slot times one plus its draft +// tokens. Under exact mode the CUDA column policy has to cover this, and derives its bound from it. +int common_exact_decode_width(const common_params & params); + +// report that width to the CUDA backend, refusing a smaller explicit +// GGML_CUDA_BATCH_INVARIANT_MAX_COLS; false if the configuration must not run +bool common_exact_concurrency_init(const common_params & params); + struct llama_model_params common_model_params_to_llama ( common_params & params); struct llama_context_params common_context_params_to_llama(const common_params & params); diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index 1cd81eeaebc..07131df327c 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -38,6 +38,13 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); + +// [TAG_EXACT_CONCURRENCY] report the widest ubatch a decode step of this process can build: one +// column per slot times one plus its draft tokens. Under LLAMA_EXACT_CONCURRENCY the column policy +// defaults to that instead of a fixed number, so --parallel or a wider draft cannot silently push a +// decode above the bound. An explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before +// the first graph is computed; also available through ggml_backend_reg_get_proc_address(). +GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5e5541aac94..9b9213df856 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2656,6 +2656,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten return true; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return false; + } #ifdef ASCEND_310P // FA not support on 310p device return false; diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 8cece71f186..8bed84e44bb 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -474,6 +474,10 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return ggml_is_contiguous(op->src[0]); case GGML_OP_SSM_SCAN: return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; + // [TAG_EXACT_CONCURRENCY] note: FLASH_ATTN_EXT with src[5], the page table, is deliberately + // still accepted. The CPU ignores it and attends in physical order, but it is also the + // reference test-backend-ops compares the paged CUDA kernel against, and that test's mask + // selects exactly the listed cells. A KV layer cannot reach the CPU under the mode anyway. default: return true; } diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 14dd1098c97..d749b42ee48 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -49,6 +49,12 @@ #define GGML_CUDA_CC_PASCAL 600 #define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products +// [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits +int ggml_cuda_batch_invariant(); +// widest batch the split applies to, 0 = no bound; bounding it gives up prompt-phase invariance +// only, and prompt-sized batches cost far more to split +int ggml_cuda_batch_invariant_max_cols(); + #define GGML_CUDA_CC_VOLTA 700 #define GGML_CUDA_CC_TURING 750 #define GGML_CUDA_CC_AMPERE 800 diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index e67cc7fdf78..10f28eb229b 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1091,7 +1091,10 @@ void launch_fattn( // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + // [TAG_BATCH_INVARIANT] without this scan the KV loop runs to K->ne[1], which grows with the + // other sequences sharing the cache; scanning the mask bounds it by the sequence's own extent + const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0; + if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); @@ -1148,6 +1151,14 @@ void launch_fattn( if (ntiles_dst % blocks_num.x != 0) { // Fixup is only needed if the SMs work on fractional tiles. dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2))); } + } else if (dst->src[5] || ggml_cuda_batch_invariant()) { + // [TAG_BATCH_INVARIANT] the KV split between blocks, and so the order the partials combine + // in, follows K->ne[1], which grows with the other sequences: pin it to one block per tile + parallel_blocks = 1; + + blocks_num.x = ntiles_x; + blocks_num.y = parallel_blocks; + blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3]; } else { // parallel_blocks must not be larger than what the tensor size allows: parallel_blocks = std::min(parallel_blocks, ntiles_KV); @@ -1214,7 +1225,7 @@ void launch_fattn( V_data, mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, - KV_max.ptr, + dst->src[5] ? (const int *) dst->src[5]->data : KV_max.ptr, !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 69dd9368624..4005d28879e 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -16,7 +16,7 @@ static constexpr __device__ int ggml_cuda_fattn_vec_get_nthreads_device() { #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wpass-failed" #endif // __clang__ -template // D == head size +template // D == head size __launch_bounds__(ggml_cuda_fattn_vec_get_nthreads_device(), 1) static __global__ void flash_attn_ext_vec( const char * Q_ptr, @@ -247,13 +247,25 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } - const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; + // in the paged specialization KV_max carries [count, physical page IDs...] per query; the loop + // and each warp's recurrence follow logical positions, never physical addresses + static_assert(!paged || ncols == 1, "paged attention has one query per block"); + const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr; + const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11); + const char * K_base = K; + const char * V_base = V; + const half * mask_base = maskh; K += blockIdx.y*nthreads * nb11; V += blockIdx.y*nthreads * nb21; maskh += blockIdx.y*nthreads; for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, - // Increment pointers after each loop: K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) { + if constexpr (paged) { + const int physical = pages[1 + k_VKQ_0/FATTN_KQ_STRIDE]*FATTN_KQ_STRIDE + k_VKQ_0%FATTN_KQ_STRIDE; + K = K_base + int64_t(physical)*nb11; + V = V_base + int64_t(physical)*nb21; + maskh = mask_base + physical; + } // Calculate KQ tile and keep track of new maximum KQ values: float KQ_reg[ncols]; // KQ in registers. diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c0..915c2d04da8 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -457,6 +457,12 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; + // [TAG_BATCH_INVARIANT] every choice below switches on Q->ne[1] or K->ne[1], both of which grow + // with the other sequences, so pin the kernel a batch of one would use + if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) { + return BEST_FATTN_KERNEL_VEC; + } + // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel) { @@ -569,6 +575,52 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + + if (dst->src[5]) { + GGML_ASSERT(dst->src[0]->ne[0] == 256 && dst->src[2]->ne[0] == 256); + GGML_ASSERT(dst->src[1]->type == GGML_TYPE_F16 && dst->src[2]->type == GGML_TYPE_F16); + GGML_ASSERT(dst->src[3] && dst->src[0]->ne[3] == 1); + GGML_ASSERT(dst->src[5]->type == GGML_TYPE_I32 && ggml_is_contiguous(dst->src[5])); + GGML_ASSERT(dst->src[5]->ne[0] == 1 + dst->src[1]->ne[1]/FATTN_KQ_STRIDE); + GGML_ASSERT(dst->src[5]->ne[1] == dst->src[0]->ne[1]); + float softcap; + memcpy(&softcap, (const float *) dst->op_params + 2, sizeof(softcap)); + GGML_ASSERT(softcap == 0.0f); + fattn_kernel_t kernel = flash_attn_ext_vec<256, 1, GGML_TYPE_F16, GGML_TYPE_F16, false, true>; + launch_fattn<256, 1, 1>(ctx, dst, kernel, 4, 0, 128, false, false, false); + return; + } + + // [TAG_BATCH_INVARIANT] attend one query row at a time, as a batch of one would + const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols(); + if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 && + (fattn_max_cols <= 0 || dst->src[0]->ne[1] <= fattn_max_cols)) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * mask = dst->src[3]; + + for (int64_t i = 0; i < Q->ne[1]; ++i) { + ggml_tensor Q_row = *Q; + Q_row.ne[1] = 1; + Q_row.data = (char *) Q->data + i*Q->nb[1]; + + ggml_tensor mask_row; + ggml_tensor dst_row = *dst; + // ne[2] runs to the end of dst so the F16 K/V scratch behind dst stays in place + dst_row.ne[2] = dst->ne[2] - i; + dst_row.data = (char *) dst->data + i*dst->nb[2]; + dst_row.src[0] = &Q_row; + if (mask) { + mask_row = *mask; + mask_row.ne[1] = 1; + mask_row.data = (char *) mask->data + i*mask->nb[1]; + dst_row.src[3] = &mask_row; + } + + ggml_cuda_flash_attn_ext(ctx, &dst_row); + } + return; + } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc62..76fa4dc66c6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1758,6 +1758,12 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving + // it on would give a solo request a different code path from a batched one + if (ggml_cuda_batch_invariant()) { + return false; + } + ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1785,6 +1791,12 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { } static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving + // it on would give a solo request a different code path from a batched one + if (ggml_cuda_batch_invariant()) { + return false; + } + ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1813,60 +1825,303 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - GGML_TENSOR_BINARY_OP_LOCALS +// [TAG_BATCH_INVARIANT] the token count picks the matmul implementation and how its K loop is +// divided between threads, both of which change the summation order, so the same request produces +// different bits depending on how many others decode alongside it. GGML_CUDA_BATCH_INVARIANT: +// 1 - compute every destination column on its own, exactly as a batch of one would +// 2 - split off only the columns whose batch-of-one configuration differs from the batched one +static bool ggml_cuda_exact_concurrency() { + static const bool exact = []() { + const char * value = getenv("LLAMA_EXACT_CONCURRENCY"); + return value && atoi(value) != 0; + }(); + return exact; +} - const int32_t hint = ggml_get_op_params_i32(dst, 1); - if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { +int ggml_cuda_batch_invariant() { + static const int mode = []() { + if (ggml_cuda_exact_concurrency()) { return 2; } + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT"); + return val ? atoi(val) : 0; + }(); + return mode; +} + +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch the caller says it can build, 0 if it never said +static std::atomic g_exact_decode_width{0}; + +void ggml_backend_cuda_set_exact_decode_width(int n_cols) { + // monotonic: the widest figure ever reported stays, whatever order the reports arrive in + int cur = g_exact_decode_width.load(std::memory_order_relaxed); + + while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + } +} + +int ggml_cuda_batch_invariant_max_cols() { + // [TAG_EXACT_CONCURRENCY] prompt ubatches hold one sequence, so a prefill already matches its + // solo run and needs no unbounded column policy. An explicit bound always wins. + static const int explicit_cols = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + return val ? atoi(val) : -1; + }(); + + if (explicit_cols >= 0) { + return explicit_cols; + } + + if (!ggml_cuda_exact_concurrency()) { + return 0; + } + + // the widest ubatch a decode step can build: one column per slot times one plus its draft + // tokens, as reported by ggml_backend_cuda_set_exact_decode_width(). Failing a report, 16, + // which covers --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2. Above the bound the + // column split does not fire and ggml_cuda_warn_above_exact_bound() says so once. + const int width = g_exact_decode_width.load(std::memory_order_relaxed); + + return width > 0 ? width : 16; +} + +// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so the mode does not hold +// for that op; say so once. Only when nothing reported a decode width: a reported one makes the +// batches above the bound prompt ubatches, which are exact by holding a single sequence, so +// warning on those would be crying wolf on every prefill. +static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) { + if (!ggml_cuda_exact_concurrency()) { + return; + } + + if (g_exact_decode_width.load(std::memory_order_relaxed) > 0) { return; } + static std::atomic_flag warned = ATOMIC_FLAG_INIT; + if (warned.test_and_set(std::memory_order_relaxed)) { + return; + } + + GGML_LOG_WARN("%s: LLAMA_EXACT_CONCURRENCY is set, but this %s is %d columns wide while " + "GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d, so it is left batched and its result depends " + "on the other columns in the ubatch. Raise the bound, set it to 0 for no bound, or call " + "ggml_backend_cuda_set_exact_decode_width() with the widest decode this process builds. " + "Reported once.\n", __func__, op, (int) ncols, max_cols); +} + +enum ggml_cuda_mm_path { + GGML_CUDA_MM_CUBLAS_UNSUPPORTED, + GGML_CUDA_MM_MMVF, + GGML_CUDA_MM_MMVF_TRANSPOSED, + GGML_CUDA_MM_MMF, + GGML_CUDA_MM_MMVQ, + GGML_CUDA_MM_MMQ, + GGML_CUDA_MM_CUBLAS, +}; + +// the implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns +static ggml_cuda_mm_path ggml_cuda_mul_mat_path( + int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) { // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. // Therefore, in such cases use cuBLAS. const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); - return; + return GGML_CUDA_MM_CUBLAS_UNSUPPORTED; } - - const int cc = ggml_cuda_info().devices[ctx.device].cc; - const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; - if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { // The custom F16 vector kernel can be used over batched cuBLAS GEMM. // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVF; } // A transposed vector can still use MMVQ (i.e. ne01 == 1) - if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1 + if (src0->ne[1] == 1 && ne11 > MMVF_MAX_BATCH_SIZE && dst->ne[2] == 1 && dst->ne[3] == 1 && src0->type == GGML_TYPE_F32 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && ggml_cuda_should_use_mmvf(src1->type, cc, src1->ne, src1->nb, /*ne11 =*/ 1)) { - ggml_tensor dst_vec = *dst; - dst_vec.ne[0] = ne11; - dst_vec.ne[1] = 1; - dst_vec.nb[1] = dst_vec.nb[0]*ne11; - dst_vec.nb[2] = dst_vec.nb[1]; - dst_vec.nb[3] = dst_vec.nb[1]; - ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); - return; + return GGML_CUDA_MM_MMVF_TRANSPOSED; } if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { - ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMF; } if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVQ; } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return GGML_CUDA_MM_MMQ; + } + return GGML_CUDA_MM_CUBLAS; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +// [TAG_BATCH_INVARIANT] the widest slice of columns that can be recomputed in one launch while +// every column still sums as a batch of one would. A column's result depends on the implementation +// and, for MMVQ, the launch's warp count, never on the other columns, so a slice this wide reads +// the weights once for all of them instead of once per column. Always below ncols_dst, so the +// recursive call cannot land back here with the same shape. +static int64_t ggml_cuda_mul_mat_invariant_width( + int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, + ggml_cuda_mm_path path_one, int64_t ncols_dst) { + if (path_one != GGML_CUDA_MM_MMVF && path_one != GGML_CUDA_MM_MMVQ) { + return 1; + } + const int64_t widest = path_one == GGML_CUDA_MM_MMVF ? MMVF_MAX_BATCH_SIZE : MMVQ_MAX_BATCH_SIZE; + for (int64_t w = std::min(ncols_dst - 1, widest); w > 1; --w) { + if (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, w) != path_one) { + continue; + } + if (path_one == GGML_CUDA_MM_MMVQ && !ggml_cuda_mmvq_matches_single_column(src0->type, cc, w)) { + continue; + } + return w; + } + return 1; +} + +// recompute dst in slices of columns so each column sees the batch-of-one configuration; false +// when the batched launch already gives every column that same value +static bool ggml_cuda_mul_mat_split_columns( + ggml_backend_cuda_context & ctx, int cc, int warp_size, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + // recurrent output projections broadcast one weight matrix over sequence planes: these are + // token projections too, so normalize each plane before applying the column policy + if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 && + (dst->ne[2] > 1 || dst->ne[3] > 1) && + src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) { + for (int64_t i3 = 0; i3 < dst->ne[3]; ++i3) { + for (int64_t i2 = 0; i2 < dst->ne[2]; ++i2) { + ggml_tensor src_plane = *src1; + ggml_tensor dst_plane = *dst; + src_plane.ne[2] = src_plane.ne[3] = 1; + dst_plane.ne[2] = dst_plane.ne[3] = 1; + src_plane.data = (char *) src1->data + i2*src1->nb[2] + i3*src1->nb[3]; + dst_plane.data = (char *) dst->data + i2*dst->nb[2] + i3*dst->nb[3]; + ggml_cuda_mul_mat(ctx, src0, &src_plane, &dst_plane); + } + } + return true; + } + + const int64_t ncols_dst = dst->ne[1]; + if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) { + return false; + } + // only the token dimension is split; batched matmuls (attention) keep their shape + if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { + return false; + } + const int max_cols = ggml_cuda_batch_invariant_max_cols(); + if (max_cols > 0 && ncols_dst > max_cols) { + ggml_cuda_warn_above_exact_bound("MUL_MAT", ncols_dst, max_cols); + return false; + } + + // mode 1 recomputes one column at a time; mode 2, which exact concurrency runs under, uses the + // widest slices that keep the batch-of-one arithmetic + int64_t width = 1; + if (ggml_cuda_batch_invariant() >= 2) { + const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); + const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst); + if (path_one == path_batched) { + // same implementation, but it still has to sum in the same order + if (path_batched == GGML_CUDA_MM_MMVF) { + return false; // the block size follows K alone + } + if (path_batched == GGML_CUDA_MM_MMVQ && + ggml_cuda_mmvq_matches_single_column(src0->type, cc, ncols_dst)) { + return false; + } + } + width = ggml_cuda_mul_mat_invariant_width(cc, warp_size, src0, src1, dst, path_one, ncols_dst); + if (width >= ncols_dst) { + width = 1; + } + } + + for (int64_t i = 0; i < ncols_dst; i += width) { + const int64_t n = std::min(width, ncols_dst - i); + + ggml_tensor src1_col = *src1; + ggml_tensor dst_col = *dst; + + src1_col.ne[1] = n; + src1_col.nb[2] = n*src1_col.nb[1]; + src1_col.nb[3] = n*src1_col.nb[1]; + src1_col.data = (char *) src1->data + i*src1->nb[1]; + + dst_col.ne[1] = n; + dst_col.nb[2] = n*dst_col.nb[1]; + dst_col.nb[3] = n*dst_col.nb[1]; + dst_col.data = (char *) dst->data + i*dst->nb[1]; + + ggml_cuda_mul_mat(ctx, src0, &src1_col, &dst_col); + } + return true; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS + + const int32_t hint = ggml_get_op_params_i32(dst, 1); + if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { return; } - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + + if (ggml_cuda_batch_invariant() && ggml_cuda_mul_mat_split_columns(ctx, cc, warp_size, src0, src1, dst)) { + return; + } + + switch (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ne11)) { + case GGML_CUDA_MM_CUBLAS_UNSUPPORTED: + case GGML_CUDA_MM_CUBLAS: + ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + return; + case GGML_CUDA_MM_MMVF: + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMVF_TRANSPOSED: { + ggml_tensor dst_vec = *dst; + dst_vec.ne[0] = ne11; + dst_vec.ne[1] = 1; + dst_vec.nb[1] = dst_vec.nb[0]*ne11; + dst_vec.nb[2] = dst_vec.nb[1]; + dst_vec.nb[3] = dst_vec.nb[1]; + ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); + return; + } + case GGML_CUDA_MM_MMF: + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMVQ: + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMQ: + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return; + } + GGML_ABORT("fatal error"); +} + +// [TAG_BATCH_INVARIANT] true when the policy computes this MUL_MAT_ID one token at a time, so +// every expert product reduces as it would in a batch of one +static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { + if (!ggml_cuda_batch_invariant()) { + return false; + } + const int64_t ntokens = dst->ne[2]; + if (ntokens <= 1) { + return false; + } + const int max_cols = ggml_cuda_batch_invariant_max_cols(); + if (max_cols > 0 && ntokens > max_cols) { + ggml_cuda_warn_above_exact_bound("MUL_MAT_ID", ntokens, max_cols); + return false; + } + return true; } // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization @@ -1879,9 +2134,13 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c return true; } - if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) { + // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path that + // decides whether the stream is synchronized is the single-token one + const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2]; + + if (ntokens <= MMVQ_MAX_BATCH_SIZE) { if (ggml_is_quantized(src0->type)) { - if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) { + if (ntokens <= get_mmvq_mmid_max_batch(src0->type, cc)) { return false; } } else if (GGML_CUDA_CC_IS_AMD(cc)) { @@ -1889,17 +2148,54 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c } } - if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) { + if (ggml_cuda_should_use_mmq(src0->type, cc, ntokens, /*n_experts=*/src0->ne[2])) { return false; } - if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, ntokens, /*mul_mat_id=*/true)) { return false; } return true; } +static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// [TAG_BATCH_INVARIANT] recompute dst one token at a time. Every implementation below groups the +// ubatch's tokens by the expert they routed to, so shapes depend on what the other tokens picked; +// one call per token makes the callee see the shapes a batch of one has. +static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * ids = dst->src[2]; + + const int64_t ntokens = dst->ne[2]; + + for (int64_t i = 0; i < ntokens; ++i) { + ggml_tensor src1_token = *src1; + ggml_tensor ids_token = *ids; + ggml_tensor dst_token = *dst; + + // src1 is [ne10, ne11, ntokens], one expert list per token in ids [n_expert_used, ntokens] + src1_token.ne[2] = 1; + src1_token.nb[3] = src1_token.nb[2]; + src1_token.data = (char *) src1->data + i*src1->nb[2]; + + ids_token.ne[1] = 1; + ids_token.nb[2] = ids_token.nb[1]; + ids_token.nb[3] = ids_token.nb[1]; + ids_token.data = (char *) ids->data + i*ids->nb[1]; + + dst_token.ne[2] = 1; + dst_token.nb[3] = dst_token.nb[2]; + dst_token.data = (char *) dst->data + i*dst->nb[2]; + + dst_token.src[1] = &src1_token; + dst_token.src[2] = &ids_token; + + ggml_cuda_mul_mat_id(ctx, &dst_token); + } +} + static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -1912,6 +2208,19 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + // [TAG_BATCH_INVARIANT] + if (ggml_cuda_mul_mat_id_splits_tokens(dst)) { + GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1); + // a quantized expert matrix takes the single-token MMVQ path at every token count, and that + // path can put the tokens on its sample axis in one launch; anything else goes token by token + if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); + return; + } + ggml_cuda_mul_mat_id_split_tokens(ctx, dst); + return; + } + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); @@ -3297,9 +3606,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } - //topk-moe - if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || - cgraph->nodes[i]->op == GGML_OP_ARGSORT) { + // topk-moe + // [TAG_BATCH_INVARIANT] the routing fusion passes its memory-range check only for a one-token + // ubatch, so a solo request takes the fused top-k kernel and a batched one takes the softmax, + // argsort and normalize chain: two algorithms for one set of routing weights + if (!ggml_cuda_batch_invariant() && + (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || + cgraph->nodes[i]->op == GGML_OP_ARGSORT)) { ggml_cuda_topk_moe_args args; const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); std::vector ops; @@ -5492,6 +5805,10 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + // [TAG_EXACT_CONCURRENCY] + if (strcmp(name, "ggml_backend_cuda_set_exact_decode_width") == 0) { + return (void *)ggml_backend_cuda_set_exact_decode_width; + } return nullptr; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 97053480980..30f37c7722e 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -541,6 +541,21 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } +// [TAG_BATCH_INVARIANT] +bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst) { + if (ncols_dst < 1 || ncols_dst > MMVQ_MAX_BATCH_SIZE) { + return false; + } + const mmvq_parameter_table_id table_id = get_device_table_id(cc); + if (table_id == MMVQ_PARAMETERS_GB10) { + // There nwarps also depends on the K loop trip count, which the caller does not pass in. + return ncols_dst == 1; + } + // blocks_per_iter, which assigns K blocks to threads, is proportional to nwarps; + // rows_per_cuda_block only changes which rows a block owns, not the order within a row + return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id); +} + template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( @@ -577,9 +592,12 @@ static __global__ void mul_mat_vec_q( uint32_t sample_dst; ggml_cuda_pdl_sync(); - channel_x = ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); - channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; sample_dst = blockIdx.z; + // [TAG_BATCH_INVARIANT] with ids, a sample is a token: the batch-invariant MUL_MAT_ID launch + // puts every token on the z axis of one single-column launch, so each (token, expert slot) + // block runs the single-token configuration. The stock launch has one sample, as before. + channel_x = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio); + channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); const uint32_t sample_y = sample_dst; @@ -1266,7 +1284,13 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); - GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE); + // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token + // configuration with the tokens on the sample axis, so every (token, expert slot) block + // reduces as the token alone would and the count is not bounded by the column templates + const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant(); + + GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples); + GGML_ASSERT(!tokens_as_samples || !fusion); const float * src1_d = (const float *) src1->data; const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; @@ -1354,6 +1378,17 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; + if (tokens_as_samples) { + GGML_ASSERT(ne03 == 1 && ne13 == 1 && ne3 == 1); + // one column, one sample per token: y advances by s12 per token, dst by s2, x not at all + mul_mat_vec_q_switch_type( + src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, + ne01, 1, s01, stride_col_y, stride_col_dst, + ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, + 1, ne2, s03, s12, s2, ids_stride, stream); + return; + } + mul_mat_vec_q_switch_type( src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4e6..67d69b1415d 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -4,6 +4,10 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); +// [TAG_BATCH_INVARIANT] true when an MMVQ launch of ncols_dst columns sums each destination element +// in the same order as a single-column launch, i.e. when the column count leaves nwarps unchanged +bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst); + // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, // based on the quantization type and GPU architecture (compute capability). int get_mmvq_mmid_max_batch(ggml_type type, int cc); diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b87b189a57a..b787108a7aa 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1266,6 +1266,11 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm (op->src[1]->ne[1] % op->src[4]->ne[1] == 0); break; case GGML_OP_FLASH_ATTN_EXT: + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + supported = false; + break; + } if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] && (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] && (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr && diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e8a5009b381..574276fed15 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4157,7 +4157,8 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons break; case GGML_OP_FLASH_ATTN_EXT: - supp = ggml_hexagon_supported_flash_attn_ext(sess, op); + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + supp = op->src[5] == nullptr && ggml_hexagon_supported_flash_attn_ext(sess, op); break; case GGML_OP_SET_ROWS: diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 19c57820e85..e5fc6a8063b 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1592,6 +1592,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_ROLL: return ggml_is_contiguous(op->src[0]); case GGML_OP_FLASH_ATTN_EXT: + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; + // walking the pool in physical order here would be silently wrong + if (op->src[5] != NULL) { + return false; + } // for new head sizes, add checks here if (op->src[0]->ne[0] != 32 && op->src[0]->ne[0] != 40 && diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 64f3325b2a5..9c2362c8c4b 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7842,6 +7842,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_MEAN: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return false; + } // The E17 compilers segfault while building FA kernels, skip E17 for now if (adreno_e17_compiler_quirks(backend_ctx)) { return false; diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index e299e16c778..1cec1583d51 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1128,6 +1128,10 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { break; } case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return true; + } float scale = 1.0f; float max_bias = 0.0f; float logit_softcap = 0.0f; diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..51b552acf75 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1915,7 +1915,11 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { GGML_UNUSED(dev); - GGML_UNUSED(op); + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; + // the remote end is not asked, so it is not claimed here + if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) { + return false; + } //TODO: call the remote backend and cache the results return true; } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834..a31a6a41ca8 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6342,7 +6342,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K; case GGML_OP_FLASH_ATTN_EXT: - return ggml_sycl_flash_attn_ext_supported(device, op); + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op); default: return false; } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1d86aaac5c..eb855f24d17 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18192,6 +18192,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return false; + } bool coopmat2 = device->coopmat2; uint32_t HSK = op->src[1]->ne[0]; uint32_t HSV = op->src[2]->ne[0]; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 2434848a55a..26717d8804d 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4408,6 +4408,12 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + supports_op = false; + break; + } + // conservative support checks for whether the more resource-intensive shader paths // can be used, to avoid cases where flash_attn is assigned to the CPU later on supports_op = src0->type == GGML_TYPE_F32 && diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..077f17ecabf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -795,6 +795,29 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); + // [TAG_EXACT_CONCURRENCY] cells the memory allocates in one indivisible unit: 1 ordinarily, + // larger where a mode places cells in blocks, and then n tokens occupy round_up(n, granularity) + // cells. A caller deciding whether the pool has room must round the same way. + LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); + + // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or 1 + // plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence + // set with more left to place is a prompt and is prefilled in a ubatch of its own; one at or + // below stays grouped with the other decodes. Process-wide, default 1, never lowered. Raising + // it widens every existing context's decode step and re-reports their width; returns false, + // and changes nothing, when an explicit column bound cannot cover that width. + LLAMA_API bool llama_set_exact_decode_tokens(uint32_t n_tokens); + LLAMA_API uint32_t llama_exact_decode_tokens(void); + + // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the + // sequences a context holds times the tokens each contributes. Every context reports its own + // at creation and a backend keeps the widest it has heard. A caller that builds wider steps + // reports the width itself, before the context or the first decode. Never lowered. Returns + // false, reporting nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is positive and below the + // width: that bound wins in the backend, so decodes above it would be left batched. + LLAMA_API bool llama_set_exact_decode_width(uint32_t n_cols); + LLAMA_API uint32_t llama_exact_decode_width(void); + // // State / sessions // diff --git a/scripts/batchinv/README.md b/scripts/batchinv/README.md new file mode 100644 index 00000000000..82b769b6d04 --- /dev/null +++ b/scripts/batchinv/README.md @@ -0,0 +1,71 @@ +# Exact concurrency experiment p + +Opt in before loading the model with `LLAMA_EXACT_CONCURRENCY=1`. This also forces +`GGML_CUDA_BATCH_INVARIANT=2` and gives `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` a +default. The column policy only has to cover the widest ubatch a decode step can +build, one column per slot times one plus the draft length, because a prompt +ubatch is kept to one sequence and gets its exactness from that instead. Tools +built on `common` report that width, so the default is `--parallel` times one plus +`--spec-draft-n-max`, and an explicitly set `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` +smaller than it is refused at startup. Nothing reported a width, the default is 16 +and the dispatcher warns once the first time a `MUL_MAT` or `MUL_MAT_ID` above the +bound is left unsplit. Set the variable to `0` for no bound; above the bound the +column policy does not fire, including during prefill. + +The experimental policy supports unified, offloaded F16 K/V, causal flash attention, +256-dimensional K and V heads, no attention soft cap, and no sliding window. +Shared-weight matmuls over multiple sequence planes are normalized to one plane +before the inherited selective column dispatcher. Without this, the recurrent +output projection bypasses batch invariance during concurrent prefill. +It is measured on text prompts with Qwen3.5-4B on one B200. Every KV layer has to +be on the CUDA backend, since no other backend reads the page table; a partial or +absent offload fails the load naming the layer. The V-less attention layouts have +no page table either, and a model on one of those is refused at context creation. +Context shifting, position division, cross-sequence prefix copies, shared-prefix +input tokens, and whole-context state loading are unsupported. Per-sequence state +save and restore is supported. Unsupported cache transformations are refused with +a logged error and leave the cells untouched; `--context-shift` and +`--cache-reuse` are reported as unsupported at load and disabled there. + +The allocator owns pages of 256 cells on behalf of one (sequence, position/256). +Position modulo 256 fixes the cell offset. Empty pages remain in the unified pool +and can be allocated by any sequence. The metadata is derived from live cells so +allocation rollback, tail removal, and sequence removal do not need another +transaction log. Restoring a sequence allocates free pages from this same pool. +The cost is up to 255 reserved cells per active sequence tail, plus holes introduced +by partial range removal. + +Attention receives an I32 page table in source 5: `[count, physical page IDs...]` +for each query, sorted by logical position. The physical K/V view and mask span +the pool, but the attention loop only visits the query's logical pages. The +final page is padded to 256 cells using the existing causal mask. Wholly future +pages in a prefill ubatch are excluded from the query's table. + +The vector attention specialization runs one query per block, four warps, and +`parallel_blocks=1`. It reads K/V directly from the physical pages, with two +128-cell softmax iterations per page, in logical page order. There is no K/V +gather and no split-K combine. The default vector specialization has no page +lookup. The ordinary path allocates no page metadata and launches no extra kernels. + +`FATTN_KQ_STRIDE=256` is a mask-scan stride, not the actual MMA rescaling tile. +For K/V head size 256, the Ampere-or-newer MMA configurations use 64 KV rows at +8 query/head columns and 32 KV rows at 16/32/64 columns. The retained vector path +has a 128-cell iteration. Both divide the 256-cell placement page. + +The probe is adapted from the existing batch-invariant harness and rejects +nonfinite logits and attention. `PROBE_B_REVERSE=1` fills neighbours before P0; +`PROBE_RESTORE=1` parks P0, releases a neighbour, restores P0, then rebuilds the +neighbour. Compute rows can be compared across this relocation; physical cache +views and index tensors must not be mistaken for sequence-0 compute outputs. + +`divergence.py --reference FILE` compares against an existing unparked solo token +reference. `bench.py --modes 0,1 --pairs 3` measures default off against exact mode +on, with 256 predicted tokens. Set `UNSLOTH_WORKSPACE` to the model parent workspace +and `LD_LIBRARY_PATH` to this build's bin directory. The harness uses GPU 3; +select a port in 9601-9610 explicitly. + +A concurrent request that fails now fails the run instead of being dropped from +the result, and every run records the environment the server actually inherited, +including `LLAMA_EXACT_CONCURRENCY`, under `env` in the JSON and in the server log +header, so a run labelled as the mode-off reference can be checked rather than +trusted. Teardown is POSIX only. diff --git a/scripts/batchinv/bench.py b/scripts/batchinv/bench.py new file mode 100644 index 00000000000..64d8b32244b --- /dev/null +++ b/scripts/batchinv/bench.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Cost of the knob: solo tok/s and four-chat aggregate tok/s, knob off and on, back to back.""" +import argparse, json, os, sys, time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from divergence import Server, completion, run_concurrent +from prompts import PROMPTS + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--binary", required=True) + ap.add_argument("--spec", default="none") + ap.add_argument("--n-predict", type=int, default=256) + ap.add_argument("--pairs", type=int, default=3) + ap.add_argument("--port", type=int, default=9602) + ap.add_argument("--modes", default="0,1") + ap.add_argument("--out", required=True) + a = ap.parse_args() + + modes = a.modes.split(",") + rows = [] + for pair in range(a.pairs): + for mode in modes: + env = {"LLAMA_EXACT_CONCURRENCY": mode, "GGML_CUDA_BATCH_INVARIANT": "0" if mode == "0" else "2"} + with Server(a.port, a.binary, [], env, a.out + ".server.log", a.spec) as s: + completion(a.port, PROMPTS["P0"], 32) # warm + solo = completion(a.port, PROMPTS["P0"], a.n_predict) + outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict) + row = { + "pair": pair, "mode": mode, "spec": a.spec, + "solo_tok_per_s": solo["timings"]["predicted_per_second"], + "solo_prompt_tok_per_s": solo["timings"]["prompt_per_second"], + "four_aggregate_tok_per_s": sum(o["timings"]["predicted_per_second"] for o in outs.values()), + "four_wall_s": wall, + "four_total_tokens": sum(len(o["tokens"]) for o in outs.values()), + } + row["four_wall_tok_per_s"] = row["four_total_tokens"] / wall + rows.append(row) + print(json.dumps(row), flush=True) + with open(a.out, "w") as f: + json.dump(rows, f, indent=2) + + print("\n=== summary ===", flush=True) + for mode in modes: + rs = [r for r in rows if r["mode"] == mode] + for k in ("solo_tok_per_s", "four_aggregate_tok_per_s", "four_wall_tok_per_s", "solo_prompt_tok_per_s"): + vals = sorted(r[k] for r in rs) + print(f"mode={mode} {k}: median {vals[len(vals)//2]:.1f} values {[round(v,1) for v in vals]}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py new file mode 100644 index 00000000000..198c2bdf881 --- /dev/null +++ b/scripts/batchinv/divergence.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +"""Baseline / patched divergence harness: solo P0 vs P0 sharing batches with P1..P3.""" +import argparse, json, os, signal, subprocess, sys, threading, time, urllib.request, urllib.error + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from prompts import PROMPTS + +# recorded with every run; LLAMA_EXACT_CONCURRENCY inherited from the shell decides whether a run +# labelled as the mode-off reference actually was one, so it is not optional +RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT", + "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY", + "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES") + +MODEL_REL = "models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf" + + +def model_path(): + """Resolved when the server args are built, so --help works without the variable set.""" + ws = os.environ.get("UNSLOTH_WORKSPACE") + if not ws: + raise RuntimeError("UNSLOTH_WORKSPACE is not set; it must point at the workspace holding " + + MODEL_REL) + return os.path.join(ws, MODEL_REL) + + +def post(port, path, payload, timeout=1800): + req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read().decode()) + + +def get(port, path, timeout=10): + with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=timeout) as r: + return json.loads(r.read().decode()) + + +def completion(port, prompt, n_predict): + return post(port, "/completion", { + "prompt": prompt, "n_predict": n_predict, "temperature": 0.0, "top_k": 1, + "top_p": 1.0, "min_p": 0.0, "typical_p": 1.0, "seed": 0, + "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0, + "cache_prompt": False, "return_tokens": True, "samplers": ["top_k", "temperature"], + }) + + +class Server: + def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=True): + self.port, self.log_path = port, log_path + self.args = [binary, "-m", model_path(), "--port", str(port), "--host", "127.0.0.1", + "--parallel", "4", "-c", "8192", + "--flash-attn", "on", "--metrics", "-ngl", "99", "--no-warmup", + "--seed", "0", "--spec-type", spec] + if kv_unified: + self.args += ["--kv-unified"] + if spec == "draft-mtp": + self.args += ["--spec-draft-n-max", "2"] + self.args += extra + self.env = dict(os.environ) + self.env["CUDA_VISIBLE_DEVICES"] = "3" + self.env.update(env_extra) + # what the server will actually see, not what this run meant to set + self.env_resolved = {k: self.env[k] for k in RECORDED_ENV if k in self.env} + self.p = None + self.fh = None + + def __enter__(self): + self.fh = open(self.log_path, "ab") + self.fh.write(("\n=== " + " ".join(self.args) + "\n=== env " + + json.dumps(self.env_resolved) + "\n").encode()) + self.fh.flush() + self.p = subprocess.Popen(self.args, stdout=self.fh, stderr=subprocess.STDOUT, + env=self.env, start_new_session=True) + print(f"[server] pid={self.p.pid} port={self.port} log={self.log_path}", flush=True) + try: + deadline = time.time() + 600 + while time.time() < deadline: + if self.p.poll() is not None: + raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}") + try: + if get(self.port, "/health").get("status") == "ok": + print("[server] ready", flush=True) + return self + except Exception: + time.sleep(1.0) + raise RuntimeError("server did not become healthy") + except BaseException: + # __exit__ is not called when __enter__ raises, and a server that started but never + # reported healthy would keep the GPU, the port and the log handle + self.__exit__(None, None, None) + raise + + def __exit__(self, *a): + # note: POSIX only; Windows would need CREATE_NEW_PROCESS_GROUP at Popen + if self.p is not None: + print(f"[server] stopping pid={self.p.pid}", flush=True) + try: + os.killpg(os.getpgid(self.p.pid), signal.SIGTERM) + self.p.wait(timeout=60) + except Exception: + try: + os.killpg(os.getpgid(self.p.pid), signal.SIGKILL) + except Exception: + pass + try: + self.p.wait(timeout=60) + except Exception: + pass + self.p = None + if self.fh is not None: + self.fh.close() + self.fh = None + + +def run_concurrent(port, names, n_predict): + barrier = threading.Barrier(len(names)) + lock = threading.Lock() + out = {} + errors = [] + + def work(name): + try: + barrier.wait() + res = completion(port, PROMPTS[name], n_predict) + except BaseException as e: + with lock: + errors.append((name, e)) + # release the others rather than let them block on a barrier that will never fill + barrier.abort() + return + with lock: + out[name] = res + + ts = [threading.Thread(target=work, args=(n,)) for n in names] + t0 = time.time() + for t in ts: + t.start() + for t in ts: + t.join() + wall = time.time() - t0 + + # without this a run where P1..P3 failed and P0 succeeded reads as a clean four-way result + if errors: + raise RuntimeError("concurrent requests failed: " + + "; ".join(f"{n}: {type(e).__name__}: {e}" for n, e in errors)) + missing = set(names) - set(out) + if missing: + raise RuntimeError(f"concurrent requests produced no result for {sorted(missing)}") + + return out, wall + + +def first_diff(a, b): + for i, (x, y) in enumerate(zip(a, b)): + if x != y: + return i + return None if len(a) == len(b) else min(len(a), len(b)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--reference") + ap.add_argument("--label", required=True) + ap.add_argument("--port", type=int, default=9601) + ap.add_argument("--binary", required=True) + ap.add_argument("--spec", default="none") + ap.add_argument("--n-predict", type=int, default=512) + ap.add_argument("--repeats", type=int, default=3) + ap.add_argument("--env", action="append", default=[]) + ap.add_argument("--extra", action="append", default=[]) + ap.add_argument("--out", required=True) + ap.add_argument("--no-kv-unified", action="store_true") + a = ap.parse_args() + + env_extra = dict(kv.split("=", 1) for kv in a.env) + server = Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec, + kv_unified=not a.no_kv_unified) + res = {"label": a.label, "spec": a.spec, "n_predict": a.n_predict, + "env_requested": env_extra, "env": server.env_resolved, + "model": server.args[2], "args": server.args, + "extra": a.extra, "binary": a.binary, + "kv_unified": not a.no_kv_unified} + + with server as s: + solo = completion(a.port, PROMPTS["P0"], a.n_predict) + ref = json.load(open(a.reference))["tokens"] if a.reference else solo["tokens"] + res["solo_first_diff"] = first_diff(ref, solo["tokens"]) + res["reference"] = a.reference + res["solo"] = {"n_tokens": len(ref), "tok_per_s": solo["timings"]["predicted_per_second"], + "text_sha": None} + # solo repeat, to prove solo itself is stable + solo2 = completion(a.port, PROMPTS["P0"], a.n_predict) + res["solo_repeat_first_diff"] = first_diff(ref, solo2["tokens"]) + res["rounds"] = [] + for r in range(a.repeats): + outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict) + p0 = outs["P0"]["tokens"] + fd = first_diff(ref, p0) + agg = sum(outs[n]["timings"]["predicted_per_second"] for n in outs) + row = {"round": r, "first_diff": fd, "n_tokens": len(p0), + "identical": fd is None, "wall_s": wall, + "p0_tok_per_s": outs["P0"]["timings"]["predicted_per_second"], + "aggregate_tok_per_s": agg, + "per_req_n": {n: len(outs[n]["tokens"]) for n in outs}, "p0_tokens": p0} + res["rounds"].append(row) + print(f"[round {r}] first_diff={fd} identical={fd is None} wall={wall:.1f}s agg={agg:.1f} tok/s", flush=True) + with urllib.request.urlopen(f"http://127.0.0.1:{a.port}/metrics") as response: + res["metrics"] = response.read().decode() + with open(a.out + ".p0_solo.json", "w") as f: + json.dump({"tokens": ref, "content": solo["content"]}, f) + + with open(a.out, "w") as f: + json.dump(res, f, indent=2) + print(json.dumps({k: v for k, v in res.items() if k != "rounds"}, indent=2), flush=True) + print(json.dumps(res["rounds"], indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp new file mode 100644 index 00000000000..204379c423b --- /dev/null +++ b/scripts/batchinv/probe.cpp @@ -0,0 +1,370 @@ +// Locate the first graph op whose sequence-0 output changes when the decode batch holds four +// sequences instead of one. Seq 0's prompt KV is identical in both phases, so the only +// difference is the width of the final decode ubatch. +#include "llama.h" +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct node_rec { + std::string name; + std::string op; + std::string tname; + int64_t ne[4]; + int64_t gdn_tokens = 0, gdn_seqs = 0; + size_t esize = 0; // bytes per element, 0 = not byte comparable + std::vector data; // empty when skipped + bool contiguous = false; + uint64_t hash = 0; +}; + +static bool g_record = false; +static std::vector * g_sink = nullptr; + +static uint64_t fnv1a(const uint8_t * p, size_t n) { + uint64_t h = 1469598103934665603ULL; + for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ULL; } + return h; +} + +static bool eval_cb(struct ggml_tensor * t, bool ask, void * /*ud*/) { + if (!g_record) return false; + if (ask) return true; + + node_rec r; + r.name = ggml_get_name(t); + r.tname = ggml_type_name(t->type); + r.op = t->op == GGML_OP_NONE ? "LEAF" : ggml_op_name(t->op); + if (t->op == GGML_OP_UNARY) r.op = std::string("UNARY_") + ggml_unary_op_name(ggml_get_unary_op(t)); + if (t->op == GGML_OP_GLU) r.op = std::string("GLU_") + ggml_glu_op_name(ggml_get_glu_op(t)); + for (int i = 0; i < 4; ++i) r.ne[i] = t->ne[i]; + r.contiguous = ggml_is_contiguous(t); + if (r.name == "linear_attn_out-0") { + fprintf(stderr, "linear_attn_out-0: weight=%s input=[%lld,%lld,%lld,%lld]\n", + ggml_type_name(t->src[0]->type), (long long)t->src[1]->ne[0], + (long long)t->src[1]->ne[1], (long long)t->src[1]->ne[2], (long long)t->src[1]->ne[3]); + } + if (t->op == GGML_OP_GATED_DELTA_NET) { + r.gdn_tokens = t->src[2]->ne[2]; + r.gdn_seqs = t->src[2]->ne[3]; + } + + const size_t nbytes = ggml_nbytes(t); + if (r.contiguous && ggml_blck_size(t->type) == 1 && nbytes <= (256u << 20)) { + r.esize = ggml_type_size(t->type); + r.data.resize(nbytes); + ggml_backend_tensor_get(t, r.data.data(), 0, nbytes); + if (t->op == GGML_OP_FLASH_ATTN_EXT) { + for (size_t i = 0; i < nbytes/sizeof(float); ++i) { + float v; memcpy(&v, r.data.data() + i*sizeof(float), sizeof(float)); + if (!std::isfinite(v)) { fprintf(stderr, "nonfinite attention: %s\n", t->name); exit(5); } + } + } + r.hash = fnv1a(r.data.data(), nbytes); + if (nbytes > (64u << 20)) { r.data.clear(); } // keep the hash only for the big ones + } + g_sink->push_back(std::move(r)); + return true; +} + +static std::string slurp(const char * path) { + std::ifstream f(path); + std::stringstream ss; ss << f.rdbuf(); + std::string s = ss.str(); + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + return s; +} + +static std::vector tokenize(const llama_vocab * v, const std::string & s) { + std::vector out(s.size() + 16); + int n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); + if (n < 0) { out.resize(-n); n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); } + out.resize(n); + return out; +} + +struct batch_holder { + std::vector tok; + std::vector pos; + std::vector nsid; + std::vector sid; + std::vector sidp; + std::vector out; + llama_batch get() { + sidp.resize(tok.size()); + for (size_t i = 0; i < tok.size(); ++i) sidp[i] = &sid[i]; + llama_batch b{}; + b.n_tokens = (int32_t) tok.size(); + b.token = tok.data(); b.pos = pos.data(); b.n_seq_id = nsid.data(); + b.seq_id = sidp.data(); b.logits = out.data(); + return b; + } +}; + +static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) { + const float * l = llama_get_logits_ith(ctx, i); + for (int k = 0; k < n_vocab; ++k) { + if (!std::isfinite(l[k])) { fprintf(stderr, "nonfinite logits at %d\n", k); exit(4); } + } + int best = 0; + for (int k = 1; k < n_vocab; ++k) if (l[k] > l[best]) best = k; + return best; +} + +// feed a prompt as one decode call for one sequence, return the greedy next token +static llama_token feed(llama_context * ctx, const std::vector & p, llama_seq_id seq, int n_vocab) { + batch_holder h; + for (size_t i = 0; i < p.size(); ++i) { + h.tok.push_back(p[i]); h.pos.push_back((llama_pos) i); + h.nsid.push_back(1); h.sid.push_back(seq); + h.out.push_back(i + 1 == p.size()); + } + llama_batch b = h.get(); + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "decode failed\n"); exit(1); } + return greedy(ctx, (int32_t) p.size() - 1, n_vocab); +} + +int main(int argc, char ** argv) { + const bool prefill = getenv("PROBE_PREFILL") != nullptr; + const char * model_path = argv[1]; + const int n_seqs = argc > 2 ? atoi(argv[2]) : 4; // width of the probed decode batch + const char * out_path = argc > 3 ? argv[3] : nullptr; + std::vector prompts; + for (int i = 4; i < argc; ++i) prompts.push_back(slurp(argv[i])); + + llama_backend_init(); + llama_model_params mp = llama_model_default_params(); + mp.n_gpu_layers = 99; + llama_model * model = llama_model_load_from_file(model_path, mp); + if (!model) { fprintf(stderr, "model load failed\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + + std::vector> ptok; + for (auto & s : prompts) ptok.push_back(tokenize(vocab, s)); + for (size_t i = 0; i < ptok.size(); ++i) fprintf(stderr, "prompt %zu: %zu tokens\n", i, ptok[i].size()); + + auto make_ctx = [&]() { + llama_context_params cp = llama_context_default_params(); + cp.n_ctx = 8192; cp.n_batch = 2048; cp.n_ubatch = 512; + if (prefill) { cp.n_ubatch = 2048; } + cp.n_seq_max = 4; cp.kv_unified = true; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + cp.cb_eval = eval_cb; cp.cb_eval_user_data = nullptr; + cp.no_perf = true; + return llama_init_from_model(model, cp); + }; + + std::vector rec_a, rec_b; + llama_token first_tok[4] = {0, 0, 0, 0}; + + // phase A: decode ubatch width 1. PROBE_A_FILL is how many sequences are already in the + // shared KV cache, which sets K->ne[1] for attention. + const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1; + // PROBE_A_PERM reorders which prompt goes into which sequence in phase A, keeping the cache + // length but changing what the masked cells hold. It may only move the neighbours: sequence 0 + // keeps prompt 0, or the two phases would compare different sequences. + int a_perm[4] = {0, 1, 2, 3}; + if (const char * perm = getenv("PROBE_A_PERM")) { + for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0'; + if (a_perm[0] != 0) { fprintf(stderr, "PROBE_A_PERM must keep prompt 0 on sequence 0\n"); return 1; } + } + { + llama_context * ctx = make_ctx(); + g_sink = &rec_a; g_record = prefill; + for (int s = 0; s < a_fill; ++s) { + const llama_token t = feed(ctx, ptok[a_perm[s]], s, n_vocab); + if (a_perm[s] == 0) first_tok[0] = t; + } + batch_holder h; + h.tok = {first_tok[0]}; h.pos = {(llama_pos) ptok[0].size()}; + h.nsid = {1}; h.sid = {0}; h.out = {1}; + llama_batch b = h.get(); + g_sink = &rec_a; g_record = !prefill; + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "A decode failed\n"); return 1; } + g_record = false; + llama_free(ctx); + } + + // phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens + { + llama_context * ctx = make_ctx(); + if (prefill) { + batch_holder h; + for (int seq = 0; seq < n_seqs; ++seq) { + for (size_t i = 0; i < ptok[0].size(); ++i) { + h.tok.push_back(ptok[seq][i%ptok[seq].size()]); h.pos.push_back(i); + h.nsid.push_back(1); h.sid.push_back(seq); h.out.push_back(i+1 == ptok[0].size()); + } + } + auto b = h.get(); + g_sink = &rec_b; g_record = true; + if (llama_decode(ctx, b) != 0) { return 6; } + g_record = false; + for (int seq = 0; seq < n_seqs; ++seq) { + first_tok[seq] = greedy(ctx, (seq+1)*ptok[0].size()-1, n_vocab); + } + } else for (int k = 0; k < n_seqs; ++k) { + const int s = getenv("PROBE_B_REVERSE") ? n_seqs - 1 - k : k; + first_tok[s] = feed(ctx, ptok[s], s, n_vocab); + } + if (getenv("PROBE_RESTORE")) { + std::vector state(llama_state_seq_get_size(ctx, 0)); + if (llama_state_seq_get_data(ctx, state.data(), state.size(), 0) != state.size()) { return 2; } + llama_memory_seq_rm(llama_get_memory(ctx), 0, -1, -1); + llama_memory_seq_rm(llama_get_memory(ctx), 1, -1, -1); + if (llama_state_seq_set_data(ctx, state.data(), state.size(), 0) != state.size()) { return 3; } + first_tok[1] = feed(ctx, ptok[1], 1, n_vocab); + } + if (first_tok[0] != 0 && rec_a.size()) {} + batch_holder h; + for (int s = 0; s < n_seqs; ++s) { + h.tok.push_back(first_tok[s]); h.pos.push_back((llama_pos) ptok[s].size()); + h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1); + } + llama_batch b = h.get(); + g_sink = &rec_b; g_record = !prefill; + if (!prefill && llama_decode(ctx, b) != 0) { fprintf(stderr, "B decode failed\n"); return 1; } + g_record = false; + llama_free(ctx); + } + + // optional: keep decoding and report the first step at which seq 0's token differs + const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0; + int first_bad_step = -1; + if (n_steps > 0) { + std::vector tok_a, tok_b; + for (int phase = 0; phase < 2; ++phase) { + const int fill = phase == 0 ? a_fill : n_seqs; + const int width = phase == 0 ? 1 : n_seqs; + std::vector & out = phase == 0 ? tok_a : tok_b; + llama_context * ctx = make_ctx(); + std::vector next(4, 0); + std::vector pos(4, 0); + for (int s = 0; s < fill; ++s) { + const int p = phase == 0 ? a_perm[s] : s; + next[s] = feed(ctx, ptok[p], s, n_vocab); + pos[s] = (llama_pos) ptok[p].size(); + } + for (int step = 0; step < n_steps; ++step) { + batch_holder h; + for (int s = 0; s < width; ++s) { + h.tok.push_back(next[s]); h.pos.push_back(pos[s]); + h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1); + } + llama_batch b = h.get(); + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "step decode failed\n"); exit(1); } + out.push_back(next[0]); + for (int s = 0; s < width; ++s) { next[s] = greedy(ctx, s, n_vocab); pos[s] += 1; } + } + llama_free(ctx); + } + for (int i = 0; i < n_steps; ++i) { + if (tok_a[i] != tok_b[i]) { first_bad_step = i; break; } + } + fprintf(stderr, "steps: %d first differing step: %d\n", n_steps, first_bad_step); + } + + fprintf(stderr, "nodes: A=%zu B=%zu first tokens: %d %d %d %d\n", + rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]); + + // walk both node lists in order and compare seq 0's slice + FILE * out = out_path ? fopen(out_path, "w") : stdout; + fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size()); + size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size(); + int emitted = 0; + for (size_t i = 0; i < n; ++i) { + const node_rec & A = rec_a[i]; + const node_rec & B = rec_b[i]; + const char * verdict = nullptr; + double max_abs = 0.0; + size_t ndiff = 0, ncmp = 0; + + if (A.name != B.name || A.op != B.op) { + verdict = "misaligned"; + } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens && + !A.data.empty() && !B.data.empty()) { + // packed GDN outputs put all token outputs before all sequence states, so seq 0's + // state moves when the number of sequences changes + const size_t output = A.ne[0]*A.gdn_tokens; + const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output; + for (size_t k = 0; k < output + state; ++k) { + const size_t ia = k < output ? k : A.gdn_seqs*output + k-output; + const size_t ib = k < output ? k : B.gdn_seqs*output + k-output; + float va, vb; + memcpy(&va, A.data.data()+ia*4, 4); memcpy(&vb, B.data.data()+ib*4, 4); + ++ncmp; + if (memcmp(&va, &vb, 4)) { + ++ndiff; + if (std::abs(double(va)-vb) > max_abs) { max_abs = std::abs(double(va)-vb); } + } + } + verdict = ndiff ? "row-differs" : nullptr; + } else if (A.esize == 0 || B.esize == 0 || A.esize != B.esize) { + verdict = "skipped"; + } else { + int tdim = -1; bool same = true; + for (int d = 0; d < 4; ++d) { + if (A.ne[d] == B.ne[d]) continue; + same = false; + if (B.ne[d] == n_seqs*A.ne[d] && tdim < 0) tdim = d; else { tdim = -2; break; } + } + if (tdim == -2) { + verdict = "shape-incomparable"; + } else if (same) { + verdict = (A.hash == B.hash) ? nullptr : "whole-tensor-differs"; + } else if (A.data.empty() || B.data.empty()) { + verdict = "too-large"; + } else { + // compare element (.., i_tdim = 0, ..) across all other indices + int64_t st[4] = {1, A.ne[0], A.ne[0]*A.ne[1], A.ne[0]*A.ne[1]*A.ne[2]}; + int64_t stb[4] = {1, B.ne[0], B.ne[0]*B.ne[1], B.ne[0]*B.ne[1]*B.ne[2]}; + for (int64_t i3 = 0; i3 < A.ne[3]; ++i3) + for (int64_t i2 = 0; i2 < A.ne[2]; ++i2) + for (int64_t i1 = 0; i1 < A.ne[1]; ++i1) + for (int64_t i0 = 0; i0 < A.ne[0]; ++i0) { + int64_t idx[4] = {i0, i1, i2, i3}; + + size_t oa = 0, ob = 0; + for (int d = 0; d < 4; ++d) { oa += idx[d]*st[d]; ob += idx[d]*stb[d]; } + ncmp++; + const uint8_t * pa = A.data.data() + oa*A.esize; + const uint8_t * pb = B.data.data() + ob*B.esize; + if (memcmp(pa, pb, A.esize) != 0) { + ndiff++; + if (A.esize == 4) { + float fa, fb; memcpy(&fa, pa, 4); memcpy(&fb, pb, 4); + double d2 = fa - fb; if (d2 < 0) d2 = -d2; + if (d2 > max_abs) max_abs = d2; + } + } + } + verdict = ndiff ? "row-differs" : nullptr; + } + } + { + if (emitted++) fprintf(out, ","); + fprintf(out, "\n{\"i\":%zu,\"name\":\"%s\",\"op\":\"%s\",\"ne_a\":[%lld,%lld,%lld,%lld]," + "\"ne_b\":[%lld,%lld,%lld,%lld],\"type\":\"%s\",\"verdict\":\"%s\",\"ndiff\":%zu,\"ncmp\":%zu,\"max_abs\":%.6g}", + i, A.name.c_str(), A.op.c_str(), + (long long)A.ne[0],(long long)A.ne[1],(long long)A.ne[2],(long long)A.ne[3], + (long long)B.ne[0],(long long)B.ne[1],(long long)B.ne[2],(long long)B.ne[3], + A.tname.c_str(), verdict ? verdict : "same", ndiff, ncmp, max_abs); + } + } + fprintf(out, "\n]}\n"); + if (out_path) fclose(out); + + llama_model_free(model); + llama_backend_free(); + return 0; +} diff --git a/scripts/batchinv/prompts.py b/scripts/batchinv/prompts.py new file mode 100644 index 00000000000..efabc448a9d --- /dev/null +++ b/scripts/batchinv/prompts.py @@ -0,0 +1,53 @@ +# four distinct prompts, each about 300 tokens of raw text (no chat template) +_BODIES = { +"P0": """The history of numerical computing is a history of compromises between speed and exactness. +Early machines used fixed point arithmetic because it was cheap, and programmers carried scaling +factors in their heads. Floating point hardware moved the bookkeeping into silicon, but it did not +remove the compromise, it only hid it. Addition of floating point numbers is commutative but it is +not associative, so the order in which a long sum is accumulated changes the last few bits of the +result. On a single processor that order is fixed by the program text and nobody notices. On a +parallel processor the order is fixed by how the work was divided, and the division is chosen for +speed, not for reproducibility. A reduction split across two warps sums a different set of partial +products than the same reduction split across four warps, and the two answers differ in the low +bits. Nothing is wrong with either answer. Both are within a fraction of an ulp of the exact value. +The trouble begins when a downstream decision is discrete. A comparison, a rounding to an integer, +or the selection of the largest element of a vector turns a difference of one bit into a difference +of one branch, and from there the two computations walk away from each other and never come back. +Explain, carefully and at length, why this matters for a system that serves many users at once, +and what an engineer would have to give up to make the answer depend only on the request and not +on what else the machine happened to be doing at the time. Discuss the cost.""", +"P1": """Consider a public library that lends physical books and must decide how many copies of a +popular title to buy. The librarian has a fixed budget, a waiting list that grows and shrinks, and +a shelf that is already full. Every copy purchased shortens the queue for that title and lengthens +the queue for every other title, because the money and the shelf space are shared. The obvious +policy, buy copies of whatever has the longest queue, is unstable, because a title that briefly +becomes fashionable will absorb the whole budget and then sit unread for a decade. A better policy +has to weigh how long the demand is likely to last against how long the book will remain useful, +and it has to do this with almost no information. Describe in detail how you would design such a +policy, what data you would collect, how you would test it without harming readers, and how you +would know whether it was working. Consider what happens when the budget is cut in half without +warning, when a title is suddenly assigned as required reading by a local school, and when the +shelf itself must shrink because the building is being renovated. Explain the tradeoffs plainly.""", +"P2": """A small coastal town has one bridge to the mainland and it is failing. The engineers say it +has perhaps eight years left. Replacing it costs more than the town has ever spent on anything. +Repairing it buys maybe four years and costs a third as much, and the repair work closes the bridge +for two months in the summer, which is when the town earns most of its money. Doing nothing is +free until the day it is not. The town council is split, the ferry operator has opinions, and the +regional government will match funds only for a replacement, only if construction begins within +three years, and only if the town covers the first quarter of the cost itself. Write a long and +careful analysis of the options available to the council. Identify the assumptions that matter +most, the ones where being wrong changes the recommendation, and say how the council could cheaply +find out whether those assumptions hold. Then give a recommendation and state honestly what would +have to be true for the recommendation to be wrong. Do not hedge. Commit to an answer at the end.""", +"P3": """Describe the process by which a large body of water freezes over in winter, beginning with +the surface layer and working downward, and explain why the ice floats rather than sinking, why a +deep lake takes much longer to freeze than a shallow one of the same surface area, and why the +temperature at the bottom of a frozen lake settles near four degrees Celsius rather than at zero. +Then explain what this means for the animals that live there, how fish survive a winter under a +solid lid, why a heavy snowfall on top of the ice can be more dangerous to them than the cold +itself, and what happens in the spring when the whole column overturns. Use plain language and +avoid equations. Where a common explanation is wrong or incomplete, say so and give the better one. +Be thorough. Assume the reader is curious and patient but has no training in physics or biology, +and would rather understand one thing properly than be told five things quickly.""", +} +PROMPTS = {k: " ".join(v.split()) for k, v in _BODIES.items()} diff --git a/scripts/unsloth/additive_merge.py b/scripts/unsloth/additive_merge.py index 5364dc1b7c1..d930158a542 100644 --- a/scripts/unsloth/additive_merge.py +++ b/scripts/unsloth/additive_merge.py @@ -94,18 +94,11 @@ def nonblank(lines: list[str]) -> list[str]: return [ln.strip() for ln in lines if ln.strip()] -# A line that closes or opens a block and nothing else. Two INDEPENDENT case -# arms in the same switch share these by construction -- `{`, `} break;`, `}` -# are what a case arm is made of, not what makes it that case arm -- so finding -# them on both sides says nothing about whether the two sides added the same -# construct. Matching them as "shared" is what refused the real add/add of -# PROJECTOR_TYPE_KIMIK3 next to PROJECTOR_TYPE_DEEPSEEK4V in tools/mtmd/clip.cpp -# with "one change made twice: {, } break;", when the two arms had no line of -# actual content in common. -# -# Deliberately narrow: braces, brackets, parens, semicolons and commas, around -# at most one bare block-terminating keyword. `break;` matches, `return true;` -# does not, and anything naming a type, a constant or a function does not. +# A line that only opens or closes a block. Two independent case arms share these by +# construction, so finding them on both sides says nothing about the two sides adding the +# same construct: treating them as shared is what refused the real PROJECTOR_TYPE_KIMIK3 / +# PROJECTOR_TYPE_DEEPSEEK4V add/add in tools/mtmd/clip.cpp. Deliberately narrow: brackets, +# semicolons and commas around at most one bare block-terminating keyword. STRUCTURAL = re.compile(r"^[\s{}()\[\];,]*(?:break|continue|return|pass)?[\s{}()\[\];,]*$") @@ -114,8 +107,7 @@ def identifying(lines: list[str]) -> set[str]: return {ln for ln in nonblank(lines) if not STRUCTURAL.match(ln)} -# `case FOO:`, `case FOO :`, `default:`. A fallthrough label may carry no body -# at all, which is the shape the nightly hits most often. +# `case FOO:`, `case FOO :`, `default:`; a fallthrough label may carry no body at all CASE_LABEL = re.compile(r"^(?:case\s+[^:]+|default\s*):") @@ -147,32 +139,23 @@ def resolve_region(ours: list[str], base: list[str], theirs: list[str]) -> list[ return list(ours) ours_arms, theirs_arms = case_arms(ours), case_arms(theirs) if ours_arms and theirs_arms and ours_arms.isdisjoint(theirs_arms): - # Both sides added case arms, and not one label is on both sides. Two - # arms of the same switch labelled differently are two constructs, so - # any line they happen to share is body text, not a duplicate: the real - # tools/mtmd/clip.cpp collision has a KIMIK3 arm and a DEEPSEEK4V arm - # that both set `hparams.rope_theta = 10000.0f;`, and refusing on that - # coincidence is what the shared-line check is for, backwards. - # - # The same change made twice would keep its label, so it lands in the - # check below instead. This is the one place where a shared line is - # allowed, and it is allowed because the labels prove the arms are - # distinct -- a duplicated label would not even compile. + # Both sides added case arms and no label is on both, so they are two constructs + # and any line they share is body text: the real clip.cpp collision has arms that + # both set `hparams.rope_theta = 10000.0f;`. The same change made twice would keep + # its label and land in the check below, so this is the one place a shared line is + # allowed - a duplicated label would not even compile. return list(theirs) + list(ours) shared = identifying(ours) & identifying(theirs) if shared: # Overlapping content is the signature of one construct added twice, # not two independent additions. Unioning it would duplicate code. - # Scaffolding lines are excluded above, so what is left is content both - # sides genuinely wrote, which is the thing that makes this a duplicate. + # scaffolding is excluded above, so what is left is content both sides wrote raise Unresolvable( "both sides add the same line(s), so this is one change made twice: " + ", ".join(sorted(shared)[:3]) ) if not identifying(ours) or not identifying(theirs): - # Everything one side added is scaffolding, so there is no content to - # tell the two additions apart and the exclusion above has nothing left - # to work with. Refuse rather than union braces onto braces. + # one side is all scaffolding, so there is no content to tell the additions apart raise Unresolvable( "one side adds only block scaffolding, so the two additions cannot " "be told apart" diff --git a/scripts/unsloth/feature_matrix.py b/scripts/unsloth/feature_matrix.py index 00b392ca1ec..d8fad55d469 100644 --- a/scripts/unsloth/feature_matrix.py +++ b/scripts/unsloth/feature_matrix.py @@ -37,7 +37,7 @@ import sys from pathlib import Path -# Output that means "this did not run" from a process that exited 0. +# output that means "this did not run" from a process that exited 0 SKIP_RE = re.compile(r"\bSKIP\b|not supported|unsupported|no tests|0 tests", re.I) @@ -78,7 +78,7 @@ def probe_arch(check: dict, b: Path, gpu: bool) -> str: rc, out = run([str(b / "test-llama-archs"), "-a", arch, "-s", "1234"], b, gpu) if rc != 0: raise Unproven(f"test-llama-archs -a {arch} exited {rc}") - # The arch's own rows, not the header and not another arch's. + # the arch's own rows, not the header and not another arch's rows = [ln for ln in out.splitlines() if ln.strip().startswith("|") and f"|{arch:>16}|" in ln or (ln.strip().startswith("|") and ln.split("|")[1].strip() == arch)] if not rows: @@ -119,8 +119,8 @@ def probe_mtmd(check: dict, b: Path, gpu: bool) -> str: m = re.search(r"assertions\s*:\s*(\d+)", out) if not m or int(m.group(1)) == 0: raise Unproven("test_projector_registry ran no assertions; the filter matched nothing") - # The registry test walks the whole enum, so it proves the table is sound. - # That the specific projector is IN the enum is pin_contract.py's job. + # the test walks the whole enum, so it proves the table is sound; that this projector + # is IN the enum is pin_contract.py's job return f"projector registry intact over {m.group(1)} assertions" @@ -173,8 +173,7 @@ def main() -> int: print(f"ok {name}: " + "; ".join(r["evidence"] for r in entry["results"]) + (f" [{len(entry['deferred'])} needs a GPU]" if entry["deferred"] else "")) else: - # Nothing was shown either way. Not a failure here, but it must not - # read as one of the ok lines. + # nothing shown either way: not a failure, but not an ok line either print(f"-- {name}: nothing provable without a GPU " f"({len(entry['deferred'])} check(s) deferred)") @@ -188,8 +187,7 @@ def main() -> int: if failed: print(f"\n{failed} feature(s) could not be shown to work", file=sys.stderr) return 1 - # Say what was NOT proven in the same breath as what was. A run that only - # ever prints a success line teaches the reader that green means covered. + # say what was NOT proven alongside what was, or green starts to read as covered tail = f", {deferred} check(s) need a GPU and were not run" if deferred else "" print(f"\nall {len(report['features'])} features demonstrated" + (" on GPU" if args.gpu else " on CPU") + tail) diff --git a/scripts/unsloth/pin_contract.py b/scripts/unsloth/pin_contract.py index bf07d58408e..d6a60746098 100644 --- a/scripts/unsloth/pin_contract.py +++ b/scripts/unsloth/pin_contract.py @@ -57,34 +57,28 @@ r"^https://github\.com/([^/]+)/llama\.cpp/pull/(\d+)/commits/([0-9a-f]{40})/?$" ) -# Identifier families that name a FEATURE. Deliberately not "every new symbol": -# a helper function renamed by a later upstream commit is not a lost feature, -# but a missing LLM_ARCH_ entry always is. These are the tables that decide -# whether an architecture, an op, a projector or a quant type exists at all. +# Identifier families that name a FEATURE, not every new symbol: a renamed helper is not a +# lost feature, but a missing LLM_ARCH_ entry always is. SYMBOL_FAMILIES = ( "LLM_ARCH_", "LLM_TENSOR_", "LLM_KV_", "LLM_TYPE_", "PROJECTOR_TYPE_", "GGML_OP_", "GGML_TYPE_", "LLAMA_FTYPE_", ) SYMBOL_RE = re.compile(r"\b(?:" + "|".join(SYMBOL_FAMILIES) + r")[A-Z0-9_]+\b") -# The subset that names a whole feature rather than one of its tensors. Used -# only to keep --emit readable; the check itself uses all of SYMBOL_FAMILIES. +# the subset naming a whole feature; only to keep --emit readable, the check uses them all HEADLINE = ("LLM_ARCH_", "GGML_OP_", "GGML_TYPE_", "PROJECTOR_TYPE_", "LLAMA_FTYPE_") -# A line worth tracking for survival. Comments and short punctuation drift with -# every reformat and would make the check noise; a substantial code line does -# not move on its own. +# a line worth tracking for survival: comments and short punctuation drift with every +# reformat, a substantial code line does not move on its own TRIVIAL_RE = re.compile(r"^\s*(?://|/\*|\*|\*/|#\s|$)") MIN_LINE = 12 -# Comments are stripped before anything is read off a line. A pin that merely -# NAMES an arch in a comment has not registered it, and holding the comment's -# wording as a contract fails the moment upstream rewords it. Observed on -# unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to explain why it does -# NOT use it. +# Comments are stripped first: a pin that merely NAMES an arch in a comment has not +# registered it, and holding the wording as a contract fails when upstream rewords it. +# Observed on unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to say it is not used. COMMENT_RE = re.compile(r"//.*$|/\*.*?\*/|(? dict: if code: symbols[cur].update(SYMBOL_RE.findall(code)) - # Only symbols the base does not ALREADY have in that file are evidence of - # this pin. Upstream naming an arch in a file the pin also touches is not - # something the pin is owed. + # only symbols the base does not already have in that file are evidence of this pin new_symbols: dict[str, list[str]] = {} for path, names in symbols.items(): fresh = sorted(n for n in names @@ -306,9 +298,8 @@ def main() -> int: if args.emit: report["pins"].append(entry) - # Only the families that NAME a feature are printed. Every symbol - # is still checked; a new file legitimately contributes a hundred - # LLM_TENSOR_ names and listing them buries the one that matters. + # only feature-naming families are printed; all are still checked, but a + # hundred LLM_TENSOR_ names would bury the one that matters sym = sorted({s for v in contract["symbols"].values() for s in v if s.startswith(HEADLINE)}) print(f"{name:>18} {entry['line_count']:>5} lines, " @@ -345,8 +336,7 @@ def main() -> int: if args.emit: return 0 - # Notices after the verdict lines, never mixed into them: "upstream took - # this, drop the entry" is housekeeping and must not read as a failure. + # notices after the verdict lines: housekeeping must not read as a failure for n in notices: print(f"note {n}") if failed: diff --git a/scripts/unsloth/test_additive_merge.py b/scripts/unsloth/test_additive_merge.py index 0f91e9afd12..dd2b22c6cac 100644 --- a/scripts/unsloth/test_additive_merge.py +++ b/scripts/unsloth/test_additive_merge.py @@ -97,8 +97,7 @@ def run(repo, *extra): reason.endswith('twice: log("same");'), reason) # --- 3b. two independent case arms: braces are shared, content is not ------- -# The real tools/mtmd/clip.cpp shape. Refusing this on `{` and `} break;` is -# what took the 09-02 nightly's last pin down. +# the real clip.cpp shape; refusing it on `{` and `} break;` took the 09-02 nightly down base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " builder = std::make_unique(ctx, img);\n" @@ -115,8 +114,7 @@ def run(repo, *extra): txt.count("} break;") == 2 and txt.count("clip_graph_kimik3") == 1, txt) # --- 3b2. two case arms that share a body line, which is a coincidence ------ -# The clip.cpp shape after upstream landed DEEPSEEK4V: both arms set the same -# rope_theta, and refusing on that is the shared-line check backwards. +# clip.cpp after upstream landed DEEPSEEK4V: both arms set the same rope_theta base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;\n" diff --git a/scripts/unsloth/test_pin_contract.py b/scripts/unsloth/test_pin_contract.py index 3c7b58d62d5..e76586b1fc4 100644 --- a/scripts/unsloth/test_pin_contract.py +++ b/scripts/unsloth/test_pin_contract.py @@ -98,8 +98,7 @@ def run(repo, pr_set, *extra): check("intact merge reports no notices", rep["notices"] == [], rep) # --- 2. the arm is dropped from ONE file: a tree-wide grep would pass ------ -# The real shape: LLM_ARCH_INKLING survives in the enum and the dispatch arm -# that makes it do anything is gone. +# the real shape: LLM_ARCH_INKLING survives in the enum, the dispatch arm is gone repo, pr_set, sha = make_repo() p = repo / "src" / "llama-model.cpp" p.write_text(MODEL_CPP_BASE) @@ -128,8 +127,8 @@ def run(repo, pr_set, *extra): any("do_the_banded_thing" in x for x in rep["pins"][0]["problems"]), rep) # --- 5. redundancy: the base already has everything the pin adds ---------- -# Built the way it happens for real: upstream lands the same work, so the base -# tag has it and the pin is not an ancestor of anything. +# as it happens for real: upstream lands the same work, so the base tag has it and +# the pin is not an ancestor of anything d = Path(tempfile.mkdtemp(prefix="pc_")) git(d, "init", "-q", "-b", "main") (d / "src").mkdir() @@ -163,8 +162,8 @@ def run(repo, pr_set, *extra): rep["pins"][0]["added_files"] == ["src/inkling.cpp"], rep) # --- 7. a comment is not a contract --------------------------------------- -# unslothai#70 has a comment naming GGML_OP_SSM_SCAN to say it does NOT use it. -# Holding comment wording would fail the moment upstream rewords it. +# unslothai#70 names GGML_OP_SSM_SCAN in a comment to say it does NOT use it, and +# holding that wording would fail the moment upstream rewords it repo, pr_set, sha = make_repo() git(repo, "checkout", "-q", "pin") (repo / "src" / "note.cpp").write_text( diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48..4fc1c808631 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -507,7 +507,36 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) { +bool llama_batch_allocr::has_shared_tokens() const { + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] > 1) { + return true; + } + } + + return false; +} + +bool llama_batch_allocr::has_seq_wider_than(uint32_t n_tokens) const { + std::vector n_per_seq(n_seq_max, 0); + + for (int32_t i = 0; i < batch.n_tokens; ++i) { + // tokens already placed in an earlier ubatch do not make the rest of the batch a prompt + if (used[i]) { + continue; + } + + for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { + if (++n_per_seq[batch.seq_id[i][s]] > n_tokens) { + return true; + } + } + } + + return false; +} + +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -518,6 +547,10 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, llama_seq_id last_seq_id = -1; + // [TAG_EXACT_CONCURRENCY] tokens left in the first set taken, when isolating: only sets with + // the same count join it, so that every set in the ubatch finishes in this ubatch + uint32_t n_left_first = 0; + // determine the non-overlapping sequence sets participating in this ubatch for (int32_t i = 0; i < batch.n_tokens; ++i) { if (used[i]) { @@ -540,6 +573,46 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { + // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a + // prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so + // give it one of its own. Sets at or below that width are decode steps, kept exact by + // the backend's column policy, so keep grouping them or one prompt would serialize + // every concurrent decode. Grouped sets must have the same number of tokens left, or + // the equal-length expansion below would cut a longer set in two and a memory that + // reduces over a chunk of tokens would sum in a different order than the solo run. + if (isolate_seqs_above > 0) { + uint32_t n_left = 0; + + for (const auto idx : seq_set_map[seq_set[i]]) { + if (!used[idx]) { + ++n_left; + } + } + + if (n_left > isolate_seqs_above) { + if (!cur_seq_set.empty()) { + // let the sets already taken have this ubatch; the prompt gets the next one + break; + } + + cur_seq_set.push_back(seq_set[i]); + + last_seq_id = batch.seq_id[i][0]; + + break; + } + + if (cur_seq_set.empty()) { + n_left_first = n_left; + } else if (n_left != n_left_first) { + continue; + } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) { + // one more set would not finish here, and the expansion below would then cut + // every set part way: the chunking this guard exists to prevent + break; + } + } + cur_seq_set.push_back(seq_set[i]); last_seq_id = batch.seq_id[i][0]; diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a0..f0ecc9d8407 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,7 +105,17 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail); + // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this + // many tokens left to place is a prompt and gets a ubatch of its own; sets at or + // below it are decode steps and stay grouped, so a prompt does not serialize them + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0); + + // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, + // i.e. what remains of the batch holds a prompt rather than decode steps only + bool has_seq_wider_than(uint32_t n_tokens) const; + + // [TAG_EXACT_CONCURRENCY] true if some token carries more than one sequence id + bool has_shared_tokens() const; // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 66940d4fc61..16e090a0623 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -101,6 +101,20 @@ llama_context::llama_context( throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ)); } + // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build: one column per + // sequence times the tokens a sequence contributes, reported so a backend splitting columns + // covers it (a caller that builds wider steps uses llama_set_exact_decode_width). The + // sequence count is what is reported, so a later rise in the tokens figure follows it here + // too. Checked now but reported at the end of the constructor, so a construction that fails + // later does not leave a width behind that no context needs. + if (llama_exact_concurrency()) { + // an explicit column bound wins in the backend, so one below this context's width would + // leave decodes batched above it; the report refuses that, and that is an error here + if (!llama_exact_check_n_seq(cparams.n_seq_max)) { + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } + } + cparams.n_rs_seq = params.n_rs_seq; if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) { LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n", @@ -393,6 +407,13 @@ llama_context::llama_context( }; memory.reset(model.create_memory(params_mem, cparams)); + + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a non-causal context with a + // cache would assert on its first graph + if (llama_exact_concurrency() && memory && !cparams.causal_attn) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__); + throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache"); + } } // init backends @@ -476,6 +497,12 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + + // [TAG_EXACT_CONCURRENCY] nothing above can fail now, so publish the width; already checked + // against the explicit bound at the top, so a refusal here means the bound moved + if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) { + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } } llama_context::~llama_context() { @@ -1188,6 +1215,13 @@ void llama_context::set_causal_attn(bool value) { return; } + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a context with a cache keeps + // causal attention rather than asserting in the next graph + if (!value && memory && llama_exact_concurrency()) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__); + return; + } + cparams.causal_attn = value; sched_need_reserve = true; @@ -2625,17 +2659,16 @@ class llama_io_read_host : public llama_io_read_i { } const size_t tensor_bytes = ggml_nbytes(tensor); auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - // A fragmented sequence can require thousands of synchronous device - // transfers per layer. For bounded tensors, stage the tensor once and - // preserve every byte belonging to other sequences. Bound scratch RAM - // and leave ordinary contiguous transfers on their original fast path. + // a fragmented sequence can need thousands of synchronous device transfers per + // layer: stage a bounded tensor once instead, preserving other sequences' bytes and + // leaving ordinary contiguous transfers on their fast path if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && !ggml_backend_buffer_is_host(buffer)) { std::vector staging; try { staging.resize(tensor_bytes); } catch (const std::bad_alloc &) { - // Fall back to the individual transfers below. + // fall back to the individual transfers below } if (!staging.empty()) { ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); @@ -3226,6 +3259,13 @@ size_t llama_context::state_write_data(llama_io_write_i & io) { } size_t llama_context::state_read_data(llama_io_read_i & io) { + // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical + // index, which the paged pool owns. Refused before anything is parsed, so the caller's cache + // is left as it was: the generic restore path clears it on failure. + if (memory && memory->alloc_granularity() > 1) { + throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence"); + } + LLAMA_LOG_DEBUG("%s: reading state\n", __func__); // read model info @@ -4030,6 +4070,14 @@ bool llama_memory_can_shift(llama_memory_t mem) { return mem->get_can_shift(); } +uint32_t llama_memory_alloc_granularity(llama_memory_t mem) { + if (!mem) { + return 1; + } + + return mem->alloc_granularity(); +} + // llama state API // deprecated diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..dc02560fc66 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -21,11 +21,29 @@ #include #include #include +#include #include #include // dedup helpers +// [TAG_EXACT_CONCURRENCY] the page table is wired into llm_graph_input_attn_kv only, so a V-less +// layout would have its cells placed in pages and then attend in physical order anyway, with the +// mode reporting itself as on. Refuse the context instead: wiring self_pages into +// llm_graph_input_attn_k alone fixes one of the four V-less classes, and DeepSeek 3.2 uses two of +// them, so the model would end up half paged, which is worse than refused. +static void llm_graph_reject_exact_concurrency(const char * layout) { + if (!llama_exact_concurrency()) { + return; + } + + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, but this model uses the %s attention " + "layout, which carries no page table and would attend in physical cell order\n", + __func__, layout); + + throw std::runtime_error("exact concurrency: unsupported attention layout"); +} + static ggml_tensor * build_attn_inp_kq_mask( ggml_context * ctx, const llama_kv_cache_context * mctx, @@ -468,6 +486,7 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { + if (self_pages && self_pages->buffer) { mctx->set_input_pages(self_pages, ubatch); } mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch); @@ -1084,6 +1103,7 @@ void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { + if (inp_attn->self_pages) { mctx->get_attn()->set_input_pages(inp_attn->self_pages, ubatch); } mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); @@ -2547,7 +2567,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, - int il) const { + int il, + ggml_tensor * pages) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -2580,6 +2601,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + cur->src[5] = pages; res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); @@ -2769,6 +2791,8 @@ static std::unique_ptr build_attn_inp_kv_impl( inp->self_kq_mask_cnv = inp->self_kq_mask; } + inp->self_pages = mctx_cur->build_input_pages(ctx0, ubatch); + GGML_ASSERT(!inp->self_pages || (cparams.flash_attn && cparams.causal_attn)); inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0); inp->self_v_rot = mctx_cur->build_input_v_rot(ctx0); @@ -2831,7 +2855,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il, inp->self_pages); cb(cur, "kqv_out", il); if (inp->self_v_rot) { @@ -2865,6 +2889,8 @@ static std::unique_ptr build_attn_inp_k_impl( const llama_cparams & cparams, const llama_kv_cache_context * mctx_cur) { + llm_graph_reject_exact_concurrency("V-less KV (attn_k)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { @@ -3241,6 +3267,8 @@ static std::unique_ptr build_attn_inp_k_dsa_impl( const llama_cparams & cparams, const llama_kv_cache_dsa_context * mctx_cur) { + llm_graph_reject_exact_concurrency("sparse V-less KV (attn_k_dsa)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { @@ -3358,6 +3386,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { const auto * mctx_cur = static_cast(mctx); + llm_graph_reject_exact_concurrency("V-less sliding window KV (attn_k_iswa)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb5..26f2169532d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -319,6 +319,7 @@ class llm_graph_input_attn_no_cache : public llm_graph_input_i { class llm_graph_input_attn_kv : public llm_graph_input_i { public: + ggml_tensor * self_pages = nullptr; // I32 [1 + physical pages, n_tokens] llm_graph_input_attn_kv( const llama_hparams & hparams, const llama_cparams & cparams, @@ -1172,7 +1173,8 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, - int il) const; + int il, + ggml_tensor * pages = nullptr) const; llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index b3a94b946d2..f577414a51c 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -1,11 +1,15 @@ #include "llama-impl.h" +#include "ggml-backend.h" #include "gguf.h" #include "llama.h" #include #include +#include +#include #include +#include #include #include #include @@ -169,3 +173,154 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) { return gguf_data_to_str(type, gguf_get_val_data(ctx_gguf, i), 0); } } + +// [TAG_EXACT_CONCURRENCY] +bool llama_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + +// [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h +static std::atomic g_exact_decode_tokens{1}; + +// one lock for the token figure, the sequence count and the width, since the three move together +// and a report interleaved with a change of figure could leave the backend with a width that +// covers neither; recursive, since the setters call each other +static std::recursive_mutex g_exact_mutex; + +// the most sequences any context was created with. The tokens figure is process wide, so raising +// it widens every existing context's decode step and their width is re-reported with it. +static std::atomic g_exact_max_n_seq{0}; + +static bool llama_exact_width_within_explicit_bound(uint32_t n_cols); + +// sequences times tokens, handed to a backend as an int; a product that overflows is refused +static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) { + const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens; + + if (w > (uint64_t) INT32_MAX) { + LLAMA_LOG_ERROR("%s: a decode step of %u sequences with %u tokens each is too wide to report\n", __func__, n_seq, n_tokens); + return false; + } + + n_cols = (uint32_t) w; + + return true; +} + +bool llama_exact_check_n_seq(uint32_t n_seq) { + std::lock_guard lock(g_exact_mutex); + + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); + + uint32_t n_cols = 0; + + return llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) && llama_exact_width_within_explicit_bound(n_cols); +} + +bool llama_exact_report_n_seq(uint32_t n_seq) { + std::lock_guard lock(g_exact_mutex); + + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); + + uint32_t n_cols = 0; + + if (!llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) || !llama_set_exact_decode_width(n_cols)) { + return false; + } + + uint32_t cur = g_exact_max_n_seq.load(std::memory_order_relaxed); + + while (n_seq > cur && !g_exact_max_n_seq.compare_exchange_weak(cur, n_seq, std::memory_order_relaxed)) { + } + + return true; +} + +bool llama_set_exact_decode_tokens(uint32_t n_tokens) { + n_tokens = n_tokens > 0 ? n_tokens : 1; + + std::lock_guard lock(g_exact_mutex); + + // never lowered: a narrower context set up later would turn an existing speculative + // context's verify steps into prompts and serialise them + if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) { + return true; + } + + // every context widens with the figure, so report the width first; one the explicit bound + // cannot cover leaves the old figure in place + const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); + + uint32_t n_cols = 0; + + if (n_seq > 0 && (!llama_exact_width_of(n_seq, n_tokens, n_cols) || !llama_set_exact_decode_width(n_cols))) { + return false; + } + + g_exact_decode_tokens.store(n_tokens, std::memory_order_relaxed); + + return true; +} + +uint32_t llama_exact_decode_tokens(void) { + return g_exact_decode_tokens.load(std::memory_order_relaxed); +} + +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. Backends read it +// through ggml_backend_cuda_set_exact_decode_width, reached through the registry so an absent or +// late-loaded backend costs nothing. +static std::atomic g_exact_decode_width{0}; + +// an explicit column bound wins in the CUDA backend, so a width above it would leave decodes +// batched past the bound: refuse such a width instead of storing it +static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) { + static const int explicit_cols = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + return val ? atoi(val) : -1; + }(); + + if (explicit_cols > 0 && (uint32_t) explicit_cols < n_cols) { + LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u columns for the decode step just requested; raise it, set it to 0 for no bound, or unset it\n", + __func__, explicit_cols, n_cols); + return false; + } + + return true; +} + +bool llama_set_exact_decode_width(uint32_t n_cols) { + if (!llama_exact_width_within_explicit_bound(n_cols)) { + return false; + } + + std::lock_guard lock(g_exact_mutex); + + uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); + + while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + } + + // the widest figure goes to every backend on every call, not only when it grew, or a width + // reported before a backend was loaded would never reach it + const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed); + + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + ggml_backend_reg_t reg = ggml_backend_reg_get(i); + + auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); + if (fn) { + fn((int) widest); + } + } + + return true; +} + +uint32_t llama_exact_decode_width(void) { + return g_exact_decode_width.load(std::memory_order_relaxed); +} diff --git a/src/llama-impl.h b/src/llama-impl.h index 4988b06d2ca..d0ab2e5cf1c 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -103,3 +103,15 @@ std::string llama_format_tensor_shape(const std::vector & ne); std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); + +// [TAG_EXACT_CONCURRENCY] opt-in mode under which a sequence's attention depends only on its own +// cells, in position order, so its output does not change when others share the KV cache. Off by +// default; reads the same variable as the paged KV cache and the CUDA backend. +bool llama_exact_concurrency(); + +// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so the backend +// knows the width every context needs and it follows llama_set_exact_decode_tokens +bool llama_exact_report_n_seq(uint32_t n_seq); + +// the same check without the report, for a constructor that may still fail after asking +bool llama_exact_check_n_seq(uint32_t n_seq); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a75314..428fbc6b567 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include static bool ggml_is_power_of_2(int n) { @@ -61,6 +62,79 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { // llama_kv_cache // +// [TAG_EXACT_CONCURRENCY] the paged attention specialization lives in the CUDA sources, which are +// also built as ROCm and MUSA. Every other backend ignores src[5] and walks the pool in physical +// cell order, so a KV layer placed there would silently lose the mode's guarantee. +static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { + if (!dev) { + return false; + } + + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (!reg) { + return false; + } + + const char * name = ggml_backend_reg_name(reg); + if (!name) { + return false; + } + + return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0; +} + +// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a layer +// of this shape. The registry name only says which backends carry the kernels, not that the build +// has them or that this architecture, head width and K/V types land on one; otherwise the op +// falls to the CPU, which ignores the page table. So build the op as the graph does, at the +// widths a decode step, a verify step and a prompt chunk use, and ask the device. +static bool llama_dev_supports_paged_attn( + ggml_backend_dev_t dev, + ggml_type type_k, ggml_type type_v, + uint32_t n_embd_head_k, uint32_t n_embd_head_v, + uint32_t n_head, uint32_t n_head_kv, + uint32_t n_cells, uint32_t page_size) { + if (!llama_dev_has_paged_attn(dev)) { + return false; + } + + ggml_init_params ip = { + /*.mem_size =*/ ggml_tensor_overhead()*16 + ggml_graph_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + return false; + } + + bool res = true; + + const int64_t n_kv = page_size; + + for (const int64_t n_tokens : { (int64_t) 1, (int64_t) 4, (int64_t) 16, (int64_t) 512 }) { + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd_head_k, n_tokens, n_head, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, n_embd_head_k, n_kv, n_head_kv, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, n_embd_head_v, n_kv, n_head_kv, 1); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_tokens, 1, 1); + + ggml_tensor * op = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf((float) n_embd_head_k), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(op, GGML_PREC_F32); + + op->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + n_cells/page_size, n_tokens); + + if (!ggml_backend_dev_supports_op(dev, op)) { + res = false; + break; + } + } + + ggml_free(ctx); + + return res; +} + llama_kv_cache::llama_kv_cache( const llama_model & model, const llama_hparams & hparams, @@ -84,6 +158,10 @@ llama_kv_cache::llama_kv_cache( v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { + // [TAG_EXACT_CONCURRENCY] read the knob through the same cached reader the graph and the CUDA + // dispatcher use, so a mid-process change cannot leave the two disagreeing + exact_pages = llama_exact_concurrency(); + // shared cells view the source cache's K/V tensors, so the cell count // follows the source allocation: a fitted target can be smaller than the // draft default and oversized views would overflow the source tensors @@ -97,6 +175,29 @@ llama_kv_cache::llama_kv_cache( GGML_ASSERT(kv_size % n_pad == 0); + // [TAG_EXACT_CONCURRENCY] all of these are reachable from the command line, so name the one + // that failed instead of aborting on a bare assert + if (exact_pages) { + const char * unsupported = nullptr; + + if (!unified) { + unsupported = "it needs a unified KV cache (pass --kv-unified)"; + } else if (v_trans) { + unsupported = "it needs a non-transposed V cache (pass --flash-attn on)"; + } else if (n_swa != 0) { + unsupported = "the paged pool does not support sliding window attention"; + } else if (type_k != GGML_TYPE_F16 || type_v != GGML_TYPE_F16) { + unsupported = "it needs an F16 KV cache (do not pass --cache-type-k or --cache-type-v)"; + } else if (kv_size % exact_page_size != 0) { + unsupported = "the context size must be a multiple of 256 (pass -c as a multiple of 256)"; + } + + if (unsupported) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s\n", __func__, unsupported); + throw std::runtime_error("exact concurrency: unsupported KV cache configuration"); + } + } + const uint32_t n_layer = hparams.n_layer_all; // define a comparator for the buft -> ctx map to ensure that the order is well-defined: @@ -220,6 +321,43 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); + // [TAG_EXACT_CONCURRENCY] the paged kernel handles 256-wide K and V heads only; any other + // width would run unpaged while the mode reports itself as on + if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, " + "and the paged attention kernel supports 256-wide K and V heads only\n", + __func__, il, hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), is_mla ? " (MLA)" : ""); + throw std::runtime_error("exact concurrency: unsupported attention head size"); + } + + // [TAG_EXACT_CONCURRENCY] the paged kernel has no soft-capped variant and would assert + if (exact_pages && hparams.attn_soft_cap) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), " + "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping); + throw std::runtime_error("exact concurrency: attention soft cap is not supported"); + } + + // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order + if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " + "which has no paged attention: every layer must be offloaded to the CUDA backend " + "(pass -ngl to offload all layers and do not pass --no-kv-offload)\n", + __func__, il, dev_name); + throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend"); + } + + // [TAG_EXACT_CONCURRENCY] right backend; ask whether this layer's attention, with the + // page table attached, lands on one of its kernels at all + if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v, + hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), + hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s cannot run the paged attention for layer %d " + "(K %s, V %s, %u-wide heads): the build or the device has no flash attention kernel for it, " + "and the op would fall to the CPU, which ignores the page table\n", + __func__, dev_name, il, ggml_type_name(type_k), ggml_type_name(type_v), hparams.n_embd_head_k(il)); + throw std::runtime_error("exact concurrency: the device cannot run the paged attention"); + } + ggml_context * ctx = ctx_for_buft(buft); if (!ctx) { throw std::runtime_error("failed to create ggml context for kv cache"); @@ -364,7 +502,76 @@ llama_kv_cache::llama_kv_cache( debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0; } +// [TAG_EXACT_CONCURRENCY] +void llama_kv_cache::exact_pages_rebuild() const { + const auto & cells = v_cells[0]; + + exact_page_owner.assign(cells.size()/exact_page_size, exact_page{}); + + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { + continue; + } + + GGML_ASSERT(cells.seq_count(i) == 1); + + const auto pos = cells.pos_get(i); + + GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size); + + const exact_page cur { cells.seq_get(i), llama_pos(pos/(llama_pos) exact_page_size) }; + + auto & owner = exact_page_owner[i/exact_page_size]; + + GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg)); + + owner = cur; + } + + exact_page_owner_dirty = false; +} + +// [TAG_EXACT_CONCURRENCY] +void llama_kv_cache::exact_pages_sync() const { + if (exact_page_owner_dirty) { + exact_pages_rebuild(); + + return; + } + + if (debug > 0) { + // the incrementally maintained ownership has to say what the cells say + const auto kept = exact_page_owner; + + exact_pages_rebuild(); + + GGML_ASSERT(kept.size() == exact_page_owner.size()); + + for (size_t p = 0; p < kept.size(); ++p) { + GGML_ASSERT(kept[p].seq == exact_page_owner[p].seq && kept[p].lpg == exact_page_owner[p].lpg); + } + } +} + +// [TAG_EXACT_CONCURRENCY] +void llama_kv_cache::exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos) { + if (exact_page_owner_dirty || exact_page_owner.empty()) { + // the next sync rebuilds from the cells anyway + return; + } + + const exact_page cur { seq, llama_pos(pos/(llama_pos) exact_page_size) }; + + auto & owner = exact_page_owner[idx/exact_page_size]; + + GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg)); + + owner = cur; +} + void llama_kv_cache::clear(bool data) { + exact_page_owner_dirty = true; + for (uint32_t s = 0; s < n_stream; ++s) { v_cells[s].reset(); v_heads[s] = 0; @@ -383,6 +590,9 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { return true; } + // [TAG_EXACT_CONCURRENCY] a removal can empty a page, which only the cells know + exact_page_owner_dirty = true; + // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); @@ -452,6 +662,16 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll return; } + // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared between + // two. Refuse rather than abort the process, so an uncovered caller gets a failed copy it can + // report. After the shared-cells return, so a draft cache is unaffected. + if (exact_pages && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between " + "sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size()); GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size()); @@ -544,6 +764,9 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { return; } + // [TAG_EXACT_CONCURRENCY] as in seq_rm, this can empty pages + exact_page_owner_dirty = true; + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); auto & cells = v_cells[seq_to_stream[seq_id]]; @@ -571,6 +794,15 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll return; } + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo the page size, + // so shifting positions would misplace every cell; say so rather than abort + if (exact_pages && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions " + "(seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1"); @@ -621,6 +853,14 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in return; } + // [TAG_EXACT_CONCURRENCY] as in seq_add: dividing positions breaks the position/offset identity + if (exact_pages && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions " + "(seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1"); @@ -704,11 +944,26 @@ llama_memory_context_ptr llama_kv_cache::init_batch( GGML_UNUSED(embd_all); do { + // [TAG_EXACT_CONCURRENCY] a token shared by several sequences would be one cell in a page + // that belongs to one sequence; the placement asserts on it later, so refuse it here + if (exact_pages && balloc.has_shared_tokens()) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " + "give every token exactly one sequence id\n", __func__); + break; + } + balloc.split_reset(); std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); + // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt into one ubatch, + // so a prefill would run at a width its solo run never sees. The sequence-set split + // gives each prompt its own ubatch; a plain decode step keeps taking split_simple. + const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; + + auto ubatch = n_stream == 1 && !isolate + ? balloc.split_simple(n_ubatch) + : balloc.split_equal(n_ubatch, n_stream > 1, 0, isolate); if (ubatch.n_tokens == 0) { break; @@ -755,6 +1010,10 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector v_heads_old; // old positions of the heads, before placing the ubatch std::vector v_cells; // copy of the old cells, before placing the ubatch + + // [TAG_EXACT_CONCURRENCY] page ownership before the ubatch, so undoing a speculative + // placement does not force a rebuild from every cell + std::vector exact_page_owner_old; }; // remember the old state of the cells so we can restore it in the end @@ -775,7 +1034,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vectorv_cells[s]); head = it->v_heads_old[s]; } + + // [TAG_EXACT_CONCURRENCY] put back what the allocator knew, unless the placement also + // removed cells, in which case only the cells can say what is left + if (!exact_page_owner_dirty) { + exact_page_owner = it->exact_page_owner_old; + } } if (!success) { @@ -961,6 +1226,48 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } } + if (exact_pages) { + // ownership is maintained as cells are placed, so this reads one entry per page rather + // than scanning every cell; the claims are local and prepare() can still roll them back + const auto & cells = v_cells[0]; + + exact_pages_sync(); + + using page_key = std::pair; + + std::vector owner = exact_page_owner; + std::map pages; + + for (uint32_t p = 0; p < owner.size(); ++p) { + if (owner[p].seq >= 0) { + pages.emplace(page_key {owner[p].seq, owner[p].lpg}, p); + } + } + + std::set assigned; + slot_info res {0, 0, {0}, {{}}}; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1 && ubatch.pos[i] >= 0); + const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size}; + auto it = pages.find(key); + if (it == pages.end()) { + // round-robin free-page search, deliberately nonmonotonic in physical order + uint32_t page = v_heads[0]/exact_page_size; + uint32_t tested = 0; + while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; } + if (tested == owner.size()) { return {}; } + page %= owner.size(); + owner[page] = exact_page {key.first, key.second}; + it = pages.emplace(key, page).first; + } + const uint32_t idx = it->second*exact_page_size + ubatch.pos[i]%exact_page_size; + if (!cells.is_empty(idx) || !assigned.insert(idx).second) { return {}; } + res.idxs[0].push_back(idx); + } + if (cont && !res.is_contiguous()) { return {}; } + return res; + } + uint32_t n_tokens = ubatch.n_tokens; uint32_t n_seqs = 1; @@ -1139,6 +1446,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) { cells.seq_add(idx, ubatch.seq_id[i][s]); } + + // [TAG_EXACT_CONCURRENCY] the page this cell belongs to is now owned by its sequence + if (exact_pages) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1); + + exact_pages_claim(idx, ubatch.seq_id[i][0], ubatch.pos[i]); + } } } @@ -1170,7 +1484,19 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } } +uint32_t llama_kv_cache::alloc_granularity() const { + // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so n + // tokens hold round_up(n, exact_page_size) cells: the tail page is charged in full + return exact_pages ? exact_page_size : 1; +} + bool llama_kv_cache::get_can_shift() const { + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo 256, so the pool + // cannot shift positions. Reporting it here is what disables --context-shift and + // --cache-reuse at load rather than failing on the first request that needs them. + if (exact_pages) { + return false; + } // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot. if (model.arch == LLM_ARCH_STEP35) { return false; @@ -1232,7 +1558,53 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const { return v_cells[seq_to_stream[seq_id]]; } +ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const { + if (!exact_pages) { return nullptr; } + auto * pages = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + get_size()/exact_page_size, ubatch.n_tokens); + ggml_set_input(pages); + ggml_set_name(pages, "attn_logical_pages"); + return pages; +} + +void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { + GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens); + + // [TAG_EXACT_CONCURRENCY] one entry per physical page, not one per cell + exact_pages_sync(); + + std::map> pages; + for (uint32_t p = 0; p < exact_page_owner.size(); ++p) { + const auto & owner = exact_page_owner[p]; + if (owner.seq >= 0) { + pages[owner.seq][owner.lpg] = p; + } + } + std::vector data(ggml_nelements(dst), -1); + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + GGML_ASSERT(ubatch->n_seq_id[i] == 1); + auto * row = data.data() + i*dst->ne[0]; + row[0] = 0; + for (const auto & page : pages[ubatch->seq_id[i][0]]) { + // exclude wholly future pages even when prefill includes later query rows + if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; } + row[++row[0]] = page.second; + } + } + ggml_backend_tensor_set(dst, data.data(), 0, data.size()*sizeof(int32_t)); +} + +ggml_tensor * llama_kv_cache_context::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const { + return kv->build_input_pages(ctx, ubatch); +} + +void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { + kv->set_input_pages(dst, ubatch); +} + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { + // the physical view spans the pool; the per-query page map is the only loop bound for exact + // attention, so neighbours cannot extend it + if (exact_pages) { return get_size(); } uint32_t result = 0; // pad the n_kv value so that the graph remains constant across batches and can be reused @@ -2037,6 +2409,14 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical index, + // which the paged pool owns. Refused before a byte is read, so the clearing failure path below + // is never entered for it. + if (exact_pages && seq_id == -1) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__); + throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY"); + } + // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -2285,6 +2665,9 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 } else { // whole KV cache restore + // [TAG_EXACT_CONCURRENCY] refused at the top of state_read(), before anything is read + GGML_ASSERT(!exact_pages); + if (cell_count > cells.size()) { LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__); return false; diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f98..c139ede57e0 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -131,6 +131,9 @@ class llama_kv_cache : public llama_memory_i { bool get_can_shift() const override; + // [TAG_EXACT_CONCURRENCY] the page size under exact mode, 1 otherwise + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; @@ -171,6 +174,8 @@ class llama_kv_cache : public llama_memory_i { // uint32_t get_n_kv(const slot_info & sinfo) const; + ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const; + void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const; // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; @@ -235,6 +240,30 @@ class llama_kv_cache : public llama_memory_i { std::vector v_stream; }; + static constexpr uint32_t exact_page_size = 256; + bool exact_pages = false; + + // [TAG_EXACT_CONCURRENCY] which (sequence, logical page) owns each physical page; seq < 0 means + // free. Kept current as cells are placed and marked dirty by removals, so find_slot() and + // set_input_pages() read one entry per page instead of rebuilding from every live cell twice per + // ubatch. Mutable because set_input_pages() is const. + struct exact_page { + llama_seq_id seq = -1; + llama_pos lpg = -1; + }; + + mutable std::vector exact_page_owner; + mutable bool exact_page_owner_dirty = true; + + // bring exact_page_owner up to date; rebuilds only when a removal marked it dirty + void exact_pages_sync() const; + + // recompute it from the live cells + void exact_pages_rebuild() const; + + // record that a cell of (seq, pos) now lives at physical cell idx + void exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos); + bool v_trans = true; // the value tensor is transposed const uint32_t n_seq_max = 1; @@ -365,6 +394,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // uint32_t get_n_kv() const; + ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const; + void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const; ggml_type type_k() const; ggml_type type_v() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..d9e609fee20 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -66,6 +66,13 @@ llama_memory_hybrid::llama_memory_hybrid( llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { + // [TAG_EXACT_CONCURRENCY] refused before the attention half asserts on it, see llama_kv_cache::init_batch + if (llama_exact_concurrency() && balloc.has_shared_tokens()) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " + "give every token exactly one sequence id\n", __func__); + break; + } + balloc.split_reset(); // follow the recurrent pattern for creating the ubatch splits @@ -86,7 +93,13 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // so that the rollback snapshots remain valid const uint32_t n_rs_seq = mem_recr->n_rs_seq; - ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + // [TAG_EXACT_CONCURRENCY] the recurrent half is not invariant to the ubatch shape: + // a prompt processed next to other prompts leaves a different gated delta net state, + // so it gets a ubatch of its own while plain decode steps stay batched. Passed + // whenever the mode is on, since it also keeps sets of unequal token counts apart. + const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } if (ubatch.n_tokens == 0) { @@ -135,6 +148,12 @@ bool llama_memory_hybrid::get_can_shift() const { return mem_attn->get_can_shift(); } +uint32_t llama_memory_hybrid::alloc_granularity() const { + // the recurrent half holds one state per sequence, so the attention half is the one whose + // cells a caller is planning capacity for + return mem_attn->alloc_granularity(); +} + void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); mem_recr->clear(data); @@ -150,6 +169,14 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + // [TAG_EXACT_CONCURRENCY] the attention half refuses this, so refuse before either half is + // touched or the two could end up describing different states + if (llama_exact_concurrency() && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); } @@ -160,11 +187,23 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + if (llama_exact_concurrency() && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions (seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + mem_attn->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + if (llama_exact_concurrency() && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions (seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + mem_attn->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); } diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb7499..70ba19ca323 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -58,6 +58,8 @@ class llama_memory_hybrid : public llama_memory_i { bool get_can_shift() const override; + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7..dadc0661344 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -431,7 +431,13 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // [TAG_RECURRENT_ROLLBACK_SPLITS] // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch // so that the rollback snapshots remain valid - ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: the state a prompt leaves + // behind depends on what shared its ubatch, so isolate prompts. Passed whenever the + // mode is on, since it also keeps sets of unequal token counts apart, which would + // otherwise be reduced in chunks the solo run never had. + const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; + + ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } if (ubatch.n_tokens == 0) { diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645..c0bc2ad7e30 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,6 +100,12 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode + // that allocates in larger blocks is on, and then n tokens occupy round_up(n, granularity) + // cells, so a caller planning capacity in tokens would see room that does not exist. Not pure, + // so a module that has never heard of this inherits the answer that was always true of it. + virtual uint32_t alloc_granularity() const { return 1; } + // // ops // diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 53e93a1448d..3f08a060ef1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7193,6 +7193,41 @@ struct test_flash_attn_ext : public test_case { } }; +// same attention as the CPU mask reference, but visiting nonadjacent pages in a different order; +// covers a partial tail and different page counts per query +struct test_flash_attn_ext_pages : public test_flash_attn_ext { + test_flash_attn_ext_pages(int64_t batch) : + test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {} + + std::string vars() override { return test_flash_attn_ext::vars() + ",exact_pages=1"; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + auto * out = test_flash_attn_ext::build_graph(ctx); + out->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 5, nb); + ggml_set_name(out->src[5], "pages"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + test_flash_attn_ext::initialize_tensors(ctx); + auto * pages = ggml_get_tensor(ctx, "pages"); + auto * mask = ggml_get_tensor(ctx, "m"); + std::vector ids(5*nb, -1); + std::vector values(1024*nb, ggml_fp32_to_fp16(-INFINITY)); + for (int64_t q = 0; q < nb; ++q) { + ids[5*q] = q%2 ? 1 : 2; + ids[5*q + 1] = 2; + ids[5*q + 2] = 0; + for (int j = 0; j < 256; ++j) { values[1024*q + 512 + j] = ggml_fp32_to_fp16(0.0f); } + if (q%2 == 0) { + for (int j = 0; j < 17; ++j) { values[1024*q + j] = ggml_fp32_to_fp16(0.0f); } + } + } + ggml_backend_tensor_set(pages, ids.data(), 0, ids.size()*sizeof(int32_t)); + ggml_backend_tensor_set(mask, values.data(), 0, values.size()*sizeof(ggml_fp16_t)); + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; @@ -9170,6 +9205,24 @@ static std::vector> make_test_cases_eval() { } } + // shared weights over sequence planes, as in a recurrent-model output projection + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) { + for (int n : {1, 17, 307}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1})); + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {1, 4})); + } + } + + // MoE projections at the token counts a decode ubatch forms: gate and up broadcast one + // activation row over the expert list, down carries one row per expert, and a real MoE gguf + // puts different types on the two. 17 tokens is past the width exact concurrency pins. + for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) { + for (int n : {1, 2, 4, 8, 17}) { + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true, 512, n, 2048)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, false, 2048, n, 512)); + } + } + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); @@ -9937,6 +9990,9 @@ static std::vector> make_test_cases_eval() { } // mixed quant and Q1_0 test cases + for (int64_t batch : {1, 4, 12}) { + test_cases.emplace_back(new test_flash_attn_ext_pages(batch)); + } test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)); diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index 428a9252981..24b6359b05f 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,8 +73,8 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); - // A fragmented restore may stage a whole device tensor. Check every - // sequence byte-for-byte, including the neighbours that must be preserved. + // a fragmented restore may stage a whole device tensor, so check every sequence + // byte-for-byte, including the neighbours that must be preserved std::vector> before(params.n_parallel); for (int s = 0; s < params.n_parallel; ++s) { before[s].resize(llama_state_seq_get_size(ctx, s)); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b18fa4e2f23..4f02ef9aa79 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -38,6 +38,17 @@ constexpr int HTTP_POLLING_SECONDS = 1; +// [TAG_EXACT_CONCURRENCY] read from the env like the KV cache, batch splitter and CUDA +// backend do: the answer is needed before a context exists. +static bool server_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + static common_speculative_output_limits server_output_limits(const common_params & params) { if (params.embedding || (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { @@ -65,17 +76,11 @@ enum slot_state { // [TAG_PREEMPT] server-side request preemption // -// With --kv-unified the cells are one pool shared by every slot, and each slot believes it -// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves -// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size -// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the -// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and -// continue with the rest". -// -// Nothing is terminated here. The cells of one slot are taken back and given to it again -// later: its sequence is copied to host RAM, its cells are released, and when the pool has -// room the copy goes back and the slot carries on with the same sampler, the same generated -// text and the same open stream. A streaming client sees a pause, not an error. +// With --kv-unified one full pool ends EVERY conversation in flight with "Context size has +// been exceeded", including the ones nowhere near their own limit. Instead of terminating, +// one slot's sequence is copied to host RAM and its cells released; when there is room the +// copy goes back and the slot carries on with the same sampler, text and open stream, so a +// streaming client sees a pause rather than an error. constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected @@ -91,6 +96,37 @@ constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long +// [TAG_EXACT_CONCURRENCY] the planner counts cells, not tokens: llama_memory_alloc_granularity() +// is 1 ordinarily but the page size under exact concurrency, where a sequence of n tokens +// occupies round_up(n, page) cells. A planner counting tokens would see room find_slot cannot +// find in pages, never park anybody, and end every request in the context error instead. + +// cells a run of n_tokens occupies when the pool allocates g at a time +static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) { + return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g; +} + +// cells a step of n_step more costs: nothing until it crosses a page boundary, a page when it does +static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) { + return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g); +} + +// at a granularity of 1 both are the identity, so nothing changes in a configuration that does not page +static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 && + preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3, + "at a granularity of 1 a run of n tokens has to cost exactly n cells"); +static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(8191, 1, 1) == 1 && + preempt_n_cells_step_g(1000, 512, 1) == 512, + "at a granularity of 1 a step of n tokens has to cost exactly n cells"); + +// and the page arithmetic itself, so the rounding cannot change by accident +static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 && + preempt_n_cells_g(257, 256) == 512, + "a tail page is charged in full"); +static_assert(preempt_n_cells_step_g(255, 1, 256) == 0 && preempt_n_cells_step_g(256, 1, 256) == 256 && + preempt_n_cells_step_g(256, 257, 256) == 512, + "a step is free until it crosses a page boundary and costs whole pages when it does"); + struct server_slot; // forward declaration struct server_batch { @@ -323,16 +359,13 @@ struct server_slot { prompt.clear(); } - // [TAG_PREEMPT] state of a slot whose cells were taken back - // - // Only the KV cells leave. The task, the sampler, the generated text and the position - // the stream has reached stay on the slot, so a resume is a memcpy and not a new - // request: no retokenisation, no replayed prompt, no seam in the output. + // [TAG_PREEMPT] state of a slot whose cells were taken back. Only the KV cells leave; + // the task, sampler, generated text and stream position stay, so a resume is a memcpy. slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; int32_t n_preempt = 0; // times the CURRENT task has been preempted - int32_t n_ctx_shift = 0; // context shifts the CURRENT task has made: it is at the pool's limit and cycling + int32_t n_ctx_shift = 0; // context shifts it has made: it is at the pool's limit and cycling int32_t n_preempt_fail = 0; // consecutive failed restores int64_t t_preempt_us = 0; // when it was parked bool preempt_rotation_refused = false; // this park has logged a rotation refused for budget @@ -382,10 +415,8 @@ struct server_slot { return false; } - // The draft is a prediction, not a result, so it goes with the cells. Preemption - // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already - // holds exactly the tokens the state above covers -- including the rollback done by - // the checkpoint path when a draft was only partially accepted. + // the draft is a prediction, not a result, so it goes with the cells; prompt.tokens + // already covers exactly what the state above holds spec_draft.clear(); spec_i_batch.clear(); spec_ckpt.clear(); @@ -393,8 +424,7 @@ struct server_slot { i_batch = -1; - // note: prompt.tokens is deliberately kept. It is the mirror of the state just - // copied out, and the resume needs it to know how many cells to ask for. + // note: prompt.tokens is deliberately kept - the resume sizes its request from it mem.seq_rm(id, -1, -1); state_before_preempt = state; @@ -432,10 +462,8 @@ struct server_slot { state = state_before_preempt; - // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks - // that the draft context is where it should be, which the restore above ensures. - // A slot parked while still processing its prompt makes that transition itself - // once the prompt is done. + // same call the DONE_PROMPT -> GENERATING transition makes; a slot parked mid-prompt + // makes that transition itself once the prompt is done if (state == SLOT_STATE_GENERATING && can_speculate()) { common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } @@ -443,12 +471,9 @@ struct server_slot { return true; } - // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds for this sequence. - // For a batch that is given up after it was built: the tokens added for this slot - // that were never decoded come off, the sampled token stays in `sampled` and goes into - // the next batch the way it went into this one, and a draft is a prediction that goes - // with them. A chunk that failed to decode left nothing in the cache, so the cache is - // the boundary. + // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up + // after it was built: never-decoded tokens and the draft come off, `sampled` is kept for + // the next batch. A failed chunk left nothing in the cache, so the cache is the boundary. void rewind_to_cache() { const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; @@ -456,8 +481,7 @@ struct server_slot { prompt.tokens.keep_first(n_cached); } - // a prompt whose last chunk was in the batch was marked done when the chunk was - // built; the chunk never ran, so the prompt is not done + // the last chunk was marked done when it was built but never ran, so it is not done if (state == SLOT_STATE_DONE_PROMPT && task && prompt.n_tokens() < task->n_tokens()) { state = SLOT_STATE_PROCESSING_PROMPT; } @@ -688,9 +712,8 @@ struct server_slot { t_last_used = ggml_time_us(); - // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be - // released while parked), so the mirror of them must not outlive them: the next - // task on this slot would otherwise take a prefix match against an empty cache + // [TAG_PREEMPT] the cells are already gone, so the mirror of them must not outlive + // them: the next task would take a prefix match against an empty cache if (state == SLOT_STATE_PREEMPTED) { preempt_state_free(); prompt_clear(); @@ -1444,6 +1467,26 @@ struct server_context_impl { } } + // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates rather than assume a cell per + // token; 1 unless a mode that places cells in blocks is on + { + preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); + + // test knob: the paged kernel needs a head size of 256, so a harness model cannot + // turn exact concurrency on and this is the only way to reach the paged arithmetic + const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY"); + + if (LLAMA_SERVER_PREEMPT_GRANULARITY) { + preempt_alloc_granularity = std::max(1, atoi(LLAMA_SERVER_PREEMPT_GRANULARITY)); + + SRV_WRN("LLAMA_SERVER_PREEMPT_GRANULARITY = %d (test knob: planning the kv pool in blocks of %d cells)\n", + preempt_alloc_granularity, preempt_alloc_granularity); + } else if (preempt_alloc_granularity > 1) { + SRV_INF("preemption: the kv pool allocates %d cells at a time, planning in pages\n", + preempt_alloc_granularity); + } + } + { // read on every load and kept on this context, so a reload after the variable // changed, or another context loaded in the same process, has an order of its own @@ -1464,9 +1507,8 @@ struct server_context_impl { preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks, for comparing - // policies against each other on the same workload. smallest (the default and the - // shipped one), largest, youngest (the most recent task, as vLLM's scheduler - // preempts), oldest. The leader is kept and the starvation guard applies under all. + // policies on the same workload: smallest (default), largest, youngest, oldest. + // The leader is kept and the starvation guard applies under all. const char * LLAMA_SERVER_PREEMPT_POLICY = getenv("LLAMA_SERVER_PREEMPT_POLICY"); preempt_test_policy = LLAMA_SERVER_PREEMPT_POLICY ? LLAMA_SERVER_PREEMPT_POLICY : "smallest"; @@ -2889,9 +2931,8 @@ struct server_context_impl { void abort_all_slots(const std::string & reason) { for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in - // host RAM, not in the cache, and it comes back when there is room, the same as - // in the decode error sweep + // [TAG_PREEMPT] a parked slot took no part in what failed and comes back when + // there is room, the same as in the decode error sweep if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); @@ -2935,16 +2976,34 @@ struct server_context_impl { // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, - // whether or not the pool is under pressure. It exists to answer the only question that - // matters about a resume: with one request on an idle server the batch has the same - // shape at every step, so a preempted continuation that is not byte-identical to an - // uninterrupted one is the preemption's fault and nothing else's. + // under pressure or not: on an idle server the batch shape is fixed, so a continuation + // that is not byte-identical to an uninterrupted one is the preemption's fault. int32_t preempt_test_every = 0; std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model - // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so - // the KV-full retry ladder and its last resort are the only thing between a full pool - // and the context error + // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load: 1 + // ordinarily, the page size under exact concurrency. Everything below plans in cells + // because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides it for the harness. + int32_t preempt_alloc_granularity = 1; + + // cells a slot holding n_tokens actually occupies + int32_t preempt_n_cells(int32_t n_tokens) const { + return preempt_n_cells_g(n_tokens, preempt_alloc_granularity); + } + + // cells a slot holding n_tokens has to be given for a step of n_step more + int32_t preempt_n_cells_step(int32_t n_tokens, int32_t n_step) const { + return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity); + } + + // cells kept spare on top of the reservation, rounded up to a page since a boundary + // crossing costs a whole one; PREEMPT_N_MARGIN at a granularity of 1 + int32_t preempt_n_margin() const { + return preempt_n_cells(PREEMPT_N_MARGIN); + } + + // LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, leaving + // only the KV-full retry ladder and its last resort bool preempt_planner_off = false; // LLAMA_SERVER_PREEMPT_RESUME: head (the default) puts parked slots back in the order they @@ -2964,8 +3023,8 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } - // draft tokens this slot's next step can actually carry: the configured maximum, cut to - // what its context and its prediction budget leave, the way get_n_draft_max() cuts it + // draft tokens the next step can carry: the maximum cut to what context and prediction + // budget leave, the way get_n_draft_max() cuts it int32_t preempt_n_spec(const server_slot & slot) const { int32_t res = preempt_n_spec_max(); @@ -3045,19 +3104,18 @@ struct server_context_impl { res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); } - return res; + // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in + // full; undercounting here admits a resume that find_slot cannot satisfy + return preempt_n_cells(res); } - // Cells the pool is holding right now. A released slot keeps its prompt in the cache - // for the next request to reuse as a prefix, so idle slots count too: the first version - // of this counted only the running ones, decided a pool holding 8185 cached cells was - // empty, and every resume failed against a cache that was actually full. + // cells the pool is holding right now. A released slot keeps its prompt in the cache as a + // prefix for the next request, so idle slots count too or a full pool looks empty. int32_t preempt_kv_used() const { int32_t res = 0; - // n_cmpl > 1: the parent and its children share the prompt's cells through seq_cp, so - // the prompt is charged once per family, to whichever resident member comes first; - // the others are charged only what they generated on top of it + // n_cmpl > 1: a family shares the prompt's cells through seq_cp, so the prompt is + // charged once, to the first resident member; the others only for what they added std::vector charged; for (const auto & slot : slots) { @@ -3065,11 +3123,13 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } - // a child waiting for its parent's prompt does not share anything yet: until - // copy_state_to() runs it still holds whatever the previous request left in its - // cells, so it is charged that on its own, outside the family + // [TAG_EXACT_CONCURRENCY] the tail page is charged in full: it cannot be given to + // anybody else, however little of it is used + + // a child waiting for its parent shares nothing until copy_state_to() runs, so it + // is charged its own stale cells, outside the family if (slot.state == SLOT_STATE_WAIT_OTHER) { - res += slot.prompt.n_tokens(); + res += preempt_n_cells(slot.prompt.n_tokens()); continue; } @@ -3077,7 +3137,7 @@ struct server_context_impl { const int family = slot.task->is_parent() ? slot.task->id : slot.task->id_parent; if (std::find(charged.begin(), charged.end(), family) != charged.end()) { - res += std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens()); + res += preempt_n_cells(std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens())); continue; } @@ -3090,7 +3150,7 @@ struct server_context_impl { // can be a long time behind a running generation. Measured by the prefix, a // restore was found to fit and attempted against cells still occupied. Under // pressure the planner trims such slots itself, see preempt_normalize_started_all() - res += slot.prompt.n_tokens(); + res += preempt_n_cells(slot.prompt.n_tokens()); } return res; @@ -3103,42 +3163,53 @@ struct server_context_impl { int32_t res = 0; int32_t res_pmt = 0; + // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the + // tokens: preempt_kv_used() already rounds up the tail page, so reserving tokens on + // top of it would miss the boundary crossing, the only moment the pool can run out for (const auto & slot : slots) { + const int32_t n_cur = slot.prompt.n_tokens(); + switch (slot.state) { case SLOT_STATE_GENERATING: case SLOT_STATE_DONE_PROMPT: { - res += 1 + preempt_n_spec(slot); + res += preempt_n_cells_step(n_cur, 1 + preempt_n_spec(slot)); } break; case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { // from the prefix a started slot keeps, not from the prompt it still // mirrors: measured by the mirror, a request shorter than the last one - // reserved one cell for a chunk of hundreds + // reserved one cell for a chunk of hundreds. The page arithmetic below + // still starts from n_cur, the cells preempt_kv_used() charges for it, + // so a trim that has not happened yet cannot make the step look free. const int32_t n_left = slot.task ? slot.task->n_tokens() - preempt_n_retained(slot) : 0; - res_pmt += std::max(1, std::min(n_batch, n_left)); + res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left))); } break; default: break; } } - // one batch is all the prompt slots get between them, however many are waiting - return res + std::min(res_pmt, n_batch); + // one batch is all the prompt slots get between them, however many are waiting; under + // page allocation each can still cross a boundary of its own, so the cap allows one + // boundary per prompt slot on top of the batch + int32_t n_pmt = 0; + + for (const auto & slot : slots) { + if (slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT) { + n_pmt++; + } + } + + return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } - // Keep the slot that is furthest along -- it is the closest to finishing and to giving - // its cells back -- and among the rest prefer one that has not been preempted - // PREEMPT_N_STARVED times already, then the smallest. - // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt - // until the batch builder keeps the prefix the two share and drops the rest (see the - // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out, - // charged and sized by the old prompt, and a short unrelated request could exceed the - // budget or stay parked for room it will never use. Keeping only the shared prefix now - // is what the batch builder does anyway; the chunk reuse it can add on top is given up - // for a slot the planner has to touch, which is rare. + // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt until + // the batch builder drops it (see the SLOT_STATE_STARTED block of update_slots). Parked as + // it is, it would be copied out and sized by the old prompt. Keeping only the shared prefix + // now is what the batch builder does anyway, at the cost of the chunk reuse it can add. // every started slot, when the pool is short: true when any of them gave cells up bool preempt_normalize_started_all() { bool res = false; @@ -3173,10 +3244,8 @@ struct server_context_impl { return; } - // a memory that cannot remove part of a sequence (a recurrent state without rollback - // room for the stale suffix) aborts on a partial removal; for it the whole stale - // sequence goes, and the prompt is processed from the start on resume, as it would be - // without a usable checkpoint + // a memory that cannot remove part of a sequence aborts on a partial removal; for it + // the whole stale sequence goes and the prompt is reprocessed from the start const bool partial_ok = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART && (!ctx_dft || ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART); @@ -3194,10 +3263,8 @@ struct server_context_impl { server_slot * leader = nullptr; int32_t n_running = 0; - // a slot just given a task is measured by the prefix it keeps, not by the previous - // request's prompt it still mirrors: measured by the mirror, a short request over a - // large stale cache would be the never-parked leader while the longest live - // conversation was parked in its place + // measure a just-started slot by the prefix it keeps, not by the stale prompt it + // mirrors, or a short request over a large stale cache becomes the leader for (auto & slot : slots) { preempt_normalize_started(slot); } @@ -3213,19 +3280,16 @@ struct server_context_impl { } if (n_running < 2) { - // a single conversation that does not fit the pool on its own is a real context - // overflow and not a scheduling problem - leave it to the existing error path + // one conversation that does not fit alone is a real overflow, not a scheduling + // problem - leave it to the existing error path return nullptr; } server_slot * victim = nullptr; for (auto & slot : slots) { - // Before the batch is built every one of these is at a token boundary: a - // generating slot between two sampled tokens, a prompt-processing slot between - // two chunks of its prompt, a started slot with only a cached prefix (or - // nothing) in the pool. A slot holding no cells is still worth parking - it - // is about to ask for a whole batch of them. + // before the batch is built every one of these is at a token boundary. A slot + // holding no cells is still worth parking - it is about to ask for a batch. if (slot.state != SLOT_STATE_GENERATING && slot.state != SLOT_STATE_PROCESSING_PROMPT && slot.state != SLOT_STATE_STARTED) { @@ -3257,9 +3321,8 @@ struct server_context_impl { return victim; } - // is a the better victim of the two? the smallest slot under the shipped policy: it - // gives up the least work and its restore is the cheapest (see the PR's simulation); - // the other choices exist for the comparison runs behind LLAMA_SERVER_PREEMPT_POLICY + // is a the better victim? the smallest under the shipped policy, since it gives up the + // least work; the rest exist for comparison runs behind LLAMA_SERVER_PREEMPT_POLICY bool preempt_better_victim(const server_slot & a, const server_slot & b) const { if (preempt_test_policy == "largest") { return a.prompt.n_tokens() > b.prompt.n_tokens(); @@ -3276,10 +3339,8 @@ struct server_context_impl { return a.prompt.n_tokens() < b.prompt.n_tokens(); } - // called once per update_slots(), before the batch is built: at that point every slot is - // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no - // draft is in flight, so a slot can be removed from the picture without unpicking a - // half-decoded batch + // called once per update_slots(), before the batch is built: every slot is then at a token + // boundary with no draft in flight, so one can be removed without unpicking a batch void update_preemption() { if (!params_base.kv_unified || slots.size() < 2) { return; // with a cache per slot, no slot can take another one's cells @@ -3295,10 +3356,7 @@ struct server_context_impl { const int32_t n_cells = n_ctx; - // Put back what fits, in the order preempt_resume_head describes: by default - // the slot parked longest, and only that one until it fits; under - // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked - // longest, and a smaller slot may pass a head that does not fit. + // put back what fits, in the order preempt_resume_head describes const bool head_of_line = preempt_resume_head; for (;;) { @@ -3328,12 +3386,9 @@ struct server_context_impl { server_slot * best = nullptr; - // A parked slot whose sequence plus its next step would not fit an empty pool can - // never be restored, and would otherwise sit at the head of the line for ever - // without a restore ever being attempted: a prompt within n_ctx that was parked - // before it took any cells, but too close to n_ctx to leave room for its first - // batch. That is the single-conversation overflow the KV-full path reports, so - // report it the same way and rescan the line without it. + // a parked slot that would not fit an empty pool can never be restored and would + // sit at the head of the line for ever. That is the single-conversation overflow + // the KV-full path reports, so report it the same way and rescan without it. { server_slot * impossible = nullptr; @@ -3353,15 +3408,13 @@ struct server_context_impl { } } - // Room for the sequence AND for the next step of everything already running, - // so that a resume cannot immediately trigger the preemption of someone else. - // The margin is headroom for the others; with nothing resident there is nobody - // to keep it for, so a sequence that fits the pool exactly is let back in. - // A cached prompt on an idle slot is worth less than a conversation waiting to - // continue, so give those cells up first - same call the KV-full path makes. + // room for the sequence AND for the next step of everything running, so a resume + // cannot immediately preempt someone else. The margin is headroom for the others, + // so with nothing resident a sequence that fits exactly is let back in. Cached + // prompts on idle slots are worth less than a waiting conversation, so go first. for (;;) { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); - const int32_t margin = occupied == 0 ? 0 : PREEMPT_N_MARGIN; + const int32_t margin = occupied == 0 ? 0 : preempt_n_margin(); for (auto * slot : parked) { if (occupied + preempt_n_need(*slot) + margin <= n_cells) { @@ -3386,22 +3439,21 @@ struct server_context_impl { } } - // Nothing fits. A resident that has reached the pool's limit and is cycling - // through context shifts holds the room for as long as it likes to generate, - // and the head behind it would wait for ever. After the head has waited its - // turn, that resident is parked in its place: it is at a token boundary like - // any other park, and when it comes back it is the one waiting, so the two - // take turns instead of one taking everything. + // nothing fits. A resident cycling through context shifts holds the room for as + // long as it generates, so once the head has waited its turn that resident is + // parked in its place and the two take turns. if (!best) { server_slot * head = parked.front(); if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { - // the resident whose cells let the head in, the smallest of those; failing - // one that does so alone, the largest, since it makes the most room. Taking - // the first shifting resident in slot order could park one too small to - // matter, spend the park budget on it, and leave the head waiting anyway. + // the smallest resident whose cells let the head in; failing one that does + // so alone, the largest, since it makes the most room + // [TAG_EXACT_CONCURRENCY] the rotation arrived after the planner moved to + // cells and was still asking in tokens: the margin unrounded, which makes + // the requirement smaller than a page-allocated pool can meet, and the + // resident's holding unrounded, which undercounts what parking it frees const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); - const int32_t need = preempt_n_need(*head) + PREEMPT_N_MARGIN; + const int32_t need = preempt_n_need(*head) + preempt_n_margin(); server_slot * pick = nullptr; bool pick_enough = false; @@ -3426,7 +3478,7 @@ struct server_context_impl { continue; } - const bool enough = occupied - slot.prompt.n_tokens() + need <= n_cells; + const bool enough = occupied - preempt_n_cells(slot.prompt.n_tokens()) + need <= n_cells; if (!pick || (enough && !pick_enough) || @@ -3470,9 +3522,8 @@ struct server_context_impl { const int64_t t_start = ggml_time_us(); if (!best->preempt_restore()) { - // update_slots() runs in a tight loop while tasks are pending, so a counter - // alone burns its whole budget in a couple of milliseconds. Give up only on - // a slot that has been failing for a while, and keep the log quiet. + // update_slots() loops tightly, so a counter alone burns its budget in + // milliseconds: give up only on a slot failing for a while, and log quietly if (best->n_preempt_fail % 64 == 1) { SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); @@ -3520,7 +3571,7 @@ struct server_context_impl { for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); - if (n_used + PREEMPT_N_MARGIN <= n_cells) { + if (n_used + preempt_n_margin() <= n_cells) { break; } @@ -3652,8 +3703,7 @@ struct server_context_impl { #endif if (preempt_batch_abandoned) { - // [TAG_PREEMPT] the rest of this batch was never decoded and the slots no - // longer describe it; the next pass builds a new one + // [TAG_PREEMPT] the rest of this batch never ran; the next pass rebuilds it preempt_batch_abandoned = false; break; } @@ -3687,8 +3737,7 @@ struct server_context_impl { // apply context-shift if needed // TODO: simplify and improve - // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift, - // not with the cells the shift is about to give back + // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { @@ -3901,8 +3950,7 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to - // batch; it takes no part in this pass until it is restored + // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to batch if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { return; } @@ -4429,14 +4477,11 @@ struct server_context_impl { } } - // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is - // the context error for every slot in the batch. With a park budget the batch is given - // up instead: every resident slot is rewound to the token boundary the cache is at (a - // batch is applied one chunk at a time, and the chunk that failed left nothing behind), - // the smallest are parked until the planner's own bound holds again, and the next - // update_slots() rebuilds the batch from the survivors. The planner brings the parked - // ones back as cells free up. A multimodal prompt has no boundary the cache can name, - // so it keeps the old path. + // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell, which upstream is + // the context error for every slot in the batch. With a park budget the batch is given up + // instead: resident slots are rewound to the token boundary the cache is at, the smallest + // are parked until the planner's bound holds, and the next update_slots() rebuilds the + // batch. A multimodal prompt has no boundary the cache can name, so it keeps the old path. bool preempt_last_resort_possible() const { return params_base.kv_unified && params_base.preempt_ram_mib != 0 && !preempt_recurrent && slots.size() >= 2 && llama_get_memory(ctx_tgt); } @@ -4476,7 +4521,7 @@ struct server_context_impl { for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); - if (n_parked > 0 && n_used + PREEMPT_N_MARGIN <= n_cells) { + if (n_parked > 0 && n_used + preempt_n_margin() <= n_cells) { break; } @@ -4572,11 +4617,9 @@ struct server_context_impl { { std::string err; - // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have - // to stay in one view: a narrower view splits the group and the verify step - // throws for the slot whose tokens straddle it. Halving is no help there, so - // after the idle slots the ladder goes to its last resort straight away. With - // no budget to park into the ladder is what it always was. + // [TAG_PREEMPT] a slot's sampled token and its draft have to stay in one view, + // so halving would split the group and make the verify step throw: after the + // idle slots the ladder goes straight to its last resort if (ret == 1 && n_batch > 1 && preempt_last_resort_possible() && batch_has_spec_groups()) { if (try_clear_idle_slots()) { SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n"); @@ -4613,8 +4656,7 @@ struct server_context_impl { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the - // cache; it is not part of this failure and comes back when there is room + // [TAG_PREEMPT] a parked slot is not part of this failure if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, err); slot.release(); @@ -5222,6 +5264,16 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_cmpl_id = completion_id; task.params.oaicompat_model = meta->model_name; + // [TAG_EXACT_CONCURRENCY] children of an n_cmpl > 1 task are served by copying the + // parent's cells to another sequence id, and exact mode gives a page to a single + // sequence, so refuse here where it becomes a 400 rather than at seq_cp + if (task.params.n_cmpl > 1 && server_exact_concurrency()) { + throw std::runtime_error( + "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " + "completion needs its own sequence, and in exact mode a KV page belongs " + "to a single sequence. Send n separate requests, or unset the variable."); + } + // prepare child tasks if (task.params.n_cmpl > 1) { int n_children = task.params.n_cmpl - 1; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 0c8f5dc5f29..fb50a2b1687 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -1,13 +1,13 @@ import os +import re import time import tempfile import pytest from utils import * -# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked -# (its sequence copied to host RAM, its cells released) instead of every slot being -# terminated. Both tests need more than one slot and --kv-unified, which is the only -# configuration where one slot can take another one's cells. +# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked (its +# sequence copied to host RAM, its cells released) instead of every slot being terminated. Needs +# more than one slot and --kv-unified, the only configuration where slots share cells. server = ServerPreset.tinyllama2() @@ -38,6 +38,7 @@ def create_server(): os.close(fd) yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_SERVER_PREEMPT_GRANULARITY", None) os.environ.pop("LLAMA_SERVER_PREEMPT_PLANNER", None) os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) @@ -55,9 +56,8 @@ def _complete(n_predict: int, prompt: str = "Hi how are you"): def test_forced_preemption_does_not_change_the_output(): - # Park and restore the only running slot every 8 tokens. With one request the batch - # has the same shape at every step whether or not the slot was parked in between, so - # any difference in the output is the preemption's fault and nothing else's. + # park and restore the only running slot every 8 tokens: the batch shape is the same at every + # step, so any difference in the output is the preemption's fault global server server.n_ctx = 512 server.start() @@ -84,11 +84,9 @@ def test_forced_preemption_does_not_change_the_output(): def test_two_slots_that_overflow_the_pool_together_both_finish(): - # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well - # under 256. Together they do not, 336 against 256. Without preemption the retry - # ladder ends with "Context size has been exceeded" on every processing slot; with it - # the smaller slot is parked until the leader finishes and its cells are purged, and - # then it resumes from the token it was parked on. + # each request fits the pool alone (168 of 256 cells) but not together (336). Without + # preemption both end with "Context size has been exceeded"; with it the smaller is parked + # until the leader finishes, then resumes from the token it was parked on. global server server.n_ctx = 256 server.start() @@ -112,6 +110,43 @@ def test_two_slots_that_overflow_the_pool_together_both_finish(): assert len(res.body["tokens"]) == n_predict +def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): + # a pool that allocates in blocks gives a whole block to one sequence, so n tokens occupy + # round_up(n, block) cells and the planner has to count cells: counting tokens it sees room + # the allocator cannot find, never parks anybody, and the retry ladder ends every request. + # LLAMA_SERVER_PREEMPT_GRANULARITY injects the block size, since the only mode that reports + # one needs a head size this model does not have; the arithmetic is the same at 64 as at 256. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_GRANULARITY = 64" in log.drain() + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool holds + # and "(wanted N)" is that plus the next decode's reservation + held = [int(n) for n in re.findall(r"kv (\d+)/256", text)] + wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)] + assert held and wanted, f"the planner logged no figures:\n{text}" + assert all(n % 64 == 0 for n in held + wanted), f"not whole blocks: {held} {wanted}" + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + _WORDS = ( "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor " @@ -140,9 +175,8 @@ def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: def test_two_prompts_that_overflow_the_pool_together_both_finish(): - # Neither slot ever generates before the pool is full: both are still processing their - # prompts. A prompt-processing slot is between two chunks of its prompt, which is as - # clean a boundary as the one between two sampled tokens, so it is parked the same way. + # neither slot generates before the pool is full: a slot between two chunks of its prompt is + # as clean a boundary as one between two sampled tokens, so it is parked the same way global server server.n_ctx = 256 server.start() @@ -171,20 +205,17 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish(): def test_a_generating_slot_and_a_large_prompt_both_finish(): - # One slot is generating a long answer to a short prompt when a large prompt arrives - # beside it. Together they need far more than the pool has. The prompt is admitted - # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish. - # This model produces a thousand tokens a second, so the second request is sent right - # behind the first rather than after a delay: its prompt takes several batches to - # process, which is enough for the two to overlap however fast the first one runs. + # a slot generating a long answer to a short prompt meets a large prompt arriving beside it, + # needing far more than the pool has: the prompt is admitted chunk by chunk, whoever is + # smaller is parked, and both finish. The second request follows immediately, since its + # prompt takes several batches and that is enough overlap however fast the first one runs. global server server.n_ctx = 256 server.start() log = LogReader(server.log_path) prompt_b, n_b = _prompt_of_about(150, "Charlie") - # b lives long enough for the two to collide: the first run of this used 16 tokens - # and b was finished and purged before a had grown into it + # b has to live long enough for the two to collide n_predict_a = 230 n_predict_b = 90 assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 @@ -210,8 +241,8 @@ def _late(n_predict, prompt): def test_preempt_ram_zero_disables_preemption(): - # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the - # KV-full path ends the requests the way it always did. + # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path + # ends the requests the way it always did global server server.n_ctx = 256 os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" @@ -231,9 +262,8 @@ def test_preempt_ram_zero_disables_preemption(): def test_metrics_and_slots_report_the_parked_state(): - # A client that wants to tell a parked chat from a slow one reads /slots, and an - # operator reads /metrics. Both must show the preemption happening, and the counters - # must survive the requests finishing. + # /slots tells a parked chat from a slow one and /metrics reports it to an operator; both + # must show the preemption, and the counters must survive the requests finishing global server server.n_ctx = 256 server.server_metrics = True @@ -271,11 +301,9 @@ def test_metrics_and_slots_report_the_parked_state(): def test_two_prompts_near_the_context_size_both_complete(): - # Two prompts that each fit the context alone but not together. The second one is - # parked before it takes any cells, and it is close enough to n_ctx that its sequence - # plus its first batch would not leave the usual scheduling margin. It must still be - # restored once the first one finishes: with nothing resident there is nobody to keep - # the margin for. Before the fix it was parked for ever, with no restore ever tried. + # two prompts that each fit the context alone but not together. The second is parked before + # it takes any cells and is too close to n_ctx to leave the usual margin, but must still be + # restored once the first finishes: with nothing resident there is nobody to keep it for. global server server.n_ctx = 256 # the whole prompt in one batch, so the parked slot's first step is the whole prompt @@ -285,7 +313,7 @@ def test_two_prompts_near_the_context_size_both_complete(): # sized in tokens, not words: the prompt is the token ids of a short sentence repeated base = server.make_request("POST", "/tokenize", data={"content": "Once upon a time there was a little girl"}).body["tokens"] - long_prompt = (base * 64)[:250] + long_prompt = (base * 64)[:240] n_predict = 4 together = parallel_function_calls([(_complete, (n_predict, long_prompt)) for _ in range(2)]) @@ -298,11 +326,9 @@ def test_two_prompts_near_the_context_size_both_complete(): def test_the_last_resort_parks_instead_of_ending_everyone(): - # With the planner off nothing is parked ahead of the decode, so two generations that - # fit alone but not together fill the pool until a single token finds no cell. That - # is where upstream ends every slot with the context error. Instead the batch is - # given up, the smaller slot is parked, the larger one finishes, and the parked one - # comes back and finishes too. + # with the planner off, two generations that fit alone but not together fill the pool until a + # single token finds no cell, where upstream ends every slot with the context error. Instead + # the batch is given up, the smaller slot is parked, and both finish. global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -356,10 +382,8 @@ def test_the_last_resort_works_with_an_unlimited_budget(): def test_the_last_resort_rewinds_a_prompt_in_flight(): - # Same, with a prompt being processed when the pool runs out: the chunk that failed - # is taken back off the slot's tokens and processed again after the resume, so the - # prompt is neither skipped nor fed twice. The prompt is far longer than a batch, so - # the failing chunk is a chunk of it, not its last token. + # same, with a prompt being processed when the pool runs out: the failed chunk comes back off + # the slot's tokens and is processed again after the resume, neither skipped nor fed twice global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -395,12 +419,10 @@ def _late(n_predict, prompt): def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): - # Two generations that each outgrow the pool on their own, with context shift on. The - # resident reaches the limit, shifts, keeps about half the pool and would keep going - # for as long as it has tokens to make, while the parked one never fits beside it. - # After the head has waited its turn the resident is parked in its place, and the two - # take turns until both finish. Long enough that the resident is still going when the - # head's turn comes: this model makes a couple of thousand tokens a second. + # two generations that each outgrow the pool, with context shift on: the resident shifts and + # would hold half the pool for as long as it generates, while the parked one never fits + # beside it. After the head has waited its turn the resident is parked in its place and the + # two take turns. n_predict is large enough that the resident is still going by then. global server server.n_ctx = 256 server.enable_ctx_shift = True @@ -424,9 +446,8 @@ def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_hea def test_the_rotation_parks_a_resident_that_lets_the_head_in(): - # Three generations with no end in a 256-cell pool with context shift on: two residents - # cycle through shifts while the third waits parked. Every rotation must let the head - # in, so all three keep finishing their tokens and no stream ends short. + # three endless generations with context shift on: two residents cycle through shifts while + # the third waits parked, and every rotation must let the head in so no stream ends short global server server.n_slots = 3 server.n_ctx = 384 @@ -452,11 +473,9 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in(): def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): - # One request asking for two completions is one conversation in two slots: a parent - # and a child sharing the prompt. When the two together do not fit the pool there is - # nobody else to park, since the family is charged once and a member of it is not a - # victim for the other, so the request gets the context error it would get alone, and - # the server carries on serving. + # a two-completion request is one conversation in two slots, and a family member is not a + # victim for the other, so with nobody else to park it gets the context error it would get + # alone and the server carries on serving global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off"