File size: 1,417 Bytes
3fd1a35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | #include "ling3/execution_plan.h"
namespace ling3 {
ExecutionPlan BuildExecutionPlan(std::size_t max_context) {
ExecutionPlan plan;
std::size_t kda_layers = 0;
std::size_t mla_layers = 0;
for (std::size_t layer = 0; layer < plan.layers.size(); ++layer) {
const bool mla = (layer + 1) % 4 == 0;
plan.layers[layer] = {
static_cast<std::uint8_t>(layer),
mla ? AttentionKind::kMla : AttentionKind::kKda,
layer == 0 ? FfnKind::kDense : FfnKind::kMoe,
};
kda_layers += !mla;
mla_layers += mla;
}
constexpr std::size_t routed_expert_contexts = 128 * 2;
constexpr std::size_t shared_expert_contexts = 2;
constexpr std::size_t kda_projection_contexts = 7;
constexpr std::size_t mla_projection_contexts = 6;
plan.decode_matmul_contexts =
23 * (routed_expert_contexts + shared_expert_contexts) +
kda_layers * kda_projection_contexts + mla_layers * mla_projection_contexts + 3;
plan.rknn_island_contexts = kda_layers * 3 + 23 * 3 + 24;
plan.kda_state_bytes_fp16 = kda_layers * 16 * 128 * 128 * sizeof(std::uint16_t);
plan.mla_cache_bytes_fp16 = mla_layers * max_context * (512 + 64) * sizeof(std::uint16_t);
plan.activation_arena_bytes =
2 * 1536 + 3 * (1024 * sizeof(std::int32_t) + 1536 * sizeof(std::int32_t)) + 16 * 128 * 4;
return plan;
}
} // namespace ling3
|