diff --git a/docs/abot_world.md b/docs/abot_world.md index 78d0f6a0e..d6a2dc556 100644 --- a/docs/abot_world.md +++ b/docs/abot_world.md @@ -29,7 +29,16 @@ a rolling attention window and a 4-step distilled sampler. image (text-only, `first_frame_mask = 0`) are supported at the format level but gated: the distilled checkpoint cannot bootstrap a coherent first frame from noise, so front-ends should require an image until a T2V-capable - checkpoint ships. + checkpoint ships. Scene creation zeroes every prompt-embedding row past the + last real token (mirroring the reference text encoder's `u[v:] = 0`); the + encoder's attention mask only masks attention *inside* the encoder, so + without this a pack carries live pad-token embeddings in all 512 context + rows and the walk washes out from the first generated block. +- **Scene-pack diagnostics**: loading a pack logs `scene pack: prompt rows N + live / M`, where `N` is the prompt's real token count and `M` the fixed + 512-row context. `N == M` warns loudly — the pack's padding was not zeroed + (see above), so the walk is conditioned on pad embeddings and output will be + washed out; regenerate the pack with an engine that zeroes the padding. **Not supported:** the batch `generate_image()`/`generate_video()` paths — those are one-shot, whereas ABot needs the stateful causal session. Both batch diff --git a/src/abot_world.hpp b/src/abot_world.hpp index b24152826..1b7124ef9 100644 --- a/src/abot_world.hpp +++ b/src/abot_world.hpp @@ -101,6 +101,8 @@ struct AbotScenePack { sd::Tensor ref_latents; // {32, 32, C, K} (T=1 squeezed) sd::Tensor ref_mask; // {K} int ref_slots = 0; + // prompt rows before the zeroed padding (diagnostic; 0 = not computed) + int64_t text_rows_live = 0; // false = text-only scene: block-0 frame 0 is generated from noise // instead of being pinned to first_frame_latents bool has_first_frame = true; @@ -233,6 +235,41 @@ struct AbotScenePack { return false; } prompt_embeds.resize({sh[2], sh[1], 1, 1}); + // Real prompt rows: the producer zeroes everything past the last token + // (the reference's `u[v:] = 0`), so trailing all-zero rows are padding. + // Reported because a pack whose padding is NOT zeroed conditions the + // walk on pad-token embeddings and degrades generation from the first + // block - a silent failure that is otherwise only visible in the + // output pixels. + { + const int64_t emb = prompt_embeds.shape()[0]; + const int64_t rows = prompt_embeds.shape()[1]; + const float* pd = prompt_embeds.data(); + int64_t live = 0; + for (int64_t r = rows - 1; r >= 0; r--) { + bool nonzero = false; + for (int64_t i = 0; i < emb; i++) { + if (pd[r * emb + i] != 0.0f) { + nonzero = true; + break; + } + } + if (nonzero) { + live = r + 1; + break; + } + } + text_rows_live = live; + if (live == rows) { + LOG_WARN( + "scene pack: all %lld prompt rows are non-zero - padding is not zeroed, so the " + "walk is conditioned on pad embeddings (expect washed-out output; the pack " + "producer is missing the reference's zero-padding step)", + (long long)rows); + } else { + LOG_INFO("scene pack: prompt rows %lld live / %lld", (long long)live, (long long)rows); + } + } if (!fetch("first_frame_latents", first_frame_latents, sh) || // [1,1,C,H,W] !expect("first_frame_latents", sh, 5, {0, 1})) { return false; @@ -542,6 +579,34 @@ struct AbotWorldRunner : public GGMLRunner { // 8-key vector to 8 channels, repeat_interleaves x4 -> 32 channels constant // over HxW, then PixelUnshuffle(16): output channel c corresponds to input // channel c / (16*16) -> value = key[(c / 256) / 4]. + // The action planes are the largest host-built input (~51 MB per frame) and + // are identical for every graph of a block: same keys held, same frame + // count, same latent size. Refilling them per graph put ~150 MB of memset + // on the critical path before each denoise step, so keep the filled buffer + // and rebuild it only when the key really changes. + sd::Tensor act_planes; + uint8_t act_planes_mask = 0; + int act_planes_frames = -1; + int64_t act_planes_w = 0; + int64_t act_planes_h = 0; + + sd::Tensor& action_planes(uint8_t action_mask, int F_cur, int64_t lat_w, int64_t lat_h, int c_unsh) { + if (act_planes_frames == F_cur && act_planes_mask == action_mask && + act_planes_w == lat_w && act_planes_h == lat_h && !act_planes.empty()) { + return act_planes; + } + act_planes = sd::zeros({lat_w, lat_h, c_unsh, F_cur}); + for (int f = 0; f < F_cur; f++) { + fill_act_plane(act_planes.data() + static_cast(f) * c_unsh * lat_w * lat_h, + action_mask, static_cast(lat_w), static_cast(lat_h), c_unsh); + } + act_planes_mask = action_mask; + act_planes_frames = F_cur; + act_planes_w = lat_w; + act_planes_h = lat_h; + return act_planes; + } + void fill_act_plane(float* dst, uint8_t action_mask, int w_in, int h_in, int c_unsh) { for (int c = 0; c < c_unsh; c++) { int key = (c / (cfg.act_downscale_factor * cfg.act_downscale_factor)) / 4; @@ -641,7 +706,11 @@ struct AbotWorldRunner : public GGMLRunner { return gf; }; - auto result = GGMLRunner::compute(get_graph, n_threads, false); + // keep the compute buffer and its graph allocator across steps: the walk + // runs 5-6 graphs per block forever, and the defaults would free and + // re-reserve multi-GB of VRAM on every one of them (ggml re-reserves + // automatically when the graph shape changes between denoise and append) + auto result = GGMLRunner::compute(get_graph, n_threads, false, false, false); if (!result.has_value()) { return {}; } @@ -747,11 +816,7 @@ struct AbotWorldRunner : public GGMLRunner { memcpy(dst, src, static_cast(lat_w) * lat_h * sizeof(float)); } } - sd::Tensor act({lat_w, lat_h, c_unsh, F_cur}); - for (int f = 0; f < F_cur; f++) { - fill_act_plane(act.data() + static_cast(f) * c_unsh * lat_w * lat_h, - action_mask, static_cast(lat_w), static_cast(lat_h), c_unsh); - } + sd::Tensor& act = action_planes(action_mask, F_cur, lat_w, lat_h, c_unsh); sd::Tensor tvec({F_cur + 1}); for (int f = 0; f < F_cur; f++) { tvec.data()[f] = frame_timesteps[static_cast(f)]; @@ -886,7 +951,11 @@ struct AbotWorldRunner : public GGMLRunner { return gf; }; - auto result = GGMLRunner::compute(get_graph, n_threads, false); + // keep the compute buffer and its graph allocator across steps: the walk + // runs 5-6 graphs per block forever, and the defaults would free and + // re-reserve multi-GB of VRAM on every one of them (ggml re-reserves + // automatically when the graph shape changes between denoise and append) + auto result = GGMLRunner::compute(get_graph, n_threads, false, false, false); if (prof) { const int64_t prof_t2 = ggml_time_ms(); const char* mode_s = mode == KvMode::INIT_CAPTURE ? "init" : mode == KvMode::APPEND ? "append" : "denoise"; @@ -1132,6 +1201,10 @@ class AbotWalkSession { true, VERSION_ABOT_WORLD, model_manager); + // Direct convolution for the pixel decoder: im2col+GEMM materializes a + // large intermediate per conv, and the decoder is all small 3x3 convs. + // Measured 148 -> 68 ms per block decode, bit-identical output. + tae->set_conv2d_direct_enabled(true); if (!model_manager->register_runner_params("ABot-World DiT", *runner, "model.diffusion_model", diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 39157869d..ce77efc7c 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -216,12 +216,20 @@ namespace WAN { const float scale = 1.0f / sqrtf(static_cast(head_dim)); ggml_tensor* kq = ggml_mul_mat(gctx, K, q_r); // {T_kv, n_token, n_head*N} - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - kq = ggml_scale_inplace(gctx, kq, scale); - if (mask != nullptr) { - kq = ggml_add_inplace(gctx, kq, mask); // {T_kv, n_token} broadcast over heads - } - kq = ggml_soft_max_inplace(gctx, kq); + // No GGML_PREC_F32 here. On a coopmat2 device an F32 x F32 matmul + // asking for F32 precision has no cooperative-matrix pipeline to + // fall back on (pipeline_matmul_f32_cm1 is only built in the + // coopmat1 branch), so it lands on a scalar shader with no tensor + // cores; the default precision converts both operands to F16 and + // uses matmul_f16_f16acc_cm2 instead. Measured 241 -> 207 ms per + // denoise step on an RTX 5090. Scores feed a softmax, so the + // accumulator's dynamic range is not the limit here. + // One fused pass instead of scale -> add -> soft_max: the scores are + // the largest tensor in the walk graph (T_kv x n_token x n_head), so + // each separate elementwise pass costs a full read+write of it. The + // mask is {T_kv, n_token} and broadcasts over heads, which + // ggml_soft_max_ext supports on CPU, CUDA and Vulkan. + kq = ggml_soft_max_ext(gctx, kq, mask, scale, 0.0f); ggml_tensor* kqv = ggml_mul_mat(gctx, V, kq); // {d_head, n_token, n_head*N} kqv = ggml_reshape_4d(gctx, kqv, head_dim, n_token, num_heads, N); diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 5005af2e1..ab72a0335 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -7565,6 +7565,22 @@ bool sd_abot_scene_create(const sd_abot_scene_params_t* p) { sd::Tensor::from_vector(std::get<0>(tokens)), sd::Tensor::from_vector(std::get<2>(tokens))); if (prompt.empty()) { LOG_ERROR("sd_abot_scene_create: prompt encoding failed"); return false; } + // The reference WanTextEncoder zeroes embeddings past the real tokens + // (`u[v:] = 0`). The encoder's attention mask only affects attention + // inside the encoder, not its output rows, so without this the pack + // carries live pad-token embeddings in all 512 context rows and the + // walk collapses into blur from the first generated block. The + // tokenizer mask is additive: 0.0 = real token, -inf = padding. + { + const std::vector& attn_mask = std::get<2>(tokens); + const int64_t emb_dim = prompt.shape()[0]; + float* pd = prompt.data(); + for (size_t i = 0; i < attn_mask.size(); i++) { + if (attn_mask[i] != 0.0f) { + memset(pd + static_cast(emb_dim) * i, 0, static_cast(emb_dim) * sizeof(float)); + } + } + } if (prompt.shape().size() == 2) prompt = prompt.unsqueeze(2); sd::Tensor first; if (has_image) {