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
681 lines
30 KiB
C++
681 lines
30 KiB
C++
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/debug/profiler.h"
|
|
#include "engine/framework/modules/transformers/causal_decoder.h"
|
|
#include "engine/framework/modules/transformers/decoder.h"
|
|
#include "engine/framework/modules/optimizations/fast_kv_modules.h"
|
|
#include "engine/framework/modules/optimizations/fast_projection_modules.h"
|
|
|
|
#include <cmath>
|
|
#include <initializer_list>
|
|
#include <iostream>
|
|
#include <limits>
|
|
#include <optional>
|
|
#include <sstream>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
constexpr size_t kGraphBytes = 16 * 1024 * 1024;
|
|
constexpr size_t kGraphNodes = 4096;
|
|
|
|
std::vector<float> patterned(size_t count, float phase, float scale) {
|
|
std::vector<float> values(count);
|
|
for (size_t i = 0; i < count; ++i) {
|
|
const float x = static_cast<float>(i);
|
|
values[i] = scale * (std::sin(phase + 0.19f * x) + 0.35f * std::cos(phase + 0.07f * x));
|
|
}
|
|
return values;
|
|
}
|
|
|
|
void require_allclose(
|
|
const std::vector<float> & actual,
|
|
const std::vector<float> & expected,
|
|
float tolerance,
|
|
const std::string & label) {
|
|
if (actual.size() != expected.size()) {
|
|
throw std::runtime_error(label + " size mismatch");
|
|
}
|
|
for (size_t i = 0; i < actual.size(); ++i) {
|
|
if (!std::isfinite(actual[i]) || !std::isfinite(expected[i])) {
|
|
throw std::runtime_error(label + " contains non-finite values");
|
|
}
|
|
const float diff = std::fabs(actual[i] - expected[i]);
|
|
if (diff > tolerance) {
|
|
std::ostringstream message;
|
|
message << label << " mismatch at " << i << ": expected " << expected[i]
|
|
<< ", got " << actual[i] << ", diff=" << diff;
|
|
throw std::runtime_error(message.str());
|
|
}
|
|
}
|
|
}
|
|
|
|
struct LayerResult {
|
|
std::vector<float> output;
|
|
std::vector<float> key;
|
|
std::vector<float> value;
|
|
};
|
|
|
|
LayerResult run_layer(bool packed,
|
|
engine::core::BackendType backend_type = engine::core::BackendType::Cpu,
|
|
bool batched_decode = false,
|
|
bool fused_mlp = false) {
|
|
const int64_t batch = batched_decode ? 2 : 1;
|
|
const int64_t steps = batched_decode ? 1 : 3;
|
|
constexpr int64_t cache_steps = 8;
|
|
constexpr int64_t hidden = 8;
|
|
constexpr int64_t heads = 2;
|
|
constexpr int64_t kv_heads = 1;
|
|
constexpr int64_t head_dim = 4;
|
|
constexpr int64_t intermediate = 12;
|
|
constexpr int64_t q_out = heads * head_dim;
|
|
constexpr int64_t kv_out = kv_heads * head_dim;
|
|
|
|
engine::core::BackendConfig backend_config{backend_type, 0, 8};
|
|
ggml_backend_t backend = engine::core::init_backend(backend_config);
|
|
if (backend == nullptr) {
|
|
throw std::runtime_error("failed to initialize test backend");
|
|
}
|
|
|
|
ggml_init_params params{kGraphBytes, nullptr, true};
|
|
ggml_context * ggml = ggml_init(params);
|
|
if (ggml == nullptr) {
|
|
ggml_backend_free(backend);
|
|
throw std::runtime_error("failed to initialize GGML context");
|
|
}
|
|
|
|
ggml_backend_buffer_t buffer = nullptr;
|
|
try {
|
|
engine::core::ModuleBuildContext ctx{ggml, "qwen_packed_projection_test", backend_type};
|
|
ctx.elide_identity_reshapes = fused_mlp;
|
|
auto make_f32 = [&](std::initializer_list<int64_t> dims) {
|
|
return engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims(dims));
|
|
};
|
|
|
|
auto input = make_f32({batch, steps, hidden});
|
|
auto positions = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_I32,
|
|
engine::core::TensorShape::from_dims({batched_decode ? batch : steps}));
|
|
|
|
engine::modules::DecoderLayerWeights weights;
|
|
weights.input_norm = {make_f32({hidden}), std::nullopt};
|
|
weights.post_norm = {make_f32({hidden}), std::nullopt};
|
|
weights.self_attention.out_weight = make_f32({hidden, hidden});
|
|
weights.mlp.down_proj = {make_f32({hidden, intermediate}), std::nullopt};
|
|
|
|
const auto q_values = patterned(static_cast<size_t>(q_out * hidden), 0.1f, 0.12f);
|
|
const auto k_values = patterned(static_cast<size_t>(kv_out * hidden), 0.5f, 0.10f);
|
|
const auto v_values = patterned(static_cast<size_t>(kv_out * hidden), 0.9f, 0.08f);
|
|
const auto gate_values = patterned(static_cast<size_t>(intermediate * hidden), 1.3f, 0.11f);
|
|
const auto up_values = patterned(static_cast<size_t>(intermediate * hidden), 1.7f, 0.09f);
|
|
|
|
if (packed) {
|
|
weights.self_attention.qkv_weight = make_f32({q_out + 2 * kv_out, hidden});
|
|
weights.mlp.gate_up_proj = engine::modules::LinearWeights{
|
|
make_f32({intermediate * 2, hidden}),
|
|
std::nullopt,
|
|
};
|
|
} else {
|
|
weights.self_attention.q_weight = make_f32({q_out, hidden});
|
|
weights.self_attention.k_weight = make_f32({kv_out, hidden});
|
|
weights.self_attention.v_weight = make_f32({kv_out, hidden});
|
|
weights.mlp.gate_proj = {make_f32({intermediate, hidden}), std::nullopt};
|
|
weights.mlp.up_proj = {make_f32({intermediate, hidden}), std::nullopt};
|
|
}
|
|
|
|
engine::modules::DecoderLayerConfig config;
|
|
config.hidden_size = hidden;
|
|
config.num_attention_heads = heads;
|
|
config.num_key_value_heads = kv_heads;
|
|
config.head_dim = head_dim;
|
|
config.intermediate_size = intermediate;
|
|
config.rms_norm_eps = 1e-5f;
|
|
config.qkv_layout = packed
|
|
? engine::modules::DecoderQKVLayout::PackedQKV
|
|
: engine::modules::DecoderQKVLayout::Separate;
|
|
config.runtime.mlp.mode = packed
|
|
? engine::modules::DecoderMLPMode::PackedGateUp
|
|
: (fused_mlp ? engine::modules::DecoderMLPMode::FusedSwiGLU : engine::modules::DecoderMLPMode::Exact);
|
|
config.use_qk_norm = false;
|
|
config.runtime.attention.prefill_mode = engine::modules::DecoderAttentionMode::ManualRepeat;
|
|
|
|
ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false);
|
|
engine::modules::DecoderLayerOutputs outputs;
|
|
engine::core::TensorValue cache_key, cache_value, cache_slot, mask;
|
|
if (batched_decode) {
|
|
cache_key = make_f32({batch, cache_steps, kv_heads, head_dim});
|
|
cache_value = make_f32({batch, cache_steps, kv_heads, head_dim});
|
|
cache_slot = engine::core::make_tensor(ctx, GGML_TYPE_I32,
|
|
engine::core::TensorShape::from_dims({batch}));
|
|
mask = make_f32({batch, 1, 1, cache_steps});
|
|
config.runtime.static_cache.update_mode = engine::modules::DecoderStaticCacheUpdateMode::DirectSetRows;
|
|
config.runtime.attention.static_mode = engine::modules::DecoderAttentionMode::ManualRepeat;
|
|
outputs = engine::modules::DecoderLayerModule(config).build_with_static_cache_tail_batched(
|
|
ctx, graph, input, positions, weights, cache_key, cache_value, cache_slot, mask);
|
|
} else {
|
|
outputs = engine::modules::DecoderLayerModule(config).build(ctx, input, positions, weights);
|
|
}
|
|
ggml_build_forward_expand(graph, outputs.output.tensor);
|
|
if (batched_decode) {
|
|
int rope_nodes = 0;
|
|
int copied_positions = 0;
|
|
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
|
|
const auto * node = ggml_graph_node(graph, i);
|
|
if (node->op != GGML_OP_ROPE) {
|
|
continue;
|
|
}
|
|
++rope_nodes;
|
|
if (node->src[1]->op == GGML_OP_CONT) {
|
|
++copied_positions;
|
|
}
|
|
if (backend_type == engine::core::BackendType::Vulkan && node->src[1]->view_offs != 0) {
|
|
throw std::runtime_error("Vulkan batched RoPE retains an offset position view");
|
|
}
|
|
}
|
|
const int expected_copies = backend_type == engine::core::BackendType::Vulkan ? 2 : 0;
|
|
if (rope_nodes != 4 || copied_positions != expected_copies) {
|
|
throw std::runtime_error("Unexpected batched RoPE position layout");
|
|
}
|
|
}
|
|
buffer = ggml_backend_alloc_ctx_tensors(ggml, backend);
|
|
if (buffer == nullptr) {
|
|
throw std::runtime_error("failed to allocate test tensors");
|
|
}
|
|
|
|
engine::core::write_tensor_f32(input, patterned(static_cast<size_t>(batch * steps * hidden), 2.1f, 0.20f));
|
|
engine::core::write_tensor_i32(positions, batched_decode ? std::vector<int32_t>{2, 5} : std::vector<int32_t>{0, 1, 2});
|
|
if (batched_decode) {
|
|
engine::core::write_tensor_f32(cache_key, patterned(batch * cache_steps * kv_out, 0.4f, 0.1f));
|
|
engine::core::write_tensor_f32(cache_value, patterned(batch * cache_steps * kv_out, 0.8f, 0.1f));
|
|
engine::core::write_tensor_i32(cache_slot, {2, cache_steps + 5});
|
|
std::vector<float> mask_values(batch * cache_steps, -std::numeric_limits<float>::infinity());
|
|
for (int64_t b = 0; b < batch; ++b) {
|
|
for (int64_t t = 0; t <= (b == 0 ? 2 : 5); ++t) {
|
|
mask_values[b * cache_steps + t] = 0.0f;
|
|
}
|
|
}
|
|
engine::core::write_tensor_f32(mask, mask_values);
|
|
}
|
|
engine::core::write_tensor_f32(*weights.input_norm.weight, patterned(hidden, 0.3f, 0.7f));
|
|
engine::core::write_tensor_f32(*weights.post_norm.weight, patterned(hidden, 0.7f, 0.8f));
|
|
engine::core::write_tensor_f32(
|
|
weights.self_attention.out_weight,
|
|
patterned(static_cast<size_t>(hidden * hidden), 1.1f, 0.10f));
|
|
engine::core::write_tensor_f32(
|
|
weights.mlp.down_proj.weight,
|
|
patterned(static_cast<size_t>(hidden * intermediate), 1.9f, 0.10f));
|
|
|
|
if (packed) {
|
|
std::vector<float> qkv_values;
|
|
qkv_values.reserve(q_values.size() + k_values.size() + v_values.size());
|
|
qkv_values.insert(qkv_values.end(), q_values.begin(), q_values.end());
|
|
qkv_values.insert(qkv_values.end(), k_values.begin(), k_values.end());
|
|
qkv_values.insert(qkv_values.end(), v_values.begin(), v_values.end());
|
|
engine::core::write_tensor_f32(*weights.self_attention.qkv_weight, qkv_values);
|
|
|
|
std::vector<float> gate_up_values;
|
|
gate_up_values.reserve(gate_values.size() + up_values.size());
|
|
gate_up_values.insert(gate_up_values.end(), gate_values.begin(), gate_values.end());
|
|
gate_up_values.insert(gate_up_values.end(), up_values.begin(), up_values.end());
|
|
engine::core::write_tensor_f32(weights.mlp.gate_up_proj->weight, gate_up_values);
|
|
} else {
|
|
engine::core::write_tensor_f32(weights.self_attention.q_weight, q_values);
|
|
engine::core::write_tensor_f32(weights.self_attention.k_weight, k_values);
|
|
engine::core::write_tensor_f32(weights.self_attention.v_weight, v_values);
|
|
engine::core::write_tensor_f32(weights.mlp.gate_proj.weight, gate_values);
|
|
engine::core::write_tensor_f32(weights.mlp.up_proj.weight, up_values);
|
|
}
|
|
|
|
if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) {
|
|
throw std::runtime_error("test graph execution failed");
|
|
}
|
|
LayerResult result;
|
|
engine::core::read_tensor_f32_into(outputs.output.tensor, result.output);
|
|
engine::core::read_tensor_f32_into(outputs.key.tensor, result.key);
|
|
engine::core::read_tensor_f32_into(outputs.value.tensor, result.value);
|
|
|
|
ggml_backend_buffer_free(buffer);
|
|
buffer = nullptr;
|
|
ggml_free(ggml);
|
|
ggml_backend_free(backend);
|
|
return result;
|
|
} catch (...) {
|
|
if (buffer != nullptr) {
|
|
ggml_backend_buffer_free(buffer);
|
|
}
|
|
ggml_free(ggml);
|
|
ggml_backend_free(backend);
|
|
throw;
|
|
}
|
|
}
|
|
|
|
void test_packed_qkv_and_gate_up_match_separate_projections() {
|
|
const auto separate = run_layer(false);
|
|
const auto packed = run_layer(true);
|
|
require_allclose(packed.output, separate.output, 2.0e-5f, "decoder output");
|
|
require_allclose(packed.key, separate.key, 2.0e-5f, "decoder key");
|
|
require_allclose(packed.value, separate.value, 2.0e-5f, "decoder value");
|
|
}
|
|
|
|
void test_suffix_causal_mask() {
|
|
const auto values = engine::modules::causal_suffix_mask_values(2, 3, 2);
|
|
if (values.size() != 30) {
|
|
throw std::runtime_error("suffix causal mask size mismatch");
|
|
}
|
|
const std::vector<bool> expected{
|
|
true, true, true, false, false,
|
|
true, true, true, true, false,
|
|
true, true, true, true, true,
|
|
};
|
|
for (int batch = 0; batch < 2; ++batch) {
|
|
for (size_t i = 0; i < expected.size(); ++i) {
|
|
const float actual = ggml_fp16_to_fp32(values[static_cast<size_t>(batch) * expected.size() + i]);
|
|
if ((expected[i] && actual != 0.0F) || (!expected[i] && !std::isinf(actual))) {
|
|
throw std::runtime_error("suffix causal mask visibility mismatch");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void test_f16_kv_set_rows(engine::modules::FastKVSetRowsMode mode) {
|
|
engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4};
|
|
ggml_backend_t backend = engine::core::init_backend(backend_config);
|
|
if (backend == nullptr) {
|
|
throw std::runtime_error("failed to initialize CPU backend");
|
|
}
|
|
|
|
ggml_init_params params{kGraphBytes, nullptr, true};
|
|
ggml_context * ggml = ggml_init(params);
|
|
if (ggml == nullptr) {
|
|
ggml_backend_free(backend);
|
|
throw std::runtime_error("failed to initialize GGML context");
|
|
}
|
|
|
|
ggml_backend_buffer_t buffer = nullptr;
|
|
try {
|
|
engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_test", engine::core::BackendType::Cpu};
|
|
const auto cache = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F16,
|
|
engine::core::TensorShape::from_dims({1, 3, 1, 2}));
|
|
const auto row = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({1, 1, 1, 2}));
|
|
const auto row_index = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_I32,
|
|
engine::core::TensorShape::from_dims({1}));
|
|
const auto output = engine::modules::FastKVSetRowsModule({
|
|
mode,
|
|
}).build(ctx, cache, row, row_index);
|
|
|
|
ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false);
|
|
ggml_build_forward_expand(graph, output.tensor);
|
|
buffer = ggml_backend_alloc_ctx_tensors(ggml, backend);
|
|
if (buffer == nullptr) {
|
|
throw std::runtime_error("failed to allocate f16 KV test tensors");
|
|
}
|
|
|
|
engine::core::write_tensor_f16(cache, std::vector<float>(6, 0.0F));
|
|
engine::core::write_tensor_f32(row, {1.25F, -2.5F});
|
|
const int32_t index = 1;
|
|
ggml_backend_tensor_set(row_index.tensor, &index, 0, sizeof(index));
|
|
if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) {
|
|
throw std::runtime_error("f16 KV set-rows graph compute failed");
|
|
}
|
|
|
|
const auto values = engine::core::read_tensor_f16(output.tensor);
|
|
require_allclose(values, {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F}, 1.0e-3F, "f16 KV cache");
|
|
|
|
ggml_backend_buffer_free(buffer);
|
|
buffer = nullptr;
|
|
ggml_free(ggml);
|
|
ggml_backend_free(backend);
|
|
} catch (...) {
|
|
if (buffer != nullptr) {
|
|
ggml_backend_buffer_free(buffer);
|
|
}
|
|
ggml_free(ggml);
|
|
ggml_backend_free(backend);
|
|
throw;
|
|
}
|
|
}
|
|
|
|
void test_f16_kv_set_rows_batched(engine::modules::FastKVSetRowsMode mode) {
|
|
engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4};
|
|
ggml_backend_t backend = engine::core::init_backend(backend_config);
|
|
if (backend == nullptr) {
|
|
throw std::runtime_error("failed to initialize CPU backend");
|
|
}
|
|
|
|
ggml_init_params params{kGraphBytes, nullptr, true};
|
|
ggml_context * ggml = ggml_init(params);
|
|
if (ggml == nullptr) {
|
|
ggml_backend_free(backend);
|
|
throw std::runtime_error("failed to initialize GGML context");
|
|
}
|
|
|
|
ggml_backend_buffer_t buffer = nullptr;
|
|
try {
|
|
engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_batched_test", engine::core::BackendType::Cpu};
|
|
const auto cache = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F16,
|
|
engine::core::TensorShape::from_dims({2, 3, 1, 2}));
|
|
const auto rows = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({2, 1, 1, 2}));
|
|
const auto row_indices = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_I64,
|
|
engine::core::TensorShape::from_dims({2}));
|
|
const auto output = engine::modules::FastKVSetRowsModule({
|
|
mode,
|
|
}).build(ctx, cache, rows, row_indices);
|
|
|
|
ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false);
|
|
ggml_build_forward_expand(graph, output.tensor);
|
|
buffer = ggml_backend_alloc_ctx_tensors(ggml, backend);
|
|
if (buffer == nullptr) {
|
|
throw std::runtime_error("failed to allocate batched F16 KV test tensors");
|
|
}
|
|
|
|
engine::core::write_tensor_f16(cache, std::vector<float>(12, 0.0F));
|
|
engine::core::write_tensor_f32(rows, {1.25F, -2.5F, 3.5F, -4.5F});
|
|
const std::vector<int64_t> indices{1, 4};
|
|
ggml_backend_tensor_set(row_indices.tensor, indices.data(), 0, indices.size() * sizeof(int64_t));
|
|
if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) {
|
|
throw std::runtime_error("batched F16 KV set-rows graph compute failed");
|
|
}
|
|
|
|
const auto values = engine::core::read_tensor_f16(output.tensor);
|
|
require_allclose(
|
|
values,
|
|
{0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F,
|
|
0.0F, 0.0F, 3.5F, -4.5F, 0.0F, 0.0F},
|
|
1.0e-3F,
|
|
"batched f16 KV cache");
|
|
|
|
ggml_backend_buffer_free(buffer);
|
|
buffer = nullptr;
|
|
ggml_free(ggml);
|
|
ggml_backend_free(backend);
|
|
} catch (...) {
|
|
if (buffer != nullptr) {
|
|
ggml_backend_buffer_free(buffer);
|
|
}
|
|
ggml_free(ggml);
|
|
ggml_backend_free(backend);
|
|
throw;
|
|
}
|
|
}
|
|
|
|
int count_graph_op(ggml_cgraph * graph, ggml_op op) {
|
|
int count = 0;
|
|
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
|
|
const ggml_tensor * node = ggml_graph_node(graph, i);
|
|
count += node != nullptr && node->op == op ? 1 : 0;
|
|
}
|
|
return count;
|
|
}
|
|
|
|
bool graph_contains_sequence(ggml_cgraph * graph, std::initializer_list<ggml_op> ops) {
|
|
if (ops.size() == 0 || static_cast<int>(ops.size()) > ggml_graph_n_nodes(graph)) {
|
|
return false;
|
|
}
|
|
for (int start = 0; start + static_cast<int>(ops.size()) <= ggml_graph_n_nodes(graph); ++start) {
|
|
bool matches = true;
|
|
int offset = 0;
|
|
for (const ggml_op op : ops) {
|
|
const ggml_tensor * node = ggml_graph_node(graph, start + offset++);
|
|
matches = matches && node != nullptr && node->op == op;
|
|
}
|
|
if (matches) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
std::string graph_ops(ggml_cgraph * graph) {
|
|
std::ostringstream out;
|
|
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
|
|
const ggml_tensor * node = ggml_graph_node(graph, i);
|
|
if (i != 0) {
|
|
out << ',';
|
|
}
|
|
out << (node != nullptr ? ggml_op_name(node->op) : "null");
|
|
}
|
|
return out.str();
|
|
}
|
|
|
|
void test_fast_projection_accepts_cuda_or_hip_backends() {
|
|
constexpr int64_t in_features = 8;
|
|
constexpr int64_t out_features = 12;
|
|
for (const auto backend_type : {
|
|
engine::core::BackendType::Cuda,
|
|
engine::core::BackendType::Hip,
|
|
}) {
|
|
ggml_init_params params{kGraphBytes, nullptr, true};
|
|
ggml_context * ggml = ggml_init(params);
|
|
if (ggml == nullptr) {
|
|
throw std::runtime_error("failed to initialize fast projection graph test context");
|
|
}
|
|
try {
|
|
engine::core::ModuleBuildContext ctx{ggml, "fast_projection_cuda_or_hip_test", backend_type};
|
|
const auto input = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({1, in_features}));
|
|
const auto weight = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({out_features, in_features}));
|
|
const auto output = engine::modules::FastPackedProjection4Module({in_features, out_features})
|
|
.build(ctx, input, {weight, std::nullopt});
|
|
if (output.shape.rank != 2 || output.shape.dims[0] != 1 || output.shape.dims[1] != out_features) {
|
|
throw std::runtime_error("fast projection CUDA/HIP output shape mismatch");
|
|
}
|
|
ggml_free(ggml);
|
|
} catch (...) {
|
|
ggml_free(ggml);
|
|
throw;
|
|
}
|
|
}
|
|
|
|
for (const auto rejected_backend : {engine::core::BackendType::Cpu, engine::core::BackendType::Metal}) {
|
|
ggml_init_params params{kGraphBytes, nullptr, true};
|
|
ggml_context * ggml = ggml_init(params);
|
|
if (ggml == nullptr) {
|
|
throw std::runtime_error("failed to initialize fast projection rejection test context");
|
|
}
|
|
bool rejected = false;
|
|
try {
|
|
engine::core::ModuleBuildContext ctx{ggml, "fast_projection_rejection_test", rejected_backend};
|
|
const auto input = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({1, in_features}));
|
|
const auto weight = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({out_features, in_features}));
|
|
try {
|
|
(void) engine::modules::FastPackedProjection4Module({in_features, out_features})
|
|
.build(ctx, input, {weight, std::nullopt});
|
|
} catch (const std::runtime_error &) {
|
|
rejected = true;
|
|
}
|
|
ggml_free(ggml);
|
|
} catch (...) {
|
|
ggml_free(ggml);
|
|
throw;
|
|
}
|
|
if (!rejected) {
|
|
throw std::runtime_error("fast projection must reject non-CUDA/HIP backends");
|
|
}
|
|
}
|
|
}
|
|
|
|
void test_higgs_decode_graph_exposes_cuda_fast_paths() {
|
|
constexpr int64_t hidden = 8;
|
|
constexpr int64_t heads = 2;
|
|
constexpr int64_t kv_heads = 1;
|
|
constexpr int64_t head_dim = 4;
|
|
constexpr int64_t intermediate = 12;
|
|
constexpr int64_t cache_steps = 8;
|
|
constexpr int64_t qkv_out = heads * head_dim + 2 * kv_heads * head_dim;
|
|
|
|
ggml_init_params params{kGraphBytes, nullptr, true};
|
|
ggml_context * ggml = ggml_init(params);
|
|
if (ggml == nullptr) {
|
|
throw std::runtime_error("failed to initialize Higgs decode graph test context");
|
|
}
|
|
|
|
try {
|
|
engine::core::ModuleBuildContext ctx{ggml, "higgs_decode_fast_path_test", engine::core::BackendType::Cuda};
|
|
auto make_tensor = [&](ggml_type type, std::initializer_list<int64_t> dims) {
|
|
return engine::core::make_tensor(ctx, type, engine::core::TensorShape::from_dims(dims));
|
|
};
|
|
|
|
const auto input = make_tensor(GGML_TYPE_F32, {1, 1, hidden});
|
|
const auto positions = make_tensor(GGML_TYPE_I32, {1});
|
|
const auto cache_key = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim});
|
|
const auto cache_value = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim});
|
|
const auto cache_slot = make_tensor(GGML_TYPE_I64, {1});
|
|
const auto attention_mask = make_tensor(GGML_TYPE_F16, {1, 1, 1, cache_steps});
|
|
|
|
engine::modules::DecoderLayerWeights weights;
|
|
weights.input_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt};
|
|
weights.q_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt};
|
|
weights.k_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt};
|
|
weights.post_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt};
|
|
weights.self_attention.qkv_weight = make_tensor(GGML_TYPE_F32, {qkv_out, hidden});
|
|
weights.self_attention.out_weight = make_tensor(GGML_TYPE_F32, {hidden, hidden});
|
|
weights.mlp.gate_up_proj = engine::modules::LinearWeights{
|
|
make_tensor(GGML_TYPE_F32, {intermediate * 2, hidden}),
|
|
std::nullopt,
|
|
};
|
|
weights.mlp.down_proj = {
|
|
make_tensor(GGML_TYPE_F32, {hidden, intermediate}),
|
|
std::nullopt,
|
|
};
|
|
|
|
engine::modules::DecoderLayerConfig config;
|
|
config.hidden_size = hidden;
|
|
config.num_attention_heads = heads;
|
|
config.num_key_value_heads = kv_heads;
|
|
config.head_dim = head_dim;
|
|
config.intermediate_size = intermediate;
|
|
config.qkv_layout = engine::modules::DecoderQKVLayout::PackedQKV;
|
|
config.use_qk_norm = true;
|
|
config.runtime.attention.static_mode =
|
|
engine::modules::DecoderAttentionMode::FlashGroupedViewKV;
|
|
config.runtime.static_cache.update_mode =
|
|
engine::modules::DecoderStaticCacheUpdateMode::DirectSetRows;
|
|
config.runtime.static_cache.set_rows_mode =
|
|
engine::modules::DecoderStaticCacheSetRowsMode::BackendViewOptimized;
|
|
config.runtime.mlp.mode = engine::modules::DecoderMLPMode::PackedGateUp;
|
|
|
|
ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false);
|
|
const auto outputs = engine::modules::DecoderLayerModule(config).build_with_static_cache_tail(
|
|
ctx,
|
|
graph,
|
|
input,
|
|
positions,
|
|
weights,
|
|
cache_key,
|
|
cache_value,
|
|
cache_slot,
|
|
attention_mask);
|
|
ggml_build_forward_expand(graph, outputs.output.tensor);
|
|
|
|
if (count_graph_op(graph, GGML_OP_FLASH_ATTN_EXT) != 1) {
|
|
throw std::runtime_error("Higgs decode graph must contain one grouped FlashAttention op");
|
|
}
|
|
if (count_graph_op(graph, GGML_OP_SET_ROWS) != 2) {
|
|
throw std::runtime_error("Higgs decode graph must update both F16 KV caches with set-rows");
|
|
}
|
|
if (count_graph_op(graph, GGML_OP_GLU) != 1) {
|
|
throw std::runtime_error("Higgs decode graph must contain one packed SwiGLU op");
|
|
}
|
|
if (count_graph_op(graph, GGML_OP_REPEAT) != 0) {
|
|
throw std::runtime_error("grouped FlashAttention must not materialize repeated KV heads");
|
|
}
|
|
if (!graph_contains_sequence(graph, {GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS})) {
|
|
throw std::runtime_error(
|
|
"Higgs key-cache update must expose CUDA RoPE/view/set-rows fusion; graph=" +
|
|
graph_ops(graph));
|
|
}
|
|
|
|
ggml_free(ggml);
|
|
} catch (...) {
|
|
ggml_free(ggml);
|
|
throw;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char ** argv) {
|
|
const auto started = std::chrono::steady_clock::now();
|
|
try {
|
|
bool vulkan = false;
|
|
bool cuda = false;
|
|
for (int i = 1; i < argc; ++i) {
|
|
const std::string arg = argv[i];
|
|
if (arg == "--vulkan") {
|
|
vulkan = true;
|
|
} else if (arg == "--cuda") {
|
|
cuda = true;
|
|
} else if (arg == "--log") {
|
|
engine::debug::configure_logging({true, ""});
|
|
} else {
|
|
throw std::runtime_error("Unknown argument: " + arg);
|
|
}
|
|
}
|
|
test_packed_qkv_and_gate_up_match_separate_projections();
|
|
for (const bool batched : {false, true}) {
|
|
const auto reference = run_layer(false, engine::core::BackendType::Cpu, batched);
|
|
const auto fused = run_layer(false, engine::core::BackendType::Cpu, batched, true);
|
|
require_allclose(fused.output, reference.output, 2.0e-5f, "fused MLP decoder output");
|
|
if (cuda) {
|
|
const auto actual = run_layer(false, engine::core::BackendType::Cuda, batched, true);
|
|
require_allclose(actual.output, reference.output, 2.0e-5f, "CUDA fused MLP decoder output");
|
|
}
|
|
if (vulkan) {
|
|
const auto backend_reference = run_layer(false, engine::core::BackendType::Vulkan, batched);
|
|
const auto actual = run_layer(false, engine::core::BackendType::Vulkan, batched, true);
|
|
require_allclose(actual.output, backend_reference.output, 2.0e-5f, "Vulkan fused MLP decoder output");
|
|
}
|
|
}
|
|
for (bool packed : {false, true}) {
|
|
const auto reference = run_layer(packed, engine::core::BackendType::Cpu, true);
|
|
if (vulkan) {
|
|
const auto actual = run_layer(packed, engine::core::BackendType::Vulkan, true);
|
|
require_allclose(actual.output, reference.output, 2.0e-5f, "Vulkan batched decoder output");
|
|
require_allclose(actual.key, reference.key, 2.0e-5f, "Vulkan batched decoder key");
|
|
require_allclose(actual.value, reference.value, 2.0e-5f, "Vulkan batched decoder value");
|
|
}
|
|
}
|
|
test_suffix_causal_mask();
|
|
for (const auto mode : {engine::modules::FastKVSetRowsMode::Exact,
|
|
engine::modules::FastKVSetRowsMode::BackendViewOptimized}) {
|
|
test_f16_kv_set_rows(mode);
|
|
test_f16_kv_set_rows_batched(mode);
|
|
}
|
|
test_fast_projection_accepts_cuda_or_hip_backends();
|
|
test_higgs_decode_graph_exposes_cuda_fast_paths();
|
|
std::cout << "decoder_packed_projection_test: ok\n";
|
|
engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(started));
|
|
return 0;
|
|
} catch (const std::exception & ex) {
|
|
engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(started));
|
|
std::cerr << "decoder_packed_projection_test: failed: " << ex.what() << "\n";
|
|
return 1;
|
|
}
|
|
}
|