From 28a31e560eeaeccfaadcfbf7442838781aa63da9 Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Mon, 19 Oct 2026 05:04:52 +0300 Subject: [PATCH 1/8] perf(qwen4_exp): rank threads, pipeline windows and decode groups for stages without peer access Split Flash-Next models on GPUs without peer access (Windows/WDDM) queued every stage from one host thread, which a cross-device wait blocks, so the stages took turns. Opt-in rank threads give each stage its own host thread; pipeline windows run the guessed next MTP window on the front stages from shadow states while the head verifies; decode groups step several requests as round-robin groups. The head stage now reads its MTP cells' embeddings from a mapped host copy of the table, so split models' draft chains replay their graph. --- docs/qwen3-8-flash-next.md | 76 ++ src/models/qwen4_exp/executor.cpp | 1477 ++++++++++++++++++++++--- src/models/qwen4_exp/executor.h | 26 + src/runtime/engine/qwen4_exp_core.cpp | 327 +++++- 4 files changed, 1778 insertions(+), 128 deletions(-) diff --git a/docs/qwen3-8-flash-next.md b/docs/qwen3-8-flash-next.md index 4a090df034..9201446d0b 100644 --- a/docs/qwen3-8-flash-next.md +++ b/docs/qwen3-8-flash-next.md @@ -628,6 +628,48 @@ draft-chain graph and each width its own verification graph, captured at their s - A pass whose layers may wait on the device for their missing experts ends synchronized: CUDA loads a kernel lazily at its first launch, a load may wait for the device to go idle, and a launch during such a wait would deadlock with the service it waits for. +- A split model looks up its MTP cells' token embeddings on the head stage from a copy of the + table's payload in mapped pinned host memory (a row of the bus per cell) instead of asking stage 0 + and copying the rows back, so the draft chain stays on one device and replays its graph there. + `NINFER_QWEN4_HEAD_EMBEDDING=0` keeps the round trip. + +### Stages without peer access (rank threads, pipeline windows) + +Between GPUs without peer access (every GeForce pair under Windows/WDDM), a stream that waits on +another device's event, or a staged peer copy, blocks the host thread that queues it until the other +device gets there. One host thread queuing every stage therefore walks each pass through the +pipeline at device speed, and the stages take turns. Three opt-in modes, for device experts with +CUDA Graphs, change that: + +- **Rank threads** (`NINFER_QWEN4_RANK_THREADS=1`). The host thread prepares a pass's inputs + (positions, n-gram rows) into a staging entry and posts it; each stage's own thread uploads them, + waits on the host for its stack from the stage before (the only thread that blocks), runs its + segments and hands the stack on through pinned host buffers counted per boundary. Prompt chunks + then fill every stage, and the head stage's MTP catch-up takes the pass's embeddings from host + memory instead of a round trip to stage 0 that would order stage 0's next chunk after the head. + Calls that read results or state first wait until the threads have queued every posted pass. +- **Pipeline windows** (`NINFER_QWEN4_PIPELINE_WINDOWS=1`, with rank threads; one request at a time, + MTP drafts, no grammar or logprobs). A round's verification goes to the rank threads, and before + it the head stage runs the MTP chain on past the round's drafts: the bonus token the window would + end with if it were kept whole, and the drafts after it. The front stages run that guessed window + at once, while the head stage verifies the current one, from shadow states: a second slot of + each Gated DeltaNet pool and copies of the indexer tails and PLE histories, advanced over the + window in flight as if it were kept whole (the fold of a commit into the shadow slot), with a + second set of replay records so the window in flight can still commit any prefix. When the verdict + keeps the window whole with the guessed bonus, the head stage completes the guessed window from + the front stages' stack and the next guess chain runs through its tokens; otherwise the guess is + dropped and the next round drafts as before. Commits run on the rank threads too. Every token is + still the verification's, and the weighted expert outputs are summed in fixed point, so the output + is the serial loop's byte for byte. `NINFER_QWEN4_PIPELINE_FORCE_MISS=1` drops every guess (the + shadow work still runs), a test of the shadow states. +- **Decode groups** (`NINFER_QWEN4_DECODE_GROUPS=2..8`, with rank threads and two or more stages). + Several decoding requests step as up to that many groups in round robin, each group's plain step + posted while the others are in the pipeline, with its sampling queued on the head stage's thread; + a request alone keeps MTP. `NINFER_QWEN4_GROUP_STEPS` sets the steps a group runs before the + others get a turn (default 8). + +`NINFER_QWEN4_ROUND_TIMING=1` prints the host's time per MTP round phase every 100 rounds (with +the guesses made and held), and for host-ordered verifications each stage's device time in a window. Without CUDA Graphs the rank threads stay off. ## Measurements @@ -896,6 +938,40 @@ With host experts a verification's four columns route to up to four times as man missing one crossing the bus, so three drafts gained about 12% here; after the October 9 host-expert work three drafts decode 90-123 tok/s against 87-94 plain (table above). +### Stages without peer access on Windows + +IQ3_S with the IQ4_NL table read from its file, device experts, Windows 11 (WDDM, no peer access +between any pair, every card on PCIe x1 risers), CUDA 12.8 build for `75;86`, GPU clocks held at +1500-2100 MHz, the built-in device profiles. A is a TITAN RTX + four RTX 3060 12 GB (`--stage-layers +8,8,8,8,16`, the TITAN last with the head, rk8v4 KV); B is two RTX 3060 + a TITAN RTX + an RTX 3090 +(`8,8,17,15` with int8 KV, `8,8,16,16` with rk8v4). A 206,053-token prompt (the repository's docs), +262,144-token context, 2,048-token chunks, greedy, 512 new tokens, one run per row. + +| Layout, KV | Mode | Prefill tok/s | Decode tok/s | +|---|---|---:|---:| +| A, rk8v4 | host-ordered, MTP3 | 516 | 37.1 | +| A, rk8v4 | rank threads, MTP2 | 1,720 | 36.5 | +| A, rk8v4 | rank threads + pipeline windows, MTP2 (26% of guesses held) | 1,720 | 45.7 | +| A, rk8v4 | rank threads + pipeline windows, MTP3 (20% held) | 1,720 | 40.6 | +| B, rk8v4 | host-ordered, MTP3 | 683 | 42.2 | +| B, rk8v4 | rank threads + pipeline windows, MTP2 (32% held) | 1,640 | 60.4 | +| B, int8 | rank threads, MTP2 | 1,590 | 46.3 | +| B, int8 | rank threads + pipeline windows, MTP2 (31% held) | 1,590 | 62.2 | + +With every guess dropped (`NINFER_QWEN4_PIPELINE_FORCE_MISS=1`) B with rk8v4 and MTP2 still decoded +at 53.4 tok/s against 43.8 host-ordered: the verifications and commits on the rank threads, which no +longer wait for each other on the host, are about half of the gain, and the held guesses the rest. +Every pipeline-window run generated the host-ordered run's text at the same draft count. A round of +B (rk8v4, MTP2) took 45.5 ms host-ordered (draft 4.8, verification 39.9) and 33.9 ms with pipeline +windows. Through `ninfer-serve` (one slot of 262,144 tokens, a 15-token chat prompt, MTP2) A decoded +at 65.1 tok/s against 49.3 and B (int8) at 92.0 against 63.0, with the same text; a sampled request +(temperature 0.7), a request cancelled after 3 s and three requests in a row then ran normally. The +head stage's embedding copy alone took B's draft chain from 8.5 to 6.5 ms a round (MTP3). + +Four concurrent chat requests (512 tokens each, four slots of 65,536 tokens, rk8v4): B 139.7 tok/s +in all with decode groups (`NINFER_QWEN4_DECODE_GROUPS=4`) against 45.7 with MTP3 rounds, A 112.4 +against 40.9, each request's text unchanged; one request alone kept MTP (61.0 and 49.1 tok/s). + ### Hyper-connections in Q8_0 The hyper-connection projections (97 down/up pairs and 96 inject matrices) are stored as ggml Q8_0, diff --git a/src/models/qwen4_exp/executor.cpp b/src/models/qwen4_exp/executor.cpp index a0099ee1b8..aef247e72d 100644 --- a/src/models/qwen4_exp/executor.cpp +++ b/src/models/qwen4_exp/executor.cpp @@ -44,9 +44,18 @@ #include #include +#include #include +#include +#include +#include #include +#include +#include #include +#include +#include +#include #include #include #include @@ -489,8 +498,12 @@ struct RankState { DeviceBuffer activations; std::unique_ptr workspace; DeviceBuffer block_table; - std::unique_ptr staging; - cudaEvent_t staged = nullptr; + // A ring of host staging buffers, one per pass in flight (Impl::staging_slot picks the + // pass's), each with the event that marks its uploads done: a pass rewrites a buffer only + // once the pass that used it ring-size passes ago has uploaded it, so prompt chunks can fill + // every stage of a pipeline instead of the host waiting for the previous pass on all ranks. + std::vector> staging; + std::vector staged; cudaEvent_t done = nullptr; cudaEvent_t copied = nullptr; // its stack has been copied in from the previous rank cudaEvent_t routes = nullptr; // the last pass's routes reached the host @@ -527,10 +540,12 @@ struct RankState { // DeltaNet layers and their fold, each sparse-attention layer's indexer projections (BF16 // [640, columns] per layer), the PLE layer's normalised convolution input (FP32 // [10240, columns]), and scratch for the states a verification reads but must not advance. - DeviceBuffer record_backing; - std::optional records; - std::optional fold; - DeviceBuffer index_records, ple_records; + // Pipeline windows keep a second record set, so a guessed window can record while the window + // before it still awaits its commit, and each set's positions for that commit. + std::array record_backing; + std::array, 2> records; + std::array, 2> fold; + std::array index_records, ple_records, window_positions; DeviceBuffer conv_scratch, tail_scratch; // Streamed sparse-attention layers: a prompt chunk reads one layer's whole KV from this // device copy (the four planes, every page), which holds `kv_staged`'s positions before @@ -585,6 +600,8 @@ struct Executor::Impl { std::size_t rows_rank = 0; bool rows_waited = true; bool warming = false; // warm_up() runs: the expert cache counts no routes + // The staging buffer of every rank the current pass uses (RankState::staging). + std::size_t staging_slot = 0; std::array stall_events{}; std::size_t stall_next = 0; std::uint64_t stalls = 0; @@ -608,6 +625,19 @@ struct Executor::Impl { std::int32_t domain = 0; bool verifying = false; std::vector verified; // the last verify()'s sequences, in order + // What the segments a thread queues compute: a verification (its record set, verified rows and + // width), or a guessed window ahead of the verdict on the front ranks (pipeline windows), which + // reads the sequence's shadow states. The host thread sets it around verify(); a rank thread + // takes it from its task. + struct PassMode { + bool verify = false; + bool ahead = false; + std::uint32_t set = 0; + std::int32_t rows = 0; + std::uint32_t width = 0; + }; + static inline thread_local PassMode pass_mode{}; + std::uint32_t verified_set = 0; // the record set of the last verification // The MTP chain on the head rank: token ids and positions of a pass's cells (a draft step's at // step * sequences), staged from the host. DeviceBuffer mtp_ids, mtp_positions, ngram_prefetch_tokens; @@ -812,6 +842,7 @@ struct Executor::Impl { throw std::invalid_argument("--expert-prefill-buffers needs GGUF host experts"); } plan_segments(); + pipeline = pipeline_requested(); allocate_ranks(); allocate_vision(); allocate_sequences(); @@ -828,6 +859,9 @@ struct Executor::Impl { } ~Impl() { + // Every queued pass reaches its device before the buffers it reads go away. + drain_quietly(); + stop_rank_threads(); draft_prefetch.reset(); // The copier's queued copies finish before the events and buffers they use go. for (auto& prefetch : prefetches) { @@ -839,6 +873,12 @@ struct Executor::Impl { } } pool_copier.reset(); + for (std::size_t i = 0; i < window_timing.events.size(); ++i) { + RankBinding bind(device, segments[i].rank); + for (cudaEvent_t event : window_timing.events[i]) { + if (event != nullptr) { cudaEventDestroy(event); } + } + } // A pass that failed between starting the row reads and waiting for them leaves them // writing into a staging buffer. if (table) { @@ -866,7 +906,7 @@ struct Executor::Impl { } for (auto& rank : ranks) { RankBinding bind(device, rank.rank); - if (rank.staged != nullptr) { cudaEventDestroy(rank.staged); } + for (cudaEvent_t event : rank.staged) { cudaEventDestroy(event); } if (rank.done != nullptr) { cudaEventDestroy(rank.done); } if (rank.copied != nullptr) { cudaEventDestroy(rank.copied); } if (rank.routes != nullptr) { cudaEventDestroy(rank.routes); } @@ -878,6 +918,10 @@ struct Executor::Impl { RankBinding bind(device, model.head_rank()); cudaEventDestroy(mtp_staged); } + for (const CrossRing& ring : crossings) { + for (cudaEvent_t event : ring.written) { cudaEventDestroy(event); } + for (cudaEvent_t event : ring.read) { cudaEventDestroy(event); } + } for (const StallEvents& events : stall_events) { RankBinding bind(device, rows_rank); if (events.wanted != nullptr) { cudaEventDestroy(events.wanted); } @@ -1180,16 +1224,29 @@ struct Executor::Impl { offset += round_up(bytes); } rank.workspace = std::make_unique(workspace_bytes(r)); - std::vector identity(std::max(pages, options.sequences)); + // Pool slots: one per sequence, and with pipeline windows the shadow slot after them. + const std::uint32_t slots = options.sequences + (pipeline ? 1 : 0); + std::vector identity(std::max(pages, slots)); std::iota(identity.begin(), identity.end(), 0); rank.block_table = DeviceBuffer(std::size_t(pages) * 4); rank.block_table.copy_from_host(identity.data(), std::size_t(pages) * 4); - rank.slot_ids = DeviceBuffer(std::size_t(options.sequences) * 4); - rank.slot_ids.copy_from_host(identity.data(), std::size_t(options.sequences) * 4); + rank.slot_ids = DeviceBuffer(std::size_t(slots) * 4); + rank.slot_ids.copy_from_host(identity.data(), std::size_t(slots) * 4); // ids, positions, RoPE positions, scatter columns, three-axis RoPE positions, rows. - rank.staging = std::make_unique( - round_up(t * 28) + round_up(std::uint64_t(config.ngram_heads()) * t * row_b)); - CUDA_CHECK(cudaEventCreateWithFlags(&rank.staged, cudaEventDisableTiming)); + // One more buffer than there are stages: every stage can hold a pass while the host + // stages the next. NINFER_QWEN4_STAGING_RING=N sets the ring (1: the host waits for + // the previous pass on every rank, as before the ring), for A/B measurements. + std::size_t ring = device.size() + 1; + if (const char* value = std::getenv("NINFER_QWEN4_STAGING_RING")) { + ring = std::clamp(std::strtoull(value, nullptr, 10), 1, 16); + } + for (std::size_t s = 0; s < ring; ++s) { + rank.staging.push_back(std::make_unique( + round_up(t * 28) + round_up(std::uint64_t(config.ngram_heads()) * t * row_b))); + cudaEvent_t event = nullptr; + CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + rank.staged.push_back(event); + } CUDA_CHECK(cudaEventCreateWithFlags(&rank.done, cudaEventDisableTiming)); CUDA_CHECK(cudaEventCreateWithFlags(&rank.copied, cudaEventDisableTiming)); CUDA_CHECK(cudaEventCreateWithFlags(&rank.routes, cudaEventDisableTiming)); @@ -1413,7 +1470,7 @@ struct Executor::Impl { .value_heads = static_cast(config.linear_num_value_heads), .value_head_dim = static_cast(config.linear_value_head_dim), .key_head_dim = static_cast(config.linear_key_head_dim), - .slot_count = static_cast(options.sequences), + .slot_count = static_cast(options.sequences + (pipeline ? 1 : 0)), }; LayoutBuilder builder; const auto layout = plan_linear_attention_state_pool(builder, spec); @@ -1457,6 +1514,49 @@ struct Executor::Impl { sequences.push_back(std::move(sequence)); reset(s); } + if (!pipeline) { return; } + // The shadow of the sequence's states on the front ranks: its states as if the window in + // flight were kept whole, which a guessed window reads. Its KV and pooled keys are the + // sequence's own (a guessed window writes positions after the verified one). + ahead_layers.resize(layers.size()); + const auto shadow = static_cast(options.sequences); + for (std::size_t i = 0; i < layers.size(); ++i) { + const LayerPlan& plan = layers[i]; + if (plan.rank == model.head_rank()) { continue; } + RankBinding bind(device, plan.rank); + LayerState& state = ahead_layers[i]; + if (plan.gdn) { + state.ssm = ranks[plan.rank].gdn->recurrent_slot(gdn_local[i], shadow); + state.conv = ranks[plan.rank].gdn->conv_slot(gdn_local[i], shadow); + } + if (plan.qsa) { + state.tail = DeviceBuffer(sequences.front().layers[i].tail.bytes); + memory.ranks[plan.rank].state_bytes += state.tail.bytes; + } + if (plan.ple) { + state.history = DeviceBuffer(sequences.front().layers[i].history.bytes); + memory.ranks[plan.rank].state_bytes += state.history.bytes; + } + } + } + + // Pipeline windows (NINFER_QWEN4_PIPELINE_WINDOWS=1, with rank threads): while the head rank + // verifies a window, the front ranks run the next one, guessed by the MTP block on the + // assumption that this window is kept whole with the block's bonus; a wrong guess is dropped. + // One sequence, MTP drafts, device experts and CUDA Graphs. + bool pipeline = false; + std::vector ahead_layers; // the shadow states, by layer (front ranks) + [[nodiscard]] bool pipeline_requested() const { + const auto on = [](const char* name) { + const char* value = std::getenv(name); + return value != nullptr && value[0] == '1'; + }; + const bool split = std::any_of(segments.begin(), segments.end(), [&](const Segment& s) { + return s.rank != model.head_rank(); + }); + return on("NINFER_QWEN4_PIPELINE_WINDOWS") && on("NINFER_QWEN4_RANK_THREADS") && split && + verify_width > 0 && options.sequences == 1 && options.cuda_graphs && + model.options().experts == ExpertResidency::Device; } // What a speculative round needs beyond the sequences' state: on every rank the replay records @@ -1473,43 +1573,60 @@ struct Executor::Impl { for (std::size_t i = 0; i < layers.size(); ++i) { if (layers[i].qsa) { qsa_local[i] = attention[layers[i].rank]++; } } + const std::size_t sets = pipeline ? 2 : 1; for (RankState& rank : ranks) { RankBinding bind(device, rank.rank); std::uint64_t bytes = 0; + const bool ple_here = + std::any_of(layers.begin(), layers.end(), + [&](const LayerPlan& p) { return p.ple && p.rank == rank.rank; }); + for (std::size_t set = 0; set < sets; ++set) { + if (rank.gdn) { + const auto& pool = rank.gdn->spec(); + const GdnReplayRecordSpec spec{ + .layers = static_cast(pool.layers), + .record_capacity = static_cast(options.sequences), + .width = static_cast(verify_width), + .conv_channels = pool.conv_channels, + .qk_heads = static_cast(config.linear_num_key_heads), + .value_heads = pool.value_heads, + .key_dim = pool.key_head_dim, + .value_dim = pool.value_head_dim, + }; + LayoutBuilder builder; + const auto layout = plan_gdn_replay_records(builder, spec); + rank.record_backing[set] = + DeviceBuffer(builder.finish(256, "GDN replay records")); + rank.records[set].emplace(DeviceSpan{rank.record_backing[set].p, + rank.record_backing[set].bytes}, + layout); + rank.fold[set].emplace(*rank.records[set], rank.gdn->all_layers_view()); + bytes += rank.record_backing[set].bytes; + } + if (attention[rank.rank] > 0) { + rank.index_records[set] = + DeviceBuffer(attention[rank.rank] * index_rows * columns * 2); + bytes += rank.index_records[set].bytes; + } + if (ple_here) { + rank.ple_records[set] = DeviceBuffer(stack_width * columns * 4); + bytes += rank.ple_records[set].bytes; + } + if (pipeline) { + rank.window_positions[set] = DeviceBuffer(columns * 4); + bytes += rank.window_positions[set].bytes; + } + } if (rank.gdn) { - const auto& pool = rank.gdn->spec(); - const GdnReplayRecordSpec spec{ - .layers = static_cast(pool.layers), - .record_capacity = static_cast(options.sequences), - .width = static_cast(verify_width), - .conv_channels = pool.conv_channels, - .qk_heads = static_cast(config.linear_num_key_heads), - .value_heads = pool.value_heads, - .key_dim = pool.key_head_dim, - .value_dim = pool.value_head_dim, - }; - LayoutBuilder builder; - const auto layout = plan_gdn_replay_records(builder, spec); - rank.record_backing = DeviceBuffer(builder.finish(256, "GDN replay records")); - rank.records.emplace(DeviceSpan{rank.record_backing.p, rank.record_backing.bytes}, - layout); - rank.fold.emplace(*rank.records, rank.gdn->all_layers_view()); + const auto& pool = rank.gdn->spec(); rank.conv_scratch = DeviceBuffer(std::uint64_t(pool.conv_channels) * (config.linear_conv_kernel_dim - 1) * 2); - bytes += rank.record_backing.bytes + rank.conv_scratch.bytes; + bytes += rank.conv_scratch.bytes; } if (attention[rank.rank] > 0) { - rank.index_records = DeviceBuffer(attention[rank.rank] * index_rows * columns * 2); - rank.tail_scratch = DeviceBuffer(std::uint64_t(config.indexer_head_dim) * - (config.indexer_compress_ratio - 1) * 4); - bytes += rank.index_records.bytes + rank.tail_scratch.bytes; - } - const bool ple_here = - std::any_of(layers.begin(), layers.end(), - [&](const LayerPlan& p) { return p.ple && p.rank == rank.rank; }); - if (ple_here) { - rank.ple_records = DeviceBuffer(stack_width * columns * 4); - bytes += rank.ple_records.bytes; + rank.tail_scratch = DeviceBuffer(std::uint64_t(config.indexer_head_dim) * + (config.indexer_compress_ratio - 1) * 4); + bytes += rank.tail_scratch.bytes; } memory.ranks[rank.rank].workspace_bytes += bytes; } @@ -1523,7 +1640,8 @@ struct Executor::Impl { if (options.draft_min_p > 0) { mtp_logprobs = DeviceBuffer(std::uint64_t(options.draft_tokens) * options.sequences * 4); } - if (table && options.ngram_read.residency != NgramResidency::Ram) { + // Pipeline windows read a verification's rows on the host thread as they stage it. + if (table && options.ngram_read.residency != NgramResidency::Ram && !pipeline) { draft_prefetch = std::make_unique( device.rank(head), *table, ngram, config.eos_token_id, config.vocab_size, options.sequences, options.draft_tokens); @@ -1537,6 +1655,62 @@ struct Executor::Impl { } memory.ranks[head].workspace_bytes += mtp_ids.bytes + mtp_positions.bytes + mtp_logprobs.bytes + ngram_prefetch_tokens.bytes + sequences.size() * sequences.front().mtp.tail.bytes; + mirror_embedding_for_head(); + } + + // A split model's MTP cells take their token embeddings from rank 0's table; without peer + // access each draft step's lookup is a round trip head -> rank 0 -> head that the host waits + // on twice. The head instead reads the rows from a mapped copy of the table's payload in pinned + // host memory (the same bytes through the same lookup, a row of the bus per cell), which also + // keeps the draft chain on one device, where it can be captured. + std::unique_ptr head_embedding_host; + Weight head_embedding; // payload == nullptr: look the rows up on rank 0 + void mirror_embedding_for_head() { + const std::size_t head = model.head_rank(); + const Weight& table = embedding_table; + if (head == 0 || table.payload == nullptr || table.payload_bytes == 0) { return; } + if (const char* value = std::getenv("NINFER_QWEN4_HEAD_EMBEDDING"); value && value[0] == '0') { + return; + } + const auto* begin = static_cast(table.payload); + const auto* end = begin + table.payload_bytes; + const auto inside = [&](const void* p) { + return p == nullptr || (static_cast(p) >= begin && + static_cast(p) < end); + }; + if (!inside(table.qdata) || !inside(table.qhigh) || !inside(table.scales) || + !inside(table.weight_divisors) || !inside(table.input_columns)) { + return; + } + auto host = std::make_unique(table.payload_bytes); + { + RankBinding bind(device, 0); + CUDA_CHECK(cudaMemcpy(host->data(), table.payload, table.payload_bytes, + cudaMemcpyDeviceToHost)); + } + void* mapped = nullptr; + { + RankBinding bind(device, head); + if (cudaHostGetDevicePointer(&mapped, host->data(), 0) != cudaSuccess) { + (void)cudaGetLastError(); + return; + } + } + const auto rebase = [&](const void* p) -> const void* { + if (p == nullptr) { return nullptr; } + return static_cast(mapped) + (static_cast(p) - begin); + }; + head_embedding = table; + head_embedding.payload = mapped; + head_embedding.qdata = rebase(table.qdata); + head_embedding.qhigh = rebase(table.qhigh); + head_embedding.scales = rebase(table.scales); + head_embedding.weight_divisors = rebase(table.weight_divisors); + head_embedding.input_columns = + static_cast(rebase(table.input_columns)); + head_embedding_host = std::move(host); + std::fprintf(stderr, "qwen4_exp: MTP embeddings read from a %.1f MiB host copy of the table\n", + double(table.payload_bytes) / (1024.0 * 1024.0)); } // The layers whose experts the route records, the expert cache and the expert stream serve: @@ -2009,6 +2183,8 @@ struct Executor::Impl { } void reset(std::uint32_t s) { + reset_pipeline(); + drain(); auto& sequence = sequences.at(s); if (std::find(verified.begin(), verified.end(), &sequence) != verified.end()) { verified.clear(); } sequence.position = 0; @@ -2054,6 +2230,526 @@ struct Executor::Impl { if (mtp) { forget(sequence.mtp, model.head_rank()); } } + // A stage boundary's host bounce buffers (rank threads): the source's thread copies its stack + // into one on its stream, the target's thread uploads it on its own, each copy ordered by + // events. A buffer is written again only once the target has read it. + struct CrossRing { + std::vector> buffers; + std::vector written; // on the source's stream + std::vector read; // on the target's stream + }; + std::vector crossings; // by target rank + static constexpr std::size_t kCrossRing = 2; + + // ---- Rank threads (NINFER_QWEN4_RANK_THREADS=1) ---- + // + // Between devices without peer access, a stream waiting on another device's event or a staged + // peer copy blocks the calling host thread until the other device gets there (Windows/WDDM). + // One host thread queuing every rank's work therefore walks a pass through the pipeline at + // device speed, and the next pass starts only once the last rank has it: the stages take + // turns. With rank threads the main thread prepares a pass's inputs and posts it; each rank's + // thread uploads them, waits on the host for its stack from the rank before (the only thread + // that blocks), runs its segments and hands the stack on through pinned host memory. Prompt + // chunks then fill every stage. The head rank's thread runs the MTP catch-up, whose cells take + // the pass's embeddings from host memory (rank 0 copied them there) instead of a round trip to + // rank 0, which would order rank 0's next chunk after the head's work. + // + // Prompt chunks and plain decode passes of device-resident experts without media take this + // route; anything else, and every call that reads results or state, first waits until the + // threads have queued all posted passes (drain()). + struct RankTask { + std::uint64_t id = 0; + std::size_t slot = 0; // staging ring entry + std::int32_t t = 0; // 0: no pass, only `work` + std::uint32_t logit_rows = 0; + std::size_t row_bytes = 0; + std::vector parts; + SequenceState* graph_sequence = nullptr; // replay its decode graphs + bool catch_up = false; // run the MTP catch-up on the head rank + std::vector committed; // its tokens point into `tokens` + std::vector tokens; + Executor::AfterHead after_head; // on the head rank's thread, after the head + // Pipeline windows: the pass's kind, the segments it runs, per-segment graphs to replay, + // whether its stack crosses to the head rank through the guessed window's buffer (a + // guessed window) or the head rank's segments take it from there (its verification), and + // work each rank's thread does first (a commit, the shadow states, the guess chain). + PassMode mode; + std::size_t first_segment = 0, end_segment = ~std::size_t{0}; + std::vector* graphs = nullptr; + bool to_ahead = false, from_ahead = false; + std::uint64_t ahead_number = 0; + std::function work; + }; + struct RankWorker { + std::thread thread; + std::deque> queue; + }; + bool rank_threads = false; + std::vector workers; + std::mutex worker_mutex; + std::condition_variable worker_cv; + std::vector done_upto; // by rank, task ids + // Stage boundaries, by target rank: stacks sent and received so far. Both ends count the + // passes that cross the boundary in the same order, so a pass's crossing number picks its + // buffer without the ids of passes that do not cross there. + std::vector cross_sent, cross_received; + std::vector send_count, receive_count; // each touched by one thread only + std::vector slot_owner; // the pass task that last staged a slot + std::uint64_t staged_passes = 0; + std::uint64_t posted = 0; + bool workers_stop = false; + std::exception_ptr worker_error; + std::vector> embedded_host; // rank 0's, by staging entry + DeviceBuffer embed_scratch; // rank 0: graph passes' embeddings for the MTP catch-up + + void start_rank_threads() { + const char* value = std::getenv("NINFER_QWEN4_RANK_THREADS"); + if (value == nullptr || value[0] != '1' || ranks.size() < 2) { return; } + if (model.options().experts != ExpertResidency::Device || hybrid || misses || stream) { + return; + } + // Without CUDA Graphs the threads' eager decode steps hit an illegal address during the + // warm-up (Windows, 4 stages, unresolved), so that mode keeps the host-ordered passes. + if (!options.cuda_graphs) { return; } + rank_threads = true; + const std::size_t ring = ranks.front().staging.size(); + const std::size_t capacity = + std::size_t(config.hc_count) * config.hidden_size * options.prefill_chunk * 4; + crossings.resize(ranks.size()); + for (std::size_t to = 1; to < ranks.size(); ++to) { + CrossRing& cross = crossings[to]; + for (std::size_t i = 0; i < kCrossRing; ++i) { + cross.buffers.push_back(std::make_unique(capacity)); + cudaEvent_t written = nullptr, read = nullptr; + { + RankBinding bind(device, to - 1); + CUDA_CHECK(cudaEventCreateWithFlags(&written, cudaEventDisableTiming)); + } + { + RankBinding bind(device, to); + CUDA_CHECK(cudaEventCreateWithFlags(&read, cudaEventDisableTiming)); + } + cross.written.push_back(written); + cross.read.push_back(read); + } + } + for (std::size_t s = 0; s < ring; ++s) { + embedded_host.push_back(std::make_unique( + std::size_t(config.hidden_size) * options.prefill_chunk * 2)); + } + if (mtp && options.cuda_graphs) { + RankBinding bind(device, 0); + embed_scratch = DeviceBuffer(std::size_t(config.hidden_size) * kVectorTokens * 2); + } + done_upto.assign(ranks.size(), 0); + cross_sent.assign(ranks.size(), 0); + cross_received.assign(ranks.size(), 0); + send_count.assign(ranks.size(), 0); + receive_count.assign(ranks.size(), 0); + slot_owner.assign(ring, 0); + if (pipeline) { start_pipeline(); } + workers = std::vector(ranks.size()); + for (std::size_t r = 0; r < ranks.size(); ++r) { + workers[r].thread = std::thread([this, r] { worker_loop(r); }); + } + } + + void stop_rank_threads() noexcept { + if (!rank_threads) { return; } + { + std::lock_guard lock(worker_mutex); + workers_stop = true; + } + worker_cv.notify_all(); + for (RankWorker& worker : workers) { + if (worker.thread.joinable()) { worker.thread.join(); } + } + } + + // Waits until `counter[rank]` reaches `value`; a failed thread ends every wait with its error. + void wait_for(const std::vector& counter, std::size_t rank, std::uint64_t value) { + std::unique_lock lock(worker_mutex); + worker_cv.wait(lock, [&] { return counter[rank] >= value || worker_error != nullptr; }); + if (worker_error) { std::rethrow_exception(worker_error); } + } + + // Waits until the rank threads have queued every posted pass on their devices. + void drain() { + if (!rank_threads) { return; } + for (std::size_t r = 0; r < ranks.size(); ++r) { wait_for(done_upto, r, posted); } + } + // The same for the head rank's thread alone: with pipeline windows the front ranks may still + // be queuing a guessed window while the host reads the verdict or drafts on the head rank. + void drain_head() { + if (!rank_threads) { return; } + if (!pipeline) { + drain(); + return; + } + wait_for(done_upto, model.head_rank(), posted); + } + void drain_quietly() noexcept { + try { + drain(); + } catch (...) {} + } + void drain_head_quietly() noexcept { + try { + drain_head(); + } catch (...) {} + } + + void worker_loop(std::size_t r) { + for (;;) { + std::shared_ptr task; + { + std::unique_lock lock(worker_mutex); + worker_cv.wait(lock, [&] { return workers_stop || !workers[r].queue.empty(); }); + if (workers[r].queue.empty()) { return; } + task = workers[r].queue.front(); + workers[r].queue.pop_front(); + } + try { + bool failed = false; + { + std::lock_guard lock(worker_mutex); + failed = worker_error != nullptr; + } + if (!failed) { run_rank_task(r, *task); } + } catch (...) { + std::lock_guard lock(worker_mutex); + if (!worker_error) { worker_error = std::current_exception(); } + } + { + std::lock_guard lock(worker_mutex); + done_upto[r] = task->id; + } + worker_cv.notify_all(); + } + } + + [[nodiscard]] std::size_t stack_bytes(std::int32_t t) const { + return std::size_t(config.hc_count) * config.hidden_size * std::size_t(t) * 4; + } + + // The source half of a stage boundary, on the source's thread: the stack goes to the + // boundary's host buffer for this task once the target has read what it held before. + void send_stack(std::size_t from, std::size_t to, const RankTask& task) { + CrossRing& ring = crossings[to]; + const std::uint64_t n = ++send_count[to]; + const std::size_t s = n % kCrossRing; + if (n > kCrossRing) { + wait_for(cross_received, to, n - kCrossRing); + CUDA_CHECK(cudaEventSynchronize(ring.read[s])); + } + RankState& source = ranks[from]; + CUDA_CHECK(cudaMemcpyAsync(ring.buffers[s]->data(), source.stack, stack_bytes(task.t), + cudaMemcpyDeviceToHost, source.stream)); + CUDA_CHECK(cudaEventRecord(ring.written[s], source.stream)); + { + std::lock_guard lock(worker_mutex); + cross_sent[to] = n; + } + worker_cv.notify_all(); + } + + // The target half, on the target's thread: waits on the host for the source's copy, then + // uploads the stack on its own stream. + void receive_stack(std::size_t to, const RankTask& task) { + CrossRing& ring = crossings[to]; + const std::uint64_t n = ++receive_count[to]; + const std::size_t s = n % kCrossRing; + wait_for(cross_sent, to, n); + CUDA_CHECK(cudaEventSynchronize(ring.written[s])); + RankState& target = ranks[to]; + CUDA_CHECK(cudaMemcpyAsync(target.stack, ring.buffers[s]->data(), stack_bytes(task.t), + cudaMemcpyHostToDevice, target.stream)); + CUDA_CHECK(cudaEventRecord(ring.read[s], target.stream)); + { + std::lock_guard lock(worker_mutex); + cross_received[to] = n; + } + worker_cv.notify_all(); + } + + // The guessed window's crossing to the head rank (pipeline windows): one buffer, written by + // the last front rank once the head rank has read, or the host has dropped, the guess before; + // read by the head rank only when the guess holds. + std::unique_ptr ahead_buffer; + cudaEvent_t ahead_written = nullptr, ahead_read = nullptr; + std::uint64_t ahead_sent = 0, ahead_resolved = 0; // guessed windows (worker_mutex) + std::uint64_t ahead_count = 0; // the host's count of guessed windows + + void send_ahead(std::size_t from, const RankTask& task, std::uint64_t number) { + { + std::unique_lock lock(worker_mutex); + worker_cv.wait(lock, [&] { return ahead_resolved + 1 >= number || worker_error; }); + if (worker_error) { std::rethrow_exception(worker_error); } + } + RankState& source = ranks[from]; + CUDA_CHECK(cudaMemcpyAsync(ahead_buffer->data(), source.stack, stack_bytes(task.t), + cudaMemcpyDeviceToHost, source.stream)); + CUDA_CHECK(cudaEventRecord(ahead_written, source.stream)); + { + std::lock_guard lock(worker_mutex); + ahead_sent = number; + } + worker_cv.notify_all(); + } + + void receive_ahead(std::size_t to, const RankTask& task, std::uint64_t number) { + { + std::unique_lock lock(worker_mutex); + worker_cv.wait(lock, [&] { return ahead_sent >= number || worker_error; }); + if (worker_error) { std::rethrow_exception(worker_error); } + } + CUDA_CHECK(cudaEventSynchronize(ahead_written)); + RankState& target = ranks[to]; + CUDA_CHECK(cudaMemcpyAsync(target.stack, ahead_buffer->data(), stack_bytes(task.t), + cudaMemcpyHostToDevice, target.stream)); + // The host waits for this upload before the next guess overwrites the buffer. + CUDA_CHECK(cudaEventRecord(ahead_read, target.stream)); + CUDA_CHECK(cudaEventSynchronize(ahead_read)); + { + std::lock_guard lock(worker_mutex); + ahead_resolved = std::max(ahead_resolved, number); + } + worker_cv.notify_all(); + } + + void run_rank_task(std::size_t r, const RankTask& task) { + RankBinding bind(device, r); + RankState& rank = ranks[r]; + pass_mode = task.mode; + struct ModeReset { + ~ModeReset() { pass_mode = {}; } + } mode_reset; + if (task.work) { task.work(r); } + if (task.t == 0) { return; } + const std::size_t end = std::min(task.end_segment, segments.size()); + bool runs = false; + for (std::size_t i = task.first_segment; i < end; ++i) { runs = runs || segments[i].rank == r; } + if (!runs) { return; } + const std::size_t c = options.prefill_chunk; + auto* base = static_cast(rank.staging[task.slot]->data()); + const std::size_t bytes = std::size_t(task.t) * 4; + CUDA_CHECK(cudaMemcpyAsync(rank.ids, base, bytes, cudaMemcpyHostToDevice, rank.stream)); + CUDA_CHECK(cudaMemcpyAsync(rank.positions, base + 4 * c, bytes, cudaMemcpyHostToDevice, + rank.stream)); + CUDA_CHECK(cudaMemcpyAsync(rank.rope, base + 8 * c, bytes, cudaMemcpyHostToDevice, + rank.stream)); + if (task.mode.verify && pipeline) { + // The window's positions for its commit, which later passes' staging overwrites. + CUDA_CHECK(cudaMemcpyAsync(rank.window_positions[task.mode.set].p, rank.positions, bytes, + cudaMemcpyDeviceToDevice, rank.stream)); + } + for (std::size_t i = task.first_segment; i < end; ++i) { + const Segment& segment = segments[i]; + if (segment.rank != r) { continue; } + if (i > 0 && segments[i - 1].rank != r) { + if (task.from_ahead && i == task.first_segment) { + receive_ahead(r, task, task.ahead_number); + } else { + receive_stack(r, task); + } + } + if (segment.rows && table) { + CUDA_CHECK(cudaMemcpyAsync(rank.rows, staged_rows(rank, task.slot, task.row_bytes).data(), + task.row_bytes, cudaMemcpyHostToDevice, rank.stream)); + } + if (task.graphs != nullptr) { + (*task.graphs)[i].launch(rank.stream); + } else if (task.graph_sequence != nullptr) { + if (segment.embed && task.catch_up) { + // The graph's embedding plane is rewritten by its layers; the MTP catch-up + // takes the tokens' embeddings from a lookup of their own. + const auto h = static_cast(config.hidden_size); + Tensor rows(embed_scratch.p, DType::BF16, {h, task.t}); + ops::embedding(Tensor(rank.ids, DType::I32, {task.t}), embedding_table, rows, + rank.stream); + CUDA_CHECK(cudaMemcpyAsync(embedded_host[task.slot]->data(), embed_scratch.p, + std::size_t(h) * task.t * 2, cudaMemcpyDeviceToHost, + rank.stream)); + } + task.graph_sequence->decode[i].launch(rank.stream); + } else { + void* embedded = segment.embed && task.catch_up ? embedded_host[task.slot]->data() + : nullptr; + run_segment(task.parts, segment, task.t, task.logit_rows, embedded); + } + if (segment.head && task.catch_up) { + mtp_catch_up(task.committed, + static_cast(embedded_host[task.slot]->data())); + } + if (segment.head && task.after_head) { + task.after_head(rank.stream, head_logits(std::size_t(task.logit_rows))); + } + if (i + 1 < segments.size() && segments[i + 1].rank != r) { + if (task.to_ahead && i + 1 == end) { + send_ahead(r, task, task.ahead_number); + } else if (i + 1 < end) { + send_stack(r, segments[i + 1].rank, task); + } + } + } + // The staging entry (and rank 0's embeddings in it) may be rewritten once this is done. + CUDA_CHECK(cudaEventRecord(rank.staged[task.slot], rank.stream)); + } + + // Whether a pass over `parts` can go to the rank threads. + [[nodiscard]] bool threadable(std::span parts, std::int32_t t) const { + if (!rank_threads || verifying) { return false; } + for (const Part& part : parts) { + if (!part.sequence->media_rope.empty() && + part.sequence->position < part.sequence->media_rope.size() / 3) { + return false; + } + } + (void)t; + return true; + } + + // Prepares a pass's inputs on the host and posts it to the rank threads; the sequences' + // positions advance at once (the threads read positions from the staged planes). + void threaded_pass(std::span parts, std::span tokens, + std::uint32_t logit_rows, Executor::AfterHead after_head = {}) { + const auto t = static_cast(tokens.size()); + for (const Part& part : parts) { + if (part.sequence->position + std::uint32_t(part.count) > options.max_context) { + throw std::invalid_argument("qwen4_exp forward: the sequence exceeds max_context"); + } + } + std::vector contexts; + for (const Part& part : parts) { contexts.push_back(&part.sequence->context); } + auto task = stage_task(parts, tokens, logit_rows, contexts, 0, segments.size()); + task->after_head = std::move(after_head); + // Decode graphs: the first single-token pass runs eagerly, the second captures them (with + // the threads drained, since a capture forbids other threads' launches). + SequenceState& first = *parts.front().sequence; + if (options.cuda_graphs && t == 1 && parts.size() == 1) { + if (first.decode.empty() && first.decode_steps++ > 0) { + drain(); + capture_decode(first); + } + if (!first.decode.empty()) { task->graph_sequence = &first; } + } + if (mtp) { + task->catch_up = true; + task->tokens.assign(tokens.begin(), tokens.end()); + for (const Part& part : parts) { + task->committed.push_back( + {part.sequence, part.column, part.count, part.sequence->position, + std::span(task->tokens) + .subspan(std::size_t(part.column), std::size_t(part.count))}); + } + } + post(task); + for (const Part& part : parts) { + part.sequence->position += static_cast(part.count); + } + } + + // A task id for the next post. + std::shared_ptr new_task() { + auto task = std::make_shared(); + task->id = posted + 1; + return task; + } + + void post(const std::shared_ptr& task) { + { + std::lock_guard lock(worker_mutex); + if (worker_error) { std::rethrow_exception(worker_error); } + if (task->id != posted + 1) { throw std::logic_error("qwen4_exp: rank tasks out of order"); } + for (RankWorker& worker : workers) { worker.queue.push_back(task); } + posted = task->id; + } + worker_cv.notify_all(); + } + + // The host half of a pass to the rank threads: a staging entry (once every rank has queued the + // pass that staged it last) with the tokens, their positions from each part's sequence position + // plus `offset`, their RoPE positions and, when a segment in [first, end) reads them, their + // n-gram rows, hashed from `contexts` (one per part, advanced over the part's tokens) and read + // on this thread. + std::shared_ptr stage_task(std::span parts, + std::span tokens, + std::uint32_t logit_rows, + std::span contexts, + std::uint32_t offset, std::size_t end, + std::size_t first_segment = 0) { + const std::size_t ring = ranks.front().staging.size(); + auto task = new_task(); + const std::size_t slot = ++staged_passes % ring; + if (slot_owner[slot] != 0) { + for (std::size_t r = 0; r < ranks.size(); ++r) { wait_for(done_upto, r, slot_owner[slot]); } + for (auto& rank : ranks) { CUDA_CHECK(cudaEventSynchronize(rank.staged[slot])); } + } + slot_owner[slot] = task->id; + task->slot = slot; + task->t = static_cast(tokens.size()); + task->logit_rows = logit_rows; + task->first_segment = first_segment; + task->end_segment = end; + task->parts.assign(parts.begin(), parts.end()); + std::vector positions(tokens.size()), rope(tokens.size()); + const std::size_t heads = config.ngram_heads(); + bool rows_read = false; + for (std::size_t i = first_segment; i < end; ++i) { rows_read = rows_read || segments[i].rows; } + rows_read = rows_read && table; + std::size_t row_b = 0; + if (rows_read) { + row_ids.resize(tokens.size() * heads); + row_b = ops::ngram_row_bytes(options.ngram->format); + } + for (std::size_t k = 0; k < parts.size(); ++k) { + const Part& part = parts[k]; + const SequenceState& sequence = *part.sequence; + const auto start = static_cast(sequence.position + offset); + const auto at = positions.begin() + part.column; + std::iota(at, at + part.count, start); + for (std::int32_t i = 0; i < part.count; ++i) { + rope[std::size_t(part.column + i)] = start + i + sequence.rope_delta; + } + const auto own = tokens.subspan(std::size_t(part.column), std::size_t(part.count)); + if (rows_read) { + ngram_row_ids(ngram, own, config.eos_token_id, config.vocab_size, *contexts[k], + std::span(row_ids).subspan(std::size_t(part.column) * heads, + std::size_t(part.count) * heads)); + } else if (table) { + std::vector scratch(own.size() * heads); + ngram_row_ids(ngram, own, config.eos_token_id, config.vocab_size, *contexts[k], + scratch); + } + } + const std::size_t c = options.prefill_chunk; + if (rows_read) { + task->row_bytes = row_ids.size() * row_b; + // Read on this thread: the threads only upload the rows. + table->submit(row_ids, staged_rows(ranks[rows_rank], slot, task->row_bytes)); + table->wait(); + const auto rows = staged_rows(ranks[rows_rank], slot, task->row_bytes); + for (auto& rank : ranks) { + if (&rank == &ranks[rows_rank]) { continue; } + const bool consumes = std::any_of(segments.begin(), segments.end(), [&](const Segment& s) { + return s.rows && s.rank == rank.rank; + }); + if (consumes) { + std::memcpy(staged_rows(rank, slot, task->row_bytes).data(), rows.data(), + task->row_bytes); + } + } + } + for (auto& rank : ranks) { + auto* base = static_cast(rank.staging[slot]->data()); + std::memcpy(base, tokens.data(), tokens.size() * 4); + std::memcpy(base + 4 * c, positions.data(), positions.size() * 4); + std::memcpy(base + 8 * c, rope.data(), rope.size() * 4); + } + return task; + } + + // Moves the stack's first `tokens` columns from one rank to another, ordered after the // source's work. The source's next pass (the next prompt chunk, queued before the target gets // to this copy) rewrites the source's stack, so it waits until the copy has read it. @@ -2168,11 +2864,13 @@ struct Executor::Impl { table->wait(); } catch (...) {} } + staging_slot = (staging_slot + 1) % ranks.front().staging.size(); for (auto& rank : ranks) { RankBinding bind(device, rank.rank); - // The staging buffer is rewritten only once the previous upload has left it. - CUDA_CHECK(cudaEventSynchronize(rank.staged)); - auto* base = static_cast(rank.staging->data()); + // The staging buffer is rewritten only once its last upload, ring-size passes ago, has + // left it. + CUDA_CHECK(cudaEventSynchronize(rank.staged[staging_slot])); + auto* base = static_cast(rank.staging[staging_slot]->data()); const auto upload = [&](void* target, std::size_t at, const std::vector& values) { std::memcpy(base + at, values.data(), values.size() * 4); @@ -2195,13 +2893,16 @@ struct Executor::Impl { const bool ple_here = std::any_of(layers.begin(), layers.end(), [&](const LayerPlan& p) { return p.ple && p.rank == rank.rank; }); - if (!ple_here || !read_rows) { CUDA_CHECK(cudaEventRecord(rank.staged, rank.stream)); } + if (!ple_here || !read_rows) { CUDA_CHECK(cudaEventRecord(rank.staged[staging_slot], rank.stream)); } } } // The n-gram rows' place in a rank's staging buffer. std::span staged_rows(RankState& rank, std::size_t bytes) const { - auto* base = static_cast(rank.staging->data()); + return staged_rows(rank, staging_slot, bytes); + } + std::span staged_rows(RankState& rank, std::size_t slot, std::size_t bytes) const { + auto* base = static_cast(rank.staging[slot]->data()); return {reinterpret_cast( base + round_up(28 * std::uint64_t(options.prefill_chunk))), bytes}; @@ -2222,7 +2923,7 @@ struct Executor::Impl { } CUDA_CHECK( cudaMemcpyAsync(rank.rows, rows.data(), bytes, cudaMemcpyHostToDevice, rank.stream)); - CUDA_CHECK(cudaEventRecord(rank.staged, rank.stream)); + CUDA_CHECK(cudaEventRecord(rank.staged[staging_slot], rank.stream)); } // Waits for the rows being read, unless already waited for, between a pair of stall events on @@ -2280,8 +2981,8 @@ struct Executor::Impl { } // A verification's record row of the part: its position among the verified sequences. - [[nodiscard]] std::int32_t record_row(const Part& part) const { - return part.column / static_cast(active_verify_width); + [[nodiscard]] static std::int32_t record_row(const Part& part) { + return part.column / static_cast(pass_mode.width); } void run_gdn(const LayerPlan& plan, std::size_t index, LayerState& state, RankState& rank, @@ -2302,14 +3003,18 @@ struct Executor::Impl { // A verification records its raw projection, the transitions and the gates for the fold // that commits a prefix, and leaves the sequence's states as they were. std::optional record; - if (verifying) { + if (pass_mode.verify) { const auto row = record_row(part); - const auto all = rank.records->layer(static_cast(gdn_local[index]), - static_cast(verified.size()), - static_cast(active_verify_width)); + const auto all = rank.records[pass_mode.set]->layer( + static_cast(gdn_local[index]), pass_mode.rows, + static_cast(pass_mode.width)); record = GdnReplayRecordLayer{all.conv.slice(2, row, 1), all.key.slice(3, row, 1), all.value.slice(3, row, 1), all.gate.slice(3, row, 1)}; } + // A guessed window reads the shadow slot. + const LayerState& source = pass_mode.ahead ? ahead_layers[index] : state; + const std::uint32_t source_slot = + pass_mode.ahead ? options.sequences : part.sequence->slot; Tensor qkv = record ? Tensor(record->conv.data, DType::BF16, {g.qkv.rows, t}) : Tensor(rank.a, DType::BF16, {g.qkv.rows, t}); Tensor z(rank.b, DType::BF16, {vd, t}); @@ -2326,7 +3031,7 @@ struct Executor::Impl { Tensor k(conv_out + round_up(std::uint64_t(kd) * t * 2), DType::BF16, {kd, t}); Tensor v(rank.e, DType::BF16, {vd, t}); const auto window = static_cast(config.linear_conv_kernel_dim - 1); - Tensor conv_state(state.conv.data, DType::BF16, {g.qkv.rows, window}); + Tensor conv_state(source.conv.data, DType::BF16, {g.qkv.rows, window}); Tensor conv_next = record ? Tensor(rank.conv_scratch.p, DType::BF16, {g.qkv.rows, window}) : conv_state; ops::causal_conv1d_silu_split(qkv, g.convolution, conv_state, conv_next, q, k, v, s); @@ -2340,7 +3045,7 @@ struct Executor::Impl { const float scale = 1.0F / std::sqrt(float(dh)); if (record) { const Tensor states = rank.gdn->layer_view(gdn_local[index]).recurrent; - const Tensor slot(static_cast(rank.slot_ids.p) + part.sequence->slot, + const Tensor slot(static_cast(rank.slot_ids.p) + source_slot, DType::I32, {1}); Tensor o4 = o.view({dh, nv, t, 1}); ops::gated_delta_net_replay_record(q3.view({dh, nk, t, 1}), k3.view({dh, nk, t, 1}), @@ -2381,13 +3086,14 @@ struct Executor::Impl { .index = nullptr, .tail = state.tail.p, .first = part.sequence->position}; - if (verifying) { + if (pass_mode.verify) { const std::uint64_t rows = std::uint64_t(config.indexer_n_heads + 1) * config.indexer_head_dim; const std::uint64_t columns = std::uint64_t(verify_width) * options.sequences; - in.index = static_cast(rank.index_records.p) + + in.index = static_cast(rank.index_records[pass_mode.set].p) + (qsa_local[index] * columns + std::uint64_t(part.column)) * rows * 2; - CUDA_CHECK(cudaMemcpyAsync(rank.tail_scratch.p, state.tail.p, state.tail.bytes, + const DeviceBuffer& tail = pass_mode.ahead ? ahead_layers[index].tail : state.tail; + CUDA_CHECK(cudaMemcpyAsync(rank.tail_scratch.p, tail.p, tail.bytes, cudaMemcpyDeviceToDevice, rank.stream)); in.tail = rank.tail_scratch.p; } @@ -2498,7 +3204,8 @@ struct Executor::Impl { project(a.output, flat, y, ws, s); } - void run_ple(const LayerPlan& plan, LayerState& state, RankState& rank, const Part& part) { + void run_ple(const LayerPlan& plan, std::size_t index, LayerState& state, RankState& rank, + const Part& part) { const PlePlan& p = *plan.ple; WorkspaceArena& ws = *rank.workspace; const cudaStream_t s = rank.stream; @@ -2517,15 +3224,16 @@ struct Executor::Impl { Tensor value(rank.b, DType::BF16, {h, t}); project(p.key, embedding, key, ws, s); project(p.value, embedding, value, ws, s); - Tensor history(state.history.p, DType::FP32, + const DeviceBuffer& own = pass_mode.ahead ? ahead_layers[index].history : state.history; + Tensor history(own.p, DType::FP32, {width, static_cast((config.ple_conv_kernel_size - 1) * config.ngram.ngram_size)}); const ops::PleInjectWeights weights{&p.norm_key, &p.norm_query, &p.norm_conv, &p.convolution}; auto scope = ws.scope(); - if (verifying) { + if (pass_mode.verify) { // The history stays; commit() advances it over the kept positions' recorded input. - Tensor normalized(static_cast(rank.ple_records.p) + + Tensor normalized(static_cast(rank.ple_records[pass_mode.set].p) + std::size_t(width) * std::size_t(part.column), DType::FP32, {width, t}); ops::ple_inject_record(stack, key, value, weights, config.rms_norm_eps, history, @@ -2645,7 +3353,7 @@ struct Executor::Impl { if (m.native) { Tensor y(rank.y, DType::FP32, {h, t}); if (hybrid) { - hybrid->run(hybrid_layers[index], mixed, ids, weights, shared, m.native->banks(), t == 1 || verifying, + hybrid->run(hybrid_layers[index], mixed, ids, weights, shared, m.native->banks(), t == 1 || pass_mode.verify, ws, y); return; } @@ -2727,8 +3435,10 @@ struct Executor::Impl { // Queues one segment of the pass for the parts' `t` tokens on its rank's stream; with `head`, // the logits of the last `logit_rows` tokens. + // `embedded_host` (rank threads with MTP): the pass's token embeddings are also copied there, + // pinned host memory the head rank reads for the MTP catch-up instead of asking rank 0 again. void run_segment(std::span parts, const Segment& segment, std::int32_t t, - std::uint32_t logit_rows) { + std::uint32_t logit_rows, void* embedded_host = nullptr) { RankState& rank = ranks[segment.rank]; const auto h = static_cast(config.hidden_size); const auto hc = static_cast(config.hc_count); @@ -2745,6 +3455,10 @@ struct Executor::Impl { const Tensor columns(rank.scatter + run.offset, DType::I32, {run.count}); ops::scatter(source, columns, mixed, rank.stream); } + if (embedded_host != nullptr) { + CUDA_CHECK(cudaMemcpyAsync(embedded_host, rank.mixed, std::size_t(h) * t * 2, + cudaMemcpyDeviceToHost, rank.stream)); + } ops::hyper_connection_expand(mixed, stack, rank.stream); } // A layer's MoE output is written into the stack by the next layer's read, which takes @@ -2758,7 +3472,7 @@ struct Executor::Impl { if (pending) { ops::hyper_connection_write(stack, moe_out, inject, rank.stream); } pending = false; for (const Part& part : parts) { - run_ple(plan, part.sequence->layers[i], rank, part); + run_ple(plan, i, part.sequence->layers[i], rank, part); } } { @@ -2806,9 +3520,17 @@ struct Executor::Impl { // A graph per segment of a pass over one part of `t` tokens of `sequence` (a decode step, or a // verification while `verifying`). - std::vector capture_segments(SequenceState& sequence, std::int32_t t) { + // Pipeline windows capture segments [first, end) only; the others stay empty. + std::vector capture_segments(SequenceState& sequence, std::int32_t t, + std::size_t first = 0, + std::size_t end = ~std::size_t{0}) { std::vector graphs; - for (const Segment& segment : segments) { + for (std::size_t i = 0; i < segments.size(); ++i) { + const Segment& segment = segments[i]; + if (i < first || i >= end) { + graphs.emplace_back(); + continue; + } RankBinding bind(device, segment.rank); DecodeGraphDefinition definition; const Part part{.sequence = &sequence, .column = 0, .count = t}; @@ -2872,12 +3594,15 @@ struct Executor::Impl { if (verification && verification->segments.empty() && verification->runs++ > 0) { verification->segments = capture_segments(first, t); } + const bool timed_window = verifying && window_timing_enabled(); + if (timed_window) { collect_window_timing(); } for (std::size_t i = 0; i < segments.size(); ++i) { const Segment& segment = segments[i]; if (i > 0 && segments[i - 1].rank != segment.rank) { cross(segments[i - 1].rank, segment.rank, t); } RankBinding bind(device, segment.rank); + if (timed_window) { record_window_event(i, 0, ranks[segment.rank].stream); } if (segment.rows && table) { stage_rows(ranks[segment.rank]); } if (graph && !first.decode.empty()) { first.decode[i].launch(ranks[segment.rank].stream); @@ -2886,7 +3611,9 @@ struct Executor::Impl { } else { run_segment(parts, segment, t, logit_rows); } + if (timed_window) { record_window_event(i, 1, ranks[segment.rank].stream); } } + if (timed_window) { window_timing.pending = true; } if (hybrid) { hybrid->finish(); } // A verification commits its positions later (commit()). if (!verifying) { @@ -2902,6 +3629,52 @@ struct Executor::Impl { if (misses && t <= kVectorTokens) { synchronize_ranks(); } } + // NINFER_QWEN4_ROUND_TIMING=1: each stage's device time in a verification window (events around + // its segments, read at the next window), printed every 100 windows. + struct WindowTiming { + std::vector> events; // by segment + std::array stage{}; + std::uint64_t windows = 0; + bool pending = false; + }; + WindowTiming window_timing; + static bool window_timing_enabled() { + static const bool enabled = [] { + const char* value = std::getenv("NINFER_QWEN4_ROUND_TIMING"); + return value != nullptr && value[0] == '1'; + }(); + return enabled; + } + void record_window_event(std::size_t segment, int end, cudaStream_t stream) { + auto& events = window_timing.events; + if (events.size() != segments.size()) { + events.resize(segments.size(), {nullptr, nullptr}); + } + auto& event = events[segment][std::size_t(end)]; + if (event == nullptr) { CUDA_CHECK(cudaEventCreate(&event)); } + CUDA_CHECK(cudaEventRecord(event, stream)); + } + void collect_window_timing() { + WindowTiming& w = window_timing; + if (!w.pending) { return; } + w.pending = false; + for (std::size_t i = 0; i < w.events.size(); ++i) { + RankBinding bind(device, segments[i].rank); + CUDA_CHECK(cudaEventSynchronize(w.events[i][1])); + float ms = 0; + CUDA_CHECK(cudaEventElapsedTime(&ms, w.events[i][0], w.events[i][1])); + w.stage[std::min(segments[i].rank, 7)] += ms; + } + if (++w.windows % 100 != 0) { return; } + std::fprintf(stderr, "window timing | %llu windows | device ms/window:", + static_cast(w.windows)); + for (std::size_t r = 0; r < ranks.size() && r < w.stage.size(); ++r) { + std::fprintf(stderr, " r%zu %.2f", r, w.stage[r] / 100.0); + } + std::fprintf(stderr, "\n"); + w.stage = {}; + } + void synchronize_ranks() { for (auto& rank : ranks) { RankBinding bind(device, rank.rank); @@ -3134,7 +3907,7 @@ struct Executor::Impl { Tensor stack(rank.stack, DType::FP32, {h, hc, t}); Tensor mixed(rank.mixed, DType::BF16, {h, t}); Tensor inject(rank.inject, DType::FP32, {hc, t}); - if (plan.ple) { run_ple(plan, state, rank, part); } + if (plan.ple) { run_ple(plan, i, state, rank, part); } { auto scope = rank.workspace->scope(); ops::hyper_connection_read(stack, plan.attn_hc.weights(), @@ -3238,6 +4011,7 @@ struct Executor::Impl { throw std::invalid_argument("qwen4_exp forward: more next tokens than a prompt step"); } if (tokens.size() > options.prefill_chunk) { + drain(); forward_span(s, tokens, logit_rows, next); return; } @@ -3246,6 +4020,12 @@ struct Executor::Impl { const std::uint32_t position = sequence.position; const Part part{ .sequence = &sequence, .column = 0, .count = static_cast(tokens.size())}; + if (threadable(std::span(&part, 1), part.count)) { + threaded_pass(std::span(&part, 1), tokens, logit_rows); + hint_next(sequence, next); + return; + } + drain(); pass(std::span(&part, 1), tokens, logit_rows); hint_next(sequence, next); if (mtp) { @@ -3254,7 +4034,8 @@ struct Executor::Impl { } } - void decode(std::span batch, std::span tokens) { + void decode(std::span batch, std::span tokens, + Executor::AfterHead after_head) { if (batch.empty() || batch.size() != tokens.size()) { throw std::invalid_argument("qwen4_exp decode: one token per sequence"); } @@ -3270,6 +4051,12 @@ struct Executor::Impl { } parts.push_back({.sequence = sequence, .column = std::int32_t(j), .count = 1}); } + if (threadable(parts, static_cast(tokens.size()))) { + threaded_pass(parts, tokens, static_cast(tokens.size()), + std::move(after_head)); + return; + } + drain(); std::vector committed; for (const Part& part : parts) { committed.push_back({part.sequence, part.column, 1, part.sequence->position, @@ -3277,6 +4064,15 @@ struct Executor::Impl { } pass(parts, tokens, static_cast(tokens.size())); if (mtp) { mtp_catch_up(committed); } + if (after_head) { + RankBinding bind(device, model.head_rank()); + after_head(ranks[model.head_rank()].stream, head_logits(tokens.size())); + } + } + + [[nodiscard]] Tensor head_logits(std::size_t rows) const { + return Tensor(ranks[model.head_rank()].logits, DType::BF16, + {static_cast(config.vocab_size), static_cast(rows)}); } void warm_up() { @@ -3300,7 +4096,7 @@ struct Executor::Impl { prefill(2); const std::uint32_t sequence = 0; if (sequences.front().position + 1 < options.max_context) { - decode(std::span(&sequence, 1), std::span(tokens).first(1)); + decode(std::span(&sequence, 1), std::span(tokens).first(1), {}); } if (mtp && can_draft(0) && sequences.front().position + verify_width <= options.max_context) { @@ -3407,10 +4203,10 @@ struct Executor::Impl { const std::size_t head = model.head_rank(); const auto h = static_cast(config.hidden_size); RankState& last = ranks[head]; - if (head == 0) { + if (head == 0 || head_embedding.payload != nullptr) { Tensor rows(out, DType::BF16, {h, n}); - ops::embedding(Tensor(const_cast(ids), DType::I32, {n}), embedding_table, - rows, last.stream); + ops::embedding(Tensor(const_cast(ids), DType::I32, {n}), + head == 0 ? embedding_table : head_embedding, rows, last.stream); return; } RankState& first = ranks[0]; @@ -3451,7 +4247,19 @@ struct Executor::Impl { const auto h = static_cast(config.hidden_size); const auto hc = static_cast(config.hc_count); const float eps = config.rms_norm_eps; - embed_cells(ids, n, rank.mixed); + if (catch_up_host != nullptr) { + // A rank-thread catch-up: the cells' embeddings are the pass's own, which rank 0 left + // in pinned host memory before it handed the stack on. + for (const CatchUpRows& rows : catch_up_rows) { + CUDA_CHECK(cudaMemcpyAsync(static_cast(rank.mixed) + + std::size_t(rows.cell) * h * 2, + catch_up_host + std::size_t(rows.column) * h * 2, + std::size_t(rows.count) * h * 2, + cudaMemcpyHostToDevice, s)); + } + } else { + embed_cells(ids, n, rank.mixed); + } Tensor embedded(rank.mixed, DType::BF16, {h, n}); Tensor embedded_norm(rank.d, DType::BF16, {h, n}); ops::rmsnorm(embedded, mtp->embedding_norm, eps, true, embedded_norm, s); @@ -3550,13 +4358,16 @@ struct Executor::Impl { // position, each pairing that position's stack with the next token, and keeps the last // committed position's stack for the next cell. The target's stacks are still in the head // rank's stack plane. - void mtp_catch_up(std::span committed) { + // `embedded` (rank threads): the pass's token embeddings in pinned host memory, one column per + // pass column, which the cells take instead of a round trip to rank 0. + void mtp_catch_up(std::span committed, const std::byte* embedded = nullptr) { RankBinding bind(device, model.head_rank()); RankState& rank = ranks[model.head_rank()]; const auto width = std::size_t(config.hc_count) * config.hidden_size; std::vector parts; std::vector ids, positions; std::int32_t n = 0; + catch_up_rows.clear(); for (const Committed& c : committed) { SequenceState& sequence = *c.sequence; if (!sequence.mtp_follows || c.count == 0) { continue; } @@ -3565,6 +4376,8 @@ struct Executor::Impl { const std::int32_t cells = c.count - (kept ? 0 : 1); const std::uint32_t first = kept ? c.position - 1 : 0; if (cells == 0) { continue; } + catch_up_rows.push_back( + {.cell = n, .column = c.column + (kept ? 0 : 1), .count = cells}); if (kept) { stage_hidden(sequence.mtp_pending.p, 1, n); } if (c.count > 1) { stage_hidden(rank.stack + width * std::size_t(c.column), c.count - 1, @@ -3588,11 +4401,25 @@ struct Executor::Impl { if (hybrid) { hybrid->begin(cancelled, !warming); } stage_mtp(ids, 0, positions, 0); if (misses) { misses->keep_alive(); } - mtp_pass(parts, n, static_cast(mtp_ids.p), - static_cast(mtp_positions.p), false); + catch_up_host = embedded; + try { + mtp_pass(parts, n, static_cast(mtp_ids.p), + static_cast(mtp_positions.p), false); + } catch (...) { + catch_up_host = nullptr; + throw; + } + catch_up_host = nullptr; if (hybrid) { hybrid->finish(); } } + // The rows of a rank-thread catch-up's cells in its pass's host embeddings. + struct CatchUpRows { + std::int32_t cell = 0, column = 0, count = 0; + }; + std::vector catch_up_rows; + const std::byte* catch_up_host = nullptr; + void draft(std::span batch, std::span anchors, std::span out, std::span extents = {}, std::uint32_t steps = 0) { @@ -3622,19 +4449,24 @@ struct Executor::Impl { } drafting.push_back(sequence); } + drain_head(); RankBinding bind(device, model.head_rank()); RankState& rank = ranks[model.head_rank()]; // Step i's cell of a sequence sits at its position - 1 + i and pairs the stack before it - // (the kept one at step 0) with the step's token (the anchor at step 0). - std::vector positions(k * b); + // (the kept one at step 0) with the step's token (the anchor at step 0). Pipeline windows + // stage the positions of the guess chain's steps after them too. + const std::size_t staged = pipeline && b == 1 ? 2 * k + 1 : k; + std::vector positions(staged * b); std::vector parts; for (std::size_t j = 0; j < b; ++j) { const std::uint32_t first = drafting[j]->position - 1; - for (std::size_t i = 0; i < k; ++i) { + for (std::size_t i = 0; i < staged; ++i) { positions[i * b + j] = static_cast(first + i); } parts.push_back({drafting[j], std::int32_t(j), 1, first}); } + draft_first = drafting.front()->position - 1; + draft_steps = static_cast(k); if (hybrid) { hybrid->begin(cancelled, !warming); } stage_mtp(anchors, 0, positions, 0); if (draft_prefetch) { @@ -3669,7 +4501,9 @@ struct Executor::Impl { // captured exchange as text decode; GGUF host experts replay too, without counting the // draft steps' routes. auto& only = drafting.front()->draft.at(k); - const bool graph = options.cuda_graphs && b == 1 && device.size() == 1 && !stream; + // A split model's chain stays on the head rank once its embeddings come from the host copy. + const bool graph = options.cuda_graphs && b == 1 && + (device.size() == 1 || head_embedding.payload != nullptr) && !stream; if (misses) { misses->keep_alive(); } if (graph && !only.chain.ready() && only.runs++ > 0) { DecodeGraphDefinition definition; @@ -3744,7 +4578,28 @@ struct Executor::Impl { tokens.begin() + std::ptrdiff_t((j + 1) * w)); verified.push_back(sequence); } + // The host's pass orders itself after every rank thread's queued work. + if (!pipeline) { drain(); } verifying = true; + verified_set = pipeline ? next_set() : 0; + pass_mode = {.verify = true, + .ahead = false, + .set = verified_set, + .rows = static_cast(b), + .width = static_cast(w)}; + if (pipeline) { + try { + threaded_verify(parts, tokens); + } catch (...) { + pass_mode = {}; + verifying = false; + verified.clear(); + throw; + } + pass_mode = {}; + verifying = false; + return; + } try { pass(parts, tokens, static_cast(tokens.size())); if (draft_prefetch) { @@ -3764,55 +4619,52 @@ struct Executor::Impl { draft_prefetch->join(rank.stream); } } catch (...) { + pass_mode = {}; verifying = false; verified.clear(); throw; } + pass_mode = {}; verifying = false; } - void commit(std::span batch, std::span kept) { - require_mtp("commit"); - const std::uint32_t w = active_verify_width; - if (batch.size() != verified.size() || kept.size() != batch.size()) { - throw std::invalid_argument("qwen4_exp commit: the last verification's sequences"); - } - for (std::size_t j = 0; j < batch.size(); ++j) { - if (&sequences.at(batch[j]) != verified[j] || kept[j] == 0 || kept[j] > w) { - throw std::invalid_argument("qwen4_exp commit: the last verification's sequences, " - "1 to draft_tokens + 1 tokens each"); - } - } - if (draft_prefetch) { - RankBinding bind(device, model.head_rank()); - // Captured events encode graph dependencies; wait on the joined stream before - // reading host callback state. Host synchronization of a captured event is invalid. - CUDA_CHECK(cudaStreamSynchronize(ranks[model.head_rank()].stream)); - draft_prefetch->check(); - } - const auto di = static_cast(config.indexer_head_dim); + // One rank's part of a commit, on its stream: its sparse-attention layers' indexer tails and + // its PLE histories advance over the kept positions' records of `set` (a window of width `w`), + // and its Gated DeltaNet layers replay the kept transitions. `shadow` (pipeline windows): the + // window is taken whole into the shadow states, which start from the sequence's own. + void commit_rank(std::size_t r, std::span rows_of, + std::span kept, std::uint32_t set, std::uint32_t w, + bool shadow) { + RankState& rank = ranks[r]; + const auto di = static_cast(config.indexer_head_dim); const std::uint64_t index_rows = std::uint64_t(config.indexer_n_heads + 1) * config.indexer_head_dim; const std::uint64_t columns = std::uint64_t(verify_width) * options.sequences; const auto stack_width = std::int32_t(config.hc_count * config.hidden_size); + // The window's positions: still staged, unless a later pass has been staged since. + const auto* positions = pipeline ? static_cast(rank.window_positions[set].p) + : rank.positions; for (std::size_t i = 0; i < layers.size(); ++i) { const LayerPlan& plan = layers[i]; - RankBinding bind(device, plan.rank); - RankState& rank = ranks[plan.rank]; - for (std::size_t j = 0; j < verified.size(); ++j) { - LayerState& state = verified[j]->layers[i]; - const auto m = static_cast(kept[j]); + if (plan.rank != r) { continue; } + for (std::size_t j = 0; j < rows_of.size(); ++j) { + LayerState& state = rows_of[j]->layers[i]; + LayerState& target = shadow ? ahead_layers[i] : state; + const auto m = static_cast(kept[j]); if (plan.qsa) { - // The indexer's tail advances over the kept positions' recorded projections; - // the verification's first position is still staged. + // The indexer's tail advances over the kept positions' recorded projections. + if (shadow) { + CUDA_CHECK(cudaMemcpyAsync(target.tail.p, state.tail.p, state.tail.bytes, + cudaMemcpyDeviceToDevice, rank.stream)); + } const QsaPlan& a = *plan.qsa; - const Tensor projection(static_cast(rank.index_records.p) + + const Tensor projection(static_cast(rank.index_records[set].p) + (qsa_local[i] * columns + j * w) * index_rows * 2, DType::BF16, {a.index.rows, m}); - const Tensor first(rank.positions + j * w, DType::I32, {1}); + const Tensor first(const_cast(positions) + j * w, DType::I32, {1}); Tensor pooled(state.pooled.p, DType::FP32, {di, static_cast(pooled_slots)}); - Tensor tail(state.tail.p, DType::FP32, + Tensor tail(target.tail.p, DType::FP32, {di, static_cast(config.indexer_compress_ratio - 1)}); const ops::QsaIndexerWeights iw{&a.index_query_norm, &a.index_key_norm}; ops::qsa_indexer_append(projection, first, iw, config.rms_norm_eps, pooled, @@ -3820,11 +4672,16 @@ struct Executor::Impl { } if (plan.ple) { // The PLE history advances over the kept positions' recorded input. + if (shadow) { + CUDA_CHECK(cudaMemcpyAsync(target.history.p, state.history.p, + state.history.bytes, cudaMemcpyDeviceToDevice, + rank.stream)); + } Tensor history( - state.history.p, DType::FP32, + target.history.p, DType::FP32, {stack_width, static_cast((config.ple_conv_kernel_size - 1) * config.ngram.ngram_size)}); - const Tensor normalized(static_cast(rank.ple_records.p) + + const Tensor normalized(static_cast(rank.ple_records[set].p) + std::size_t(stack_width) * j * w, DType::FP32, {stack_width, m}); ops::ple_history_advance(history, normalized, rank.stream); @@ -3832,15 +4689,43 @@ struct Executor::Impl { } } // The Gated DeltaNet layers replay the kept transitions into each sequence's slot. + if (!rank.fold[set]) { return; } std::vector rows; - for (std::size_t j = 0; j < verified.size(); ++j) { - const auto slot = static_cast(verified[j]->slot); - rows.push_back({slot, slot, static_cast(kept[j])}); + for (std::size_t j = 0; j < rows_of.size(); ++j) { + const auto slot = static_cast(rows_of[j]->slot); + const auto to = shadow ? static_cast(options.sequences) : slot; + rows.push_back({slot, to, static_cast(kept[j])}); } - for (RankState& rank : ranks) { - if (!rank.fold) { continue; } - RankBinding bind(device, rank.rank); - rank.fold->execute(rows, static_cast(w), rank.stream); + rank.fold[set]->execute(rows, static_cast(w), rank.stream); + } + + void commit(std::span batch, std::span kept) { + require_mtp("commit"); + const std::uint32_t w = active_verify_width; + if (batch.size() != verified.size() || kept.size() != batch.size()) { + throw std::invalid_argument("qwen4_exp commit: the last verification's sequences"); + } + for (std::size_t j = 0; j < batch.size(); ++j) { + if (&sequences.at(batch[j]) != verified[j] || kept[j] == 0 || kept[j] > w) { + throw std::invalid_argument("qwen4_exp commit: the last verification's sequences, " + "1 to draft_tokens + 1 tokens each"); + } + } + if (pipeline) { + pipelined_commit(kept); + return; + } + drain(); + if (draft_prefetch) { + RankBinding bind(device, model.head_rank()); + // Captured events encode graph dependencies; wait on the joined stream before + // reading host callback state. Host synchronization of a captured event is invalid. + CUDA_CHECK(cudaStreamSynchronize(ranks[model.head_rank()].stream)); + draft_prefetch->check(); + } + for (std::size_t r = 0; r < ranks.size(); ++r) { + RankBinding bind(device, r); + commit_rank(r, verified, kept, verified_set, w, false); } std::vector committed; std::vector scratch; @@ -3860,6 +4745,308 @@ struct Executor::Impl { mtp_catch_up(committed); } + // ---- Pipeline windows ---- + // A round's verification goes to the rank threads. Before it, the head rank's thread runs the + // MTP chain on past the round's drafts to the window that would follow if this one is kept + // whole: its bonus token and drafts after it (the guess). The host posts that window to the + // front ranks at once; they run it on shadow states (theirs advanced over the window in flight, + // taken whole) while the head rank verifies. When the verdict keeps the window whole and its + // bonus is the guessed one, the head rank completes the guessed window from the front ranks' + // stack; otherwise the guess is dropped and the next round drafts as before. Every token is + // still the verification's, so the output is the serial loop's. + std::size_t head_first = 0; // the first segment of the head rank + std::uint32_t set_flip = 0; // the record set of the last window + struct Window { // the last window posted to the front ranks + std::uint32_t set = 0, width = 0, position = 0; + std::vector tokens; + NgramContext context; // the n-gram context before it + } window; + struct Ahead { // the guessed window the front ranks run + bool live = false; + std::uint64_t number = 0; + std::uint32_t set = 0, width = 0, position = 0; + std::vector tokens; + } ahead; + std::unique_ptr guess_host; + cudaEvent_t guess_event = nullptr; + std::uint64_t guess_posted = 0; // the host's count of guess chains + std::uint64_t guess_recorded = 0; // the head thread's (worker_mutex) + std::uint32_t draft_first = 0, draft_steps = 0; // the last draft chain + // Graphs per record set and width: verification (every segment) and guessed windows (the + // front ranks' segments), for the one sequence. + std::array, 2> verify_graphs, ahead_graphs; + static constexpr std::size_t kDiscardCell = 512; // where forced chain steps write their argmax + + void start_pipeline() { + const std::size_t head = model.head_rank(); + head_first = 0; + while (head_first < segments.size() && segments[head_first].rank != head) { ++head_first; } + bool layout = head_first > 0 && head_first < segments.size(); + for (std::size_t i = head_first; layout && i < segments.size(); ++i) { + layout = segments[i].rank == head; + } + if (!layout) { + pipeline = false; + return; + } + ahead_buffer = std::make_unique(stack_bytes(std::int32_t(verify_width))); + { + RankBinding bind(device, segments[head_first - 1].rank); + CUDA_CHECK(cudaEventCreateWithFlags(&ahead_written, cudaEventDisableTiming)); + } + RankBinding bind(device, head); + CUDA_CHECK(cudaEventCreateWithFlags(&ahead_read, cudaEventDisableTiming)); + CUDA_CHECK(cudaEventCreateWithFlags(&guess_event, cudaEventDisableTiming)); + guess_host = std::make_unique(64 * sizeof(std::int32_t)); + std::fprintf(stderr, "qwen4_exp: pipeline windows on (%zu front segments)\n", head_first); + } + + std::uint32_t next_set() { + set_flip ^= 1U; + return set_flip; + } + + // A verification posted to the rank threads (the host's pass_mode describes it). + void threaded_verify(std::span parts, std::span tokens) { + SequenceState& sequence = *parts.front().sequence; + const auto w = static_cast(tokens.size()); + NgramContext context = sequence.context; + NgramContext* contexts[] = {&context}; + auto task = stage_task(parts, tokens, w, contexts, 0, segments.size()); + task->mode = pass_mode; + auto& graphs = verify_graphs[pass_mode.set][w]; + if (graphs.segments.empty() && graphs.runs++ > 0) { + drain(); + graphs.segments = capture_segments(sequence, std::int32_t(w), 0, segments.size()); + } + if (!graphs.segments.empty()) { task->graphs = &graphs.segments; } + post(task); + window = {.set = pass_mode.set, + .width = w, + .position = sequence.position, + .tokens = std::vector(tokens.begin(), tokens.end()), + .context = sequence.context}; + } + + // The MTP chain on the head rank's thread from step `from` to step `to` of a chain whose cells + // start at position `first`; steps before `forced` take the staged tokens (their argmax is + // dropped). The argmax of steps forced - 1 .. to - the guessed bonus and drafts - goes to the + // host. + void run_guess_chain(SequenceState* sequence, std::uint32_t first, std::uint32_t from, + std::uint32_t to, std::uint32_t forced, std::uint64_t number) { + RankState& rank = ranks[model.head_rank()]; + auto* ids = static_cast(mtp_ids.p); + const auto* at = static_cast(mtp_positions.p); + const MtpPart part{sequence, 0, 1, first}; + for (std::uint32_t i = from; i <= to; ++i) { + if (i == 0) { + CUDA_CHECK(cudaMemcpyAsync(sequence->mtp_scratch_tail.p, sequence->mtp.tail.p, + sequence->mtp.tail.bytes, cudaMemcpyDeviceToDevice, + rank.stream)); + stage_hidden(sequence->mtp_pending.p, 1, 0); + } else { + stage_hidden(rank.stack, 1, 0); + } + mtp_pass(std::span(&part, 1), 1, ids + i, at + i, true); + mtp_head(1, i + 1 < forced ? ids + kDiscardCell : ids + i + 1); + } + const std::uint32_t k = to / 2; // the guess: steps k .. 2k write cells k + 1 .. 2k + 1 + CUDA_CHECK(cudaMemcpyAsync(guess_host->data(), ids + k + 1, (k + 1) * sizeof(std::int32_t), + cudaMemcpyDeviceToHost, rank.stream)); + CUDA_CHECK(cudaEventRecord(guess_event, rank.stream)); + { + std::lock_guard lock(worker_mutex); + guess_recorded = number; + } + worker_cv.notify_all(); + } + + void post_head_work(std::function work) { + auto task = new_task(); + task->work = [this, work = std::move(work)](std::size_t r) { + if (r == model.head_rank()) { work(); } + }; + post(task); + } + + // After draft(): the head rank's thread runs the chain on from the drafts (before it takes + // the verification's stack). + void guess_after_draft() { + const std::uint64_t number = ++guess_posted; + SequenceState* sequence = &sequences.front(); + const std::uint32_t k = draft_steps, first = draft_first; + post_head_work([=, this] { run_guess_chain(sequence, first, k, 2 * k, 0, number); }); + } + + // The guess of the last chain: the bonus token and the drafts after it (width tokens). + void take_guess(std::span out) { + { + std::unique_lock lock(worker_mutex); + worker_cv.wait(lock, [&] { return guess_recorded >= guess_posted || worker_error; }); + if (worker_error) { std::rethrow_exception(worker_error); } + } + RankBinding bind(device, model.head_rank()); + CUDA_CHECK(cudaEventSynchronize(guess_event)); + std::memcpy(out.data(), guess_host->data(), out.size() * sizeof(std::int32_t)); + } + + // The front ranks run `tokens` as the window after the one in flight, taken whole. + void verify_ahead(std::span tokens) { + SequenceState& sequence = sequences.front(); + const auto w = static_cast(tokens.size()); + NgramContext start = window.context; + if (table) { + std::vector scratch(window.tokens.size() * config.ngram_heads()); + ngram_row_ids(ngram, window.tokens, config.eos_token_id, config.vocab_size, start, + scratch); + } + NgramContext context = start; + NgramContext* contexts[] = {&context}; + const std::uint32_t position = window.position + window.width; + if (position + w > options.max_context) { + throw std::invalid_argument("qwen4_exp: a guessed window past max_context"); + } + const Part part{.sequence = &sequence, .column = 0, .count = std::int32_t(w)}; + const std::uint32_t set = next_set(); + auto task = stage_task(std::span(&part, 1), tokens, w, contexts, + position - sequence.position, head_first); + task->mode = {.verify = true, .ahead = true, .set = set, .rows = 1, .width = w}; + task->to_ahead = true; + task->ahead_number = ++ahead_count; + const std::uint32_t source = window.set, width = window.width; + SequenceState* rows_of = &sequence; + task->work = [this, source, width, rows_of](std::size_t r) { + if (r == model.head_rank()) { return; } + SequenceState* const of[] = {rows_of}; + const std::uint32_t kept[] = {width}; + commit_rank(r, of, kept, source, width, true); + }; + auto& graphs = ahead_graphs[set][w]; + if (graphs.segments.empty() && graphs.runs++ > 0) { + drain(); + const PassMode host = pass_mode; + pass_mode = task->mode; + try { + graphs.segments = capture_segments(sequence, std::int32_t(w), 0, head_first); + } catch (...) { + pass_mode = host; + throw; + } + pass_mode = host; + } + if (!graphs.segments.empty()) { task->graphs = &graphs.segments; } + post(task); + ahead = {.live = true, + .number = task->ahead_number, + .set = set, + .width = w, + .position = position, + .tokens = std::vector(tokens.begin(), tokens.end())}; + window = {.set = set, + .width = w, + .position = position, + .tokens = ahead.tokens, + .context = start}; + } + + [[nodiscard]] bool guess_live() const noexcept { return ahead.live; } + + // The verdict dropped the guessed window. + void drop_guess() { + if (!ahead.live) { return; } + { + std::lock_guard lock(worker_mutex); + ahead_resolved = std::max(ahead_resolved, ahead.number); + } + worker_cv.notify_all(); + ahead.live = false; + } + + // The verdict kept the window in flight whole with the guessed bonus: the guessed window is + // the next verification. The head rank's thread first runs the next guess chain (through the + // guessed window's tokens, then past them), then completes the window from the front ranks' + // stack. + void verify_guessed(std::span tokens) { + SequenceState& sequence = sequences.front(); + const auto w = static_cast(tokens.size()); + if (!ahead.live || w != ahead.width || sequence.position != ahead.position || + !std::equal(tokens.begin(), tokens.end(), ahead.tokens.begin())) { + throw std::logic_error("qwen4_exp: the verification is not the guessed window"); + } + verified.assign(1, &sequence); + sequence.verify_tokens.assign(tokens.begin(), tokens.end()); + sequence.verify_context = sequence.context; + active_verify_width = w; + verified_set = ahead.set; + const std::uint32_t k = w - 1; + const std::uint32_t first = sequence.position - 1; + const std::uint64_t number = ++guess_posted; + std::vector forced(tokens.begin(), tokens.end()); + std::vector positions(2 * k + 1); + std::iota(positions.begin(), positions.end(), static_cast(first)); + SequenceState* of = &sequence; + post_head_work([=, this] { + stage_mtp(forced, 0, positions, 0); + run_guess_chain(of, first, 0, 2 * k, k + 1, number); + }); + NgramContext context = sequence.context; + NgramContext* contexts[] = {&context}; + const Part part{.sequence = &sequence, .column = 0, .count = std::int32_t(w)}; + auto task = stage_task(std::span(&part, 1), tokens, w, contexts, 0, segments.size(), + head_first); + task->mode = {.verify = true, .ahead = false, .set = ahead.set, .rows = 1, .width = w}; + task->from_ahead = true; + task->ahead_number = ahead.number; + auto& graphs = verify_graphs[ahead.set][w]; + if (!graphs.segments.empty()) { task->graphs = &graphs.segments; } + post(task); + ahead.live = false; + } + + // A commit on the rank threads: each rank's part, and on the head rank the MTP catch-up. + void pipelined_commit(std::span kept) { + const std::uint32_t w = active_verify_width; + auto task = new_task(); + std::vector rows_of(verified.begin(), verified.end()); + std::vector counts(kept.begin(), kept.end()); + for (SequenceState* sequence : rows_of) { + task->tokens.insert(task->tokens.end(), sequence->verify_tokens.begin(), + sequence->verify_tokens.end()); + } + std::vector scratch; + for (std::size_t j = 0; j < rows_of.size(); ++j) { + SequenceState& sequence = *rows_of[j]; + const std::span tokens(task->tokens.data() + j * w, kept[j]); + if (table) { + scratch.resize(tokens.size() * config.ngram_heads()); + ngram_row_ids(ngram, tokens, config.eos_token_id, config.vocab_size, + sequence.context, scratch); + } + task->committed.push_back({&sequence, static_cast(j * w), + static_cast(kept[j]), sequence.position, + tokens}); + sequence.position += kept[j]; + } + const std::uint32_t set = verified_set; + RankTask* own = task.get(); + task->work = [this, rows_of, counts, set, w, own](std::size_t r) { + commit_rank(r, rows_of, counts, set, w, false); + if (r == model.head_rank()) { mtp_catch_up(own->committed); } + }; + verified.clear(); + post(task); + } + + void reset_pipeline() { + if (!pipeline) { return; } + { + std::lock_guard lock(worker_mutex); + ahead_resolved = ahead_count; + } + worker_cv.notify_all(); + ahead.live = false; + } + std::uint64_t snapshot_bytes() const { const SequenceState& sequence = sequences.at(0); std::uint64_t bytes = 0; @@ -4059,7 +5246,11 @@ struct Executor::Impl { }; Executor::Executor(const Model& model, DeviceContext& device, ExecutorOptions options) - : impl_(std::make_unique(model, device, std::move(options))) {} + : impl_(std::make_unique(model, device, std::move(options))) { + impl_->start_rank_threads(); + // Pipeline windows run on the rank threads. + if (!impl_->rank_threads) { impl_->pipeline = false; } +} Executor::~Executor() = default; @@ -4075,9 +5266,14 @@ std::uint64_t Executor::ngram_resident_bytes() const noexcept { void Executor::reset(std::uint32_t sequence) { impl_->reset(sequence); } -void Executor::warm_up() { impl_->warm_up(); } +void Executor::warm_up() { + impl_->drain(); + impl_->warm_up(); + impl_->drain(); +} std::uint32_t Executor::prompt_step() const noexcept { return impl_->prompt_step(); } void Executor::grow_expert_cache(std::uint64_t keep_free_bytes) { + impl_->drain(); impl_->grow_expert_cache(keep_free_bytes); } @@ -4092,16 +5288,23 @@ void Executor::forward(std::uint32_t sequence, std::span tok void Executor::decode(std::span sequences, std::span tokens) { - impl_->decode(sequences, tokens); + impl_->decode(sequences, tokens, {}); +} + +void Executor::decode(std::span sequences, + std::span tokens, AfterHead after_head) { + impl_->decode(sequences, tokens, std::move(after_head)); } void Executor::snapshot(std::uint32_t sequence, SequenceSnapshot& out) { + impl_->drain(); impl_->snapshot(sequence, out); } std::uint64_t Executor::snapshot_bytes() const { return impl_->snapshot_bytes(); } void Executor::restore(std::uint32_t sequence, const SequenceSnapshot& from) { + impl_->drain(); impl_->restore(sequence, from); } @@ -4111,16 +5314,19 @@ std::uint64_t Executor::image_bytes(std::uint32_t position) const { SequenceImage Executor::save_image(std::uint32_t sequence, const SequenceSnapshot* at, std::span pages, std::size_t page_bytes) { + impl_->drain(); return impl_->save_image(sequence, at, pages, page_bytes); } void Executor::load_image(std::uint32_t sequence, const SequenceImage& image, std::span pages, std::size_t page_bytes) { + impl_->drain(); impl_->load_image(sequence, image, pages, page_bytes); } void Executor::set_media(std::uint32_t sequence, std::span items, std::vector rope_positions, std::int32_t rope_delta) { + impl_->drain(); impl_->set_media(sequence, items, std::move(rope_positions), rope_delta); } @@ -4148,7 +5354,24 @@ void Executor::commit(std::span sequences, impl_->commit(sequences, columns); } +bool Executor::pipeline_windows() const noexcept { return impl_->pipeline; } + +void Executor::guess_after_draft() { impl_->guess_after_draft(); } + +void Executor::take_guess(std::span out) { impl_->take_guess(out); } + +void Executor::verify_ahead(std::span tokens) { impl_->verify_ahead(tokens); } + +bool Executor::guess_live() const noexcept { return impl_->guess_live(); } + +void Executor::drop_guess() { impl_->drop_guess(); } + +void Executor::verify_guessed(std::span tokens) { + impl_->verify_guessed(tokens); +} + Tensor Executor::logits(std::uint32_t rows) const { + impl_->drain_head_quietly(); return Tensor( impl_->ranks[impl_->model.head_rank()].logits, DType::BF16, {static_cast(impl_->config.vocab_size), static_cast(rows)}); @@ -4176,6 +5399,7 @@ ExpertCacheStats Executor::expert_cache_stats() const noexcept { } cudaStream_t Executor::head_stream() const noexcept { + impl_->drain_head_quietly(); return impl_->ranks[impl_->model.head_rank()].stream; } @@ -4211,6 +5435,7 @@ void Executor::raise_service_failure() { } void Executor::abort(std::uint32_t sequence) { + impl_->drain_quietly(); if (impl_->table) { impl_->table->wait(); } impl_->rows_waited = true; impl_->device.synchronize(); diff --git a/src/models/qwen4_exp/executor.h b/src/models/qwen4_exp/executor.h index 087879c122..5f45f7f891 100644 --- a/src/models/qwen4_exp/executor.h +++ b/src/models/qwen4_exp/executor.h @@ -29,6 +29,7 @@ #include "models/qwen3_5/program/vision_control.h" #include +#include #include #include #include @@ -171,6 +172,13 @@ class Executor { // leaves their logits in logits(), one column per sequence in that order. The experts read // their weights once for the whole batch; each sequence's mixers run on its own state. void decode(std::span sequences, std::span tokens); + // The same, and `after_head` runs once the pass's head is queued, with the head stream and + // the pass's logits: work it queues there reads them before any later pass overwrites them. + // With rank threads it runs on the head rank's thread (the caller must not touch the head + // stream or logits() before it has run); otherwise before decode() returns. + using AfterHead = std::function; + void decode(std::span sequences, std::span tokens, + AfterHead after_head); // Copies the sequence's recurrent state into `out`, and back; both wait for the copies. void snapshot(std::uint32_t sequence, SequenceSnapshot& out); // Device bytes one SequenceSnapshot's buffers take once its first snapshot sizes them. @@ -221,6 +229,24 @@ class Executor { // then what feeding those tokens would have left, and its MTP block has run their cells. void commit(std::span sequences, std::span columns); + // Pipeline windows (NINFER_QWEN4_PIPELINE_WINDOWS=1 with rank threads, one sequence): while + // the head rank verifies a window, the front ranks run the guessed next one. + [[nodiscard]] bool pipeline_windows() const noexcept; + // After draft() and before verify(): the MTP chain runs on past the drafts to the guess. + void guess_after_draft(); + // The last guess: the bonus token the window in flight would end with if kept whole, then + // the drafts after it (draft_tokens() + 1 tokens). Waits for the head rank's chain. + void take_guess(std::span out); + // The front ranks run `tokens` (a guess) as the window after the one in flight. + void verify_ahead(std::span tokens); + [[nodiscard]] bool guess_live() const noexcept; + // The verdict did not keep the window in flight whole with the guessed bonus. + void drop_guess(); + // After commit() kept the window in flight whole with the guessed bonus: verifies the guessed + // window (`tokens`, as guessed), whose front ranks' part has run, and starts the next guess + // chain (take_guess()). + void verify_guessed(std::span tokens); + [[nodiscard]] Tensor logits(std::uint32_t rows) const; [[nodiscard]] std::size_t head_rank() const noexcept; [[nodiscard]] ExpertCacheStats expert_cache_stats() const noexcept; diff --git a/src/runtime/engine/qwen4_exp_core.cpp b/src/runtime/engine/qwen4_exp_core.cpp index f39c179c32..be73e64488 100644 --- a/src/runtime/engine/qwen4_exp_core.cpp +++ b/src/runtime/engine/qwen4_exp_core.cpp @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -668,6 +669,10 @@ struct Qwen4ExpCore::Impl { } queue_cv.notify_all(); if (worker.joinable()) { worker.join(); } + if (!group_done.empty()) { + RankBinding bind(device, instance.executor->head_rank()); + for (cudaEvent_t event : group_done) { cudaEventDestroy(event); } + } } [[nodiscard]] std::size_t mask_words() const noexcept { @@ -1848,6 +1853,19 @@ struct Qwen4ExpCore::Impl { } } if (batch.empty() && alone.empty() && speculative.empty() && deferred.empty()) { return; } + // With decode groups over a pipeline, several decoding requests run plain steps in groups + // that keep every stage busy, which outruns their MTP rounds; a request alone keeps its + // drafts (the rounds of one sequence cannot overlap stages). + if (decode_groups() >= 2 && device.size() >= 2 && speculative.size() + batch.size() >= 2) { + for (const auto& request : speculative) { + if (!request->constrained && !request->options.execution.logprobs) { + batch.push_back(request); + } + } + std::erase_if(speculative, [](const auto& r) { + return !r->constrained && !r->options.execution.logprobs; + }); + } const auto start = Clock::now(); if (!deferred.empty()) { sample_deferred(deferred); } for (const auto& request : alone) { @@ -1928,10 +1946,210 @@ struct Qwen4ExpCore::Impl { } } + // Decode groups (NINFER_QWEN4_DECODE_GROUPS=G, at least 2): over a layer pipeline, a plain step + // splits its requests into up to G groups and queues one decode pass and the sampling of each + // group after the other without waiting, so the first stage runs a group while later stages run + // the groups before it; the host waits once for all of them. 0 or 1 keeps one pass for the batch. + [[nodiscard]] static std::size_t decode_groups() { + static const std::size_t groups = [] { + const char* value = std::getenv("NINFER_QWEN4_DECODE_GROUPS"); + return value ? std::size_t(std::clamp(std::strtol(value, nullptr, 10), 0, 8)) : 0; + }(); + return groups; + } + + // Queues the sampling of `rows` from the first rows.size() columns of `logits` on `stream` + // (the head rank's, current) into entries [base, base + rows.size()) of the sampling planes, + // the tokens copied to the host; no grammar masks or logprobs. collect_plain() reads them once + // the stream has finished. It may run on the executor's head-rank thread (decode groups). + void enqueue_plain_sample(std::span rows, std::size_t base, + cudaStream_t stream, const Tensor& head_logits) { + const std::size_t n = rows.size(); + auto* configs = static_cast(host_sample->data()); + auto* positions = reinterpret_cast(configs + slots.size()); + auto* tokens = positions + slots.size(); + for (std::size_t b = 0; b < n; ++b) { + const SampleRow& row = rows[b]; + ops::SamplingConfig config; + config.temperature = row.params->temperature; + config.top_k = row.params->top_k; + config.top_p = row.params->top_p; + config.min_p = row.params->min_p; + config.presence_penalty = row.params->presence_penalty; + config.frequency_penalty = row.params->frequency_penalty; + config.seed = row.params->seed; + config.token_counts = + row.counts ? static_cast(token_counts.p) + std::size_t(row.slot) * domain + : nullptr; + configs[base + b] = config; + positions[base + b] = static_cast(row.position); + } + auto* device_configs = static_cast(sample_config.p) + base; + auto* device_positions = static_cast(sample_position.p) + base; + auto* device_out = static_cast(sample_out.p) + base; + CUDA_CHECK(cudaMemcpyAsync(device_configs, configs + base, n * sizeof(ops::SamplingConfig), + cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(device_positions, positions + base, n * sizeof(std::int32_t), + cudaMemcpyHostToDevice, stream)); + const auto width = static_cast(n); + Tensor sampled(device_out, DType::I32, {width}); + const Tensor logical(device_positions, DType::I32, {width}); + const Tensor logits(head_logits.data, DType::BF16, {head_logits.ne[0], width}); + { + auto scope = sample_workspace->scope(); + ops::sample(logits, sampled, static_cast(domain), device_configs, logical, + ops::kSamplePurposeDecode, *sample_workspace, stream); + } + CUDA_CHECK(cudaMemcpyAsync(tokens + base, device_out, n * sizeof(std::int32_t), + cudaMemcpyDeviceToHost, stream)); + } + + // The token enqueue_plain_sample() drew for entry `index`, once the head stream is done. + [[nodiscard]] TokenId collect_plain(std::size_t index) const { + const auto* configs = static_cast(host_sample->data()); + const auto* tokens = reinterpret_cast(configs + slots.size()) + slots.size(); + const std::int32_t token = tokens[index]; + if (token == ops::kSamplerNonFiniteToken || token < 0 || std::uint32_t(token) >= domain) { + throw std::runtime_error("Qwen3.8-Flash-Next produced non-finite logits"); + } + return token; + } + + // Steps a grouped call runs per request (NINFER_QWEN4_GROUP_STEPS, default 8): each group is + // queued again as soon as its own token is back, while the other groups are still in the + // pipeline, so the stages drain once per call instead of once per step. + [[nodiscard]] static std::uint32_t group_steps() { + static const std::uint32_t steps = [] { + const char* value = std::getenv("NINFER_QWEN4_GROUP_STEPS"); + return value ? std::uint32_t(std::clamp(std::strtol(value, nullptr, 10), 1, 64)) : 8U; + }(); + return steps; + } + + struct DecodeGroup { + std::vector> requests; + std::vector rows; // the step in flight's + std::size_t base = 0; // its entries of the sampling planes + std::uint32_t steps = 0; // left in this call + bool in_flight = false; + std::uint64_t issued = 0; // steps queued + }; + std::vector group_done; // on the head stream, one per group + // Steps whose sampling the head-rank callback has queued, by group: the event above belongs + // to a step only once its count has caught up. + std::mutex group_mutex; + std::condition_variable group_cv; + std::vector group_sampled; + + // Whether a request of a decode group takes another step in this call. + [[nodiscard]] bool keeps_stepping(const std::shared_ptr& request) const { + return active(request) && request->decoding && request->feed.size() == 1 && + !request->awaiting_sample && !request->constrained && + !request->cancelled.load(std::memory_order_acquire) && + request->position + 1 <= max_context; + } + + // Queues one decode pass of the group and the sampling of its tokens. + void issue_group(DecodeGroup& group, std::size_t index) { + std::erase_if(group.requests, [&](const auto& r) { return !keeps_stepping(r); }); + if (group.requests.empty() || group.steps == 0) { + group.in_flight = false; + return; + } + auto& executor = *instance.executor; + std::vector sequences; + std::vector tokens; + group.rows.clear(); + for (const auto& request : group.requests) { + sequences.push_back(request->slot); + tokens.push_back(request->feed.front()); + group.rows.push_back(sample_row(*request)); + } + const std::uint64_t step = ++group.issued; + // The sampling follows the head on the head stream, queued by whoever queues the head (the + // head rank's thread with rank threads), so a later group's head cannot overwrite these + // logits before they are read. + executor.decode(sequences, tokens, + [this, &group, index, step](cudaStream_t stream, const Tensor& logits) { + enqueue_plain_sample(group.rows, group.base, stream, logits); + CUDA_CHECK(cudaEventRecord(group_done[index], stream)); + { + std::lock_guard lock(group_mutex); + group_sampled[index] = step; + } + group_cv.notify_all(); + }); + group.in_flight = true; + --group.steps; + } + + // A plain step in decode groups (decode_groups()); requests with grammar masks or logprobs + // take the one-pass step. + [[nodiscard]] bool grouped_step(const std::vector>& batch) { + const std::size_t count = std::min({decode_groups(), batch.size(), device.size()}); + if (count < 2) { return false; } + for (const auto& request : batch) { + if (request->constrained || request->options.execution.logprobs) { return false; } + } + auto& executor = *instance.executor; + { + RankBinding bind(device, executor.head_rank()); + while (group_done.size() < count) { + cudaEvent_t event = nullptr; + CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + group_done.push_back(event); + } + } + { + std::lock_guard lock(group_mutex); + group_sampled.assign(count, 0); + } + // Contiguous groups whose sizes differ by at most one. + std::vector groups(count); + std::size_t begin = 0; + for (std::size_t g = 0; g < count; ++g) { + const std::size_t end = begin + (batch.size() - begin) / (count - g); + groups[g].requests.assign(batch.begin() + std::ptrdiff_t(begin), + batch.begin() + std::ptrdiff_t(end)); + groups[g].base = begin; + groups[g].steps = group_steps(); + begin = end; + } + for (std::size_t g = 0; g < count; ++g) { issue_group(groups[g], g); } + // Round robin: each group's tokens are taken as soon as they are back, and the group is + // queued again behind the others. + for (bool any = true; any;) { + any = false; + for (std::size_t g = 0; g < count; ++g) { + DecodeGroup& group = groups[g]; + if (!group.in_flight) { continue; } + any = true; + { + std::unique_lock lock(group_mutex); + group_cv.wait(lock, [&] { return group_sampled[g] >= group.issued; }); + } + CUDA_CHECK(cudaEventSynchronize(group_done[g])); + for (std::size_t b = 0; b < group.requests.size(); ++b) { + const auto& request = group.requests[b]; + slots[request->slot].fed.push_back(request->feed.front()); + request->position += 1; + if (drafts > 0) { ++request->speculative.fallback_steps; } + try { + accept(request, collect_plain(group.base + b), nullptr); + next_mask(request); + } catch (...) { fail(request, std::current_exception()); } + } + issue_group(group, g); + } + } + return true; + } + // One batched decode step of the requests that feed one token, whose experts read their weights // once. A grammar row whose mask is still being computed keeps its logits for later; the step // itself does not wait for the mask. void plain_step(const std::vector>& batch) { + if (grouped_step(batch)) { return; } auto& executor = *instance.executor; std::vector sequences; std::vector tokens; @@ -1990,6 +2208,49 @@ struct Qwen4ExpCore::Impl { // accepted drafts and a correction or bonus, sampled as the request samples, so the output // follows the target's distribution), the output policy takes what it accepts of the run, and // the executor commits the tokens fed up to the last of them. + // NINFER_QWEN4_ROUND_TIMING=1: the host's wall time of an MTP round's phases (the draft with its + // synchronization, the verification until the acceptance is on the host, the commit and the + // bookkeeping after it), printed every 100 rounds. A diagnostic of pipelined rounds. + struct RoundTiming { + double draft = 0, verify = 0, accept = 0, commit = 0, total = 0; + std::uint64_t rounds = 0, tokens = 0, full = 0, guessed = 0, held = 0; + }; + RoundTiming round_timing; + static bool round_timing_enabled() { + static const bool enabled = [] { + const char* value = std::getenv("NINFER_QWEN4_ROUND_TIMING"); + return value != nullptr && value[0] == '1'; + }(); + return enabled; + } + void report_round_timing() { + RoundTiming& t = round_timing; + if (t.rounds % 100 != 0) { return; } + const double n = double(t.rounds); + std::fprintf(stderr, + "round timing | %llu rounds, %.2f tokens/round, %.0f%% whole windows, guesses " + "%llu held %llu | ms/round: total %.2f draft %.2f verify %.2f accept %.2f " + "commit %.2f\n", + static_cast(t.rounds), double(t.tokens) / n, + 100.0 * double(t.full) / n, static_cast(t.guessed), + static_cast(t.held), 1e3 * t.total / n, 1e3 * t.draft / n, + 1e3 * t.verify / n, 1e3 * t.accept / n, 1e3 * t.commit / n); + t = RoundTiming{}; + } + + // Pipeline windows (the executor's): a lone unconstrained request without logprobs, MTP drafts + // of the full width. A round whose verdict keeps its window whole with the guessed bonus + // leaves the guessed window as the next round's, already through the front ranks. + bool guessed_next = false; + std::uint32_t guessed_slot = 0; + std::vector guessed_tokens; + [[nodiscard]] bool pipelined_round(const std::vector>& batch) const { + const auto& executor = *instance.executor; + return executor.pipeline_windows() && batch.size() == 1 && lookup_ngram == 0 && + !mtp_controller && executor.options().draft_min_p == 0 && + !batch.front()->constrained && !batch.front()->options.execution.logprobs; + } + void speculative_round(const std::vector>& batch) { auto& executor = *instance.executor; const std::size_t b = batch.size(); @@ -1999,6 +2260,14 @@ struct Qwen4ExpCore::Impl { sequences.push_back(request->slot); anchors.push_back(request->feed.front()); } + const bool pipelined = pipelined_round(batch); + const bool continuing = pipelined && guessed_next && guessed_slot == batch.front()->slot && + guessed_tokens.size() == std::size_t(drafts) + 1 && + anchors.front() == guessed_tokens.front(); + if (!continuing && executor.pipeline_windows() && executor.guess_live()) { + executor.drop_guess(); + } + guessed_next = false; const auto round_start = Clock::now(); // With --adaptive-mtp the controller picks the round's steps from each request's draft // survival and the round times it has measured; a new set of requests restarts its probe. @@ -2047,7 +2316,10 @@ struct Qwen4ExpCore::Impl { } } } - if (lookups < b) { + if (continuing) { + // The guessed window: its drafts are the guess's. + std::copy(guessed_tokens.begin() + 1, guessed_tokens.end(), all_drafts.begin()); + } else if (lookups < b) { std::vector mtp_sequences; std::vector mtp_anchors; for (std::size_t j = 0; j < b; ++j) { @@ -2058,6 +2330,7 @@ struct Qwen4ExpCore::Impl { std::vector mtp_drafts(mtp_sequences.size() * drafts); std::vector mtp_extents(mtp_sequences.size(), steps); executor.draft(mtp_sequences, mtp_anchors, mtp_drafts, mtp_extents, steps); + if (round_timing_enabled()) { round_timing.draft += seconds(round_start, Clock::now()); } for (std::size_t j = 0, m = 0; j < b; ++j) { if (looked_up[j]) { continue; } std::copy_n(mtp_drafts.begin() + std::ptrdiff_t(m * drafts), drafts, @@ -2091,7 +2364,21 @@ struct Qwen4ExpCore::Impl { std::vector(proposed.begin() + std::ptrdiff_t(j * k), proposed.begin() + std::ptrdiff_t(j * k + own))); } - executor.verify(sequences, tokens); + const auto verify_start = Clock::now(); + if (continuing) { + executor.verify_guessed(tokens); + } else { + if (pipelined) { executor.guess_after_draft(); } + executor.verify(sequences, tokens); + } + // The guess of the window after this one goes to the front ranks while the head rank + // verifies this one. + std::vector guess; + if (pipelined && batch.front()->position + 2 * w <= max_context) { + guess.resize(w); + executor.take_guess(guess); + executor.verify_ahead(guess); + } RankBinding bind(device, executor.head_rank()); const cudaStream_t stream = executor.head_stream(); @@ -2215,6 +2502,7 @@ struct Qwen4ExpCore::Impl { stream)); CUDA_CHECK(cudaStreamSynchronize(stream)); instance.executor->raise_service_failure(); + const auto accept_start = Clock::now(); // The output policy takes its part of each licensed run; the executor then keeps the tokens // fed up to the last one taken. @@ -2298,7 +2586,42 @@ struct Qwen4ExpCore::Impl { ops::increment_token_counts(ids, row_counts, stream); } } + const auto commit_start = Clock::now(); executor.commit(sequences, kept); + if (!guess.empty()) { + // The guess holds when the request took the whole window and its bonus was guessed. + // NINFER_QWEN4_PIPELINE_FORCE_MISS=1 drops every guess (a test of the shadow states). + static const bool force_miss = [] { + const char* value = std::getenv("NINFER_QWEN4_PIPELINE_FORCE_MISS"); + return value != nullptr && value[0] == '1'; + }(); + const std::int32_t licensed_n = host[b * w]; + const bool hold = !force_miss && !failed.front() && active(batch.front()) && + kept.front() == w && licensed_n == columns && + host[w - 1] == guess.front(); + if (hold) { + guessed_next = true; + guessed_slot = batch.front()->slot; + guessed_tokens = std::move(guess); + if (round_timing_enabled()) { round_timing.held += 1; } + } else { + executor.drop_guess(); + } + if (round_timing_enabled()) { round_timing.guessed += 1; } + } + if (round_timing_enabled()) { + const auto now = Clock::now(); + round_timing.verify += seconds(verify_start, accept_start); + round_timing.accept += seconds(accept_start, commit_start); + round_timing.commit += seconds(commit_start, now); + round_timing.total += seconds(round_start, now); + round_timing.rounds += 1; + for (std::size_t j = 0; j < b; ++j) { + round_timing.tokens += kept[j]; + if (kept[j] == w) { round_timing.full += 1; } + } + report_round_timing(); + } // The controller measures rounds of MTP drafts only: a lookup changes the round's width. if (mtp_controller && lookups == 0) { mtp_controller->observe_execution(static_cast(b), steps, From 5d530b501a9297b0a9d02f22efe8ebde7f6f4d6b Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Mon, 19 Oct 2026 06:08:49 +0300 Subject: [PATCH 2/8] feat(qwen4_exp): sequences share one kv pool below a window each An explicit --kv-capacity of at least --max-context but below a window per sequence makes Flash-Next's sequences share one paged KV pool: one set of planes per sparse-attention layer and the MTP layer, a block table per sequence filled as it grows, pages returned at reset and taken from idle sequences when short. Prompts wait while the running requests hold the pages they need; a decode step without pages finishes at context capacity. Pipeline windows now serve whichever request decodes alone. Also: a threaded verification of several sequences hashed every part from the first one's n-gram context; rank threads run without CUDA Graphs again, since the draft/verify ordering fix removed the warm-up's illegal access. --- docs/qwen3-8-flash-next.md | 19 +- src/models/qwen4_exp/executor.cpp | 247 ++++++++++++++++++++++---- src/models/qwen4_exp/executor.h | 19 ++ src/runtime/engine/qwen4_exp_core.cpp | 87 ++++++++- 4 files changed, 331 insertions(+), 41 deletions(-) diff --git a/docs/qwen3-8-flash-next.md b/docs/qwen3-8-flash-next.md index 9201446d0b..bb77b809f0 100644 --- a/docs/qwen3-8-flash-next.md +++ b/docs/qwen3-8-flash-next.md @@ -451,6 +451,19 @@ one round and the streams keep their share during a long one; each round is one whose experts read their weights once for the whole batch, so the batch costs little more than one token while the experts dominate the step. +An explicit `--kv-capacity` of at least `--max-context` but below `--max-concurrency` windows makes +the sequences share one KV pool of that many positions instead: every sparse-attention layer (and +the MTP layer) keeps one set of KV planes, and a sequence's block table maps its positions' pages +into them, taken as it grows and returned when its sequence is reset. So four sequences of 262,144 +tokens can share 262,144 positions: one long request and short ones beside it, or several of any +length while they fit together, each up to `--max-context`. A request takes pages for its prompt +when it is admitted and for a few rounds ahead before each decode step; when the pool is short, a +free sequence gives up its pages (and its cached prefix), least recently used first. A prompt that +does not fit beside the running requests waits in the queue; a decoding request that finds no +pages left finishes at its context capacity. A sequence still costs its own pooled keys (1.5 KiB a +position of `--max-context`) and recurrent state: about 0.45 GiB a sequence at 262,144 tokens. The pool keeps no host images, so it needs +`--host-state-slots 0 --host-kv-mib 0`. + With the context cache on (the default), a sequence keeps its state when its request ends, and a snapshot of its recurrent state where the prompt's last user turn closes (or at the prompt's end when the template marks no turn), 74 MiB on the device: a later prompt that continues what the @@ -638,8 +651,8 @@ draft-chain graph and each width its own verification graph, captured at their s Between GPUs without peer access (every GeForce pair under Windows/WDDM), a stream that waits on another device's event, or a staged peer copy, blocks the host thread that queues it until the other device gets there. One host thread queuing every stage therefore walks each pass through the -pipeline at device speed, and the stages take turns. Three opt-in modes, for device experts with -CUDA Graphs, change that: +pipeline at device speed, and the stages take turns. Three opt-in modes, for device experts, change +that: - **Rank threads** (`NINFER_QWEN4_RANK_THREADS=1`). The host thread prepares a pass's inputs (positions, n-gram rows) into a staging entry and posts it; each stage's own thread uploads them, @@ -669,7 +682,7 @@ CUDA Graphs, change that: others get a turn (default 8). `NINFER_QWEN4_ROUND_TIMING=1` prints the host's time per MTP round phase every 100 rounds (with -the guesses made and held), and for host-ordered verifications each stage's device time in a window. Without CUDA Graphs the rank threads stay off. +the guesses made and held), and for host-ordered verifications each stage's device time in a window. ## Measurements diff --git a/src/models/qwen4_exp/executor.cpp b/src/models/qwen4_exp/executor.cpp index aef247e72d..06765a6f4d 100644 --- a/src/models/qwen4_exp/executor.cpp +++ b/src/models/qwen4_exp/executor.cpp @@ -451,6 +451,9 @@ struct SequenceState { }; // The captured record strides and kernels depend on the actual verification width. std::array verify; + // Pipeline windows: verifications on the rank threads (every segment) and guessed windows (the + // front stages' segments), per record set and width. + std::array, 2> pipelined_verify, pipelined_ahead; // Draft chains by step count (adaptive MTP drafts fewer than draft_tokens). struct DraftGraph { DecodeGraphExecutable chain; @@ -789,6 +792,14 @@ struct Executor::Impl { domain = static_cast(model.resources().public_token_count); } pages = (options.max_context + kPagedKVPageSize - 1) / kPagedKVPageSize; + if (options.kv_pool_tokens != 0) { + const std::uint64_t asked = + (std::uint64_t(options.kv_pool_tokens) + kPagedKVPageSize - 1) / kPagedKVPageSize; + // A pool that holds every sequence's whole window is no pool. + if (asked < std::uint64_t(pages) * options.sequences) { + pool_pages = static_cast(asked); + } + } pooled_slots = options.max_context / config.indexer_compress_ratio + 1; kv_layout = paged_kv_storage_layout(options.kv_cache, static_cast(config.head_dim)); @@ -803,6 +814,12 @@ struct Executor::Impl { stream_width = width; } } + if (stream_slots != 0 && pool_pages != 0) { + // Streamed layers keep each sequence's own host copy of its KV. + std::fprintf(stderr, "qwen4_exp: streamed KV (--kv-resident) keeps a window per sequence; " + "no shared KV pool\n"); + pool_pages = 0; + } memory.ranks.resize(device.size()); plan_weights(); if (verify_width > 0) { plan_mtp(); } @@ -1228,8 +1245,15 @@ struct Executor::Impl { const std::uint32_t slots = options.sequences + (pipeline ? 1 : 0); std::vector identity(std::max(pages, slots)); std::iota(identity.begin(), identity.end(), 0); - rank.block_table = DeviceBuffer(std::size_t(pages) * 4); - rank.block_table.copy_from_host(identity.data(), std::size_t(pages) * 4); + if (pool_pages == 0) { + rank.block_table = DeviceBuffer(std::size_t(pages) * 4); + rank.block_table.copy_from_host(identity.data(), std::size_t(pages) * 4); + } else { + // A table per sequence into the shared pool, filled as its pages are taken. + const std::vector zeros(std::size_t(pages) * options.sequences, 0); + rank.block_table = DeviceBuffer(zeros.size() * 4); + rank.block_table.copy_from_host(zeros.data(), zeros.size() * 4); + } rank.slot_ids = DeviceBuffer(std::size_t(slots) * 4); rank.slot_ids.copy_from_host(identity.data(), std::size_t(slots) * 4); // ids, positions, RoPE positions, scatter columns, three-axis RoPE positions, rows. @@ -1376,13 +1400,30 @@ struct Executor::Impl { return cache; } + // A sparse-attention layer's KV planes of `count` pages in device memory (a shared pool's). + void allocate_device_kv(LayerState& state, std::size_t rank, std::uint32_t count) { + std::uint64_t device_kv = 0; + const auto planes = kv_planes(); + for (std::size_t i = 0; i < planes.size(); ++i) { + if (planes[i].first == 0) { continue; } + KvPlane& plane = state.kv[i]; + plane.bytes = kv_plane_bytes(i, kPagedKVPageSize, count); + plane.device = DeviceBuffer(plane.bytes); + plane.p = plane.device.p; + device_kv += plane.bytes; + } + memory.ranks[rank].state_bytes += device_kv; + memory.kv_bytes += device_kv; + } + // A sparse-attention layer's per-sequence state: paged KV (one plane per K or V vector and // per scale, each [leading, page, head, page]), pooled keys and the indexer's tail. A // streamed layer keeps the paged KV in pinned host memory, and its pool and map on the device. void allocate_attention_state(LayerState& state, std::size_t rank) { std::uint64_t device_kv = 0, bytes = 0; const auto planes = kv_planes(); - for (std::size_t i = 0; i < planes.size(); ++i) { + // A shared KV pool keeps the planes once, in pool_kv. + for (std::size_t i = 0; pool_pages == 0 && i < planes.size(); ++i) { if (planes[i].first == 0) { continue; } KvPlane& plane = state.kv[i]; plane.bytes = kv_plane_bytes(i, kPagedKVPageSize, pages); @@ -1446,7 +1487,6 @@ struct Executor::Impl { CUDA_CHECK(cudaMemcpyAsync(rank.kv_staging[i].p, plane.p, plane.bytes / pages * used, cudaMemcpyDefault, rank.stream)); } - rank.kv_staged = &state; } // Each rank's Gated DeltaNet layers keep their state in one pool with a slot per sequence, @@ -1481,9 +1521,96 @@ struct Executor::Impl { } } + // ---- Shared KV pool (ExecutorOptions::kv_pool_tokens) ---- + // One set of KV planes per sparse-attention layer (and the MTP layer) of pool_pages pages; a + // sequence's block table maps its positions' pages into them, taken from the free list as the + // sequence grows and returned at reset. The host keeps the tables in pinned memory and copies + // a sequence's new entries to every rank's table on that rank's stream before the pass that + // first writes them. + std::uint32_t pool_pages = 0; + std::vector pool_kv; // by layer; index layers.size() is the MTP layer + std::vector free_pages; + std::vector> page_of; // by sequence + std::unique_ptr table_host; // [sequences][pages] + + void allocate_pool() { + pool_kv.resize(layers.size() + 1); + for (std::size_t i = 0; i < layers.size(); ++i) { + if (!layers[i].qsa) { continue; } + RankBinding bind(device, layers[i].rank); + allocate_device_kv(pool_kv[i], layers[i].rank, pool_pages); + } + if (mtp) { + RankBinding bind(device, model.head_rank()); + allocate_device_kv(pool_kv[layers.size()], model.head_rank(), pool_pages); + } + free_pages.resize(pool_pages); + // Taken from the back: the lowest pages first. + for (std::uint32_t p = 0; p < pool_pages; ++p) { + free_pages[p] = static_cast(pool_pages - 1 - p); + } + page_of.assign(options.sequences, {}); + table_host = std::make_unique(std::size_t(pages) * options.sequences * 4); + std::memset(table_host->data(), 0, table_host->size()); + std::fprintf(stderr, "qwen4_exp: %u sequences share a KV pool of %llu tokens\n", + options.sequences, + static_cast(pool_pages) * kPagedKVPageSize); + } + + // Pages for `sequence`'s positions up to `tokens` (at most max_context): true when it holds + // them now; false, taking none, when the pool lacks some. + bool reserve_pages(const SequenceState& sequence, std::uint64_t tokens) { + if (pool_pages == 0) { return true; } + const std::uint32_t need = static_cast(std::min( + (tokens + kPagedKVPageSize - 1) / kPagedKVPageSize, pages)); + auto& held = page_of[sequence.slot]; + if (need <= held.size()) { return true; } + const std::size_t more = need - held.size(); + if (free_pages.size() < more) { return false; } + auto* table = static_cast(table_host->data()) + + std::size_t(sequence.slot) * pages; + const std::size_t first = held.size(); + for (std::size_t k = 0; k < more; ++k) { + held.push_back(free_pages.back()); + free_pages.pop_back(); + table[first + k] = held.back(); + } + for (RankState& rank : ranks) { + RankBinding bind(device, rank.rank); + CUDA_CHECK(cudaMemcpyAsync(static_cast(rank.block_table.p) + + std::size_t(sequence.slot) * pages + first, + table + first, more * 4, cudaMemcpyHostToDevice, + rank.stream)); + } + return true; + } + + void require_pages(const SequenceState& sequence, std::uint64_t tokens) { + if (!reserve_pages(sequence, tokens)) { + throw KvPoolExhausted("qwen4_exp: the shared KV pool has no pages left for sequence " + + std::to_string(sequence.slot)); + } + } + + // Returns a sequence's pages; the caller has waited for every pass that reads them. + void release_pages(const SequenceState& sequence) { + if (pool_pages == 0) { return; } + auto& held = page_of[sequence.slot]; + free_pages.insert(free_pages.end(), held.rbegin(), held.rend()); + held.clear(); + } + + [[nodiscard]] std::uint64_t free_pool_tokens() const noexcept { + return std::uint64_t(free_pages.size()) * kPagedKVPageSize; + } + [[nodiscard]] std::uint64_t held_pool_tokens(std::uint32_t s) const noexcept { + return pool_pages == 0 ? 0 : std::uint64_t(page_of.at(s).size()) * kPagedKVPageSize; + } + void allocate_sequences() { const std::uint64_t width = std::uint64_t(config.hc_count) * config.hidden_size; allocate_gdn_pools(); + if (pool_pages != 0) { allocate_pool(); } for (std::uint32_t s = 0; s < options.sequences; ++s) { SequenceState sequence; sequence.slot = s; @@ -1555,7 +1682,7 @@ struct Executor::Impl { return s.rank != model.head_rank(); }); return on("NINFER_QWEN4_PIPELINE_WINDOWS") && on("NINFER_QWEN4_RANK_THREADS") && split && - verify_width > 0 && options.sequences == 1 && options.cuda_graphs && + verify_width > 0 && options.cuda_graphs && model.options().experts == ExpertResidency::Device; } @@ -2186,6 +2313,11 @@ struct Executor::Impl { reset_pipeline(); drain(); auto& sequence = sequences.at(s); + if (pool_pages != 0 && !page_of[s].empty()) { + // Another sequence may take the pages (and their table entries) next. + synchronize_ranks(); + release_pages(sequence); + } if (std::find(verified.begin(), verified.end(), &sequence) != verified.end()) { verified.clear(); } sequence.position = 0; sequence.context = NgramContext::sequence_start(ngram, config.eos_token_id); @@ -2308,9 +2440,6 @@ struct Executor::Impl { if (model.options().experts != ExpertResidency::Device || hybrid || misses || stream) { return; } - // Without CUDA Graphs the threads' eager decode steps hit an illegal address during the - // warm-up (Windows, 4 stages, unresolved), so that mode keeps the host-ordered passes. - if (!options.cuda_graphs) { return; } rank_threads = true; const std::size_t ring = ranks.front().staging.size(); const std::size_t capacity = @@ -2678,6 +2807,10 @@ struct Executor::Impl { std::span contexts, std::uint32_t offset, std::size_t end, std::size_t first_segment = 0) { + for (const Part& part : parts) { + require_pages(*part.sequence, std::uint64_t(part.sequence->position) + offset + + std::uint64_t(part.count)); + } const std::size_t ring = ranks.front().staging.size(); auto task = new_task(); const std::size_t slot = ++staged_passes % ring; @@ -3100,7 +3233,8 @@ struct Executor::Impl { return in; } - void run_qsa(const LayerPlan& plan, LayerState& state, RankState& rank, const Part& part, + void run_qsa(const LayerPlan& plan, std::size_t kv_index, LayerState& state, RankState& rank, + const Part& part, const AttentionInputs& in) { const QsaPlan& a = *plan.qsa; WorkspaceArena& ws = *rank.workspace; @@ -3141,14 +3275,18 @@ struct Executor::Impl { const Tensor rope(const_cast(in.rope), DType::I32, {t}); ops::rmsnorm_rope(rope, a.query_norm, a.key_norm, q3, k3, qo, ko, s); } - // The paged KV: device pages, or a streamed layer's host copy. A streamed layer's call of - // at most stream_width tokens reads its device pool once kv_stream_resolve has made the - // call's blocks resident there; a wider call (a prompt chunk) reads the rank's staging - // copy, into which it appends as well, and then drops the blocks it wrote from the pool. + // The paged KV: device pages (a shared KV pool's through the sequence's block table), or a + // streamed layer's host copy. A streamed layer's call of at most stream_width tokens reads its + // device pool once kv_stream_resolve has made the call's blocks resident there; a wider call + // (a prompt chunk) reads the rank's staging copy, into which it appends as well, and then drops + // the blocks it wrote from the pool. const auto paged_count = static_cast(pages); - const PagedKVLayerView cache = - kv_view([&](std::size_t i) { return state.kv[i].p; }, kPagedKVPageSize, paged_count, - rank.block_table.p, paged_count); + const LayerState& planes_of = pool_pages != 0 ? pool_kv[kv_index] : state; + auto* const table = static_cast(rank.block_table.p) + + (pool_pages != 0 ? std::size_t(part.sequence->slot) * pages : 0); + const PagedKVLayerView cache = kv_view( + [&](std::size_t i) { return planes_of.kv[i].p; }, kPagedKVPageSize, + static_cast(pool_pages != 0 ? pool_pages : pages), table, paged_count); const bool streamed = state.map.table != nullptr; const bool pooled_read = streamed && t <= static_cast(stream_width); const Tensor v3(v.data, DType::BF16, {d, nk, t}); @@ -3492,7 +3630,7 @@ struct Executor::Impl { if (plan.gdn) { run_gdn(plan, i, state, rank, part); } else { - run_qsa(plan, state, rank, part, target_inputs(i, state, rank, part)); + run_qsa(plan, i, state, rank, part, target_inputs(i, state, rank, part)); } } { @@ -3571,6 +3709,7 @@ struct Executor::Impl { if (part.sequence->position + std::uint32_t(part.count) > options.max_context) { throw std::invalid_argument("qwen4_exp forward: the sequence exceeds max_context"); } + require_pages(*part.sequence, part.sequence->position + std::uint32_t(part.count)); } settle_routes(); // A copy queued for an earlier pass may predate the cache's last swaps. @@ -3807,6 +3946,7 @@ struct Executor::Impl { if (start + n > options.max_context) { throw std::invalid_argument("qwen4_exp forward: the sequence exceeds max_context"); } + require_pages(sequence, std::uint64_t(start) + n); const std::size_t prompt = sequence.media_rope.size() / 3; if (prompt != 0 && start < prompt) { throw std::invalid_argument("qwen4_exp forward: a media prompt prefills chunk by chunk"); @@ -3920,7 +4060,7 @@ struct Executor::Impl { // The span holds the sequence at its start; chunk k begins k chunks on. AttentionInputs in = target_inputs(i, state, rank, part); in.first = start + std::uint32_t(k * c); - run_qsa(plan, state, rank, part, in); + run_qsa(plan, i, state, rank, part, in); } { auto scope = rank.workspace->scope(); @@ -4285,7 +4425,7 @@ struct Executor::Impl { for (const MtpPart& part : parts) { SequenceState& sequence = *part.sequence; const Part cells{.sequence = &sequence, .column = part.column, .count = part.count}; - run_qsa(plan, sequence.mtp, rank, cells, + run_qsa(plan, layers.size(), sequence.mtp, rank, cells, {.positions = positions + part.column, .rope = positions + part.column, .index = nullptr, @@ -4465,8 +4605,13 @@ struct Executor::Impl { } parts.push_back({drafting[j], std::int32_t(j), 1, first}); } + draft_sequence = drafting.front(); draft_first = drafting.front()->position - 1; draft_steps = static_cast(k); + // The MTP layer writes its cells' KV at the steps' positions. + for (const SequenceState* sequence : drafting) { + require_pages(*sequence, std::uint64_t(sequence->position) + staged); + } if (hybrid) { hybrid->begin(cancelled, !warming); } stage_mtp(anchors, 0, positions, 0); if (draft_prefetch) { @@ -4757,11 +4902,13 @@ struct Executor::Impl { std::size_t head_first = 0; // the first segment of the head rank std::uint32_t set_flip = 0; // the record set of the last window struct Window { // the last window posted to the front ranks + SequenceState* sequence = nullptr; std::uint32_t set = 0, width = 0, position = 0; std::vector tokens; NgramContext context; // the n-gram context before it } window; struct Ahead { // the guessed window the front ranks run + SequenceState* sequence = nullptr; bool live = false; std::uint64_t number = 0; std::uint32_t set = 0, width = 0, position = 0; @@ -4772,9 +4919,7 @@ struct Executor::Impl { std::uint64_t guess_posted = 0; // the host's count of guess chains std::uint64_t guess_recorded = 0; // the head thread's (worker_mutex) std::uint32_t draft_first = 0, draft_steps = 0; // the last draft chain - // Graphs per record set and width: verification (every segment) and guessed windows (the - // front ranks' segments), for the one sequence. - std::array, 2> verify_graphs, ahead_graphs; + SequenceState* draft_sequence = nullptr; static constexpr std::size_t kDiscardCell = 512; // where forced chain steps write their argmax void start_pipeline() { @@ -4809,19 +4954,25 @@ struct Executor::Impl { // A verification posted to the rank threads (the host's pass_mode describes it). void threaded_verify(std::span parts, std::span tokens) { SequenceState& sequence = *parts.front().sequence; - const auto w = static_cast(tokens.size()); - NgramContext context = sequence.context; - NgramContext* contexts[] = {&context}; - auto task = stage_task(parts, tokens, w, contexts, 0, segments.size()); + const auto w = static_cast(tokens.size() / parts.size()); + // Each part hashes from a copy of its context: the commit advances the live one. + std::vector copies; + for (const Part& part : parts) { copies.push_back(part.sequence->context); } + std::vector contexts; + for (NgramContext& copy : copies) { contexts.push_back(©); } + auto task = stage_task(parts, tokens, static_cast(tokens.size()), contexts, + 0, segments.size()); task->mode = pass_mode; - auto& graphs = verify_graphs[pass_mode.set][w]; - if (graphs.segments.empty() && graphs.runs++ > 0) { + auto& graphs = sequence.pipelined_verify[pass_mode.set][w]; + if (parts.size() == 1 && graphs.segments.empty() && graphs.runs++ > 0) { drain(); graphs.segments = capture_segments(sequence, std::int32_t(w), 0, segments.size()); } - if (!graphs.segments.empty()) { task->graphs = &graphs.segments; } + if (parts.size() == 1 && !graphs.segments.empty()) { task->graphs = &graphs.segments; } post(task); - window = {.set = pass_mode.set, + if (parts.size() != 1) { return; } + window = {.sequence = &sequence, + .set = pass_mode.set, .width = w, .position = sequence.position, .tokens = std::vector(tokens.begin(), tokens.end()), @@ -4873,7 +5024,7 @@ struct Executor::Impl { // the verification's stack). void guess_after_draft() { const std::uint64_t number = ++guess_posted; - SequenceState* sequence = &sequences.front(); + SequenceState* sequence = draft_sequence; const std::uint32_t k = draft_steps, first = draft_first; post_head_work([=, this] { run_guess_chain(sequence, first, k, 2 * k, 0, number); }); } @@ -4892,7 +5043,8 @@ struct Executor::Impl { // The front ranks run `tokens` as the window after the one in flight, taken whole. void verify_ahead(std::span tokens) { - SequenceState& sequence = sequences.front(); + if (window.sequence == nullptr) { throw std::logic_error("qwen4_exp: no window in flight"); } + SequenceState& sequence = *window.sequence; const auto w = static_cast(tokens.size()); NgramContext start = window.context; if (table) { @@ -4921,7 +5073,7 @@ struct Executor::Impl { const std::uint32_t kept[] = {width}; commit_rank(r, of, kept, source, width, true); }; - auto& graphs = ahead_graphs[set][w]; + auto& graphs = sequence.pipelined_ahead[set][w]; if (graphs.segments.empty() && graphs.runs++ > 0) { drain(); const PassMode host = pass_mode; @@ -4936,13 +5088,15 @@ struct Executor::Impl { } if (!graphs.segments.empty()) { task->graphs = &graphs.segments; } post(task); - ahead = {.live = true, + ahead = {.sequence = &sequence, + .live = true, .number = task->ahead_number, .set = set, .width = w, .position = position, .tokens = std::vector(tokens.begin(), tokens.end())}; - window = {.set = set, + window = {.sequence = &sequence, + .set = set, .width = w, .position = position, .tokens = ahead.tokens, @@ -4967,7 +5121,8 @@ struct Executor::Impl { // guessed window's tokens, then past them), then completes the window from the front ranks' // stack. void verify_guessed(std::span tokens) { - SequenceState& sequence = sequences.front(); + if (ahead.sequence == nullptr) { throw std::logic_error("qwen4_exp: no guessed window"); } + SequenceState& sequence = *ahead.sequence; const auto w = static_cast(tokens.size()); if (!ahead.live || w != ahead.width || sequence.position != ahead.position || !std::equal(tokens.begin(), tokens.end(), ahead.tokens.begin())) { @@ -4980,6 +5135,8 @@ struct Executor::Impl { verified_set = ahead.set; const std::uint32_t k = w - 1; const std::uint32_t first = sequence.position - 1; + // The chain's MTP cells reach position first + 2k. + require_pages(sequence, std::uint64_t(first) + 2 * k + 1); const std::uint64_t number = ++guess_posted; std::vector forced(tokens.begin(), tokens.end()); std::vector positions(2 * k + 1); @@ -4997,7 +5154,7 @@ struct Executor::Impl { task->mode = {.verify = true, .ahead = false, .set = ahead.set, .rows = 1, .width = w}; task->from_ahead = true; task->ahead_number = ahead.number; - auto& graphs = verify_graphs[ahead.set][w]; + auto& graphs = sequence.pipelined_verify[ahead.set][w]; if (!graphs.segments.empty()) { task->graphs = &graphs.segments; } post(task); ahead.live = false; @@ -5113,6 +5270,10 @@ struct Executor::Impl { [[nodiscard]] std::vector image_regions(const SequenceState& sequence, const SequenceSnapshot* at, std::uint32_t position) const { + if (pool_pages != 0) { + throw std::invalid_argument("qwen4_exp: sequence images need per-sequence KV windows, " + "not a shared KV pool"); + } std::vector out; const std::uint64_t used = (std::uint64_t(position) + kPagedKVPageSize - 1) / kPagedKVPageSize; @@ -5354,6 +5515,18 @@ void Executor::commit(std::span sequences, impl_->commit(sequences, columns); } +bool Executor::kv_pooled() const noexcept { return impl_->pool_pages != 0; } + +std::uint64_t Executor::kv_free_tokens() const noexcept { return impl_->free_pool_tokens(); } + +std::uint64_t Executor::kv_held_tokens(std::uint32_t sequence) const noexcept { + return impl_->held_pool_tokens(sequence); +} + +bool Executor::reserve_kv(std::uint32_t sequence, std::uint32_t tokens) { + return impl_->reserve_pages(impl_->sequences.at(sequence), tokens); +} + bool Executor::pipeline_windows() const noexcept { return impl_->pipeline; } void Executor::guess_after_draft() { impl_->guess_after_draft(); } diff --git a/src/models/qwen4_exp/executor.h b/src/models/qwen4_exp/executor.h index 5f45f7f891..34e7544959 100644 --- a/src/models/qwen4_exp/executor.h +++ b/src/models/qwen4_exp/executor.h @@ -33,15 +33,25 @@ #include #include #include +#include #include namespace ninfer::models::qwen4_exp { +// A pass needed pages a shared KV pool did not have; nothing of it was queued. +struct KvPoolExhausted : std::runtime_error { + using std::runtime_error::runtime_error; +}; + struct ExecutorOptions { std::uint32_t max_context = 32768; // Tokens per forward call: a prompt is fed in calls of at most this many. std::uint32_t prefill_chunk = 2048; std::uint32_t sequences = 1; + // A KV pool shared by the sequences (tokens, page-aligned up): each sequence takes pages from + // it as its positions grow, up to max_context, and returns them at reset(). Zero gives every + // sequence its own max_context window. + std::uint32_t kv_pool_tokens = 0; // The n-gram table of the PLE layer. Empty runs the model without it: the PLE injection is // skipped, which is what an all-zero table would give. std::optional ngram; @@ -231,6 +241,15 @@ class Executor { // Pipeline windows (NINFER_QWEN4_PIPELINE_WINDOWS=1 with rank threads, one sequence): while // the head rank verifies a window, the front ranks run the guessed next one. + // A shared KV pool (ExecutorOptions::kv_pool_tokens): the tokens its free pages hold, and the + // tokens of pages `sequence` holds. reserve_kv() takes pages for `sequence` up to `tokens` + // positions and returns false, taking none, when the pool lacks them; a pass past a sequence's + // reserved positions throws KvPoolExhausted. + [[nodiscard]] bool kv_pooled() const noexcept; + [[nodiscard]] std::uint64_t kv_free_tokens() const noexcept; + [[nodiscard]] std::uint64_t kv_held_tokens(std::uint32_t sequence) const noexcept; + bool reserve_kv(std::uint32_t sequence, std::uint32_t tokens); + [[nodiscard]] bool pipeline_windows() const noexcept; // After draft() and before verify(): the MTP chain runs on past the drafts to the guess. void guess_after_draft(); diff --git a/src/runtime/engine/qwen4_exp_core.cpp b/src/runtime/engine/qwen4_exp_core.cpp index be73e64488..3202abafb5 100644 --- a/src/runtime/engine/qwen4_exp_core.cpp +++ b/src/runtime/engine/qwen4_exp_core.cpp @@ -241,6 +241,23 @@ ConstructedQwen4Exp construct_qwen4_exp(const EngineOptions& options, DeviceCont models::qwen4_exp::ExecutorOptions executor; executor.max_context = options.max_context; executor.sequences = options.max_concurrency; + // An explicit --kv-capacity of at least one window but less than a window per sequence makes + // the sequences share one KV pool of that size (what the requests hold at once). + if (options.kv_capacity.mode == KvCapacityMode::Explicit && options.max_concurrency > 1 && + options.kv_capacity.explicit_tokens >= options.max_context && + std::uint64_t(options.kv_capacity.explicit_tokens) < + std::uint64_t(options.max_concurrency) * options.max_context) { + const auto& cache = options.context_cache; + if (cache.enabled && (cache.host_state_slots != 0 || cache.host_kv_capacity_bytes != 0 || + cache.host_cache_budget_bytes.value_or(0) != 0 || + cache.host_cache_auto || !cache.disk_kv_path.empty())) { + throw std::invalid_argument( + "Qwen3.8-Flash-Next's shared KV pool (--kv-capacity below max-concurrency x " + "max-context) keeps no host or disk context cache: pass --host-state-slots 0 " + "--host-kv-mib 0"); + } + executor.kv_pool_tokens = options.kv_capacity.explicit_tokens; + } executor.draft_tokens = mtp ? speculative.draft_tokens : 0; executor.draft_min_p = speculative.draft_min_p; // A chunk also holds a verification of every sequence at once. @@ -1229,6 +1246,51 @@ struct Qwen4ExpCore::Impl { return true; } + // ---- A shared KV pool (an explicit --kv-capacity below a window per sequence) ---- + // A request takes pages for its prompt when it is admitted and for a few rounds ahead before + // each decode step; an idle sequence's pages (and with them its cached prefix) go to whoever + // needs them, least recently used first. A prompt waits in the queue while running requests + // hold the pages it needs; a decoding request that finds none finishes at its context capacity. + struct KvPoolBusy {}; + + // Positions beyond a request's next feed that a step may write: a speculative round with its + // guessed window, or a decode group's steps in flight. + [[nodiscard]] std::uint32_t kv_margin() const noexcept { + return 2 * (drafts + 1) + group_steps() + 1; + } + + bool make_room(std::uint32_t s, std::uint64_t tokens) { + auto& executor = *instance.executor; + if (!executor.kv_pooled()) { return true; } + const auto want = static_cast(std::min(tokens, max_context)); + if (executor.reserve_kv(s, want)) { return true; } + std::vector idle; + for (std::uint32_t o = 0; o < slots.size(); ++o) { + if (o != s && !slots[o].request && executor.kv_held_tokens(o) > 0) { idle.push_back(o); } + } + std::sort(idle.begin(), idle.end(), + [&](std::uint32_t a, std::uint32_t b) { return slots[a].last_used < slots[b].last_used; }); + for (const std::uint32_t o : idle) { + executor.reset(o); + slots[o].fed.clear(); + slots[o].anchor.clear(); + if (executor.reserve_kv(s, want)) { return true; } + } + return false; + } + + // Whether the pool could hold `tokens` more positions for a new request: its free pages and + // the idle sequences' pages. + [[nodiscard]] bool pool_could_hold(std::uint64_t tokens) const { + const auto& executor = *instance.executor; + if (!executor.kv_pooled()) { return true; } + std::uint64_t available = executor.kv_free_tokens(); + for (std::uint32_t o = 0; o < slots.size(); ++o) { + if (!slots[o].request) { available += executor.kv_held_tokens(o); } + } + return tokens <= available; + } + // Moves queued requests into free sequences, oldest first: each into the free sequence whose // cached state serves the most of its prompt, else the least recently used. void admit() { @@ -1252,6 +1314,11 @@ struct Qwen4ExpCore::Impl { } } if (!best) { return; } + // With a shared KV pool, a prompt the running requests leave no room for waits. + if (any_active() && + !pool_could_hold(pending.front()->prompt_tokens.size() + kv_margin())) { + return; + } request = std::move(pending.front()); pending.pop_front(); } @@ -1263,6 +1330,10 @@ struct Qwen4ExpCore::Impl { } try { begin(request, *best); + } catch (const KvPoolBusy&) { + std::lock_guard lock(queue_mutex); + pending.push_front(std::move(request)); + return; } catch (...) { fail(request, std::current_exception()); } } } @@ -1315,6 +1386,11 @@ struct Qwen4ExpCore::Impl { // The snapshot's prefix is about to be overwritten. slot.anchor.clear(); } + if (!make_room(s, std::uint64_t(prompt_n) + kv_margin())) { + if (any_active()) { throw KvPoolBusy{}; } + throw RequestError(RequestErrorKind::ContextLengthExceeded, + "prepared prompt exceeds the shared KV pool"); + } slot.request = request; slot.last_used = ++use_clock; r.reused = r.prefilled = reused; @@ -1837,6 +1913,13 @@ struct Qwen4ExpCore::Impl { // A copy: finishing or failing the request clears its slot. const std::shared_ptr request = slot.request; if (!request || !request->decoding) { continue; } + // A shared KV pool: pages for the step's positions, or the request ends at its capacity. + if (executor.kv_pooled() && !request->awaiting_sample && + !make_room(request->slot, std::uint64_t(request->position) + request->feed.size() + + kv_margin())) { + finish_now(request, FinishReason::ContextCapacity); + continue; + } if (request->awaiting_sample) { // Its logits wait for its grammar mask; a cancellation need not. if (mask_ready(*request)) { @@ -2914,7 +2997,9 @@ MemorySummary Qwen4ExpCore::memory_summary() const { MemorySummary out; out.device = devices.front().device; out.max_context = impl_->max_context; - out.kv_capacity = impl_->max_context; + out.kv_capacity = impl_->instance.executor->kv_pooled() + ? impl_->instance.executor->options().kv_pool_tokens + : impl_->max_context; out.kv_cache = impl_->instance.executor->options().kv_cache; out.weights = devices.front().weights; out.sequence = devices.front().sequence; From d2cb4bd931b3d472c02aaba1a2e703b54fcc8a94 Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Mon, 19 Oct 2026 06:42:53 +0300 Subject: [PATCH 3/8] docs(qwen4_exp): shared kv pool measurements --- docs/qwen3-8-flash-next.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/qwen3-8-flash-next.md b/docs/qwen3-8-flash-next.md index bb77b809f0..8352fe1874 100644 --- a/docs/qwen3-8-flash-next.md +++ b/docs/qwen3-8-flash-next.md @@ -985,6 +985,16 @@ Four concurrent chat requests (512 tokens each, four slots of 65,536 tokens, rk8 in all with decode groups (`NINFER_QWEN4_DECODE_GROUPS=4`) against 45.7 with MTP3 rounds, A 112.4 against 40.9, each request's text unchanged; one request alone kept MTP (61.0 and 49.1 tok/s). +A shared KV pool of 262,144 positions with pipeline windows and decode groups: B (int8) with four +slots of 262,144 tokens left 819 MiB free at startup against 1.06 GiB with one slot. A chat alone +decoded at 90.7 tok/s and the 206K prompt alone at 61.7 (1.59k tok/s prefill), the one-slot server's +text both times; the 206K request with three chats beside it decoded at 40.7 tok/s while the chats +ran at 41.0-43.8. A (rk8v4) holds two such slots (four ran out of memory at startup): the 206K prompt +alone at 45.4 tok/s, the one-slot text, and beside a chat at 35.9 while the chat ran at 46.5. A chat +alone three times, one slot against the pool, same build: B 91.0/93.5/93.7 against 91.2/92.4/92.6 +tok/s (four slots), A 64.4/64.4/64.4 against 64.2/64.4/64.5 (two slots). Without decode groups four +pooled chats at once ran at about 10.5 tok/s each. + ### Hyper-connections in Q8_0 The hyper-connection projections (97 down/up pairs and 96 inject matrices) are stored as ggml Q8_0, From 2ca1117e63a8604286ec876f73c2ba2f15282d4f Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Wed, 21 Oct 2026 03:10:21 +0300 Subject: [PATCH 4/8] fix(qwen4_exp): start a shared kv pool without the image store Master pins the host image store's first pages at startup by the size of a whole-context image, which a shared KV pool refuses to give (its sequences keep no images), so a pooled server failed at startup. The pool already requires --host-state-slots 0 --host-kv-mib 0; it now runs without the store. --- src/runtime/engine/qwen4_exp_core.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/runtime/engine/qwen4_exp_core.cpp b/src/runtime/engine/qwen4_exp_core.cpp index 3202abafb5..14ef23255e 100644 --- a/src/runtime/engine/qwen4_exp_core.cpp +++ b/src/runtime/engine/qwen4_exp_core.cpp @@ -647,7 +647,9 @@ struct Qwen4ExpCore::Impl { } sample_workspace = std::make_unique(workspace); const ContextCacheOptions& cache = options.context_cache; - if (reuse_prefixes) { + // Sequences of a shared KV pool keep no images (their pages go back to the pool), so the + // pool runs without the host and disk image store. + if (reuse_prefixes && !i.executor->kv_pooled()) { // An automatic budget is sized here, once the experts are pinned. const ContextCacheOptions sized = resolve_auto_host_cache_now( cache, options.enable_vision ? std::uint64_t(options.media_cache_bytes) + From c6cc108fbc1a32bbdaa935ded65ab5709bf66e2f Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Wed, 21 Oct 2026 05:50:20 +0300 Subject: [PATCH 5/8] fix(ops): convert streamed cpu halves without the scalar f16c intrinsic MSVC does not declare _cvtsh_ss, so the Windows build stopped at streamed_cpu_avx2.cpp; the conversion takes the vector F16C form, as gguf_expert_cpu_avx2.cpp does. The change is the one build/windows-release makes, line for line. --- src/ops/sparse_moe/streamed_cpu_avx2.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ops/sparse_moe/streamed_cpu_avx2.cpp b/src/ops/sparse_moe/streamed_cpu_avx2.cpp index 6162b74e33..9345c8cedd 100644 --- a/src/ops/sparse_moe/streamed_cpu_avx2.cpp +++ b/src/ops/sparse_moe/streamed_cpu_avx2.cpp @@ -20,7 +20,10 @@ std::uint16_t load16(const std::uint8_t* p) noexcept { return value; } -float half(std::uint16_t bits) noexcept { return _cvtsh_ss(bits); } +float half(std::uint16_t bits) noexcept { + // MSVC exposes the vector F16C conversion, with the same low-lane result. + return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128(bits))); +} // Signed integer codes to weights: (code ^ bias) - bias, times the scale. __m256 weight(__m256i code, int bias, __m256 scale) noexcept { From f6e623706313ddc4d185c70b19d16e907da7d28b Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Wed, 21 Oct 2026 07:42:49 +0300 Subject: [PATCH 6/8] feat(qwen4_exp): keep a snapshot at the chunk boundary before a prompt's last message Agent clients rewrite the last message of the previous turn's prompt (Cline-family clients drop an environment block from it), which breaks the turn closure's prefix, so every such turn of a long session prefilled its whole prompt again: an API log showed 50 turns of 214-228K tokens at 0% reuse, 2 min 10 s to each first token. A prompt of 8,192 tokens or more to its anchor now also keeps a snapshot at the last chunk boundary of its prefill before its last message, and a prompt that rewrites that message resumes from it (the long-anchor reuse path). The snapshot is taken only when the prefill started at position 0 or at another such snapshot, so it lies on the chunks a prefill from position 0 runs: taking it splits no chunk, and a prompt resumed from it computes what a prefill from scratch computes. It costs a second 74 MiB snapshot per sequence; the expert cache leaves room for both. Measured on B (int8, four pooled slots) over a 67K docs prompt: the two rewritten turns resumed 65,536 tokens and answered in 3.5 and 3.9 s instead of 1 min 31 s, and every turn's text was master's byte for byte. A snapshot at the message's own start (an earlier form of this change) split a chunk, and a 206K turn resumed from it kept a cold prefill's first 85 of 128 tokens only. --- docs/qwen3-8-flash-next.md | 11 ++- src/runtime/engine/qwen4_exp_core.cpp | 104 +++++++++++++++++++++----- 2 files changed, 94 insertions(+), 21 deletions(-) diff --git a/docs/qwen3-8-flash-next.md b/docs/qwen3-8-flash-next.md index 8352fe1874..461ba97f83 100644 --- a/docs/qwen3-8-flash-next.md +++ b/docs/qwen3-8-flash-next.md @@ -469,7 +469,16 @@ snapshot of its recurrent state where the prompt's last user turn closes (or at when the template marks no turn), 74 MiB on the device: a later prompt that continues what the sequence holds resumes from its live state, and one that repeats the prompt up to that point (the next turn of a chat, which renders the previous answer without its reasoning) resumes from the -snapshot. +snapshot. A prompt of 8,192 tokens or more to that point keeps a second snapshot at the last prompt +chunk boundary before its last message starts (when that lies past half of 8,192), if its prefill +started at position 0 or at another such snapshot: agent clients rewrite that message in the next +turn's prompt (Cline-family clients drop an environment block from it), which breaks the turn +closure's prefix, and that prompt resumes from this snapshot instead ("long anchor" in the request +log). The snapshot lies on the chunks of a prefill from position 0, so taking it splits no chunk +and a prompt resumed from it computes exactly what a prefill from scratch computes; a prompt +resumed from the live state or the turn closure matches one up to rounding (its answer tokens were +decoded, and its chunks start elsewhere). Each costs another 74 MiB on the device, taken at the +sequence's first such prompt. Either way only the new tokens are prefilled; the response's prompt summary reports how many were reused. A request goes to the free sequence that holds the longest such prefix of its prompt. `--no-prefix-reuse` (ninfer-serve) prefills every prompt from scratch. diff --git a/src/runtime/engine/qwen4_exp_core.cpp b/src/runtime/engine/qwen4_exp_core.cpp index 14ef23255e..5fff8a8343 100644 --- a/src/runtime/engine/qwen4_exp_core.cpp +++ b/src/runtime/engine/qwen4_exp_core.cpp @@ -52,6 +52,11 @@ namespace { // Prompt chunks in one prefill step while other requests decode (see prefill_step). constexpr std::uint32_t kPeerSpanChunks = 2; +// Prompts from this many tokens to their anchor also keep a snapshot at the last chunk boundary of +// their prefill before their last message (Request::checkpoint_at), when it lies at least half as +// far in. +constexpr std::uint32_t kCheckpointMinPrompt = 8192; + using Clock = std::chrono::steady_clock; // Device memory kept free when the host expert cache grows after warm-up. @@ -287,9 +292,11 @@ ConstructedQwen4Exp construct_qwen4_exp(const EngineOptions& options, DeviceCont // the context cache's per-sequence snapshots take when their first prompts reach an anchor // (they are sized then); an explicit size is the cache's size. if (!options.expert_cache_bytes) { + // Two per sequence: the turn closure's and the message checkpoint's + // (Request::checkpoint_at). const std::uint64_t snapshots = options.context_cache.enabled - ? std::uint64_t(options.max_concurrency) * instance->executor->snapshot_bytes() + ? 2 * std::uint64_t(options.max_concurrency) * instance->executor->snapshot_bytes() : 0; instance->executor->grow_expert_cache(kGrowthMargin + snapshots); } @@ -430,6 +437,14 @@ struct Qwen4ExpCore::Request { // last user turn, which the next turn's prompt repeats although it renders this turn's answer // differently), else the prompt's end; and whether the prompt may be reused at all. std::uint32_t anchor_at = 0; + // The start of the prompt's last message before the anchor (0: none): agent clients often + // rewrite that message in the next turn's prompt (an environment block dropped from it). + std::uint32_t last_message_at = 0; + // Where the sequence keeps a snapshot for such a prompt (0: none): the last chunk boundary of + // the prefill at or before last_message_at, set when the prefill starts where a prefill from + // position 0 would pass (Impl::begin). Taking it cuts no chunk, and a prompt that resumes from + // it computes what a prefill from position 0 computes. + std::uint32_t checkpoint_at = 0; bool reusable = true; bool media = false; // images or video: encoded at the prompt's first chunk @@ -470,7 +485,16 @@ struct Qwen4ExpCore::Impl { std::vector fed; std::vector anchor; models::qwen4_exp::SequenceSnapshot snapshot; + // The same at that prompt's message checkpoint (Request::checkpoint_at); the slot's KV + // holds its positions while it is a prefix of `fed` no longer than the live state kept. + std::vector checkpoint; + models::qwen4_exp::SequenceSnapshot checkpoint_snapshot; std::uint64_t last_used = 0; + void forget_prefixes() { + fed.clear(); + anchor.clear(); + checkpoint.clear(); + } }; // One row of a sampling call: the request's parameters at its position. @@ -802,8 +826,7 @@ struct Qwen4ExpCore::Impl { try { instance.executor->abort(s); } catch (const FileReadError&) { instance.executor->abort(s); } - slot.fed.clear(); - slot.anchor.clear(); + slot.forget_prefixes(); complete(std::exchange(slot.request, {}), unavailable); ++failed_requests; } @@ -863,8 +886,7 @@ struct Qwen4ExpCore::Impl { // A partially advanced layer stack is discarded; it must never enter the prefix cache. try { auto request = slots.front().request; - slots.front().fed.clear(); - slots.front().anchor.clear(); + slots.front().forget_prefixes(); instance.executor->abort(0); // MTP can finish the response before its catch-up observes cancellation. // The sequence still needs cleanup, but no second terminal. @@ -931,8 +953,7 @@ struct Qwen4ExpCore::Impl { Slot& slot = slots.at(request->slot); if (slot.request == request) { slot.request = nullptr; - slot.fed.clear(); - slot.anchor.clear(); + slot.forget_prefixes(); } complete(request, std::move(error)); } @@ -949,7 +970,8 @@ struct Qwen4ExpCore::Impl { [[nodiscard]] std::uint32_t reuse(const Slot& slot, const std::vector& prompt) const { if (!reuse_prefixes) { return 0; } - return std::max(prefix_reuse(slot.fed, prompt), prefix_reuse(slot.anchor, prompt)); + return std::max({prefix_reuse(slot.fed, prompt), prefix_reuse(slot.anchor, prompt), + prefix_reuse(slot.checkpoint, prompt)}); } // ---- The context cache's store ------------------------------------------------------------ @@ -1274,8 +1296,7 @@ struct Qwen4ExpCore::Impl { [&](std::uint32_t a, std::uint32_t b) { return slots[a].last_used < slots[b].last_used; }); for (const std::uint32_t o : idle) { executor.reset(o); - slots[o].fed.clear(); - slots[o].anchor.clear(); + slots[o].forget_prefixes(); if (executor.reserve_kv(s, want)) { return true; } } return false; @@ -1356,26 +1377,39 @@ struct Qwen4ExpCore::Impl { // snapshot at the end of its last prompt when the prompt continues that, else from the // longest prefix the store keeps; what the sequence held goes to the store first. std::uint32_t reused = 0; + bool from_checkpoint = false; if (reuse_prefixes && r.reusable) { const std::uint32_t live = prefix_reuse(slot.fed, r.prompt_tokens); const std::uint32_t anchored = prefix_reuse(slot.anchor, r.prompt_tokens); + const std::uint32_t checked = prefix_reuse(slot.checkpoint, r.prompt_tokens); const auto [image, from_store] = find_stored(r.prompt_tokens); - if (from_store > std::max(live, anchored)) { + if (from_store > std::max({live, anchored, checked})) { const std::vector tokens = stored[image].tokens; const PrefixReusePath path = stored[image].path; retire(s); if (restore_stored(s, tokens)) { slot.fed = tokens; slot.anchor.clear(); + slot.checkpoint.clear(); reused = from_store; r.reuse_path = path; } - } else if (anchored > live) { + } else if (anchored > live && anchored >= checked) { retire(s); executor.restore(s, slot.snapshot); slot.fed = slot.anchor; reused = anchored; r.reuse_path = PrefixReusePath::PrivateTurnClosure; + } else if (checked > live) { + // The prompt rewrote the last message of the prompt before it. + retire(s); + executor.restore(s, slot.checkpoint_snapshot); + slot.fed = slot.checkpoint; + // The turn closure lay past the checkpoint, in positions the prompt now rewrites. + slot.anchor.clear(); + reused = checked; + from_checkpoint = true; + r.reuse_path = PrefixReusePath::PrivateLongAnchor; } else if (live > 0) { reused = live; r.reuse_path = PrefixReusePath::PrivateEndpoint; @@ -1384,9 +1418,25 @@ struct Qwen4ExpCore::Impl { if (reused == 0) { if (reuse_prefixes) { retire(s); } executor.reset(s); - slot.fed.clear(); - // The snapshot's prefix is about to be overwritten. - slot.anchor.clear(); + // The snapshots' prefixes are about to be overwritten. + slot.forget_prefixes(); + } + // The message checkpoint stays while the KV still holds its positions: a prefix of what the + // sequence resumes from, no longer than it. + if (slot.checkpoint.size() > reused || + !std::equal(slot.checkpoint.begin(), slot.checkpoint.end(), slot.fed.begin())) { + slot.checkpoint.clear(); + } + // This prompt's message checkpoint, where its prefill passes a chunk boundary that a + // prefill from position 0 passes too: from position 0 or from another such checkpoint + // (the live state, the turn closure and the store hold positions off that grid, or tokens + // decoded rather than prefilled). + r.checkpoint_at = 0; + if (reuse_prefixes && r.reusable && (reused == 0 || from_checkpoint) && + r.last_message_at > reused) { + const std::uint32_t chunk = executor.options().prefill_chunk; + const std::uint32_t at = reused + (r.last_message_at - reused) / chunk * chunk; + if (at > reused && at >= kCheckpointMinPrompt / 2) { r.checkpoint_at = at; } } if (!make_room(s, std::uint64_t(prompt_n) + kv_margin())) { if (any_active()) { throw KvPoolBusy{}; } @@ -1508,7 +1558,8 @@ struct Qwen4ExpCore::Impl { [[nodiscard]] std::uint32_t prompt_step_tokens(const Request& r) const { const auto& executor = *instance.executor; const auto prompt_n = static_cast(r.prompt_tokens.size()); - const std::uint32_t until = r.prefilled < r.anchor_at ? r.anchor_at : prompt_n; + std::uint32_t until = r.prefilled < r.anchor_at ? r.anchor_at : prompt_n; + if (r.checkpoint_at > r.prefilled && r.checkpoint_at < until) { until = r.checkpoint_at; } const std::uint32_t chunk_tokens = executor.options().prefill_chunk; std::uint32_t step = r.media ? chunk_tokens : executor.prompt_step(); const bool peers_decoding = @@ -1704,8 +1755,7 @@ struct Qwen4ExpCore::Impl { slot.request = nullptr; if (r.media) { // Its state depends on media its tokens do not identify. - slot.fed.clear(); - slot.anchor.clear(); + slot.forget_prefixes(); } complete(request, std::move(result)); } @@ -1752,6 +1802,10 @@ struct Qwen4ExpCore::Impl { executor.forward(r.slot, chunk, 1, next); slot.fed.insert(slot.fed.end(), chunk.begin(), chunk.end()); r.prefilled += n; + if (reuse_prefixes && r.reusable && r.checkpoint_at != 0 && r.prefilled == r.checkpoint_at) { + executor.snapshot(r.slot, slot.checkpoint_snapshot); + slot.checkpoint = slot.fed; + } if (reuse_prefixes && r.reusable && r.prefilled == r.anchor_at) { executor.snapshot(r.slot, slot.snapshot); slot.anchor = slot.fed; @@ -2736,8 +2790,7 @@ struct Qwen4ExpCore::Impl { // Logits for position p predict token p + 1; the executor keeps at most 512 logit rows. const std::uint32_t chunk = std::min(executor.options().prefill_chunk, 512); executor.reset(0); - slots[0].fed.clear(); - slots[0].anchor.clear(); + slots[0].forget_prefixes(); std::vector out; out.reserve(n - first_target); RankBinding bind(device, executor.head_rank()); @@ -2912,6 +2965,17 @@ Qwen4ExpCore::Submission Qwen4ExpCore::submit(models::qwen3_5::PreparedPrompt pr identity.rewrite_checkpoint->frontier < prompt_n) { request->anchor_at = identity.rewrite_checkpoint->frontier; } + // A long prompt also notes the start of its last message before the anchor. + if (request->reusable && request->anchor_at >= kCheckpointMinPrompt) { + std::uint32_t last = 0; + for (const auto& hint : data.tap_hints.hints) { + if (hint.kind == runtime::prefix_cache::TapHintKind::MessageBoundary && + hint.position > last && hint.position < request->anchor_at) { + last = hint.position; + } + } + if (last >= kCheckpointMinPrompt / 2) { request->last_message_at = last; } + } } request->prompt_summary = prompt_summary; request->prepare_seconds = prepare_seconds; From 193220f84fe579aa61d6f960aceb506ea4a088db Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Wed, 21 Oct 2026 03:51:37 +0300 Subject: [PATCH 7/8] docs(qwen4_exp): measure the windows modes on master f6ec88bb Each mode added in turn on layout A, the 206K prompt on A and B, four pooled slots beside chats, what the RTX 3060s do without the clock floor, and the message checkpoint against a cold prefill. --- docs/qwen3-8-flash-next.md | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/docs/qwen3-8-flash-next.md b/docs/qwen3-8-flash-next.md index 461ba97f83..9735f5a6d3 100644 --- a/docs/qwen3-8-flash-next.md +++ b/docs/qwen3-8-flash-next.md @@ -1004,6 +1004,36 @@ alone three times, one slot against the pool, same build: B 91.0/93.5/93.7 again tok/s (four slots), A 64.4/64.4/64.4 against 64.2/64.4/64.5 (two slots). Without decode groups four pooled chats at once ran at about 10.5 tok/s each. +On master f6ec88bb (graph segments of at most six layers), A with rk8v4, each row adding one mode to +the row before, a 33,009-token prompt and a 1,024-token chat, through `ninfer-serve`: + +| Modes | Prefill tok/s | Decode tok/s | +|---|---:|---:| +| none, MTP3, `NINFER_QWEN4_HEAD_EMBEDDING=0` | 573 | 49.9 | +| the head stage's embedding copy | 703 | 51.3 | +| rank threads | 1,644 | 51.3 | +| pipeline windows, MTP2 | 1,641 | 67.3 | +| decode groups, two slots sharing a 262,144-token pool | 1,648 | 67.2 | +| the same, a build that kept each stage's segment whole | 1,646 | 67.2 | + +Every text was the first row's. On the 206K prompt A decoded at 33.4 tok/s (510 tok/s prefill) +without the modes and 45.8 (1,754) with the last row's, B with int8 at 42.4 (699) and 61.8 (1,617, +four pooled slots), the same text in each pair; B's four pooled slots ran the 206K prompt beside +three chats at 40.7 tok/s and the chats at 40.4-42.5. Without the clock floor the RTX 3060s sank to +210 MHz between their stages' bursts: A then decoded the chat at 9.6 tok/s, each 3060's share of a +verification taking 60-77 ms instead of 8. + +The message checkpoint on the same build, B with int8 and four pooled slots: the 206K docs prompt as +a user message, then a turn that repeats it with the next user message's environment block dropped, +as Cline-family clients send it. The build before it prefilled that turn's 206,299 tokens again +(2 min 12 s to the first token); this one resumed from the checkpoint, 206,080 tokens reused, and +answered in 1.1 s. A third turn that only appended was reused whole by both (0.6 s). Against a cold +prefill of the same prompts (`--no-prefix-reuse`, top-2 logprobs), the turn resumed from the +checkpoint kept the cold run's first 85 of 128 tokens and then took the other side of a near tie +(" block" and " " 0.125 nats apart cold, equal resumed), its chosen tokens' logprobs up to 0.133 nats +from the cold run's before that; the third turn, resumed from the live state as before, kept all 128 +tokens with logprobs up to 0.153 nats apart. + ### Hyper-connections in Q8_0 The hyper-connection projections (97 down/up pairs and 96 inject matrices) are stored as ggml Q8_0, From 0d28854158503c595b1caa64e680fbff506b8c18 Mon Sep 17 00:00:00 2001 From: "A. Ugryumov" Date: Wed, 21 Oct 2026 09:02:55 +0300 Subject: [PATCH 8/8] docs(qwen4_exp): measure the modes against master as pulled, with their power Master f6ec88bb as pulled against this branch as the API runs it, on both layouts: 34 of 34 texts byte for byte (single prompts, agent turns that rewrite or append, requests at once), the speed of each kind of request, the GPU power and energy a token of each mode, and master's default prompt chunk. --- docs/qwen3-8-flash-next.md | 87 +++++++++++++++++++++++++++++++++----- 1 file changed, 77 insertions(+), 10 deletions(-) diff --git a/docs/qwen3-8-flash-next.md b/docs/qwen3-8-flash-next.md index 9735f5a6d3..93809795ea 100644 --- a/docs/qwen3-8-flash-next.md +++ b/docs/qwen3-8-flash-next.md @@ -1023,16 +1023,83 @@ three chats at 40.7 tok/s and the chats at 40.4-42.5. Without the clock floor th 210 MHz between their stages' bursts: A then decoded the chat at 9.6 tok/s, each 3060's share of a verification taking 60-77 ms instead of 8. -The message checkpoint on the same build, B with int8 and four pooled slots: the 206K docs prompt as -a user message, then a turn that repeats it with the next user message's environment block dropped, -as Cline-family clients send it. The build before it prefilled that turn's 206,299 tokens again -(2 min 12 s to the first token); this one resumed from the checkpoint, 206,080 tokens reused, and -answered in 1.1 s. A third turn that only appended was reused whole by both (0.6 s). Against a cold -prefill of the same prompts (`--no-prefix-reuse`, top-2 logprobs), the turn resumed from the -checkpoint kept the cold run's first 85 of 128 tokens and then took the other side of a near tie -(" block" and " " 0.125 nats apart cold, equal resumed), its chosen tokens' logprobs up to 0.133 nats -from the cold run's before that; the third turn, resumed from the live state as before, kept all 128 -tokens with logprobs up to 0.153 nats apart. +Against master f6ec88bb as pulled (with only the MSVC line of `build/windows-release`), run as an API +ran before these modes (MTP3, one slot), and this branch as the API runs it (rank threads, pipeline +windows, decode groups of four, MTP2, four slots sharing a 262,144-token pool on B, two on A), both +with 2,048-token chunks, greedy and without logprobs (a request with logprobs runs no pipeline +windows): on A and on B every one of 17 texts was master's byte for byte. They were the 206K prompt +(335 tokens on A, 422 on B), the 33K prompt, a code prompt and two 1,024-token chats, one at a time; +four turns of an agent session over a 67K docs prompt, the second and third dropping the +environment block of the user message before them as Cline-family clients send them, the fourth only +appending; four chats at once; and the 33K prompt beside three chats. Master's plain decode (`--spec +none`) gave its MTP3 texts on B's five single prompts. + +| A, rk8v4 | Master, MTP3 | This branch | +|---|---:|---:| +| 206K prompt: prefill / decode, tok/s | 531 / 35.1 | 1,750 / 46.0 | +| 33K prompt: prefill / decode, tok/s | 569 / 40.3 | 1,600 / 56.7 | +| code prompt, decode tok/s | 63.7 | 96.3 | +| chats of 1,024 tokens, decode tok/s | 50.0, 58.0 | 69.5, 90.0 | +| agent turns 1-4, time to the first token | 2 min 0 s, 2 min 0 s, 2 min 0 s, 0.56 s | 42 s, 4.5 s, 4.9 s, 0.54 s | +| four chats: one after another (master), two at once (this branch) | 38.6 s | 25.9 s | + +On A's two slots a chat beside the 33K prompt's prefill decoded at 12.6 tok/s until the prompt was +done (a prompt chunk between decode steps), against 37-46 beside another chat. + +What each mode costs in GPU power, one server per row adding one mode to the row before (A) or +switching one (B): the 33K prompt with one output token (prefill) and a 1,024-token chat (decode), +nvidia-smi sampling every card each 100 ms; the energy is the layout's cards' mean power over the +request's wall time, times that time, per token. Master's rows are master as pulled. + +| A | Prefill tok/s | W | J/token | Decode tok/s | W | J/token | +|---|---:|---:|---:|---:|---:|---:| +| master, plain decode | 761 | 384 | 0.505 | 34.7 | 254 | 7.41 | +| master, MTP3 | 573 | 345 | 0.603 | 48.0 | 278 | 5.89 | +| the head stage's embedding copy | 702 | 377 | 0.539 | 49.2 | 280 | 5.79 | +| rank threads | 1,653 | 549 | 0.333 | 49.0 | 274 | 5.69 | +| MTP2 | 1,650 | 550 | 0.335 | 47.2 | 258 | 5.58 | +| pipeline windows | 1,652 | 551 | 0.335 | 67.6 | 423 | 6.43 | +| decode groups, two pooled slots | 1,653 | 553 | 0.336 | 67.7 | 424 | 6.45 | +| the same without the RTX 3060 clock floor | 1,620 | 546 | 0.338 | 59.2 | 369 | 6.63 | + +| B | Prefill tok/s | W | J/token | Decode tok/s | W | J/token | +|---|---:|---:|---:|---:|---:|---:| +| master, MTP3 | 761 | 428 | 0.563 | 55.6 | 307 | 5.62 | +| rank threads, decode groups, MTP2, four pooled slots | 1,570 | 548 | 0.350 | 57.6 | 301 | 5.32 | +| the same with pipeline windows | 1,572 | 554 | 0.354 | 90.4 | 593 | 6.74 | + +Rank threads run the same passes, only at once: prefill 2.4 times as fast for 1.46 times the +power, 38% less energy a prompt token. Pipeline windows are the one mode that computes what may be +thrown away: 197 of the 437 guessed windows of A's chat held (45%), 203 of 428 on B (47%), and the +dropped ones cost 15% (A) and 27% (B) more energy a generated token for 43% and 57% more speed. +Four chats at once on B ran at 127.8 tok/s in all with decode groups (718 W, 5.62 J/token), against +53.7 one after another on master (290 W, 5.41 J/token) and 41.3 in the engine's batched rounds +without decode groups (305 W, 7.40 J/token); the decode groups' texts were the ones master gave +one at a time, the batched rounds' were not (four of four differed). Idle, A's cards drew 64 W with +the RTX 3060s held at 1,500 MHz and 48 W without; `--gpu-clock-floor` holds the floor only while +requests run. + +At master's default prompt chunk (1,024 tokens) this branch gave master's 1,024-token texts on B's +206K, code and chat prompts too, and prefilled the 206K prompt at 1,457 tok/s on B and 1,557 on A, +against 1,614 and 1,753 with 2,048-token chunks (master: 665 against 700). Master's own 206K answer +differs between the two chunk sizes; the short prompts, one chunk either way, do not. + +| B, int8 | Master, plain | Master, MTP3 | This branch | +|---|---:|---:|---:| +| 206K prompt: prefill / decode, tok/s | 760 / 37.3 | 700 / 43.8 | 1,610 / 64.7 | +| 33K prompt: prefill / decode, tok/s | 810 / 44.2 | 751 / 50.4 | 1,540 / 70.5 | +| code prompt, decode tok/s | 42.6 | 74.2 | 115.9 | +| chats of 1,024 tokens, decode tok/s | 40.9, 40.9 | 56.6, 67.8 | 90.9, 118.1 | +| agent turns 1-4, time to the first token | | 1 min 30 s, 1 min 31 s, 1 min 31 s, 0.51 s | 44 s, 3.5 s, 3.9 s, 0.44 s | +| four chats: one after another (master), at once (this branch) | | 32.8 s | 14.0 s | + +Master prefilled each of the first three agent turns from scratch. This branch resumed the second +and third from the message checkpoint at position 65,536 (2,048 tokens times 32, the boundary before +the message's start) and the fourth from the live state. The checkpoint used to sit at the message's +start, which split a prompt chunk: a 206K turn resumed from it answered in 1.1 s instead of 2 min +12 s but kept a cold prefill's first 85 of 128 tokens only, then took the other side of a 0.125-nat +tie; a turn resumed from the live state still matches a prefill from scratch only up to rounding (it +kept all 128 tokens, with logprobs up to 0.153 nats apart), as on master. ### Hyper-connections in Q8_0