diff --git a/CMakeLists.txt b/CMakeLists.txt index 227bd3e8..5fd53e10 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2982,6 +2982,8 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST qwen_decoder_packed_projection_test tests/unittests/test_qwen_decoder_packed_projections.cpp ) + add_engine_unittest(qwen_chunked_prefill_test tests/unittests/test_qwen_chunked_prefill.cpp) + add_test(NAME qwen_chunked_prefill_test COMMAND qwen_chunked_prefill_test) add_test( NAME qwen_decoder_packed_projection_test COMMAND qwen_decoder_packed_projection_test diff --git a/include/engine/framework/modules/optimizations/fast_kv_modules.h b/include/engine/framework/modules/optimizations/fast_kv_modules.h index 247a8cc2..7ee1185d 100644 --- a/include/engine/framework/modules/optimizations/fast_kv_modules.h +++ b/include/engine/framework/modules/optimizations/fast_kv_modules.h @@ -25,6 +25,13 @@ class FastKVSetRowsModule { const core::TensorValue & row, const core::TensorValue & row_index) const; + // Explicit single-sequence, multi-token cache update. + core::TensorValue build_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & cache, + const core::TensorValue & rows, + const core::TensorValue & indices) const; + static const core::ModuleSchema & static_schema() noexcept; private: diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index 15063d25..37c2f04f 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -69,6 +69,11 @@ class QwenCausalDecodeRuntime { QwenCausalPrefillResult prefill_tokens(const std::vector & token_ids); QwenCausalPrefillResult prefill_embeddings(const std::vector & embeddings, int64_t steps); + // Prefill bounded blocks directly into the token-decode cache on the backend. + // No host KV export/import; subsequent decode_token calls continue this state. + QwenCausalDecodeStepResult prefill_embeddings_into_cache( + const std::vector & embeddings, int64_t steps, int64_t cache_steps, int64_t chunk_steps); + QwenCausalBatchedPrefillResult prefill_tokens_batched( const std::vector & token_ids, int64_t batch_size, diff --git a/include/engine/framework/modules/transformers/qwen_decoder.h b/include/engine/framework/modules/transformers/qwen_decoder.h index fa1c5ac9..0560371e 100644 --- a/include/engine/framework/modules/transformers/qwen_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_decoder.h @@ -167,6 +167,17 @@ class QwenDecoderLayerModule { const std::optional & cache_slot, const core::TensorValue & attention_mask) const; + QwenDecoderLayerOutputs build_with_static_cache_block( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const std::optional & cache_slot, + const core::TensorValue & attention_mask) const; + QwenDecoderLayerOutputs build_with_static_cache_tail_batched( core::ModuleBuildContext & ctx, ggml_cgraph * graph, @@ -181,6 +192,17 @@ class QwenDecoderLayerModule { static const core::ModuleSchema & static_schema() noexcept; private: + QwenDecoderLayerOutputs build_static_cache_impl( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const std::optional & cache_slot, + const core::TensorValue & attention_mask, + bool block) const; QwenDecoderLayerConfig config_; }; diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index 55c43711..f7775916 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/include/engine/framework/runtime/kv_cache.h @@ -55,6 +55,7 @@ class TransformerKVCache { TransformerKVCacheOptions options); void import_state(const TransformerKVState & state); + void clear_on_backend(); TransformerKVState export_state() const; void advance_after_direct_append(int64_t steps); diff --git a/src/framework/modules/optimizations/fast_kv_modules.cpp b/src/framework/modules/optimizations/fast_kv_modules.cpp index aa879b0e..734490c7 100644 --- a/src/framework/modules/optimizations/fast_kv_modules.cpp +++ b/src/framework/modules/optimizations/fast_kv_modules.cpp @@ -119,6 +119,38 @@ core::TensorValue FastKVSetRowsModule::build( return core::reshape_tensor(ctx, flat_updated, cache.shape); } + +core::TensorValue FastKVSetRowsModule::build_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & cache, + const core::TensorValue & rows, + const core::TensorValue & indices) const { + core::validate_rank_between(cache, 4, 4, "cache"); + core::validate_rank_between(rows, 4, 4, "rows"); + const int64_t queries = rows.shape.dims[1]; + core::validate_shape(rows, core::TensorShape::from_dims( + {1, queries, cache.shape.dims[2], cache.shape.dims[3]}), "rows"); + core::validate_shape(indices, core::TensorShape::from_dims({queries}), "indices"); + if (ctx.ggml == nullptr || cache.shape.dims[0] != 1 || queries <= 0 || + queries > cache.shape.dims[1] || rows.type != GGML_TYPE_F32 || + (indices.type != GGML_TYPE_I32 && indices.type != GGML_TYPE_I64)) { + throw std::runtime_error("FastKVSetRowsModule block requires one sequence, f32 rows and integer indices"); + } + const bool optimized = config_.mode == FastKVSetRowsMode::BackendViewOptimized; + if ((!optimized && cache.type != GGML_TYPE_F32) || + (optimized && cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16 && cache.type != GGML_TYPE_BF16) || + !core::has_backend_addressable_layout(cache.tensor)) { + throw std::runtime_error("FastKVSetRowsModule block cache type or layout is unsupported"); + } + const int64_t width = cache.shape.dims[2] * cache.shape.dims[3]; + auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({cache.shape.dims[1], width})); + auto contiguous_rows = tensor_layout::ensure_contiguous_layout_if_needed(ctx, rows); + auto flat_rows = core::reshape_tensor(ctx, contiguous_rows, core::TensorShape::from_dims({queries, width})); + auto * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_rows.tensor, indices.tensor); + if (optimized) { updated->src[2] = cache.tensor; } + return core::reshape_tensor(ctx, core::wrap_tensor(updated, flat_cache.shape, cache.type), cache.shape); +} + const core::ModuleSchema & FastKVSetRowsModule::static_schema() noexcept { return kFastKVSetRowsSchema; } diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 9df457ae..811985a6 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -526,6 +526,64 @@ class QwenCausalDecodeRuntime::Impl { return run_batched_prefill(); } + QwenCausalDecodeStepResult prefill_embeddings_into_cache( + const std::vector & embeddings, int64_t steps, int64_t cache_steps, int64_t chunk_steps) { + if (steps <= 0 || chunk_steps <= 0 || cache_steps < steps || + embeddings.size() != static_cast(steps * config_.decoder.stack.hidden_size)) { + throw std::runtime_error("Qwen chunked prefill has invalid dimensions"); + } + if (config_.decoder.stack.runtime.static_cache.update_mode != QwenDecoderStaticCacheUpdateMode::DirectSetRows || + config_.decoder.logits_mode != QwenCausalDecoderLogitsMode::LastStep || + !config_.logits_readback_token_ids.empty()) { + throw std::runtime_error("Qwen chunked prefill requires DirectSetRows and full last-token readback"); + } + const int64_t chunk = std::min(steps, chunk_steps); + // Padding writes only unused future slots; those slots stay masked until overwritten. + const int64_t capacity = std::max(cache_steps, ((steps + chunk - 1) / chunk) * chunk); + ensure_decode_token_graph(capacity); + if (block_steps_ != chunk) { + release_block_graph(); + build_block_graph(chunk); + } + decode_cache_.clear_on_backend(); + const int64_t width = config_.decoder.stack.hidden_size; + std::vector input(static_cast(chunk * width)); + std::vector mask(static_cast(chunk * decode_cache_steps_)); + const auto masked = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); + for (int64_t offset = 0; offset < steps; offset += chunk) { + const int64_t count = std::min(chunk, steps - offset); + std::fill(input.begin(), input.end(), 0.f); + std::copy_n(embeddings.data() + offset * width, count * width, input.data()); + auto positions = qwen_position_ids(chunk, offset); + std::fill(mask.begin(), mask.end(), masked); + for (int64_t q = 0; q < chunk; ++q) { + const int64_t begin = config_.sliding_window > 0 + ? std::max(0, offset + q - config_.sliding_window + 1) : 0; + std::fill(mask.begin() + q * decode_cache_steps_ + begin, + mask.begin() + q * decode_cache_steps_ + offset + q + 1, + ggml_fp32_to_fp16(0.f)); + } + const int32_t last = static_cast(count - 1); + ggml_backend_tensor_set(block_input_, input.data(), 0, input.size() * sizeof(float)); + ggml_backend_tensor_set(block_positions_, positions.data(), 0, positions.size() * sizeof(int32_t)); + ggml_backend_tensor_set(block_mask_, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(block_last_, &last, 0, sizeof(last)); + core::set_backend_threads(backend_, threads_); + if (core::compute_backend_graph(backend_, block_graph_) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Qwen chunked prefill compute failed"); + } + decode_cache_.advance_after_direct_append(count); + } + ggml_backend_synchronize(backend_); + QwenCausalDecodeStepResult result; + if (block_logits_) { result.logits = core::read_tensor_f32(block_logits_); } + if (block_hidden_) { + result.hidden = core::read_tensor_f32(block_hidden_); + round_readback(result.hidden, config_); + } + return result; + } + void start_decode_tokens(const runtime::TransformerKVState & state, int64_t required_cache_steps) { if (required_cache_steps <= 0) { throw std::runtime_error("QwenCausalDecodeRuntime decode requires positive cache capacity"); @@ -1407,7 +1465,71 @@ class QwenCausalDecodeRuntime::Impl { batched_prefill_input_kind_ = InputKind::None; } + void build_block_graph(int64_t chunk) { + block_ctx_.reset(ggml_init({config_.prefill_graph_arena_bytes, nullptr, true})); + if (!block_ctx_) { throw std::runtime_error("Qwen chunked prefill context allocation failed"); } + core::ModuleBuildContext ctx{block_ctx_.get(), config_.trace_name.c_str(), backend_type_}; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({1, chunk, config_.decoder.stack.hidden_size})); + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({chunk})); + auto mask = core::make_tensor(ctx, GGML_TYPE_F16, + core::TensorShape::from_dims({1, 1, chunk, decode_cache_steps_})); + block_input_ = input.tensor; + block_positions_ = positions.tensor; + block_mask_ = mask.tensor; + block_last_ = ggml_new_tensor_1d(ctx.ggml, GGML_TYPE_I32, 1); + for (auto * tensor : {block_input_, block_positions_, block_mask_, block_last_}) { ggml_set_input(tensor); } + block_graph_ = ggml_new_graph_custom(ctx.ggml, 65536, false); + const QwenDecoderLayerModule layer(qwen_decoder_layer_config_from_stack(config_.decoder.stack)); + auto hidden = input; + for (size_t i = 0; i < weights_.stack.layers.size(); ++i) { + hidden = layer.build_with_static_cache_block(ctx, block_graph_, hidden, positions, + weights_.stack.layers[i], decode_cache_.key_tensor(i), decode_cache_.value_tensor(i), + positions, mask).output; + } + hidden = core::reshape_tensor(ctx, hidden, + core::TensorShape::from_dims({chunk, config_.decoder.stack.hidden_size})); + hidden = core::wrap_tensor(ggml_get_rows(ctx.ggml, hidden.tensor, block_last_), + core::TensorShape::from_dims({1, 1, config_.decoder.stack.hidden_size}), GGML_TYPE_F32); + hidden = RMSNormModule({config_.decoder.stack.hidden_size, config_.decoder.stack.rms_norm_eps, true, false}) + .build(ctx, hidden, weights_.final_norm); + if (config_.return_hidden || config_.output_mode == QwenCausalDecodeOutputMode::Hidden) { + block_hidden_ = hidden.tensor; + ggml_set_output(block_hidden_); + ggml_build_forward_expand(block_graph_, block_hidden_); + } + if (config_.output_mode == QwenCausalDecodeOutputMode::Logits) { + if (config_.decoder.lm_head_input_type) { + hidden = core::wrap_tensor(ggml_cast(ctx.ggml, hidden.tensor, *config_.decoder.lm_head_input_type), + hidden.shape, *config_.decoder.lm_head_input_type); + } + block_logits_ = LinearModule({config_.decoder.stack.hidden_size, config_.decoder.logits_size, + config_.decoder.use_lm_head_bias, config_.decoder.lm_head_precision}) + .build(ctx, hidden, *weights_.lm_head).tensor; + ggml_set_output(block_logits_); + ggml_build_forward_expand(block_graph_, block_logits_); + } + block_allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (!ggml_gallocr_alloc_graph(block_allocator_, block_graph_)) { + throw std::runtime_error("Qwen chunked prefill graph allocation failed"); + } + block_steps_ = chunk; + } + + void release_block_graph() { + if (block_graph_) { + core::release_backend_graph_resources(backend_, block_graph_, config_.evict_cuda_graph_cache_on_release); + } + ggml_gallocr_free(block_allocator_); + block_allocator_ = nullptr; + block_ctx_.reset(); + block_graph_ = nullptr; + block_logits_ = block_hidden_ = nullptr; + block_steps_ = 0; + } + void release_decode_graph() { + release_block_graph(); if (decode_graph_ != nullptr) { core::release_backend_graph_resources( backend_, decode_graph_, config_.evict_cuda_graph_cache_on_release); @@ -1500,6 +1622,16 @@ class QwenCausalDecodeRuntime::Impl { int64_t batched_prefill_steps_ = 0; InputKind batched_prefill_input_kind_ = InputKind::None; + std::unique_ptr block_ctx_; + ggml_tensor * block_input_ = nullptr; + ggml_tensor * block_positions_ = nullptr; + ggml_tensor * block_mask_ = nullptr; + ggml_tensor * block_last_ = nullptr; + ggml_tensor * block_logits_ = nullptr; + ggml_tensor * block_hidden_ = nullptr; + ggml_cgraph * block_graph_ = nullptr; + ggml_gallocr_t block_allocator_ = nullptr; + int64_t block_steps_ = 0; std::unique_ptr decode_ctx_; ggml_tensor * decode_input_ = nullptr; ggml_tensor * decode_positions_ = nullptr; @@ -1565,6 +1697,11 @@ QwenCausalBatchedPrefillResult QwenCausalDecodeRuntime::prefill_embeddings_batch return impl_->prefill_embeddings_batched(embeddings, batch_size, steps); } +QwenCausalDecodeStepResult QwenCausalDecodeRuntime::prefill_embeddings_into_cache( + const std::vector & embeddings, int64_t steps, int64_t cache_steps, int64_t chunk_steps) { + return impl_->prefill_embeddings_into_cache(embeddings, steps, cache_steps, chunk_steps); +} + void QwenCausalDecodeRuntime::start_decode_tokens( const runtime::TransformerKVState & state, int64_t required_cache_steps) { diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index 743f5aaa..aa4b60c2 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -680,7 +680,40 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( const core::TensorValue & cache_value, const std::optional & cache_slot, const core::TensorValue & attention_mask) const { + return build_static_cache_impl(ctx, graph, input, positions, weights, cache_key, cache_value, + cache_slot, attention_mask, false); +} + +QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_block( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const std::optional & cache_slot, + const core::TensorValue & attention_mask) const { + return build_static_cache_impl(ctx, graph, input, positions, weights, cache_key, cache_value, + cache_slot, attention_mask, true); +} + +QwenDecoderLayerOutputs QwenDecoderLayerModule::build_static_cache_impl( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const std::optional & cache_slot, + const core::TensorValue & attention_mask, + bool block) const { validate_sequence_input(input, config_.hidden_size, "input"); + if (block && (input.shape.dims[0] != 1 || + config_.runtime.static_cache.update_mode != QwenDecoderStaticCacheUpdateMode::DirectSetRows)) { + throw std::runtime_error("Qwen static-cache blocks require a single sequence and DirectSetRows"); + } const int64_t dim = require_head_dim(config_); const int64_t kv_repeats = config_.num_attention_heads / config_.num_key_value_heads; @@ -735,8 +768,10 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( ? FastKVSetRowsMode::BackendViewOptimized : FastKVSetRowsMode::Exact, }); - attention_key_cache = set_rows.build(ctx, cache_key, k, *cache_slot); - attention_value_cache = set_rows.build(ctx, cache_value, v, *cache_slot); + attention_key_cache = block ? set_rows.build_block(ctx, cache_key, k, *cache_slot) + : set_rows.build(ctx, cache_key, k, *cache_slot); + attention_value_cache = block ? set_rows.build_block(ctx, cache_value, v, *cache_slot) + : set_rows.build(ctx, cache_value, v, *cache_slot); if (config_.activation_cast.enabled && config_.activation_cast.after_static_cache_update) { attention_key_cache = activation_cast(ctx, attention_key_cache, config_.activation_cast); attention_value_cache = activation_cast(ctx, attention_value_cache, config_.activation_cast); @@ -806,7 +841,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( context = core::reshape_tensor( ctx, context, - core::TensorShape::from_dims({1, 1, config_.num_attention_heads * dim})); + core::TensorShape::from_dims({1, block ? input.shape.dims[1] : 1, config_.num_attention_heads * dim})); auto attn_out = LinearModule( { diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 05b7b7ea..4a8450a0 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -122,6 +122,15 @@ TransformerKVCache::TransformerKVCache( } } +void TransformerKVCache::clear_on_backend() { + for (auto & layer : layers_) { + ggml_backend_tensor_memset(layer.key_tensor.tensor, 0, 0, ggml_nbytes(layer.key_tensor.tensor)); + ggml_backend_tensor_memset(layer.value_tensor.tensor, 0, 0, ggml_nbytes(layer.value_tensor.tensor)); + } + current_end_ = 0; + valid_steps_ = 0; +} + void TransformerKVCache::import_state(const TransformerKVState & state) { current_end_ = state.current_end; if (layers_.empty()) { diff --git a/tests/unittests/test_qwen_chunked_prefill.cpp b/tests/unittests/test_qwen_chunked_prefill.cpp new file mode 100644 index 00000000..18e35c5f --- /dev/null +++ b/tests/unittests/test_qwen_chunked_prefill.cpp @@ -0,0 +1,166 @@ +#include "engine/framework/modules/transformers/qwen_causal_decode_runtime.h" +#include "engine/framework/debug/trace.h" + +#include +#include +#include +#include +#include + +namespace { +std::vector pattern(size_t size, float scale) { + std::vector result(size); + for (size_t i = 0; i < size; ++i) { result[i] = scale * std::sin(static_cast(i) * .17f + .3f); } + return result; +} +void close(const std::vector & actual, const std::vector & expected) { + if (actual.size() != expected.size()) { throw std::runtime_error("logit size mismatch"); } + float maximum = 0; + for (size_t i = 0; i < actual.size(); ++i) { + if (!std::isfinite(actual[i]) || !std::isfinite(expected[i])) { throw std::runtime_error("nonfinite logits"); } + maximum = std::max(maximum, std::abs(actual[i] - expected[i])); + } + engine::debug::trace_log_scalar("qwen_chunked_prefill_test.max_error", maximum); + if (maximum > .003f) { throw std::runtime_error("chunked prefill logits differ from reference"); } +} +} + +int main(int argc, char ** argv) { + try { + engine::core::BackendConfig backend; + backend.threads = 8; + engine::debug::LoggingConfig logging; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--backend" && i + 1 < argc) { + const std::string value(argv[++i]); + if (value == "cuda") { backend.type = engine::core::BackendType::Cuda; } + else if (value == "cpu") { backend.type = engine::core::BackendType::Cpu; } + else { throw std::runtime_error("unsupported test backend"); } + } else if (arg == "--log" && i + 1 < argc) { + logging.enabled = true; + logging.file_path = argv[++i]; + } else { throw std::runtime_error("unknown test argument"); } + } + engine::debug::configure_logging(logging); + engine::core::ExecutionContext execution(backend); + auto * context = ggml_init({2 * 1024 * 1024, nullptr, true}); + if (!context) { throw std::runtime_error("weight context allocation failed"); } + struct ContextGuard { ggml_context * p; ~ContextGuard() { ggml_free(p); } } context_guard{context}; + engine::core::ModuleBuildContext ctx{context, "qwen_chunked_prefill_test", backend.type}; + std::vector tensors; + auto tensor = [&](std::initializer_list shape) { + auto value = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims(shape)); + tensors.push_back(value); + return value; + }; + engine::modules::QwenCausalDecodeRuntimeWeights weights; + weights.token_embedding = tensor({97, 64}); + weights.final_norm = {tensor({64}), std::nullopt}; + weights.lm_head = engine::modules::LinearWeights{weights.token_embedding, std::nullopt}; + for (int layer = 0; layer < 2; ++layer) { + engine::modules::QwenDecoderLayerWeights w; + w.input_norm = {tensor({64}), std::nullopt}; + w.post_norm = {tensor({64}), std::nullopt}; + w.q_norm = {tensor({64}), std::nullopt}; + w.k_norm = {tensor({64}), std::nullopt}; + w.self_attention.q_weight = tensor({128, 64}); + w.self_attention.k_weight = tensor({64, 64}); + w.self_attention.v_weight = tensor({64, 64}); + w.self_attention.out_weight = tensor({64, 128}); + w.mlp.gate_proj = {tensor({128, 64}), std::nullopt}; + w.mlp.up_proj = {tensor({128, 64}), std::nullopt}; + w.mlp.down_proj = {tensor({64, 128}), std::nullopt}; + weights.stack.layers.push_back(w); + } + std::array, 3> import_tensors; + const std::array import_types{GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}; + for (size_t type = 0; type < import_types.size(); ++type) { + for (auto & value : import_tensors[type]) { + value = engine::core::make_tensor(ctx, import_types[type], engine::core::TensorShape::from_dims({1, 9, 1, 4})); + } + } + auto * buffer = ggml_backend_alloc_ctx_tensors(context, execution.backend()); + if (!buffer) { throw std::runtime_error("weight buffer allocation failed"); } + struct BufferGuard { ggml_backend_buffer_t p; ~BufferGuard() { ggml_backend_buffer_free(p); } } buffer_guard{buffer}; + for (const auto & value : tensors) { + auto values = pattern(ggml_nelements(value.tensor), .04f); + if (value.shape.rank == 1) { std::fill(values.begin(), values.end(), 1.f); } + engine::core::write_tensor_f32(value, values); + } + for (auto & values : import_tensors) { + engine::runtime::TransformerKVCacheOptions options; + options.allow_f16_storage = true; + options.allow_bf16_storage = true; + engine::runtime::TransformerKVCache legacy(9, 4, {values[0]}, {values[1]}, options); + engine::runtime::TransformerKVCache device_zero(9, 4, {values[2]}, {values[3]}, options); + for (int steps : {5, 2, 0}) { + engine::runtime::TransformerKVState state; + state.current_end = steps; + state.layers.resize(1); + state.layers[0].valid_steps = steps; + state.layers[0].key = pattern(steps * 4, .4f); + state.layers[0].value = pattern(steps * 4, -.7f); + legacy.import_state(state); + device_zero.import_state(state); + device_zero.clear_on_backend(); + for (size_t kind = 0; kind < 2; ++kind) { + const auto bytes = ggml_nbytes(values[kind].tensor); + std::vector expected(bytes), actual(bytes); + if (device_zero.current_end() != 0 || device_zero.valid_steps() != 0) { + throw std::runtime_error("cache clear did not reset positions"); + } + ggml_backend_tensor_get(values[kind + 2].tensor, actual.data(), 0, bytes); + if (actual != expected) { throw std::runtime_error("cache clear left nonzero bytes"); } + } + } + } + engine::modules::QwenCausalDecodeRuntimeConfig config; + auto & stack = config.decoder.stack; + stack.hidden_size = 64; + stack.num_attention_heads = 2; + stack.num_key_value_heads = 1; + stack.head_dim = 64; + stack.intermediate_size = 128; + stack.layers = 2; + stack.runtime.static_cache.update_mode = engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + stack.runtime.static_cache.set_rows_mode = engine::modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + stack.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + stack.runtime.attention.static_mode = engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + config.decoder.logits_size = 97; + config.decoder.static_cache_type = GGML_TYPE_F16; + config.prefill_graph_arena_bytes = 4 * 1024 * 1024; + config.decode_graph_arena_bytes = 4 * 1024 * 1024; + config.evict_cuda_graph_cache_on_release = true; + engine::modules::QwenCausalDecodeRuntime reference(execution, config, weights); + engine::modules::QwenCausalDecodeRuntime saver(execution, config, weights); + for (int steps : {11, 5, 1, 11, 28}) { + const auto embeddings = pattern(steps * 64, .3f); + auto expected = reference.prefill_embeddings(embeddings, steps); + auto actual = saver.prefill_embeddings_into_cache(embeddings, steps, 32, 4); + close(actual.logits, expected.logits); + reference.start_decode_tokens(expected.state, 32); + for (int token = 0; token < 4; ++token) { + const auto reference_logits = reference.decode_token(token + 7).logits; + close(saver.decode_token(token + 7).logits, reference_logits); + } + if (saver.decode_current_end() != steps + 4) { throw std::runtime_error("cache position mismatch"); } + } + bool rejected = false; + try { saver.decode_token(7); } catch (const std::runtime_error &) { rejected = true; } + if (!rejected) { throw std::runtime_error("context overflow accepted"); } + const auto grown_embeddings = pattern(41 * 64, .3f); + auto grown_reference = reference.prefill_embeddings(grown_embeddings, 41); + close(saver.prefill_embeddings_into_cache(grown_embeddings, 41, 64, 8).logits, + grown_reference.logits); + reference.start_decode_tokens(grown_reference.state, 64); + close(saver.decode_token(7).logits, reference.decode_token(7).logits); + saver.release_runtime_graphs(); + if (saver.decode_current_end() != 0) { throw std::runtime_error("reset failed"); } + std::cout << "PASS chunk boundaries, repeated prefill, decode, reset, capacity, reference parity\n"; + return 0; + } catch (const std::exception & error) { + std::cerr << error.what() << '\n'; + return 1; + } +}