tts: reuse decode scratch buffers across streaming steps
Keep the talker causal mask in each static attention window graph and reuse positions, KV rows and sliding attention mask buffers in each codec stream graph class. This removes the remaining host allocations from the talker and codec streaming hot loops without changing graph execution.
This commit is contained in:
@@ -492,18 +492,17 @@ bool pipeline_codec_decode_stream(PipelineCodec * pc, const int32_t * codes, int
|
|||||||
|
|
||||||
ggml_backend_tensor_set(sg->codes, codes, 0, (size_t) T * (size_t) K * sizeof(int32_t));
|
ggml_backend_tensor_set(sg->codes, codes, 0, (size_t) T * (size_t) K * sizeof(int32_t));
|
||||||
|
|
||||||
std::vector<int32_t> pos((size_t) T);
|
sg->pos_buf.resize((size_t) T);
|
||||||
std::vector<int64_t> rows((size_t) T);
|
sg->rows_buf.resize((size_t) T);
|
||||||
for (int t = 0; t < T; t++) {
|
for (int t = 0; t < T; t++) {
|
||||||
pos[(size_t) t] = pc->stream_pos + t;
|
sg->pos_buf[(size_t) t] = pc->stream_pos + t;
|
||||||
rows[(size_t) t] = (int64_t) ((pc->stream_pos + t) % ring);
|
sg->rows_buf[(size_t) t] = (int64_t) ((pc->stream_pos + t) % ring);
|
||||||
}
|
}
|
||||||
ggml_backend_tensor_set(sg->pos, pos.data(), 0, (size_t) T * sizeof(int32_t));
|
ggml_backend_tensor_set(sg->pos, sg->pos_buf.data(), 0, (size_t) T * sizeof(int32_t));
|
||||||
ggml_backend_tensor_set(sg->rows, rows.data(), 0, (size_t) T * sizeof(int64_t));
|
ggml_backend_tensor_set(sg->rows, sg->rows_buf.data(), 0, (size_t) T * sizeof(int64_t));
|
||||||
|
|
||||||
std::vector<float> mask_buf;
|
tok_trans_build_stream_mask(pc->stream_pos, T, ring, pc->transformer.sliding_window, sg->mask_buf);
|
||||||
tok_trans_build_stream_mask(pc->stream_pos, T, ring, pc->transformer.sliding_window, mask_buf);
|
ggml_backend_tensor_set(sg->mask, sg->mask_buf.data(), 0, sg->mask_buf.size() * sizeof(float));
|
||||||
ggml_backend_tensor_set(sg->mask, mask_buf.data(), 0, mask_buf.size() * sizeof(float));
|
|
||||||
|
|
||||||
enum ggml_status st = ggml_backend_graph_compute(pc->backend, sg->gf);
|
enum ggml_status st = ggml_backend_graph_compute(pc->backend, sg->gf);
|
||||||
if (st != GGML_STATUS_SUCCESS) {
|
if (st != GGML_STATUS_SUCCESS) {
|
||||||
|
|||||||
@@ -64,6 +64,9 @@ struct CodecStreamGraph {
|
|||||||
struct ggml_tensor * rows = nullptr;
|
struct ggml_tensor * rows = nullptr;
|
||||||
struct ggml_tensor * mask = nullptr;
|
struct ggml_tensor * mask = nullptr;
|
||||||
struct ggml_tensor * out = nullptr;
|
struct ggml_tensor * out = nullptr;
|
||||||
|
std::vector<int32_t> pos_buf;
|
||||||
|
std::vector<int64_t> rows_buf;
|
||||||
|
std::vector<float> mask_buf;
|
||||||
};
|
};
|
||||||
|
|
||||||
// One primed stream state snapshot: a mirror of every conv context and
|
// One primed stream state snapshot: a mirror of every conv context and
|
||||||
|
|||||||
@@ -9,6 +9,8 @@
|
|||||||
#include "ggml-alloc.h"
|
#include "ggml-alloc.h"
|
||||||
#include "ggml.h"
|
#include "ggml.h"
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
struct TalkerDecodeGraph {
|
struct TalkerDecodeGraph {
|
||||||
struct ggml_context * ctx = nullptr;
|
struct ggml_context * ctx = nullptr;
|
||||||
struct ggml_cgraph * gf = nullptr;
|
struct ggml_cgraph * gf = nullptr;
|
||||||
@@ -19,6 +21,7 @@ struct TalkerDecodeGraph {
|
|||||||
struct ggml_tensor * rows_in = nullptr; // [1] i64
|
struct ggml_tensor * rows_in = nullptr; // [1] i64
|
||||||
struct ggml_tensor * mask_in = nullptr; // [n_kv_pad, 1] f16
|
struct ggml_tensor * mask_in = nullptr; // [n_kv_pad, 1] f16
|
||||||
struct ggml_tensor * logits = nullptr; // [vocab, 1] f32
|
struct ggml_tensor * logits = nullptr; // [vocab, 1] f32
|
||||||
|
std::vector<ggml_fp16_t> mask; // [n_kv_pad] f16
|
||||||
int n_kv_pad = 0; // window class width, 0 marks an empty slot
|
int n_kv_pad = 0; // window class width, 0 marks an empty slot
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -38,5 +41,6 @@ static void talker_decode_graph_free(TalkerDecodeGraph * tg) {
|
|||||||
tg->rows_in = nullptr;
|
tg->rows_in = nullptr;
|
||||||
tg->mask_in = nullptr;
|
tg->mask_in = nullptr;
|
||||||
tg->logits = nullptr;
|
tg->logits = nullptr;
|
||||||
|
tg->mask.clear();
|
||||||
tg->n_kv_pad = 0;
|
tg->n_kv_pad = 0;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -571,6 +571,7 @@ static bool talker_decode_graph_build(const TalkerWeights * tw,
|
|||||||
tg->rows_in = rows_in;
|
tg->rows_in = rows_in;
|
||||||
tg->mask_in = mask_in;
|
tg->mask_in = mask_in;
|
||||||
tg->logits = logits;
|
tg->logits = logits;
|
||||||
|
tg->mask.resize((size_t) n_kv_pad);
|
||||||
tg->n_kv_pad = n_kv_pad;
|
tg->n_kv_pad = n_kv_pad;
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
@@ -620,7 +621,7 @@ static bool talker_forward_decode(const TalkerWeights * tw,
|
|||||||
|
|
||||||
// Causal mask: keys [0, n_past] carry 0, the padded tail neg inf.
|
// Causal mask: keys [0, n_past] carry 0, the padded tail neg inf.
|
||||||
{
|
{
|
||||||
std::vector<ggml_fp16_t> mask((size_t) tg->n_kv_pad);
|
std::vector<ggml_fp16_t> & mask = tg->mask;
|
||||||
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
|
const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f);
|
||||||
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
|
const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-INFINITY);
|
||||||
for (size_t i = 0; i < mask.size(); i++) {
|
for (size_t i = 0; i < mask.size(); i++) {
|
||||||
|
|||||||
Reference in New Issue
Block a user