| 1892 | class DecodeGraph { |
| 1893 | public: |
| 1894 | DecodeGraph( |
| 1895 | std::shared_ptr<PlannerWeightsRuntime> runtime, |
| 1896 | int64_t cache_steps, |
| 1897 | size_t graph_arena_bytes) |
| 1898 | : runtime_(std::move(runtime)), |
| 1899 | cache_steps_(cache_steps) { |
| 1900 | if (cache_steps_ <= 0) { |
| 1901 | throw std::runtime_error("ACE-Step planner decode requires positive cache length"); |
| 1902 | } |
| 1903 | ggml_init_params params{graph_arena_bytes, nullptr, true}; |
| 1904 | ctx_.reset(ggml_init(params)); |
| 1905 | if (ctx_ == nullptr) { |
| 1906 | throw std::runtime_error("failed to initialize ACE-Step planner decode graph context"); |
| 1907 | } |
| 1908 | const auto & config = runtime_->assets().config.planner; |
| 1909 | const auto & weights = runtime_->weights(); |
| 1910 | core::ModuleBuildContext ctx{ctx_.get(), "ace_step.planner.decode", runtime_->backend_type()}; |
| 1911 | token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); |
| 1912 | auto token = core::wrap_tensor(token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); |
| 1913 | auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) |
| 1914 | .build(ctx, token, weights.token_embedding); |
| 1915 | x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 1, config.hidden_size})); |
| 1916 | positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); |
| 1917 | auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); |
| 1918 | attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_ + 1, 1, 1, 1); |
| 1919 | auto attention_mask = core::wrap_tensor( |
| 1920 | attention_mask_, |
| 1921 | core::TensorShape::from_dims({1, 1, 1, cache_steps_ + 1}), |
| 1922 | GGML_TYPE_F16); |
| 1923 | graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); |
| 1924 | |
| 1925 | std::vector<core::TensorValue> cache_keys; |
| 1926 | std::vector<core::TensorValue> cache_values; |
| 1927 | for (const auto & layer : weights.layers.layers) { |
| 1928 | cache_keys.push_back(core::make_tensor( |
| 1929 | ctx, |
| 1930 | GGML_TYPE_F32, |
| 1931 | core::TensorShape::from_dims({1, cache_steps_ + 1, config.num_key_value_heads, config.head_dim}))); |
| 1932 | cache_values.push_back(core::make_tensor( |
| 1933 | ctx, |
| 1934 | GGML_TYPE_F32, |
| 1935 | core::TensorShape::from_dims({1, cache_steps_ + 1, config.num_key_value_heads, config.head_dim}))); |
| 1936 | auto out = planner_decoder_layer_with_static_cache_tail_batched( |
| 1937 | ctx, |
| 1938 | graph_, |
| 1939 | x, |
| 1940 | positions, |
| 1941 | layer, |
| 1942 | config, |
| 1943 | cache_keys.back(), |
| 1944 | cache_values.back(), |
| 1945 | attention_mask, |
| 1946 | GGML_TYPE_F32); |
| 1947 | x = out.output; |
| 1948 | key_sources_.push_back(ggml_view_1d(ctx_.get(), out.key.tensor, config.num_key_value_heads * config.head_dim, 0)); |
| 1949 | value_sources_.push_back(ggml_view_1d(ctx_.get(), out.value.tensor, config.num_key_value_heads * config.head_dim, 0)); |
| 1950 | } |
| 1951 | step_cache_ = runtime::TransformerKVCache( |
nothing calls this directly
no test coverage detected