Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
327 lines
20 KiB
C++
327 lines
20 KiB
C++
#include "engine/framework/modules/transformers/causal_decoder_runtime.h"
|
|
#include "engine/framework/debug/profiler.h"
|
|
|
|
#include <cmath>
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <iostream>
|
|
#include <stdexcept>
|
|
|
|
namespace {
|
|
std::vector<float> pattern(size_t size, float scale) {
|
|
std::vector<float> result(size);
|
|
for (size_t i = 0; i < size; ++i) { result[i] = scale * std::sin(static_cast<float>(i) * .17f + .3f); }
|
|
return result;
|
|
}
|
|
void close(const std::vector<float> & actual, const std::vector<float> & 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) {
|
|
const auto started = std::chrono::steady_clock::now();
|
|
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<engine::core::TensorValue> tensors;
|
|
auto tensor = [&](std::initializer_list<int64_t> 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::CausalDecoderRuntimeWeights 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::DecoderLayerWeights 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<std::array<engine::core::TensorValue, 6>, 3> import_tensors;
|
|
const std::array<ggml_type, 3> 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);
|
|
options.lazy_import_scratch = true;
|
|
engine::runtime::TransformerKVCache lazy(9, 4, {values[4]}, {values[5]}, options);
|
|
for (int steps : {0, 5, 2, 0, 9}) {
|
|
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);
|
|
lazy.import_state(state);
|
|
if (steps > 0) {
|
|
ggml_tensor key = *values[0].tensor;
|
|
ggml_tensor value = *values[1].tensor;
|
|
key.ne[2] = value.ne[2] = steps;
|
|
key.nb[3] = key.nb[2] * steps;
|
|
value.nb[3] = value.nb[2] * steps;
|
|
const auto shape = engine::core::TensorShape::from_dims({1, steps, 1, 4});
|
|
device_zero.import_prefix(
|
|
{engine::core::wrap_tensor(&key, shape, values[0].type)},
|
|
{engine::core::wrap_tensor(&value, shape, values[1].type)}, steps);
|
|
for (size_t kind = 0; kind < 2; ++kind) {
|
|
const auto bytes = ggml_nbytes(values[kind].tensor);
|
|
std::vector<unsigned char> expected(bytes), actual(bytes);
|
|
ggml_backend_tensor_get(values[kind].tensor, expected.data(), 0, bytes);
|
|
ggml_backend_tensor_get(values[kind + 2].tensor, actual.data(), 0, bytes);
|
|
if (actual != expected) throw std::runtime_error("backend prefix import changed values or tail");
|
|
}
|
|
if (device_zero.current_end() != steps || device_zero.valid_steps() != steps) {
|
|
throw std::runtime_error("backend prefix import changed positions");
|
|
}
|
|
}
|
|
if (lazy.current_end() != legacy.current_end() || lazy.valid_steps() != legacy.valid_steps()) {
|
|
throw std::runtime_error("lazy cache import changed positions");
|
|
}
|
|
for (size_t kind = 0; kind < 2; ++kind) {
|
|
const auto bytes = ggml_nbytes(values[kind].tensor);
|
|
std::vector<unsigned char> expected(bytes), actual(bytes);
|
|
ggml_backend_tensor_get(values[kind].tensor, expected.data(), 0, bytes);
|
|
ggml_backend_tensor_get(values[kind + 4].tensor, actual.data(), 0, bytes);
|
|
if (actual != expected) { throw std::runtime_error("lazy cache import changed values or tail"); }
|
|
}
|
|
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<unsigned char> 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::CausalDecoderRuntimeConfig 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::DecoderStaticCacheUpdateMode::DirectSetRows;
|
|
stack.runtime.static_cache.set_rows_mode = engine::modules::DecoderStaticCacheSetRowsMode::BackendViewOptimized;
|
|
stack.runtime.attention.prefill_mode = engine::modules::DecoderAttentionMode::FlashGroupedViewKV;
|
|
stack.runtime.attention.static_mode = engine::modules::DecoderAttentionMode::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::CausalDecoderRuntime reference(execution, config, weights);
|
|
engine::modules::CausalDecoderRuntime saver(execution, config, weights);
|
|
engine::modules::CausalDecoderRuntime direct(execution, config, weights);
|
|
for (const int steps : {11, 5, 1, 11, 41, 5}) {
|
|
const auto embeddings = pattern(steps * 64, .3f);
|
|
const auto expected = reference.prefill_embeddings(embeddings, steps);
|
|
const auto actual = direct.prefill_embeddings_into_decode_cache(embeddings, steps, 64);
|
|
close(actual.logits, expected.logits);
|
|
direct.release_prefill_graph();
|
|
reference.start_decode_embeddings(expected.state, 64);
|
|
for (int step = 0; step < 4; ++step) {
|
|
const auto next = pattern(64, .1f * (step + 1));
|
|
close(direct.decode_embedding(next).logits, reference.decode_embedding(next).logits);
|
|
}
|
|
if (direct.decode_current_end() != steps + 4) {
|
|
throw std::runtime_error("direct embedding prefill cache position mismatch");
|
|
}
|
|
// Switching back must not reuse a graph that writes the decode cache.
|
|
close(direct.prefill_embeddings(embeddings, steps).logits, expected.logits);
|
|
}
|
|
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);
|
|
// Kept prefix: a prompt sharing its first 24 rows with the previous
|
|
// one recomputes only the remainder, including after decode steps
|
|
// appended rows past the prefix, and across a block-size change.
|
|
for (const auto & [steps, keep, chunk] : std::array<std::array<int, 3>, 4>{{{45, 24, 8}, {50, 24, 8}, {48, 47, 16}, {49, 40, 5}}}) {
|
|
auto kept_embeddings = pattern(steps * 64, .3f);
|
|
std::fill(kept_embeddings.begin() + 24 * 64, kept_embeddings.end(), 0.f);
|
|
const auto tail = pattern((steps - 24) * 64, -.2f + .01f * steps);
|
|
std::copy(tail.begin(), tail.end(), kept_embeddings.begin() + 24 * 64);
|
|
auto kept_reference = reference.prefill_embeddings(kept_embeddings, steps);
|
|
if (saver.retainable_prefix_steps(steps, 64, chunk, keep) != std::min(keep, steps - 1)) {
|
|
throw std::runtime_error("resident prefix not retainable");
|
|
}
|
|
close(saver.prefill_embeddings_into_cache(kept_embeddings, steps, 64, chunk, keep).logits,
|
|
kept_reference.logits);
|
|
if (saver.decode_current_end() != steps) { throw std::runtime_error("kept prefix cache position mismatch"); }
|
|
reference.start_decode_tokens(kept_reference.state, 64);
|
|
for (int token = 0; token < 3; ++token) {
|
|
close(saver.decode_token(token + 9).logits, reference.decode_token(token + 9).logits);
|
|
}
|
|
}
|
|
// A prefix longer than the resident rows, or a capacity the current
|
|
// graph cannot hold, falls back to a full prefill.
|
|
if (saver.retainable_prefix_steps(60, 64, 8, 55) != 0) { throw std::runtime_error("unresident prefix retained"); }
|
|
if (saver.retainable_prefix_steps(40, 128, 8, 20) != 0) { throw std::runtime_error("prefix retained across capacity growth"); }
|
|
const auto regrown = pattern(40 * 64, .3f);
|
|
auto regrown_reference = reference.prefill_embeddings(regrown, 40);
|
|
close(saver.prefill_embeddings_into_cache(regrown, 40, 128, 8, 20).logits, regrown_reference.logits);
|
|
saver.release_runtime_graphs();
|
|
if (saver.decode_current_end() != 0) { throw std::runtime_error("reset failed"); }
|
|
// Unequal-length embedding prompts prefilled into separate lanes of the
|
|
// batched cache and then decoded together match two single-lane runtimes
|
|
// (the CFG pattern: both lanes receive the same embedding every step).
|
|
engine::modules::CausalDecoderRuntime lanes(execution, config, weights);
|
|
engine::modules::CausalDecoderRuntime lane_a(execution, config, weights);
|
|
engine::modules::CausalDecoderRuntime lane_b(execution, config, weights);
|
|
for (const auto & [steps_a, steps_b] : std::array<std::array<int, 2>, 4>{{{11, 5}, {3, 17}, {9, 9}, {11, 5}}}) {
|
|
const auto prompt_a = pattern(steps_a * 64, .3f);
|
|
const auto prompt_b = pattern(steps_b * 64, -.25f);
|
|
const auto prefilled = lanes.prefill_embeddings_into_batched_decode_cache({prompt_a, prompt_b}, 32);
|
|
close(prefilled[0].logits, lane_a.prefill_embeddings_into_decode_cache(prompt_a, steps_a, 32).logits);
|
|
close(prefilled[1].logits, lane_b.prefill_embeddings_into_decode_cache(prompt_b, steps_b, 32).logits);
|
|
for (int step = 0; step < 4; ++step) {
|
|
const auto next = pattern(64, .1f * (step + 1));
|
|
auto paired = next;
|
|
paired.insert(paired.end(), next.begin(), next.end());
|
|
const auto batched = lanes.decode_embeddings_batched(paired, 2);
|
|
const auto row = static_cast<std::ptrdiff_t>(batched.logits.size() / 2);
|
|
close({batched.logits.begin(), batched.logits.begin() + row}, lane_a.decode_embedding(next).logits);
|
|
close({batched.logits.begin() + row, batched.logits.end()}, lane_b.decode_embedding(next).logits);
|
|
}
|
|
}
|
|
// Token prompts take the same lane path (Yue2 uses it for CFG). Running them on
|
|
// the same runtime after the embedding lanes, and switching back afterwards,
|
|
// also covers changing the input kind between requests.
|
|
engine::modules::CausalDecoderRuntime token_a(execution, config, weights);
|
|
engine::modules::CausalDecoderRuntime token_b(execution, config, weights);
|
|
const auto token_prompt = [](int count, int offset) {
|
|
std::vector<int32_t> ids(static_cast<size_t>(count));
|
|
for (int i = 0; i < count; ++i) { ids[static_cast<size_t>(i)] = (i * 7 + offset) % 97; }
|
|
return ids;
|
|
};
|
|
for (const auto & [steps_a, steps_b] : std::array<std::array<int, 2>, 3>{{{7, 13}, {6, 6}, {7, 13}}}) {
|
|
const auto prompt_a = token_prompt(steps_a, 3);
|
|
const auto prompt_b = token_prompt(steps_b, 11);
|
|
const auto prefilled = lanes.prefill_tokens_into_batched_decode_cache({prompt_a, prompt_b}, 32);
|
|
close(prefilled[0].logits, token_a.prefill_tokens_into_decode_cache(prompt_a, 32).logits);
|
|
close(prefilled[1].logits, token_b.prefill_tokens_into_decode_cache(prompt_b, 32).logits);
|
|
for (int step = 0; step < 4; ++step) {
|
|
const int32_t next = (step * 13 + 5) % 97;
|
|
const auto batched = lanes.decode_tokens_batched({next, next});
|
|
const auto row = static_cast<std::ptrdiff_t>(batched.logits.size() / 2);
|
|
close({batched.logits.begin(), batched.logits.begin() + row}, token_a.decode_token(next).logits);
|
|
close({batched.logits.begin() + row, batched.logits.end()}, token_b.decode_token(next).logits);
|
|
}
|
|
}
|
|
{
|
|
const auto prompt_a = pattern(5 * 64, .2f);
|
|
const auto prompt_b = pattern(9 * 64, -.35f);
|
|
const auto prefilled = lanes.prefill_embeddings_into_batched_decode_cache({prompt_a, prompt_b}, 32);
|
|
close(prefilled[0].logits, lane_a.prefill_embeddings_into_decode_cache(prompt_a, 5, 32).logits);
|
|
close(prefilled[1].logits, lane_b.prefill_embeddings_into_decode_cache(prompt_b, 9, 32).logits);
|
|
const auto next = pattern(64, .45f);
|
|
auto paired = next;
|
|
paired.insert(paired.end(), next.begin(), next.end());
|
|
const auto batched = lanes.decode_embeddings_batched(paired, 2);
|
|
const auto row = static_cast<std::ptrdiff_t>(batched.logits.size() / 2);
|
|
close({batched.logits.begin(), batched.logits.begin() + row}, lane_a.decode_embedding(next).logits);
|
|
close({batched.logits.begin() + row, batched.logits.end()}, lane_b.decode_embedding(next).logits);
|
|
}
|
|
for (int64_t alignment : {int64_t{0}, int64_t{8}}) {
|
|
auto scratch_config = config;
|
|
scratch_config.reuse_batched_decode_scratch = true;
|
|
scratch_config.batched_attention_alignment = alignment;
|
|
engine::modules::CausalDecoderRuntime scratch(execution, scratch_config, weights);
|
|
for (int request = 0; request < 2; ++request) {
|
|
const std::vector<std::vector<int32_t>> prompts{
|
|
token_prompt(request == 0 ? 7 : 6, 3), token_prompt(request == 0 ? 13 : 6, 11)};
|
|
const auto expected = lanes.prefill_tokens_into_batched_decode_cache(prompts, 32);
|
|
const auto actual = scratch.prefill_tokens_into_batched_decode_cache(prompts, 32);
|
|
close(actual[0].logits, expected[0].logits);
|
|
close(actual[1].logits, expected[1].logits);
|
|
for (int step = 0; step < 18; ++step) {
|
|
const std::vector<int32_t> ids{step % 97, (step + 13) % 97};
|
|
close(scratch.decode_tokens_batched(ids).logits, lanes.decode_tokens_batched(ids).logits);
|
|
}
|
|
}
|
|
}
|
|
std::cout << "PASS chunk boundaries, repeated prefill, kept prefix, decode, reset, capacity, reference parity, "
|
|
"batched embedding and token lanes\n";
|
|
engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(started));
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(started));
|
|
std::cerr << error.what() << '\n';
|
|
return 1;
|
|
}
|
|
}
|