From 5311997581c2c05ca434990cd620b3cb73a7c75f Mon Sep 17 00:00:00 2001 From: Concedo <39025047+LostRuins@users.noreply.github.com> Date: Mon, 23 Feb 2026 23:01:10 +0800 Subject: [PATCH] updated ace step cpp --- otherarch/acestep/ace-qwen3.cpp | 595 ++++++++++++++++++-------------- otherarch/acestep/cond.h | 8 +- otherarch/acestep/dit-vae.cpp | 24 +- otherarch/acestep/dit.h | 14 +- otherarch/acestep/qwen3-lm.h | 123 +++---- otherarch/acestep/qwen3.h | 112 ++++-- otherarch/acestep/request.cpp | 15 +- otherarch/acestep/request.h | 3 +- otherarch/acestep/tokenizer.h | 1 + otherarch/acestep/vae.h | 5 +- 10 files changed, 508 insertions(+), 392 deletions(-) diff --git a/otherarch/acestep/ace-qwen3.cpp b/otherarch/acestep/ace-qwen3.cpp index 5c566c09f..b9e1643d7 100644 --- a/otherarch/acestep/ace-qwen3.cpp +++ b/otherarch/acestep/ace-qwen3.cpp @@ -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 tmp(logits, logits + V); + std::nth_element(tmp.begin(), tmp.begin() + (top_k - 1), tmp.end(), std::greater()); + 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 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 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 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 dist(0.0f, renorm_sum); + + std::uniform_real_distribution 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 & codes) { // false for partial mode (user provided lyrics). static void parse_phase1_into_aces( const std::vector & texts, const AcePrompt & base, - std::vector & aces, int base_seed, + std::vector & 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 generate_phase1_batch( Qwen3LM * m, BPETokenizer * bpe, const std::vector & 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 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 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 generate_phase1_batch( seqs[i].last_token = tok; } - // KV set arrays + // KV set arrays + merged CFG arrays std::vector cond_sets(N), uncond_sets(N); for (int i = 0; i < N; i++) { cond_sets[i] = i; @@ -818,6 +852,17 @@ static std::vector generate_phase1_batch( std::vector logits_uncond(V * N); std::vector tokens(N); + // CFG: single forward with 2*N (cond + uncond) + int N2 = use_cfg ? 2 * N : N; + std::vector tokens_2n(N2), sets_2n(N2); + std::vector 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 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 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 generate_phase1_batch( std::vector 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 generate_phase1_batch( // Returns N code strings. Seeds = base_seed + 0, 1, ..., N-1. static std::vector run_phase2_batch( Qwen3LM * m, BPETokenizer & bpe, const std::vector & 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 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 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 lg(prefill_logits_vec[i]); // copy @@ -987,7 +1041,7 @@ static std::vector 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 run_phase2_batch( std::vector logits_uncond(V * N); std::vector tokens(N); + // CFG: single forward with 2*N (cond + uncond) + int N2 = use_cfg ? 2 * N : N; + std::vector tokens_2n(N2), sets_2n(N2); + std::vector 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 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 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 run_phase2_batch( std::vector 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 --model [options]\n" "\n" "Required:\n" - " --request Request JSON (read, enriched, overwritten)\n" - " --model Model GGUF file (from convert.py)\n" + " --request Input request JSON\n" + " --model Model GGUF file\n" "\n" - "Infra:\n" - " --max-seq KV cache size (default: 8192)\n" + "Batch:\n" " --batch 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 KV cache size (default: 8192)\n" + " --no-fsm Disable FSM constrained decoding\n" " --dump-logits Dump prefill logits (binary f32)\n" " --dump-tokens 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 prompt; - std::vector 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 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 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 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 prompt; + std::vector 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 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 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); +} diff --git a/otherarch/acestep/cond.h b/otherarch/acestep/cond.h index 53e8e8307..9b76131e8 100644 --- a/otherarch/acestep/cond.h +++ b/otherarch/acestep/cond.h @@ -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); diff --git a/otherarch/acestep/dit-vae.cpp b/otherarch/acestep/dit-vae.cpp index a360f15ae..15d47cccb 100644 --- a/otherarch/acestep/dit-vae.cpp +++ b/otherarch/acestep/dit-vae.cpp @@ -73,17 +73,16 @@ static void print_usage(const char * prog) { "Required:\n" " --request One or more request JSONs (from ace-qwen3 --request)\n" " --text-encoder Text encoder GGUF file\n" - " --dit DiT GGUF file (from convert.py)\n" + " --dit DiT GGUF file\n" " --vae VAE GGUF file\n\n" "Batch:\n" - " --batch DiT variations per request (default: 1, max 9)\n\n" - "Audio:\n" - " --noise-file Load noise from bf16 file (Philox RNG dump, batch=1 only)\n\n" - "VAE tiling (memory control):\n" - " --vae-chunk Latent frames per tile (default: 256)\n" - " --vae-overlap Overlap frames per side (default: 64)\n\n" + " --batch 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 Latent frames per tile (default: 256)\n" + " --vae-overlap Overlap frames per side (default: 64)\n\n" "Debug:\n" + " --noise-file Load noise from bf16 file (Philox RNG dump, batch=1 only)\n" " --dump 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 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); } } diff --git a/otherarch/acestep/dit.h b/otherarch/acestep/dit.h index 64c714d90..4210aebb6 100644 --- a/otherarch/acestep/dit.h +++ b/otherarch/acestep/dit.h @@ -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 ctx_buf(ctx_size); diff --git a/otherarch/acestep/qwen3-lm.h b/otherarch/acestep/qwen3-lm.h index 445db4702..6cb3f402c 100644 --- a/otherarch/acestep/qwen3-lm.h +++ b/otherarch/acestep/qwen3-lm.h @@ -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); diff --git a/otherarch/acestep/qwen3.h b/otherarch/acestep/qwen3.h index 4cc895b6d..0b7c6b725 100644 --- a/otherarch/acestep/qwen3.h +++ b/otherarch/acestep/qwen3.h @@ -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); diff --git a/otherarch/acestep/request.cpp b/otherarch/acestep/request.cpp index 9c5b430e0..7a57ca0e8 100644 --- a/otherarch/acestep/request.cpp +++ b/otherarch/acestep/request.cpp @@ -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", diff --git a/otherarch/acestep/request.h b/otherarch/acestep/request.h index 1e9d5d723..9dae39ecf 100644 --- a/otherarch/acestep/request.h +++ b/otherarch/acestep/request.h @@ -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,...") diff --git a/otherarch/acestep/tokenizer.h b/otherarch/acestep/tokenizer.h index eb69d6a1b..640f0d203 100644 --- a/otherarch/acestep/tokenizer.h +++ b/otherarch/acestep/tokenizer.h @@ -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 ctx_buf(ctx_size); struct ggml_init_params p = { ctx_size, ctx_buf.data(), true }; diff --git a/otherarch/acestep/vae.h b/otherarch/acestep/vae.h index 310fc0a39..51ce0b9ad 100644 --- a/otherarch/acestep/vae.h +++ b/otherarch/acestep/vae.h @@ -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 };