diff --git a/.github/ci_config.yaml b/.github/ci_config.yaml index cda0593e3..d15fa998a 100644 --- a/.github/ci_config.yaml +++ b/.github/ci_config.yaml @@ -46,7 +46,7 @@ platforms: shm_size: 64g timeout: 3600 env: - TEST_PARAM: ['default', '--enable-paged-attn', '--enable-paged-attn --enable-graph', '--enable-paged-attn --enable-graph --attn=flash-attn'] + TEST_PARAM: ['default', '--enable-paged-attn', '--enable-paged-attn --enable-graph', '--enable-paged-attn --enable-graph --attn=flash-attn', '--enable-paged-attn --enable-graph --attn=hybrid'] stages: - name: test run: python InfiniLM/examples/bench.py --device nvidia --model=/data-aisoft/mechdancer/models/9g_8b_thinking/ --input-len=256,1024 --output-len=256,1024 --batch-size=8 diff --git a/.gitignore b/.gitignore index 149d4bc38..4101a9d88 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,9 @@ python/infinilm/lib/*.so # Vscode .vscode/ +# JetBrains +.idea/ + *.sh model_weight/ @@ -34,3 +37,4 @@ __pycache__/ *.http *.nsys-rep +dev_perf/results/ diff --git a/csrc/backends/attention_backends.hpp b/csrc/backends/attention_backends.hpp index b274aacc7..6c85734ae 100644 --- a/csrc/backends/attention_backends.hpp +++ b/csrc/backends/attention_backends.hpp @@ -8,12 +8,21 @@ namespace infinilm::backends { /** * @brief Enumeration of all supported attention backends. + * + * 各后端说明： + * - STATIC_ATTN：静态 attention，prefill/decode 走同一套实现（默认） + * - PAGED_ATTN：自研 paged-attention，prefill 为 PagedAttentionPrefill，decode 为 splitkv + * - FLASH_ATTN：FlashAttention-2（mha_varlen_fwd / mha_fwd_kvcache） + * - FLASHINFER：FlashInfer 后端 + * - HYBRID：按阶段分离路由——prefill 走 FA2 varlen，decode 走自研 paged kernel + * （见 HybridAttentionImpl） */ enum class AttentionBackend { STATIC_ATTN, PAGED_ATTN, FLASH_ATTN, FLASHINFER, + HYBRID, // prefill → FlashAttention (FA2 varlen), decode → PagedAttention Default = STATIC_ATTN }; @@ -27,6 +36,8 @@ inline std::ostream &operator<<(std::ostream &os, AttentionBackend backend) { return os << "AttentionBackend::FLASH_ATTN"; case AttentionBackend::FLASHINFER: return os << "AttentionBackend::FLASHINFER"; + case AttentionBackend::HYBRID: + return os << "AttentionBackend::HYBRID"; default: throw std::invalid_argument("infinilm::backends: invalid attention backend: " + std::to_string(static_cast(backend))); break; @@ -49,9 +60,13 @@ inline AttentionBackend parse_attention_backend(const std::string &backend) { if (backend == "flashinfer") { return AttentionBackend::FLASHINFER; } + if (backend == "hybrid") { + // "hybrid"：prefill→FA2 varlen，decode→自研 paged-attention（splitkv） + return AttentionBackend::HYBRID; + } throw std::invalid_argument( - "Invalid attention_backend: " + backend + ". Valid options are: static-attn, paged-attn, flash-attn, flashinfer"); + "Invalid attention_backend: " + backend + ". Valid options are: static-attn, paged-attn, flash-attn, flashinfer, hybrid"); } } // namespace infinilm::backends diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp index 2a7780090..457ad2193 100644 --- a/csrc/cache/kv_cache.cpp +++ b/csrc/cache/kv_cache.cpp @@ -131,8 +131,13 @@ infinicore::Tensor create_layer_kv_cache( size_t block_size = config.block_size(); infinicore::Shape kv_shape; - if (global_state::get_infinilm_config().attention_backend == backends::AttentionBackend::FLASH_ATTN) { - // FLASH_ATTN kernel expects BSHD layout + const auto cache_attn_backend = global_state::get_infinilm_config().attention_backend; + if (cache_attn_backend == backends::AttentionBackend::FLASH_ATTN || cache_attn_backend == backends::AttentionBackend::HYBRID) { + // FLASH_ATTN kernel expects BSHD layout. HYBRID uses the same layout; + // its decode paged-attention kernel reads it via strides. + // FLASH_ATTN 的 kernel 需要 BSHD 布局（块 → 块内 token → head → dim）； + // HYBRID 沿用同一布局——decode 阶段自研 paged kernel 通过 stride + // 以 BHSD 逻辑视图零拷贝读取，无需切换布局。 kv_shape = {2, num_blocks_per_layer, block_size, num_rank_k_heads, k_dim}; } else { kv_shape = {2, num_blocks_per_layer, num_rank_k_heads, block_size, k_dim}; diff --git a/csrc/engine/compiler/general_compiler.cpp b/csrc/engine/compiler/general_compiler.cpp index 84ee670d4..de6fcd1a0 100644 --- a/csrc/engine/compiler/general_compiler.cpp +++ b/csrc/engine/compiler/general_compiler.cpp @@ -23,4 +23,13 @@ GeneralCompiler::Compiled GeneralCompiler::get_compiled(const InfinilmModel::Inp return result; } +std::pair, infinicore::Tensor> +GeneralCompiler::get_sampling_compiled(size_t batch_size) { + auto result = paged_compiler_->get_sampling_compiled(batch_size); + if (result.first != nullptr) { + return result; + } + return static_batching_compiler_->get_sampling_compiled(batch_size); +} + } // namespace infinilm::engine diff --git a/csrc/engine/compiler/general_compiler.hpp b/csrc/engine/compiler/general_compiler.hpp index e8b84b5d9..b63353677 100644 --- a/csrc/engine/compiler/general_compiler.hpp +++ b/csrc/engine/compiler/general_compiler.hpp @@ -12,6 +12,9 @@ class GeneralCompiler : public GraphCompiler { Compiled get_compiled(const InfinilmModel::Input &input) override; + std::pair, infinicore::Tensor> + get_sampling_compiled(size_t batch_size) override; + private: std::unique_ptr static_batching_compiler_; std::unique_ptr paged_compiler_; diff --git a/csrc/engine/compiler/graph_compiler.hpp b/csrc/engine/compiler/graph_compiler.hpp index 5173994fd..ff1bae8b6 100644 --- a/csrc/engine/compiler/graph_compiler.hpp +++ b/csrc/engine/compiler/graph_compiler.hpp @@ -17,6 +17,16 @@ class GraphCompiler { virtual void compile() = 0; virtual Compiled get_compiled(const InfinilmModel::Input &input) = 0; + // Optional companion graph that replays captured greedy-sampling kernels + // (per-request argmax) for a decode batch size. Only valid right after a + // successful get_compiled()+run() for the same batch size, since it reads + // that graph's logits blob. Default: no captured sampling. + virtual std::pair, infinicore::Tensor> + get_sampling_compiled(size_t batch_size) { + (void)batch_size; + return {nullptr, {}}; + } + protected: std::shared_ptr model_; RankBarrier *barrier_; diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index dee3123c9..63537aea5 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -1,9 +1,15 @@ #include "paged_compiler.hpp" #include "../../global_state/global_state.hpp" +#include "../../layers/attention/backends/attention_layer.hpp" #include "../../utils.hpp" +#include + #include #include +#include +#include +#include #include #include @@ -23,6 +29,19 @@ bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_conte return has_state(forward_context.conv_state_vec) || has_state(forward_context.ssm_state_vec); } +// Wraps the per-request greedy-sampling loop as one recordable operator. +// infinicore::op::random_sample_ has no graph-recording hook of its own (it +// runs eagerly even while recording), so without this wrapper a companion +// sampling graph would record an empty op list and replay as a no-op. +class SamplingLoopOperator : public infinicore::graph::GraphOperator { +public: + explicit SamplingLoopOperator(std::function fn) : fn_(std::move(fn)) {} + void run() const override { fn_(); } + +private: + std::function fn_; +}; + } // namespace PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBarrier *barrier) @@ -76,32 +95,37 @@ void PagedCompiler::compile() { {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); set_zeros(block_tables_holder_); - auto make_decode_input = [&](size_t b) { + auto make_decode_input = [&](size_t b, size_t fake_max_seq_len, CompiledResult *cr) { InfinilmModel::Input input; - input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty( - position_id_axes > 1 - ? std::vector{position_id_axes, b} - : std::vector{b}, - infinicore::DataType::I64, infinicore::context::getDevice()); - input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); - set_zeros(input.input_ids.value()); - set_zeros(input.position_ids.value()); - set_zeros(input.total_sequence_lengths.value()); - std::vector total_sequence_lengths_vec(b, 1); - infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false); - input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); - std::vector input_offsets_vec(b + 1, 0); - for (size_t i = 0; i <= b; i++) { - input_offsets_vec[i] = i; + // The small per-step inputs live as views into two contiguous + // device buffers, so get_compiled() can refresh them with two + // packed H2D copies instead of one copy per tensor. + // pack_i64: input_ids (b) | position_ids (axes*b) | slot_mapping (b) + // pack_i32: total_seq_lens (b) | input_offsets (b+1) | cu_seqlens (b+1) + auto pack_i64 = infinicore::Tensor::empty({(position_id_axes + 2) * b}, infinicore::DataType::I64, infinicore::context::getDevice()); + auto pack_i32 = infinicore::Tensor::empty({3 * b + 2}, infinicore::DataType::I32, infinicore::context::getDevice()); + set_zeros(pack_i64); + { + std::vector init_i32(3 * b + 2, 0); + for (size_t i = 0; i < b; i++) { + init_i32[i] = 1; // total_sequence_lengths + init_i32[b + i] = i; // input_offsets + init_i32[2 * b + 1 + i] = i; // cu_seqlens + } + init_i32[2 * b] = b; // input_offsets[b] + init_i32[3 * b + 1] = b; // cu_seqlens[b] + infinicore::context::memcpyH2D(pack_i32->data(), init_i32.data(), init_i32.size() * sizeof(int32_t), false); } - infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); - infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); + input.input_ids = pack_i64->narrow({{0, 0, b}})->view({1, b}); + input.position_ids = position_id_axes > 1 + ? pack_i64->narrow({{0, b, position_id_axes * b}})->view({position_id_axes, b}) + : pack_i64->narrow({{0, b, b}}); + input.slot_mapping = pack_i64->narrow({{0, (position_id_axes + 1) * b, b}}); + input.total_sequence_lengths = pack_i32->narrow({{0, 0, b}}); + input.input_offsets = pack_i32->narrow({{0, b, b + 1}}); + input.cu_seqlens = pack_i32->narrow({{0, 2 * b + 1, b + 1}}); const size_t block_per_req = nblocks; input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1}); - input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); - set_zeros(input.slot_mapping.value()); if (has_mamba_state) { input.mamba_init_state_indices = infinicore::Tensor::empty( @@ -131,6 +155,11 @@ void PagedCompiler::compile() { input.block_tables, input.slot_mapping, }; + // Decode-kernel routing hint for the hybrid attention layer. The + // recorded op list freezes whichever kernel this selects, so the + // long-ctx capture passes a value above the routing threshold + // (1 = always the splitkv variant). + forward_context.attn_metadata.max_sequence_length = fake_max_seq_len; // Hybrid linear-attention layers read cache indices from the same // thread-local context. These tensors remain alive in CompiledResult // and are updated in place before every graph replay. @@ -139,12 +168,19 @@ void PagedCompiler::compile() { input.mamba_init_state_indices, input.mamba_final_state_indices, }; + if (cr != nullptr) { + cr->pack_i64 = pack_i64; + cr->pack_i32 = pack_i32; + cr->stage_i64.assign((position_id_axes + 2) * b, 0); + cr->stage_i32.assign(3 * b + 2, 0); + cr->position_id_axes = position_id_axes; + } return input; }; { const size_t warmup_batch_size = std::min(max_batch_size, static_cast(64)); - auto input = make_decode_input(warmup_batch_size); + auto input = make_decode_input(warmup_batch_size, 1, nullptr); model_->forward(input); infinicore::context::syncStream(); // Warmup runs the eager Marlin path and may leave per-layer lock @@ -154,27 +190,115 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); } + // Only HYBRID routes decode between two kernels by ctx length; for + // every other backend the long-ctx variant would be a bit-identical + // copy of the short one, so skip the second capture and its memory. + if (infinilm::global_state::get_infinilm_config().attention_backend == ::infinilm::backends::AttentionBackend::HYBRID) { + decode_ctx_threshold_ = layers::attention::backends::decode_fa_ctx_threshold(); + } + for (size_t b : decode_batch_sizes_) { - auto input = make_decode_input(b); + DecodeVariants variants; - barrier_->wait(); - (void)model_->forward(input); - infinicore::context::syncStream(); - // Capture must not start with stale Marlin locks from previous - // warmup/capture attempts. This reset is intentionally outside - // graph capture; the current implementation still pays a memset - // before every graph replay in get_compiled(). - model_->reset_runtime_state(); - infinicore::context::syncStream(); - infinicore::context::startGraphRecording(); - auto output = model_->forward(input); - auto graph = infinicore::context::stopGraphRecording(); - barrier_->wait(); + // fake_max_seq_len only steers the hybrid layer's decode-kernel + // routing while recording; the captured kernels read the real + // per-request lengths from pack_i32 at replay. + auto capture_variant = [&](size_t fake_max_seq_len, CompiledResult &cr) { + // Both barrier waits must stay unconditional: with TP>1 the + // ranks are threads sharing this host-side barrier, so a rank + // that fails mid-capture has to arrive at both waits anyway + // before the error propagates, or the other ranks block + // forever at mismatched barrier generations. + barrier_->wait(); + bool recording = false; + try { + auto input = make_decode_input(b, fake_max_seq_len, &cr); + + (void)model_->forward(input); + infinicore::context::syncStream(); + // Capture must not start with stale Marlin locks from previous + // warmup/capture attempts. This reset is intentionally outside + // graph capture; the current implementation still pays a memset + // before every graph replay in get_compiled(). + model_->reset_runtime_state(); + infinicore::context::syncStream(); + infinicore::context::startGraphRecording(); + recording = true; + auto output = model_->forward(input); + auto graph = infinicore::context::stopGraphRecording(); + recording = false; + + auto shared_output = std::shared_ptr( + new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); - auto shared_output = std::shared_ptr( - new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); + cr.input = std::move(input); + cr.compiled = std::make_tuple(graph, shared_output); + + // Capture the greedy-sampling kernels (one cub ArgMax + index + // cast per request, reading the decode graph's logits blob) + // into a companion graph. Greedy batches replay this instead + // of issuing ~3 kernel launches per request per step. + // Best-effort: any capture failure just falls back to the + // eager sampling loop. Limited to small batch sizes to bound + // capture time at load. + if (b <= 64) { + try { + const size_t vocab_size = output.logits->shape().back(); + auto logits2d = output.logits->view({b, vocab_size}); + auto sampling_out = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + // top_k == 1 / temperature == 0 both take the argmax branch + // in the sampling op, so the captured kernels match greedy + // semantics and never consume random_val. Captured by + // value: the operator outlives this scope and keeps both + // tensors alive. + auto run_sampling = [logits2d, sampling_out, b, vocab_size]() { + for (size_t i = 0; i < b; ++i) { + auto score = logits2d->narrow({{0, i, 1}})->view({vocab_size}); + auto out = sampling_out->narrow({{0, i, 1}})->view({}); + infinicore::op::random_sample_(out, score, 0.0f, 1.0f, 1, 0.0f); + } + }; + // Recording alone does not run the op; instantiate() warms + // the loop (settling descriptor/workspace allocations) + // before capturing it into a device graph segment. + infinicore::context::startGraphRecording(); + infinicore::context::addGraphOperator(std::make_shared(run_sampling)); + cr.sampling_graph = infinicore::context::stopGraphRecording(); + cr.sampling_out = sampling_out; + } catch (const std::exception &e) { + spdlog::warn("PagedCompiler: sampling graph capture failed for batch {}: {}", b, e.what()); + cr.sampling_graph = nullptr; + cr.sampling_out = {}; + } + } + } catch (...) { + if (recording) { + // Discard the dangling partial recording so later + // captures start from a clean state. + try { + (void)infinicore::context::stopGraphRecording(); + } catch (...) { + } + } + barrier_->wait(); + throw; + } + barrier_->wait(); + }; - compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)}; + capture_variant(1, variants.short_ctx); + if (decode_ctx_threshold_ != std::numeric_limits::max()) { + // Second capture with the FA kvcache decode kernel recorded. + // Best-effort: the short-ctx variant alone is still correct. + try { + capture_variant(decode_ctx_threshold_ + 1, variants.long_ctx); + variants.has_long = true; + } catch (const std::exception &e) { + spdlog::warn("PagedCompiler: long-ctx decode graph capture failed for batch {}: {}", b, e.what()); + variants.has_long = false; + } + } + compiled_map_decode_[b] = std::move(variants); } } } @@ -192,13 +316,80 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & if (result == compiled_map_decode_.end()) { return {nullptr, nullptr}; } - auto &graph_input = result->second.input; + auto &variants = result->second; + + const auto &rt_input_ids = input.input_ids.value(); + const auto &rt_position_ids = input.position_ids.value(); + const auto &rt_total_lens = input.total_sequence_lengths.value(); + const auto &rt_offsets = input.input_offsets.value(); + const auto &rt_cu_seqlens = input.cu_seqlens.value(); + const auto &rt_slots = input.slot_mapping.value(); + + // Pick the decode-kernel variant by the batch's longest context. + // The lengths tensor is expected on the host (graph replay + // requires a CPU int32 tensor below anyway); anything else + // conservatively routes to the short-ctx (splitkv) variant. + size_t max_ctx = 0; + if (variants.has_long + && rt_total_lens->device().getType() == infinicore::Device::Type::CPU + && rt_total_lens->dtype() == infinicore::DataType::I32 + && rt_total_lens->is_contiguous() + && rt_total_lens->shape().size() == 1 + && rt_total_lens->shape()[0] == batch_size) { + const auto *lens = reinterpret_cast(rt_total_lens->data()); + for (size_t i = 0; i < batch_size; ++i) { + max_ctx = std::max(max_ctx, static_cast(lens[i])); + } + } + auto &cr = (variants.has_long && max_ctx > decode_ctx_threshold_) + ? variants.long_ctx + : variants.short_ctx; + variants.last_served = &cr; + auto &graph_input = cr.input; + + // Fast path: pack the six small inputs host-side and refresh the + // graph inputs with two H2D copies. Host pointers only: anything + // living on a device must go through copy_from below. + const bool packable = + rt_input_ids->device().getType() == infinicore::Device::Type::CPU + && rt_position_ids->device().getType() == infinicore::Device::Type::CPU + && rt_total_lens->device().getType() == infinicore::Device::Type::CPU + && rt_offsets->device().getType() == infinicore::Device::Type::CPU + && rt_cu_seqlens->device().getType() == infinicore::Device::Type::CPU + && rt_slots->device().getType() == infinicore::Device::Type::CPU + && rt_input_ids->is_contiguous() && rt_position_ids->is_contiguous() + && rt_total_lens->is_contiguous() && rt_offsets->is_contiguous() + && rt_cu_seqlens->is_contiguous() && rt_slots->is_contiguous() + && rt_input_ids->dtype() == infinicore::DataType::I64 + && rt_position_ids->dtype() == infinicore::DataType::I64 + && rt_slots->dtype() == infinicore::DataType::I64 + && rt_total_lens->dtype() == infinicore::DataType::I32 + && rt_offsets->dtype() == infinicore::DataType::I32 + && rt_cu_seqlens->dtype() == infinicore::DataType::I32 + && rt_input_ids->numel() == batch_size + && rt_position_ids->numel() == cr.position_id_axes * batch_size + && rt_total_lens->numel() == batch_size + && rt_offsets->numel() == batch_size + 1 + && rt_cu_seqlens->numel() == batch_size + 1 + && rt_slots->numel() == batch_size; + if (packable) { + std::memcpy(cr.stage_i64.data(), rt_input_ids->data(), batch_size * sizeof(int64_t)); + std::memcpy(cr.stage_i64.data() + batch_size, rt_position_ids->data(), cr.position_id_axes * batch_size * sizeof(int64_t)); + std::memcpy(cr.stage_i64.data() + (cr.position_id_axes + 1) * batch_size, rt_slots->data(), batch_size * sizeof(int64_t)); + infinicore::context::memcpyH2D(cr.pack_i64->data(), cr.stage_i64.data(), cr.stage_i64.size() * sizeof(int64_t)); - graph_input.input_ids.value()->copy_from(input.input_ids.value()); - graph_input.position_ids.value()->copy_from(input.position_ids.value()); - graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value()); - graph_input.input_offsets.value()->copy_from(input.input_offsets.value()); - graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value()); + std::memcpy(cr.stage_i32.data(), rt_total_lens->data(), batch_size * sizeof(int32_t)); + std::memcpy(cr.stage_i32.data() + batch_size, rt_offsets->data(), (batch_size + 1) * sizeof(int32_t)); + std::memcpy(cr.stage_i32.data() + 2 * batch_size + 1, rt_cu_seqlens->data(), (batch_size + 1) * sizeof(int32_t)); + infinicore::context::memcpyH2D(cr.pack_i32->data(), cr.stage_i32.data(), cr.stage_i32.size() * sizeof(int32_t)); + } else { + graph_input.input_ids.value()->copy_from(rt_input_ids); + graph_input.position_ids.value()->copy_from(rt_position_ids); + graph_input.total_sequence_lengths.value()->copy_from(rt_total_lens); + graph_input.input_offsets.value()->copy_from(rt_offsets); + graph_input.cu_seqlens.value()->copy_from(rt_cu_seqlens); + graph_input.slot_mapping.value()->copy_from(rt_slots); + } const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); if (block_per_req > compiled_block_per_req) { @@ -208,11 +399,14 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & // Initialize only the active graph rows to -1, then overwrite the // runtime logical region. Avoid clearing the full preallocated - // holder on every decode token. + // holder on every decode token. When the runtime block table + // already covers the compiled width, the copy below overwrites + // every column and the fill is redundant. auto &graph_block_tables = graph_input.block_tables.value(); - set_minus_one_device_async(graph_block_tables); + if (block_per_req < compiled_block_per_req) { + set_minus_one_device_async(graph_block_tables); + } graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value()); - graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value()); const bool graph_has_mamba_indices = graph_input.mamba_init_state_indices.has_value() && graph_input.mamba_final_state_indices.has_value(); const bool input_has_mamba_indices = input.mamba_init_state_indices.has_value() && input.mamba_final_state_indices.has_value(); @@ -232,7 +426,7 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & // zero workspace/lock buffer shared by all Marlin layers. model_->reset_runtime_state(); - auto graph = std::get<0>(result->second.compiled); + auto graph = std::get<0>(cr.compiled); if (graph != nullptr) { const auto &runtime_seq_lens = input.total_sequence_lengths.value(); if (runtime_seq_lens->device().getType() @@ -250,7 +444,7 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & runtime_seq_lens->data()), batch_size); } - auto shared_output = std::shared_ptr(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()}); + auto shared_output = std::shared_ptr(new InfinilmModel::Output{std::get<1>(cr.compiled)->logits->resume_from_blob_()}); return std::make_tuple(graph, shared_output); } @@ -259,4 +453,27 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & } } +std::pair, infinicore::Tensor> +PagedCompiler::get_sampling_compiled(size_t batch_size) { + // Kill switch for A/B measurement: force the eager per-request sampling loop. + static const bool disabled = std::getenv("INFINILM_DISABLE_SAMPLING_GRAPH") != nullptr; + if (disabled) { + return {nullptr, {}}; + } + auto it = compiled_map_decode_.find(batch_size); + if (it == compiled_map_decode_.end()) { + return {nullptr, {}}; + } + // Serve the sampling graph of whichever decode variant get_compiled() + // most recently handed out for this batch size (they read different + // logits blobs). + const CompiledResult *cr = it->second.last_served != nullptr + ? it->second.last_served + : &it->second.short_ctx; + if (cr->sampling_graph == nullptr) { + return {nullptr, {}}; + } + return {cr->sampling_graph, cr->sampling_out}; +} + } // namespace infinilm::engine diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864d..fb4d54d1a 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -2,6 +2,7 @@ #include "graph_compiler.hpp" +#include #include namespace infinilm::engine { @@ -13,6 +14,9 @@ class PagedCompiler : public GraphCompiler { Compiled get_compiled(const InfinilmModel::Input &input) override; + std::pair, infinicore::Tensor> + get_sampling_compiled(size_t batch_size) override; + private: std::vector decode_batch_sizes_; @@ -21,11 +25,34 @@ class PagedCompiler : public GraphCompiler { struct CompiledResult { InfinilmModel::Input input; Compiled compiled; + // Small per-step inputs live as views into two contiguous device + // buffers so each replay needs only two H2D copies. Host staging is + // reused across steps. + infinicore::Tensor pack_i64; // input_ids | position_ids | slot_mapping + infinicore::Tensor pack_i32; // total_seq_lens | input_offsets | cu_seqlens + std::vector stage_i64; + std::vector stage_i32; + size_t position_id_axes = 1; + // Captured per-request argmax over the decode graph's logits blob, + // replayed for greedy batches to avoid ~3 kernel launches per request. + std::shared_ptr sampling_graph; + infinicore::Tensor sampling_out; }; - std::unordered_map< - size_t, // num_requests - CompiledResult> - compiled_map_decode_; + // Decode-kernel ctx routing: when decode_fa_ctx_threshold() is finite, + // each batch size gets a second capture whose attention layers recorded + // the FA kvcache kernel (routed by the fake max_sequence_length used at + // capture time). get_compiled() picks the variant by the runtime batch's + // max context length, which it already has on the host. + struct DecodeVariants { + CompiledResult short_ctx; + CompiledResult long_ctx; + bool has_long = false; + // Set by get_compiled(); get_sampling_compiled() serves the sampling + // graph of the variant it most recently handed out. + CompiledResult *last_served = nullptr; + }; + std::unordered_map compiled_map_decode_; + size_t decode_ctx_threshold_ = std::numeric_limits::max(); }; } // namespace infinilm::engine diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 422b5df73..23905b0fe 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -44,6 +44,28 @@ size_t max_length_from_offsets( return max_length; } +// Max entry of a per-request length tensor. Lenient by design: anything that +// is not already a one-dimensional CPU int32 tensor yields 0 (the decode +// path must never pay a device sync just to fill a routing hint). +size_t max_length_from_lengths( + const std::optional &lengths) { + if (!lengths.has_value()) { + return 0; + } + const auto &t = lengths.value(); + if (t->device().getType() != infinicore::Device::Type::CPU + || t->dtype() != infinicore::DataType::I32 + || t->shape().size() != 1) { + return 0; + } + const auto *values = reinterpret_cast(t->data()); + size_t max_length = 0; + for (size_t i = 0; i < t->shape()[0]; ++i) { + max_length = std::max(max_length, static_cast(values[i])); + } + return max_length; +} + } // namespace //------------------------------------------------------ @@ -210,7 +232,10 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { && input_ids.value()->numel() != total_sequence_lengths.value()->numel(); const size_t max_query_length = is_prefill ? max_length_from_offsets(input_offsets, "input_offsets") : 0; - const size_t max_sequence_length = is_prefill ? max_length_from_offsets(cu_seqlens, "cu_seqlens") : 0; + // Decode steps fill max_sequence_length too: the hybrid attention layer + // uses it to route long-context decode to FA's kvcache kernel. It is a + // host-side max over an already-CPU tensor, so this costs nothing. + const size_t max_sequence_length = is_prefill ? max_length_from_offsets(cu_seqlens, "cu_seqlens") : max_length_from_lengths(total_sequence_lengths); // MACA maps a registered user pointer to only one node. Serialize H2D // copies so TP ranks never access the same host registration concurrently. diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 0fa0a84cf..77677d22e 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -3,6 +3,7 @@ #include "infinicore/ops.hpp" #include "infinicore/ops/distributed/send_recv.hpp" #include +#include #include namespace infinilm::engine { @@ -418,6 +419,10 @@ void RankWorker::thread_loop() { infinicore::Tensor logits; infinicore::Tensor hidden_states; + // Batch size of the decode graph that produced logits + // (0 = eager fallback). Used to match the captured + // greedy-sampling graph below. + size_t graph_batch_size = 0; // All-position speculative/MTP runs need eager mode because // hidden states are not part of compiled graph outputs. if (!local_args.sample_all_positions && compiler_ != nullptr && rank_info_.pp_size == 1) { @@ -425,6 +430,7 @@ void RankWorker::thread_loop() { if (graph != nullptr && output != nullptr) { graph->run(); logits = output->logits; + graph_batch_size = local_args.input_offsets.value()->size(0) - 1; } } // Fall back to eager mode @@ -482,28 +488,51 @@ void RankWorker::thread_loop() { const size_t logits_positions = batch_size * total_len; const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req; const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; - auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; - for (size_t i{0}; i < n_out; ++i) { - size_t score_idx = i; - if (!sample_all_positions && !logits_are_last_token_only) { - score_idx = static_cast(input_offsets[i + 1] - 1); + infinicore::Tensor output_ids_dev; + // Greedy decode batches that ran through a compiled + // graph can replay the captured per-request argmax + // kernels instead of launching ~3 kernels per + // request here. top_k == 1 or temperature == 0 both + // map to the op's argmax branch (no RNG consumed). + std::shared_ptr sampling_graph; + infinicore::Tensor sampling_out; + if (graph_batch_size == n_req && logits_are_last_token_only + && (top_k == 1 || temperature == 0.0f)) { + std::tie(sampling_graph, sampling_out) = compiler_->get_sampling_compiled(n_req); + } + if (sampling_graph != nullptr && sampling_out && sampling_out->size(0) == n_out) { + static bool sampling_graph_logged = false; + if (!sampling_graph_logged && std::getenv("INFINILM_DEBUG_SAMPLING") != nullptr) { + sampling_graph_logged = true; + spdlog::info("sampling graph replay engaged (batch={})", n_req); + } + sampling_graph->run(); + output_ids_dev = sampling_out; + } else { + output_ids_dev = infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device); + + for (size_t i{0}; i < n_out; ++i) { + size_t score_idx = i; + if (!sample_all_positions && !logits_are_last_token_only) { + score_idx = static_cast(input_offsets[i + 1] - 1); + } + auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; + auto out{output_ids_dev->narrow({{0, i, 1}})->view({})}; + float random_val = std::uniform_real_distribution(0, 1)(rng_); + infinicore::op::random_sample_( + out, score, random_val, top_p, top_k, temperature); } - auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; - auto out{output_ids->narrow({{0, i, 1}})->view({})}; - float random_val = std::uniform_real_distribution(0, 1)(rng_); - infinicore::op::random_sample_( - out, score, random_val, top_p, top_k, temperature); } if (rank_info_.pp_size > 1) { infinicore::op::distributed::send( - output_ids, + output_ids_dev, 0, rank_info_.world_comm); } - output_ids = output_ids->to(infinicore::Device::cpu()); + auto output_ids{output_ids_dev->to(infinicore::Device::cpu())}; infinicore::context::syncStream(); diff --git a/csrc/layers/attention/backends/attention_layer.cpp b/csrc/layers/attention/backends/attention_layer.cpp index fcaefa292..2765c510c 100644 --- a/csrc/layers/attention/backends/attention_layer.cpp +++ b/csrc/layers/attention/backends/attention_layer.cpp @@ -1,7 +1,14 @@ #include "attention_layer.hpp" +#include "infinicore/ops.hpp" + +#include +#include namespace infinilm::layers::attention { +// AttentionLayer 构造：按所选 attention backend 实例化对应的实现类， +// 统一存入 attn_backend_impl_（std::variant），forward 时用 std::visit 分发。 + AttentionLayer::AttentionLayer(size_t num_heads, size_t head_size, float scale, @@ -20,6 +27,10 @@ AttentionLayer::AttentionLayer(size_t num_heads, case ::infinilm::backends::AttentionBackend::FLASH_ATTN: attn_backend_impl_ = std::make_shared(num_heads, head_size, scale, num_kv_heads, layer_idx); break; + case ::infinilm::backends::AttentionBackend::HYBRID: + // HYBRID：prefill 走 FA2，decode 走自研 paged-attention kernel（见下方 HybridAttentionImpl::forward） + attn_backend_impl_ = std::make_shared(num_heads, head_size, scale, num_kv_heads, layer_idx); + break; default: throw std::runtime_error("infinilm::layers::attention::AttentionLayer: unsupported attention backend"); } @@ -28,6 +39,8 @@ AttentionLayer::AttentionLayer(size_t num_heads, infinicore::Tensor AttentionLayer::forward(infinicore::Tensor &query, infinicore::Tensor &key, infinicore::Tensor &value) const { + // 从全局 forward 上下文中取出 attention 元数据与当前层的 KV cache， + // 再按 attn_backend_impl_ 的实际类型分发到对应后端的 forward。 auto &forward_context = infinilm::global_state::get_forward_context(); auto &attn_metadata = forward_context.attn_metadata; auto &kv_cache = forward_context.kv_cache_vec[layer_idx_]; @@ -40,3 +53,92 @@ infinicore::Tensor AttentionLayer::forward(infinicore::Tensor &query, } } // namespace infinilm::layers::attention + +namespace infinilm::layers::attention::backends { + +size_t decode_fa_ctx_threshold() { + static const size_t threshold = []() { + if (const char *env = std::getenv("INFINILM_DECODE_CTX_THRESHOLD")) { + const long long v = std::strtoll(env, nullptr, 10); + return v > 0 ? static_cast(v) : std::numeric_limits::max(); + } + const auto &model_config = infinilm::global_state::get_infinilm_config().model_config; + if (model_config == nullptr) { + return std::numeric_limits::max(); + } + const size_t hidden = model_config->get_or("hidden_size", 0); + const size_t layers = model_config->get_or("num_hidden_layers", 0); + const size_t kv_heads = model_config->get_or("num_key_value_heads", 0); + // Measured on RTX 5090, bf16, single-request decode (dev_perf + // --ctx-sweep, v12): splitkv wins below, FA kvcache wins above. + if (layers == 28 && kv_heads == 8) { + if (hidden == 1024) { // Qwen3-0.6B: splitkv +9% @3.2k, tie @3.9k + return static_cast(3400); + } + if (hidden == 2048) { // Qwen3-1.7B: splitkv +2% @1.3k, -7% @1.9k + return static_cast(1500); + } + } + return std::numeric_limits::max(); + }(); + return threshold; +} + +// HybridAttentionImpl 构造：内部持有一个 FlashAttentionImpl 实例， +// prefill 阶段的 mha_varlen_fwd 直接复用它；decode 阶段只用它的 +// do_kv_cache_update 写缓存（BSHD 布局），注意力计算走自研 paged kernel。 +HybridAttentionImpl::HybridAttentionImpl(size_t num_heads, + size_t head_size, + float scale, + size_t num_kv_heads, + size_t layer_idx) + : flash_(std::make_shared(num_heads, head_size, scale, num_kv_heads, layer_idx)), + num_heads_(num_heads), + head_size_(head_size), + scale_(scale) {} + +infinicore::Tensor HybridAttentionImpl::forward(const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const { + // 与各独立 impl 相同的 prefill 判定：flattened paged 模式下， + // 纯 decode 步每个序列恰好只有一个 query token， + // 因此 query 行数 ≠ 序列数时即为 prefill（或含 prefill 的混合 batch）。 + const size_t seq_len = query->shape()[0]; + const bool is_prefill = (seq_len != attn_metadata.total_sequence_lengths.value()->shape()[0]); + if (is_prefill) { + // prefill / 混合 batch：走 FA2 varlen（mha_varlen_fwd），prefill 性能最优。 + return flash_->forward(layer, query, key, value, kv_cache, attn_metadata); + } + // 长上下文 decode：FA 的 kvcache kernel 反超自研 splitkv（交叉点随模型 + // 几何变化，见 decode_fa_ctx_threshold）。flash_->forward 会重新判定 + // decode 分支并在同一份 BSHD cache 上跑 mha_kvcache_，切换无数据搬运。 + if (attn_metadata.max_sequence_length > decode_fa_ctx_threshold()) { + return flash_->forward(layer, query, key, value, kv_cache, attn_metadata); + } + // 纯 decode：复用 FA 的 cache update（BSHD 布局 + permuted paged_caching_ 写入）， + // 然后直接调 stride 感知的 paged-attention decode kernel（splitkv）。 + // FA2 的 mha_fwd_kvcache 是 sm80 时代的 kernel，在 Blackwell 上 decode 显著偏慢， + // 短/中上下文下自研 paged kernel 更快——这正是 HYBRID 的动机。 + auto [k_total, v_total] = flash_->do_kv_cache_update(layer, key, value, kv_cache, attn_metadata.slot_mapping.value()); + const size_t value_head_dim = value->size(value->ndim() - 1); + auto attn_output = infinicore::Tensor::empty({seq_len, num_heads_, value_head_dim}, query->dtype(), query->device()); + // paged kernel 期望的 shape 是 [num_blocks, num_kv_heads, block_size, head_dim] + // （BHSD），但 descriptor 逐维取 stride、kernel 按 row_stride 寻址， + // 因此对 BSHD 的 cache 做 permute({0,2,1,3}) 逻辑视图即可零拷贝直读。 + infinicore::op::paged_attention_( + attn_output, + query, + k_total->permute({0, 2, 1, 3}), + v_total->permute({0, 2, 1, 3}), + attn_metadata.block_tables.value(), + attn_metadata.total_sequence_lengths.value(), + std::nullopt, + scale_); + // 展平成模型层期望的 [1, seq_len, hidden] 输出。 + return attn_output->view({1, seq_len, num_heads_ * value_head_dim}); +} + +} // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/attention/backends/attention_layer.hpp b/csrc/layers/attention/backends/attention_layer.hpp index 874110629..594e84e50 100644 --- a/csrc/layers/attention/backends/attention_layer.hpp +++ b/csrc/layers/attention/backends/attention_layer.hpp @@ -10,7 +10,61 @@ #include namespace infinilm::layers::attention { -using AttentionImpl = std::variant, std::shared_ptr, std::shared_ptr>; + +class AttentionLayer; + +namespace backends { + +/** + * @brief Decode-kernel crossover: pure-decode steps whose longest context + * exceeds this threshold are routed to FA's kvcache kernel instead of the + * paged splitkv kernel (the two read the same BSHD paged cache, so switching + * costs no data movement). Measured on RTX 5090 via the dev_perf ctx sweep; + * the crossover moves with model geometry, so it is keyed on + * (hidden_size, num_hidden_layers, num_key_value_heads) and unmeasured + * geometries get SIZE_MAX (= never route, the status quo). + * Override with INFINILM_DECODE_CTX_THRESHOLD= (<=0 disables routing). + */ +size_t decode_fa_ctx_threshold(); + +/** + * @brief Hybrid attention: prefill (and mixed batches) go to FlashAttention-2 + * varlen; pure decode steps reuse FA's paged cache update (BSHD layout) but + * run the paged-attention decode kernel, which reads the BSHD cache via + * strides and is faster than FA2's kvcache path at short/medium contexts. + * Long-context decode steps (max_sequence_length > decode_fa_ctx_threshold()) + * go to FA's kvcache kernel instead. + * + * 动机：5090（Blackwell）上 FA2 decode 显著慢于自研 splitkv（v8 实测 + * -35%~-50%），FA 只赢 prefill，故按阶段分离路由；长上下文时 FA 的 + * kvcache kernel 反超，交叉点随模型几何变化，见 decode_fa_ctx_threshold()。 + */ +class HybridAttentionImpl { +public: + HybridAttentionImpl(size_t num_heads, + size_t head_size, + float scale, + size_t num_kv_heads, + size_t layer_idx); + + infinicore::Tensor forward(const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const; + +private: + std::shared_ptr flash_; // 复用的 FA2 实现：prefill 计算 + decode 阶段写 KV cache + size_t num_heads_; + size_t head_size_; + float scale_; +}; + +} // namespace backends + +// 后端实现的类型集合：static / paged / flash / hybrid，forward 时按实际类型分发 +using AttentionImpl = std::variant, std::shared_ptr, std::shared_ptr, std::shared_ptr>; /** * @brief Attention layer. diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index 5d284a316..e07743958 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -65,6 +65,12 @@ std::vector InfinilmModel::default_allocate_kv_cache_tensors case backends::AttentionBackend::FLASH_ATTN: { ; } + case backends::AttentionBackend::HYBRID: { + // HYBRID 的 KV cache 与 PAGED_ATTN 同构（分块 paged 分配）， + // 只是块内布局为 BSHD（见 PagedKVCache::create_layer_kv_cache）， + // 因此直接落入下面的 paged 分配分支。 + ; + } case backends::AttentionBackend::PAGED_ATTN: { auto paged_kv_cache_config = dynamic_cast(cache_config); if (nullptr == paged_kv_cache_config) { diff --git a/csrc/models/qwen3/qwen3_attention.cpp b/csrc/models/qwen3/qwen3_attention.cpp index 7d9beb043..1455b132a 100644 --- a/csrc/models/qwen3/qwen3_attention.cpp +++ b/csrc/models/qwen3/qwen3_attention.cpp @@ -2,9 +2,31 @@ #include "../../global_state/global_state.hpp" #include "../../layers/attention/attention.hpp" #include "../../utils.hpp" +#include "infinicore/ops/rms_norm_rope.hpp" namespace infinilm::models::qwen3 { +namespace { + +// Mirrors the rms_norm_rope device dispatch in InfiniCore +// (src/infiniop/ops/rms_norm_rope/operator.cc); devices without a backend +// (Cambricon, Ascend, Metax, Moore, Kunlun) use the unfused chain instead. +bool rms_norm_rope_supported(infinicore::Device::Type type) { + switch (type) { + case infinicore::Device::Type::CPU: + case infinicore::Device::Type::NVIDIA: + case infinicore::Device::Type::ILUVATAR: + case infinicore::Device::Type::ALI: + case infinicore::Device::Type::QY: + case infinicore::Device::Type::HYGON: + return true; + default: + return false; + } +} + +} // namespace + Qwen3Attention::Qwen3Attention(std::shared_ptr model_config, size_t layer_idx, const infinicore::Device &device) { @@ -20,6 +42,9 @@ Qwen3Attention::Qwen3Attention(std::shared_ptr mo double rms_norm_eps = model_config->get("rms_norm_eps"); attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + // The fused op is full-rotary only (head_dim == 2 * table_dim). + use_fused_norm_rope_ = model_config->get_rotary_dim() == head_dim_ + && rms_norm_rope_supported(device.getType()); const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); @@ -121,8 +146,6 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); - q_reshaped = q_norm_->forward(q_reshaped); - k_reshaped = k_norm_->forward(k_reshaped); // 3. Prepare position_ids for RoPE auto pos_shape = position_ids->shape(); @@ -136,9 +159,21 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi throw std::runtime_error("Unexpected position_ids shape"); } - // 4. Apply RoPE to QK - rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); - rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + // 4. Per-head RMSNorm + RoPE on q/k (full rotary): fused in-place op when + // the device has an rms_norm_rope backend, unfused chain otherwise. + if (use_fused_norm_rope_) { + infinicore::op::rms_norm_rope_(q_reshaped, q_norm_->weight(), pos_ids_for_rope, + rotary_emb_->sin_cache(), rotary_emb_->cos_cache(), + static_cast(q_norm_->eps()), rotary_emb_->algo()); + infinicore::op::rms_norm_rope_(k_reshaped, k_norm_->weight(), pos_ids_for_rope, + rotary_emb_->sin_cache(), rotary_emb_->cos_cache(), + static_cast(k_norm_->eps()), rotary_emb_->algo()); + } else { + q_reshaped = q_norm_->forward(q_reshaped); + k_reshaped = k_norm_->forward(k_reshaped); + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + } // 5. Attn Backend calculate auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); diff --git a/csrc/models/qwen3/qwen3_attention.hpp b/csrc/models/qwen3/qwen3_attention.hpp index b82f01fed..9404062c8 100644 --- a/csrc/models/qwen3/qwen3_attention.hpp +++ b/csrc/models/qwen3/qwen3_attention.hpp @@ -42,6 +42,7 @@ class Qwen3Attention : public infinicore::nn::Module { std::shared_ptr attn_; ::infinilm::backends::AttentionBackend attention_backend_; + bool use_fused_norm_rope_; size_t layer_idx_; size_t num_attention_heads_; size_t num_key_value_heads_; diff --git a/dev_perf/README.md b/dev_perf/README.md new file mode 100644 index 000000000..4b4f830b7 --- /dev/null +++ b/dev_perf/README.md @@ -0,0 +1,93 @@ +# dev_perf — 项目 #2(性能优化)基线工具 + +离线双后端基线压测:同一负载矩阵跑 InfiniLM 与 vLLM,产出差距清单,作为 #2 立项依据。 + +## 运行前提 + +1. InfiniCore 已构建安装（`INFINI_ROOT` 指向安装前缀，默认 `~/.infini`；flash-attn 后端需要 ATen+FA 版，见 gap_analysis.md v4），且当前 checkout 已构建 C++ 扩展（`xmake build _infinilm && xmake install`，产物在 `python/infinilm/lib/`）。运行时动态库路径由 bench.py 自动 re-exec 注入（LD_LIBRARY_PATH ← InfiniCore/InfiniLM 的 lib 目录），无需手动设置 +2. vLLM 独立环境：自建 venv（uv 创建，cu128 torch；**勿装进项目 venv**，vllm 会锁定 torch 版本）。运行时需带齐 `PATH=/bin:$CUDA_HOME/bin:$PATH` 和 `CUDA_HOME`（否则 EngineCore 找不到 ninja 起不来） +3. GPU 空闲（压测前停掉常驻的 GPU 进程）；如本机有显存看门狗，vLLM 侧 `--gpu-mem-util` 不要超过 ~0.8 +4. WSL2 下 vLLM 需 `VLLM_WSL2_ENABLE_PIN_MEMORY=1`（bench.py 已自动设置；否则 vLLM 禁用 pinned memory 导致 UvaBuffer 报 "UVA is not available"） + +## 运行 + +```bash +# InfiniLM（任何带 torch 的 venv，通过 sys.path 引用当前 checkout 的构建产物） +python dev_perf/bench.py --engine infinilm --model Qwen/Qwen3-1.7B + +# vLLM（独立 venv） +/path/to/vllm-venv/bin/python dev_perf/bench.py --engine vllm --model Qwen/Qwen3-1.7B +``` + +结果 JSON 写入 `dev_perf/results/`。 + +InfiniLM 侧常用参数： + +- `--num-blocks 64`：低显存模式。paged cache 预分配从 ~13GB 降到 ~1.8GB + （Qwen3-1.7B 口径），w1-w4 在 16k token 容量内仍可完整运行；num_blocks + 会记入结果 JSON。注意 num_blocks 对性能有一阶影响——512 在 16GB 卡上会 + 把显存占满并导致全负载严重劣化（见 gap_analysis.md v3），跨轮次对比必须 + 用相同 num_blocks。 +- `--attn-backend flash-attn`：prefill/decode 改用 FlashAttention-2（需 + ATen+FA 版 InfiniCore，见 gap_analysis.md v4 的构建与 `INFINI_ROOT` + 用法）。长 prefill 负载收益 ~20~30%。 +- `--attn-backend hybrid`：prefill 走 FA2 varlen、decode 走自研 paged + kernel（经 strides 直读 FA 的 BSHD cache）的分离路由。5090 上为最优 + 单配置（见 gap_analysis.md v9）。 +- `--only w2_long_prefill`：只跑指定负载（逗号分隔）。 +- `--concurrent-prefill-n N`：w5_concurrent_prefill 的并发长 prompt 条数 + （默认 8，复用 w2 的 prompt 构造、逐条加不同前缀以避开 prefix caching + 去重），chunked prefill 验收负载，`--only w5_concurrent_prefill` 单跑。 +- `--decode-stall-n N`：w6_decode_stall 的 decode 长流并发数（默认 8）。 + w6 是引擎级负载：N 条短 prompt 各生成 512 tok，第 64 步时注入一条 + ~6.5k tok 长 prompt（带 nonce 头避开 prefix caching），记录逐步耗时， + 考察注入前后 decode 流的最大/p90 ITL 尖峰与注入请求 TTFT——FCFS 下 + 整条 decode 流被整段 prefill 堵住，chunked prefill 应把尖峰摊平。 + 仅 infinilm 引擎支持（vLLM 跳过）。 +- `--speculative-method prompt_lookup`：开启投机采样（零训练 n-gram + draft，模型无关；`--num-draft-tokens K` 控制每步验证的 draft 数，默认 + 4；eagle 方法需 `--draft-model` 指向 MiniCPM Eagle 权重）。每个负载的 + JSON 会带 spec_accept_rate / spec_avg_tokens_per_step 增量。相关 env： + `INFINILM_PROMPT_LOOKUP_MIN/MAX_NGRAM`（默认 2/4，匹配质量调参）、 + `INFINILM_SPEC_MAX_BATCH_SIZE`（默认 32，超过则回退常规前向）、 + `INFINILM_SPEC_MIN_AVG_TOKENS` / `_GATE_WINDOW` / `_GATE_COOLDOWN` + （默认 2.0/32/64，自适应收益门控：窗口内平均每步产出低于阈值则回退 + 常规前向一段再重试，低命中负载开投机不致亏）。 + 贪心下输出与非投机数学等价（分布无损），但**不保证逐位一致**： + verify 前向的 batch 形状与基线 decode 不同，logit 近平局的位置 + argmax 可能翻转（实测分歧点 top-2 间隙 0~0.125）；高命中负载 + （w5/w7）保持逐 token 相同。用 `--dump-outputs` + compare_outputs + 对拍验收；w7_repetitive_copy 是接受率演示负载（见 gap_analysis.md v16）。 +- `--dump-outputs`：把每个请求的输出 token ids 记入 JSON，配合 + `compare_outputs.py a.json b.json ...` 做跨引擎/跨 backend 的贪心解码 + 逐 token 对拍（exact match 数 + 最早分叉位置）。注意：入库归档的 + results JSON 不保留完整 ids，只有 `output_token_ids_sha256`（逐请求 + 哈希）与 `output_token_ids_head`（前 32 token）；哈希相等即序列相等， + 对拍时报 exact match 数、无分叉位置。 + +## 负载矩阵 + +| 负载 | 请求 | 输出 | 考察点 | +|---|---|---|---| +| w1_short_decode | 1 × 短 prompt | 128 tok | 单请求 decode 延迟(ms/tok) | +| w2_long_prefill | 1 × ~2k tok prompt | 128 tok | prefill 吞吐 | +| w3_batch32 | 32 × 短 prompt | 128 tok | 批处理总吞吐 | +| w4_long_decode | 1 × 短 prompt | 1024 tok | 长生成 decode 稳定性 | +| w5_concurrent_prefill | 8 × ~2k tok prompt 并发 | 128 tok | chunked prefill：并发长 prefill 不阻塞 decode（e2e wall time） | +| w6_decode_stall | 8 × 短 prompt 长 decode + 中途注入 1 × ~6.5k tok prompt | 512/32 tok | chunked prefill 收益场景：decode 流 ITL 尖峰、注入请求 TTFT | +| w7_repetitive_copy | 1 × 重复段落 pattern 续写 | 512 tok | prompt-lookup 投机采样：n-gram 命中率高，接受率/收益演示 | + +## 公平性约定 + +- 贪心解码(temperature=0, top_k=1)、ignore_eos=True、逐字节相同的 prompt 与 max_tokens +- vLLM 用 v1 默认(CUDA graph 开、prefix caching 开);InfiniLM 用 enable_graph=False、paged-attn、prefix caching 开——双方配置记入结果 JSON,差距解读时先考虑这两项 +- 每个引擎先跑一轮 w1 预热(不计时),再正式计时 + +## 产出 + +差距清单（哪个负载、差多少、可能原因）→ 据此决定 #2 是否转正立项，以及立项的可度量目标。 + +最新进展见 `dev_perf/gap_analysis.md`（v8：5090 全栈合流——FA2+融合 +ABBA -5%~-13%，并发现 FA decode 在 Blackwell 上慢于自研 splitkv，默认 +后端结论修正；v7 elementwise 链 kernel 级闭环；含对比表、归因、复现 +命令和立项书模板）。 diff --git a/dev_perf/bench.py b/dev_perf/bench.py new file mode 100644 index 000000000..30e1024ac --- /dev/null +++ b/dev_perf/bench.py @@ -0,0 +1,491 @@ +"""Offline dual-engine perf baseline (project #2 gap analysis). + +One script, two interpreters: + + InfiniLM (any venv with torch; imports the built package from this checkout): + python dev_perf/bench.py --engine infinilm --model Qwen/Qwen3-1.7B + + vLLM (isolated venv): + /path/to/vllm-venv/bin/python dev_perf/bench.py --engine vllm --model Qwen/Qwen3-1.7B + +Fairness contract: greedy decoding (temperature=0, top_k=1), ignore_eos=True, +identical prompts and max_tokens from workload.py. vLLM runs with its v1 +defaults (CUDA graphs + prefix caching on); InfiniLM runs with +enable_graph=False + prefix caching on. Both facts are recorded in the output. +""" + +import argparse +import json +import os +import sys +import time + +from workload import ( + SHORT_PROMPTS, + WORKLOADS, + MemSampler, + concurrent_prefill_workload, + ctx_sweep_workloads, + decode_stall_workload, + resolve_model_path, +) + +# Checkout whose built infinilm package we benchmark. Defaults to this repo's +# own python/ dir (where `xmake install _infinilm` lands its .so); override +# with INF_MAIN_PYTHON to benchmark a different checkout's build. +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +MAIN_CHECKOUT_PYTHON = os.environ.get("INF_MAIN_PYTHON", os.path.join(_REPO_ROOT, "python")) +# INFINI_ROOT selects the InfiniCore install (e.g. ~/.infini-fa for the +# ATen+FA build, required by the flash-attn backend); defaults to ~/.infini. +_INFINI_CORE_LIB = os.path.join( + os.environ.get("INFINI_ROOT", os.path.expanduser("~/.infini")), "lib" +) +def _torch_lib_dir(): + """site-packages/torch/lib, needed on LD_LIBRARY_PATH for ATen-enabled + InfiniCore builds (they link libtorch).""" + for p in sys.path: + d = os.path.join(p, "torch", "lib") + if os.path.isdir(d): + return d + return None + + +LIB_DIRS = [ + _INFINI_CORE_LIB, + os.path.join(MAIN_CHECKOUT_PYTHON, "infinilm", "lib"), +] + ([_torch_lib_dir()] if _torch_lib_dir() else []) + + +def ensure_ld_library_path(): + """infinilm/infinicore extensions find their .so deps via LD_LIBRARY_PATH. + + Re-exec the interpreter with the env set if needed (the dynamic linker + reads LD_LIBRARY_PATH only at process start). + """ + cur = os.environ.get("LD_LIBRARY_PATH", "") + if all(d in cur.split(":") for d in LIB_DIRS): + return + os.environ["LD_LIBRARY_PATH"] = ":".join(LIB_DIRS + ([cur] if cur else [])) + os.execv(sys.executable, [sys.executable] + sys.argv) + + +def build_engine(args): + if args.engine == "infinilm": + ensure_ld_library_path() + sys.path.insert(0, MAIN_CHECKOUT_PYTHON) + from infinilm.llm.llm import LLM + from infinilm.llm.sampling_params import SamplingParams + + speculative_method = ( + None if args.speculative_method == "none" else args.speculative_method + ) + llm = LLM( + model_path=args.model, + device="cuda", + dtype="bfloat16", + cache_type="paged", + attn_backend=args.attn_backend, + max_batch_size=64, + num_blocks=args.num_blocks, + enable_prefix_caching=True, + enable_graph=args.enable_graph, + draft_model_path=args.draft_model, + speculative_method=speculative_method, + num_draft_tokens=args.num_draft_tokens, + ) + + def generate(prompts, max_tokens): + return llm.generate( + prompts=prompts, + sampling_params=SamplingParams( + temperature=0.0, top_k=1, max_tokens=max_tokens, ignore_eos=True + ), + use_tqdm=False, + ) + + def close(): + llm.close() + + # w6_decode_stall drives the engine directly (mid-stream add_request) + generate.llm = llm + + def spec_stats(): + """投机采样累计计数快照；未启用投机时返回 None。""" + return llm.engine.model_runner.get_speculative_stats() + + generate.spec_stats = spec_stats + + engine_notes = { + "engine": "infinilm", + "cuda_graph": bool(args.enable_graph), + "prefix_caching": True, + "attn_backend": args.attn_backend, + "num_blocks": args.num_blocks, + "speculative_method": speculative_method, + "num_draft_tokens": ( + args.num_draft_tokens if speculative_method else None + ), + } + + elif args.engine == "vllm": + from vllm import LLM, SamplingParams + + llm = LLM( + model=args.model, + dtype="bfloat16", + gpu_memory_utilization=args.gpu_mem_util, + max_model_len=4096, + seed=0, + ) + + def generate(prompts, max_tokens): + return llm.generate( + prompts, + SamplingParams( + temperature=0.0, top_k=1, max_tokens=max_tokens, ignore_eos=True + ), + use_tqdm=False, + ) + + def close(): + try: + llm.llm_engine.engine_core.shutdown() + except Exception: + pass + + engine_notes = {"engine": "vllm", "cuda_graph": True, "prefix_caching": True, "gpu_memory_utilization": args.gpu_mem_util} + + else: + raise ValueError(f"unknown engine: {args.engine}") + + return generate, close, engine_notes + + +def run_decode_stall(llm, params, dump_outputs=False): + """w6 driver: steady decode stream + one long prompt injected mid-flight. + + n_decode short-prompt requests generate 512 tokens each; after + INJECT_AFTER_STEPS engine steps one long prompt is added. Per-step wall + times are recorded so the injected prefill shows up as an ITL spike for + the decode stream (plain FCFS) or as a few mildly-longer mixed steps + (chunked prefill). Returns a bench record dict. + """ + from infinilm.llm.request import InferenceRequest + from infinilm.llm.sampling_params import SamplingParams + + n_decode, inject_prompt = params + decode_max_tokens = 512 + inject_after_steps = 64 + inject_max_tokens = 32 + + engine = llm.engine + + def add(prompt, max_tokens, tag): + req = InferenceRequest( + request_id=f"w6-{tag}-{time.time_ns()}", + prompt=prompt, + prompt_token_ids=engine.tokenize(prompt), + sampling_params=SamplingParams( + temperature=0.0, top_k=1, max_tokens=max_tokens, ignore_eos=True + ), + eos_token_ids=engine.eos_token_ids, + ) + engine.add_request(req) + return req + + decode_reqs = [ + add(f"w6问题 {i + 1}：{SHORT_PROMPTS[i % len(SHORT_PROMPTS)]}", + decode_max_tokens, f"d{i}") + for i in range(n_decode) + ] + injected = None + t_inject = None + inject_ttft = None + step_times = [] + inject_step_idx = None + + t_start = time.perf_counter() + while True: + t0 = time.perf_counter() + did_work, _ = engine.step() + dt = time.perf_counter() - t0 + if did_work: + step_times.append(dt) + if injected is None and len(step_times) >= inject_after_steps: + injected = add(inject_prompt, inject_max_tokens, "p") + t_inject = time.perf_counter() + inject_step_idx = len(step_times) + if ( + injected is not None + and inject_ttft is None + and injected.get_num_generated_tokens() >= 1 + ): + inject_ttft = time.perf_counter() - t_inject + if all(r.is_finished() for r in decode_reqs) and ( + injected is None or injected.is_finished() + ): + break + e2e = time.perf_counter() - t_start + + pre = step_times[:inject_step_idx] + post = step_times[inject_step_idx:] + post_sorted = sorted(post) + p90 = post_sorted[int(len(post_sorted) * 0.9)] if post_sorted else 0.0 + + all_reqs = decode_reqs + ([injected] if injected else []) + out_token_counts = [r.get_num_generated_tokens() for r in all_reqs] + total_out = sum(out_token_counts) + rec = { + "workload": "w6_decode_stall", + "num_requests": len(all_reqs), + "prompt_tokens_total": sum(len(r.prompt_token_ids) for r in all_reqs), + "output_tokens_total": total_out, + "e2e_seconds": round(e2e, 3), + "output_tokens_per_sec": round(total_out / e2e, 2), + "ms_per_output_token": round(1000.0 * e2e / max(total_out, 1), 3), + "min_output_tokens": min(out_token_counts), + "max_output_tokens": max(out_token_counts), + # decode-stream smoothness: mean step time before injection vs the + # worst/p90 step after the long prompt lands + "decode_step_ms_pre_inject": round(1000.0 * sum(pre) / max(len(pre), 1), 3), + "step_ms_post_inject_max": round(1000.0 * max(post), 2) if post else None, + "step_ms_post_inject_p90": round(1000.0 * p90, 2), + "inject_prompt_tokens": len(injected.prompt_token_ids) if injected else 0, + "inject_ttft_ms": round(1000.0 * inject_ttft, 1) if inject_ttft else None, + } + if dump_outputs: + rec["output_token_ids"] = [list(r.generated_token_ids) for r in all_reqs] + return rec + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--engine", choices=["infinilm", "vllm"], required=True) + parser.add_argument("--model", default="Qwen/Qwen3-1.7B") + parser.add_argument( + "--out-dir", + default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "results"), + ) + parser.add_argument( + "--enable-graph", + action="store_true", + help="infinilm only: enable graph compiling (CUDA-graph-like capture)", + ) + parser.add_argument( + "--num-blocks", + type=int, + default=512, + help="infinilm only: paged KV cache blocks (x256 tokens); lower it for a low-VRAM run", + ) + parser.add_argument( + "--gpu-mem-util", + type=float, + default=0.85, + help="vllm only: gpu_memory_utilization; lower it to stay under VRAM watchdogs", + ) + parser.add_argument( + "--only", + default=None, + help="comma-separated workload names to run (default: all)", + ) + parser.add_argument( + "--concurrent-prefill-n", + type=int, + default=8, + help="number of concurrent long prompts in w5_concurrent_prefill", + ) + parser.add_argument( + "--decode-stall-n", + type=int, + default=8, + help="number of streaming decode requests in w6_decode_stall", + ) + parser.add_argument( + "--attn-backend", + default="paged-attn", + choices=["default", "static-attn", "paged-attn", "flash-attn", "flashinfer", "hybrid"], + help="infinilm only: attention backend", + ) + parser.add_argument( + "--speculative-method", + default="none", + choices=["none", "eagle", "prompt_lookup"], + help="infinilm only: speculative decoding method (prompt_lookup needs no draft model)", + ) + parser.add_argument( + "--draft-model", + default=None, + help="infinilm only: Eagle/MTP draft model directory (speculative-method=eagle)", + ) + parser.add_argument( + "--num-draft-tokens", + type=int, + default=4, + help="infinilm only: draft tokens verified per speculative step", + ) + parser.add_argument( + "--ctx-sweep", + action="store_true", + help="replace the workload list with a decode-vs-context-length sweep " + "(single request, 128 decode tokens, ~0.2k~5k ctx)", + ) + parser.add_argument( + "--dump-outputs", + action="store_true", + help="record per-request output token ids in the JSON (for cross-engine correctness diffing)", + ) + args = parser.parse_args() + + # WSL2: vLLM disables pinned memory by default, which makes its V2 model + # runner's UvaBuffer raise "UVA is not available". Pinned memory works on + # WSL2 kernels >= 4.19.121 (verified locally), so opt back in. + os.environ.setdefault("VLLM_WSL2_ENABLE_PIN_MEMORY", "1") + + args.model = resolve_model_path(args.model) + print(f"[bench] engine={args.engine} model={args.model}", flush=True) + + sampler = MemSampler() + sampler.start() + + t_load0 = time.perf_counter() + generate, close, engine_notes = build_engine(args) + load_seconds = time.perf_counter() - t_load0 + mem_after_load = sampler.latest + print(f"[bench] load took {load_seconds:.1f}s, mem={mem_after_load} MiB", flush=True) + + def _spec_snapshot(): + fn = getattr(generate, "spec_stats", None) + return fn() if fn is not None else None + + def _spec_rec(before, after): + """本 workload 区间的投机采样增量统计；未启用投机时返回 None。""" + if after is None: + return None + before = before or {} + drafted = after["drafted_tokens"] - before.get("drafted_tokens", 0) + accepted = after["accepted_tokens"] - before.get("accepted_tokens", 0) + steps = after["spec_steps"] - before.get("spec_steps", 0) + emitted = after["emitted_tokens"] - before.get("emitted_tokens", 0) + alloc_fail = after["verify_alloc_failures"] - before.get( + "verify_alloc_failures", 0 + ) + return { + "spec_drafted_tokens": drafted, + "spec_accepted_tokens": accepted, + "spec_accept_rate": round(accepted / drafted, 4) if drafted else None, + "spec_avg_tokens_per_step": round(emitted / steps, 3) if steps else None, + "spec_verify_alloc_failures": alloc_fail, + } + + workloads = ( + ctx_sweep_workloads() + if args.ctx_sweep + else [ + concurrent_prefill_workload(args.concurrent_prefill_n) + if name == "w5_concurrent_prefill" + else decode_stall_workload(args.decode_stall_n) + if name == "w6_decode_stall" + else (name, prompts, max_tokens) + for name, prompts, max_tokens in WORKLOADS + ] + ) + + # Warmup (untimed): runs the first workload once to trigger lazy init / + # graph capture. + w1_name, w1_prompts, w1_mt = workloads[0] + generate(w1_prompts, w1_mt) + + only = set(args.only.split(",")) if args.only else None + results = [] + for name, prompts, max_tokens in workloads: + if only and name not in only: + continue + if name == "w6_decode_stall": + llm = getattr(generate, "llm", None) + if llm is None: + print( + f"[bench] {name}: skipped (engine exposes no raw llm handle)", + flush=True, + ) + continue + spec_before = _spec_snapshot() + rec = run_decode_stall(llm, prompts, dump_outputs=args.dump_outputs) + rec.update(_spec_rec(spec_before, _spec_snapshot()) or {}) + results.append(rec) + print( + f"[bench] {name}: e2e={rec['e2e_seconds']:.2f}s " + f"out={rec['output_tokens_total']} tok " + f"({rec['output_tokens_per_sec']} tok/s) " + f"step pre={rec['decode_step_ms_pre_inject']}ms " + f"post_max={rec['step_ms_post_inject_max']}ms " + f"post_p90={rec['step_ms_post_inject_p90']}ms " + f"inject_ttft={rec['inject_ttft_ms']}ms", + flush=True, + ) + continue + spec_before = _spec_snapshot() + t0 = time.perf_counter() + outputs = generate(prompts, max_tokens) + e2e = time.perf_counter() - t0 + + prompt_tokens = sum(len(o.prompt_token_ids or []) for o in outputs) + out_token_counts = [len(o.outputs[0].token_ids) for o in outputs] + total_out = sum(out_token_counts) + rec = { + "workload": name, + "num_requests": len(prompts), + "prompt_tokens_total": prompt_tokens, + "output_tokens_total": total_out, + "e2e_seconds": round(e2e, 3), + "output_tokens_per_sec": round(total_out / e2e, 2), + "ms_per_output_token": round(1000.0 * e2e / max(total_out, 1), 3), + "min_output_tokens": min(out_token_counts), + "max_output_tokens": max(out_token_counts), + } + rec.update(_spec_rec(spec_before, _spec_snapshot()) or {}) + if args.dump_outputs: + rec["output_token_ids"] = [list(o.outputs[0].token_ids) for o in outputs] + results.append(rec) + spec_note = "" + if rec.get("spec_avg_tokens_per_step") is not None: + spec_note = ( + f" spec_accept={rec['spec_accept_rate']}" + f" avg_tok/step={rec['spec_avg_tokens_per_step']}" + ) + print( + f"[bench] {name}: e2e={e2e:.2f}s out={total_out} tok " + f"({rec['output_tokens_per_sec']} tok/s, " + f"{rec['ms_per_output_token']} ms/tok){spec_note}", + flush=True, + ) + + sampler.stop() + + report = { + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), + "model": args.model, + **engine_notes, + "load_seconds": round(load_seconds, 2), + "mem_after_load_mib": mem_after_load, + "mem_peak_mib": sampler.peak, + "workloads": results, + } + os.makedirs(args.out_dir, exist_ok=True) + tag = os.path.basename(args.model.rstrip("/")) + out_path = os.path.join( + args.out_dir, f"{args.engine}_{tag}_{time.strftime('%m%d_%H%M%S')}.json" + ) + with open(out_path, "w") as f: + json.dump(report, f, indent=2, ensure_ascii=False) + print(f"[bench] wrote {out_path}", flush=True) + + # Shut down only after results are persisted; engine teardown must not + # be allowed to take the report down with it. + try: + close() + except Exception as e: + print(f"[bench] close() failed (results already saved): {e}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/dev_perf/compare_outputs.py b/dev_perf/compare_outputs.py new file mode 100644 index 000000000..8a97d8ebc --- /dev/null +++ b/dev_perf/compare_outputs.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Compare output_token_ids across bench JSON files (greedy decoding). + +Usage: compare_outputs.py file_a.json file_b.json [file_c.json ...] + +For each workload present in all files, reports per-request token match: +exact match count, and for mismatches the first divergence position. +Slimmed files carry output_token_ids_sha256 instead of full ids; those are +compared by hash (exact match only, no divergence position). A full-ids +file compared against a hash-only one is hashed on the fly with the same +digest (sha256 of the compact JSON encoding), so archived files stay +comparable with fresh --dump-outputs runs. +""" +import hashlib +import json +import sys + + +def digest(ids): + payload = json.dumps(ids, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + +def load(path): + d = json.load(open(path)) + label = f"{d['engine']}/{d.get('attn_backend', '-')}" + wl = {} + for w in d["workloads"]: + ids = w.get("output_token_ids") + wl[w["workload"]] = ids if ids is not None else w.get("output_token_ids_sha256") + return label, wl + + +def main(paths): + loaded = [load(p) for p in paths] + base_label, base_wl = loaded[0] + print(f"reference: {base_label} ({paths[0]})") + for label, wl in loaded[1:]: + print(f"\n=== {label} vs {base_label} ===") + for name, ref_ids in base_wl.items(): + ids = wl.get(name) + if ref_ids is None or ids is None: + print(f" {name}: missing token ids, skipped") + continue + if not ref_ids or not ids: + print(f" {name}: empty request list, skipped") + continue + same_count = len(ref_ids) == len(ids) + if not same_count: + print( + f" {name}: request count differs " + f"({len(ref_ids)} vs {len(ids)}), " + f"comparing first {min(len(ref_ids), len(ids))}" + ) + if isinstance(ref_ids[0], str) != isinstance(ids[0], str): + # mixed full-vs-hash: hash the full side + if isinstance(ids[0], str): + ref_ids = [digest(r) for r in ref_ids] + else: + ids = [digest(r) for r in ids] + hashed = isinstance(ref_ids[0], str) + suffix = " (sha256)" if hashed else "" + n_exact = 0 + worst_div = None + for a, b in zip(ref_ids, ids): + if a == b: + n_exact += 1 + continue + if hashed: + continue + div = next( + (j for j, (x, y) in enumerate(zip(a, b)) if x != y), + min(len(a), len(b)), + ) + if worst_div is None or div < worst_div: + worst_div = div + total = len(ref_ids) + if n_exact == total and same_count: + print(f" {name}: {total}/{total} requests exact match{suffix}") + elif hashed: + print(f" {name}: {n_exact}/{total} exact{suffix}") + elif worst_div is not None: + print( + f" {name}: {n_exact}/{total} exact, " + f"earliest divergence at token {worst_div}" + ) + else: + # mismatch purely from the request-count difference + print(f" {name}: {n_exact}/{total} exact (compared prefixes identical)") + for name in wl: + if name not in base_wl: + print(f" {name}: only in {label}, not compared") + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/dev_perf/gap_analysis.md b/dev_perf/gap_analysis.md new file mode 100644 index 000000000..0f1fcca30 --- /dev/null +++ b/dev_perf/gap_analysis.md @@ -0,0 +1,1219 @@ +# 项目 #2 基线差距清单 + +最新进展见 v16（2026-09-01）：投机采样方向开工并已在 5090 验收—— +prompt-lookup（零训练、模型无关的 n-gram draft）+ spec×chunked 融合 + +融合单前向（主采样与 draft 验证一次完成，纯 decode 批次形状固定 +b×(k+1)，为 verify 图化备好可录制形状）+ 自适应收益门控（低命中负载 +开投机不致亏）。纯 Python 实现，stub 测试 11 用例全绿；5090 实测 +w4/w5/w7 提速 2.1~2.4×、w6 吞吐 +57%、w1 打平，高命中负载输出与非 +投机逐 token 相同（开放生成的分歧定位为 near-tie argmax 翻转，非逻辑 +bug，详见 v16 验收结果节）。v15 的 kernel 级归因结论不变：小模型 +decode 的常数项（小 kernel 延迟 + 图外 CPU）决定了投机必须走融合 +单前向才有净收益。 + +--- + +## v16（2026-09-01）：prompt-lookup 投机采样 + spec×chunked 融合 + +动机：v15 归因认定 kernel 微优化天花板清晰，数量级杠杆在 FP8 与投机 +采样。仓库已有投机链路（speculative_runner.py）但三个硬限制：draft +只支持 minicpm_eagle（Qwen3-0.6B/1.7B 没有现成 EAGLE 头）、与 v13 +chunked prefill 在引擎 gate 互斥、且投机模式全程 eager（target 主前向 +走 forward_raw 拿 hidden states，被 rank_worker 的 +!sample_all_positions 门禁挡在图外）。v16 按分析的第一、三步落地 +（跳过需要训练的 EAGLE 头）：① prompt-lookup 打通链路，③ spec× +chunked 融合 + verify 形状固定化。 + +### 改动清单（全部纯 Python） + +- `config/engine_config.py`：新增 `speculative_method`（None/"eagle"/ + "prompt_lookup"）；归一化——给了 `draft_model_path` 未指定方法时默认 + "eagle"（旧调用方行为不变）；校验组合合法性。 +- `llm/model_runner/speculative_runner.py`： + - prompt_lookup draft：`_draft_prompt_lookup_tokens` 在请求自身 + prompt+已生成序列中找当前后缀的最近一次出现，取其后 k 个 token + （n-gram 范围由 `INFINILM_PROMPT_LOOKUP_MIN/MAX_NGRAM` 控制，默认 + 2/4；只影响命中质量，正确性由 verify 保证）。 + - **融合单前向**（`_forward_fused_prompt_lookup`）：纯 decode 批次时 + 每请求一行固定 k+1 个 token——[已提交尾 token x, d1..dk]（不足 k + 个用尾 token 补齐，验证逻辑对任意 draft 内容自校正），一次前向同时 + 拿到主采样 o[0]（=target_token）与逐位置验证输出 o[1..k]，接受最长 + 匹配前缀+1 修正 token，未接受槽位 rollback。与两前向流程逐 token + 等价，但省掉独立 verify 前向的 launch/CPU 常数项——这是 0.6B 级别 + 小模型上投机有没有净收益的分水岭（v15：eager 前向常数项与图化 + decode 同量级）。 + - 逐请求相位判定：`num_scheduled_tokens` 非空时按请求判断，中段 + chunk 请求不产 token、不进 draft/verify（否则 append_verify_slots + 会破坏块表不变量）；混排批次回退两前向流程。 + - batch 阈值：`len(requests) > INFINILM_SPEC_MAX_BATCH_SIZE`（默认 + 32）回退常规前向——大 batch decode 转向计算受限，投机放大每步 + 计算量反而降吞吐（w3 bs=32 类负载的保护）。 + - **自适应收益门控**：滑动窗口（`INFINILM_SPEC_GATE_WINDOW`，默认 + 32 请求步）内平均每步产出低于 `INFINILM_SPEC_MIN_AVG_TOKENS` + （默认 2.0，实测盈亏平衡点：eager 投机步成本 ≈ 2× 图化 decode + 步）时回退常规前向 `INFINILM_SPEC_GATE_COOLDOWN` 步（默认 64）， + 期满自动重试。开放生成等低命中负载开投机不致亏（见验收结果）。 + - verify 槽位分配失败（KV 块不足）退化为本请求不投机，不再抛异常。 + - 接受率埋点：`get_acceptance_stats()`（drafted/accepted/accept_rate/ + avg_tokens_per_step/verify_alloc_failures/gate_triggered），经 + `ModelRunner.get_speculative_stats()` 暴露。 +- `llm/model_runner/model_runner.py`：投机开关从 `draft_model_path is + not None` 改为 `speculative_method is not None`。 +- `llm/llm.py`：LLM/AsyncLLMEngine 新增 `speculative_method` 参数； + **拆除 spec×chunked 互斥 gate**（`_chunked_prefill_supported` 不再 + 排除投机路径；mamba/多模态处理器的排除保留）。 +- `dev_perf/bench.py`：`--speculative-method {none,eagle,prompt_lookup}` + / `--draft-model` / `--num-draft-tokens`；每个 workload 记录 + spec_accept_rate 与 spec_avg_tokens_per_step 增量并随 JSON 落盘。 +- `dev_perf/workload.py`：新增 w7_repetitive_copy（pattern 续写负载， + prompt-lookup 接受率上限的演示；不依赖指令遵循能力）。 + +### 本机验证（无 GPU，stub 模式） + +`test/test_prompt_lookup_spec.py`（复用 test_chunked_prefill.py 的 stub +harness，target 模型用确定性伪模型替代）7 用例 + 既有 4 用例全绿： + +- lookup 辅助函数语义（最近出现/min-max 边界/k 截断/未命中）； +- 高命中场景输出与非投机真值逐 token 相同，accept_rate>0.9、 + avg_tokens_per_step>1.4，且 decode 批次的 forward_raw 调用形状全是 + b×(k+1)（证明走了融合路径）； +- draft 整体拒绝/部分接受时输出仍精确、verify 槽位回滚、块无泄漏； +- spec×chunked 混排：中段 chunk 产出 []、完成 chunk 即投机、混排批次 + 走两前向、纯 decode 批次走融合（按调用形状断言）； +- 旧式整段 prefill 批次投机、非 greedy 回退； +- 自适应门控：零命中负载攒满 32 步窗口后触发回退，冷静期内不再产生 + 融合前向调用，输出仍逐 token 精确。 + +### 5090 验收步骤 + +```bash +# 基线（无投机） +python dev_perf/bench.py --engine infinilm --model Qwen/Qwen3-0.6B \ + --enable-graph --attn-backend hybrid \ + --only w1_short_decode,w4_long_decode,w7_repetitive_copy --dump-outputs +# 投机（prompt_lookup, k=4） +python dev_perf/bench.py --engine infinilm --model Qwen/Qwen3-0.6B \ + --enable-graph --attn-backend hybrid \ + --speculative-method prompt_lookup --num-draft-tokens 4 \ + --only w1_short_decode,w4_long_decode,w7_repetitive_copy --dump-outputs +# 正确性：贪心输出必须逐 token 相同 +python dev_perf/compare_outputs.py results/infinilm_*baseline*.json results/infinilm_*spec*.json +``` + +看点：w7 的 spec_accept_rate 与 avg_tokens_per_step（预期接近 k）、 +w1/w4 的 ms/tok 差（开放生成接受率低，可能接近打平——这本身就是 +prompt-lookup 的已知边界，EAGLE 头是后续解）、w3 不回归（batch 阈值 +保护）。spec×chunked 组合验收：`INFINILM_ENABLE_CHUNKED_PREFILL=1` + +`--speculative-method prompt_lookup` 跑 w5/w6 对拍。 + +### 5090 验收结果（2026-09-01，Qwen3-0.6B，hybrid + graph，k=4） + +**性能**（该 VM 当天有分时 CPU 争抢，绝对 ms/tok 在观测窗口内漂过 +2~3×；表中为紧邻交错对拍的数字，比值才是可靠信号。接受率是确定性 +计数，多轮逐位相同）： + +| workload | 基线 | prompt_lookup | 收益 | accept | avg tok/step | +|---|---|---|---|---|---| +| w1_short_decode | 2.844 / 2.683 ms/tok | 2.791 / 2.692 | ≈打平（门控生效） | 0.781 | 1.68 | +| w4_long_decode | 3.174 / 3.144 | 1.360 / 1.288 | **2.3~2.4×** | 0.982 | 4.49 | +| w7_repetitive_copy | 3.169 / 3.176 | 1.427 / 1.401 | **2.2~2.3×** | 0.958 | 4.03 | +| w5_concurrent_prefill（chunked） | 2.364 | 1.114 | **2.1×** | 0.926 | 4.57 | +| w6_decode_stall（chunked） | 1564.6 tok/s | 2451.3 tok/s | **+57%** | — | — | + +门控的价值有直接对照：无门控的首轮（机器较空闲窗口）w1 从 1.744 +回归到 2.551 ms/tok（-46%，accept 仅 0.55）；加门控后 w1 与基线 +打平（128 个 token 里约 32+32 步投机探测，其余回退图化 decode）， +高命中负载不受影响。 + +**正确性**： + +- 基线确定性成立：同配置连跑两次 w1/w4 输出逐 token 相同。 +- 全 exact：w7（512 tok）、w5（8/8 请求）——高命中负载里正确 token + 的 logit 遥遥领先，数值噪声翻不动 argmax。 +- 非 exact：w1（token 9 分歧）、w4（token 32）、w6（7/9 exact， + 最早 token 1）。定位为 **near-tie argmax 翻转**，非逻辑 bug： + - HF 探针实测分歧位置的 top-2 logit 间隙：w1 = 0.125、w4 = 0.000 + （完全平局）；分歧两侧文本都通顺且语义等价（"对猫的误解" vs + "的好奇心"、"如何进行交流" vs "如何交流"）。 + - 机制：spec 的 verify/融合前向与基线 decode 的 batch 形状不同 + （b×(k+1) eager vs b×1 图化），归约顺序差异产生 O(0.01~0.1) + 的数值噪声，足以翻转平局。w6 的 token-1 分歧是同一机制上移 + 一层：spec 改变各请求进度 → 混排批次组成不同 → 主前向数值微差。 + - 结论：投机采样"数学无损"指分布等价，不是逐位等价（vLLM 同样 + 不保证逐位一致）。验收门槛应理解为：高命中负载必须 exact + （w5/w7 已满足），开放生成负载看分歧点 logit 间隙是否近平局。 +- 反面边界不变：大 batch（w3 bs=32 类）由 batch 阈值直接回退，不回归。 + +**w1 类开放生成要真正获益，需要 EAGLE 头**（接受率 0.55 → 2.5~3.5 +才有净收益），这是后续项；prompt-lookup 的定位是把链路和验收跑通 + +覆盖重复性负载（代码、总结、agent 循环输出）。 + +### 遗留：verify/融合前向的 CUDA graph 录制（③b 的 C++ 半） + +融合前向目前 eager（sample_all_positions=True 被 +`csrc/engine/rank_worker.cpp:428` 的门禁挡在图外）。录图改动点 +（已定位，待有 GPU 构建环境时实施）： + +1. `rank_worker.cpp:428` 放宽 `!sample_all_positions` 门禁，让 + all-position 输入可进编译器； +2. `paged_compiler.cpp` 新增按 (batch, k+1) 键的 verify 图表（tuple + key 仿 static_batching_compiler.hpp:30），compile() 新增录制分支： + pack_i64/pack_i32 缓冲按 k+1 放大，录制输入置 + `sample_all_positions=true`； +3. `get_compiled` 的 decode-only 检查（`input_ids->size(1) != batch`） + 改为分支：命中 (b,k+1) 表走 verify 图； +4. hidden_states 不进图输出时 `forward_raw`（infer_engine.py）需容忍 + null（runner 只消费 output_ids）； +5. 隐性依赖：录制时冻结的 host 标量 max_query_length=k+1 没问题， + max_sequence_length 随 ctx 变化——需确认 FA varlen 路径不消费它做 + 内容相关分支； +6. 配置链：k 需在编译期可知（LLM→EngineConfig→InferEngine→pybind→ + RankWorker→PagedCompiler），或在首次遇到该形状时 lazy 录制。 + +预期收益：融合前向图化后，0.6B decode 每步成本回到图化单前向量级， +投机的理论收益（w1/w4 类小 batch 延迟敏感负载 2× 级）才完全兑现。 + +--- + +## v15（2026-09-01）：bs=1 decode step 的 kernel 级归因 + +动机：w1 实测（0.6B 1.74ms/tok）对带宽 roofline（~1.2GB/step ÷ +1.79TB/s ≈ 0.67ms）看似只有 39% 达成率，需定位缺口再定优化方向。 + +### 方法 + +nsys `-t cuda --cuda-graph-trace=node` 采集 w1（graph 模式）； +`graphNodeId IS NOT NULL` 过滤图内 kernel，按 >100µs 时间间隙切 +replay（两模式 342 kernel/replay，稳态 254 个 replay），逐步中位数 +统计。分析脚本 `/root/step_breakdown{,2..5}.py`（5090），profile： +`/root/prof_w1_06b_n.nsys-rep`、`/root/prof_w1_17b.nsys-rep`。 + +### 0.6B decode step（wall 1390µs，busy 1357µs，图内 idle 仅 2%） + +| kernel | µs/step | n/step | µs/call | 推断调用点 | 字节 | 带宽达成 | +|---|---|---|---|---|---|---| +| gemvx g(768) | 277.7 | 28 | 9.92 | gate_up [1024→6144] | 12.6MB | 1.27 TB/s (71%) | +| gemvx g(256) | 272.4 | 56 | 4.86 | o + down (N=1024) | 4.2/6.3MB | ~1.08 TB/s (60%) | +| gemvx g(512) | 203.6 | 28 | 7.27 | qkv [1024→4096] | 8.4MB | 1.15 TB/s (64%) | +| gemvx g(18992) | 189.9 | 1 | 189.9 | lm_head [1024→151936] | 311MB | **1.64 TB/s (92%)** | +| add_rmsnorm | 104.7 | 56 | 1.87 | residual+norm | ~KB 级 | 纯延迟 | +| paged splitkv cta+combine | 155.7 | 56 | — | decode attention | KV 小 | — | +| pagedCaching | 42.3 | 28 | 1.51 | KV 写入 | 小 | 纯延迟 | +| rmsNormRope ×2 | 70.7 | 56 | 1.2~1.3 | q/k norm+rope | 小 | 纯延迟 | +| SwiGLU | 33.8 | 28 | 1.21 | MLP 激活 | 小 | 纯延迟 | +| sampling/embedding | ~6.5 | — | — | 含伴随图 argmax | — | — | + +### 1.7B decode step（wall 2815µs，busy 2785µs） + +GEMV 合计 2353µs / 3.44GB = **1.46 TB/s (82%)**：gate_up 89%、 +qkv 75%、o+down 71%、lm_head 92%。小 kernel 合计 ~432µs (15%)。 + +### 归因结论（修正粗估） + +1. **GEMV 并不烂**：decode 用 cuBLAS gemvx，大矩阵贴近实测峰值 + （lm_head 92% ≈ 实用上限 1.65TB/s）；小矩阵（o/down/qkv）60~75%， + 是 launch ramp + 矩阵太小的物理下限，可榨空间 ~10~16%（0.6B）。 +2. **0.6B 的三个真实资金池**（按大小）： + - 小 kernel 延迟 ~413µs（30%）：225 个 1~2µs 级 kernel 的 dispatch + 地板。融合方向：add_rmsnorm 进 GEMV epilogue、pagedCaching 进 + rope/attention、SwiGLU 并进 down epilogue。预期回收 ~200µs。 + - 图外 CPU ~0.35ms/step（20%，bench 1.74ms − 图 1.39ms）：打包 + H2D + graphLaunch + Python 调度。结构性解法 = async scheduling + （step N 执行时跑 step N+1 的调度/更新，vLLM 0.9 同款）。 + - GEMV 70→85~92%：~150~220µs，需自研 split-K GEMV 或 cuBLASLt + 启发式调优，收益上限最小。 +3. **1.7B 空间更小**：GEMV 已 82%，小 kernel 15%——bs=1 小模型的 + 常数项随模型变大被摊薄，**kernel 微优化路线天花板清晰可见**。 +4. 真正的数量级杠杆仍在 **FP8 权重/KV**（带宽减半，与上述乘算； + lm_head 一项即占 14% step，FP8 后立省 ~95µs@0.6B）与**投机采样** + （decode 2× 级）。两者基建已在仓库（quantization_method / + kv_cache_k_scale / draft_model_path / speculative_cache_ops）。 + +--- + +## v14（2026-09-01）：v13 5090 复核 + w6 decode-stall 负载 + +动机：v13（chunked prefill + prefill/decode 混排）提交时仅有本机 +Python 单测，GPU 侧未验证；且验收负载 w5 是「同时到达」形态，测 +不出混排的真实收益（decode 流被长 prefill 队头阻塞的场景）。 + +### 部署与正确性（5090，hybrid+graph，64blk，prefix caching 开） + +远端 `/root/src/InfiniLM-hybrid` 为 v12 等价树（md5 逐文件比对）， +v13 纯 Python，推 6 文件复用 C++ binary。llm.py 上 v11 遗留的 +stepprof 插桩 patch 被 v13 版覆盖（已备份 llm.py.stepprof.bak）。 + +- gate 探针确认：`INFINILM_ENABLE_CHUNKED_PREFILL=1` 时 + `enable_chunked_prefill=True`（注意 bench 下 logging 不出 INFO， + 「Chunked prefill enabled」日志不可见，只能用探针确证）。 +- 0.6B 全负载矩阵（w1~w5，budget=1024 强制切块 vs 关 vs 默认 + budget）：三组两两对拍均 43/43 请求逐 token exact。 +- 1.7B w2+w5（budget=1024，on vs off）：9/9 exact。 +- 远端 test_chunked_prefill.py 4 用例全绿。 + +### w6_decode_stall：decode 长流 + 中途注入长 prefill + +形态：8 条短 prompt 各 decode 512 tok，第 64 步注入 1 条 ~6.5k tok +prompt（nonce 头避开 prefix caching），逐步计时。budget=1024。 + +| 指标 | 0.6B 关 | 0.6B 开 | 1.7B 关 | 1.7B 开 | +|---|---|---|---|---| +| decode 步均值（注入前） | 2.4~2.5ms | 2.6~2.8ms | 4.0~4.1ms | 4.1ms | +| 注入后最大步长 | 75.9/74.9ms | 21.8/19.4ms | 138.6ms | 31.3ms | +| 注入后 p90 步长 | 2.9ms | 3.2ms | 4.4ms | 4.4ms | +| 注入请求 TTFT | 75.9ms | 106.1ms | 138.6ms | 185.2ms | +| e2e | 1.56s | 1.61/1.62s | 2.40s | 2.41s | + +结论：chunked 把 decode 流最坏 ITL 尖峰削 **3.5×（0.6B）/ 4.4× +（1.7B）**，p90 几乎不动；代价是注入请求 TTFT +40%/+34% 与 e2e +~3%（切块后小 kernel + 混排 eager 步的开销）。w5（同时到达形态） +ABBA 下 e2e 反而亏 ~5%（无 decode 流量可保护）——chunked 的价值 +在在线服务的延迟平稳性，不在离线批量吞吐。 + +### 1.7B w6 的 on/off 分歧归因（非逻辑 bug） + +on vs off 出现 5~6/9 exact、token 80~108 分叉。排查链：off×3 两两 +9/9 exact；on×2（同配置）9/9 exact（on 模式确定性）；关 v12 ctx +路由（INFINILM_DECODE_CTX_THRESHOLD=999999）后 on vs off 仍 +5/9@80；路由开/关的 on 两跑互对 5/9@80。结论：混排步中 decode +token 走 eager varlen（而非 decode 图），叠加 6.5k ctx 触发 v12 +FA 路由，kernel 归约顺序差的 epsilon 被贪心放大——与 v9/v11 已 +接受的 kernel 切换分歧同类。0.6B 全 exact；1.7B w5（纯 prefill +切块，无 decode 混排）9/9 exact 说明切块逻辑本身无数值偏差。 + +数据归档：`results/5090_v13_chunked/`（23 份：全矩阵×4、w5 +ABBA×4、1.7B w2/w5×2、w6 0.6B×4、w6 1.7B×5）。远端遗留 +run_bench.py / probe_gate.py / llm.py.stepprof.bak。 + +--- + +## v12（2026-08-29）：decode 按 ctx 长度自适应路由 —— w2@1.7B 倒退根治 + +动机：v9 遗留「hybrid 在长 ctx decode 输给 FA kvcache」（w2@1.7B 0.75 +vs 0.62，交叉点 1k~3.3k 未定）。本轮先做 ctx 扫描把交叉点定死，再落 +路由。 + +### ctx 扫描（新增 `bench.py --ctx-sweep`，单请求 decode 128 tok，eager） + +`_BASE_PARA×k` ≈ 81k token 提示词，两后端 prefill 同为 FA2，e2e 差即 +decode kernel 差。ms/tok（越低越好）： + +| ctx | 0.6B hybrid | 0.6B FA | 1.7B hybrid | 1.7B FA | +|---|---|---|---|---| +| 162 | **2.79** | 4.25 | **3.39** | 4.24 | +| 1296 | **2.80** | 4.24 | **4.25** | 4.33 | +| 1944 | **3.13** | 4.35 | 4.70 | **4.38** | +| 2592 | **3.51** | 4.32 | 5.11 | **4.45** | +| 3240 | **3.94** | 4.28 | 5.57 | **4.55** | +| 3888 | 4.36 | **4.29** | 5.96 | **4.58** | +| 5184 | 5.26 | **4.42** | 6.92 | **4.73** | + +两个结构性发现： + +1. **FA kvcache 随 ctx 几乎平坦**（1.7B 斜率 ~0.1µs/tok），splitkv 线性 + 增长（0.6B ~0.49、1.7B ~0.70 µs/tok）——FA kernel 内对 KV 长度方向 + 并行得好，splitkv 在 bs=1 并行度不够。交叉点 **0.6B≈3.4k / + 1.7B≈1.5k**，随模型几何移动，不能写死全局常数。 +2. **FA decode 的固定底线与模型大小无关**（0.6B/1.7B 同為 4.24ms）—— + 短 ctx 下 FA 路径是固定开销（eager launch/边界）主导而非带宽主导。 + 上表是 eager 数据；graph 模式吃掉固定开销后交叉点会移动，表内阈值 + 用于 graph 选图是偏保守的（1.7B w2 实测改善即在 graph 模式拿到）。 + +### 实现（本分支工作区） + +- `attention_layer.{hpp,cpp}`：`decode_fa_ctx_threshold()`——env + `INFINILM_DECODE_CTX_THRESHOLD` 优先；否则按 (hidden, layers, + kv_heads) 查实测表：Qwen3-0.6B(1024/28/8)→3400、 + Qwen3-1.7B(2048/28/8)→1500；未测几何 SIZE_MAX（**不路由=现状**）。 + `HybridAttentionImpl::forward` decode 分支： + `max_sequence_length > 阈值` → `flash_->forward`（mha_kvcache_）； + 两 kernel 读同一块 BSHD paged cache（splitkv 走 permute 视图）， + **切换零拷贝**。 +- `infer_engine.cpp`：decode 步也填 `max_sequence_length`（host 侧 max + over CPU int32 tensor，无同步开销；非 CPU/I32 返回 0=不路由）。 +- `paged_compiler.{hpp,cpp}`：路由启用时每个 batch 档**录双图**——假 + msl=1 录 splitkv 版、假 msl=阈值+1 录 FA 版（录制时算子只登记不执行， + 假值只冻结 kernel 选择，replay 读真实 pack_i32）；`get_compiled` 按 + host 侧 max ctx 选图（非 CPU/I32 保守走短 ctx 版）；采样伴随图每变体 + 各录一份，`get_sampling_compiled` 跟随 `last_served`。两变体共享 + block_tables_holder_，block table 更新一处生效。 + +### 正确性与性能（5090，hybrid+graph，--dump-outputs） + +- **0.6B 全矩阵对 v11 构建逐 token 全 exact**（阈值 3400，w2 峰值 ctx + 3368 不触发路由；双图重构未改变短路径）。 +- **1.7B**：w1/w2/w4 对 v10 dump 全 exact、w3 31/32@109（既有批噪声底 + 同类）；**w2 路由后输出与 flash-attn 后端逐 token 全 exact**（graph + 与 eager 均验证）——路由前后 kernel 归约顺序差在该 prompt 的 128 + token 内未翻牌，且路由结果确实落在 FA kernel 上。 +- **性能**：1.7B w2 e2e **0.74→0.55s（-26%**，graph）/ 0.75→0.66s + （eager），优于 flash-attn 后端自身的 0.64s（eager）；w1/w3/w4 与 + 0.6B 全负载持平 v11（噪声内）。 + +### 边界与后续 + +- 阈值是 bs=1 eager 定的；**bs×ctx 第三象限（大 batch × 长 ctx）无数据** + ——bs=32 时 batch 方向自带 32 倍并行度，splitkv 短板被补，交叉点可能 + 大幅后移，agent 高并发场景「谁赢」待网格扫描（注意显存：两模型 KV + 均 112KB/token，32GB 上 bs32×16k≈57GB 不可行，网格上限 bs32×8k 或 + bs8×16k）。 +- 与 vLLM 的对照最大只测到 3.4k ctx；agent 主战场（10k+ 多并发）的 + 立项差距完全未知。 +- 中途切 kernel 引入与「换后端」同量级的 ulp 扰动（v9 已记录），greedy + 可能翻接近的 top-2；对未测几何默认关闭，零回退风险。 + +### 待办更新 + +- [x] ~~decode 按 ctx 长度自适应路由~~ —— 本轮落地（w2@1.7B 0.74→ + 0.55s），阈值为 bs=1 eager 实测表 + env 覆盖，未测几何不路由。 +- [ ] bs×ctx 网格扫描（bs 1/8/32 × ctx 1k/4k/8k + bs8×16k，graph 模式） + ——回答 agent 象限 splitkv/FA 谁赢，并校准 graph 模式阈值。 +- [ ] 长 ctx 多并发对 vLLM 对照（agent 主战场的立项差距）。 +- [ ] hybrid+graph 设为默认的评估（沿用 v10）。 +- [ ] （可选）5060 Ti 复核；hybrid 推广其他模型族；flashinfer 采样器。 + +数据归档：`results/5090_graphopt/`（0.6B `..._0829_154455`、1.7B +`..._154513`、1.7B FA 参考 `..._154523`、eager 路由 `..._154645`）。 + +--- + +## v11（2026-08-29）：图外开销优化落地 —— 打包 H2D + 采样伴随图 + +动机：v10 归因指出 graph 模式每步仍有图外开销——6 个小输入各一次 +H2D copy、采样段 32 请求 × 3 launches（cub ArgMax×2 + cast）。本轮把 +这两处削掉并实测收益。 + +### 改动（本地未提交，csrc/engine/{compiler/*,rank_worker.cpp}） + +1. **打包 H2D**：`make_decode_input` 把 input_ids / position_ids / + slot_mapping 打进一条连续 int64 缓冲（pack_i64），total_seq_lens / + input_offsets / cu_seqlens 打进一条连续 int32 缓冲（pack_i32）， + 图输入做成两条缓冲的视图；`get_compiled` 快速路径两次 memcpyH2D + 代替原 5+1 次 copy_from；block table 在 block_per_req==编译宽度时 + 跳过 -1 填充。 +2. **采样伴随图**：`compile()` 对 b≤64 的每个 decode 图录一个 + companion graph（逐请求 argmax：cub DeviceReduce ArgMax + 索引 + cast），greedy（top_k==1 或 temperature==0）且 decode 图命中时 + replay 它，代替每步 32×3 次 launch。kill switch： + `INFINILM_DISABLE_SAMPLING_GRAPH=1`；调试： + `INFINILM_DEBUG_SAMPLING=1`（命中时打印一次）。 + +### 关键坑（排查记录，后续接图的人必读） + +infinicore 的 `Graph` **不是裸 CUDA stream capture，是算子序列录制**： +只有经 `INFINICORE_GRAPH_OP_REGISTER_*` 注册的算子类才会在录制时进入 +`op_list_`（录制模式下算子只登记不执行；`instantiate()` 先跑 5 遍 +warmup 再按 capture 安全性切段捕获）。`random_sample` 是纯 dispatcher +注册（`random_sample_infiniop.cc` 无 graph 钩子），greedy 路径又被 +`tryGreedyWithInfiniOps`（infiniops Argmax）在 dispatcher 之前拦截—— +直接 `startGraphRecording() + random_sample_` 录到的是 **operators=0 +的空图**，replay 等于空操作，读到的是编译期陈旧 logits 的 argmax 结果 +（编译用全零退化输入，total_seq_lens=0 → logits 为 NaN 垃圾 → 越界 +token id，`tokenizer.decode` 抛 OverflowError）。修复：自定义 +`SamplingLoopOperator`（`GraphOperator` 子类）包装整个采样循环， +`context::addGraphOperator` 手动入列，由 `instantiate()` 统一 warmup+ +捕获。第二坑：包装 lambda 必须**按值捕获**张量（操作符比 compile() +作用域活得久，且持锁张量内存）。 + +### 正确性 + +hybrid+graph 全矩阵 `--dump-outputs`（`..._0829_145012.json`）与修复前 +已验证版本（`5090_graphopt/..._0829_141752.json`）**w1/w2/w3/w4 逐 +token 全 exact**（32/32 请求全对）。采样图与 eager argmax 语义一致。 + +1.7B 复验（w3_batch32，同构建采样图 ON vs kill switch）：31/32 exact +@token90；同配置两个进程互跑（采样图 ON vs ON）也是 30/32 @token90 +——偏差来自 forward 自身的运行间抖动（v10 记录的 1.7B bs32 批噪声底 +同类），与采样路径无关。1.7B w3 吞吐 6214~6269 tok/s，与 v10 graph +基线（6211）持平。 + +### 性能（同构建 ABBA，kill switch 交替，0.6B w3_batch32） + +| 轮次 | A=采样图 ON | B=eager 采样 | +|---|---|---| +| 1 | 11214 | 9530 | +| 2 | 9547 | 9499 | +| 3 | 9526 | 10275 | + +中位数 9547 vs 9530 tok/s —— **差 ~0.2%，在 5090 热双态噪声内**。 +STEP_PROF 分段（A vs B，step 64/128 均值）：forward 段 1.67 vs 1.68ms +（稳定省 ~10µs/step），sched/post/gap 不变。 + +### 结论 + +1. **图外开销在 graph 模式本就被异步隐藏**：采样 launch 在 forward 图 + 仍在 GPU 执行时就已发出（gap≈0），削掉它 e2e 不动。本轮改动保留 + （每步省 ~10µs CPU + 消除采样段分配搅动，对更大 batch 或 CPU-bound + 场景仍有意义；正确性已验证，且有 kill switch）。 +2. **瓶颈回到 forward 图本身**：bs32 每步 ~1.6~2.1ms 是图内 GPU 时间。 + 下一步方向：图内 kernel 归因（GEMM / hybrid attention / mamba 段在 + bs=32 的占比）、以及与 vLLM 的调度层差距（continuous batching 的 + step 内重叠），图外已无可削。 + +--- + +## v10（2026-08-29）：5090 全栈对照 —— InfiniLM(hybrid) vs vLLM 0.28 + +动机：v8/v9 遗留的「立项决策最后一块」——此前全部 vLLM 数据出自 +5060 Ti（WSL2）与 8GB 病态机，vLLM 从未在健康平台跑过。本轮在 5090 +建起 vLLM 独立 venv 并跑全负载矩阵，InfiniLM 侧用 v9 的 hybrid 基线。 + +环境： + +- **vLLM 0.28.0 + torch 2.13.0+cu130**（`/root/.venv-vllm`，uv 0.9.9 + 经阿里云镜像安装）：v1 默认（CUDA graph on、prefix caching on、 + chunked prefill on），gpu_memory_utilization=0.85，attention 自动选 + **FLASH_ATTN**（vllm 内置 FA2，sm_120 候选序列首位）。两个坑见文末 + 「5090 装 vLLM 工程记录」。 +- **InfiniLM hybrid**：64blk、no-graph（沿用 v1~v9 对比惯例），另加测 + 一组 `--enable-graph`（hybrid×graph 首次实测，顺带验证兼容性）。 +- 同树同机同日：hybrid ×2 轮、hybrid+graph ×1 轮、vLLM ×2 轮，全部 + `--dump-outputs`；hybrid 同日两轮与 v9（0827）数值互验一致。 + +### 性能（两轮均值；vLLM Δ 列为对 hybrid / 对 hybrid+graph） + +**0.6B**（hybrid 峰值 4.1GB，vLLM 28GB@0.85 档）： + +| 负载 | hybrid | hybrid+graph | vLLM | Δ vs hybrid | Δ vs hy+graph | +|---|---|---|---|---|---| +| w1 单请求 decode | 2.92 ms/tok | 1.76 | **1.63** | **-44%** | -7% | +| w2 长 prefill + 128 decode | 0.53 s | 0.53 | **0.26 s** | **-50%** | **-50%** | +| w3 batch32 | 8621 tok/s | 9479 | **15510** | **+80%** | **+64%** | +| w4 长 decode | 2.93 ms/tok | 2.03 | **1.60** | **-45%** | **-21%** | + +**1.7B**（hybrid 峰值 7.2GB，vLLM 27.6~28.2GB）： + +| 负载 | hybrid | hybrid+graph | vLLM | Δ vs hybrid | Δ vs hy+graph | +|---|---|---|---|---|---| +| w1 | 3.30 ms/tok | 3.20 | **3.10** | **-6%** | -3% | +| w2 | 0.75 s | 0.74 | **0.48 s** | **-36%** | -35% | +| w3 | 6162 tok/s | 6211 | **8767** | **+42%** | +41% | +| w4 | 3.45 ms/tok | 3.50 | **3.06** | **-11%** | -12% | + +（加载耗时参考：hybrid ~6s；vLLM 冷启 46s / 编译缓存命中后 18~22s。） + +### 结论 + +1. **健康平台上 vLLM 全面领先**，5060 Ti 的「decode 持平、仅 prefill + 差 1.3~1.5×」结论不能外推。差距结构： + - **w3（batch32 吞吐）差距最大**（+42%~+80%），其次 **w2** + （-36%~-50%）——指向批处理调度、chunked prefill 与全图捕获的 + decode 段，而非单个 kernel； + - 单请求 decode（w1/w4）1.7B 仅差 6%~11%（hybrid 无 graph 对 + vLLM 有 graph），0.6B 差 44%~45%（小模型 launch-bound 占比高）。 +2. **graph 不是差距主因**：hybrid+graph 在 0.6B decode 上 -31%~-40% + （0.6B launch 开销占比确实大），但补不回 w3/w2 的差距（仍落后 + 41%~64%）；1.7B 上 graph 收益在噪声内（±3%）。hybrid×graph 首次 + 实测无回退，输出与 no-graph 对照 0.6B 四负载全 exact、1.7B 仅 + w3 30/32（@93，既有批噪声底同类），作为默认配置可行。 +3. **v9 的内部结论不受影响**：hybrid 仍是 InfiniLM 三后端中最优单配置； + 但对 vLLM 的立项差距在健康平台上重新打开，下一轮优化目标应转向 + w3/w2 的归因（见待办）。 + +### w3 差距归因（nsys，0.6B，同日补抓） + +方法：nsys 2024.6 抓 `--only w3_batch32` 全 trace，sqlite 导出后按 w3 +窗口切片聚合（脚本 VM `/root/slice_w3.py`，trace 存 `/root/prof/v10/`）。 +w3 = 32×~15tok prefill + 128 步 batch32 decode，**decode 段主导**。 +注意：nsys 对 CUDA graph replay 内的 kernel 记录不全（vLLM 侧 window +busy 仅 9%，物理上不可能——0.6B×bs32 单权重读取就需 ~0.7ms/step）， +vLLM kernel 级数字仅作参考，wall/launch 数可信；hybrid 侧记录完整。 + +**hybrid no-graph**（w3 窗口 570ms）： + +- **每 decode step ~435 次 kernel launch**：gemm 4 次/层（qkv/o/gate_up/ + down 已是融合 gemm）+ 2 次 cublas splitk 伴随 + add_rms_norm 2 + + rms_norm_rope 2 + swiglu 1 + cache 1 + attn 2 + ~4 次小 elementwise； +- **GPU busy 仅 306ms（54%）**；未 profile 口径 wall 3.7ms/step vs busy + ~2.4ms/step → **~35% 的 wall 是 launch 空泡**； +- kernel 时间大头是 gemm（窗口内 248ms）：bs=32 的 gemm 平均 13.7µs， + 已在 latency floor，kernel 本身无多少油水。 + +**hybrid+graph**（e2e 0.43s，仅 +10%）： + +- `PagedCompiler` 对 bs=1..64 逐档捕获 decode-only 图，bs=32 在列—— + w3 decode 确实走了 graph replay，launch 空泡大头已消除； +- 但每步图外仍有残余开销：5×D2D 输入拷贝 + 每次 replay 前的 reset + memset（`paged_compiler.cpp:165` 注释自承「still pays a memset before + every graph replay」）+ 采样 kernel + Python 调度，合计 ~0.7ms/step； +- 图内 GPU 时间 ~2.4ms/step 不变 → wall 3.1ms/step，**从 launch-bound + 转为 kernel-time-bound**。 + +**vLLM**（e2e 0.26s，2.0ms/step）：整步单次 graph replay + inductor +combo/fused kernel（triton 融合 norm/silu、combo gemm），图外仅少量 +采样 kernel。 + +**归因结论**：w3 的 +80% 差距 = ①launch 空泡（hybrid 有 graph 但图外 +残余未削）+ ②图内 kernel 时间本身（bs=32 gemm 贴 latency floor，靠 +数量取胜——vLLM 融合度更高）。后续方向：a) 削图外 per-step 开销（合并 +5 次输入拷贝、去掉 per-replay memset、采样入图）；b) 评估跨 gemm 的 +进一步融合/grouped gemm。w2 的账独立于本次（v9 已定位：3.3k ctx 下 +decode 段变慢 + 无 graph），需要补 trace 时另抓。 + +### 正确性（`compare_outputs.py`，双轮互验） + +- **同引擎跨轮噪声底**：hybrid 0.6B 全 exact、1.7B w3 30/32（@91， + 与 v9 底同位置）；vLLM 自身 w3 也非确定（0.6B 27/32、1.7B 25/32， + 批处理数值噪声）。 +- **跨引擎**：1.7B w1/w2 全 exact，w4 @190 与 v4/v5 历史分叉位置 + 完全重合；0.6B w2 exact，w1 @0 首 token 翻牌（hybrid「这个问题可能 + 引发一些人对猫的误解…」连贯，vLLM 落入重复循环——0.6B 贪心退化的 + 常见形态，双侧均为模型自身质量范围内输出），w4 @32 与 v9 + hybrid-vs-fa 的 @32 重合；w3 22~27/32 落在双方噪声底交集内。 +- 判定：**无功能性错误**，全部为 bf16 并列翻牌 × 自回归放大。 + +### 5090 装 vLLM 工程记录（复现要点） + +- **无外网**：NGC 镜像自带 `/etc/pip.conf` + `/etc/xdg/pip/pip.conf` + 配了 pypi.org 主 index + `pypi.ngc.nvidia.com` extra-index，`-i` 只 + 覆盖主 index，包查询全卡在 ngc 的 TCP 超时——已将两个文件挪为 + `.bak`。pip 单连接仍被限速 ~200KB/s（同文件 curl 直连 11MB/s， + aliyun/tuna 同速），改用 **uv**（并发下载）安装。 +- vLLM 0.28 运行需 `CUDA_HOME=/usr/local/cuda-12.8` + venv 内 ninja + （README 已有此条）。 +- **flashinfer 采样器 JIT 在 sm_120 上要求 nvcc≥12.9**（编译 + `compute_120f`），VM 工具包为 12.8 → 引擎 warmup 直接报 + 「FlashInfer requires GPUs with sm75 or higher」（TARGET_CUDA_ARCHS + 为空后的误导性报错）。对策：`VLLM_USE_FLASHINFER_SAMPLER=0` 走 + 原生采样器；贪心解码（temperature=0/top_k=1）不受影响。attention + 不受影响——Qwen3 在 sm_120 默认 FLASH_ATTN。 +- 堡垒机 scp 逐文件重新鉴权，批量拉取需在 VM 侧先打 tar 包。 + +### 待办更新 + +- [x] ~~w3/w2 对 vLLM 差距的 nsys 归因~~ —— w3 已由 v10 补抓完成： + launch 空泡（图外残余 per-step 开销 ~0.7ms/step）+ 图内 kernel + 时间（bs=32 gemm 贴 latency floor）。派生优化项：削图外开销 + （合并输入拷贝/去 per-replay memset/采样入图）与 gemm 融合度。 + w2 归因沿用 v9 结论（长 ctx decode + 无 graph），需要时补 trace。 +- [ ] hybrid+graph 设为默认的评估（0.6B decode -31%~-40%、四负载无 + 回退；需补更多模型的正确性对拍）。 +- [x] ~~decode 按 ctx 长度自适应路由~~ —— 已由 v12 落地（w2@1.7B + 0.74→0.55s，阈值查表 + env 覆盖，未测几何不路由）。 +- [ ] （可选）5060 Ti 复核 hybrid；hybrid 推广到其他模型族（沿用 v9）。 +- [ ] （可选）flashinfer 采样器恢复：pip 侧装 `nvidia-cuda-nvcc-cu13` + 或升级工具包至 ≥12.9；当前 FLASH_ATTN + 原生采样器已够用。 + +数据归档：`results/5090_vllm/`（10 份：hybrid×2、hybrid+graph×1、 +vLLM×2，双模型）；v9 的 15 份同步拉回 `results/5090_hybrid/`。VM 侧 +venv `/root/.venv-vllm`，运行环境变量： +`HF_HUB_OFFLINE=1 CUDA_HOME=/usr/local/cuda-12.8 VLLM_USE_FLASHINFER_SAMPLER=0`。 + +--- + +## v9（2026-08-27）：hybrid 后端落地 —— prefill 走 FA2 / decode 走自研 splitkv + +动机：v8 发现 5090 上 FA2 只赢 prefill、decode 慢自研 splitkv 23%~54%， +「flash-attn 设为默认」不成立。本轮把分离路由做成一个后端： +`--attn-backend hybrid`。 + +实现（全部在本分支工作区，零新增文件）： + +- `attention_backends.hpp`：新增 `AttentionBackend::HYBRID` + 字符串解析 + `"hybrid"`；Python 侧 `llm.py`/`bench.py` 透传（pybind 原样走 + `parse_attention_backend`，无需其他改动）。 +- `attention_layer.{hpp,cpp}`：新增 `HybridAttentionImpl`（内含一个 + `FlashAttentionImpl`）。`is_prefill` 判定与既有 impl 相同（展平 paged + 模式下纯 decode 步恰为每序列 1 个 query token）；prefill/混合 batch → + `flash_->forward`（mha_varlen），纯 decode → 复用 + `flash_->do_kv_cache_update`（paged_caching_ 经 permuted view 写 BSHD + cache）后，以 `k_total->permute({0,2,1,3})` 的逻辑 BHSD 视图直接调 + `paged_attention_`。 +- `kv_cache.cpp`：HYBRID 与 FLASH_ATTN 同走 BSHD 物理布局； + `infinilm_model.cpp`：HYBRID 落入 paged cache 分配分支（Qwen3 因此 + 走 `forward_paged_`，自动带上 v5 的 rms_norm_rope 融合）。 + +strides 安全性（hybrid 成立的命门，InfiniCore 侧核实）： + +- `paged_attention` descriptor（`info.h`）从 tensor descriptor 逐维取 + stride，仅要求 `stride(3)==1`（head_dim 连续）——permuted BSHD 视图 + 满足（stride(0)=BS·H·D, stride(1)=D, stride(2)=H·D, stride(3)=1）。 +- 实际使用的 `kernel_v2.cuh`（所有 nvidia launcher 都 include 它）按 + `k_row_stride` 逐 token 寻址；旧 `kernel.cuh` 里 `token*HEAD_SIZE` 的 + 硬编码路径未被任何 launcher 引用。 +- infinicore wrapper 的 FA 快速通道（`canUseFlashAttention`）两条后端 + 条件一致，且该通道本来就总被喂 permuted 视图；本平台实证 decode 走 + 的是 kernel_v2（见下，w1/w4 hybrid≡paged 而远快于 FA）。 + +性能（5090，64blk，no-graph，同树同 binary 交错 ABBA）： + +0.6B hybrid(A) vs flash-attn(B)，e2e 秒： + +| 负载 | A1 | A2 | B1 | B2 | e2e Δ | +|---|---|---|---|---|---| +| w1 单请求 decode | 0.33 | 0.36 | 0.51 | 0.53 | **-33.7%** | +| w2 长 prefill + 128 decode | 0.51 | 0.53 | 0.56 | 0.57 | **-8.0%** | +| w3 batch32 | 0.39 | 0.42 | 0.65 | 0.64 | **-37.2%** | +| w4 长 decode | 2.68 | 2.81 | 4.27 | 4.33 | **-36.2%** | + +0.6B hybrid(C) vs paged-attn(D)，e2e 秒： + +| 负载 | C1 | C2 | D1 | D2 | e2e Δ | +|---|---|---|---|---|---| +| w1 | 0.35 | 0.36 | 0.35 | 0.36 | 持平 | +| w2 | 0.53 | 0.53 | 0.64 | 0.64 | **-17.2%** | +| w3 | 0.45 | 0.46 | 0.44 | 0.47 | 持平 | +| w4 | 2.85 | 2.93 | 2.87 | 2.89 | 持平 | + +1.7B（HY/FA 为 ABBA，PG 单跑 + v8 复测值 0.86s 佐证）： + +| 负载 | HY A1/A2 | FA B1/B2 | PG | HY vs FA | HY vs PG | +|---|---|---|---|---|---| +| w1 decode | 3.26/3.26 ms | 4.20/4.37 ms | 3.23 ms | **-22.5%** | 持平 | +| w2 e2e | 0.75/0.75 s | 0.61/0.63 s | 0.84 s | **+19%（倒退）** | **-10.7%** | +| w3 e2e | 0.66/0.66 s | 0.61/0.68 s | 0.66 s | 持平 | 持平 | +| w4 decode | 3.54/3.57 ms | 4.24/4.38 ms | 3.55 ms | **-18.4%** | 持平 | + +结论： + +1. **设计目标达成**：hybrid decode ≡ paged decode（w1/w4 逐项持平）， + prefill 保留 FA2（w2 对 paged -17%）；0.6B 上对 flash-attn 全面 + -8%~-37%，把 v8 发现的 FA2 decode 回退全部吃回。 +2. **w2@1.7B 是唯一倒退项**（0.75 vs FA 0.62，可复现、非漂移： + A1≡A2/B1≡B2）。分解：5090 上 1.7B 3240-token prefill 仅 ~70ms， + w2 e2e 由 decode 段（ctx 3240→3368）主导；该 ctx 下 paged kernel + 不再赢 FA kvcache——而 w4（ctx≤~1k）hybrid≡paged 仍快 FA 18%。 + **交叉点在 1k~3.3k ctx 之间**（0.6B 在 3.3k 处 hybrid 仍赢 FA， + 与 kv-head 数/几何相关）。v8 的「FA decode 慢 23~54%」因此应限定为 + 短/中 ctx。精确分相计时留待 nsys。 +3. 后续方向：decode 按 ctx 长度自适应选 kernel（短 ctx→paged splitkv， + 长 ctx→FA kvcache），vLLM 式调度；当前 hybrid 已是严格优于 + paged-attn 全负载、优于 flash-attn 3/4 负载的单配置。 + +正确性： + +- **决定性**：0.6B A1≡A2、C1≡C2 四负载全 exact；1.7B A1≡A2 仅 w3 + 31/32（@91，batch 调度时序噪声，与 v5 观察到的 batch 噪声底同类）。 +- 对两参考后端的交叉对照：0.6B w2 三方全 exact；w1 hybrid-vs-fa @8 + = fa-vs-pg 噪声底 @8；w3 27~28/32 ≈ 底 27/32；w4 单请求 @32 并列 + 翻牌（「人们如何交流」vs「人们如何进行交流」，两侧 1024 token 全程 + 连贯枚举）——hybrid 对 fa 与 pg 同 token 分叉而 fa≡pg，说明 strided + 读引入 ~ulp 级 logit 差（kernel 内不同访存路径的归约顺序差）， + 量级与换后端同。1.7B：w1/w2 对 paged 全 exact；w3 24~26/32 vs 底 + 25/32；w4 @160/@190 = 底 @160。 + +数据归档：`results/5090_hybrid/`（15 份：0.6B ABBA×8、1.7B 单跑×3、 +1.7B ABBA×4）。远程树 `/root/src/InfiniLM-hybrid`（本工作区快照， +非 git）；构建 `XMAKE_ROOT=y INFINI_ROOT=/root/.infini-fa xmake f -c +-m release && xmake build -j32 _infinilm && xmake install _infinilm`； +运行脚本 `/root/run_hybrid.sh`、`/root/run_hybrid17.log`。 + +### 待办更新 + +- [x] ~~decode 按 ctx 长度自适应路由~~ —— 已由 v12 落地（交叉点经 + ctx 扫描实测：0.6B≈3.4k / 1.7B≈1.5k，bs=1 eager）。 +- [x] ~~InfiniLM-vs-vLLM 健康平台全栈对照~~ —— 已由 v10 完成：5090 上 + vLLM 0.28 全面领先 hybrid（0.6B decode -44%、w3 +80%；1.7B + decode -6%~-11%、w3 +42%），新归因待办见 v10。 +- [ ] （可选）5060 Ti 上复核 hybrid（v4 平台上 FA decode 无回退， + hybrid 预期与 paged 持平）。 +- [ ] hybrid 推广到其他模型族（当前仅 Qwen3 paged 路径带融合； + hybrid 本身对任意走 forward_paged_ 的模型可用）。 + +--- + +## v8（2026-08-27）：5090 全栈合流 —— FA2 + rms_norm_rope 融合 + +动机：v4（FA2）与 v5（rms_norm_rope 融合）此前分别在 16GB 5060 Ti 和 +8GB 病态机上验证，从未在同一健康平台上叠加。本轮在 5090 上构建 +FA 版 InfiniCore（aten=y，装到 `/root/.infini-fa`），两份 InfiniLM +扩展（fused/unfused）重建对准该库，跑 flash-attn 后端的 ABBA。 + +构建要点（在 v6 复现要点之上新增）： + +- FA 源码用 **pypi sdist**（`pip download flash-attn==2.7.4.post1 + --no-binary :all:`，阿里云镜像可达）：自带裁剪版 cutlass，够 FA 用。 +- FA 的 cute 头文件包含路径：VM 侧给 `xmake/nvidia.lua` 的 + flash-attn-nvidia target 补了一行 `FLASH_ATTN_ROOT/csrc/cutlass/include` + （上游 FA target 只加 csrc/flash_attn，此前依赖外部 CUTLASS_ROOT—— + 值得上游化）。 +- **不要** export 空的 `CUTLASS_ROOT`（`os.getenv` 返回 "" ≠ nil，会误开 + ENABLE_CUTLASS_API，scaled_mm 在裁剪版 cutlass 上编译不过）。 +- FA bwd kernel 单文件 nvcc 峰值内存大：-j24 触发 OOM（cicc signal 9）， + **-j6** 通过；全量约 25min（含 FA 84 .cu）。 +- aten=y 后 InfiniLM 两份扩展需对准 `/root/.infini-fa` 重建（C++ ABI + 一致性）。 + +### 0.6B flash-attn 交错 ABBA（e2e 秒） + +| 负载 | A1 未融合 | A2 未融合 | B1 融合 | B2 融合 | e2e Δ | +|---|---|---|---|---|---| +| w1 单请求 decode | 0.59 | 0.60 | 0.53 | 0.54 | **-10.1%** | +| w2 长 prefill + 128 decode | 0.65 | 0.68 | 0.58 | 0.58 | **-12.8%** | +| w3 batch32 | 0.68 | 0.69 | 0.65 | 0.65 | **-5.1%** | +| w4 长 decode | 4.73 | 4.79 | 4.34 | 4.38 | **-8.4%** | + +### 1.7B flash-attn 单跑对照 + +| 负载 | 未融合 | 融合 | Δ | +|---|---|---|---| +| w1 单请求 decode | 4.77 ms/tok | 4.33 ms/tok | **-9.3%** | +| w2 长 prefill + 128 decode | 0.71s | 0.63s | **-11.3%** | +| w3 batch32 | 5745 tok/s | 6006 tok/s | **+4.5%** | +| w4 长 decode | 4.81 ms/tok | 4.38 ms/tok | **-8.8%** | + +**融合收益在 FA 栈下比在 paged-attn 栈下更大**（对比 v6：0.6B +-4%~-8% → 本轮 -5%~-13%；1.7B -1%~-4% → -9%~-11%）。机制自洽： +attention 被 FA 压快后，elementwise 链占比上升，融合的相对收益放大。 + +### 同树双后端对照（fused 树，5090 新现象） + +| 负载 | 0.6B paged | 0.6B flash | 1.7B paged | 1.7B flash | +|---|---|---|---|---| +| w1 decode | **2.76 ms/tok** | 4.17 | **3.30** | 4.33 | +| w2 prefill e2e | 0.64s | **0.58s** | 0.86s | **0.63s** | +| w3 batch32 | **9049 tok/s** | 6285 | **5859** | 6006（≈持平） | +| w4 decode | **2.77 ms/tok** | 4.24 | **3.57** | 4.38 | + +**FA2 在 5090 上 decode 慢 23%~54%**（mha_fwd_kvcache 的 sm80 时代 +kernel 在 Blackwell 上效率不佳；v4 在 5060 Ti 上两者基本持平），只赢 +prefill（1.7B w2 -27%）。含义："flash-attn 设为默认后端"在本平台 +不成立；合理方向是 **prefill 走 FA、decode 走自研 splitkv 的分离路由** +（vLLM 即此类设计）。5060 Ti 上该结论需复核（v4 数据是融合前的）。 + +### 正确性（--dump-outputs + compare_outputs.py） + +- FA 栈同 binary 噪声底 = 0（A1 vs A2 四负载全 exact）。 +- fused vs unfused（FA 栈）：0.6B w2/w4 全 exact，w1 @8、w3 28/32； + 1.7B w2 exact，w1 @40、w3 28/32、w4 @190。 +- 关键对照：同树换后端（paged↔flash，代码不变）的噪声底分叉位置 + **与 fused-vs-unfused 重合**（0.6B w1@8、1.7B w1@40、w4@160~190）—— + 融合引入的数值扰动与换一个 attention 实现同量级，非功能错误。 + +数据归档：`results/*_rtx5090_fa_{fused,unfused}.json`（FA ABBA 八轮中的 +四+1.7B 两轮）与 `*_rtx5090_paged_fused.json`（同树 paged 参考）。 + +### 待办更新 + +- [ ] InfiniLM-vs-vLLM 健康平台全栈对照（5090：vLLM 独立 venv 待建）—— + 立项决策表的最后一块。 +- [ ] prefill=FA / decode=splitkv 分离路由的引擎支持评估（v8 新方向）。 +- [ ] （可选）5060 Ti 上复核 FA decode 回退现象。 + +--- + +## v7（2026-08-27）：elementwise 链闭环 —— nsys kernel 级证据 + +动机：v3 的 nsys 归因（"rmsnorm+rope+swiglu 未融合，~157ms/prefill， +~590 次小 kernel"）留下的第二优化项，在 v5 融合 rms_norm_rope 后还剩 +多少？v6 期间代码走读发现 **swiglu 与 add_rms_norm 其实早已融合**： + +- `Qwen3MLP = layers::MLP`（`qwen3_for_causal_lm.hpp:7`）→ + `csrc/layers/mlp/mlp.cpp:34` 调 `infinicore::op::swiglu`（InfiniCore + 的 NVIDIA 融合 kernel，2025-07 起就在上游 main）——v3 旧树同样如此； +- paged 路径 `TextDecoderLayer::forward(positions, hidden, residual)` + 走 `RMSNorm::forward_inplace(x, residual)` → NVIDIA 上 + `op::add_rms_norm_inplace`（`InfiniCore src/infinicore/nn/rmsnorm.cc:37`）。 + +即 v3 口径里的"swiglu 未融合"不成立（当时已是单 kernel），剩余项只有 +q/k norm+rope（v5 已融合）。本轮在 5090 上用 nsys 对 w2 +（3240 tok prefill + 128 decode，含一轮 warmup，即两次前向）做 +kernel 级 A/B 实证： + +| 成分 | 未融合（n / GPU 时间） | 融合（n / GPU 时间） | +|---|---|---| +| q/k norm+rope | rmsnormKernel 14592 / 39.4ms + ropeThreadPerItem 14336 / 28.4ms | **rmsNormRopeKernel 14336 / 21.3ms** | +| 残差+norm（两侧均融合） | add_rmsnormKernel 14336 / 31.5ms | 14336 / 31.1ms | +| MLP swiglu（两侧均融合） | SwiGLUCuda 7168 / 10.5ms | 7168 / 10.4ms | +| **全 trace kernel 总数** | **101,834** | **87,498（-14%）** | + +（两侧 trace 均含 warmup+计时两次 w2；每次前向的 q/k norm+rope 从 +28 层 × 4 launch 降到 ×2。） + +对账：q/k norm+rope GPU 时间 67.8ms→21.3ms（两次前向合计省 ~46ms， +单次 ~23ms），叠加 launch 延迟节省，与 v6 的 w2 e2e -4.5%（0.67→ +0.64s，省 ~30ms）量级吻合。 + +结论与待办更新： + +1. **v3 的 elementwise 链项至此闭环**：paged 路径三项（残差+norm、 + q/k norm+rope、swiglu）全部单 kernel 化。"swiglu 融合"待办销项—— + 无需新算子，上游既有实现。 +2. trace 中剩余的大头：decode 段 paged-attn splitkv（286.7ms/两次）与 + gemm（259ms，28,704 次小 gemm —— decode 每 token 每层 5 个投影 + gemm，微 batched 化是潜在方向但收益待估）；prefill 段自研 + PagedAttentionPrefill 在 0.6B 上 53ms/次（3240 tok），v4 已证 FA2 + 可再压一个量级——5090 上落 FA 是下一个候选动作。 +3. 分析脚本：profile 采集与聚合命令见 v6 复现要点 + 本轮 + `nsys profile -t cuda` + 自研聚合脚本（分类统计 kernel 名）。 + +--- + +## v6（2026-08-27）：5090 复测 —— rms_norm_rope 收益确认，量级 -4%~-8% + +动机：v5 在 8GB WSL2 病态平台（显存驻留超 ~3~4GB 后带宽崩塌至 +1~11 GB/s）测得 -4%~-23%，需在健康平台复测量级。本轮在租用 +RTX 5090 32GB（Gitee AI 容器，CUDA 12.8 / driver 610.43.02 / +384 核 x86_64）完成。 + +**backend 差异注意**：本轮为 **paged-attn**（5090 上尚未构建 FA 版 +InfiniCore；v5 主数据为 flash-attn）。融合点在 attention 之前的 q/k +norm+rope，与 attention 后端无关，但绝对数字不可与 v5 直接比较。 +配置：64 blocks、no-graph、贪心解码、逐字节同 prompt（同 v5）。 + +### 0.6B 交错 ABBA（paged-attn，e2e 秒） + +| 负载 | A1 未融合 | A2 未融合 | B1 融合 | B2 融合 | e2e Δ | +|---|---|---|---|---|---| +| w1 单请求 decode | 0.36 | 0.39 | 0.36 | 0.36 | -4.0%（A 侧自身波动同量级，边际） | +| w2 长 prefill + 128 decode | 0.67 | 0.67 | 0.64 | 0.64 | **-4.5%** | +| w3 batch32 | 0.49 | 0.48 | 0.45 | 0.45 | **-7.2%** | +| w4 长 decode | 3.14 | 3.08 | 2.87 | 2.85 | **-8.0%** | + +本机无"越跑越慢"漂移（A1≈A2；ABBA 仅作保险）。方向与 v5 一致， +量级收窄——launch 开销在强 CPU + 健康显存平台上占比下降。 + +### 1.7B 单跑对照（paged-attn，各一轮，量级仅作参考） + +| 负载 | 未融合 | 融合 | Δ | +|---|---|---|---| +| w1 单请求 decode | 3.32 ms/tok | 3.28 ms/tok | -1.2% | +| w2 长 prefill + 128 decode | 0.90s | 0.86s | -4.4% | +| w3 batch32 | 5795 tok/s | 5637 tok/s | -2.7%（疑噪声，未复跑） | +| w4 长 decode | 3.66 ms/tok | 3.60 ms/tok | -1.7% | + +### 正确性（--dump-outputs + compare_outputs.py） + +- **噪声底 = 0**：同 binary 跨轮（A1 vs A2）四负载全 exact——本机上 + unfused 完全确定，因此下述分叉全部可归因于融合 kernel 的 fp32 归约 + 顺序差异（算子级 ≤3ulp 已在 v5 证明）。 +- 0.6B fused vs unfused：w2、w4 全 exact；w1 @8 分叉；w3 29/32 + exact。分叉处两侧文本均连贯（如 w3 req9："描述一个没有重力的世界" + vs "描述一个有重力的世界"），属 bf16 logit 并列翻牌。 +- 1.7B fused vs unfused：w1、w2 全 exact；w3 23/32（含一处 @0 首 + token 翻牌）；w4 @160 分叉（"从人类学角度分析" vs "从科技角度分析"， + 两侧连贯）。分叉形态与 v5（8GB 机）一致。 + +### 结论 + +1. rms_norm_rope 融合在健康平台确认有效：0.6B ABBA e2e **-4%~-8%** + （prefill/批处理越重收益越大），1.7B -1%~-4%（单跑）。v5 的 -23% + 量级含 8GB 机病态放大；收益随平台算力/CPU 性能上升而收窄，数据中心 + 卡上预期也是这个量级。 +2. 正确性证据链闭环：算子级 ≤3ulp（v5）+ 双模型 e2e 并列翻牌形态 + 双平台一致（v5/v6）。 +3. 数据归档：`results/*_rtx5090_{fused,unfused}.json`（0.6B ABBA 四轮 + + 1.7B 各一轮，均含 dump-outputs）。 + +### 5090 租用机复现要点（Gitee AI 容器） + +- SSH 经堡垒机：原 `dev_perf/vm_ssh.py` / `vm_scp.py`（pexpect 状态机， + 密码从 `VM_PASSWORD` 环境变量读取）因脚本内含明文 endpoint + （IP+端口+账号）已从仓库移除，复现时需自备等价脚本。 +- GitHub 不可达：xmake 从 gitee 源码构建（`gitee.com/tboox/xmake`， + 子模块为相对 URL，clone 时自动落在 gitee）；xmake-repo 预置 + `gitee.com/tboox/xmake-repo` 到 `~/.xmake/repositories`（apt 的 + xmake 2.8.7 与现版仓库不兼容：on_source nil）。git 在 pty 下会开 + pager 卡住自动化，需 `git --no-pager`。 +- pip 用 `-i https://mirrors.aliyun.com/pypi/simple/`（pypi.org DNS + 只回 IPv6）；HF 下载需 `HF_HUB_DISABLE_XET=1`（hf-mirror 的 xet + 通道 401）。 +- 容器仅 /root、/data 持久化；macOS 侧打包后需 + `find -name '._*' -delete`（bsdtar 的 AppleDouble 文件会混进编译）。 +- InfiniCore 构建：`--cudnn=y`（cudnn=n 在本 HEAD 上 avg_pool3d 编译 + 不过）；删除空的 third_party/cutlass 目录（否则 ENABLE_CUTLASS_API + 打开后 scaled_mm 找不到 cute 头）；`xmake build` 后需显式 build+install + `infiniccl`、`infinicore_cpp_api`、`_infinicore`（非默认 target）。 +- unfused 快照（2366377）的 bench.py 有 /home/yyy 硬编码路径，VM 上 + sed 成自身 checkout 路径（v5 已去硬编码，仅影响旧快照）。 + +--- + +## v5（2026-08-24）：第二优化项之一落地 —— 融合 per-head RMSNorm+RoPE + +动机：v4 遗留的第二优化项——未融合 elementwise 链（~137ms/prefill， +1.7B 口径；v3 nsys：rmsnorm+rope+swiglu 合计 ~157ms / ~590 次小 +kernel）。v5 先融合 attention 内的 q/k per-head RMSNorm + RoPE：每个 +张量 2 次 kernel launch + 2 遍显存往返 → 1 次，decode 每 token 省 +28 层 × 4 = 112 次 launch。 + +**平台注意**：本轮实验在另一台机器（8GB 卡 + WSL2；实测显存驻留超 +~3~4GB 后带宽崩塌至 1~11 GB/s）完成，与 v1~v4 的 16GB 5060 Ti 不是 +同一平台，绝对数值不可直接比较。安排：计时 A/B 用 Qwen3-0.6B（小模型 +launch 开销占比更高，对融合更敏感），正确性对拍用 Qwen3-1.7B（慢速区 +不影响数值）。**收益量级需回 16GB 机复测确认。** + +做法： + +- InfiniCore 新增 `rms_norm_rope` 融合算子（11 个新文件）：C 层 + `src/infiniop/ops/rms_norm_rope/`（CUDA kernel `cuda/kernel.cuh`： + fp32 归约 + 逐变体复刻 rope 舍入路径，GPT_J/NEOX × half/bf16）； + C++ 桥接 `include/infinicore/ops/rms_norm_rope.hpp`。约束：full-rotary + （head_dim = 2 × table_dim）、pos I32/I64、sin/cos F32。 +- InfiniLM 接线（`csrc/models/qwen3/qwen3_attention.cpp`）： + `forward_paged_` 的 q_norm/k_norm + 两次 rope 共 4 处调用 → 两次 + in-place `rms_norm_rope_`（prefill/decode 共用路径）； + `forward_static_` 保持未融合链不变。当前仅覆盖 Qwen3 paged 路径。 + +性能（0.6B，flash-attn，交错 ABBA 控漂移——本机存在"越跑越慢"漂移， +A2 全面慢于 A1，单次先后 A/B 会系统性扭曲差值）： + +| 负载 | A1 未融合 | A2 未融合 | B1 融合 | B2 融合 | e2e Δ | +|---|---|---|---|---|---| +| w1 单请求 decode | 1.00s | 1.05s | 0.92s | 1.05s | **-3.9%** | +| w2 长 prefill + 128 decode | 1.49s | 1.72s | 1.20s | 1.26s | **-23.4%** | +| w3 batch32 | 1.56s | 1.77s | 1.33s | 1.35s | **-19.5%** | +| w4 长 decode | 9.06s | 9.65s | 8.01s | 8.14s | **-13.7%** | + +机制自洽性：decode 为 launch-bound，每 token 省 112 次 launch × ~10µs +≈ 1.1ms，与 w4 的 -13.7%（8.8→7.9 ms/tok）吻合；w2/w3 的更大收益叠加 +了 prefill/批处理段 elementwise 链融合。首轮未控漂移数据方向一致 +（paged-attn：w1 -12.9%、w2 -7.9%、w3 -20.0%、w4 -10.8%；flash-attn +w4 -25.0%）。 + +正确性（三层证据）： + +1. **算子级**（`dev_perf/op_check_rms_norm_rope.cpp`：bf16 输入下融合 + kernel vs `rms_norm_` + `rope_` 参考链逐元素对比）：**>99.999% 逐位 + 一致，最差 ~3ulp**——差异仅来自融合 kernel 内 fp32 归约的分块顺序。 +2. **0.6B e2e 逐 token 对拍**（fused vs unfused，双后端， + `compare_outputs.py`）：分叉率与"同二进制换后端"噪声底同量级 + （w2 两侧均 exact；w3 28~29/32 vs 底 30/32；w4 分叉为单请求 + late-token 并列翻牌）。 +3. **1.7B e2e 逐 token 对拍**（双后端）：分叉位置与噪声底互有先后—— + w2-paged @59、w4-flash @190 与底完全重合，w1-paged 与 w2-flash 的 + fused 全 exact 而底自身分叉（@92 / @59）；目检 w4 分叉处两侧文本 + 均为连贯枚举（"从文学角度分析" vs "从人类社会角度分析"），1024 + token 全程连贯。结论：bf16 logit 并列翻牌经自回归放大，非功能错误。 + +结论与剩余项： + +1. rms_norm_rope 融合在 0.6B 上 e2e 收益 -4%~-23%（负载相关，prefill/ + 批处理越重收益越大），正确性三层证据齐备；**16GB 机上量级待复测**。 +2. 未覆盖：swiglu 融合（v3 口径中 elementwise 链的另一半）、Qwen3 + 以外模型、`forward_static_` 路径。 +3. 工程配套：bench.py/README 去硬编码路径（INFINI_ROOT 默认 + `~/.infini`，INF_MAIN_PYTHON 可指向其他 checkout 的构建产物）。 + +复现： + +```bash +# 前提：FA 版 InfiniCore（见 v4，rms_norm_rope 已在库中）；增量重建 InfiniLM 扩展 +cd && xmake build _infinilm && xmake install +HF_HUB_OFFLINE=1 INFINI_ROOT=$HOME/.infini-fa PYTHONPATH=/python \ + python dev_perf/bench.py --engine infinilm --num-blocks 64 \ + --attn-backend flash-attn --dump-outputs +python dev_perf/compare_outputs.py +# 算子级校验的构建/运行命令见 dev_perf/op_check_rms_norm_rope.cpp 头部注释 +``` + +--- + +## v4（2026-08-24）：B 方向落地 —— prefill 接 FlashAttention-2 + +按 v3 末尾的 kernel 定位（自研 `PagedAttentionPrefill` 比 FA2 慢 13.5×）， +走"直接调 FA2"路线：重建 InfiniCore（`aten=y` + +`--flash-attn=`，FA 取 **v2.7.4.post1**，其 `mha_varlen_fwd` / +`mha_fwd_kvcache` 与 InfiniCore `flash_attention_adaptor.hpp` 逐参数匹配， +原生支持 paged block_table），装到 `~/.infini-fa`；InfiniLM 侧用已有的 +`FlashAttentionImpl`（`--attn-backend flash-attn`：prefill 走 +`flash::mha_varlen_fwd`，decode 走 `flash::mha_fwd_kvcache`）。 + +同机同批对照（Qwen3-1.7B，64blk，no-graph；paged 与 flash 跑在同一个 +FA 版 InfiniCore 库上，对照干净）： + +| 负载 | paged-attn | flash-attn | vLLM（0.85 档） | +|---|---|---|---| +| w1 单请求 decode | 10.86 ms/tok | 11.62 ms/tok | 8.95 ms/tok | +| w2 长 prefill + 128 decode | 2.48s | **2.14s**（单跑复测 1.95s） | 1.50s | +| w3 batch32 总吞吐 | 2191 tok/s | 2127 tok/s | 2957 tok/s | +| w4 长 decode | 11.07 ms/tok | 11.16 ms/tok | 9.71 ms/tok | + +Qwen3-4B（64blk，no-graph）w2：paged 5.41s → flash **4.04s**（-25%）， +vLLM 参考值 3.84s → 差距收敛到 ~1.05×。 + +正确性（贪心解码逐 token 对拍，`compare_outputs.py`）： + +- **w2（FA varlen paged 路径压力最大的负载）输出与 paged-attn 完全一致， + 也与 vLLM 完全一致**——1.7B、4B 均如此。 +- 其余负载的分叉率与 vLLM-vs-paged 的分叉率同量级（w3：24/32 vs 23/32 + exact；w4：两家同在一处 late-token 分叉；目检文本均连贯）——属 bf16 + 规约顺序差异，非功能错误。 + +结论： + +1. **prefill 主差距已被 FA2 消化**：w2 e2e 1.7B -14~30%、4B -25%；4B 上 + 与 vLLM 基本持平（1.05×）。与 v3 预测（attention 491ms→~36ms）吻合。 +2. 剩余差距（1.7B w2 flash 2.14s vs vLLM 1.50s）主要在：未融合 + elementwise 链（~137ms/prefill）与 decode 段（~10%），即 v3 已列的 + 第二优化项。 +3. 注意：vLLM 0.85 档本次实测峰值 15.9GB，贴 97% 看门狗线，复测建议 + 0.72~0.8 档。 + +复现： + +```bash +# 一次性：重建 InfiniCore（约 28min，FA 84 个 .cu 全量编译） +cd /home/yyy/src/InfiniCore # flash-attention repo 在 /home/yyy/src/flash-attention @ v2.7.4.post1 +xmake f -c --aten=y --flash-attn=/home/yyy/src/flash-attention --graph=y \ + --cudnn=n --ccl=n --nv-gpu=y --cpu=y --omp=y \ + --cuda=$HOME/.local/cuda-13.2 --cuda_arch=sm_120 -m release -k shared +xmake build -j6 && xmake install -o ~/.infini-fa + +# 运行（INFINI_ROOT 指向 FA 版库，勿覆盖 V4 线在用的 ~/.infini-dsv4） +INFINI_ROOT=$HOME/.infini-fa HF_HUB_OFFLINE=1 \ + /home/yyy/src/InfiniLM/.venv/bin/python dev_perf/bench.py \ + --engine infinilm --model Qwen/Qwen3-1.7B --num-blocks 64 \ + --attn-backend flash-attn --dump-outputs +``` + +--- + +## v3（2026-08-24）：num_blocks 扫描与膝点定位 + +模型 Qwen3-1.7B（bf16），RTX 5060 Ti 16GB（WSL2），全部 no-graph。 + +| num_blocks | 加载后显存 | w1 ms/tok | w2 e2e | w3 tok/s | w4 ms/tok | +|---|---|---|---|---|---| +| 64 | 7.9GB（48%） | 11.9 / 12.6（两次） | 2.70 / 2.80 | 1995 / 1870 | 12.3 / 14.1 | +| 128 | 9.8GB（60%） | 13.3 | 2.81 | 1852 | 13.9 | +| 256 | 13.3GB（82%） | 13.3 | 2.82 | 1816 | 13.8 | +| 320 | 15.4GB（94%） | 13.5 | 2.81 | 1850 | 13.5 | +| 512（v1 数据） | 16.0GB（98%） | **33.2** | **42.0s** | **42** | **47.9** | + +run 间波动约 ±10%（64blk 两次复测所得），64~320 之间的差异在波动范围内。 + +### v3 结论 + +1. **劣化是 98% 极端饱和处的悬崖，不是渐变**：48%~94% 全区间性能持平， + 只有 512blk（98%）坠崖。排除"engine 内 O(num_blocks) 的 per-step 开销" + 假设，指向显存近满时分配慢路径/驱动行为（WSL2）。具体机制需 profile + 512blk 配置确认，但该配置会触发本机 97% 显存看门狗，暂被阻塞。 +2. **CUDA graph 收益 ~10%**（v2 的 64blk 对照：w1 11.9→10.7、w4 12.3→11.1、 + w3 +7%、w2 -5%），与 num_blocks 无关。值得默认开启，但不是量级差距。 +3. **工程缺陷确认**：默认 num_blocks=512 在 16GB 卡上必踩悬崖，且全程无 + 告警。可立项方向：cache 预分配按显存自适应（预留 ≥5~10% headroom）或 + 近饱和时显式告警。 +4. v1 的全部四条差距假设（launch 开销、prefill 路径、批处理串行、decode + 随长度劣化）均为该悬崖的表现，逐条推翻，详见 v1 存档节。 + +### 机制分析（代码侧，主 checkout 调查结论） + +- 分配链：`llm.py` num_blocks → `PagedKVCacheConfig` → 逐层 + `Tensor::zeros({2, num_blocks, 256, kv_heads, head_dim})` + （`csrc/cache/kv_cache.cpp:142`）。Qwen3-1.7B 为每层 512MiB × 28 = 14GiB， + 与实测吻合。底层是 InfiniCore `PinnableBlockAllocator`——裸 cudaMalloc + + 尺寸分级 free-list 缓存。 +- **稳态 decode 每步零新设备分配**：8 个 CPU 输入 tensor 逐个 H2D、各层 + attention workspace、采样 workspace 全部命中分配器缓存；没有任何 + per-step 开销随 num_blocks 增长。512 vs 320 的 3~50× 劣化**不可能是引擎 + 算法开销**——代码侧排除了引擎内因素。 +- 分配器失败行为是"报错即死"（throw → exit(137)），无重试、无碎片整理、 + 无自动 trim；近饱和**全程无任何告警**，只有启动时一行 + `Using Paged KV Cache with num_blocks=512`。 +- 加载耗时 59.5s（512blk）vs 3.9s（64blk）跑的是完全相同的代码路径 + （28 次 cudaMalloc(512MB) + 设备清零 + 权重 H2D），15× 差距只能来自 + cudaMalloc/kernel 执行本身，即驱动/内存子系统。 +- 综合判断：悬崖在引擎之下——WSL2（dxgkrnl）显存近满时的 + paging/eviction/residency 抖动是最可疑机制，可统一解释"含稳态零新分配 + 的 decode 在内全部变慢"。最终确认需 nsys/driver 计数器（被看门狗阻塞）。 + +### Qwen3-4B 三配置复测（2026-08-24，与 v3 同机） + +InfiniLM 均为 64blk；vLLM 为 `--gpu-mem-util 0.72`（13.97GB，86%）。 +带宽口径：4B bf16 权重 ~8GB，5060 Ti 理论 decode 上限 ~56 tok/s。 + +| 负载 | InfiniLM no-graph | InfiniLM graph | vLLM | +|---|---|---|---| +| w1 单请求 decode | 25.6 ms/tok（39.1 tok/s） | 24.3（41.2） | 23.0（43.4） | +| w2 长 prefill + 128 decode | 5.94s | 5.79s | **3.84s** | +| w3 batch32 总吞吐 | 891 tok/s | 944 tok/s | 841 tok/s | +| w4 长 decode | 26.7 ms/tok | 25.0 | 23.7 | +| 加载耗时 / 显存 | 18.0s / 14.86GB（91%） | 37.0s / 14.82GB | 43.3s / 13.97GB | + +结论： + +1. **"持平"在 4B 上成立**：decode 三家都在带宽上限的 74~77%，w3 InfiniLM + 略优，w4 基本持平。 +2. **唯一持续存在的真实差距是长 prefill**：vLLM 比 InfiniLM 快 ~1.5× + （1.7B 时 ~1.3×）。vLLM 开了 chunked prefill（max_num_batched_tokens= + 8192），这是下一个值得 profile 的点，但量级是 1.5× 而非数量级。 +3. graph 在 4B 上收益收窄到 ~5%（w1 25.6→24.3，w4 26.7→25.0）。 +4. 4B 64blk 已占 91% 显存仍无悬崖，再次印证悬崖只在 ~98% 极端饱和处； + vLLM 0.85 档在 16GB 卡上会撞 97% 看门狗（w4 中途被 SIGTERM），0.72 正常。 + +### prefill 差距的 kernel 级定位（nsys，2026-08-24） + +对 w2（3240 tok prefill + 128 decode）在 1.7B / 64blk 下分别抓 InfiniLM +（no-graph）与 vLLM 的 CUDA 轨迹（`results/prof/*.nsys-rep`，分析脚本 +`results/prof/slice.py`）。prefill 窗口内 GPU busy 均 ~100%——差距在 +kernel 内部，不在调度/Python 开销。 + +prefill 前向一次的 kernel 时间构成（3240 tokens）： + +| 成分 | InfiniLM | vLLM | 倍数 | +|---|---|---|---| +| prefill attention | **491ms**（PagedAttentionPrefillHd128WarpCta8Pipe，26 层 × ~19ms） | 36ms（FA2 splitkv，28 层 × 1.3ms） | **13.5×** | +| rmsnorm + rope + swiglu | ~157ms（未融合，~590 次小 kernel） | ~20ms（triton 融合 kernel） | 7.8× | +| gemm（qkv/o/gate/up/down） | 185ms | 212ms | 持平（略快） | +| **合计** | **888ms / 1655 次调用** | **305ms / 706 次调用** | **2.9×** | + +结论：prefill 差距的第一来源是自研 `PagedAttentionPrefill` kernel 比 +FlashAttention-2 慢一个数量级（每层 19ms vs 1.3ms），第二来源是 +elementwise 链未融合。这是边界清晰、可度量的 kernel 优化目标： +把 prefill attention 换成/优化到 FA2 量级，w2 的 prefill 段理论上可从 +~1.2s 压到 ~0.6s（1.7B 口径），e2e 差距从 1.3~1.5× 收敛到接近 1。 + +### 待办 + +- [ ] （需看门狗临时放宽）nsys 抓 512blk 的 w1，直接观察 98% 悬崖机制。 +- [x] ~~prefill attention kernel 优化立项~~ —— 已由 v4 完成：接 FA2 + （`--attn-backend flash-attn` + FA 版 InfiniCore），w2 的 prefill 段 + 差距收敛到 ~1.05~1.3×，正确性逐 token 对拍通过。 +- [x] ~~（可选）elementwise 融合（rmsnorm/rope）作为第二优化项~~ —— + 已由 v5 完成：`rms_norm_rope` 算子接入 Qwen3 paged 路径，0.6B + ABBA e2e -4%~-23%，三层正确性证据齐备；跨平台复测已由 v6 完成 + （5090 上 0.6B ABBA -4%~-8%，正确性形态一致）。 +- [x] ~~（可选）swiglu 融合~~ —— 销项：上游 swiglu/add_rms_norm 早已 + 融合，Qwen3 paged 路径一直在用；v7 nsys 实证 elementwise 链已 + 全部单 kernel 化。 +- [ ] （可选）flash-attn 设为默认 attention backend 的评估：需在更多模型 + 上补正确性对拍，并确认 FP8/滑窗/softcap 模型的回退路径。 +- [ ] Qwen3-8B-FP8 下载完成后，可作为 8B 级 + FP8 口径的复测对象。 + +--- + +## v2 存档（2026-08-24）：64blk 对照实验原始表 + +| 负载 | InfiniLM 512blk no-graph | InfiniLM 64blk no-graph | InfiniLM 64blk graph | vLLM（v1 默认） | +|---|---|---|---|---| +| w1 单请求 decode | 33.2 ms/tok | 11.9 ms/tok | **10.7 ms/tok** | 12.0 ms/tok | +| w2 长 prefill + 128 decode | 42.0s | 2.70s | 2.56s | 2.0s | +| w3 batch32 总吞吐 | 42 tok/s | 1995 tok/s | 2136 tok/s | 2166 tok/s | +| w4 长 decode | 47.9 ms/tok | 12.3 ms/tok | 11.1 ms/tok | 12.2 ms/tok | +| 加载耗时 | 59.5s | 3.9s | 11.6s（含 graph 编译） | 24.0s | +| 加载后显存 | 15.96GB | 7.89GB | 8.07GB | 15.89GB | + +--- + +## v1 存档（2026-08-23/24，已被推翻） + +数据：`results/infinilm_..._0823_235223.json`、`results/vllm_..._0824_002522.json`。 +配置：InfiniLM no-graph、paged-attn、prefix caching on、**num_blocks=512**； +vLLM v1 默认（CUDA graph on、prefix caching on）、gpu_memory_utilization=0.85。 + +| 负载 | InfiniLM | vLLM | 差距 | +|---|---|---|---| +| w1 单请求 decode | 33.2 ms/tok（30.2 tok/s） | 12.0 ms/tok（83.2 tok/s） | 2.8× | +| w2 长 prefill（3240 tok）+ 128 decode | 42.0s e2e | 2.0s e2e | 21× | +| w3 batch32 × 128 tok | 42.0 tok/s | 2166.0 tok/s | 52× | +| w4 长 decode 1024 tok | 47.9 ms/tok（20.9 tok/s） | 12.2 ms/tok（82.3 tok/s） | 3.9× | + +v1 归因假设与判决： + +1. ~~单请求 decode 低是缺 CUDA graph 的 launch 开销~~ —— 推翻：no-graph + 64blk 已达 11.9 ms/tok，graph 只再提速 ~10%。 +2. ~~w2 是 prefill 路径异常（纯 prefill 估算差距 ~75×）~~ —— 推翻： + 64blk 下 w2 e2e 2.7s，与 vLLM 同量级。 +3. ~~w3 批处理存在 per-request 串行开销~~ —— 推翻：64blk 下 w3 + 1850~2136 tok/s，扩展性正常。 +4. ~~w4 decode 随上下文长度劣化~~ —— 推翻：64blk 下 w4 与 w1 持平。 + +## 立项书模板（方向定稿后填写） + +"在 [模型] + [硬件] + [负载矩阵] 下，[TTFT/TPOT/吞吐] 从 X 提升到 Y（≥Z%）， +精度（ceval/mmlu/ppl）不降级，单测与 CI 全绿，benchmark 脚本入库、 +他人可复现。" + +风险提示：5060 Ti 的瓶颈结构 ≠ 数据中心卡，立项报告中必须注明平台， +收尾时应在数据中心卡（或训练营国产平台）上复测。 diff --git a/dev_perf/op_check_rms_norm_rope.cpp b/dev_perf/op_check_rms_norm_rope.cpp new file mode 100644 index 000000000..dae5f3173 --- /dev/null +++ b/dev_perf/op_check_rms_norm_rope.cpp @@ -0,0 +1,123 @@ +// Op-level numeric check for the fused rms_norm_rope operator. +// +// Compares, elementwise on GPU: +// reference: op::rms_norm_ into tmp, then op::rope_ in-place on tmp +// fused: op::rms_norm_rope_ in-place on a copy of the same input +// Both paths use identical bf16 inputs / weights / sin-cos tables, so any +// difference comes from the fused kernel's reduction order (fp32 accumulation +// in a smaller thread block) rather than from input rounding. +// +// Build: +// g++ -std=c++17 -I$HOME/.infini-fa/include dev_perf/op_check_rms_norm_rope.cpp \ +// -L$HOME/.infini-fa/lib -linfinicore_cpp_api -linfiniop -linfinirt -linfiniccl \ +// -Wl,-rpath,$HOME/.infini-fa/lib -o /tmp/op_check_rms_norm_rope +// Run: +// LD_LIBRARY_PATH=$HOME/.infini-fa/lib:$HOME/.local/cuda-13.2/lib64 /tmp/op_check_rms_norm_rope + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/rms_norm.hpp" +#include "infinicore/ops/rms_norm_rope.hpp" +#include "infinicore/ops/rope.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include +#include + +static uint16_t f2bf(float f) { // round-to-nearest-even + uint32_t u; + std::memcpy(&u, &f, 4); + uint32_t bias = 0x7FFFu + ((u >> 16) & 1u); + return static_cast((u + bias) >> 16); +} +static float bf2f(uint16_t b) { + uint32_t u = static_cast(b) << 16; + float f; + std::memcpy(&f, &u, 4); + return f; +} + +// deterministic pseudo-random in [-range, range] +static float frand(uint32_t &s, float range) { + s = s * 1664525u + 1013904223u; + return (static_cast((s >> 8) & 0xFFFF) / 32768.0f - 1.0f) * range; +} + +static int run_case(size_t tokens, size_t heads, size_t head_dim, + infinicore::nn::RoPE::Algo algo, const char *tag) { + using namespace infinicore; + const size_t table_len = 4096, table_dim = head_dim / 2; + const float eps = 1e-6f; + + // ---- host buffers (bf16) ---- + uint32_t seed = 42; + std::vector hx(tokens * heads * head_dim), hw(head_dim); + for (auto &v : hx) v = f2bf(frand(seed, 2.0f)); + for (auto &v : hw) v = f2bf(0.5f + std::fabs(frand(seed, 1.0f))); + std::vector hsin(table_len * table_dim), hcos(table_len * table_dim); + for (size_t p = 0; p < table_len; ++p) + for (size_t i = 0; i < table_dim; ++i) { + float freq = 1.0f / std::pow(10000.0f, 2.0f * i / head_dim); + hsin[p * table_dim + i] = f2bf(std::sin(p * freq)); + hcos[p * table_dim + i] = f2bf(std::cos(p * freq)); + } + std::vector hpos(tokens); + for (size_t t = 0; t < tokens; ++t) hpos[t] = static_cast((t * 37 + 11) % table_len); + + // ---- device tensors ---- + auto dev = context::getDevice(); + auto x_fused = Tensor::empty({tokens, heads, head_dim}, DataType::BF16, dev); + auto x_ref = Tensor::empty({tokens, heads, head_dim}, DataType::BF16, dev); + auto tmp = Tensor::empty({tokens, heads, head_dim}, DataType::BF16, dev); + auto w = Tensor::empty({head_dim}, DataType::BF16, dev); + auto sin = Tensor::empty({table_len, table_dim}, DataType::BF16, dev); + auto cos = Tensor::empty({table_len, table_dim}, DataType::BF16, dev); + auto pos = Tensor::empty({tokens}, DataType::I64, dev); + + context::memcpyH2D(x_fused->data(), hx.data(), hx.size() * 2, false); + context::memcpyH2D(x_ref->data(), hx.data(), hx.size() * 2, false); + context::memcpyH2D(w->data(), hw.data(), hw.size() * 2, false); + context::memcpyH2D(sin->data(), hsin.data(), hsin.size() * 2, false); + context::memcpyH2D(cos->data(), hcos.data(), hcos.size() * 2, false); + context::memcpyH2D(pos->data(), hpos.data(), hpos.size() * 8, false); + + // reference chain: rms_norm out-of-place, then rope in-place (engine order) + op::rms_norm_(tmp, x_ref, w, eps); + op::rope_(tmp, tmp, pos, sin, cos, algo); + // fused + op::rms_norm_rope_(x_fused, w, pos, sin, cos, eps, algo); + + std::vector out_ref(hx.size()), out_fused(hx.size()); + context::memcpyD2H(out_ref.data(), tmp->data(), hx.size() * 2); + context::memcpyD2H(out_fused.data(), x_fused->data(), hx.size() * 2); + + size_t n_exact = 0, n_1ulp = 0, n_2ulp = 0, n_worse = 0; + float max_rel = 0.0f; + for (size_t i = 0; i < hx.size(); ++i) { + float a = bf2f(out_ref[i]), b = bf2f(out_fused[i]); + float diff = std::fabs(a - b); + float rel = diff / std::max(std::fabs(a), 1e-3f); + max_rel = std::max(max_rel, rel); + if (diff == 0.0f) ++n_exact; + else if (rel <= 0.0078f) ++n_1ulp; // bf16 ulp ~= 2^-7 relative + else if (rel <= 0.0156f) ++n_2ulp; + else ++n_worse; + } + std::printf("[%s] tokens=%zu heads=%zu dim=%zu exact=%.2f%% 1ulp=%zu 2ulp=%zu worse=%zu max_rel=%.4f\n", + tag, tokens, heads, head_dim, + 100.0 * n_exact / hx.size(), n_1ulp, n_2ulp, n_worse, max_rel); + return n_worse == 0 ? 0 : 1; +} + +int main() { + infinicore::context::setDevice(infinicore::Device(infinicore::Device::Type::NVIDIA, 0)); + int rc = 0; + // Qwen3-0.6B/1.7B per-layer shape: heads=16(q)/8(k), head_dim=128, GPT_NEOX + rc |= run_case(97, 16, 128, infinicore::nn::RoPE::Algo::GPT_NEOX, "neox-q"); + rc |= run_case(97, 8, 128, infinicore::nn::RoPE::Algo::GPT_NEOX, "neox-k"); + rc |= run_case(97, 16, 128, infinicore::nn::RoPE::Algo::GPT_J, "gptj-q"); + rc |= run_case(1, 16, 128, infinicore::nn::RoPE::Algo::GPT_NEOX, "neox-single-token"); + std::puts(rc == 0 ? "PASS: all diffs within 2 bf16 ulp" : "FAIL: >2ulp diffs found"); + return rc; +} diff --git a/dev_perf/step_breakdown.py b/dev_perf/step_breakdown.py new file mode 100644 index 000000000..8d5e5f191 --- /dev/null +++ b/dev_perf/step_breakdown.py @@ -0,0 +1,79 @@ +"""Group graph-node kernels into replays by time gaps, report median +per-replay time+count per (kernel, grid), with analytic bytes/GBs for the +known Qwen3-0.6B GEMV call sites. + +Usage: python3 step_breakdown5.py /root/prof_w1_06b_n.sqlite +""" +import sqlite3 +import sys +from collections import Counter, defaultdict + +db = sqlite3.connect(sys.argv[1]) +rows = db.execute( + "SELECT k.start, k.end, s.value, k.gridX, k.gridY, k.gridZ " + "FROM CUPTI_ACTIVITY_KIND_KERNEL k " + "JOIN StringIds s ON k.demangledName = s.id " + "WHERE k.graphNodeId IS NOT NULL ORDER BY k.start" +).fetchall() +print(f"graph kernels: {len(rows)}") + +GAP = 100_000 # 100us: inter-step gap is ~0.5ms+, intra-step gaps are ~us +replays, cur = [], [] +for r in rows: + if cur and r[0] - cur[-1][1] > GAP: + replays.append(cur) + cur = [] + cur.append(r) +if cur: + replays.append(cur) + +sizes = Counter(len(r) for r in replays) +print("replay sizes (kernels: count):", sizes.most_common(6)) +mode_size = sizes.most_common(1)[0][0] +steady = [r for r in replays if len(r) == mode_size] +print(f"mode replay size: {mode_size}, steady replays: {len(steady)}") + + +def short(name): + n = name + if "internal::gemvx::kernel" in n: + return "gemvx" + if n.startswith("void cutlass::Kernel2")[0] + if "flashAttentionDecodeHd128SplitKvCta" in n: + return "paged_splitkv_cta" + if "flashAttentionDecodeHd128SplitKvCombine" in n: + return "paged_splitkv_combine" + if "DeviceReduceSingleTileKernel" in n: + return "cub_reduce_single" + if "DeviceReduceKernel" in n: + return "cub_reduce" + return n.split("<")[0].split("(")[0].replace("void ", "")[:56] + + +agg = defaultdict(lambda: [[], []]) +walls = [] +for rep in steady: + walls.append(rep[-1][1] - rep[0][0]) + per = defaultdict(lambda: [0, 0]) + for start, end, name, gx, gy, gz in rep: + key = f"{short(name)} g({gx},{gy},{gz})" + per[key][0] += end - start + per[key][1] += 1 + for k, (t, c) in per.items(): + agg[k][0].append(t) + agg[k][1].append(c) + +walls.sort() +med_wall = walls[len(walls) // 2] +items = [] +for k, (ts, cs) in agg.items(): + ts.sort() + cs.sort() + items.append((ts[len(ts) // 2], cs[len(cs) // 2], k)) +items.sort(reverse=True) +tot = sum(t for t, _, _ in items) +print(f"\nmedian step wall: {med_wall/1e3:.1f} us, GPU-busy: {tot/1e3:.1f} us, idle: {(med_wall-tot)/1e3:.1f} us ({100.0*(med_wall-tot)/med_wall:.0f}%)") +print(f"{'kernel @ grid':<56} {'us/step':>8} {'n/step':>6} {'us/call':>8} {'%':>5}") +for t, c, k in items[:22]: + print(f"{k:<56} {t/1e3:>8.1f} {c:>6} {t/c/1e3:>8.2f} {100.0*t/tot:>4.1f}") diff --git a/dev_perf/workload.py b/dev_perf/workload.py new file mode 100644 index 000000000..44c9708ca --- /dev/null +++ b/dev_perf/workload.py @@ -0,0 +1,145 @@ +"""Shared workload definition for the offline dual-engine perf baseline. + +Both bench entry points (infinilm / vllm) import from here so the workload +is byte-identical across engines. Stdlib + huggingface_hub only. +""" + +import os +import subprocess +import threading +import time + +# --------------------------------------------------------------------------- +# Prompts +# --------------------------------------------------------------------------- + +SHORT_PROMPTS = [ + "如果猫能写诗，它们会写些什么？", + "描述一个没有重力的世界。", + "如果地球停止自转，会发生什么？", + "假设你是一只会飞的鲸鱼，描述你的日常生活。", +] + +_BASE_PARA = ( + "大语言模型推理引擎的核心挑战在于如何高效利用 GPU 的算力与显存带宽。" + "预填充阶段是计算密集型的，而解码阶段则主要受限于显存带宽，" + "因为每一步都需要读取全部的模型权重与不断增长的 KV 缓存。" + "连续的批处理调度与前缀缓存是提升服务端吞吐的关键技术。" +) + +LONG_PROMPT = _BASE_PARA * 40 # ~2k tokens; actual count recorded at runtime + + +def _batch_prompts(n: int) -> list[str]: + return [f"问题 {i + 1}：{SHORT_PROMPTS[i % len(SHORT_PROMPTS)]}" for i in range(n)] + + +def concurrent_prefill_workload(n: int = 8) -> tuple: + """w5: n long prompts (w2's LONG_PROMPT construction) submitted in one + batch — the chunked-prefill acceptance workload. With plain FCFS + scheduling the n-1 queued prefills block everyone's decode; chunked + prefill should interleave them and cut e2e wall time. + + Each prompt gets a distinct 问题 i： header so they diverge from the + very first tokens — prefix caching (on for both engines) can then + never dedupe them, and every request pays its own full prefill. + """ + prompts = [f"问题 {i + 1}：{LONG_PROMPT}" for i in range(n)] + return ("w5_concurrent_prefill", prompts, 128) + + +def decode_stall_workload(n_decode: int = 8, inject_reps: int = 80) -> tuple: + """w6: the chunked-prefill *benefit* scenario — a steady decode stream + (n_decode short prompts, long generation) with one long prompt injected + mid-flight. Plain FCFS stalls every decode token for the whole injected + prefill; chunked prefill should spread it into mixed steps and shrink the + ITL spike. bench.py drives this engine-level (add_request/step); the + prompts/max_tokens fields are placeholders, parameters ride in the tuple. + + The injected prompt gets a unique nonce header so prefix caching (w2/w5 + ran earlier in the same engine) can never dedupe its blocks: the nonce + shifts every 256-token block boundary vs any earlier LONG_PROMPT use. + """ + inject_prompt = f"w6注入{time.time_ns()}：" + _BASE_PARA * inject_reps + return ("w6_decode_stall", (n_decode, inject_prompt), 0) + + +def repetitive_copy_workload() -> tuple: + """w7: prompt-lookup 投机采样的收益负载——pattern 续写使输出大段 + 重复 prompt 内容，n-gram 后缀匹配持续命中（接受率上限的演示）。 + 用纯模式重复而非指令（"请重复 N 遍"），贪心小模型必然续写， + 不依赖指令遵循能力。正确性由 verify 保证，与是否命中无关。 + """ + prompt = "下面这段文字会不断重复：\n" + _BASE_PARA * 3 + return ("w7_repetitive_copy", [prompt], 512) + + +# name, prompts, max_tokens +WORKLOADS = [ + ("w1_short_decode", SHORT_PROMPTS[:1], 128), + ("w2_long_prefill", [LONG_PROMPT], 128), + ("w3_batch32", _batch_prompts(32), 128), + ("w4_long_decode", SHORT_PROMPTS[1:2], 1024), + concurrent_prefill_workload(), + decode_stall_workload(), + repetitive_copy_workload(), +] + + +def ctx_sweep_workloads() -> list: + """Single-request decode-vs-context-length sweep for the decode-kernel + crossover study (v9/v11: paged splitkv wins short ctx, FA kvcache wins + long ctx; crossover is model-geometry dependent). + + `_BASE_PARA` is ~81 tokens/repeat on the Qwen3 tokenizer (40 reps + tokenize to ~3240), so the multiplier list below spans ~0.2k~5k ctx. + The label is only approximate — the exact prompt token count is recorded + in the report at runtime. Prefill is identical across the two backends + (both use FA2), so the e2e delta at each ctx isolates the decode kernel. + """ + return [(f"ctx81x{k}", [_BASE_PARA * k], 128) for k in (2, 4, 8, 12, 16, 24, 32, 40, 48, 64)] + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def resolve_model_path(name_or_path: str) -> str: + """Accept a local dir or an HF repo id; return a local directory path.""" + if os.path.isdir(name_or_path): + return os.path.abspath(name_or_path) + from huggingface_hub import snapshot_download + + return snapshot_download(name_or_path) + + +class MemSampler(threading.Thread): + """Sample `nvidia-smi` memory.used in the background; track peak.""" + + def __init__(self, interval: float = 0.2): + super().__init__(daemon=True) + self.interval = interval + self.peak = 0 + self.latest = 0 + self._stop_event = threading.Event() + + def run(self): + while not self._stop_event.is_set(): + try: + out = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=memory.used", + "--format=csv,noheader,nounits", + ], + text=True, + ).strip() + self.latest = int(out.splitlines()[0]) + self.peak = max(self.peak, self.latest) + except Exception: + pass + time.sleep(self.interval) + + def stop(self): + self._stop_event.set() + self.join(timeout=2) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c0..a27b46a45 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -249,7 +249,7 @@ def _add_common_args(self): "--attn", type=str, default="default", - choices=["default", "paged-attn", "flash-attn"], + choices=["default", "paged-attn", "flash-attn", "hybrid"], ) self.parser.add_argument("--enable-graph", action="store_true") self.parser.add_argument( diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index bccb7e758..dab1e9826 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -11,7 +11,11 @@ class EngineConfig: Attributes: model_path: Path to the model directory. draft_model_path: Optional Eagle/MTP draft model directory. - num_draft_tokens: Number of Eagle draft tokens to verify per step. + speculative_method: Speculative decoding method. None disables speculation + unless draft_model_path is set (which implies "eagle"). "prompt_lookup" + needs no draft model: draft tokens come from n-gram suffix matching + against the request's own prompt+output. + num_draft_tokens: Number of draft tokens to verify per step. device: Device type string ('cpu', 'cuda', 'mlu', etc.). dtype: Data type string ('float16', 'bfloat16', 'float32'). tensor_parallel_size: Number of devices for tensor parallelism. @@ -32,7 +36,7 @@ class EngineConfig: top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. - attn_backend: Attention backend to use ('default', 'flash-attn'). + attn_backend: Attention backend to use ('default', 'paged-attn', 'flash-attn', 'hybrid'). use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. skip_load: Whether to skip loading model weights (for testing). @@ -41,6 +45,7 @@ class EngineConfig: model_path: str draft_model_path: Optional[str] = None + speculative_method: Optional[str] = None # None / "eagle" / "prompt_lookup" num_draft_tokens: int = 4 device: str = "cuda" dtype: str = "float16" @@ -73,6 +78,25 @@ class EngineConfig: def __post_init__(self) -> None: if self.num_draft_tokens < 1: raise ValueError("num_draft_tokens must be >= 1") + # 归一化投机方法：给了 draft_model_path 而未指定方法时默认 eagle， + # 保持旧调用方行为不变。 + if self.speculative_method is None and self.draft_model_path is not None: + self.speculative_method = "eagle" + if self.speculative_method is not None: + if self.speculative_method not in ("eagle", "prompt_lookup"): + raise ValueError( + "speculative_method must be one of: eagle, prompt_lookup" + ) + if self.speculative_method == "eagle" and self.draft_model_path is None: + raise ValueError("speculative_method='eagle' requires draft_model_path") + if ( + self.speculative_method == "prompt_lookup" + and self.draft_model_path is not None + ): + raise ValueError( + "speculative_method='prompt_lookup' takes no draft model; " + "leave draft_model_path unset" + ) if self.pipeline_parallel_size < 1: raise ValueError("pipeline_parallel_size must be >= 1") if not 0 <= self.pipeline_parallel_stage < self.pipeline_parallel_size: diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..a5102dffa 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -25,12 +25,14 @@ FinishReason, InferenceRequest, RequestOutput, + RequestStatus, TokenOutput, ) from infinilm.llm.sampling_params import SamplingParams from infinilm.llm.scheduler import Scheduler from infinilm.llm.static_scheduler import StaticScheduler from infinilm.multimodal.multimodal import resolve_multimodal_inputs +from infinilm.processors.basic_llm_processor import BasicLLMProcessor logger = logging.getLogger(__name__) @@ -101,6 +103,10 @@ def __init__(self, config: EngineConfig): ) assert 1024 <= max_num_batched_tokens <= max_position_embeddings + enable_chunked_prefill = self._chunked_prefill_supported( + config, has_mamba_cache + ) + self.scheduler = Scheduler( max_batch_size=config.max_batch_size, num_blocks=config.num_blocks, @@ -110,8 +116,14 @@ def __init__(self, config: EngineConfig): has_mamba_cache=has_mamba_cache, num_mamba_cache_blocks=num_mamba_cache_blocks, enable_prefix_caching=config.enable_prefix_caching, + enable_chunked_prefill=enable_chunked_prefill, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") + if enable_chunked_prefill: + logger.info( + "Chunked prefill enabled with max_num_batched_tokens=%s", + max_num_batched_tokens, + ) if has_mamba_cache: logger.info( "Using Mamba cache with num_blocks=%s, zero_state_index=0", @@ -137,6 +149,36 @@ def add_request(self, request: InferenceRequest): """Add a request to the scheduler.""" self.scheduler.add_request(request) + def _chunked_prefill_supported( + self, config: EngineConfig, has_mamba_cache: bool + ) -> bool: + """是否启用 chunked prefill + prefill/decode 混排。 + + 由 env INFINILM_ENABLE_CHUNKED_PREFILL 打开（默认关闭，关闭时调度 + 行为与旧逻辑一致）。以下路径保持整段 prefill 旧行为： + mamba/线性注意力模型，以及按 batch 级 is_prefill 构建输入的处理 + 器（多模态模型，混排会破坏其输入构建）。 + + 投机采样（eagle / prompt_lookup）不再互斥：SpeculativeRunner 按 + SchedulerOutput.num_scheduled_tokens 逐请求判定阶段，混排批次里 + 的中段 chunk 请求不参与 draft/verify、也不产出 token。 + """ + if os.getenv("INFINILM_ENABLE_CHUNKED_PREFILL", "").lower() not in ( + "1", + "true", + "yes", + "on", + ): + return False + if has_mamba_cache: + return False + processor_cls = type(self.processor) + return ( + processor_cls.build_model_inputs is BasicLLMProcessor.build_model_inputs + and processor_cls._build_model_input_from_batch_scheduler_output + is BasicLLMProcessor._build_model_input_from_batch_scheduler_output + ) + def close(self): self.model_runner.close() @@ -160,6 +202,7 @@ def step(self) -> tuple[bool, list[tuple]]: pending = self._update_requests( scheduler_output.scheduled_requests, sampled_token_ids, + getattr(scheduler_output, "num_scheduled_tokens", None), ) # Return False (no immediate work) only when no requests were scheduled @@ -179,6 +222,7 @@ def _update_requests( self, requests: List[InferenceRequest], sampled_token_ids: list[int | list[int]], + num_scheduled_tokens: Optional[List[int]] = None, ) -> List[tuple]: """Apply sampled tokens and publish their target-model KV boundary.""" if len(requests) != len(sampled_token_ids): @@ -186,8 +230,37 @@ def _update_requests( "model output count does not match the scheduled request count: " f"requests={len(requests)}, outputs={len(sampled_token_ids)}" ) + if num_scheduled_tokens is not None and len(num_scheduled_tokens) != len( + requests + ): + raise RuntimeError( + "num_scheduled_tokens count does not match the scheduled request " + f"count: requests={len(requests)}, " + f"num_scheduled_tokens={len(num_scheduled_tokens)}" + ) pending = [] - for req, token_ids in zip(requests, sampled_token_ids): + for req_idx, (req, token_ids) in enumerate(zip(requests, sampled_token_ids)): + # chunked prefill 的中段 chunk：prompt 尚未算完，本步采样到的是废 + # token，直接丢弃；只推进 num_computed_tokens 并逐 chunk 发布已 + # 完整的块，跳过 token append / tokenizer.decode / finish 判定 + if num_scheduled_tokens is not None: + chunk_end = req.num_local_cached_tokens + num_scheduled_tokens[req_idx] + if chunk_end < req.get_prompt_length(): + req.num_computed_tokens = chunk_end + self.scheduler.commit_computed_tokens(req, chunk_end) + if req.is_aborted(): + logger.info( + f"Request {req.request_id} aborted by client, skipping update" + ) + # close() may have set _aborted=True without setting a terminal status + # (status still RUNNING). + if not req.is_finished(): + req.mark_canceled() + else: + # prompt 未算完，标记为 WAITING，由 complete_requests + # 放回 waiting 队列等待下一个 chunk + req.status = RequestStatus.WAITING + continue # The model successfully consumed the request's current logical tokens. # Commit this boundary before observing a concurrent client abort. pre_output_computed_tokens = req.get_total_length() @@ -339,6 +412,7 @@ def __init__( self, model_path: str, draft_model_path: Optional[str] = None, + speculative_method: Optional[str] = None, num_draft_tokens: int = 4, device: str = "cuda", dtype: str = "float16", @@ -371,6 +445,10 @@ def __init__( Args: model_path: Path to the model directory. + draft_model_path: Optional Eagle/MTP draft model directory. + speculative_method: Speculative decoding method ('eagle' or + 'prompt_lookup'). 'prompt_lookup' needs no draft model. + Defaults to 'eagle' when draft_model_path is set. device: Device type ('cpu', 'cuda', 'mlu', 'moore'). dtype: Data type ('float16', 'bfloat16', 'float32'). tensor_parallel_size: Number of devices for tensor parallelism. @@ -384,13 +462,14 @@ def __init__( top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. - attn_backend: Attention backend to use ('default', 'flash-attn'). + attn_backend: Attention backend to use ('default', 'static-attn', 'paged-attn', 'flash-attn', 'flashinfer', 'hybrid'). 'hybrid' = FA2 for prefill, paged kernel for decode. use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. """ config = EngineConfig( model_path=model_path, draft_model_path=draft_model_path, + speculative_method=speculative_method, num_draft_tokens=num_draft_tokens, device=device, dtype=dtype, @@ -566,6 +645,7 @@ def __init__( self, model_path: str, draft_model_path: Optional[str] = None, + speculative_method: Optional[str] = None, num_draft_tokens: int = 4, device: str = "cuda", dtype: str = "float16", @@ -599,6 +679,10 @@ def __init__( Args: model_path: Path to the model directory. + draft_model_path: Optional Eagle/MTP draft model directory. + speculative_method: Speculative decoding method ('eagle' or + 'prompt_lookup'). 'prompt_lookup' needs no draft model. + Defaults to 'eagle' when draft_model_path is set. device: Device type ('cpu', 'cuda', 'mlu', 'moore'). dtype: Data type ('float16', 'bfloat16', 'float32'). tensor_parallel_size: Number of devices for tensor parallelism. @@ -612,7 +696,7 @@ def __init__( top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. - attn_backend: Attention backend to use ('default', 'flash-attn'). + attn_backend: Attention backend to use ('default', 'static-attn', 'paged-attn', 'flash-attn', 'flashinfer', 'hybrid'). 'hybrid' = FA2 for prefill, paged kernel for decode. kv_connector: KV connector type ('MooncakeConnector'). kv_role: Role in KV connector ('kv_producer' or 'kv_consumer'). kv_connector_extra_config: Extra config dict for KV connector. @@ -623,6 +707,7 @@ def __init__( config = EngineConfig( model_path=model_path, draft_model_path=draft_model_path, + speculative_method=speculative_method, num_draft_tokens=num_draft_tokens, device=device, dtype=dtype, diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index a1696f848..7b9131214 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -107,7 +107,7 @@ def __init__(self, config: EngineConfig, initialize_processor: bool = True): ) self.speculative_runner = None - if config.draft_model_path is not None: + if config.speculative_method is not None: self.speculative_runner = SpeculativeRunner( config, self.model_engine, self.device ) @@ -244,6 +244,12 @@ def _model_forward(self, scheduler_output): def _model_forward_with_speculative(self, scheduler_output, model_input): return self.speculative_runner.forward(scheduler_output, model_input) + def get_speculative_stats(self): + """投机采样统计快照；未启用投机时返回 None。""" + if self.speculative_runner is None: + return None + return self.speculative_runner.get_acceptance_stats() + @contextmanager def maybe_get_kv_connector_output( self, scheduler_output: Any diff --git a/python/infinilm/llm/model_runner/speculative_runner.py b/python/infinilm/llm/model_runner/speculative_runner.py index d3c5211d4..ead6ce88d 100644 --- a/python/infinilm/llm/model_runner/speculative_runner.py +++ b/python/infinilm/llm/model_runner/speculative_runner.py @@ -1,18 +1,94 @@ +import logging +import os + import infinicore from infinilm.cache.cache import StaticKVCacheConfig from infinilm.distributed import DistConfig from infinilm.infer_engine import InferEngine from infinilm.modeling_utils import load_model_state_dict_by_file +logger = logging.getLogger(__name__) + class SpeculativeRunner: + """Speculative decoding runner: draft → paged verify → accept/rollback. + + 两种 draft 来源（config.speculative_method）： + - "eagle": MiniCPM Eagle draft 头逐步自回归猜测（需要 draft_model_path， + 依赖 target 前向的 hidden states，主前向只能走 eager 的 forward_raw）。 + - "prompt_lookup": 零训练的 n-gram 后缀匹配，从请求自身的 + prompt+已生成序列里复制后续片段作为 draft。不依赖 hidden states； + 纯 decode 批次走融合单前向（主采样 + 验证一次完成，形状固定为 + b × (k+1)），混排/prefill 批次走两前向流程。 + + 两种方法共用同一条精确验证链路：append_verify_slots 申请临时 KV 槽 → + 全位置采样前向（融合或独立 verify）→ 接受最长匹配前缀 + 1 个修正 + token → rollback_to_length 回收未接受的槽位。贪心解码下输出与非投机 + 数学等价（分布无损）；注意不保证逐位一致——verify 前向的 batch 形状 + 与基线 decode 不同，logit 近平局的位置 argmax 可能翻转（5090 实测 + 分歧点 top-2 间隙 0~0.125，见 dev_perf/gap_analysis.md v16）。大 + batch（> INFINILM_SPEC_MAX_BATCH_SIZE，默认 32）回退常规前向。 + + 自适应收益门控：滑动窗口（默认 32 个请求步）内平均每步产出低于 + INFINILM_SPEC_MIN_AVG_TOKENS（默认 2.0，eager 投机步成本 ≈ 2× 常规 + decode 步的盈亏平衡点）时，回退常规前向 INFINILM_SPEC_GATE_COOLDOWN + 步（默认 64）再自动重试——开放生成等低命中负载开投机不致亏。 + """ + def __init__(self, config, target_model_engine, device): self.config = config self.target_model_engine = target_model_engine + self.speculative_method = config.speculative_method self.num_draft_tokens = config.num_draft_tokens self.draft_max_batch_size = config.max_batch_size - self.eagle_accept_count = 0 - self.eagle_total_count = 0 + # 接受率埋点：drafted/accepted 衡量 draft 质量，emitted_tokens/steps + # 衡量端到端收益（非投机基线恒为 1.0 token/step）。 + self.total_drafted = 0 + self.total_accepted = 0 + self.total_emitted_tokens = 0 + self.total_emitted_steps = 0 + self.total_verify_alloc_failures = 0 + + # 大 batch 下 decode 转向计算受限，投机把每步计算量放大 k 倍反而 + # 降吞吐：batch 超过阈值时回退常规前向（env 可调） + self.spec_max_batch_size = int( + os.getenv("INFINILM_SPEC_MAX_BATCH_SIZE", "32") + ) + + # 自适应收益门控状态（窗口/冷静期均可 env 调，见类 docstring） + self._gate_min_avg = float( + os.getenv("INFINILM_SPEC_MIN_AVG_TOKENS", "2.0") + ) + self._gate_window_size = int(os.getenv("INFINILM_SPEC_GATE_WINDOW", "32")) + if self._gate_window_size < 1: + raise ValueError("INFINILM_SPEC_GATE_WINDOW must be >= 1") + self._gate_cooldown_len = int( + os.getenv("INFINILM_SPEC_GATE_COOLDOWN", "64") + ) + self._gate_window_tokens = 0 + self._gate_window_steps = 0 + self._gate_cooldown = 0 + self.gate_triggered_count = 0 + + self.draft_model_engine = None + if self.speculative_method == "prompt_lookup": + # n-gram 后缀匹配的长度范围（仅影响命中质量，可调参不影响正确性） + self.prompt_lookup_max_ngram = int( + os.getenv("INFINILM_PROMPT_LOOKUP_MAX_NGRAM", "4") + ) + self.prompt_lookup_min_ngram = int( + os.getenv("INFINILM_PROMPT_LOOKUP_MIN_NGRAM", "2") + ) + if ( + self.prompt_lookup_min_ngram < 1 + or self.prompt_lookup_max_ngram < self.prompt_lookup_min_ngram + ): + raise ValueError( + "prompt-lookup n-gram bounds must satisfy " + "1 <= INFINILM_PROMPT_LOOKUP_MIN_NGRAM " + "<= INFINILM_PROMPT_LOOKUP_MAX_NGRAM" + ) + return draft_cache_config = StaticKVCacheConfig( max_batch_size=config.max_batch_size, max_cache_len=config.max_cache_len @@ -39,6 +115,29 @@ def __init__(self, config, target_model_engine, device): dtype=self.draft_model_engine.dtype, ) + def get_acceptance_stats(self) -> dict: + """接受率/收益统计快照（计数自引擎启动起累计）。""" + return { + "method": self.speculative_method, + "num_draft_tokens": self.num_draft_tokens, + "drafted_tokens": self.total_drafted, + "accepted_tokens": self.total_accepted, + "accept_rate": ( + self.total_accepted / self.total_drafted + if self.total_drafted + else None + ), + "spec_steps": self.total_emitted_steps, + "emitted_tokens": self.total_emitted_tokens, + "avg_tokens_per_step": ( + self.total_emitted_tokens / self.total_emitted_steps + if self.total_emitted_steps + else None + ), + "verify_alloc_failures": self.total_verify_alloc_failures, + "gate_triggered": self.gate_triggered_count, + } + def forward(self, scheduler_output, model_input): cache_ops = getattr(scheduler_output, "speculative_cache_ops", None) if cache_ops is None: @@ -46,7 +145,7 @@ def forward(self, scheduler_output, model_input): return sampled_tokens.to_numpy().tolist() # Keep non-greedy sampling on the established target path. Correct stochastic - # speculative sampling needs distribution-level acceptance, while current MTP + # speculative sampling needs distribution-level acceptance, while current # verification is exact for greedy decoding. if self.config.top_k != 1 or self.config.temperature != 1.0: sampled_tokens = self.target_model_engine.forward(**model_input) @@ -56,19 +155,70 @@ def forward(self, scheduler_output, model_input): if not requests: return [] - target_output = self.target_model_engine.forward_raw(**model_input) - target_token_ids = target_output["output_ids"].to_numpy().tolist() + # 大 batch 下 decode 转向计算受限，投机把每步计算量放大 k 倍反而 + # 降吞吐：超过阈值直接走常规前向（decode 形状仍可命中 CUDA graph） + if len(requests) > self.spec_max_batch_size: + sampled_tokens = self.target_model_engine.forward(**model_input) + return sampled_tokens.to_numpy().tolist() + + # 收益门控冷静期：回退常规前向（decode 形状仍可命中 CUDA graph）， + # 冷静期结束后自动重试投机 + if self._gate_cooldown > 0: + self._gate_cooldown -= 1 + sampled_tokens = self.target_model_engine.forward(**model_input) + return sampled_tokens.to_numpy().tolist() + + use_prompt_lookup = self.speculative_method == "prompt_lookup" + # 纯 decode 批次的 prompt_lookup 走融合单前向：一次前向同时完成 + # 主采样与 draft 验证，省掉独立 verify 前向的 launch/CPU 开销。 + # 含 prefill chunk 的混排批次仍走两前向流程（draft 在主前向之后）。 + if use_prompt_lookup and self._is_pure_decode_batch( + requests, scheduler_output + ): + return self._forward_fused_prompt_lookup( + scheduler_output, cache_ops, requests + ) + + if use_prompt_lookup: + # 主前向只取每请求最后位置的采样结果 + sampled = self.target_model_engine.forward(**model_input) + target_token_ids = sampled.to_numpy().tolist() + hidden_states = None + else: + target_output = self.target_model_engine.forward_raw(**model_input) + target_token_ids = target_output["output_ids"].to_numpy().tolist() + hidden_states = target_output["hidden_states"] if not target_token_ids: return target_token_ids input_offsets = model_input["input_offsets"].to_numpy().tolist() - hidden_states = target_output["hidden_states"] + num_scheduled = getattr(scheduler_output, "num_scheduled_tokens", None) output_tokens_by_req: list[list[int]] = [[] for _ in requests] draft_jobs = [] for req_idx, req in enumerate(requests): - last_input_idx = int(input_offsets[req_idx + 1]) - 1 - target_token = int(target_token_ids[last_input_idx]) + # 逐请求阶段判定（chunked prefill 混排批次中 batch 级 is_prefill + # 已退化为"是否包含 chunk"）：prompt 未算完的中段 chunk 请求本步 + # 不产生 token（废 token 由 llm.py 的 _update_requests 丢弃）， + # 也绝不能进入 draft/verify——其序列末端不是真实生成位置， + # 追加 verify 槽位会破坏块表不变量。 + if num_scheduled is not None: + chunk_end = req.num_local_cached_tokens + num_scheduled[req_idx] + if chunk_end < req.get_prompt_length(): + output_tokens_by_req[req_idx] = [] + continue + is_prefill_req = ( + req.num_local_cached_tokens < req.get_prompt_length() + ) + else: + is_prefill_req = scheduler_output.is_prefill + + if use_prompt_lookup: + target_token = int(target_token_ids[req_idx]) + else: + last_input_idx = int(input_offsets[req_idx + 1]) - 1 + target_token = int(target_token_ids[last_input_idx]) + max_tokens = req.sampling_params.max_tokens remaining = ( None @@ -86,23 +236,31 @@ def forward(self, scheduler_output, model_input): output_tokens_by_req[req_idx] = [target_token] continue - source_token, source_position = self._get_last_input_token_and_position( - req, scheduler_output.is_prefill - ) - draft_jobs.append( - { - "req_idx": req_idx, - "req": req, - "target_token": target_token, - "remaining": remaining, - "source_token": source_token, - "source_position": source_position, - "target_hidden": hidden_states.narrow(1, last_input_idx, 1), - "num_tokens": draft_budget, - } - ) + job = { + "req_idx": req_idx, + "req": req, + "target_token": target_token, + "remaining": remaining, + "num_tokens": draft_budget, + } + if not use_prompt_lookup: + ( + source_token, + source_position, + ) = self._get_last_input_token_and_position(req, is_prefill_req) + job["source_token"] = source_token + job["source_position"] = source_position + job["target_hidden"] = hidden_states.narrow(1, last_input_idx, 1) + draft_jobs.append(job) + + if use_prompt_lookup: + draft_results = [ + self._draft_prompt_lookup_tokens(job["req"], job["num_tokens"]) + for job in draft_jobs + ] + else: + draft_results = self._draft_eagle_tokens_batch(draft_jobs) - draft_results = self._draft_eagle_tokens_batch(draft_jobs) verify_candidates = [] for job, draft_tokens in zip(draft_jobs, draft_results): req_idx = job["req_idx"] @@ -112,17 +270,24 @@ def forward(self, scheduler_output, model_input): output_tokens_by_req[req_idx] = [target_token] continue - self.eagle_total_count += len(draft_tokens) + self.total_drafted += len(draft_tokens) if draft_tokens[0] != target_token: output_tokens_by_req[req_idx] = [target_token] continue base_len = req.get_total_length() - verify_block_table, verify_slots = cache_ops.append_verify_slots( - list(req.block_table), - base_len + 1, - len(draft_tokens), - ) + try: + verify_block_table, verify_slots = cache_ops.append_verify_slots( + list(req.block_table), + base_len + 1, + len(draft_tokens), + ) + except RuntimeError: + # 临时验证槽位分配失败（KV 块不足）：退化为只产出 target + # token，等价于非投机的一步，不阻塞正常 decode。 + self.total_verify_alloc_failures += 1 + output_tokens_by_req[req_idx] = [target_token] + continue req.block_table = verify_block_table req.num_blocks = len(req.block_table) verify_candidates.append( @@ -166,7 +331,7 @@ def forward(self, scheduler_output, model_input): if correction is None: correction = int(segment[len(draft_tokens) - 1]) - self.eagle_accept_count += accepted + self.total_accepted += accepted keep_tokens = candidate["base_len"] + accepted req.block_table = cache_ops.rollback_to_length( req.block_table, keep_tokens @@ -180,8 +345,227 @@ def forward(self, scheduler_output, model_input): output_tokens = output_tokens[:remaining] output_tokens_by_req[req_idx] = output_tokens + self._record_step_stats(output_tokens_by_req) + return output_tokens_by_req + def _record_step_stats(self, output_tokens_by_req): + """累积端到端收益统计，并维护自适应门控的滑动窗口。 + + 窗口攒满后结算一次：平均每步产出低于盈亏平衡阈值则进入冷静期 + （forward 开头回退常规前向），窗口清零、冷静期满后自动重试。 + """ + for output_tokens in output_tokens_by_req: + if output_tokens: + self.total_emitted_steps += 1 + self.total_emitted_tokens += len(output_tokens) + self._gate_window_steps += 1 + self._gate_window_tokens += len(output_tokens) + if self._gate_window_steps >= self._gate_window_size: + avg = self._gate_window_tokens / self._gate_window_steps + self._gate_window_steps = 0 + self._gate_window_tokens = 0 + if avg < self._gate_min_avg: + self._gate_cooldown = self._gate_cooldown_len + self.gate_triggered_count += 1 + logger.info( + "spec decoding gated: avg %.2f tok/step < %.2f over the " + "last window, fallback to plain decode for %d steps", + avg, + self._gate_min_avg, + self._gate_cooldown_len, + ) + + def _is_pure_decode_batch(self, requests, scheduler_output) -> bool: + """批次内所有请求都处于 decode 相位（本步调度的 token 位于序列尾部）。 + + chunked 混排中的中段/完成 chunk（num_local_cached_tokens < + prompt_length）与旧式整段 prefill 批次都不算纯 decode。 + """ + num_scheduled = getattr(scheduler_output, "num_scheduled_tokens", None) + if num_scheduled is None: + return not scheduler_output.is_prefill + return all( + req.num_local_cached_tokens >= req.get_prompt_length() + for req in requests + ) + + def _forward_fused_prompt_lookup( + self, scheduler_output, cache_ops, requests + ) -> list[list[int]]: + """prompt_lookup 专用融合前向：主采样 + draft 验证一次完成。 + + 每个 decode 请求贡献固定 k+1 个 token 的一行：[已提交尾 token x, + draft_1..draft_k]（draft 不足 k 个时用尾 token 补齐——验证逻辑对 + 任意 draft 内容自校正，pad 最多造成巧合接受，不影响正确性）。 + 批次形状固定为 b × (k+1)，是后续 verify 图化要录制的形状。 + + 输出与两前向流程逐 token 等价：行输出 o[0] 即 target_token（位置 + T 的预测）；o[j] 是 draft_j 之后位置的预测。接受 drafts 的最长 + 匹配前缀 + 1 个修正 token，未接受的 KV 槽位回滚。 + """ + k = self.num_draft_tokens + candidates = [] + for req_idx, req in enumerate(requests): + base_len = req.get_total_length() + max_tokens = req.sampling_params.max_tokens + remaining = ( + None + if max_tokens is None + else max_tokens - req.get_num_generated_tokens() + ) + last_token = ( + req.generated_token_ids[-1] + if req.generated_token_ids + else req.prompt_token_ids[-1] + ) + real_drafts = ( + self._draft_prompt_lookup_tokens(req, k) + if remaining is None or remaining > 1 + else [] + ) + drafts = list(real_drafts) + [int(last_token)] * ( + k - len(real_drafts) + ) + + # remaining<=1 时与两前向路径一样短路：本步只跑尾 token + # （1 token 行），不分配 verify 槽位——既不浪费 k+1 行前向， + # 也不会把本不需要的分配计进 verify_alloc_failures。 + if remaining is not None and remaining <= 1: + verify_slots = None + else: + # x（位置 base_len-1）的槽位已由调度器分配；draft 占用 + # base_len..base_len+k-1 的临时槽位，验证后按接受长度回滚 + try: + block_table, verify_slots = cache_ops.append_verify_slots( + list(req.block_table), base_len + 1, k + ) + except RuntimeError: + # KV 块不足：本请求只跑尾 token（1 token 行），不投机 + self.total_verify_alloc_failures += 1 + verify_slots = None + else: + req.block_table = block_table + req.num_blocks = len(req.block_table) + + candidates.append( + { + "req_idx": req_idx, + "req": req, + "base_len": base_len, + "remaining": remaining, + "last_token": int(last_token), + "real_draft_len": len(real_drafts), + "drafts": drafts, + "verify_slots": verify_slots, + } + ) + + tokens = [] + position_ids = [] + past_lens = [] + seq_lens = [] + input_offsets = [0] + cu_seqlens = [0] + slot_mapping = [] + block_tables = [] + max_block_table_len = max(len(c["req"].block_table) for c in candidates) + + for c in candidates: + req = c["req"] + base_len = c["base_len"] + row = ( + [c["last_token"]] + c["drafts"] + if c["verify_slots"] is not None + else [c["last_token"]] + ) + tokens.extend(row) + position_ids.extend(range(base_len - 1, base_len - 1 + len(row))) + past_lens.append(base_len - 1) + seq_lens.append(base_len - 1 + len(row)) + input_offsets.append(input_offsets[-1] + len(row)) + cu_seqlens.append(cu_seqlens[-1] + base_len - 1 + len(row)) + slot_mapping.extend(req.slot_mapping) + if c["verify_slots"] is not None: + slot_mapping.extend(c["verify_slots"]) + block_tables.append( + req.block_table + + [-1] * (max_block_table_len - len(req.block_table)) + ) + + fused_output = self.target_model_engine.forward_raw( + input_ids=infinicore.from_list([tokens], dtype=infinicore.int64), + position_ids=infinicore.from_list(position_ids, dtype=infinicore.int64), + past_kv_lengths=infinicore.from_list(past_lens, dtype=infinicore.int32), + total_kv_lengths=infinicore.from_list(seq_lens, dtype=infinicore.int32), + input_offsets=infinicore.from_list(input_offsets, dtype=infinicore.int32), + cu_seqlens=infinicore.from_list(cu_seqlens, dtype=infinicore.int32), + block_tables=infinicore.from_list(block_tables, dtype=infinicore.int32), + slot_mapping=infinicore.from_list(slot_mapping, dtype=infinicore.int64), + temperature=1.0, + top_k=1, + top_p=1.0, + ) + out_ids = fused_output["output_ids"].to_numpy().tolist() + + output_tokens_by_req: list[list[int]] = [] + for c in candidates: + req = c["req"] + row_len = 1 + ( + k if c["verify_slots"] is not None else 0 + ) + segment = out_ids[input_offsets[c["req_idx"]] : input_offsets[c["req_idx"]] + row_len] + target_token = int(segment[0]) + + if c["verify_slots"] is None: + output_tokens_by_req.append([target_token]) + continue + + m = 0 + for j in range(k): + if c["drafts"][j] != int(segment[j]): + break + m += 1 + correction = int(segment[m]) + + self.total_drafted += c["real_draft_len"] + self.total_accepted += min(m, c["real_draft_len"]) + + keep_tokens = c["base_len"] + m + req.block_table = cache_ops.rollback_to_length( + req.block_table, keep_tokens + ) + req.num_blocks = len(req.block_table) + req.slot_mapping = [] + + output_tokens = c["drafts"][:m] + [correction] + if c["remaining"] is not None: + output_tokens = output_tokens[: c["remaining"]] + output_tokens_by_req.append(output_tokens) + + self._record_step_stats(output_tokens_by_req) + + return output_tokens_by_req + + def _draft_prompt_lookup_tokens(self, req, num_tokens: int) -> list[int]: + """Prompt-lookup draft：n-gram 后缀匹配。 + + 在请求自身序列（prompt + 已生成）中查找当前后缀的最近一次出现， + 取该出现之后的最多 num_tokens 个 token 作为 draft；命中不了返回 []。 + 正确性不依赖命中质量——所有 draft 都会经过 target 模型的精确验证。 + """ + context = list(req.prompt_token_ids) + list(req.generated_token_ids) + n = len(context) + max_ngram = min(self.prompt_lookup_max_ngram, n - 1) + for size in range(max_ngram, self.prompt_lookup_min_ngram - 1, -1): + pattern = context[n - size :] + # 从后往前找最近一次出现（排除末尾的 pattern 自身） + for i in range(n - size - 1, -1, -1): + if context[i : i + size] == pattern: + start = i + size + return context[start : start + num_tokens] + return [] + def _get_last_input_token_and_position(self, req, is_prefill): if is_prefill: return req.prompt_token_ids[-1], req.prompt_length - 1 diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index c10b55f2f..11aae8ba5 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -44,11 +44,16 @@ def __init__( scheduled_requests: List[InferenceRequest], is_prefill: bool = False, speculative_cache_ops: Optional[SpeculativeCacheOps] = None, + num_scheduled_tokens: Optional[List[int]] = None, ): self.scheduled_requests = scheduled_requests self.num_requests = len(scheduled_requests) self.is_prefill = is_prefill self.speculative_cache_ops = speculative_cache_ops + # 与 scheduled_requests 逐请求对齐的本步调度 token 数： + # decode 请求为 1，prefill chunk 为本块长度。None 表示旧语义 + # （由 is_prefill 区分整段 prefill 与单 token decode）。 + self.num_scheduled_tokens = num_scheduled_tokens self.kv_connector_metadata = None @@ -71,6 +76,7 @@ def __init__( has_mamba_cache: bool = False, num_mamba_cache_blocks: int | None = None, enable_prefix_caching: bool = True, + enable_chunked_prefill: bool = False, ): self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() @@ -94,6 +100,7 @@ def __init__( self.max_num_batched_tokens = max_num_batched_tokens self.connector = connector self.enable_prefix_caching = enable_prefix_caching + self.enable_chunked_prefill = enable_chunked_prefill def add_request(self, request: InferenceRequest): if request is not None: @@ -128,6 +135,8 @@ def _exceeds_token_budget( def schedule(self) -> Optional[SchedulerOutput]: """Schedule and return batch of requests to execute.""" + if self.enable_chunked_prefill: + return self._schedule_chunked() deferred_requests = [] scheduled_requests = [] is_prefill = False @@ -376,6 +385,287 @@ def schedule(self) -> Optional[SchedulerOutput]: return None + def _schedule_chunked(self) -> Optional[SchedulerOutput]: + """vLLM 式统一 token-budget 调度（chunked prefill + prefill/decode 混排）。 + + 每个 step 先给 running 队列的每个请求排 1 个 decode token，再把剩余的 + max_num_batched_tokens 预算切给 waiting 队列请求的 prefill chunk + （chunk = min(剩余 prompt, 剩余预算)），二者可混入同一批次。 + 多模态请求与远端 KV 加载请求不切分，保持整段 prefill 的旧行为。 + + 批次级 is_prefill 退化为派生语义（本批次是否包含 prefill chunk）， + 逐请求的阶段与调度长度由 SchedulerOutput.num_scheduled_tokens 表达。 + + 注意：prefill 首次调度时仍按整段 prompt 分配并预留全部 KV 块 + （与旧行为一致的准入/预留粒度），仅 slot_mapping 按 chunk 切片； + 这样中段 chunk 请求放回 waiting 队列期间，其后续块需求依旧被 + can_accept_request 的预留记账覆盖，续跑时不会出现块不足。 + """ + deferred_requests = [] + scheduled_requests = [] + num_scheduled_tokens = [] + is_prefill = False + current_num_batched_tokens = 0 + current_prefill_extra_blocks = 0 + + # Process Running queue (decode phase): 每请求 1 个 token，优先占用预算 + while ( + len(scheduled_requests) < self.max_batch_size + and current_num_batched_tokens < self.max_num_batched_tokens + ): + try: + req = self.running_queue.sync_q.get_nowait() + except queue.Empty: + break + # Skip requests that were already finished (e.g., timed out/canceled while running) + if req.is_finished(): + self.complete_requests([req]) + continue + + # Decode phase: allocate slot for newly generated token + req.block_table, new_slot = self.cache_manager.append_slot( + req.block_table, req.get_total_length() + ) + req.slot_mapping = [new_slot] + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = req.get_total_length() - 1 + scheduled_requests.append(req) + num_scheduled_tokens.append(1) + current_num_batched_tokens += 1 + + # Promote completed remote KV transfers (lower priority than running queue). + # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. + if self.connector is not None and self.remote_kv_requests: + for req_id in list(self.remote_kv_requests.keys()): + req = self.remote_kv_requests[req_id] + if req.is_finished(): + self.complete_requests([req]) + continue + if req_id in self.failed_receiving_kv_req_ids: + logger.warning( + f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." + ) + self.update_waiting_for_remote_kv(req) + req.status = RequestStatus.WAITING + self.waiting_queue.sync_q.put(req) + elif req_id in self.finished_receiving_kv_req_ids: + if len(scheduled_requests) < self.max_batch_size: + logger.info( + f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." + ) + self.update_waiting_for_remote_kv(req) + req.status = RequestStatus.RUNNING + scheduled_requests.append(req) + num_scheduled_tokens.append(1) + current_num_batched_tokens += 1 + else: + break # Defer promotion to next schedule() if batch is full + + # Process Waiting queue (prefill phase): 剩余预算按 chunk 切分 + while ( + len(scheduled_requests) < self.max_batch_size + and current_num_batched_tokens < self.max_num_batched_tokens + ): + try: + req = self.waiting_queue.sync_q.get_nowait() + except queue.Empty: + break + # Skip requests that were already finished (e.g., timed out/canceled while waiting) + if req.is_finished(): + self.complete_requests([req]) + continue + + if req.num_computed_tokens == 0: + if self.has_mamba_cache: + cached_block_table = [] + num_local_computed_tokens = 0 + load_kv_async = False + num_external_computed_tokens = 0 + else: + if self.enable_prefix_caching: + ( + cached_block_table, + num_local_computed_tokens, + ) = self.cache_manager.get_computed_blocks( + req.block_hashes, req.get_prompt_length() - 1 + ) + else: + cached_block_table = [] + num_local_computed_tokens = 0 + if self.connector is not None: + ext_tokens, load_kv_async = ( + self.connector.get_num_new_matched_tokens( + req, num_local_computed_tokens + ) + ) + num_external_computed_tokens = ext_tokens + else: + load_kv_async = False + num_external_computed_tokens = 0 + + available_cached_tokens = ( + num_local_computed_tokens + num_external_computed_tokens + ) + num_computed_tokens = min( + available_cached_tokens, + max(req.get_prompt_length() - 1, 0), + ) + # 整段分配（与旧行为一致的准入与块预留），仅计算量按 chunk 调度 + num_new_tokens = req.get_prompt_length() - num_computed_tokens + + # 多模态请求与远端 KV 加载请求不切分，保持整段 prefill + if load_kv_async or req.has_multimodal_inputs: + num_tokens_this_step = num_new_tokens + else: + num_tokens_this_step = min( + num_new_tokens, + self.max_num_batched_tokens - current_num_batched_tokens, + ) + + # chunk 粒度的预算检查；单请求保底放行语义不变 + if not load_kv_async and self._exceeds_token_budget( + current_num_batched_tokens, + num_tokens_this_step, + len(scheduled_requests), + ): + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + if not self.can_accept_request( + req, + num_local_computed_tokens, + current_prefill_extra_blocks, + ): + logger.warning( + "Insufficient KV cache blocks for request %s, deferring.", + req.request_id, + ) + + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + allocation = self.cache_manager.allocate_slots( + num_new_tokens, + num_computed_tokens=num_computed_tokens, + cached_block_table=cached_block_table, + ) + + if allocation is None: + logger.warning( + "Failed to allocate KV cache blocks for request: %s", + req.request_id, + ) + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + req_blocks, slot_mapping = allocation + + if self.has_mamba_cache and req.mamba_cache_index is None: + req.mamba_cache_index = self.mamba_cache_manager.allocate() + if req.mamba_cache_index is None: + self.cache_manager.free_blocks(req_blocks) + logger.warning( + "Insufficient mamba cache rows for request %s, deferring.", + req.request_id, + ) + deferred_requests.append(req) + break + + req.block_table = req_blocks + req.num_blocks = len(req_blocks) + if load_kv_async: + req.slot_mapping = slot_mapping + else: + # 只写入本 chunk 覆盖的槽位 + req.slot_mapping = slot_mapping[:num_tokens_this_step] + req.num_local_cached_tokens = ( + num_local_computed_tokens if load_kv_async else num_computed_tokens + ) + req.num_cache_indexed_blocks = len(cached_block_table) + req.num_computed_tokens = num_computed_tokens + + if self.connector is not None: + self.connector.update_state_after_alloc( + req, + req.block_table, + num_external_computed_tokens, + self.block_size, + ) + else: + # 继续未算完的 prefill：块已在首次调度时整段分配， + # 直接切出本 chunk 覆盖的槽位 + load_kv_async = False + num_tokens_this_step = req.get_prompt_length() - req.num_computed_tokens + if not req.has_multimodal_inputs: + num_tokens_this_step = min( + num_tokens_this_step, + self.max_num_batched_tokens - current_num_batched_tokens, + ) + if self._exceeds_token_budget( + current_num_batched_tokens, + num_tokens_this_step, + len(scheduled_requests), + ): + deferred_requests.append(req) + break + self.commit_computed_tokens(req, req.num_computed_tokens) + chunk_end = req.num_computed_tokens + num_tokens_this_step + req.slot_mapping = self.cache_manager.update_blocks_slot( + req.block_table, + req.num_computed_tokens, + chunk_end, + ) + req.num_local_cached_tokens = req.num_computed_tokens + + if load_kv_async: + req.status = RequestStatus.WAITING_FOR_REMOTE_KVS + self.remote_kv_requests[req.request_id] = req + self.pending_kv_decode_blocks += ( + req.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + continue + + current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) + scheduled_requests.append(req) + num_scheduled_tokens.append(num_tokens_this_step) + current_num_batched_tokens += num_tokens_this_step + is_prefill = True + + req.status = RequestStatus.RUNNING + + if deferred_requests: + for req in deferred_requests: + self.waiting_queue.sync_q.put(req) + + # Return mixed batch if any requests were scheduled + if scheduled_requests: + scheduler_output = SchedulerOutput( + scheduled_requests=scheduled_requests, + is_prefill=is_prefill, + speculative_cache_ops=self.speculative_cache_ops, + num_scheduled_tokens=num_scheduled_tokens, + ) + if self.connector is not None: + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + if self.connector is not None: + scheduler_output = SchedulerOutput( + scheduled_requests=[], + speculative_cache_ops=self.speculative_cache_ops, + ) + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + return None + def update_waiting_for_remote_kv(self, request: InferenceRequest): self.remote_kv_requests.pop(request.request_id, None) self.pending_kv_decode_blocks -= ( @@ -446,6 +736,10 @@ def complete_requests(self, requests: List[InferenceRequest]): logger.error( f"Request {req.request_id[:8]}... timed out: {req.finish_reason}" ) + elif req.status == RequestStatus.WAITING: + # chunked prefill 的中段 chunk 请求：prompt 未算完， + # 放回 waiting 队列等待下一次 chunk 调度 + self.waiting_queue.sync_q.put(req) else: # Still running, put back in running queue self.running_queue.sync_q.put(req) diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..05b0e01d8 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -180,6 +180,11 @@ def _build_model_input_from_batch_scheduler_output( - input_offsets: Offsets for each request - block_tables: Padded block_table for each request - slot_mapping: Single slot per request + + Chunked prefill（num_scheduled_tokens 非 None）时按逐请求切片： + prompt 未算完的请求是本步调度长度的 prefill chunk，其余为 decode； + total_kv_lengths/cu_seqlens 写本 chunk 结束后的可见 KV 长度， + 两种请求可混入同一批次。 """ import infinicore @@ -202,22 +207,37 @@ def _build_model_input_from_batch_scheduler_output( ) current_offset = 0 - for req in scheduler_output.scheduled_requests: + for req_idx, req in enumerate(scheduler_output.scheduled_requests): num_cached = req.num_local_cached_tokens - if scheduler_output.is_prefill: - # Prefill phase + if scheduler_output.num_scheduled_tokens is not None: + # chunked prefill：逐请求按本步调度的 token 数切片 + num_new_tokens = scheduler_output.num_scheduled_tokens[req_idx] + is_prefill_req = num_cached < req.get_prompt_length() + else: + num_new_tokens = None + is_prefill_req = scheduler_output.is_prefill + + if is_prefill_req: + # Prefill phase（可能是一个 chunk） req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + if num_new_tokens is None: + tokens_to_compute = req_tokens[num_cached:] + else: + tokens_to_compute = req_tokens[ + num_cached : num_cached + num_new_tokens + ] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) + # 本 chunk 结束后的可见 KV 长度；绝不能写整段 prompt 长， + # 否则注意力会读到尚未写入的 KV + seq_len = num_cached + compute_len seq_lens.append(seq_len) current_offset += compute_len seq_offsets.append(current_offset) - slot_mapping.extend(req.slot_mapping) + slot_mapping.extend(req.slot_mapping[:compute_len]) cached_lens.append(num_cached) position_ids.extend(range(num_cached, num_cached + compute_len)) diff --git a/test/test_chunked_prefill.py b/test/test_chunked_prefill.py new file mode 100644 index 000000000..2a19367e4 --- /dev/null +++ b/test/test_chunked_prefill.py @@ -0,0 +1,434 @@ +"""Chunked prefill + prefill/decode 混排的离线冒烟测试（无需 GPU）。 + +infinilm 包本体依赖 infinicore/torch/transformers/janus 等重型库，本测试通过 +sys.modules 预置轻量 stub，只真实加载纯 Python 的 scheduler / cache_manager / +request / basic_llm_processor / llm 模块，验证： + +1. chunked 模式下 decode 优先、prefill 按剩余预算切块，二者可混入同一批次； +2. 混排批次构建的 input_ids/position_ids/past_kv_lengths/total_kv_lengths/ + cu_seqlens/input_offsets/slot_mapping 均按 (num_cached, chunk_len) 切片； +3. 中段 chunk 请求不 append 废 token、只推进 num_computed_tokens 并逐 chunk + 发布 prefix-cache 块； +4. 多模态请求不切分，保持整段 prefill； +5. 开关关闭时（旧模式）调度与输入构建行为与旧逻辑一致。 + +直接运行: python3 test/test_chunked_prefill.py +""" + +import collections +import hashlib +import queue +import struct +import sys +import types +from pathlib import Path + +REPO_PYTHON = Path(__file__).resolve().parents[1] / "python" +sys.path.insert(0, str(REPO_PYTHON)) + + +def _install_stubs(): + """预置轻量 stub，绕过 infinicore/torch/transformers/janus 等重型依赖。""" + # 1) 包占位：避免执行各 __init__.py（会链入 torch/infinicore） + pkg_root = REPO_PYTHON / "infinilm" + for name, rel in [ + ("infinilm", ""), + ("infinilm.llm", "llm"), + ("infinilm.llm.model_runner", "llm/model_runner"), + ("infinilm.processors", "processors"), + ("infinilm.config", "config"), + ("infinilm.multimodal", "multimodal"), + ]: + mod = types.ModuleType(name) + mod.__path__ = [str(pkg_root / rel) if rel else str(pkg_root)] + sys.modules[name] = mod + + # 2) janus：调度器只用到 sync_q 的阻塞队列接口 + janus = types.ModuleType("janus") + + class _SyncQueue: + def __init__(self): + self._q = collections.deque() + + def put(self, item, block=True, timeout=None): + self._q.append(item) + + put_nowait = put + + def get(self, block=True, timeout=None): + if not self._q: + raise queue.Empty + return self._q.popleft() + + def get_nowait(self): + return self.get() + + def qsize(self): + return len(self._q) + + def empty(self): + return not self._q + + class _JanusQueue: + def __init__(self, maxsize=0): + self.sync_q = _SyncQueue() + self.async_q = None + + janus.Queue = _JanusQueue + sys.modules["janus"] = janus + + # 3) numpy / xxhash：prefix_cache 只要求确定性的 16 字节哈希，用标准库模拟 + numpy = types.ModuleType("numpy") + + class _FakeArray: + def __init__(self, token_ids): + self._token_ids = list(token_ids) + + def __len__(self): + return len(self._token_ids) + + def tobytes(self): + return struct.pack(f"<{len(self._token_ids)}i", *self._token_ids) + + numpy.asarray = lambda token_ids, dtype=None: _FakeArray(token_ids) + sys.modules["numpy"] = numpy + + xxhash = types.ModuleType("xxhash") + + class _Hasher: + def __init__(self): + self._buf = bytearray() + + def update(self, data): + self._buf += data + + def digest(self): + return hashlib.sha256(bytes(self._buf)).digest()[:16] + + xxhash.xxh3_128 = _Hasher + sys.modules["xxhash"] = xxhash + + # 4) transformers：basic_llm_processor 顶层 import 用到 + transformers = types.ModuleType("transformers") + transformers.AutoTokenizer = object + transformers.AutoProcessor = object + sys.modules["transformers"] = transformers + + # 5) infinicore：builder 内部 import，from_list 只做数据捕获 + infinicore = types.ModuleType("infinicore") + infinicore.int64 = "int64" + infinicore.int32 = "int32" + + class _FakeTensor: + def __init__(self, data, dtype): + self.data = data + self.dtype = dtype + + infinicore.from_list = lambda data, dtype=None: _FakeTensor(data, dtype) + sys.modules["infinicore"] = infinicore + + # 6) llm.py 的重型依赖（只在引擎初始化时才真正使用） + engine_config = types.ModuleType("infinilm.config.engine_config") + engine_config.EngineConfig = object + sys.modules["infinilm.config.engine_config"] = engine_config + + kv_transfer = types.ModuleType("infinilm.config.kv_transfer") + kv_transfer.KVTransferConfig = object + sys.modules["infinilm.config.kv_transfer"] = kv_transfer + + infer_engine = types.ModuleType("infinilm.infer_engine") + infer_engine.model_uses_mamba_cache = lambda hf_config: False + infer_engine.read_hf_config = lambda model_path: {} + sys.modules["infinilm.infer_engine"] = infer_engine + + kv_connector = types.ModuleType("infinilm.kv_connector") + kv_connector.KVConnectorFactory = object + kv_connector.KVConnectorRole = object + sys.modules["infinilm.kv_connector"] = kv_connector + + model_runner = types.ModuleType("infinilm.llm.model_runner.model_runner") + model_runner.ModelRunner = object + sys.modules["infinilm.llm.model_runner.model_runner"] = model_runner + + multimodal = types.ModuleType("infinilm.multimodal.multimodal") + multimodal.resolve_multimodal_inputs = lambda messages: {} + sys.modules["infinilm.multimodal.multimodal"] = multimodal + + +_install_stubs() + +from infinilm.llm.llm import LLMEngine # noqa: E402 +from infinilm.llm.request import InferenceRequest # noqa: E402 +from infinilm.llm.sampling_params import SamplingParams # noqa: E402 +from infinilm.llm.scheduler import Scheduler # noqa: E402 +from infinilm.processors.basic_llm_processor import BasicLLMProcessor # noqa: E402 + +_SAMPLED_TOKEN = 42 # 模拟 C++ 每请求采出的 token（非 eos） + + +def _make_request(req_id, prompt_token_ids, max_tokens=3, has_multimodal_inputs=False): + return InferenceRequest( + request_id=req_id, + prompt="x" * len(prompt_token_ids), + prompt_token_ids=list(prompt_token_ids), + sampling_params=SamplingParams(max_tokens=max_tokens), + eos_token_ids=[], + has_multimodal_inputs=has_multimodal_inputs, + ) + + +def _make_engine(scheduler): + engine = LLMEngine.__new__(LLMEngine) + engine.scheduler = scheduler + engine.eos_token_ids = [] + + class _FakeTokenizer: + def decode(self, token_ids): + return "" + + engine.tokenizer = _FakeTokenizer() + return engine + + +def _make_processor(): + return BasicLLMProcessor.__new__(BasicLLMProcessor) + + +def _run_one_step(engine, processor): + """跑一个调度 step：真实 schedule + 真实输入构建 + 真实 _update_requests。""" + out = engine.scheduler.schedule() + if out is None or not out.scheduled_requests: + return None, None + model_input = processor._build_model_input_from_batch_scheduler_output( + out, 1.0, 0.8, 1 + ) + # 模拟 C++：每个请求在其最后输入位置采一个 token（1:1 对齐） + sampled = [_SAMPLED_TOKEN] * out.num_requests + engine._update_requests( + out.scheduled_requests, + sampled, + getattr(out, "num_scheduled_tokens", None), + ) + return out, model_input + + +def _run_until_finished(engine, processor, requests, max_steps=50): + records = [] + for _ in range(max_steps): + out, model_input = _run_one_step(engine, processor) + if out is None: + break + records.append((out, model_input)) + if all(r.is_finished() for r in requests): + break + return records + + +def _make_scheduler(enable_chunked_prefill, enable_prefix_caching=False, budget=6): + return Scheduler( + max_batch_size=8, + num_blocks=64, + block_size=4, + max_num_batched_tokens=budget, + enable_prefix_caching=enable_prefix_caching, + enable_chunked_prefill=enable_chunked_prefill, + ) + + +def test_chunk_split_and_mixed_batch(): + sched = _make_scheduler(enable_chunked_prefill=True) + engine = _make_engine(sched) + processor = _make_processor() + + req_a = _make_request("a", range(10), max_tokens=3) + req_b = _make_request("b", range(1000, 1005), max_tokens=3) + sched.add_request(req_a) + sched.add_request(req_b) + + # step 1: 预算 6，A 切 6 个 token 的中段 chunk；B 无预算留在 waiting + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [6], out.num_scheduled_tokens + assert out.is_prefill + assert mi["input_ids"].data == [list(range(0, 6))] + assert mi["position_ids"].data == [0, 1, 2, 3, 4, 5] + assert mi["past_kv_lengths"].data == [0] + assert mi["total_kv_lengths"].data == [6] + assert mi["cu_seqlens"].data == [0, 6] + assert mi["slot_mapping"].data == [0, 1, 2, 3, 4, 5] + # 中段 chunk：不 append 废 token，只推进 num_computed_tokens + assert req_a.num_computed_tokens == 6 + assert req_a.get_num_generated_tokens() == 0 + assert req_a.generated_text == "" + + # step 2: B 整段 5 个 token 完成 prefill，剩余预算 1 给 A 续 1 个 token + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [5, 1], out.num_scheduled_tokens + assert out.is_prefill + assert mi["input_ids"].data == [[1000, 1001, 1002, 1003, 1004, 6]] + assert mi["past_kv_lengths"].data == [0, 6] + assert mi["total_kv_lengths"].data == [5, 7] + assert mi["cu_seqlens"].data == [0, 5, 12] + assert mi["input_offsets"].data == [0, 5, 6] + assert req_b.get_num_generated_tokens() == 1 # B 最后一块 chunk，产出首 token + assert req_a.num_computed_tokens == 7 + assert req_a.get_num_generated_tokens() == 0 # A 仍是中段 chunk + + # step 3: B decode 1 个 token + A 续 3 个 token 的 chunk，混入同一批次 + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [1, 3], out.num_scheduled_tokens + assert out.is_prefill # 派生语义：批次中包含 prefill chunk + assert mi["input_ids"].data == [[_SAMPLED_TOKEN, 7, 8, 9]] + assert mi["position_ids"].data == [5, 7, 8, 9] + assert mi["past_kv_lengths"].data == [5, 7] + # total_kv_lengths 是本 chunk 结束后的可见 KV 长度，不是整段 prompt 长 + assert mi["total_kv_lengths"].data == [6, 10] + assert mi["cu_seqlens"].data == [0, 6, 16] + assert mi["input_offsets"].data == [0, 1, 4] + assert mi["slot_mapping"].data == [17, 7, 8, 9] + assert mi["block_tables"].data == [[3, 4, -1], [0, 1, 2]] + assert req_a.get_num_generated_tokens() == 1 # A 完成 prefill，产出首 token + + # step 4: 纯 decode 批次 + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [1, 1], out.num_scheduled_tokens + assert not out.is_prefill + + # 后续步骤预算都不超限，直至两个请求全部完成 + while not (req_a.is_finished() and req_b.is_finished()): + out, _ = _run_one_step(engine, processor) + assert out is not None + assert sum(out.num_scheduled_tokens) <= 6 + assert list(req_a.generated_token_ids) == [_SAMPLED_TOKEN] * 3 + assert list(req_b.generated_token_ids) == [_SAMPLED_TOKEN] * 3 + # 全部结束后可用块数恢复（无泄漏） + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_legacy_mode_unchanged(): + sched = _make_scheduler(enable_chunked_prefill=False) + engine = _make_engine(sched) + processor = _make_processor() + + req_a = _make_request("a", range(10), max_tokens=3) + sched.add_request(req_a) + + # step 1: 旧逻辑整段 prefill，num_scheduled_tokens 为 None + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens is None + assert out.is_prefill + assert mi["input_ids"].data == [list(range(10))] + assert mi["total_kv_lengths"].data == [10] + assert req_a.get_num_generated_tokens() == 1 + + # step 2: A 在 decode 时来了 B，旧逻辑 prefill/decode 互斥： + # 本步只排 B 的整段 prefill，A 的 decode 停摆 + req_b = _make_request("b", range(1000, 1005), max_tokens=2) + sched.add_request(req_b) + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens is None + assert out.is_prefill + assert [r.request_id for r in out.scheduled_requests] == ["b"] + assert mi["input_ids"].data == [list(range(1000, 1005))] + assert mi["total_kv_lengths"].data == [5] + + # step 3: 纯 decode 批次 + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens is None + assert not out.is_prefill + assert {r.request_id for r in out.scheduled_requests} == {"a", "b"} + + _run_until_finished(engine, processor, [req_a, req_b]) + assert list(req_a.generated_token_ids) == [_SAMPLED_TOKEN] * 3 + assert list(req_b.generated_token_ids) == [_SAMPLED_TOKEN] * 2 + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_multimodal_request_not_chunked(): + sched = _make_scheduler(enable_chunked_prefill=True) + engine = _make_engine(sched) + processor = _make_processor() + + req_a = _make_request("a", range(4), max_tokens=3) + sched.add_request(req_a) + # step 1: A 整段 prefill（4 <= 预算 6） + out, _ = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [4] + + req_mm = _make_request( + "mm", range(1000, 1010), max_tokens=2, has_multimodal_inputs=True + ) + sched.add_request(req_mm) + + # step 2: A decode 占 1 预算，多模态请求不切分且整体超预算 -> 整请求推迟 + out, mi = _run_one_step(engine, processor) + assert [r.request_id for r in out.scheduled_requests] == ["a"] + assert out.num_scheduled_tokens == [1] + assert not out.is_prefill + + # step 3: 同上，A 生成最后一个 token 后结束 + out, _ = _run_one_step(engine, processor) + assert [r.request_id for r in out.scheduled_requests] == ["a"] + assert req_a.is_finished() + + # step 4: running 已空，多模态请求作为单请求保底放行，整段 prefill + out, mi = _run_one_step(engine, processor) + assert [r.request_id for r in out.scheduled_requests] == ["mm"] + assert out.num_scheduled_tokens == [10], out.num_scheduled_tokens + assert out.is_prefill + assert mi["total_kv_lengths"].data == [10] + assert req_mm.get_num_generated_tokens() == 1 + + _run_until_finished(engine, processor, [req_mm]) + assert list(req_mm.generated_token_ids) == [_SAMPLED_TOKEN] * 2 + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_prefix_cache_across_chunks(): + sched = _make_scheduler( + enable_chunked_prefill=True, enable_prefix_caching=True, budget=6 + ) + engine = _make_engine(sched) + processor = _make_processor() + prompt_ids = list(range(12)) # block_size=4，共 3 个完整块 + + req_a = _make_request("a", prompt_ids, max_tokens=1) + sched.add_request(req_a) + + # step 1: chunk [0,6)，commit 后只发布第 0 个完整块 + out, _ = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [6] + assert req_a.num_computed_tokens == 6 + assert req_a.num_cache_indexed_blocks == 1 + assert len(sched.cache_manager.hash_to_block_ids) == 1 + + # step 2: chunk [6,12)，prefill 完成，3 个块全部发布 + out, _ = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [6] + assert req_a.is_finished() + assert req_a.num_cache_indexed_blocks == 3 + assert len(sched.cache_manager.hash_to_block_ids) == 3 + + # step 3: 相同 prompt 的 B 命中前缀缓存（最多复用到 prompt_len-1）， + # 只需计算 [8,12) 一个 chunk + req_b = _make_request("b", prompt_ids, max_tokens=1) + sched.add_request(req_b) + out, mi = _run_one_step(engine, processor) + assert out.num_scheduled_tokens == [4] + assert req_b.num_local_cached_tokens == 8 + assert mi["input_ids"].data == [[8, 9, 10, 11]] + assert mi["past_kv_lengths"].data == [8] + assert mi["total_kv_lengths"].data == [12] + assert req_b.is_finished() + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +if __name__ == "__main__": + test_chunk_split_and_mixed_batch() + print("PASS test_chunk_split_and_mixed_batch") + test_legacy_mode_unchanged() + print("PASS test_legacy_mode_unchanged") + test_multimodal_request_not_chunked() + print("PASS test_multimodal_request_not_chunked") + test_prefix_cache_across_chunks() + print("PASS test_prefix_cache_across_chunks") + print("ALL OK") diff --git a/test/test_prompt_lookup_spec.py b/test/test_prompt_lookup_spec.py new file mode 100644 index 000000000..d9c9cdcad --- /dev/null +++ b/test/test_prompt_lookup_spec.py @@ -0,0 +1,554 @@ +"""prompt-lookup 投机采样 + spec×chunked 混排的离线冒烟测试（无需 GPU）。 + +与 test_chunked_prefill.py 相同的 stub 思路：预置轻量 sys.modules stub， +真实加载 scheduler / cache_manager / request / basic_llm_processor / llm / +speculative_runner。target 模型用确定性伪模型（next token 只由当前 token +决定的 succ 函数，可配 trap）替代，从而可以精确断言： + +1. prompt-lookup draft 辅助函数的 n-gram 后缀匹配语义（最近一次出现、 + min/max ngram 边界、k 截断、未命中返回 []）； +2. 投机路径输出与非投机贪心路径逐 token 相同（数学无损），命中时接受率 + 计数增长、平均每步产出 > 1； +3. draft 被拒绝（整体/部分）时输出仍然精确，verify 临时槽位被回滚， + 结束后 KV 块无泄漏； +4. spec × chunked prefill 混排：中段 chunk 请求不参与 draft/verify、 + 产出 []，decode 请求在同一批次里正常投机，完成 prefill 的请求在 + 混排批次里也能投机； +5. 非 greedy 配置回退到常规前向（输出形状为每请求单个 token）； +6. 自适应收益门控：零命中负载攒满窗口后回退常规前向（冷静期内无 + 融合前向调用），输出仍与非投机真值逐 token 相同。 + +直接运行: python3 test/test_prompt_lookup_spec.py +""" + +import collections +import hashlib +import queue +import struct +import sys +import types +from pathlib import Path + +REPO_PYTHON = Path(__file__).resolve().parents[1] / "python" +sys.path.insert(0, str(REPO_PYTHON)) + + +def _install_stubs(): + """预置轻量 stub，绕过 infinicore/torch/transformers/janus 等重型依赖。""" + # 1) 包占位：避免执行各 __init__.py（会链入 torch/infinicore） + pkg_root = REPO_PYTHON / "infinilm" + for name, rel in [ + ("infinilm", ""), + ("infinilm.llm", "llm"), + ("infinilm.llm.model_runner", "llm/model_runner"), + ("infinilm.processors", "processors"), + ("infinilm.config", "config"), + ("infinilm.multimodal", "multimodal"), + ("infinilm.cache", "cache"), + ]: + mod = types.ModuleType(name) + mod.__path__ = [str(pkg_root / rel) if rel else str(pkg_root)] + sys.modules[name] = mod + + # 2) janus：调度器只用到 sync_q 的阻塞队列接口 + janus = types.ModuleType("janus") + + class _SyncQueue: + def __init__(self): + self._q = collections.deque() + + def put(self, item, block=True, timeout=None): + self._q.append(item) + + put_nowait = put + + def get(self, block=True, timeout=None): + if not self._q: + raise queue.Empty + return self._q.popleft() + + def get_nowait(self): + return self.get() + + def qsize(self): + return len(self._q) + + def empty(self): + return not self._q + + class _JanusQueue: + def __init__(self, maxsize=0): + self.sync_q = _SyncQueue() + self.async_q = None + + janus.Queue = _JanusQueue + sys.modules["janus"] = janus + + # 3) numpy / xxhash：prefix_cache 只要求确定性的 16 字节哈希，用标准库模拟 + numpy = types.ModuleType("numpy") + + class _FakeArray: + def __init__(self, token_ids): + self._token_ids = list(token_ids) + + def __len__(self): + return len(self._token_ids) + + def tobytes(self): + return struct.pack(f"<{len(self._token_ids)}i", *self._token_ids) + + numpy.asarray = lambda token_ids, dtype=None: _FakeArray(token_ids) + sys.modules["numpy"] = numpy + + xxhash = types.ModuleType("xxhash") + + class _Hasher: + def __init__(self): + self._buf = bytearray() + + def update(self, data): + self._buf += data + + def digest(self): + return hashlib.sha256(bytes(self._buf)).digest()[:16] + + xxhash.xxh3_128 = _Hasher + sys.modules["xxhash"] = xxhash + + # 4) transformers：basic_llm_processor 顶层 import 用到 + transformers = types.ModuleType("transformers") + transformers.AutoTokenizer = object + transformers.AutoProcessor = object + sys.modules["transformers"] = transformers + + # 5) infinicore：from_list 只做数据捕获 + infinicore = types.ModuleType("infinicore") + infinicore.int64 = "int64" + infinicore.int32 = "int32" + + class _FakeTensor: + def __init__(self, data, dtype=None): + self.data = data + self.dtype = dtype + + def to_numpy(self): + return self + + def tolist(self): + return self.data + + infinicore.from_list = lambda data, dtype=None: _FakeTensor(data, dtype) + infinicore._FakeTensor = _FakeTensor + sys.modules["infinicore"] = infinicore + + # 6) speculative_runner / llm 的重型依赖（prompt_lookup 路径不会真正调用） + engine_config = types.ModuleType("infinilm.config.engine_config") + engine_config.EngineConfig = object + sys.modules["infinilm.config.engine_config"] = engine_config + + kv_transfer = types.ModuleType("infinilm.config.kv_transfer") + kv_transfer.KVTransferConfig = object + sys.modules["infinilm.config.kv_transfer"] = kv_transfer + + infer_engine = types.ModuleType("infinilm.infer_engine") + infer_engine.model_uses_mamba_cache = lambda hf_config: False + infer_engine.read_hf_config = lambda model_path: {} + infer_engine.InferEngine = object + sys.modules["infinilm.infer_engine"] = infer_engine + + cache_mod = types.ModuleType("infinilm.cache.cache") + cache_mod.StaticKVCacheConfig = object + sys.modules["infinilm.cache.cache"] = cache_mod + + distributed = types.ModuleType("infinilm.distributed") + distributed.DistConfig = object + sys.modules["infinilm.distributed"] = distributed + + modeling_utils = types.ModuleType("infinilm.modeling_utils") + modeling_utils.load_model_state_dict_by_file = lambda *a, **k: None + sys.modules["infinilm.modeling_utils"] = modeling_utils + + kv_connector = types.ModuleType("infinilm.kv_connector") + kv_connector.KVConnectorFactory = object + kv_connector.KVConnectorRole = object + sys.modules["infinilm.kv_connector"] = kv_connector + + model_runner = types.ModuleType("infinilm.llm.model_runner.model_runner") + model_runner.ModelRunner = object + sys.modules["infinilm.llm.model_runner.model_runner"] = model_runner + + multimodal = types.ModuleType("infinilm.multimodal.multimodal") + multimodal.resolve_multimodal_inputs = lambda messages: {} + sys.modules["infinilm.multimodal.multimodal"] = multimodal + + +_install_stubs() + +from infinilm.llm.llm import LLMEngine # noqa: E402 +from infinilm.llm.model_runner.speculative_runner import ( # noqa: E402 + SpeculativeRunner, +) +from infinilm.llm.request import InferenceRequest # noqa: E402 +from infinilm.llm.sampling_params import SamplingParams # noqa: E402 +from infinilm.llm.scheduler import Scheduler # noqa: E402 +from infinilm.processors.basic_llm_processor import ( # noqa: E402 + BasicLLMProcessor, +) + +_FakeTensor = sys.modules["infinicore"]._FakeTensor + + +class _FakeTargetEngine: + """伪 target 模型：next token 只由当前输入 token 决定。 + + succ(t) = (t+1) % vocab，traps 可覆盖个别 token 的后继（用来制造 + "prompt 撒谎"的 draft 拒绝场景，或让生成绕回 prompt 片段制造命中）。 + forward = 每请求最后输入位置采一个（对应 sample_all_positions=False）； + forward_raw = 全位置采样（verify 批次逐位置出 succ）。 + """ + + def __init__(self, traps=None, vocab=2000): + self.traps = traps or {} + self.vocab = vocab + # 记录每次 forward_raw 的输入 token 数：融合前向 = b×(k+1)， + # 两前向的 verify = 候选数×k，据此可断言走了哪条路径 + self.raw_calls = [] + + def succ(self, token): + if token in self.traps: + return self.traps[token] + return (token + 1) % self.vocab + + def forward(self, input_ids, **kwargs): + ids = input_ids.data[0] + offsets = kwargs["input_offsets"].data + outs = [ + self.succ(ids[offsets[i + 1] - 1]) for i in range(len(offsets) - 1) + ] + return _FakeTensor(outs) + + def forward_raw(self, input_ids, **kwargs): + ids = input_ids.data[0] + self.raw_calls.append(len(ids)) + return { + "output_ids": _FakeTensor([self.succ(t) for t in ids]), + "logits": None, + "hidden_states": None, + } + + +def _make_runner(fake_engine, num_draft_tokens=3, top_k=1, temperature=1.0): + config = types.SimpleNamespace( + speculative_method="prompt_lookup", + num_draft_tokens=num_draft_tokens, + max_batch_size=8, + top_k=top_k, + temperature=temperature, + ) + return SpeculativeRunner(config, fake_engine, device=None) + + +def _make_request(req_id, prompt_token_ids, max_tokens): + return InferenceRequest( + request_id=req_id, + prompt="x" * len(prompt_token_ids), + prompt_token_ids=list(prompt_token_ids), + sampling_params=SamplingParams(max_tokens=max_tokens), + eos_token_ids=[], + ) + + +def _make_engine(scheduler): + engine = LLMEngine.__new__(LLMEngine) + engine.scheduler = scheduler + engine.eos_token_ids = [] + + class _FakeTokenizer: + def decode(self, token_ids): + return "" + + engine.tokenizer = _FakeTokenizer() + return engine + + +def _make_scheduler(enable_chunked_prefill, budget=6): + return Scheduler( + max_batch_size=8, + num_blocks=64, + block_size=4, + max_num_batched_tokens=budget, + enable_prefix_caching=False, + enable_chunked_prefill=enable_chunked_prefill, + ) + + +def _run_one_step(engine, processor, runner): + """跑一个调度 step：真实 schedule + 真实输入构建 + 真实 runner + 真实 + _update_requests。返回 (scheduler_output, runner 输出)。""" + out = engine.scheduler.schedule() + if out is None or not out.scheduled_requests: + return None, None + model_input = processor._build_model_input_from_batch_scheduler_output( + out, 1.0, 0.8, 1 + ) + sampled = runner.forward(out, model_input) + engine._update_requests( + out.scheduled_requests, + sampled, + getattr(out, "num_scheduled_tokens", None), + ) + return out, sampled + + +def _run_until_finished(engine, processor, runner, requests, max_steps=500): + for _ in range(max_steps): + out, _ = _run_one_step(engine, processor, runner) + if out is None: + break + if all(r.is_finished() for r in requests): + break + + +def _truth(fake, last_prompt_token, n): + """非投机贪心路径的真值：从 prompt 尾 token 起逐次 succ。""" + out = [] + t = last_prompt_token + for _ in range(n): + t = fake.succ(t) + out.append(t) + return out + + +def test_lookup_helper(): + runner = _make_runner(_FakeTargetEngine()) + # 默认 env：max_ngram=4, min_ngram=2 + req = _make_request("u", [10, 11, 12, 13, 10, 11, 12, 55], max_tokens=8) + + # 序列 [10,11,12,13,10,11,12,55]，当前后缀指向末尾 + req._generated_token_ids.extend([13, 10, 11]) + # context = [10,11,12,13,10,11,12,55,13,10,11]，后缀 [13,10,11]： + # i=3 处 context[3:6]=[13,10,11] 命中（最近一次），其后是 [12,55,...] + drafts = runner._draft_prompt_lookup_tokens(req, 2) + assert drafts == [12, 55], drafts + + # k 截断：只取 1 个 + drafts = runner._draft_prompt_lookup_tokens(req, 1) + assert drafts == [12], drafts + + # 未命中：后缀不存在于前文 + req2 = _make_request("u2", [1, 2, 3, 4], max_tokens=8) + req2._generated_token_ids.extend([99, 100]) + assert runner._draft_prompt_lookup_tokens(req2, 3) == [] + + # min_ngram 以下不匹配（默认 min=2，单 token 后缀不匹配） + req3 = _make_request("u3", [7, 8, 9], max_tokens=8) + req3._generated_token_ids.extend([8]) + # context=[7,8,9,8]，后缀 [9,8] 无命中；size=1 不达 min，返回 [] + assert runner._draft_prompt_lookup_tokens(req3, 3) == [] + + # 序列太短（不足 min_ngram+1）直接返回 [] + req4 = _make_request("u4", [5], max_tokens=8) + assert runner._draft_prompt_lookup_tokens(req4, 3) == [] + + +def test_decode_exact_with_acceptance(): + """高命中场景：prompt 内含 wrap 片段，wrap 后 draft 持续全接受。""" + fake = _FakeTargetEngine(vocab=64) + sched = _make_scheduler(enable_chunked_prefill=False) + engine = _make_engine(sched) + processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + runner = _make_runner(fake, num_draft_tokens=3) + + prompt = [60, 61, 62, 63, 0, 1, 2, 3] + req = _make_request("a", prompt, max_tokens=130) + sched.add_request(req) + _run_until_finished(engine, processor, runner, [req]) + + expected = _truth(fake, prompt[-1], 130) + assert list(req.generated_token_ids) == expected, ( + list(req.generated_token_ids)[:16], + expected[:16], + ) + stats = runner.get_acceptance_stats() + assert stats["drafted_tokens"] > 0 + assert stats["accepted_tokens"] > 0 + # 前 60 步 suffix 是新内容（无命中）；wrap 之后 suffix 都能在 + # prompt/已生成片段里命中，每步产出 k+1=4 个 token，全程平均应显著 > 1 + assert stats["accept_rate"] > 0.9, stats + assert stats["avg_tokens_per_step"] > 1.4, stats + # 纯 decode 批次全部走融合单前向（b=1 行 × (k+1)=4 token）， + # 不应出现两前向流程的 verify 调用（1×k=3 token） + assert fake.raw_calls, "decode 批次应产生融合前向调用" + assert all(n == 4 for n in fake.raw_calls), fake.raw_calls + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_rejection_and_rollback(): + """拒绝路径：draft 整体被拒绝 / 部分接受，输出仍与真值逐 token 相同。""" + # trap 让生成绕回 prompt 中的片段；prompt 在关键位置"撒谎" + fake = _FakeTargetEngine(traps={999: 20, 888: 30}, vocab=2000) + sched = _make_scheduler(enable_chunked_prefill=False) + engine = _make_engine(sched) + processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + runner = _make_runner(fake, num_draft_tokens=2) + + # r1: prompt=[20,21,22,77,999]。生成 20,21 后 suffix [20,21] 命中， + # draft=[22,77]；d1=22 接受，d2=77 被拒绝（真值 23），修正为 23。 + r1 = _make_request("r1", [20, 21, 22, 77, 999], max_tokens=7) + # r2: prompt=[30,31,55,888]。生成 30,31 后 suffix [30,31] 命中， + # draft=[55,888]；d1=55 与 target 32 不符，整体拒绝，只产 32。 + r2 = _make_request("r2", [30, 31, 55, 888], max_tokens=5) + sched.add_request(r1) + sched.add_request(r2) + _run_until_finished(engine, processor, runner, [r1, r2]) + + assert list(r1.generated_token_ids) == [20, 21, 22, 23, 24, 25, 26] + assert list(r2.generated_token_ids) == [30, 31, 32, 33, 34] + stats = runner.get_acceptance_stats() + assert stats["drafted_tokens"] >= 4, stats # 两条 draft 链都被验证过 + assert stats["accepted_tokens"] >= 1, stats # r1 的 d1 被接受 + assert stats["avg_tokens_per_step"] > 1.0, stats + # verify 槽位回滚 + 请求结束释放后无块泄漏 + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_spec_with_chunked_mixed_batch(): + """spec × chunked prefill 混排：中段 chunk 不投机、不产 token， + decode 请求同批次正常投机，完成 prefill 的请求也能立刻投机。""" + fake = _FakeTargetEngine(traps={9: 5}, vocab=2000) + sched = _make_scheduler(enable_chunked_prefill=True, budget=6) + engine = _make_engine(sched) + processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + runner = _make_runner(fake, num_draft_tokens=3) + + # A：长 prompt 被切块（10 > 预算 6），prompt 内含自重复片段 + req_a = _make_request("a", [100, 101, 102, 103, 100, 101, 102, 103, 100, 101], max_tokens=6) + # B：trap succ(9)=5 形成周期 5 循环，decode 期持续命中 + req_b = _make_request("b", [5, 6, 7, 8, 9, 5, 6, 7, 8, 9], max_tokens=12) + sched.add_request(req_a) + sched.add_request(req_b) + + # step 1: A 切 6 token 中段 chunk，B 无预算留在 waiting + out, sampled = _run_one_step(engine, processor, runner) + assert out.num_scheduled_tokens == [6] + assert sampled == [[]], sampled # 中段 chunk：runner 产出空，不投机 + assert req_a.get_num_generated_tokens() == 0 + assert req_a.num_computed_tokens == 6 + + # step 2: A 被重新入队排在 B 后面，B 先切 6 token 中段 chunk + out, sampled = _run_one_step(engine, processor, runner) + assert out.num_scheduled_tokens == [6], out.num_scheduled_tokens + assert sampled == [[]], sampled + assert req_b.get_num_generated_tokens() == 0 + assert req_b.num_computed_tokens == 6 + + # step 3: A 续 4 token 完成 prefill 并投机产出，B 只分到 2 token + # （预算被 A 占去 4），仍是中段 chunk + out, sampled = _run_one_step(engine, processor, runner) + assert out.num_scheduled_tokens == [4, 2], out.num_scheduled_tokens + assert sampled[0] == [102, 103, 104], sampled # A 完成块即投机（d3 被拒） + assert sampled[1] == [], sampled + assert req_b.get_num_generated_tokens() == 0 + + # step 4: A decode 占 1 预算（正常投机），B 续 2 token 完成 prefill 并投机 + out, sampled = _run_one_step(engine, processor, runner) + assert out.num_scheduled_tokens == [1, 2], out.num_scheduled_tokens + assert sampled[0] == [105], sampled + assert sampled[1] == [5, 6, 7, 8], sampled # B 完成块即投机（全接受） + + # 后续步骤自由推进（纯 decode 混排）直至全部完成 + _run_until_finished(engine, processor, runner, [req_a, req_b]) + + assert list(req_a.generated_token_ids) == [102, 103, 104, 105, 106, 107] + assert list(req_b.generated_token_ids) == [5, 6, 7, 8, 9, 5, 6, 7, 8, 9, 5, 6] + stats = runner.get_acceptance_stats() + assert stats["accepted_tokens"] > 0, stats + assert stats["avg_tokens_per_step"] > 1.2, stats + # 混排批次走两前向（verify：候选×k token），纯 decode 批次走融合 + # （b×(k+1)）：两种形状的 forward_raw 调用都应出现 + assert any(n in (3, 6) for n in fake.raw_calls), fake.raw_calls + assert any(n in (4, 8) for n in fake.raw_calls), fake.raw_calls + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_legacy_prefill_batch_spec(): + """旧调度（非 chunked）的整段 prefill 批次同样能投机。""" + fake = _FakeTargetEngine(traps={9: 5}, vocab=2000) + sched = _make_scheduler(enable_chunked_prefill=False) + engine = _make_engine(sched) + processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + runner = _make_runner(fake, num_draft_tokens=3) + + req = _make_request("a", [5, 6, 7, 8, 9, 5, 6, 7, 8, 9], max_tokens=8) + sched.add_request(req) + + # step 1: 整段 prefill，prompt 内自重复使 draft 全接受 + out, sampled = _run_one_step(engine, processor, runner) + assert out.is_prefill + assert sampled == [[5, 6, 7, 8]], sampled + + _run_until_finished(engine, processor, runner, [req]) + assert list(req.generated_token_ids) == [5, 6, 7, 8, 9, 5, 6, 7] + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_nongreedy_fallback(): + """非 greedy 配置：回退到常规前向，每请求产出单个 token（平铺列表）。""" + fake = _FakeTargetEngine(traps={9: 5}, vocab=2000) + sched = _make_scheduler(enable_chunked_prefill=False) + engine = _make_engine(sched) + processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + runner = _make_runner(fake, num_draft_tokens=3, temperature=0.7) + + req = _make_request("a", [5, 6, 7, 8, 9], max_tokens=4) + sched.add_request(req) + _run_until_finished(engine, processor, runner, [req]) + + assert list(req.generated_token_ids) == [5, 6, 7, 8] + stats = runner.get_acceptance_stats() + assert stats["drafted_tokens"] == 0 # 从未进入投机路径 + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +def test_adaptive_gate(): + """低收益场景自适应回退：无 n-gram 命中时平均每步只产 1 token， + 攒满窗口（默认 32 请求步）后触发门控，冷静期（默认 64 步）内走 + 常规前向、不再有融合 verify 调用，输出仍与真值逐 token 相同。""" + fake = _FakeTargetEngine(vocab=2000) + sched = _make_scheduler(enable_chunked_prefill=False) + engine = _make_engine(sched) + processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + runner = _make_runner(fake, num_draft_tokens=3) + + # 严格递增序列（vocab 足够大不 wrap）：任何后缀在前文中都不重复， + # draft 永不命中 → 每步只产 1 token + req = _make_request("g", [1000, 1001], max_tokens=80) + sched.add_request(req) + _run_until_finished(engine, processor, runner, [req]) + + assert list(req.generated_token_ids) == _truth(fake, 1001, 80) + stats = runner.get_acceptance_stats() + assert stats["gate_triggered"] == 1, stats + # 第 32 步触发门控，之后 48 个 token 全部走常规前向。融合前向 + # （forward_raw）停留在 31 次：第 1 步是整段 prefill 批次，走两前向 + # 流程且无 draft 命中、不产生 verify 调用；第 2..32 步才是融合前向 + assert len(fake.raw_calls) == 31, len(fake.raw_calls) + assert sched.cache_manager.get_total_usable_blocks() == 64 + + +if __name__ == "__main__": + test_lookup_helper() + print("PASS test_lookup_helper") + test_decode_exact_with_acceptance() + print("PASS test_decode_exact_with_acceptance") + test_rejection_and_rollback() + print("PASS test_rejection_and_rollback") + test_spec_with_chunked_mixed_batch() + print("PASS test_spec_with_chunked_mixed_batch") + test_legacy_prefill_batch_spec() + print("PASS test_legacy_prefill_batch_spec") + test_nongreedy_fallback() + print("PASS test_nongreedy_fallback") + test_adaptive_gate() + print("PASS test_adaptive_gate") + print("ALL OK")