vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-07-10 01:18:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

updated ace step cpp

Browse source
This commit is contained in:
Concedo 2026-02-23 23:01:10 +08:00
parent 2e713cfff5
commit 5311997581
10 changed files with 508 additions and 392 deletions
Download patch file Download diff file Expand all files Collapse all files

595
otherarch/acestep/ace-qwen3.cpp
View file

@ -46,41 +46,75 @@ struct TokenProb {
float prob;
};
static int sample_top_p(float * logits, int vocab_size, float temperature, float top_p, std::mt19937 & rng) {
for (int i = 0; i < vocab_size; i++)
logits[i] /= temperature;
float max_val = *std::max_element(logits, logits + vocab_size);
// Sampling: temperature -> top_k -> top_p -> softmax -> multinomial
// Matches nano-vLLM Sampler: div_(temperature) -> apply_top_k_top_p -> softmax -> sample
static int sample_top_k_p(float * logits, int V, float temperature, float top_p, int top_k, std::mt19937 & rng) {
if (temperature <= 0.0f) {
// greedy
return (int)(std::max_element(logits, logits + V) - logits);
}
// 1. temperature (matches nano-vLLM: logits.float().div_(temperatures))
float inv_temp = 1.0f / temperature;
for (int i = 0; i < V; i++)
logits[i] *= inv_temp;
// 2. top_k: keep top K values, set rest to -inf
// nano-vLLM: topk(k) returns k-th largest as threshold, mask < threshold
if (top_k > 0 && top_k < V) {
std::vector<float> tmp(logits, logits + V);
std::nth_element(tmp.begin(), tmp.begin() + (top_k - 1), tmp.end(), std::greater<float>());
float threshold = tmp[top_k - 1];
for (int i = 0; i < V; i++)
if (logits[i] < threshold) logits[i] = -INFINITY;
}
// 3. top_p: nucleus filter on temp-scaled logits (matches nano-vLLM: softmax on scaled logits)
// nano-vLLM sorts ascending, cumsum, masks cumsum <= (1-p), keeps last element.
// Equivalent descending: mask tokens where cumsum_before >= top_p (shift-right).
if (top_p > 0.0f && top_p < 1.0f) {
std::vector<TokenProb> sorted(V);
for (int i = 0; i < V; i++) sorted[i] = {i, logits[i]};
std::sort(sorted.begin(), sorted.end(),
[](const TokenProb & a, const TokenProb & b) { return a.prob > b.prob; });
// softmax of temp-scaled logits for cumsum
float max_val = sorted[0].prob;
float sum = 0.0f;
std::vector<float> probs(V);
for (int i = 0; i < V; i++) {
probs[i] = expf(sorted[i].prob - max_val);
sum += probs[i];
}
float inv = 1.0f / sum;
// cumulative sum, test before accumulating (shift-right trick)
float cum = 0.0f;
for (int i = 0; i < V; i++) {
if (i > 0 && cum >= top_p) // i>0: always keep at least first token
logits[sorted[i].id] = -INFINITY;
cum += probs[i] * inv;
}
}
// 4. softmax -> multinomial (temperature already applied)
float max_val = -INFINITY;
for (int i = 0; i < V; i++)
if (logits[i] > max_val) max_val = logits[i];
float sum = 0.0f;
for (int i = 0; i < vocab_size; i++) {
for (int i = 0; i < V; i++) {
logits[i] = expf(logits[i] - max_val);
sum += logits[i];
}
float inv_sum = 1.0f / sum;
for (int i = 0; i < vocab_size; i++)
logits[i] *= inv_sum;
std::vector<TokenProb> candidates;
float threshold = 1.0f / (float)vocab_size * 0.01f;
for (int i = 0; i < vocab_size; i++)
if (logits[i] > threshold) candidates.push_back({i, logits[i]});
std::sort(candidates.begin(), candidates.end(),
[](const TokenProb & a, const TokenProb & b) { return a.prob > b.prob; });
float cum = 0.0f;
int n_keep = 0;
for (size_t i = 0; i < candidates.size(); i++) {
cum += candidates[i].prob;
n_keep = (int)i + 1;
if (cum >= top_p) break;
}
float renorm_sum = 0.0f;
for (int i = 0; i < n_keep; i++) renorm_sum += candidates[i].prob;
std::uniform_real_distribution<float> dist(0.0f, renorm_sum);
std::uniform_real_distribution<float> dist(0.0f, sum);
float r = dist(rng);
float acc = 0.0f;
for (int i = 0; i < n_keep; i++) {
acc += candidates[i].prob;
if (acc >= r) return candidates[i].id;
for (int i = 0; i < V; i++) {
acc += logits[i];
if (acc >= r) return i;
}
return candidates[0].id;
return 0;
}
//
@ -701,12 +735,12 @@ static std::string codes_to_string(const std::vector<int> & codes) {
// false for partial mode (user provided lyrics).
static void parse_phase1_into_aces(
const std::vector<std::string> & texts, const AcePrompt & base,
std::vector<AcePrompt> & aces, int base_seed,
std::vector<AcePrompt> & aces, long long base_seed,
const char * label, bool merge_lyrics) {
int N = (int)texts.size();
aces.resize(N);
for (int i = 0; i < N; i++) {
fprintf(stderr, "[%s Batch%d] seed=%d:\n%s\n", label, i, base_seed + i, texts[i].c_str());
fprintf(stderr, "[%s Batch%d] seed=%lld:\n%s\n", label, i, base_seed + i, texts[i].c_str());
AcePrompt parsed = {};
if (!parse_cot_and_lyrics(texts[i], &parsed))
fprintf(stderr, "WARNING: batch %d CoT parse incomplete\n", i);
@ -729,8 +763,8 @@ static void parse_phase1_into_aces(
static std::vector<std::string> generate_phase1_batch(
Qwen3LM * m, BPETokenizer * bpe,
const std::vector<int> & prompt_tokens,
int max_new_tokens, float temperature, float top_p,
int base_seed, int N,
int max_new_tokens, float temperature, float top_p, int top_k,
long long base_seed, int N,
MetadataFSM * fsm_template,
bool lyrics_mode,
float cfg_scale = 1.0f,
@ -776,7 +810,7 @@ static std::vector<std::string> generate_phase1_batch(
// Sample first token from shared prefill logits
for (int i = 0; i < N; i++) {
seqs[i].rng.seed(base_seed + i);
seqs[i].rng.seed((uint32_t)(base_seed + i));
if (fsm_template) seqs[i].fsm = *fsm_template;
seqs[i].codes_phase = false;
seqs[i].done = false;
@ -789,7 +823,7 @@ static std::vector<std::string> generate_phase1_batch(
if (fsm_template && fsm_template->enabled)
seqs[i].fsm.apply_mask(lg.data());
int tok = sample_top_p(lg.data(), V, temperature, top_p, seqs[i].rng);
int tok = sample_top_k_p(lg.data(), V, temperature, top_p, top_k, seqs[i].rng);
if (tok == TOKEN_IM_END) {
seqs[i].done = true;
@ -805,7 +839,7 @@ static std::vector<std::string> generate_phase1_batch(
seqs[i].last_token = tok;
}
// KV set arrays
// KV set arrays + merged CFG arrays
std::vector<int> cond_sets(N), uncond_sets(N);
for (int i = 0; i < N; i++) {
cond_sets[i] = i;
@ -818,6 +852,17 @@ static std::vector<std::string> generate_phase1_batch(
std::vector<float> logits_uncond(V * N);
std::vector<int> tokens(N);
// CFG: single forward with 2*N (cond + uncond)
int N2 = use_cfg ? 2 * N : N;
std::vector<int> tokens_2n(N2), sets_2n(N2);
std::vector<float> logits_2n((size_t)V * N2);
if (use_cfg) {
for (int i = 0; i < N; i++) {
sets_2n[i] = cond_sets[i];
sets_2n[N + i] = uncond_sets[i];
}
}
int n_active = N;
for (int i = 0; i < N; i++)
if (seqs[i].done) n_active--;
@ -826,9 +871,18 @@ static std::vector<std::string> generate_phase1_batch(
for (int i = 0; i < N; i++)
tokens[i] = seqs[i].last_token;
qw3lm_forward_batch(m, tokens.data(), cond_sets.data(), N, logits_cond.data());
if (use_cfg)
qw3lm_forward_batch(m, tokens.data(), uncond_sets.data(), N, logits_uncond.data());
if (use_cfg) {
// Single batched forward: cond[0..N-1] + uncond[N..2N-1]
for (int i = 0; i < N; i++) {
tokens_2n[i] = tokens[i];
tokens_2n[N + i] = tokens[i];
}
qw3lm_forward_batch(m, tokens_2n.data(), sets_2n.data(), N2, logits_2n.data());
memcpy(logits_cond.data(), logits_2n.data(), (size_t)V * N * sizeof(float));
memcpy(logits_uncond.data(), logits_2n.data() + (size_t)V * N, (size_t)V * N * sizeof(float));
} else {
qw3lm_forward_batch(m, tokens.data(), cond_sets.data(), N, logits_cond.data());
}
for (int i = 0; i < N; i++) {
if (seqs[i].done) continue;
@ -852,7 +906,7 @@ static std::vector<std::string> generate_phase1_batch(
if (v != TOKEN_IM_END) lc[v] = -1e9f;
}
int tok = sample_top_p(lc, V, temperature, top_p, seqs[i].rng);
int tok = sample_top_k_p(lc, V, temperature, top_p, top_k, seqs[i].rng);
if (tok == TOKEN_IM_END) {
seqs[i].done = true;
@ -887,7 +941,7 @@ static std::vector<std::string> generate_phase1_batch(
std::vector<std::string> results(N);
for (int i = 0; i < N; i++) {
results[i] = bpe_decode(*bpe, seqs[i].gen_tokens);
fprintf(stderr, "[Phase1 Batch%d] seed=%d, %zu tokens\n",
fprintf(stderr, "[Phase1 Batch%d] seed=%lld, %zu tokens\n",
i, base_seed + i, seqs[i].gen_tokens.size());
}
return results;
@ -900,7 +954,7 @@ static std::vector<std::string> generate_phase1_batch(
// Returns N code strings. Seeds = base_seed + 0, 1, ..., N-1.
static std::vector<std::string> run_phase2_batch(
Qwen3LM * m, BPETokenizer & bpe, const std::vector<AcePrompt> & aces,
float temperature, float top_p, int base_seed, int N,
float temperature, float top_p, int top_k, long long base_seed, int N,
float cfg_scale, const char * negative_prompt) {
int V = m->cfg.vocab_size;
@ -921,7 +975,7 @@ static std::vector<std::string> run_phase2_batch(
int mt = (int)(a.duration * 5) + 100;
if (mt > max_tokens) max_tokens = mt;
}
fprintf(stderr, "[Phase2] max_tokens: %d, CFG: %.2f, seeds: %d..%d\n",
fprintf(stderr, "[Phase2] max_tokens: %d, CFG: %.2f, seeds: %lld..%lld\n",
max_tokens, cfg_scale, base_seed, base_seed + N - 1);
// Reset all KV sets: cond [0..N-1], uncond [N..2N-1]
@ -974,7 +1028,7 @@ static std::vector<std::string> run_phase2_batch(
// Sample first token from per-element prefill logits (N different seeds)
for (int i = 0; i < N; i++) {
seqs[i].rng.seed(base_seed + i);
seqs[i].rng.seed((uint32_t)(base_seed + i));
seqs[i].done = false;
std::vector<float> lg(prefill_logits_vec[i]); // copy
@ -987,7 +1041,7 @@ static std::vector<std::string> run_phase2_batch(
for (int v = 0; v < AUDIO_CODE_BASE; v++)
if (v != TOKEN_IM_END) lg[v] = -1e9f;
int tok = sample_top_p(lg.data(), V, temperature, top_p, seqs[i].rng);
int tok = sample_top_k_p(lg.data(), V, temperature, top_p, top_k, seqs[i].rng);
seqs[i].last_token = tok;
if (tok == TOKEN_IM_END) {
@ -1010,6 +1064,17 @@ static std::vector<std::string> run_phase2_batch(
std::vector<float> logits_uncond(V * N);
std::vector<int> tokens(N);
// CFG: single forward with 2*N (cond + uncond)
int N2 = use_cfg ? 2 * N : N;
std::vector<int> tokens_2n(N2), sets_2n(N2);
std::vector<float> logits_2n((size_t)V * N2);
if (use_cfg) {
for (int i = 0; i < N; i++) {
sets_2n[i] = cond_sets[i];
sets_2n[N + i] = uncond_sets[i];
}
}
int n_active = N;
for (int i = 0; i < N; i++)
if (seqs[i].done) n_active--;
@ -1019,12 +1084,18 @@ static std::vector<std::string> run_phase2_batch(
for (int i = 0; i < N; i++)
tokens[i] = seqs[i].last_token;
// Batched forward: cond
qw3lm_forward_batch(m, tokens.data(), cond_sets.data(), N, logits_cond.data());
// Batched forward: uncond
if (use_cfg)
qw3lm_forward_batch(m, tokens.data(), uncond_sets.data(), N, logits_uncond.data());
if (use_cfg) {
// Single batched forward: cond[0..N-1] + uncond[N..2N-1]
for (int i = 0; i < N; i++) {
tokens_2n[i] = tokens[i];
tokens_2n[N + i] = tokens[i];
}
qw3lm_forward_batch(m, tokens_2n.data(), sets_2n.data(), N2, logits_2n.data());
memcpy(logits_cond.data(), logits_2n.data(), (size_t)V * N * sizeof(float));
memcpy(logits_uncond.data(), logits_2n.data() + (size_t)V * N, (size_t)V * N * sizeof(float));
} else {
qw3lm_forward_batch(m, tokens.data(), cond_sets.data(), N, logits_cond.data());
}
// Per-sequence: CFG combine + sample
for (int i = 0; i < N; i++) {
@ -1041,7 +1112,7 @@ static std::vector<std::string> run_phase2_batch(
for (int v = 0; v < AUDIO_CODE_BASE; v++)
if (v != TOKEN_IM_END) lc[v] = -1e9f;
int tok = sample_top_p(lc, V, temperature, top_p, seqs[i].rng);
int tok = sample_top_k_p(lc, V, temperature, top_p, top_k, seqs[i].rng);
seqs[i].last_token = tok;
if (tok == TOKEN_IM_END) {
@ -1069,7 +1140,7 @@ static std::vector<std::string> run_phase2_batch(
std::vector<std::string> results(N);
for (int i = 0; i < N; i++) {
results[i] = codes_to_string(seqs[i].audio_codes);
fprintf(stderr, "[Batch %d] seed=%d, %zu codes\n",
fprintf(stderr, "[Batch %d] seed=%lld, %zu codes\n",
i, base_seed + i, seqs[i].audio_codes.size());
}
return results;
@ -1084,214 +1155,22 @@ static void usage(const char * prog) {
"Usage: %s --request <json> --model <gguf> [options]\n"
"\n"
"Required:\n"
" --request <json> Request JSON (read, enriched, overwritten)\n"
" --model <gguf> Model GGUF file (from convert.py)\n"
" --request <json> Input request JSON\n"
" --model <gguf> Model GGUF file\n"
"\n"
"Infra:\n"
" --max-seq <N> KV cache size (default: 8192)\n"
"Batch:\n"
" --batch <N> Batch N sequences (default: 1)\n"
" --no-fsm Disable FSM constrained decoding\n"
"\n"
"Output naming: input.json -> input0.json, input1.json, ... (last digit = batch index)\n"
"\n"
"Debug:\n"
" --max-seq <N> KV cache size (default: 8192)\n"
" --no-fsm Disable FSM constrained decoding\n"
" --dump-logits <path> Dump prefill logits (binary f32)\n"
" --dump-tokens <path> Dump prompt token IDs (CSV)\n"
"\n", prog);
, prog);
}
//kcpp stuff
static Qwen3LM acestep_llm;
static BPETokenizer acestep_bpe;
static bool acestep_loaded = false;
bool load_acestep(std::string model_path)
{
acestep_loaded = false;
int max_seq = 8192;
const int batch_size = 1; //only bs 1 is allowed
if (!load_bpe_from_gguf(&acestep_bpe, model_path.c_str())) {
return false;
}
// Load model
int n_kv_sets = 2 * batch_size;
if (!qw3lm_load(&acestep_llm, model_path.c_str(), max_seq, n_kv_sets)) {
return false;
}
acestep_loaded = true;
return true;
}
std::string acestep_prepare_request(const music_generation_inputs inputs)
{
const int batch_size = 1;
bool use_fsm = true;
MetadataFSM fsm;
if (use_fsm) {
fsm.init(acestep_bpe, acestep_llm.cfg.vocab_size);
}
// Read request and set essentials
AceRequest req;
std::string injson = inputs.input_json;
if (!request_parse_from_str(&req, injson))
{
fprintf(stderr, "\nMusic JSON parse error\n");
return "";
}
int seed = req.seed;
if (seed <= 0 || seed==0xFFFFFFFF)
{
seed = (((uint32_t)time(NULL)) % 1000000u);
}
req.seed = seed;
// Generation params from request
float temperature = req.lm_temperature;
float top_p = req.lm_top_p;
float cfg_scale = req.lm_cfg_scale;
const char * neg_prompt = req.lm_negative_prompt.c_str();
// Copy request -> AcePrompt (internal LLM struct)
AcePrompt ace = {};
ace.caption = req.caption;
ace.lyrics = req.lyrics;
ace.duration = req.duration;
ace.bpm = req.bpm;
ace.keyscale = req.keyscale;
ace.timesignature = req.timesignature;
ace.vocal_language = req.vocal_language;
bool user_has_codes = !req.audio_codes.empty();
bool need_lm_codes = req.thinking && !user_has_codes;
bool is_simple = ace.lyrics.empty() &&
ace.bpm <= 0 && ace.duration <= 0 &&
ace.keyscale.empty() && ace.timesignature.empty();
std::vector<int> prompt;
std::vector<AcePrompt> aces; // populated by Phase 1 (simple or partial)
// Preprocessor: simple mode generates lyrics + metas from caption
if (is_simple) {
fprintf(stderr, "[Simple] Inspiration\n");
const char * sys =
"# Instruction\n"
"Expand the user's input into a more detailed"
" and specific musical description:\n";
std::string user_msg = ace.caption + "\n\ninstrumental: "
+ std::string(req.instrumental ? "true" : "false");
prompt = build_custom_prompt(acestep_bpe, sys, user_msg.c_str());
// FSM: reset then optionally force language (shared for both paths)
fsm.reset();
if (use_fsm && ace.vocal_language != "unknown" && !ace.vocal_language.empty())
fsm.force_language(acestep_bpe, ace.vocal_language);
// Phase 1: N lyrics + metadata generations (always batched, N=batch_size)
fprintf(stderr, "[Simple] %zu tokens, N=%d, seeds: %d..%d\n",
prompt.size(), batch_size, seed, seed + batch_size - 1);
auto phase1_texts = generate_phase1_batch(
&acestep_llm, &acestep_bpe, prompt, 2048, temperature, 1.0f,
seed, batch_size, use_fsm ? &fsm : nullptr, true);
parse_phase1_into_aces(phase1_texts, ace, aces, seed, "Simple", true);
for (int i = 0; i < batch_size; i++) qw3lm_reset_kv(&acestep_llm, i);
}
// Re-evaluate after possible simple enrichment
const AcePrompt & ace_ref = aces.empty() ? ace : aces[0];
bool has_all_metas = (ace_ref.bpm > 0 && ace_ref.duration > 0 &&
!ace_ref.keyscale.empty() && !ace_ref.timesignature.empty());
if (!has_all_metas) {
// Partial-metas: Phase 1 with CFG to fill missing fields
prompt = build_lm_prompt(acestep_bpe, ace);
std::vector<int> uncond;
if (cfg_scale > 1.0f)
uncond = build_lm_prompt_uncond(acestep_bpe, ace, neg_prompt);
fprintf(stderr, "[Partial] %zu tokens, CFG: %.2f, N=%d, seeds: %d..%d\n",
prompt.size(), cfg_scale, batch_size, seed, seed + batch_size - 1);
fsm.reset();
auto phase1_texts = generate_phase1_batch(
&acestep_llm, &acestep_bpe, prompt, 2048, temperature, top_p,
seed, batch_size, use_fsm ? &fsm : nullptr, false,
cfg_scale, uncond.empty() ? nullptr : &uncond, true);
parse_phase1_into_aces(phase1_texts, ace, aces, seed, "Partial", false);
for (int i = 0; i < 2 * batch_size; i++) qw3lm_reset_kv(&acestep_llm, i);
}
// Guarantee aces is populated (all-metas: single shared ace for prefill optimization)
if (aces.empty()) {
aces = { ace };
}
// Phase 2: generate audio codes (always batched, N=batch_size)
std::vector<std::string> batch_codes(batch_size);
if (need_lm_codes) {
batch_codes = run_phase2_batch(&acestep_llm, acestep_bpe, aces,
temperature, top_p, seed, batch_size, cfg_scale, neg_prompt);
} else {
fprintf(stderr, "[Skip] %s, no code generation\n",
user_has_codes ? "user codes present" : "thinking=false");
}
// only batch size 1 is allowed
AceRequest rr = req;
const AcePrompt & a = aces[0];
rr.caption = a.caption;
rr.lyrics = a.lyrics;
rr.bpm = a.bpm;
rr.duration = a.duration;
rr.keyscale = a.keyscale;
rr.timesignature = a.timesignature;
rr.vocal_language = a.vocal_language;
if (!batch_codes[0].empty()) rr.audio_codes = batch_codes[0];
rr.seed = seed;
//now convert to string
std::ostringstream oss;
oss << "{\n";
oss << " \"caption\": \"" << json_escape(rr.caption) << "\",\n";
oss << " \"lyrics\": \"" << json_escape(rr.lyrics) << "\",\n";
if (rr.instrumental) {
oss << " \"instrumental\": true,\n";
}
oss << " \"bpm\": " << rr.bpm << ",\n";
oss << " \"duration\": " << std::fixed << std::setprecision(1) << rr.duration << ",\n";
oss << " \"keyscale\": \"" << json_escape(rr.keyscale) << "\",\n";
oss << " \"timesignature\": \"" << json_escape(rr.timesignature) << "\",\n";
oss << " \"vocal_language\": \"" << json_escape(rr.vocal_language) << "\",\n";
oss << " \"task_type\": \"" << json_escape(rr.task_type) << "\",\n";
oss << " \"seed\": " << rr.seed << ",\n";
oss << " \"thinking\": " << (rr.thinking ? "true" : "false") << ",\n";
oss << " \"lm_temperature\": " << std::fixed << std::setprecision(2) << rr.lm_temperature << ",\n";
oss << " \"lm_cfg_scale\": " << std::fixed << std::setprecision(1) << rr.lm_cfg_scale << ",\n";
oss << " \"lm_top_p\": " << std::fixed << std::setprecision(2) << rr.lm_top_p << ",\n";
oss << " \"lm_negative_prompt\": \"" << json_escape(rr.lm_negative_prompt) << "\",\n";
oss << " \"inference_steps\": " << rr.inference_steps << ",\n";
oss << " \"guidance_scale\": " << std::fixed << std::setprecision(1) << rr.guidance_scale << ",\n";
oss << " \"shift\": " << std::fixed << std::setprecision(1) << rr.shift << ",\n";
oss << " \"audio_codes\": \"" << json_escape(rr.audio_codes) << "\"\n";
oss << "}\n";
std::string output_json = oss.str();
return output_json;
}
void unload_acestep()
{
qw3lm_free(&acestep_llm);
}
// int main(int argc, char ** argv) {
// const char * model_path = nullptr;
// const char * request_path = nullptr;
@ -1352,16 +1231,18 @@ void unload_acestep()
// }
// // Resolve seed
// int seed = req.seed;
// long long seed = req.seed;
// if (seed < 0) {
// std::random_device rd;
// seed = (int)(rd() & 0x7FFFFFFF);
// seed = (int64_t)rd() << 32 | rd();
// if (seed < 0) seed = -seed; // keep positive
// }
// req.seed = seed;
// // Generation params from request
// float temperature = req.lm_temperature;
// float top_p = req.lm_top_p;
// int top_k = req.lm_top_k;
// float cfg_scale = req.lm_cfg_scale;
// const char * neg_prompt = req.lm_negative_prompt.c_str();
@ -1420,11 +1301,11 @@ void unload_acestep()
// fsm.force_language(bpe, ace.vocal_language);
// // Phase 1: N lyrics + metadata generations (always batched, N=batch_size)
// fprintf(stderr, "[Simple] %zu tokens, N=%d, seeds: %d..%d\n",
// fprintf(stderr, "[Simple] %zu tokens, N=%d, seeds: %lld..%lld\n",
// prompt.size(), batch_size, seed, seed + batch_size - 1);
// auto phase1_texts = generate_phase1_batch(
// &model, &bpe, prompt, 2048, temperature, 1.0f,
// &model, &bpe, prompt, 2048, temperature, 1.0f, 0,
// seed, batch_size, use_fsm ? &fsm : nullptr, true);
// parse_phase1_into_aces(phase1_texts, ace, aces, seed, "Simple", true);
@ -1444,12 +1325,12 @@ void unload_acestep()
// if (cfg_scale > 1.0f)
// uncond = build_lm_prompt_uncond(bpe, ace, neg_prompt);
// fprintf(stderr, "[Partial] %zu tokens, CFG: %.2f, N=%d, seeds: %d..%d\n",
// fprintf(stderr, "[Partial] %zu tokens, CFG: %.2f, N=%d, seeds: %lld..%lld\n",
// prompt.size(), cfg_scale, batch_size, seed, seed + batch_size - 1);
// fsm.reset();
// auto phase1_texts = generate_phase1_batch(
// &model, &bpe, prompt, 2048, temperature, top_p,
// &model, &bpe, prompt, 2048, temperature, top_p, top_k,
// seed, batch_size, use_fsm ? &fsm : nullptr, false,
// cfg_scale, uncond.empty() ? nullptr : &uncond, true);
@ -1496,7 +1377,7 @@ void unload_acestep()
// std::vector<std::string> batch_codes(batch_size);
// if (need_lm_codes) {
// batch_codes = run_phase2_batch(&model, bpe, aces,
// temperature, top_p, seed, batch_size, cfg_scale, neg_prompt);
// temperature, top_p, top_k, seed, batch_size, cfg_scale, neg_prompt);
// } else {
// fprintf(stderr, "[Skip] %s, no code generation\n",
// user_has_codes ? "user codes present" : "thinking=false");
@ -1527,9 +1408,203 @@ void unload_acestep()
// }
// }
// fprintf(stderr, "[Ace-Qwen3] Load %.0f | Total %.0fms | seed=%d\n",
// fprintf(stderr, "[Ace-Qwen3] Load %.0f | Total %.0fms | seed=%lld\n",
// load_ms, t_total.ms(), seed);
// qw3lm_free(&model);
// return 0;
// }
//kcpp stuff
static Qwen3LM acestep_llm;
static BPETokenizer acestep_bpe;
static bool acestep_loaded = false;
bool load_acestep(std::string model_path)
{
acestep_loaded = false;
int max_seq = 8192;
const int batch_size = 1; //only bs 1 is allowed
if (!load_bpe_from_gguf(&acestep_bpe, model_path.c_str())) {
return false;
}
// Load model
int n_kv_sets = 2 * batch_size;
if (!qw3lm_load(&acestep_llm, model_path.c_str(), max_seq, n_kv_sets)) {
return false;
}
acestep_loaded = true;
return true;
}
std::string acestep_prepare_request(const music_generation_inputs inputs)
{
const int batch_size = 1;
bool use_fsm = true;
MetadataFSM fsm;
if (use_fsm) {
fsm.init(acestep_bpe, acestep_llm.cfg.vocab_size);
}
// Read request and set essentials
AceRequest req;
std::string injson = inputs.input_json;
if (!request_parse_from_str(&req, injson))
{
fprintf(stderr, "\nMusic JSON parse error\n");
return "";
}
int seed = req.seed;
if (seed <= 0 || seed==0xFFFFFFFF)
{
seed = (((uint32_t)time(NULL)) % 1000000u);
}
req.seed = seed;
// Generation params from request
float temperature = req.lm_temperature;
float top_p = req.lm_top_p;
int top_k = req.lm_top_k;
float cfg_scale = req.lm_cfg_scale;
const char * neg_prompt = req.lm_negative_prompt.c_str();
// Copy request -> AcePrompt (internal LLM struct)
AcePrompt ace = {};
ace.caption = req.caption;
ace.lyrics = req.lyrics;
ace.duration = req.duration;
ace.bpm = req.bpm;
ace.keyscale = req.keyscale;
ace.timesignature = req.timesignature;
ace.vocal_language = req.vocal_language;
bool user_has_codes = !req.audio_codes.empty();
bool need_lm_codes = req.thinking && !user_has_codes;
bool is_simple = ace.lyrics.empty() &&
ace.bpm <= 0 && ace.duration <= 0 &&
ace.keyscale.empty() && ace.timesignature.empty();
std::vector<int> prompt;
std::vector<AcePrompt> aces; // populated by Phase 1 (simple or partial)
// Preprocessor: simple mode generates lyrics + metas from caption
if (is_simple) {
fprintf(stderr, "[Simple] Inspiration\n");
const char * sys =
"# Instruction\n"
"Expand the user's input into a more detailed"
" and specific musical description:\n";
std::string user_msg = ace.caption + "\n\ninstrumental: "
+ std::string(req.instrumental ? "true" : "false");
prompt = build_custom_prompt(acestep_bpe, sys, user_msg.c_str());
// FSM: reset then optionally force language (shared for both paths)
fsm.reset();
if (use_fsm && ace.vocal_language != "unknown" && !ace.vocal_language.empty())
fsm.force_language(acestep_bpe, ace.vocal_language);
// Phase 1: N lyrics + metadata generations (always batched, N=batch_size)
fprintf(stderr, "[Simple] %zu tokens, N=%d, seeds: %lld..%lld\n",
prompt.size(), batch_size, seed, seed + batch_size - 1);
auto phase1_texts = generate_phase1_batch(
&acestep_llm, &acestep_bpe, prompt, 2048, temperature, 1.0f, 0,
seed, batch_size, use_fsm ? &fsm : nullptr, true);
parse_phase1_into_aces(phase1_texts, ace, aces, seed, "Simple", true);
for (int i = 0; i < batch_size; i++) qw3lm_reset_kv(&acestep_llm, i);
}
// Re-evaluate after possible simple enrichment
const AcePrompt & ace_ref = aces.empty() ? ace : aces[0];
bool has_all_metas = (ace_ref.bpm > 0 && ace_ref.duration > 0 &&
!ace_ref.keyscale.empty() && !ace_ref.timesignature.empty());
if (!has_all_metas) {
// Partial-metas: Phase 1 with CFG to fill missing fields
prompt = build_lm_prompt(acestep_bpe, ace);
std::vector<int> uncond;
if (cfg_scale > 1.0f)
uncond = build_lm_prompt_uncond(acestep_bpe, ace, neg_prompt);
fprintf(stderr, "[Partial] %zu tokens, CFG: %.2f, N=%d, seeds: %lld..%lld\n",
prompt.size(), cfg_scale, batch_size, seed, seed + batch_size - 1);
fsm.reset();
auto phase1_texts = generate_phase1_batch(
&acestep_llm, &acestep_bpe, prompt, 2048, temperature, top_p, top_k,
seed, batch_size, use_fsm ? &fsm : nullptr, false,
cfg_scale, uncond.empty() ? nullptr : &uncond, true);
parse_phase1_into_aces(phase1_texts, ace, aces, seed, "Partial", false);
for (int i = 0; i < 2 * batch_size; i++) qw3lm_reset_kv(&acestep_llm, i);
}
// Guarantee aces is populated (all-metas: single shared ace for prefill optimization)
if (aces.empty()) {
aces = { ace };
}
// Phase 2: generate audio codes (always batched, N=batch_size)
std::vector<std::string> batch_codes(batch_size);
if (need_lm_codes) {
batch_codes = run_phase2_batch(&acestep_llm, acestep_bpe, aces,
temperature, top_p, top_k, seed, batch_size, cfg_scale, neg_prompt);
} else {
fprintf(stderr, "[Skip] %s, no code generation\n",
user_has_codes ? "user codes present" : "thinking=false");
}
// only batch size 1 is allowed
AceRequest rr = req;
const AcePrompt & a = aces[0];
rr.caption = a.caption;
rr.lyrics = a.lyrics;
rr.bpm = a.bpm;
rr.duration = a.duration;
rr.keyscale = a.keyscale;
rr.timesignature = a.timesignature;
rr.vocal_language = a.vocal_language;
if (!batch_codes[0].empty()) rr.audio_codes = batch_codes[0];
rr.seed = seed;
//now convert to string
std::ostringstream oss;
oss << "{\n";
oss << " \"caption\": \"" << json_escape(rr.caption) << "\",\n";
oss << " \"lyrics\": \"" << json_escape(rr.lyrics) << "\",\n";
if (rr.instrumental) {
oss << " \"instrumental\": true,\n";
}
oss << " \"bpm\": " << rr.bpm << ",\n";
oss << " \"duration\": " << std::fixed << std::setprecision(1) << rr.duration << ",\n";
oss << " \"keyscale\": \"" << json_escape(rr.keyscale) << "\",\n";
oss << " \"timesignature\": \"" << json_escape(rr.timesignature) << "\",\n";
oss << " \"vocal_language\": \"" << json_escape(rr.vocal_language) << "\",\n";
oss << " \"task_type\": \"" << json_escape(rr.task_type) << "\",\n";
oss << " \"seed\": " << rr.seed << ",\n";
oss << " \"thinking\": " << (rr.thinking ? "true" : "false") << ",\n";
oss << " \"lm_temperature\": " << std::fixed << std::setprecision(2) << rr.lm_temperature << ",\n";
oss << " \"lm_cfg_scale\": " << std::fixed << std::setprecision(1) << rr.lm_cfg_scale << ",\n";
oss << " \"lm_top_p\": " << std::fixed << std::setprecision(2) << rr.lm_top_p << ",\n";
oss << " \"lm_negative_prompt\": \"" << json_escape(rr.lm_negative_prompt) << "\",\n";
oss << " \"inference_steps\": " << rr.inference_steps << ",\n";
oss << " \"guidance_scale\": " << std::fixed << std::setprecision(1) << rr.guidance_scale << ",\n";
oss << " \"shift\": " << std::fixed << std::setprecision(1) << rr.shift << ",\n";
oss << " \"audio_codes\": \"" << json_escape(rr.audio_codes) << "\"\n";
oss << "}\n";
std::string output_json = oss.str();
return output_json;
}
void unload_acestep()
{
qw3lm_free(&acestep_llm);
}

8
otherarch/acestep/cond.h
View file

@ -105,20 +105,22 @@ static bool cond_ggml_load(CondGGML * m, const char * gguf_path) {
m->lyric_embed_w = gf_load_tensor(&m->wctx, gf, "encoder.lyric_encoder.embed_tokens.weight");
m->lyric_embed_b = gf_load_tensor_f32(&m->wctx, gf, "encoder.lyric_encoder.embed_tokens.bias");
m->lyric_norm = gf_load_tensor_f32(&m->wctx, gf, "encoder.lyric_encoder.norm.weight");
fprintf(stderr, "[Load] LyricEncoder: %dL\n", m->lyric_cfg.n_layers);
for (int i = 0; i < m->lyric_cfg.n_layers; i++) {
char prefix[128];
snprintf(prefix, sizeof(prefix), "encoder.lyric_encoder.layers.%d", i);
qwen3_load_layer(&m->wctx, gf, &m->lyric_layers[i], prefix);
qwen3_load_layer(&m->wctx, gf, &m->lyric_layers[i], prefix, i);
}
// Timbre encoder
m->timbre_embed_w = gf_load_tensor(&m->wctx, gf, "encoder.timbre_encoder.embed_tokens.weight");
m->timbre_embed_b = gf_load_tensor_f32(&m->wctx, gf, "encoder.timbre_encoder.embed_tokens.bias");
m->timbre_norm = gf_load_tensor_f32(&m->wctx, gf, "encoder.timbre_encoder.norm.weight");
fprintf(stderr, "[Load] TimbreEncoder: %dL\n", m->timbre_cfg.n_layers);
for (int i = 0; i < m->timbre_cfg.n_layers; i++) {
char prefix[128];
snprintf(prefix, sizeof(prefix), "encoder.timbre_encoder.layers.%d", i);
qwen3_load_layer(&m->wctx, gf, &m->timbre_layers[i], prefix);
qwen3_load_layer(&m->wctx, gf, &m->timbre_layers[i], prefix, i);
}
// Text projector + null condition
@ -158,7 +160,7 @@ static void cond_ggml_forward(CondGGML * m,
int H = 2048;
bool has_timbre = (timbre_feats != NULL && S_ref > 0);
// Build graph
// Graph context (generous fixed allocation)
size_t ctx_size = 4096 * ggml_tensor_overhead() + ggml_graph_overhead();
struct ggml_init_params gp = { ctx_size, NULL, true };
struct ggml_context * ctx = ggml_init(gp);

24
otherarch/acestep/dit-vae.cpp
View file

@ -73,17 +73,16 @@ static void print_usage(const char * prog) {
"Required:\n"
" --request <json...> One or more request JSONs (from ace-qwen3 --request)\n"
" --text-encoder <gguf> Text encoder GGUF file\n"
" --dit <gguf> DiT GGUF file (from convert.py)\n"
" --dit <gguf> DiT GGUF file\n"
" --vae <gguf> VAE GGUF file\n\n"
"Batch:\n"
" --batch <n> DiT variations per request (default: 1, max 9)\n\n"
"Audio:\n"
" --noise-file <path> Load noise from bf16 file (Philox RNG dump, batch=1 only)\n\n"
"VAE tiling (memory control):\n"
" --vae-chunk <n> Latent frames per tile (default: 256)\n"
" --vae-overlap <n> Overlap frames per side (default: 64)\n\n"
" --batch <N> DiT variations per request (default: 1, max 9)\n\n"
"Output naming: input.json -> input0.wav, input1.wav, ... (last digit = batch index)\n\n"
"VAE tiling (memory control):\n"
" --vae-chunk <N> Latent frames per tile (default: 256)\n"
" --vae-overlap <N> Overlap frames per side (default: 64)\n\n"
"Debug:\n"
" --noise-file <path> Load noise from bf16 file (Philox RNG dump, batch=1 only)\n"
" --dump <dir> Dump intermediate tensors\n", prog);
}
@ -248,7 +247,7 @@ int main(int argc, char ** argv) {
const char * timesig = req.timesignature.empty() ? "N/A" : req.timesignature.c_str();
const char * language = req.vocal_language.empty() ? "en" : req.vocal_language.c_str();
float duration = req.duration > 0 ? req.duration : 120.0f;
int seed = req.seed;
long long seed = req.seed;
int num_steps = req.inference_steps > 0 ? req.inference_steps : 8;
float guidance_scale = req.guidance_scale > 0 ? req.guidance_scale : 7.0f;
float shift = req.shift > 0 ? req.shift : 1.0f;
@ -261,9 +260,10 @@ int main(int argc, char ** argv) {
if (seed < 0) {
std::random_device rd;
seed = (int)(rd() & 0x7FFFFFFF);
seed = (long long)rd() << 32 | rd();
if (seed < 0) seed = -seed;
}
fprintf(stderr, "[Pipeline] seed=%d, steps=%d, guidance=%.1f, shift=%.1f, duration=%.1fs\n",
fprintf(stderr, "[Pipeline] seed=%lld, steps=%d, guidance=%.1f, shift=%.1f, duration=%.1fs\n",
seed, num_steps, guidance_scale, shift, duration);
// Parse audio codes from request
@ -494,12 +494,12 @@ int main(int argc, char ** argv) {
} else {
// Generate N noise samples with seeds: seed, seed+1, ..., seed+N-1
for (int b = 0; b < batch_n; b++) {
std::mt19937 rng(seed + b);
std::mt19937 rng((uint32_t)(seed + b));
std::normal_distribution<float> normal(0.0f, 1.0f);
float * dst = noise.data() + b * Oc * T;
for (int i = 0; i < Oc * T; i++)
dst[i] = normal(rng);
fprintf(stderr, "[Context Batch%d] noise seed=%d\n", b, seed + b);
fprintf(stderr, "[Context Batch%d] noise seed=%lld\n", b, seed + b);
}
}

14
otherarch/acestep/dit.h
View file

@ -149,6 +149,10 @@ static struct ggml_tensor * dit_load_proj_in_w(
exit(1);
}
struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str());
if (!src) {
fprintf(stderr, "[GGUF] FATAL: meta tensor '%s' not found\n", name.c_str());
exit(1);
}
size_t offset = gguf_get_tensor_offset(gf.gguf, idx);
const void * raw = gf.mapping + gf.data_offset + offset;
@ -196,6 +200,10 @@ static struct ggml_tensor * dit_load_proj_out_w(
exit(1);
}
struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str());
if (!src) {
fprintf(stderr, "[GGUF] FATAL: meta tensor '%s' not found\n", name.c_str());
exit(1);
}
size_t offset = gguf_get_tensor_offset(gf.gguf, idx);
const void * raw = gf.mapping + gf.data_offset + offset;
@ -287,6 +295,8 @@ static bool dit_ggml_load(DiTGGML * m, const char * gguf_path, DiTGGMLConfig cfg
ly.sa_v_proj = gf_load_tensor(&m->wctx, gf, p + ".self_attn.v_proj.weight");
if (i == 0) fprintf(stderr, "[DiT] Self-attn: all separate (3 types differ)\n");
}
} else {
if (i == 0) fprintf(stderr, "[DiT] Self-attn: Q+K+V fused\n");
}
ly.sa_q_norm = gf_load_tensor_f32(&m->wctx, gf, p + ".self_attn.q_norm.weight");
ly.sa_k_norm = gf_load_tensor_f32(&m->wctx, gf, p + ".self_attn.k_norm.weight");
@ -311,6 +321,8 @@ static bool dit_ggml_load(DiTGGML * m, const char * gguf_path, DiTGGMLConfig cfg
ly.ca_v_proj = gf_load_tensor(&m->wctx, gf, p + ".cross_attn.v_proj.weight");
if (i == 0) fprintf(stderr, "[DiT] Cross-attn: all separate\n");
}
} else {
if (i == 0) fprintf(stderr, "[DiT] Cross-attn: Q+K+V fused\n");
}
ly.ca_q_norm = gf_load_tensor_f32(&m->wctx, gf, p + ".cross_attn.q_norm.weight");
ly.ca_k_norm = gf_load_tensor_f32(&m->wctx, gf, p + ".cross_attn.k_norm.weight");
@ -1085,7 +1097,7 @@ static void dit_ggml_generate(
fprintf(stderr, "[DiT] Batch N=%d, T=%d, S=%d, enc_S=%d\n", N, T, S, enc_S);
// Build graph once (shapes are constant across steps)
// Graph context (generous fixed allocation, shapes are constant across steps)
size_t ctx_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead_custom(8192, false);
std::vector<uint8_t> ctx_buf(ctx_size);

123
otherarch/acestep/qwen3-lm.h
View file

@ -233,7 +233,7 @@ static bool qw3lm_load(Qwen3LM * m, const char * gguf_path, int max_seq_len, int
for (int i = 0; i < c.n_layers; i++) {
char prefix[64];
snprintf(prefix, sizeof(prefix), "model.layers.%d", i);
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix);
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix, i);
}
wctx_alloc(&m->wctx, m->backend);
@ -278,10 +278,25 @@ static struct ggml_tensor * qw3lm_build_attn(
int Nkv = c.n_kv_heads;
int S = n_tokens;
// QKV projections
struct ggml_tensor * q = qwen3_linear(ctx, ly->q_proj, x); // [Nh*D, S]
struct ggml_tensor * k = qwen3_linear(ctx, ly->k_proj, x); // [Nkv*D, S]
struct ggml_tensor * v = qwen3_linear(ctx, ly->v_proj, x); // [Nkv*D, S]
// QKV projections (fused, partial, or separate)
struct ggml_tensor * q, * k, * v;
int q_dim = Nh * D;
int kv_dim = Nkv * D;
if (ly->qkv) {
struct ggml_tensor * qkv = qwen3_linear(ctx, ly->qkv, x);
q = ggml_cont(ctx, ggml_view_2d(ctx, qkv, q_dim, S, qkv->nb[1], 0));
k = ggml_cont(ctx, ggml_view_2d(ctx, qkv, kv_dim, S, qkv->nb[1], (size_t)q_dim * qkv->nb[0]));
v = ggml_cont(ctx, ggml_view_2d(ctx, qkv, kv_dim, S, qkv->nb[1], (size_t)(q_dim + kv_dim) * qkv->nb[0]));
} else if (ly->qk) {
struct ggml_tensor * qk = qwen3_linear(ctx, ly->qk, x);
q = ggml_cont(ctx, ggml_view_2d(ctx, qk, q_dim, S, qk->nb[1], 0));
k = ggml_cont(ctx, ggml_view_2d(ctx, qk, kv_dim, S, qk->nb[1], (size_t)q_dim * qk->nb[0]));
v = qwen3_linear(ctx, ly->v_proj, x);
} else {
q = qwen3_linear(ctx, ly->q_proj, x);
k = qwen3_linear(ctx, ly->k_proj, x);
v = qwen3_linear(ctx, ly->v_proj, x);
}
// Reshape to heads: [X*D, S] -> [D, X, S]
q = ggml_reshape_3d(ctx, q, D, Nh, S);
@ -351,7 +366,7 @@ static void qw3lm_forward(Qwen3LM * m, const int * token_ids, int n_tokens,
return;
}
// Graph context
// Graph context (generous fixed allocation)
size_t ctx_size = (size_t)16384 * ggml_tensor_overhead() + ggml_graph_overhead();
struct ggml_init_params gp = { ctx_size, NULL, true };
struct ggml_context * ctx = ggml_init(gp);
@ -490,76 +505,9 @@ static void qw3lm_forward_batch(Qwen3LM * m, const int * token_ids,
}
}
// Exact tensor count for context allocation
// qwen3_f32(t) creates 0 tensors if t is F32, 1 (ggml_cast) if not.
// Count conditionally based on actual weight types.
//
// GLOBAL (2):
// embed_out (new_tensor_2d) 1
// positions (new_tensor_1d) 1
//
// PER LAYER fixed (17):
// qwen3_rms_norm(input_layernorm):
// ggml_rms_norm + ggml_mul 2
// q_proj, k_proj, v_proj (qwen3_linear = mul_mat) 3
// reshape_3d (q, k, v) 3
// ggml_rms_norm(q) + ggml_mul(q, q_norm) 2
// ggml_rms_norm(k) + ggml_mul(k, k_norm) 2
// ggml_rope_ext (q, k) 2
// ggml_cont (q, k, v) 3
//
// PER LAYER * N (16 each):
// view_3d (qi, ki, vi) 3
// permute (qi, ki, vi) 3
// cont (ki, vi) 2
// view_3d (k_dst, v_dst) 2
// cpy (ki, vi) 2
// view_3d (k_full, v_full) 2
// flash_attn_ext 1
// reshape_2d 1
//
// CONCATS: N-1 (first element reuses reshape_2d)
//
// PER LAYER post (9):
// qwen3_linear(o_proj) 1
// ggml_add (residual) 1
// qwen3_rms_norm(post_attn_layernorm):
// ggml_rms_norm + ggml_mul 2
// qwen3_build_mlp:
// gate_proj, up_proj (mul_mat) 2
// ggml_swiglu_split 1
// down_proj (mul_mat) 1
// ggml_add (residual) 1
//
// PER LAYER conditional casts (0 to 4):
// qwen3_f32(input_layernorm) 0 or 1
// qwen3_f32(q_norm) 0 or 1
// qwen3_f32(k_norm) 0 or 1
// qwen3_f32(post_attn_norm) 0 or 1
//
// POST LAYERS (3):
// qwen3_rms_norm(final_norm): ggml_rms_norm + ggml_mul 2
// ggml_mul_mat (lm_head) 1
//
// POST conditional cast (0 or 1):
// qwen3_f32(final_norm) 0 or 1
//
// TOTAL = 5 + n_layers * (25 + 17*N + casts_per_layer) + 3 + cast_final
// = 8 + n_layers * (25 + 17*N + casts_per_layer) + cast_final
int casts_per_layer = 0;
if (m->layers[0].input_layernorm->type != GGML_TYPE_F32) casts_per_layer++;
if (m->layers[0].q_norm->type != GGML_TYPE_F32) casts_per_layer++;
if (m->layers[0].k_norm->type != GGML_TYPE_F32) casts_per_layer++;
if (m->layers[0].post_attn_layernorm->type != GGML_TYPE_F32) casts_per_layer++;
int cast_final = (m->final_norm->type != GGML_TYPE_F32) ? 1 : 0;
size_t n_tensors = 8
+ (size_t)c.n_layers * (25 + 17 * N + casts_per_layer)
+ cast_final;
size_t est = n_tensors * ggml_tensor_overhead()
+ ggml_graph_overhead_custom(16384, false);
struct ggml_init_params gp = { est, NULL, true };
// Graph context (generous fixed allocation, ~6 MB, negligible vs model weights)
size_t ctx_size = (size_t)16384 * ggml_tensor_overhead() + ggml_graph_overhead_custom(16384, false);
struct ggml_init_params gp = { ctx_size, NULL, true };
struct ggml_context * ctx = ggml_init(gp);
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
@ -581,10 +529,25 @@ static void qw3lm_forward_batch(Qwen3LM * m, const int * token_ids,
// Pre-attention norm [H, N]
struct ggml_tensor * norm = qwen3_rms_norm(ctx, hidden, ly->input_layernorm, c.rms_norm_eps);
// Batched QKV projections (weights read once for N tokens)
struct ggml_tensor * q = qwen3_linear(ctx, ly->q_proj, norm); // [Nh*D, N]
struct ggml_tensor * k = qwen3_linear(ctx, ly->k_proj, norm); // [Nkv*D, N]
struct ggml_tensor * v = qwen3_linear(ctx, ly->v_proj, norm); // [Nkv*D, N]
// Batched QKV projections (fused, partial, or separate)
struct ggml_tensor * q, * k, * v;
int q_dim = Nh * D;
int kv_dim = Nkv * D;
if (ly->qkv) {
struct ggml_tensor * qkv = qwen3_linear(ctx, ly->qkv, norm);
q = ggml_cont(ctx, ggml_view_2d(ctx, qkv, q_dim, N, qkv->nb[1], 0));
k = ggml_cont(ctx, ggml_view_2d(ctx, qkv, kv_dim, N, qkv->nb[1], (size_t)q_dim * qkv->nb[0]));
v = ggml_cont(ctx, ggml_view_2d(ctx, qkv, kv_dim, N, qkv->nb[1], (size_t)(q_dim + kv_dim) * qkv->nb[0]));
} else if (ly->qk) {
struct ggml_tensor * qk = qwen3_linear(ctx, ly->qk, norm);
q = ggml_cont(ctx, ggml_view_2d(ctx, qk, q_dim, N, qk->nb[1], 0));
k = ggml_cont(ctx, ggml_view_2d(ctx, qk, kv_dim, N, qk->nb[1], (size_t)q_dim * qk->nb[0]));
v = qwen3_linear(ctx, ly->v_proj, norm);
} else {
q = qwen3_linear(ctx, ly->q_proj, norm);
k = qwen3_linear(ctx, ly->k_proj, norm);
v = qwen3_linear(ctx, ly->v_proj, norm);
}
// Reshape to heads: [D, Heads, N]
q = ggml_reshape_3d(ctx, q, D, Nh, N);

112
otherarch/acestep/qwen3.h
View file

@ -40,15 +40,22 @@ struct Qwen3Config {
struct Qwen3Layer {
struct ggml_tensor * input_layernorm; // [H]
struct ggml_tensor * post_attn_layernorm; // [H]
struct ggml_tensor * q_proj; // [H, Nh*D] ggml = [Nh*D, H] PyTorch
struct ggml_tensor * k_proj; // [H, Nkv*D]
struct ggml_tensor * v_proj; // [H, Nkv*D]
struct ggml_tensor * o_proj; // [Nh*D, H]
struct ggml_tensor * q_norm; // [D]
struct ggml_tensor * k_norm; // [D]
struct ggml_tensor * gate_proj; // [H, FFN]
struct ggml_tensor * up_proj; // [H, FFN]
struct ggml_tensor * down_proj; // [FFN, H]
// Attention (fused or separate, same pattern as DiT)
struct ggml_tensor * qkv; // [H, (Nh+2*Nkv)*D] full fused (or NULL)
struct ggml_tensor * qk; // [H, (Nh+Nkv)*D] Q+K fused (or NULL)
struct ggml_tensor * q_proj; // [H, Nh*D] (NULL when fused)
struct ggml_tensor * k_proj; // [H, Nkv*D] (NULL when fused)
struct ggml_tensor * v_proj; // [H, Nkv*D] (NULL when QKV fused)
struct ggml_tensor * o_proj; // [Nh*D, H]
struct ggml_tensor * q_norm; // [D]
struct ggml_tensor * k_norm; // [D]
// MLP (fused or separate)
struct ggml_tensor * gate_up; // [H, 2*FFN] fused (or NULL)
struct ggml_tensor * gate_proj; // [H, FFN] (NULL when fused)
struct ggml_tensor * up_proj; // [H, FFN] (NULL when fused)
struct ggml_tensor * down_proj; // [FFN, H]
};
// Standalone model (text encoder)
@ -113,10 +120,25 @@ static struct ggml_tensor * qwen3_build_self_attn(
int Nh = c.n_heads;
int Nkv = c.n_kv_heads;
// 1) Q/K/V projections
struct ggml_tensor * q = qwen3_linear(ctx, ly->q_proj, x); // [Nh*D, S]
struct ggml_tensor * k = qwen3_linear(ctx, ly->k_proj, x); // [Nkv*D, S]
struct ggml_tensor * v = qwen3_linear(ctx, ly->v_proj, x); // [Nkv*D, S]
// 1) Q/K/V projections (fused, partial, or separate)
struct ggml_tensor * q, * k, * v;
int q_dim = Nh * D;
int kv_dim = Nkv * D;
if (ly->qkv) {
struct ggml_tensor * qkv = qwen3_linear(ctx, ly->qkv, x);
q = ggml_cont(ctx, ggml_view_2d(ctx, qkv, q_dim, S, qkv->nb[1], 0));
k = ggml_cont(ctx, ggml_view_2d(ctx, qkv, kv_dim, S, qkv->nb[1], (size_t)q_dim * qkv->nb[0]));
v = ggml_cont(ctx, ggml_view_2d(ctx, qkv, kv_dim, S, qkv->nb[1], (size_t)(q_dim + kv_dim) * qkv->nb[0]));
} else if (ly->qk) {
struct ggml_tensor * qk = qwen3_linear(ctx, ly->qk, x);
q = ggml_cont(ctx, ggml_view_2d(ctx, qk, q_dim, S, qk->nb[1], 0));
k = ggml_cont(ctx, ggml_view_2d(ctx, qk, kv_dim, S, qk->nb[1], (size_t)q_dim * qk->nb[0]));
v = qwen3_linear(ctx, ly->v_proj, x);
} else {
q = qwen3_linear(ctx, ly->q_proj, x);
k = qwen3_linear(ctx, ly->k_proj, x);
v = qwen3_linear(ctx, ly->v_proj, x);
}
// 2) Reshape to heads: [X*D, S] -> [D, X, S]
q = ggml_reshape_3d(ctx, q, D, Nh, S);
@ -154,16 +176,22 @@ static struct ggml_tensor * qwen3_build_self_attn(
return qwen3_linear(ctx, ly->o_proj, attn);
}
// MLP: SwiGLU
// MLP: SwiGLU (fused gate+up or separate)
static struct ggml_tensor * qwen3_build_mlp(
struct ggml_context * ctx,
Qwen3Layer * ly,
struct ggml_tensor * x, // [H, S]
int S) {
(void)S;
struct ggml_tensor * gate = qwen3_linear(ctx, ly->gate_proj, x);
struct ggml_tensor * up = qwen3_linear(ctx, ly->up_proj, x);
struct ggml_tensor * ff = ggml_swiglu_split(ctx, gate, up);
struct ggml_tensor * ff;
if (ly->gate_up) {
struct ggml_tensor * gu = qwen3_linear(ctx, ly->gate_up, x);
ff = ggml_swiglu(ctx, gu);
} else {
struct ggml_tensor * gate = qwen3_linear(ctx, ly->gate_proj, x);
struct ggml_tensor * up = qwen3_linear(ctx, ly->up_proj, x);
ff = ggml_swiglu_split(ctx, gate, up);
}
return qwen3_linear(ctx, ly->down_proj, ff);
}
@ -209,17 +237,47 @@ static struct ggml_tensor * qwen3_build_layers(
// Loading
static void qwen3_load_layer(WeightCtx * wctx, const GGUFModel & gf,
Qwen3Layer * ly, const std::string & prefix) {
Qwen3Layer * ly, const std::string & prefix, int layer_idx = -1) {
ly->input_layernorm = gf_load_tensor_f32(wctx, gf, prefix + ".input_layernorm.weight");
ly->post_attn_layernorm = gf_load_tensor_f32(wctx, gf, prefix + ".post_attention_layernorm.weight");
ly->q_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.q_proj.weight");
ly->k_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.k_proj.weight");
ly->v_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.v_proj.weight");
// Attention: try Q+K+V fused, then Q+K partial, then separate
ly->qkv = gf_load_qkv_fused(wctx, gf,
prefix + ".self_attn.q_proj.weight",
prefix + ".self_attn.k_proj.weight",
prefix + ".self_attn.v_proj.weight");
if (!ly->qkv) {
ly->qk = gf_load_pair_fused(wctx, gf,
prefix + ".self_attn.q_proj.weight",
prefix + ".self_attn.k_proj.weight");
if (ly->qk) {
ly->v_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.v_proj.weight");
if (layer_idx == 0) fprintf(stderr, "[Qwen3] Attn: Q+K fused, V separate\n");
} else {
ly->q_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.q_proj.weight");
ly->k_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.k_proj.weight");
ly->v_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.v_proj.weight");
if (layer_idx == 0) fprintf(stderr, "[Qwen3] Attn: all separate\n");
}
} else {
if (layer_idx == 0) fprintf(stderr, "[Qwen3] Attn: Q+K+V fused\n");
}
ly->o_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.o_proj.weight");
ly->q_norm = gf_load_tensor_f32(wctx, gf, prefix + ".self_attn.q_norm.weight");
ly->k_norm = gf_load_tensor_f32(wctx, gf, prefix + ".self_attn.k_norm.weight");
ly->gate_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.gate_proj.weight");
ly->up_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.up_proj.weight");
// MLP: try gate+up fused, then separate
ly->gate_up = gf_load_pair_fused(wctx, gf,
prefix + ".mlp.gate_proj.weight",
prefix + ".mlp.up_proj.weight");
if (ly->gate_up) {
if (layer_idx == 0) fprintf(stderr, "[Qwen3] MLP: gate+up fused\n");
} else {
ly->gate_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.gate_proj.weight");
ly->up_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.up_proj.weight");
if (layer_idx == 0) fprintf(stderr, "[Qwen3] MLP: gate+up separate\n");
}
ly->down_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.down_proj.weight");
}
@ -259,10 +317,12 @@ static bool qwen3_load_text_encoder(Qwen3GGML * m, const char * gguf_path) {
m->embed_tokens = gf_load_tensor(&m->wctx, gf, "embed_tokens.weight");
m->final_norm = gf_load_tensor_f32(&m->wctx, gf, "norm.weight");
fprintf(stderr, "[Load] TextEncoder: %dL, H=%d, Nh=%d/%d\n",
m->cfg.n_layers, m->cfg.hidden_size, m->cfg.n_heads, m->cfg.n_kv_heads);
for (int i = 0; i < m->cfg.n_layers; i++) {
char prefix[64];
snprintf(prefix, sizeof(prefix), "layers.%d", i);
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix);
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix, i);
}
if (!wctx_alloc(&m->wctx, m->backend)) {
@ -271,8 +331,6 @@ static bool qwen3_load_text_encoder(Qwen3GGML * m, const char * gguf_path) {
}
gf_close(&gf);
fprintf(stderr, "[Load] TextEncoder: %dL, H=%d, Nh=%d/%d\n",
m->cfg.n_layers, m->cfg.hidden_size, m->cfg.n_heads, m->cfg.n_kv_heads);
return true;
}
@ -284,7 +342,7 @@ static void qwen3_forward(Qwen3GGML * m, const int * token_ids, int S, float * o
const Qwen3Config & c = m->cfg;
int H = c.hidden_size;
// Graph context
// Graph context (generous fixed allocation)
size_t ctx_size = 2048 * ggml_tensor_overhead() + ggml_graph_overhead();
struct ggml_init_params gp = { ctx_size, NULL, true };
struct ggml_context * ctx = ggml_init(gp);

15
otherarch/acestep/request.cpp
View file

@ -26,6 +26,7 @@ void request_init(AceRequest * r) {
r->lm_temperature = 0.85f;
r->lm_cfg_scale = 2.0f;
r->lm_top_p = 0.9f;
r->lm_top_k = 0;
r->lm_negative_prompt = "NO USER INPUT";
r->audio_codes = "";
r->inference_steps = 8;
@ -233,13 +234,14 @@ bool request_parse_from_str(AceRequest * r, std::string json) {
// ints
else if (k == "bpm") r->bpm = atoi(v.c_str());
else if (k == "seed") r->seed = atoi(v.c_str());
else if (k == "seed") r->seed = strtoll(v.c_str(), nullptr, 10);
// floats
else if (k == "duration") r->duration = (float)atof(v.c_str());
else if (k == "lm_temperature") r->lm_temperature = (float)atof(v.c_str());
else if (k == "lm_cfg_scale") r->lm_cfg_scale = (float)atof(v.c_str());
else if (k == "lm_top_p") r->lm_top_p = (float)atof(v.c_str());
else if (k == "lm_top_k") r->lm_top_k = atoi(v.c_str());
else if (k == "inference_steps") r->inference_steps = atoi(v.c_str());
else if (k == "guidance_scale") r->guidance_scale = (float)atof(v.c_str());
else if (k == "shift") r->shift = (float)atof(v.c_str());
@ -272,11 +274,12 @@ bool request_write(const AceRequest * r, const char * path) {
fprintf(f, " \"timesignature\": \"%s\",\n", json_escape(r->timesignature).c_str());
fprintf(f, " \"vocal_language\": \"%s\",\n", json_escape(r->vocal_language).c_str());
fprintf(f, " \"task_type\": \"%s\",\n", json_escape(r->task_type).c_str());
fprintf(f, " \"seed\": %d,\n", r->seed);
fprintf(f, " \"seed\": %lld,\n", (long long)r->seed);
fprintf(f, " \"thinking\": %s,\n", r->thinking ? "true" : "false");
fprintf(f, " \"lm_temperature\": %.2f,\n", r->lm_temperature);
fprintf(f, " \"lm_cfg_scale\": %.1f,\n", r->lm_cfg_scale);
fprintf(f, " \"lm_top_p\": %.2f,\n", r->lm_top_p);
fprintf(f, " \"lm_top_k\": %d,\n", r->lm_top_k);
fprintf(f, " \"lm_negative_prompt\": \"%s\",\n", json_escape(r->lm_negative_prompt).c_str());
fprintf(f, " \"inference_steps\": %d,\n", r->inference_steps);
fprintf(f, " \"guidance_scale\": %.1f,\n", r->guidance_scale);
@ -291,16 +294,16 @@ bool request_write(const AceRequest * r, const char * path) {
}
void request_dump(const AceRequest * r, FILE * f) {
fprintf(f, "[Request] task=%s thinking=%s seed=%d\n",
r->task_type.c_str(), r->thinking ? "true" : "false", r->seed);
fprintf(f, "[Request] task=%s thinking=%s seed=%lld\n",
r->task_type.c_str(), r->thinking ? "true" : "false", (long long)r->seed);
fprintf(f, " caption: %.60s%s\n",
r->caption.c_str(), r->caption.size() > 60 ? "..." : "");
fprintf(f, " lyrics: %zu bytes\n", r->lyrics.size());
fprintf(f, " bpm=%d dur=%.0f key=%s ts=%s lang=%s\n",
r->bpm, r->duration, r->keyscale.c_str(),
r->timesignature.c_str(), r->vocal_language.c_str());
fprintf(f, " lm: temp=%.2f cfg=%.1f top_p=%.2f\n",
r->lm_temperature, r->lm_cfg_scale, r->lm_top_p);
fprintf(f, " lm: temp=%.2f cfg=%.1f top_p=%.2f top_k=%d\n",
r->lm_temperature, r->lm_cfg_scale, r->lm_top_p, r->lm_top_k);
fprintf(f, " dit: steps=%d guidance=%.1f shift=%.1f\n",
r->inference_steps, r->guidance_scale, r->shift);
fprintf(f, " audio_codes: %s\n",

3
otherarch/acestep/request.h
View file

@ -24,13 +24,14 @@ struct AceRequest {
// generation
std::string task_type; // "text2music"
int seed; // -1 = random
int64_t seed; // -1 = random
// LM control
bool thinking; // true
float lm_temperature; // 0.85
float lm_cfg_scale; // 2.0
float lm_top_p; // 0.9
int lm_top_k; // 0 = disabled (matches Python None)
std::string lm_negative_prompt; // "NO USER INPUT"
// codes (Python-compatible string: "3101,11837,27514,...")

1
otherarch/acestep/tokenizer.h
View file

@ -131,6 +131,7 @@ static int detok_ggml_decode(DetokGGML * m, const int * codes, int T_5Hz,
// Step 2: build ggml graph for one token
// input [6] -> project_out [2048] -> embed_tokens [2048]
// -> broadcast + special_tokens [2048, 5] -> 2L encoder -> norm -> proj_out [64, 5]
// Graph context (generous fixed allocation)
size_t ctx_size = ggml_tensor_overhead() * 512 + ggml_graph_overhead_custom(4096, false);
std::vector<uint8_t> ctx_buf(ctx_size);
struct ggml_init_params p = { ctx_size, ctx_buf.data(), true };

5
otherarch/acestep/vae.h
View file

@ -2,7 +2,7 @@
//
// Architecture: conv1(64->2048,k=7) -> 5xblock(snake+convT+3xresunit) -> snake+conv2(128->2,k=7)
// ResUnit(ch, dil): skip=x -> snake->conv(k=7,dil)->snake->conv(k=1)->+skip
// Snake: x + sin^2(e^a * x) / e^b
// Snake: x + sin^2(e^a * x) * (1/e^b)
// Weight norm fused at load: w = g*v/||v||
// Upsample: 10x6x4x4x2 = 1920x
@ -327,13 +327,14 @@ static int vae_ggml_compute(
int T_latent, // window length to decode
int win_start = 0) { // offset into latent
// Build graph only when T_latent changes
// Build graph only when T_latent changes (cached for tiled decode reuse)
if (m->graph_T != T_latent) {
if (m->graph_ctx) {
ggml_backend_sched_reset(m->sched);
ggml_free(m->graph_ctx);
free(m->graph_buf);
}
// Graph context (generous fixed allocation)
size_t ctx_size = ggml_tensor_overhead() * 1024 + ggml_graph_overhead_custom(8192, false);
m->graph_buf = (uint8_t *)malloc(ctx_size);
struct ggml_init_params p = { ctx_size, m->graph_buf, true };