diff --git a/docs/image-input.md b/docs/image-input.md index f6e88c753..6d7758c07 100644 --- a/docs/image-input.md +++ b/docs/image-input.md @@ -150,16 +150,29 @@ luce_server models/DeepSeek-V4-Flash-Vision-Exp-ROCMFPX-MIX-STRIX.gguf \ ``` Its image blocks need whole-block bidirectional prefill, which the batched -engine's 16-row step cannot run. Image requests admitted since the last step -are therefore prefilled up to their last token together, in shared -layer-major sparse passes into per-request staging caches (each layer's -experts are read once for all of them); that state is copied into each -request's paged slot and the last token prefills in the batch, so the answers -decode alongside everyone else. On the Strix Halo with the encoder on the -R9700, four concurrent image answers of 256 tokens finish in 35 s (29 tok/s in -total), two images plus two text requests at 31 tok/s; four text requests -reach 38 tok/s. Image requests beyond the free slots wait in the queue. `/props` reports the -effective capability in `capabilities.image_input_supported` after backend +engine's 16-row step cannot run. An image request is therefore prefilled up to +its last token into its slot's staging cache on the layer-major sparse path, +then copied into its paged slot, and the last token prefills in the batch, so +the answer decodes alongside everyone else. Every layer must see a whole image +block at once, so the block cannot be split by rows; it is split by layers +instead. While other requests are decoding, a staged pass takes about 256 rows +(a whole image block, which may be more) shared by the pending image requests +and runs 6 of its 43 layers per batched step, so live streams decode after +every slice. The result is identical to running the pass in one go. With +nothing decoding, a pass takes up to 1,024 rows and all its layers at once, so +the requests share each layer's expert reads. With `--mmproj-device`, the +encoder works through admitted requests on its own GPU and never stalls the +batch. Measured on the Strix Halo with the encoder on the R9700: two text +streams decoding while three one-image requests and one eight-image request +arrive pause at most 0.67 s between tokens; four concurrent image answers of +256 tokens finish in 37.5 s (27 tok/s in total). + +The staging caches, one per slot, are allocated at startup and logged +(`staging caches 286 MB` for 4 slots at `--max-ctx 8192`). Image requests +beyond the free slots wait in the queue. Under KV pressure the scheduler never +parks an image request for recompute (token ids cannot rebuild image rows); it +suspends or parks a text request instead. `/props` reports the effective +capability in `capabilities.image_input_supported` after backend initialization. ## Qwen3.5 / Qwen3.8 diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 2292f259a..33c93744d 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -333,6 +333,10 @@ class SeqEngine { error = "engine does not support KV eviction"; return false; } + // False when the slot's KV cannot be rebuilt from its token history (for + // example image rows that prefill from pixels): the scheduler then picks + // another eviction victim. + virtual bool kv_recomputable(int) const { return true; } // True when a parked slot's resume reservation fits current free pool // capacity with headroom for the resident cohort's next step — the diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 1b40aa756..0fda8c7d2 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -34,8 +34,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -60,10 +62,10 @@ class DeepSeek4ImagePrompt final : public ImagePromptPayload { for (const auto & image : prepared_.images) spans_.push_back(image.layout.span); } - bool embed_chunk(const CpuEmbedder & embedder, size_t position, - int count, float * output) const { + bool embed_chunk(const CpuEmbedder & embedder, size_t position, int count, float * output, + const std::function & cancelled = {}) const { if (!output || count <= 0 || embedder.n_embd <= 0) return false; - if (!wait_for_images(position, size_t(count))) return false; + if (!wait_for_images(position, size_t(count), cancelled)) return false; std::vector result; std::string error; const bool ok = vision::embed_image_prompt_chunk( @@ -79,26 +81,64 @@ class DeepSeek4ImagePrompt final : public ImagePromptPayload { // Streaming encode: images arrive in prompt order. A chunk waits only for // the images it contains; failure or cancellation releases every waiter. - bool wait_for_images(size_t position, size_t count) const { + size_t images_needed(size_t position, size_t count) const { size_t needed = 0; for (size_t i = 0; i < spans_.size(); ++i) { if (spans_[i].block_begin < position + count && spans_[i].block_end > position) needed = i + 1; } + return needed; + } + // `cancelled` (polled while waiting) lets a request that ends early stop + // waiting for images it no longer needs. + bool wait_for_images(size_t position, size_t count, + const std::function & cancelled = {}) const { + const size_t needed = images_needed(position, count); std::unique_lock lock(stream_mutex_); - stream_ready_.wait(lock, [&] { return stream_failed_ || ready_ >= needed; }); + while (!stream_failed_ && ready_ < needed) { + if (cancelled && cancelled()) return false; + stream_ready_.wait_for(lock, std::chrono::milliseconds(50)); + } + return ready_ >= needed; + } + // Non-blocking: every image overlapping [position, position + count) is in. + bool images_ready(size_t position, size_t count) const { + const size_t needed = images_needed(position, count); + std::lock_guard lock(stream_mutex_); return ready_ >= needed; } + void wait_images_for(size_t position, size_t count, int timeout_ms) const { + const size_t needed = images_needed(position, count); + std::unique_lock lock(stream_mutex_); + stream_ready_.wait_for(lock, std::chrono::milliseconds(timeout_ms), + [&] { return stream_failed_ || ready_ >= needed; }); + } + bool stream_failed() const { + std::lock_guard lock(stream_mutex_); + return stream_failed_; + } + // A new encode of every image; the stream settles when it ends. + void begin_stream() const { + std::lock_guard lock(stream_mutex_); + materialized_.assign(prepared_.images.size(), {}); + ready_ = 0; + stream_failed_ = false; + cancelled_.store(false); + } void publish_image(size_t index, std::vector rows) const { std::lock_guard lock(stream_mutex_); materialized_[index] = std::move(rows); ready_ = index + 1; stream_ready_.notify_all(); } - void fail_stream() const { + void settle_stream(bool ok) const { std::lock_guard lock(stream_mutex_); - stream_failed_ = true; + stream_failed_ = stream_failed_ || !ok; stream_ready_.notify_all(); } + // Stops a queued or running encode at its next image. The worker owns a + // reference to the payload, so nothing needs to wait for it. + void cancel_stream() const { cancelled_.store(true); } + bool stream_cancelled() const { return cancelled_.load(); } bool complete() const { std::lock_guard lock(stream_mutex_); return !materialized_.empty() && ready_ == materialized_.size(); @@ -113,6 +153,7 @@ class DeepSeek4ImagePrompt final : public ImagePromptPayload { mutable std::condition_variable stream_ready_; mutable size_t ready_ = 0; mutable bool stream_failed_ = false; + mutable std::atomic cancelled_{false}; }; namespace { @@ -1148,99 +1189,212 @@ ImagePrepareStatus DeepSeek4Backend::prepare_images( bool DeepSeek4Backend::encode_image_request(const std::vector & prompt, const ImagePromptHandle & handle, std::string & error) { - const auto * images = dynamic_cast(handle.get()); + auto images = std::dynamic_pointer_cast(handle); if (!images || images->owner_ != this || !images->matches(prompt)) { error = "image binding does not match this prompt"; return false; } + if (vision_backend_) { + // The encoder has its own GPU: queue and return, so admission never + // stalls the batch. Staged prefill takes each image as it lands. + images->begin_stream(); + enqueue_image_encode(std::move(images)); + return true; + } DaemonIO io; const auto t0 = Clock::now(); - const bool ok = materialize_images(*images, io, error); - join_image_stream(); - if (ok && !images->complete()) { - error = "image encoding did not complete"; + if (!materialize_images(images, io, error)) return false; + std::fprintf(stderr, "[deepseek4] batched image request encoded in %.0f ms on the target GPU\n", + elapsed_s(t0) * 1000.0); + return true; +} + +DeepSeek4Cache * DeepSeek4Backend::image_staging_cache(int slot) { + if (slot == 0) return cache_.buf ? &cache_ : nullptr; + if (slot < 1 || size_t(slot) > image_staging_caches_.size()) return nullptr; + return image_staging_caches_[size_t(slot) - 1].get(); +} + +bool DeepSeek4Backend::begin_staged_prefill(StagedPrefill & item) { + const auto * images = dynamic_cast(item.images.get()); + if (!images || !item.staging || item.prefix < DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS || + item.prefix > int(item.prompt.size()) || item.prefix > item.staging->max_ctx) { + item.error = "invalid staged prefill request"; return false; } - std::fprintf(stderr, "[deepseek4] batched image request encoded in %.0f ms\n", elapsed_s(t0) * 1000.0); - return ok; + reset_deepseek4_cache(*item.staging); + item.staging->prefill_mode = PrefillAttentionMode::Sparse; + item.done = 0; + return true; } -void DeepSeek4Backend::prefill_staged(std::vector & batch) { - // Embeddings for every request's prefix (image rows + text rows). - struct Seq { StagedPrefill * item; const DeepSeek4ImagePrompt * images; std::vector embed; int done = 0; }; - std::vector seqs; - for (auto & item : batch) { - const auto * images = dynamic_cast(item.images.get()); - if (!images || !item.prompt || !item.staging || item.prefix < DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS || - item.prefix > int(item.prompt->size()) || item.prefix > item.staging->max_ctx) { - item.error = "invalid staged prefill request"; +bool DeepSeek4Backend::begin_staged_pass(const std::vector & items, int row_budget, + DeepSeek4PrefillPass & pass, std::vector & rows) { + // One pass takes the next chunk of every ready, unfinished request while + // the budget lasts: whole image blocks only (a block may exceed the + // budget), never a chunk or a tail below the layer-major minimum, and only + // rows whose images the encoder has already published. + constexpr int min_rows = DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS; + rows.assign(items.size(), 0); + std::vector seqs; + std::vector members; + std::vector> embeds; + int budget = row_budget; + for (size_t k = 0; k < items.size(); ++k) { + StagedPrefill * item = items[k]; + if (item->finished() || budget < min_rows) continue; + const auto * images = static_cast(item->images.get()); + if (images->stream_failed()) { + item->error = "image encoding failed"; continue; } - std::vector embed(size_t(item.prefix) * size_t(w_.n_embd)); - if (!images->embed_chunk(w_.embedder, 0, item.prefix, embed.data())) { - item.error = "staged prefill embedding failed"; + const int n = vision::staged_prefill_chunk(images->spans(), uint64_t(item->done), + item->prefix - item->done, budget, min_rows); + if (n == 0) { + item->error = "staged prefill could not be chunked"; continue; } - reset_deepseek4_cache(*item.staging); - item.staging->prefill_mode = PrefillAttentionMode::Sparse; - seqs.push_back({&item, images, std::move(embed), 0}); + if (!images->images_ready(size_t(item->done), size_t(n))) continue; + std::vector embed(size_t(n) * size_t(w_.n_embd)); + if (!images->embed_chunk(w_.embedder, size_t(item->done), n, embed.data())) { + item->error = "staged prefill embedding failed"; + continue; + } + embeds.push_back(std::move(embed)); + DeepSeek4PrefillSeq seq; + seq.cache = item->staging; + seq.embed = embeds.back().data(); + seq.token_ids = item->prompt.data() + item->done; + seq.n_tokens = n; + seq.kv_start = item->done; + seq.image_spans = images->spans(); + seqs.push_back(seq); + members.push_back(k); + budget -= n; + } + if (seqs.empty()) return false; + std::string error; + if (!pass.begin(backend_, cfg_.device.gpu, w_, seqs, error)) { + for (size_t k : members) items[k]->error = error.empty() ? "staged prefill failed" : error; + return false; } - ggml_backend_synchronize(backend_); - deepseek4_release_image_scratch(cache_, moe_hybrid_.get()); - const int budget_total = std::min(1024, DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS); - constexpr int min_rows = DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS; - const auto t0 = Clock::now(); - int passes = 0, rows = 0; + for (size_t m = 0; m < members.size(); ++m) rows[members[m]] = seqs[m].n_tokens; + return true; +} + +void DeepSeek4Backend::wait_staged_ready(const StagedPrefill & item, int row_budget, + int timeout_ms) const { + const auto * images = static_cast(item.images.get()); + if (!images || item.finished()) return; + // Wake as soon as the next chunk's images are in, not all of them. + const int n = vision::staged_prefill_chunk(images->spans(), uint64_t(item.done), + item.prefix - item.done, row_budget, + DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS); + if (n > 0) images->wait_images_for(size_t(item.done), size_t(n), timeout_ms); +} + +bool DeepSeek4Backend::encode_one_image(const vision::PromptImage & image, + vision::ImageRaster & raster, std::string & error) { + std::vector patches(image.input.patches_bf16.size()); + for (size_t i = 0; i < patches.size(); ++i) { + const uint32_t bits = uint32_t(image.input.patches_bf16[i]) << 16; + std::memcpy(&patches[i], &bits, sizeof(bits)); + } + vision::VisionOutput output; + if (!vision_->encode(patches, + {int(image.input.plan.vit_rows), int(image.input.plan.vit_cols)}, + output, error)) return false; + if (output.rows <= 0 || output.columns != w_.n_embd) { + error = "vision output shape differs from decoder dimensions"; + return false; + } + raster = {size_t(output.rows), size_t(output.columns), std::move(output.embeddings)}; + return true; +} + +void DeepSeek4Backend::enqueue_image_encode(std::shared_ptr images) { + std::lock_guard lock(encode_mutex_); + if (encode_stop_) { + images->settle_stream(false); + return; + } + if (!encode_worker_.joinable()) encode_worker_ = std::thread([this] { encode_worker_loop(); }); + encode_queue_.push_back(std::move(images)); + encode_ready_.notify_one(); +} + +void DeepSeek4Backend::encode_worker_loop() { for (;;) { - // One pass takes the next chunk of every unfinished request that fits: - // whole image blocks only, and never a chunk or tail below the - // layer-major minimum. - std::vector pass; - std::vector members; - int budget = budget_total; - for (auto & s : seqs) { - const int remaining = s.item->prefix - s.done; - if (remaining <= 0 || !s.item->error.empty() || budget < min_rows) continue; - int n = std::min(remaining, budget); - if (remaining - n > 0 && remaining - n < min_rows) n = std::max(0, remaining - min_rows); - n = n >= min_rows ? vision::atomic_image_chunk(s.images->spans(), uint64_t(s.done), n, - uint64_t(remaining), budget) : 0; - if (n < min_rows) continue; - DeepSeek4PrefillSeq seq; - seq.cache = s.item->staging; - seq.embed = s.embed.data() + size_t(s.done) * size_t(w_.n_embd); - seq.token_ids = s.item->prompt->data() + s.done; - seq.n_tokens = n; - seq.kv_start = s.done; - seq.image_spans = s.images->spans(); - pass.push_back(seq); - members.push_back(&s); - budget -= n; - } - if (pass.empty()) break; + std::shared_ptr images; + { + std::unique_lock lock(encode_mutex_); + if (encode_queue_.empty()) { + // Idle: hand the scratch back before sleeping. + lock.unlock(); + vision_->release_scratch(); + lock.lock(); + } + encode_ready_.wait(lock, [&] { return encode_stop_ || !encode_queue_.empty(); }); + if (encode_stop_) { + for (auto & queued : encode_queue_) queued->settle_stream(false); + encode_queue_.clear(); + return; + } + images = std::move(encode_queue_.front()); + encode_queue_.pop_front(); + } + const auto t0 = Clock::now(); + bool ok = true; std::string error; - if (!deepseek4_prefill_multi(backend_, cfg_.device.gpu, w_, pass, error)) { - for (Seq * m : members) m->item->error = error.empty() ? "staged prefill failed" : error; - continue; + try { + const auto cancelled = [&] { return images->stream_cancelled() || encode_stop_; }; + for (size_t i = 0; ok && i < images->prepared_.images.size(); ++i) { + vision::ImageRows one; + ok = vision::materialize_image_rows({images->prepared_.images[i]}, image_sentinels_, + size_t(w_.n_embd), + [this](const vision::PromptImage & image, vision::ImageRaster & raster, + std::string & encode_error) { + return encode_one_image(image, raster, encode_error); + }, cancelled, one, error) && one.size() == 1; + if (ok) images->publish_image(i, std::move(one.front())); + } + } catch (const std::exception & e) { + // Nothing may escape the worker: fail the stream so prefill stops waiting. + ok = false; + error = e.what(); + } + images->settle_stream(ok); + if (ok) { + std::fprintf(stderr, "[deepseek4] %zu image(s) encoded in %.0f ms on the --mmproj-device GPU\n", + images->prepared_.images.size(), elapsed_s(t0) * 1000.0); + } else { + std::fprintf(stderr, "[deepseek4] image encode stopped: %s\n", + error.empty() ? "cancelled" : error.c_str()); } - for (size_t k = 0; k < members.size(); ++k) members[k]->done += pass[k].n_tokens; - ++passes; - for (const auto & p : pass) rows += p.n_tokens; } - for (auto & s : seqs) { - if (s.item->error.empty() && s.done != s.item->prefix) s.item->error = "staged prefill could not be chunked"; - s.item->ok = s.item->error.empty(); - } - std::fprintf(stderr, "[deepseek4] staged prefill: %zu requests, %d rows in %d shared passes, %.0f ms\n", - seqs.size(), rows, passes, elapsed_s(t0) * 1000.0); } -void DeepSeek4Backend::join_image_stream() { - if (image_stream_.joinable()) image_stream_.join(); +void DeepSeek4Backend::cancel_image_encode(const ImagePromptPayload & payload) const { + if (const auto * images = dynamic_cast(&payload)) images->cancel_stream(); +} + +void DeepSeek4Backend::release_vision() { + { + std::lock_guard lock(encode_mutex_); + encode_stop_ = true; + encode_ready_.notify_all(); + } + if (encode_worker_.joinable()) encode_worker_.join(); + { + std::lock_guard lock(encode_mutex_); + encode_stop_ = false; + } + vision_.reset(); } -bool DeepSeek4Backend::materialize_images(const DeepSeek4ImagePrompt & images, +bool DeepSeek4Backend::materialize_images(const std::shared_ptr & handle, const DaemonIO & io, std::string & error) { + const DeepSeek4ImagePrompt & images = *handle; if (io.is_cancelled()) return false; if (images.owner_ != this || !vision_ || parked_) { error = "image binding does not belong to the loaded backend"; @@ -1309,66 +1463,11 @@ bool DeepSeek4Backend::materialize_images(const DeepSeek4ImagePrompt & images, return false; } if (images.complete()) return true; - const auto encode_one = [this](const vision::PromptImage & image, vision::ImageRaster & raster, - std::string & encode_error) { - std::vector patches(image.input.patches_bf16.size()); - for (size_t i = 0; i < patches.size(); ++i) { - const uint32_t bits = uint32_t(image.input.patches_bf16[i]) << 16; - std::memcpy(&patches[i], &bits, sizeof(bits)); - } - vision::VisionOutput output; - if (!vision_->encode(patches, - {int(image.input.plan.vit_rows), int(image.input.plan.vit_cols)}, - output, encode_error)) return false; - if (output.rows <= 0 || output.columns != w_.n_embd) { - encode_error = "vision output shape differs from decoder dimensions"; - return false; - } - raster = {size_t(output.rows), size_t(output.columns), std::move(output.embeddings)}; - return true; - }; - vision::ImageSentinels sentinels; - if (!vision_->sentinel(vision::Sentinel::Start, sentinels.start, error) || - !vision_->sentinel(vision::Sentinel::Pad, sentinels.pad, error) || - !vision_->sentinel(vision::Sentinel::Newline, sentinels.newline, error) || - !vision_->sentinel(vision::Sentinel::End, sentinels.end, error)) return false; if (vision_backend_) { // The encoder has its own GPU: encode image k+1 there while the target // prefills image k. Prefill waits per chunk in embed_chunk. - join_image_stream(); - { - std::lock_guard lock(images.stream_mutex_); - images.materialized_.assign(images.prepared_.images.size(), {}); - images.ready_ = 0; - images.stream_failed_ = false; - } - image_stream_ = std::thread([this, &images, &io, encode_one, sentinels] { - const auto t0 = Clock::now(); - bool ok = true; - std::string stream_error; - try { - for (size_t i = 0; ok && i < images.prepared_.images.size(); ++i) { - vision::ImageRows one; - ok = vision::materialize_image_rows({images.prepared_.images[i]}, sentinels, - size_t(w_.n_embd), encode_one, [&] { return io.is_cancelled(); }, - one, stream_error) && one.size() == 1; - if (ok) images.publish_image(i, std::move(one.front())); - } - } catch (const std::exception & e) { - // Nothing may escape the thread: fail the stream so prefill stops waiting. - ok = false; - stream_error = e.what(); - } - vision_->release_scratch(); - if (!ok) { - std::fprintf(stderr, "[deepseek4] streaming image encode stopped: %s\n", - stream_error.empty() ? "cancelled" : stream_error.c_str()); - images.fail_stream(); - return; - } - std::fprintf(stderr, "[deepseek4] images encoded in %.0f ms on the --mmproj-device GPU (streamed)\n", - elapsed_s(t0) * 1000.0); - }); + images.begin_stream(); + enqueue_image_encode(handle); return true; } struct ReleaseScratch { @@ -1377,8 +1476,11 @@ bool DeepSeek4Backend::materialize_images(const DeepSeek4ImagePrompt & images, } release{*vision_}; try { vision::ImageRows rows; - if (!vision::materialize_image_rows(images.prepared_.images, sentinels, size_t(w_.n_embd), - encode_one, [&] { return io.is_cancelled(); }, rows, error)) return false; + if (!vision::materialize_image_rows(images.prepared_.images, image_sentinels_, size_t(w_.n_embd), + [this](const vision::PromptImage & image, vision::ImageRaster & raster, + std::string & encode_error) { + return encode_one_image(image, raster, encode_error); + }, [&] { return io.is_cancelled(); }, rows, error)) return false; std::lock_guard lock(images.stream_mutex_); images.materialized_ = std::move(rows); images.ready_ = images.materialized_.size(); @@ -1463,6 +1565,15 @@ bool DeepSeek4Backend::load_vision() { "[deepseek4] vision weights=%.3f GiB scratch reservation=%.3f GiB on %s\n", gib(runtime->weight_bytes()), gib(vision::SCRATCH_RESERVATION), vision_backend_ ? "the --mmproj-device GPU" : "the target GPU, before expert placement"); + vision::ImageSentinels sentinels; + if (!runtime->sentinel(vision::Sentinel::Start, sentinels.start, error) || + !runtime->sentinel(vision::Sentinel::Pad, sentinels.pad, error) || + !runtime->sentinel(vision::Sentinel::Newline, sentinels.newline, error) || + !runtime->sentinel(vision::Sentinel::End, sentinels.end, error)) { + std::fprintf(stderr, "[deepseek4] projector sentinels unavailable: %s\n", error.c_str()); + return false; + } + image_sentinels_ = std::move(sentinels); vision_ = std::move(runtime); return true; } @@ -1905,14 +2016,26 @@ bool DeepSeek4Backend::init() { } if (vision_) { // Image admissions prefill on the single-request sparse path into - // this staging cache, then copy it into their paged slot. - if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { - std::fprintf(stderr, "[deepseek4] image staging cache allocation failed (ctx=%d)\n", max_ctx); + // their slot's staging cache (slot 0 uses cache_), then copy it + // into the paged slot. All of them are allocated now, so memory is + // committed at startup rather than in the middle of a request. + bool staged = create_deepseek4_cache(backend_, w_, max_ctx, cache_); + image_staging_caches_.clear(); + for (int slot = 1; staged && slot < cfg_.max_concurrency; ++slot) { + auto cache = std::make_unique(); + staged = create_deepseek4_cache(backend_, w_, max_ctx, *cache); + if (staged) image_staging_caches_.push_back(std::move(cache)); + } + if (!staged) { + std::fprintf(stderr, "[deepseek4] image staging caches do not fit (%d x ctx=%d); " + "reduce --max-ctx or --max-concurrency\n", cfg_.max_concurrency, max_ctx); return false; } cache_.prefill_mode = PrefillAttentionMode::Sparse; - std::fprintf(stderr, "[deepseek4] batched image serving: %d slots, staging cache ctx=%d\n", - cfg_.max_concurrency, max_ctx); + std::fprintf(stderr, "[deepseek4] batched image serving: %d slots, staging caches %.0f MB " + "(ctx=%d each)\n", cfg_.max_concurrency, + double(estimate_ds4_cache_bytes(w_, max_ctx)) * cfg_.max_concurrency / (1024.0 * 1024.0), + max_ctx); } } else { if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { @@ -2336,7 +2459,7 @@ bool DeepSeek4Backend::init_hybrid_model() { std::fprintf(stderr, "[deepseek4] image prefill requires resident cold experts on the secondary HIP device\n"); return false; } - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { std::fprintf(stderr, "[deepseek4] failed to reload full model after placement: %s\n", @@ -2384,7 +2507,7 @@ bool DeepSeek4Backend::init_hybrid_model() { std::fprintf(stderr, "[deepseek4] %s experts cannot decode from hybrid/cold " "placement; falling back to monolithic full load\n", m.what); - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { std::fprintf(stderr, @@ -2678,8 +2801,7 @@ bool DeepSeek4Backend::park(ParkTarget target) { } moe_placement_ = {}; moe_decode_placement_ = {}; - join_image_stream(); - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); parked_ = true; if (spec_drafter_) { @@ -2699,7 +2821,7 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { if (want_target_model && parked_) { if (!load_model()) { std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); @@ -2718,7 +2840,7 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { "[deepseek4] unpark: failed to recreate KV cache (ctx=%d)\n", max_ctx); free_deepseek4_cache(cache_); - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); @@ -2734,7 +2856,7 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { if (env_flag_enabled("LUCE_DS4_MOE_TP") && !init_moe_tensor_parallel()) { free_deepseek4_cache(cache_); - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); expert_runtime_.reset(); stream_engine_.destroy(); @@ -2753,7 +2875,7 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { std::fflush(stdout); } if (!validate_prefill_mode()) { - vision_.reset(); + release_vision(); free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); @@ -3073,7 +3195,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, std::vector embed(w_.n_embd * n_tok); const auto embed_t0 = Clock::now(); const bool embedded = images - ? images->embed_chunk(w_.embedder, size_t(i), n_tok, embed.data()) + ? images->embed_chunk(w_.embedder, size_t(i), n_tok, embed.data(), + [&] { return io.is_cancelled(); }) : w_.embedder.embed(tokens.data() + i, n_tok, embed.data()); if (!embedded) return -1; DeepSeek4StepTelemetry step_tel; @@ -3419,12 +3542,14 @@ GenerateResult DeepSeek4Backend::generate_from_state( return result; } - const auto * images = dynamic_cast(req.images.get()); - // A streaming encode reads the payload and the IO: finish it before either ends. - struct ImageStreamJoin { - DeepSeek4Backend * backend; - ~ImageStreamJoin() { backend->join_image_stream(); } - } image_stream_join{this}; + const auto image_handle = std::dynamic_pointer_cast(req.images); + const DeepSeek4ImagePrompt * images = image_handle.get(); + // A streaming encode may still be running when generation ends early: + // stop it, so the encoder moves on to the next request. + struct ImageStreamStop { + const DeepSeek4ImagePrompt * images; + ~ImageStreamStop() { if (images) images->cancel_stream(); } + } image_stream_stop{images}; if (req.images) { if (!images || images->owner_ != this || !images->matches(req.prompt) || kv_offset != 0 || req.snap_slot >= 0 || req.snap_pos >= 0 || @@ -3436,7 +3561,7 @@ GenerateResult DeepSeek4Backend::generate_from_state( } std::string error; const auto encode_t0 = Clock::now(); - const bool encoded = materialize_images(*images, out_io, error); + const bool encoded = materialize_images(image_handle, out_io, error); if (encoded && !vision_backend_) { std::fprintf(stderr, "[deepseek4] images encoded in %.0f ms on the target GPU\n", elapsed_s(encode_t0) * 1000.0); @@ -3955,7 +4080,7 @@ void DeepSeek4Backend::maybe_save_routing_stats() { } void DeepSeek4Backend::shutdown() { - join_image_stream(); + release_vision(); maybe_save_routing_stats(); free_drafter(); for (int i = 0; i < PREFIX_SLOTS; i++) { @@ -3963,6 +4088,8 @@ void DeepSeek4Backend::shutdown() { } seq_engine_.reset(); free_deepseek4_paged_cache(paged_cache_); + for (auto & cache : image_staging_caches_) free_deepseek4_cache(*cache); + image_staging_caches_.clear(); free_deepseek4_cache(cache_); expert_runtime_.reset(); stream_engine_.destroy(); @@ -3975,7 +4102,7 @@ void DeepSeek4Backend::shutdown() { routing_stats_out_path_.clear(); moe_placement_ = {}; moe_decode_placement_ = {}; - vision_.reset(); + release_vision(); if (vision_backend_) { ggml_backend_free(vision_backend_); vision_backend_ = nullptr; } free_deepseek4_weights(w_); if (snap_backend_) { ggml_backend_free(snap_backend_); snap_backend_ = nullptr; } diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 729dd13ee..6b5bd110a 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -25,7 +25,11 @@ #include "ggml.h" #include "ggml-backend.h" +#include +#include +#include #include +#include #include #include #include @@ -135,8 +139,18 @@ class DeepSeek4Backend : public ModelBackend { // Owned backend for the vision encoder when --mmproj-device names a GPU // other than the target's; null when the encoder shares backend_. ggml_backend_t vision_backend_ = nullptr; - // Encodes images on vision_backend_ while prefill consumes them. - std::thread image_stream_; + // Encoder worker on vision_backend_: encodes queued image requests in + // order and publishes each image as it lands, so neither the scheduler + // nor prefill waits for a whole request. Started on first use. + std::thread encode_worker_; + std::mutex encode_mutex_; + std::condition_variable encode_ready_; + std::deque> encode_queue_; + std::atomic encode_stop_{false}; // also read by the worker's cancel check + vision::ImageSentinels image_sentinels_; + // Batched image serving: one single-request staging cache per slot + // (slot 0 uses cache_), allocated at startup. + std::vector> image_staging_caches_; vision::ImageAdmissionReserves image_reserves_; // Sampler @@ -203,24 +217,35 @@ class DeepSeek4Backend : public ModelBackend { int prefix_tokens = 0); bool load_vision(); bool init_single_gpu_vision(); - // Waits for a streaming image encode started by materialize_images. - void join_image_stream(); - // Batched serving. encode_image_request materializes an image request's - // rows; prefill_staged fills each request's first `prefix` tokens into - // its own staging cache in shared layer-major passes (expert weights read - // once per pass for every request in it). - struct StagedPrefill { - ImagePromptHandle images; - const std::vector * prompt = nullptr; - int prefix = 0; - DeepSeek4Cache * staging = nullptr; - bool ok = false; - std::string error; - }; + // Encodes one image with the vision runtime (caller serialises use). + bool encode_one_image(const vision::PromptImage & image, vision::ImageRaster & raster, + std::string & error); + // Queues an image request on the encoder worker (--mmproj-device only). + void enqueue_image_encode(std::shared_ptr images); + void encode_worker_loop(); + void cancel_image_encode(const ImagePromptPayload & images) const; + // Stops the encoder worker (failing queued requests), then frees vision_. + void release_vision(); + // Batched serving. A staged prefill fills one request's first `prefix` + // tokens into its slot's staging cache over several steps, in shared + // layer-major passes (expert weights read once per pass for all of them) + // that the engine advances a few layers per step. + using StagedPrefill = DeepSeek4StagedPrefill; + DeepSeek4Cache * image_staging_cache(int slot); + // Starts encoding an admitted image request: queued on the encoder + // worker with --mmproj-device, otherwise encoded here. bool encode_image_request(const std::vector & prompt, const ImagePromptHandle & images, std::string & error); - void prefill_staged(std::vector & batch); - bool materialize_images(const DeepSeek4ImagePrompt & images, + bool begin_staged_prefill(StagedPrefill & item); + // Starts one shared pass over the ready, unfinished items, about + // `row_budget` rows in total (a whole image block may exceed it); `rows` + // gets each item's share. False when no item is ready (or all failed). + bool begin_staged_pass(const std::vector & items, int row_budget, + DeepSeek4PrefillPass & pass, std::vector & rows); + // Waits up to `timeout_ms` for the next rows of an item to have their + // images encoded, so an otherwise idle scheduler does not spin. + void wait_staged_ready(const StagedPrefill & item, int row_budget, int timeout_ms) const; + bool materialize_images(const std::shared_ptr & images, const DaemonIO & io, std::string & error); // Generate after either a fresh prefill or a restored prefix. kv_offset is diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index c7a77c918..8a1d290c3 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -7744,10 +7744,18 @@ static bool initialize_layer_range_cache( DeepSeek4LayerRangeCache & runtime, ggml_backend_t backend, int device, const DeepSeek4Weights & w, int layer_begin, int layer_end, bool owns_output); -bool deepseek4_prefill_multi(ggml_backend_t backend, int device, - const DeepSeek4Weights & w, - const std::vector & seqs, - std::string & error) { +DeepSeek4PrefillPass::~DeepSeek4PrefillPass() { release(); } + +void DeepSeek4PrefillPass::release() { + if (state_buf_) { ggml_backend_buffer_free(state_buf_); state_buf_ = nullptr; } + if (state_ctx_) { ggml_free(state_ctx_); state_ctx_ = nullptr; } + state_in_ = state_out_ = nullptr; +} + +bool DeepSeek4PrefillPass::begin(ggml_backend_t backend, int device, const DeepSeek4Weights & w, + const std::vector & seqs, std::string & error) { + release(); + w_ = nullptr; const auto fail_early = [&](const char * why) { error = why; return false; }; if (!backend || seqs.empty() || w.moe_hybrid || !ds4_backend_is_gpu(backend)) return fail_early("shared prefill needs a full model on one GPU"); @@ -7769,31 +7777,27 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, !initialize_layer_range_cache(runtime, backend, device, w, 0, w.n_layer, true)) return fail_early("layer runtime initialization failed"); auto & fc = runtime.fused_decode_graph_cache; - const auto & hc_weights = runtime.hc_layer_weights; - const auto & hc_out_weights = runtime.hc_output_weights; - const auto & hash_tables = runtime.hash_routing_tables; if (fc.owner_ctx != w.ctx || fc.backend != backend) { fc.destroy(); fc.owner_ctx = w.ctx; fc.backend = backend; } - if (!ds4_fused_ensure_fn_mirrors(fc, backend, w, hc_weights, hc_out_weights)) + if (!ds4_fused_ensure_fn_mirrors(fc, backend, w, runtime.hc_layer_weights, runtime.hc_output_weights)) return fail_early("HC weight mirrors failed"); const int n_embd = w.n_embd, n_hc = w.n_hc; const int64_t hc_dim = (int64_t) n_embd * n_hc; - const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; // Row offsets, concatenated ids and "is image row" flags for the pass. - std::vector offset(seqs.size()); - std::vector ids((size_t) total); - std::vector image_row((size_t) total, 0); - bool any_image = false; + offset_.assign(seqs.size(), 0); + ids_.assign((size_t) total, 0); + image_row_.assign((size_t) total, 0); + any_image_ = false; for (size_t k = 0, off = 0; k < seqs.size(); off += (size_t) seqs[k].n_tokens, ++k) { - offset[k] = (int) off; - std::copy_n(seqs[k].token_ids, seqs[k].n_tokens, ids.begin() + (ptrdiff_t) off); + offset_[k] = (int) off; + std::copy_n(seqs[k].token_ids, seqs[k].n_tokens, ids_.begin() + (ptrdiff_t) off); for (int t = 0; t < seqs[k].n_tokens; ++t) { if (vision::image_block_at(seqs[k].image_spans, uint64_t(seqs[k].kv_start + t))) { - image_row[off + (size_t) t] = 1; - any_image = true; + image_row_[off + (size_t) t] = 1; + any_image_ = true; } } } @@ -7801,25 +7805,49 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, ggml_init_params state_params{}; state_params.mem_size = 4 * ggml_tensor_overhead() + 4096; state_params.no_alloc = true; - ggml_context * state_ctx = ggml_init(state_params); - if (!state_ctx) return fail_early("state context failed"); - ggml_tensor * state_a = ggml_new_tensor_2d(state_ctx, GGML_TYPE_F32, hc_dim, total); - ggml_tensor * state_b = ggml_new_tensor_2d(state_ctx, GGML_TYPE_F32, hc_dim, total); - ggml_backend_buffer_t state_buf = ggml_backend_alloc_ctx_tensors(state_ctx, backend); - if (!state_buf) { ggml_free(state_ctx); return fail_early("state allocation failed"); } + state_ctx_ = ggml_init(state_params); + if (!state_ctx_) return fail_early("state context failed"); + state_in_ = ggml_new_tensor_2d(state_ctx_, GGML_TYPE_F32, hc_dim, total); + state_out_ = ggml_new_tensor_2d(state_ctx_, GGML_TYPE_F32, hc_dim, total); + state_buf_ = ggml_backend_alloc_ctx_tensors(state_ctx_, backend); + if (!state_buf_) { release(); return fail_early("state allocation failed"); } { std::vector initial((size_t) hc_dim * total); for (size_t k = 0; k < seqs.size(); ++k) { for (int t = 0; t < seqs[k].n_tokens; ++t) { - float * dst = initial.data() + (size_t) (offset[k] + t) * hc_dim; + float * dst = initial.data() + (size_t) (offset_[k] + t) * hc_dim; for (int h = 0; h < n_hc; ++h) { std::memcpy(dst + (size_t) h * n_embd, seqs[k].embed + (size_t) t * n_embd, sizeof(float) * (size_t) n_embd); } } } - ggml_backend_tensor_set(state_a, initial.data(), 0, sizeof(float) * initial.size()); + ggml_backend_tensor_set(state_in_, initial.data(), 0, sizeof(float) * initial.size()); } + backend_ = backend; + seqs_ = seqs; + for (auto & s : seqs_) s.embed = nullptr; // read above only + total_ = total; + next_layer_ = 0; + w_ = &w; + return true; +} + +bool DeepSeek4PrefillPass::run_layers(int count, std::string & error) { + if (!w_ || !state_buf_ || done() || count <= 0) { + error = "no shared prefill pass to run"; + return false; + } + const DeepSeek4Weights & w = *w_; + ggml_backend_t backend = backend_; + const int total = total_; + DeepSeek4LayerRangeCache & runtime = *seqs_.front().cache->layer_range_cache; + auto & fc = runtime.fused_decode_graph_cache; + const auto & hc_weights = runtime.hc_layer_weights; + const auto & hash_tables = runtime.hash_routing_tables; + const int n_embd = w.n_embd, n_hc = w.n_hc; + const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; + ggml_gallocr_t alloc = ds4_layer_major_get_shared_alloc(w, backend); const size_t meta_bytes = 160u * 1024 * 1024; if (ds4_layer_major_meta_owner != w.ctx) { @@ -7828,19 +7856,19 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, ds4_layer_major_meta_owner = w.ctx; } if (ds4_layer_major_meta_arena.size() < meta_bytes) ds4_layer_major_meta_arena.resize(meta_bytes); + // A failed layer leaves the caches partly advanced: the pass is over. auto fail = [&](const char * what, int il) { std::fprintf(stderr, "[deepseek4-prefill-multi] %s at layer %d\n", what, il); - ggml_backend_buffer_free(state_buf); - ggml_free(state_ctx); + release(); + next_layer_ = w.n_layer; error = what; return false; }; if (!alloc) return fail("shared allocator unavailable", -1); std::vector hash_scratch; - ggml_tensor * state_in = state_a; - ggml_tensor * state_out = state_b; - for (int il = 0; il < w.n_layer; ++il) { + const int layer_end = std::min(w.n_layer, next_layer_ + count); + for (int il = next_layer_; il < layer_end; ++il) { ggml_init_params params{}; params.mem_size = ds4_layer_major_meta_arena.size(); params.mem_buffer = ds4_layer_major_meta_arena.data(); @@ -7852,11 +7880,11 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; // HC pre over every row of every sequence. - ggml_tensor * norm_hc = ggml_rms_norm(ctx, state_in, w.hc_eps); + ggml_tensor * norm_hc = ggml_rms_norm(ctx, state_in_, w.hc_eps); ggml_tensor * mix_attn = ggml_reshape_2d(ctx, ggml_mul_mat(ctx, fc.fn_attn_f16[(size_t) il], norm_hc), mix_dim, total); ggml_tensor * pre_attn = ggml_ds4_hc_pre( - ctx, mix_attn, ds4_fused_hc_base_f32(ctx, L.hc_attn_base), state_in, n_hc, + ctx, mix_attn, ds4_fused_hc_base_f32(ctx, L.hc_attn_base), state_in_, n_hc, w.n_hc_sinkhorn_iter, hlw.attn.scale_data[0], hlw.attn.scale_data[1], hlw.attn.scale_data[2]); ggml_tensor * attn_in = ggml_view_2d(ctx, pre_attn, n_embd, total, pre_attn->nb[1], 0); ggml_tensor * split_attn = ggml_view_2d(ctx, pre_attn, mix_dim, total, pre_attn->nb[1], @@ -7869,10 +7897,10 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, std::vector i64_array_inputs; std::vector f32_array_inputs; ggml_tensor * attn_out = nullptr; - for (size_t k = 0; k < seqs.size(); ++k) { - const DeepSeek4PrefillSeq & s = seqs[k]; + for (size_t k = 0; k < seqs_.size(); ++k) { + const DeepSeek4PrefillSeq & s = seqs_[k]; ggml_tensor * rows = ggml_view_2d(ctx, attn_normed, n_embd, s.n_tokens, attn_normed->nb[1], - (size_t) offset[k] * attn_normed->nb[1]); + (size_t) offset_[k] * attn_normed->nb[1]); ggml_tensor * out = build_mla_attention( ctx, gf, rows, w, L, s.cache->layers[(size_t) il], il, s.kv_start, s.n_tokens, nullptr, i32_inputs, i32_array_inputs, i64_array_inputs, &f32_array_inputs, @@ -7880,7 +7908,7 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, if (!out) { ggml_free(ctx); return fail("attention graph build failed", il); } attn_out = attn_out ? ggml_concat(ctx, attn_out, out, 1) : out; } - ggml_tensor * hc_after_attn = ggml_ds4_hc_post(ctx, state_in, attn_out, split_attn, n_hc); + ggml_tensor * hc_after_attn = ggml_ds4_hc_post(ctx, state_in_, attn_out, split_attn, n_hc); // HC pre -> one MoE FFN over all rows. norm_hc = ggml_rms_norm(ctx, hc_after_attn, w.hc_eps); @@ -7897,7 +7925,7 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, ggml_tensor * selection_bias = nullptr; ggml_tensor * hash_ids = nullptr; ggml_tensor * ffn_out = nullptr; - if (any_image) { + if (any_image_) { if (!L.ffn_gate_bias_vl) { ggml_free(ctx); return fail("image rows without an image router bias", il); } selection_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, w.n_expert, total); ggml_set_input(selection_bias); @@ -7911,7 +7939,7 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, } if (!ffn_out) { ggml_free(ctx); return fail("FFN graph build failed", il); } ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_after_attn, ffn_out, split_ffn, n_hc); - ggml_tensor * state_copy = ggml_cpy(ctx, hc_next, state_out); + ggml_tensor * state_copy = ggml_cpy(ctx, hc_next, state_out_); ggml_set_output(state_copy); ggml_build_forward_expand(gf, state_copy); @@ -7928,7 +7956,7 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, hash_scratch.resize((size_t) n_used * total); const auto & table = hash_tables[(size_t) il].ids; for (int t = 0; t < total; ++t) { - std::memcpy(hash_scratch.data() + (size_t) t * n_used, table.data() + (size_t) ids[(size_t) t] * n_used, + std::memcpy(hash_scratch.data() + (size_t) t * n_used, table.data() + (size_t) ids_[(size_t) t] * n_used, sizeof(int32_t) * (size_t) n_used); } ggml_backend_tensor_set(hash_ids, hash_scratch.data(), 0, sizeof(int32_t) * hash_scratch.size()); @@ -7946,11 +7974,11 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, std::vector bias(n_expert * (size_t) total); for (int t = 0; t < total; ++t) { float * row = bias.data() + (size_t) t * n_expert; - if (image_row[(size_t) t]) { std::copy(image_bias.begin(), image_bias.end(), row); continue; } + if (image_row_[(size_t) t]) { std::copy(image_bias.begin(), image_bias.end(), row); continue; } std::copy(text_bias.begin(), text_bias.end(), row); if (hash_routed) { const int32_t * picks = hash_tables[(size_t) il].ids.data() + - (size_t) ids[(size_t) t] * (size_t) w.n_expert_used; + (size_t) ids_[(size_t) t] * (size_t) w.n_expert_used; for (int j = 0; j < w.n_expert_used; ++j) row[picks[j]] = HASH_PICK; } } @@ -7959,18 +7987,20 @@ bool deepseek4_prefill_multi(ggml_backend_t backend, int device, if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { ggml_free(ctx); return fail("compute failed", il); } ggml_free(ctx); const int ratio = (int) w.compress_ratios[(size_t) il]; - for (const auto & s : seqs) { + for (const auto & s : seqs_) { if (ratio <= 0) continue; DeepSeek4LayerCache & lc = s.cache->layers[(size_t) il]; const int next_pos = s.kv_start + s.n_tokens; lc.n_comp = std::max(lc.n_comp, next_pos / ratio); if (ratio == 4) lc.n_index_comp = std::max(lc.n_index_comp, next_pos / ratio); } - std::swap(state_in, state_out); + std::swap(state_in_, state_out_); + } + next_layer_ = layer_end; + if (done()) { + for (const auto & s : seqs_) s.cache->cur_pos = s.kv_start + s.n_tokens; + release(); } - for (const auto & s : seqs) s.cache->cur_pos = s.kv_start + s.n_tokens; - ggml_backend_buffer_free(state_buf); - ggml_free(state_ctx); return true; } diff --git a/server/src/deepseek4/deepseek4_image_spans.h b/server/src/deepseek4/deepseek4_image_spans.h index 5dde1b6b6..925fd8a4b 100644 --- a/server/src/deepseek4/deepseek4_image_spans.h +++ b/server/src/deepseek4/deepseek4_image_spans.h @@ -13,4 +13,23 @@ inline bool valid_image_spans(ImageSpanView spans, uint64_t prompt_size) { return valid_image_spans(spans, prompt_size, DS4V_MAX_IMAGES, DS4V_MAX_IMAGE_BLOCK_TOKENS); } +// Rows the next staged prefill chunk takes at `position`, with `remaining` +// rows left in the prefix: about `budget` rows, whole image blocks only (a +// block, with any short text before it, may exceed the budget), and never a +// chunk or a leftover tail shorter than `min_rows`. 0 when no chunk exists. +inline int staged_prefill_chunk(ImageSpanView spans, uint64_t position, int remaining, + int budget, int min_rows) { + const auto valid = [&](int n) { + return n >= min_rows && (n == remaining || remaining - n >= min_rows); + }; + const auto cut = [&](int proposed) { + return atomic_image_chunk(spans, position, proposed, uint64_t(remaining), remaining); + }; + int n = cut(std::min(remaining, std::max(budget, min_rows))); + // Too short (text just before an image) or leaving a stub tail: grow to + // the next boundary that is valid. + for (int proposed = n + 1; !valid(n) && proposed <= remaining; ++proposed) n = cut(proposed); + return valid(n) ? n : 0; +} + } // namespace luce::vision diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 7e7432da7..c7d8ea909 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -39,6 +39,15 @@ inline constexpr int DS4_NUMERICAL_PREFILL_BAND = 2048; inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 10240; // Chunks of four rows or fewer take the decode-shaped path, not layer-major. inline constexpr int DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS = 5; +// Staged image prefill rows per batched step, shared by the pending requests. +// With live decoders the budget bounds how long they wait between tokens; +// with none, a larger pass shares each layer's expert reads across requests. +inline constexpr int DS4_STAGED_PREFILL_ROWS_PER_STEP = 256; +inline constexpr int DS4_STAGED_PREFILL_ROWS_WITHOUT_DECODE = 1024; +// With live decoders, a staged pass also runs only this many of its layers +// per step: an image block must go through each layer whole, but its layers +// can be spread over steps, so decoders wait for a slice, not the whole pass. +inline constexpr int DS4_STAGED_PREFILL_LAYERS_PER_STEP = 6; // Normal verification stays within one ratio-4 compressor window. Q5 is an // explicit opt-in whose fused graph models a second boundary. inline constexpr int DS4_CONSERVATIVE_VERIFY_MAX_TOKENS = 4; @@ -593,10 +602,46 @@ struct DeepSeek4PrefillSeq { // against its own cache; the HC mixing and the MoE FFN run once over all the // sequences' rows, so every layer's expert weights are read once for all of // them. Produces no logits and no feature capture. -bool deepseek4_prefill_multi(ggml_backend_t backend, int device, - const DeepSeek4Weights & w, - const std::vector & seqs, - std::string & error); +// +// The pass can run a few layers at a time, so a caller can interleave other +// work (batched decode) between slices: the hidden state stays on the GPU in +// between, and every layer still sees all rows of the pass, so whole-block +// image attention is unchanged. A pass cannot be abandoned half-way without +// leaving its caches' compressor state partly advanced. +class DeepSeek4PrefillPass { +public: + DeepSeek4PrefillPass() = default; + DeepSeek4PrefillPass(const DeepSeek4PrefillPass &) = delete; + DeepSeek4PrefillPass & operator=(const DeepSeek4PrefillPass &) = delete; + ~DeepSeek4PrefillPass(); + + // Validates the sequences and loads their embeddings; `embed` is read + // here only, `token_ids`, `image_spans` and `cache` until done(). + bool begin(ggml_backend_t backend, int device, const DeepSeek4Weights & w, + const std::vector & seqs, std::string & error); + // Runs up to `count` more layers; after the last one the caches' cur_pos + // is advanced and done() is true. + bool run_layers(int count, std::string & error); + bool done() const { return w_ && next_layer_ >= w_->n_layer; } + +private: + void release(); + + ggml_backend_t backend_ = nullptr; + const DeepSeek4Weights * w_ = nullptr; + std::vector seqs_; + std::vector offset_; + std::vector ids_; + std::vector image_row_; + bool any_image_ = false; + int total_ = 0; + int next_layer_ = 0; + ggml_context * state_ctx_ = nullptr; + ggml_backend_buffer_t state_buf_ = nullptr; + ggml_tensor * state_in_ = nullptr; + ggml_tensor * state_out_ = nullptr; +}; + bool deepseek4_validate_image_batch( const DeepSeek4Weights & w, const DeepSeek4Cache & cache, diff --git a/server/src/deepseek4/deepseek4_seq_engine.cpp b/server/src/deepseek4/deepseek4_seq_engine.cpp index f1f620248..312667254 100644 --- a/server/src/deepseek4/deepseek4_seq_engine.cpp +++ b/server/src/deepseek4/deepseek4_seq_engine.cpp @@ -33,7 +33,8 @@ DeepSeek4SeqEngine::DeepSeek4SeqEngine( offload_(slots_, pool, backend.backend_, ds4_paged_kv_planes(backend.paged_cache_)), stride_(table_stride), host_tables_((size_t)pool.max_sequences() * table_stride, -1), - reserve_growth_((size_t)pool.max_sequences(), 0) {} + reserve_growth_((size_t)pool.max_sequences(), 0), + slot_has_images_((size_t)pool.max_sequences(), 0) {} bool DeepSeek4SeqEngine::token_is_eos(int32_t token) const { return deepseek4_is_eos_tok(token, b_.w_); @@ -80,12 +81,6 @@ SeqEngine::AdmitResult DeepSeek4SeqEngine::admit( return result; } -DeepSeek4SeqEngine::~DeepSeek4SeqEngine() { - for (auto & cache : staging_caches_) { - if (cache) free_deepseek4_cache(*cache); - } -} - bool DeepSeek4SeqEngine::supports_images() const { return b_.image_capable_ && b_.vision_ && b_.cache_.buf; } @@ -96,12 +91,12 @@ SeqEngine::AdmitResult DeepSeek4SeqEngine::admit_images( AdmitResult refused; refused.status = AdmitResult::Status::failed; // DS4V image blocks need whole-block bidirectional prefill, which the - // 16-row gathered graph cannot run. Admission claims a slot, encodes the - // images and seeds the slot with every prompt token but the last; the - // next step() prefills all pending image requests together on the - // layer-major sparse path and copies their state into the slots, and the - // last (text) token then prefills in the batch, yielding the first - // sampled token as usual. + // 16-row gathered graph cannot run. Admission claims a slot, starts the + // image encode and seeds the slot with every prompt token but the last; + // later steps prefill that prefix into the slot's staging cache one shared + // layer-major pass at a time, copy it into the paged slot, and the last + // (text) token then prefills in the batch, yielding the first sampled + // token as usual. const int prefix = int(prompt.size()) - 1; if (!supports_images() || prefix < DS4_MIN_LAYER_MAJOR_PREFILL_TOKENS || prompt.size() > size_t(b_.cache_.max_ctx)) { @@ -112,65 +107,122 @@ SeqEngine::AdmitResult DeepSeek4SeqEngine::admit_images( // retries it, and encoding first would rerun the encoder on every retry. AdmitResult result = admit(request_id, prompt, sampler); if (result.status != AdmitResult::Status::admitted) return result; - std::string error; - if (!b_.encode_image_request(prompt, images, error)) { + if (staging_in_flight(result.slot)) { + // The slot's staging cache still belongs to a pass in flight (its + // previous request retired mid-pass): retry after the pass. retire(result.slot); - refused.error = error.empty() ? "image encoding failed" : error; - return refused; + result.status = AdmitResult::Status::busy; + result.slot = -1; + result.error = "image staging cache is still in use"; + return result; + } + PendingImage pending; + pending.slot = result.slot; + pending.staged.images = images; + pending.staged.prompt = prompt; + pending.staged.prefix = prefix; + pending.staged.staging = b_.image_staging_cache(result.slot); + std::string error; + bool ok = pending.staged.staging && b_.encode_image_request(prompt, images, error) && + b_.begin_staged_prefill(pending.staged); + SeqSlotManager::PrefillChunk seeded; + if (ok) { + seeded = slots_.seed_restored_prefix(result.slot, prefix); + ok = seeded.ok && seeded.rows.size() == size_t(prefix); } - SeqSlotManager::PrefillChunk seeded = slots_.seed_restored_prefix(result.slot, prefix); - bool ok = seeded.ok && seeded.rows.size() == size_t(prefix); for (size_t i = 0; ok && i < seeded.new_blocks.size(); ++i) { ok = set_block(result.slot, seeded.first_new_block + int(i), seeded.new_blocks[i]); } if (!ok) { retire(result.slot); - refused.error = "image prefix could not be seeded into the paged slot"; + refused.error = !error.empty() ? error : !pending.staged.error.empty() ? pending.staged.error + : "image request could not be staged"; return refused; } - pending_images_.push_back({result.slot, images, prompt, prefix}); + slot_has_images_[size_t(result.slot)] = 1; + pending_images_.push_back(std::move(pending)); return result; } -void DeepSeek4SeqEngine::run_pending_images(std::vector & failures, - std::vector & failed_slots) { - if (pending_images_.empty()) return; - std::vector batch(pending_images_.size()); - for (size_t k = 0; k < pending_images_.size(); ++k) { - // Staging cache k: the backend's own for the first request, then - // one more per concurrent request (allocated once, reused). - DeepSeek4Cache * staging = &b_.cache_; - if (k > 0) { - while (staging_caches_.size() < k) staging_caches_.emplace_back(); - auto & owned = staging_caches_[k - 1]; - if (!owned) { - owned = std::make_unique(); - if (!create_deepseek4_cache(b_.backend_, b_.w_, b_.cache_.max_ctx, *owned)) owned.reset(); +DeepSeek4SeqEngine::PendingImage * DeepSeek4SeqEngine::pending_image(int slot) { + for (auto & pending : pending_images_) { + if (pending.slot == slot) return &pending; + } + return nullptr; +} + +bool DeepSeek4SeqEngine::staging_in_flight(int slot) const { + if (!staged_pass_) return false; + for (const auto & member : staged_pass_->members) { + if (member.slot == slot) return true; + } + return false; +} + +void DeepSeek4SeqEngine::advance_pending_images(bool decoding, bool idle) { + if (!staged_pass_) { + std::vector items; + std::vector slots; + for (auto & pending : pending_images_) { + if (pending.staged.finished()) continue; + items.push_back(&pending.staged); + slots.push_back(pending.slot); + } + if (items.empty()) return; + auto staged = std::make_unique(); + std::vector rows; + if (b_.begin_staged_pass(items, decoding ? DS4_STAGED_PREFILL_ROWS_PER_STEP + : DS4_STAGED_PREFILL_ROWS_WITHOUT_DECODE, + staged->pass, rows)) { + for (size_t k = 0; k < items.size(); ++k) { + if (rows[k] > 0) staged->members.push_back({slots[k], items[k]->images, rows[k]}); } - staging = owned.get(); + staged->started = std::chrono::steady_clock::now(); + staged_pass_ = std::move(staged); + } else if (idle) { + // Nothing ready and nothing else to run: wait briefly for the + // encoder instead of spinning the scheduler. + b_.wait_staged_ready(*items.front(), DS4_STAGED_PREFILL_ROWS_PER_STEP, 20); } - batch[k].images = pending_images_[k].images; - batch[k].prompt = &pending_images_[k].prompt; - batch[k].prefix = pending_images_[k].prefix; - batch[k].staging = staging; - if (!staging) batch[k].error = "staging cache allocation failed"; - } - b_.prefill_staged(batch); - for (size_t k = 0; k < batch.size(); ++k) { - const PendingImage & p = pending_images_[k]; - std::string error = batch[k].error; - bool ok = batch[k].ok; - if (ok) { - ok = import_deepseek4_paged_slot( - *batch[k].staging, p.prefix, b_.paged_cache_, uint32_t(p.slot), - host_tables_.data() + size_t(p.slot) * stride_, stride_, error); + } + if (staged_pass_) { + // With live decoders, only a slice of the pass's layers: they then + // decode after each slice instead of waiting for the whole pass. + StagedPass & staged = *staged_pass_; + std::string error; + const bool ok = staged.pass.run_layers( + decoding ? DS4_STAGED_PREFILL_LAYERS_PER_STEP : b_.w_.n_layer, error); + ++staged.slices; + if (!ok || staged.pass.done()) { + int rows = 0; + for (const auto & member : staged.members) { + rows += member.rows; + PendingImage * pending = pending_image(member.slot); + // A member that retired (or whose slot now holds another + // request) is skipped. + if (!pending || pending->staged.images != member.images) continue; + if (ok) pending->staged.done += member.rows; + else pending->staged.error = error.empty() ? "staged prefill failed" : error; + } + std::fprintf(stderr, "[deepseek4] staged prefill pass: %zu requests, %d rows, %d slices, %.0f ms\n", + staged.members.size(), rows, staged.slices, + std::chrono::duration( + std::chrono::steady_clock::now() - staged.started).count()); + staged_pass_.reset(); } - if (!ok) { - fail_prefill(p.slot, failures, error.empty() ? "staged image prefill failed" : error); - if (p.slot >= 0 && p.slot < int(failed_slots.size())) failed_slots[size_t(p.slot)] = 1; + } + for (auto & pending : pending_images_) { + DeepSeek4StagedPrefill & staged = pending.staged; + if (!staged.error.empty() || staged.done < staged.prefix || !staged.staging) continue; + if (!import_deepseek4_paged_slot( + *staged.staging, staged.prefix, b_.paged_cache_, uint32_t(pending.slot), + host_tables_.data() + size_t(pending.slot) * stride_, stride_, staged.error) && + staged.error.empty()) { + staged.error = "staged image prefill could not be copied into the paged slot"; } + // Copied (or failed): the staging cache is free for the next request. + staged.staging = nullptr; } - pending_images_.clear(); } bool DeepSeek4SeqEngine::set_block(int slot, int logical, int32_t physical) { @@ -243,8 +295,11 @@ SeqEngine::StepResult DeepSeek4SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; - std::vector image_failed((size_t)n_slots, 0); - run_pending_images(result.prefills, image_failed); + if (!pending_images_.empty()) { + bool idle = inputs.empty(); + for (const PrefillSlice & slice : plan.prefills) idle = idle && pending_image(slice.slot); + advance_pending_images(!inputs.empty(), idle); + } std::vector lane_tokens; std::vector lane_positions; @@ -296,7 +351,30 @@ SeqEngine::StepResult DeepSeek4SeqEngine::step(const StepPlan & plan) { std::vector prefill_lanes; prefill_lanes.reserve(plan.prefills.size()); for (const PrefillSlice & slice : plan.prefills) { - if (image_failed[(size_t)slice.slot]) continue; + if (PendingImage * pending = pending_image(slice.slot)) { + // A staged image request reports its failure here, where its + // slot is planned; until its prefix is copied in, it sits out. + const DeepSeek4StagedPrefill & staged = pending->staged; + const bool failed = !staged.error.empty(); + const bool copied = !failed && staged.done >= staged.prefix && !staged.staging; + if (failed) { + fail_prefill(slice.slot, result.prefills, staged.error); + // Its queued encode must not keep the encoder from others. + if (staged.images) b_.cancel_image_encode(*staged.images); + } + if (failed || copied) { + pending_images_.erase(pending_images_.begin() + (pending - pending_images_.data())); + } + if (!copied) { + // Still staging: report the slot as advanced, with no row run. + if (!failed) { + PrefillOutput waiting; + waiting.slot = slice.slot; + result.prefills.push_back(std::move(waiting)); + } + continue; + } + } const SeqSlot & before = slots_.slot(slice.slot); const int remaining = before.prompt_len - before.cur_pos; const int n_rows = std::min(slice.max_tokens, remaining); @@ -454,8 +532,28 @@ bool DeepSeek4SeqEngine::restore_kv(int slot, std::string & error) { return true; } +bool DeepSeek4SeqEngine::kv_recomputable(int slot) const { + return slot >= 0 && size_t(slot) < slot_has_images_.size() && !slot_has_images_[size_t(slot)]; +} + +bool DeepSeek4SeqEngine::offload_kv(int slot, size_t bytes, std::string & error) { + // A slot still staging its image prefix has no paged KV to save yet. + for (const auto & pending : pending_images_) { + if (pending.slot == slot) { + error = "image request is still prefilling"; + return false; + } + } + return offload_.suspend(slot, bytes, error); +} + bool DeepSeek4SeqEngine::evict_kv(int slot, int32_t pending_token, std::string & error) { + if (!kv_recomputable(slot)) { + // Recompute replays token ids; image rows would come back as text. + error = "image request KV cannot be rebuilt from its tokens"; + return false; + } if (slot < 0 || slot >= slots_.slot_count() || !slots_.evict_for_recompute(slot, pending_token)) { error = "slot history cannot be re-prefilled within the paged KV pool"; @@ -466,9 +564,15 @@ bool DeepSeek4SeqEngine::evict_kv(int slot, int32_t pending_token, } void DeepSeek4SeqEngine::retire(int slot) { + for (auto & pending : pending_images_) { + if (pending.slot != slot) continue; + // Stop its encode before its payload may go away. + if (const auto * images = pending.staged.images.get()) b_.cancel_image_encode(*images); + } pending_images_.erase(std::remove_if(pending_images_.begin(), pending_images_.end(), [slot](const PendingImage & p) { return p.slot == slot; }), pending_images_.end()); + if (slot >= 0 && size_t(slot) < slot_has_images_.size()) slot_has_images_[size_t(slot)] = 0; offload_.discard(slot); if (!slots_.is_active(slot)) return; slots_.retire(slot); diff --git a/server/src/deepseek4/deepseek4_seq_engine.h b/server/src/deepseek4/deepseek4_seq_engine.h index c10126622..2cc9e4c31 100644 --- a/server/src/deepseek4/deepseek4_seq_engine.h +++ b/server/src/deepseek4/deepseek4_seq_engine.h @@ -3,7 +3,9 @@ #include "common/concurrency/seq_engine.h" #include "common/concurrency/paged_kv_offload.h" #include "common/concurrency/seq_slot_manager.h" +#include "deepseek4_internal.h" +#include #include #include #include @@ -12,7 +14,18 @@ namespace luce::common { class DeepSeek4Backend; -struct DeepSeek4Cache; + +// A batched image request's prompt prefix, prefilled into its slot's staging +// cache over several steps before it is copied into the paged slot. +struct DeepSeek4StagedPrefill { + ImagePromptHandle images; + std::vector prompt; + int prefix = 0; + DeepSeek4Cache * staging = nullptr; + int done = 0; // leading prompt rows already in `staging` + std::string error; // set once the request cannot finish + bool finished() const { return !error.empty() || done >= prefix; } +}; // Exact concurrent serving path for DeepSeek4. Model state remains in // DeepSeek4PagedCache; this class owns only scheduler-facing slot state and @@ -21,7 +34,6 @@ class DeepSeek4SeqEngine final : public SeqEngine { public: DeepSeek4SeqEngine(DeepSeek4Backend & backend, PagedKvPool & pool, int max_ctx, uint32_t table_stride); - ~DeepSeek4SeqEngine() override; int slot_count() const override { return slots_.slot_count(); } int max_context() const override { return slots_.max_context(); } @@ -32,11 +44,10 @@ class DeepSeek4SeqEngine final : public SeqEngine { bool reserve_decode(const StepPlan & plan) override; size_t kv_offload_capacity() const override { return offload_.capacity(); } KvOffloadState kv_offload_state(int slot) const override { return offload_.state(slot); } - bool offload_kv(int slot, size_t bytes, std::string & error) override { - return offload_.suspend(slot, bytes, error); - } + bool offload_kv(int slot, size_t bytes, std::string & error) override; bool restore_kv(int slot, std::string & error) override; bool evict_kv(int slot, int32_t pending_token, std::string & error) override; + bool kv_recomputable(int slot) const override; bool kv_restore_feasible(int slot) const override { return slots_.kv_restore_feasible(slot); } @@ -53,18 +64,33 @@ class DeepSeek4SeqEngine final : public SeqEngine { void fail_prefill(int slot, std::vector & outputs, const std::string & error); - // Image requests admitted since the last step: their slots hold the - // prompt minus its last token as seeded blocks, and step() fills those - // blocks from shared staged prefill passes before running the batch. + // Image requests still being staged: their slots hold the prompt minus + // its last token as seeded blocks. Every step advances them by one shared + // pass (so decode keeps running between passes) and copies each finished + // one into its paged slot; until then the slot's last-token prefill waits. struct PendingImage { int slot = -1; - ImagePromptHandle images; - std::vector prompt; - int prefix = 0; + DeepSeek4StagedPrefill staged; }; std::vector pending_images_; - std::vector> staging_caches_; - void run_pending_images(std::vector & failures, std::vector & failed_slots); + // The shared pass in flight: it runs a few layers per step and cannot be + // abandoned half-way, so members that retire meanwhile are just skipped + // when it completes, and their slots' staging caches stay busy until then. + struct StagedPass { + struct Member { + int slot = -1; + ImagePromptHandle images; // also identifies the request + int rows = 0; + }; + DeepSeek4PrefillPass pass; + std::vector members; + int slices = 0; + std::chrono::steady_clock::time_point started; + }; + std::unique_ptr staged_pass_; + bool staging_in_flight(int slot) const; + PendingImage * pending_image(int slot); + void advance_pending_images(bool decoding, bool idle); DeepSeek4Backend & b_; SeqSlotManager slots_; @@ -72,6 +98,8 @@ class DeepSeek4SeqEngine final : public SeqEngine { uint32_t stride_ = 0; std::vector host_tables_; std::vector reserve_growth_; + // Slots whose KV holds image rows: token history cannot rebuild them. + std::vector slot_has_images_; }; } // namespace luce::common diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 91705cdfe..0842ca17a 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -913,10 +913,15 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // No checkpoint fit the RAM cap: park the newest decoder for // recompute. With recovery disabled, fail only that request // before compute, then retry the remaining cohort. - int victim = -1; + // Prefer the newest decoder whose KV its token history can + // rebuild; if none can be, the newest decoder is failed below. + int victim = -1, fallback = -1; for (int candidate : residents) { - if (!slots[(size_t)candidate].prefilling) { victim = candidate; break; } + if (slots[(size_t)candidate].prefilling) continue; + if (fallback < 0) fallback = candidate; + if (engine.kv_recomputable(candidate)) { victim = candidate; break; } } + if (victim < 0) victim = fallback; if (victim < 0) break; // engines reserve prefills at admission auto & s = slots[(size_t)victim]; if (offload_budget && engine.evict_kv(victim, s.pending_tok, error)) { diff --git a/server/test/test_ds4v_image_integration.cpp b/server/test/test_ds4v_image_integration.cpp index 6e2bc3306..a79e3a149 100644 --- a/server/test/test_ds4v_image_integration.cpp +++ b/server/test/test_ds4v_image_integration.cpp @@ -77,6 +77,44 @@ void validation_and_lookup() { require(last_image_end_in(view(spans), 12, 26) == 25, "last overlapping image wins"); require(last_image_end_in(view(spans), 25, 30) == 0, "text between images"); require(last_image_end_in(view(spans), 34, 40) == 35, "chunk starting inside an image"); + { + // Staged prefill chunks: 2 text rows, a 362-row image, 2 text rows + // (the shape of a one-image chat prompt), with a 256-row budget. + const std::vector chat{{2, 3, 363, 364}}; + require(staged_prefill_chunk(view(chat), 0, 366, 256, 5) == 366, + "short text before an image grows to take the image and the stub tail"); + const std::vector two{{2, 3, 99, 100}, {300, 301, 399, 400}}; + require(staged_prefill_chunk(view(two), 0, 500, 256, 5) == 256, "text runs to the budget"); + require(staged_prefill_chunk(view(two), 256, 244, 256, 5) == 244, "an image and the tail"); + // Every layout of one or two images walks to the end in valid chunks. + const auto walk = [&](const std::vector & layout, int prefix, int budget) { + int done = 0; + while (done < prefix) { + const int n = staged_prefill_chunk(view(layout), uint64_t(done), prefix - done, budget, 5); + require(n >= 5, "staged chunk makes progress"); + require(prefix - done - n == 0 || prefix - done - n >= 5, "no stub tail"); + for (const TokenSpan & span : layout) { + require(!(uint64_t(done) < span.block_end && uint64_t(done + n) > span.block_begin && + (uint64_t(done) > span.block_begin || uint64_t(done + n) < span.block_end)), + "image block stays whole"); + } + done += n; + } + }; + for (uint64_t a = 0; a < 12; ++a) + for (uint64_t len = 5; len < 40; len += 7) + for (int budget : {5, 8, 16, 256}) { + for (int prefix = int(a + len); prefix < int(a + len) + 12; ++prefix) + walk({{a, a, a + len, a + len}}, prefix, budget); + for (uint64_t gap = 0; gap < 9; gap += 2) + for (uint64_t len2 = 5; len2 < 30; len2 += 8) { + const uint64_t b = a + len + gap; + for (int tail = 0; tail < 7; ++tail) + walk({{a, a, a + len, a + len}, {b, b, b + len2, b + len2}}, + int(b + len2) + tail, budget); + } + } + } require(valid_image_spans(view(spans), 35), "adjacent and separated blocks valid"); require(!image_block_at(view(spans), 9), "text before block excluded"); require(image_block_at(view(spans), 10) == &spans[0], "leading padding belongs to image block");