memory: bound dense 1024 decode and host postprocess allocations

include/sparse.h

@@ -23,6 +23,13 @@ ggml_tensor* sparse_submconv(ggml_context* c, const Model& m, const std::string&
 // SparseConvNeXtBlock3d: conv -> rowLN(affine,1e-6) -> Linear(C,4C)->SiLU->Linear(4C,C) -> + input.
 ggml_tensor* sparse_convnext(ggml_context* c, const Model& m, const std::string& prefix,
                              ggml_tensor* feats, ggml_tensor* nbr, int N);
+// Host-streamed ConvNeXt block for very large sparse levels. The output is identical in
+// layout to sparse_convnext, but each voxel chunk is a separate short-lived GPU graph and
+// only referenced neighbour feature columns are uploaded. This avoids the growing concat
+// chain / full sparse-table residency that can make stage-3 shape decode request 20-30+ GB.
+std::vector<float> sparse_convnext_streamed(const Model& m, const std::string& prefix,
+                                             const std::vector<float>& feats_in, int C,
+                                             const std::vector<int32_t>& nbr, int N);
 
 // SparseResBlockC2S3d up-block (channel->spatial ×2). Host-orchestrated (subdiv readback).
 // feats_in: [Cin*N] channel-major; returns new feats [Cout*M] + new coords (res ×2).

src/dit.cpp

@@ -4,6 +4,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <cstdint>
 #include <cstdlib>
 #include <string>
 
@@ -65,8 +66,8 @@ static T* apply_rope(ggml_context* c, T* x, T* cos, T* sin) {
 // An FA padding mask [Lk_pad, Lq] (F16): 0 for real keys (< Lk_real), a large negative for the
 // zero-padded tail. WITHOUT it, ggml's CUDA FlashAttention folds the (zero) padded keys into the
 // softmax; on the >=1024-token HR flow that path NaNs a subset of queries (props, <1024 tokens, dodge
-// it). WITH it the kernel masks/skips the padded KV tiles -> correct softmax, no NaN. Built once per
-// flow (same N every block) and threaded into every attention; -30000 (not -inf) so 0*mask can't NaN.
+// it). WITH it the kernel masks/skips the padded KV tiles -> correct softmax, no NaN. Large flows build
+// this per QUERY CHUNK (not full Lq) so the mask stays bounded; -30000 (not -inf) so 0*mask can't NaN.
 static T* build_pad_mask(ggml_context* c, int64_t Lk_real, int64_t Lq) {
     const int64_t KQ = 256;
     const int64_t Lk_pad = ((Lk_real + KQ - 1) / KQ) * KQ;
@@ -138,10 +139,52 @@ static T* sdpa(ggml_context* c, T* q, T* k, T* v, int d_model, T* mask = nullptr
         // ignored and the zero-padded keys are diluting the softmax (exp(0-rowmax) is only
         // negligible when rowmax >> 0), which shrinks every output toward zero.
         static const bool fa_nomask = std::getenv("TRELLIS_FA_NOMASK") != nullptr;
-        T* out = ggml_flash_attn_ext(c, qf, kf, vf, fa_nomask ? nullptr : mask, scale, 0.0f, 0.0f);  // [hd, nh, Lq]
-        if (!fa_fast) ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
-        out = ggml_scale(c, out, 1.0f / V_SCALE);
-        return ggml_reshape_2d(c, out, d_model, out->ne[2]);   // [d_model, Lq]
+
+        // IMPORTANT: the padding mask depends only on KEY index, but ggml FlashAttention
+        // requires it expanded to [Lk_pad, Lq_pad]. At 37,017 self-attention tokens that one
+        // F16 tensor is exactly 2,751,037,440 bytes -- the allocation seen in the real failure
+        // log. FlashAttention itself is tiled, so do the same for the mask: attention is
+        // independent per query, therefore query chunks are mathematically identical while
+        // reducing the mask from O(Lq*Lk) residency to O(nq*Lk).
+        const int64_t hd = qf->ne[0], Lq = qf->ne[1], nh = qf->ne[2];
+        const int64_t Lk_real = k->ne[2];
+        const int64_t Lk_pad = ((Lk_real + KQ_STRIDE - 1) / KQ_STRIDE) * KQ_STRIDE;
+        static constexpr int64_t kDefaultFaMaskChunkBytes = 256ll * 1024 * 1024;
+        int64_t mask_budget = kDefaultFaMaskChunkBytes;
+        if (const char* e = getenv("TRELLIS_FA_MASK_CHUNK_MB")) {
+            const int64_t mb = atoll(e);
+            if (mb > 0) mask_budget = mb * 1024 * 1024;
+        }
+        int64_t nq = Lq;
+        if (!fa_nomask) {
+            const int64_t bytes_per_q = std::max<int64_t>(1, Lk_pad * (int64_t)sizeof(uint16_t));
+            nq = std::max<int64_t>(1, mask_budget / bytes_per_q);
+            // FA's mask reader works in 64-query tiles. Round ordinary chunks down to a
+            // multiple of 64 so only the final chunk needs padding.
+            if (nq >= 64 && nq < Lq) nq = (nq / 64) * 64;
+            if (nq > Lq) nq = Lq;
+        }
+
+        T* out_all = nullptr;
+        for (int64_t q0 = 0; q0 < Lq; q0 += nq) {
+            const int64_t n = std::min<int64_t>(nq, Lq - q0);
+            T* qc = (n == Lq)
+                ? qf
+                : ggml_cont(c, ggml_view_3d(c, qf, hd, n, nh,
+                                             qf->nb[1], qf->nb[2], (size_t)q0 * qf->nb[1]));
+            T* cmask = nullptr;
+            if (!fa_nomask) {
+                // A caller-supplied mask is only safe to reuse when this attention was not
+                // query-chunked. build_dit_dense no longer creates the giant full mask.
+                cmask = (mask && n == Lq) ? mask : build_pad_mask(c, Lk_real, n);
+            }
+            T* o = ggml_flash_attn_ext(c, qc, kf, vf, cmask, scale, 0.0f, 0.0f); // [hd,nh,n]
+            if (!fa_fast) ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32);
+            o = ggml_scale(c, o, 1.0f / V_SCALE);
+            o = ggml_reshape_2d(c, o, d_model, o->ne[2]);          // [d_model,n]
+            out_all = out_all ? ggml_concat(c, out_all, o, 1) : o;
+        }
+        return out_all;                                           // [d_model,Lq]
     }
     // Exact SDPA, chunked over QUERIES. The whole reason FA exists here is the [Lk, Lq, nh]
     // score matrix -- at the HR flow that is 15104*15006*12*4 = 10.9 TB, so it cannot be
@@ -275,13 +318,11 @@ ggml_tensor* build_dit_dense(ggml_context* c, const Model& m, const DiTParams& p
     T* mod = lin(c, m, "adaLN_modulation.1", ggml_silu(c, te));// [6*d_model]
     keep("t_emb_mod", mod);
 
-    // Padding masks for the two attentions, built ONCE (token counts are fixed across blocks) and
-    // shared by every block — they tell the CUDA FlashAttention to exclude the zero-padded key tiles
-    // (else the >=1024-token flow NaNs a subset of queries). self: Lk=L (the latent); cross: Lk=Lc.
-    T* self_mask  = build_pad_mask(c, h0->ne[1], h0->ne[1]);
-    T* cross_mask = build_pad_mask(c, cond->ne[1], h0->ne[1]);
+    // FlashAttention padding masks used to be built once at full [Lk_pad,Lq_pad] size.
+    // That becomes a 2.75 GB single tensor at 37,017 tokens. sdpa() now creates an identical
+    // mask per QUERY CHUNK, so there is deliberately no full-flow mask tensor here.
     for (int i = 0; i < p.n_blocks; ++i) {
-        h = block(c, m, i, h, mod, cond, cos, sin, p, inter, self_mask, cross_mask);
+        h = block(c, m, i, h, mod, cond, cos, sin, p, inter, nullptr, nullptr);
         if (i == 0) keep("after_block0", h);
         if (i == 1) keep("after_block1", h);
         if (i == p.n_blocks - 1) keep("after_block29", h);

src/flow_runner.cpp

@@ -38,7 +38,12 @@ DitRunner::DitRunner(const Model& m, const DiTParams& p, int N, int n_cond,
                      const std::vector<float>& rcos, const std::vector<float>& rsin)
     : m_(m), p_(p), N_(N), Lc_(n_cond) {
     const int half = p_.head_dim / 2;
-    size_t meta = ggml_tensor_overhead() * 16384 + ggml_graph_overhead_custom(32768, false) + (1 << 20);
+    // Query-chunked FlashAttention creates several small mask/FA nodes per attention instead
+    // of one giant quadratic mask. Give the graph generous HOST metadata headroom; this does
+    // not reserve an equivalent amount of VRAM.
+    static constexpr size_t kDitGraphNodes = 65536;
+    size_t meta = ggml_tensor_overhead() * kDitGraphNodes +
+                  ggml_graph_overhead_custom(kDitGraphNodes, false) + (1 << 20);
     ctx_ = ggml_init({ meta, nullptr, true });
     gh0_  = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, p_.in_ch, N_);   ggml_set_input(gh0_);
     gtf_  = ggml_new_tensor_1d(ctx_, GGML_TYPE_F32, 256);            ggml_set_input(gtf_);
@@ -47,12 +52,21 @@ DitRunner::DitRunner(const Model& m, const DiTParams& p, int N, int n_cond,
     gsin_ = ggml_new_tensor_4d(ctx_, GGML_TYPE_F32, 1, half, 1, N_); ggml_set_input(gsin_);
     dbg_nan_ = std::getenv("TRELLIS_DBG_NAN") != nullptr;
     gout_ = build_dit_dense(ctx_, m_, p_, gh0_, gtf_, gcond_, gcos_, gsin_, dbg_nan_ ? &inter_ : nullptr);
-    g_ = ggml_new_graph_custom(ctx_, 32768, false);
+    g_ = ggml_new_graph_custom(ctx_, kDitGraphNodes, false);
     ggml_build_forward_expand(g_, gout_);
     ggml_set_output(gout_);
     if (dbg_nan_) for (auto& [nm, t] : inter_) { ggml_build_forward_expand(g_, t); ggml_set_output(t); }
     alloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m_.backend));
-    if (!ggml_gallocr_alloc_graph(alloc_, g_)) throw std::runtime_error("DitRunner: alloc failed");
+    if (!ggml_gallocr_alloc_graph(alloc_, g_)) {
+        // Constructors that throw do not run DitRunner::~DitRunner(), so explicitly release
+        // partial allocator/context state before propagating the failure.
+        ggml_gallocr_free(alloc_); alloc_ = nullptr;
+        ggml_free(ctx_); ctx_ = nullptr;
+        throw std::runtime_error("DitRunner: alloc failed");
+    }
+    if (getenv("TRELLIS_DBG_ALLOC"))
+        fprintf(stderr, "      [dit-alloc] N=%d nodes=%d gallocr buffer = %.2f GB\n",
+                N_, ggml_graph_n_nodes(g_), ggml_gallocr_get_buffer_size(alloc_, 0) / 1e9);
     rcos_ = rcos; rsin_ = rsin;   // keep; re-upload each forward (gallocr reuses input buffers across runs)
 }
 

src/remesh_dc.cpp

@@ -76,7 +76,20 @@ Mesh remesh_narrow_band_dc(const float* iverts, int64_t iV, const int32_t* iface
     {
         const int F = (int)iF;
         std::vector<std::vector<uint64_t>> parts;
-        const int nt = std::max(1u, std::thread::hardware_concurrency());
+        const int hw = (int)std::max(1u, std::thread::hardware_concurrency());
+        // Each worker previously allocated a full res^3 candidate bitset.  At
+        // res=1024 that is 128 MiB per worker, so a 32-thread CPU consumed ~4 GiB
+        // here before any geometry/BVH memory.  Cap only the replication memory;
+        // the OR result and therefore the remesh are bit-for-bit equivalent.
+        const size_t bytes_per_part = cand.size() * sizeof(uint64_t);
+        const size_t parts_budget = (size_t)1024 * 1024 * 1024; // 1 GiB
+        const int mem_workers = bytes_per_part ? (int)std::max<size_t>(1, parts_budget / bytes_per_part) : hw;
+        const int nt = std::max(1, std::min(hw, mem_workers));
+        if (nt < hw) {
+            printf("  [remesh-mem] candidate bitset %.1f MiB/worker, workers %d->%d\n",
+                   bytes_per_part / (1024.0*1024.0), hw, nt);
+            fflush(stdout);
+        }
         parts.assign(nt, {});
         std::vector<std::thread> ts;
         const int chunk = (F + nt - 1) / nt;

src/shape_decoder.cpp

@@ -6,6 +6,7 @@
 #include "ggml-alloc.h"
 
 #include <algorithm>
+#include <cstdlib>
 #include <string>
 #include <stdexcept>
 
@@ -37,12 +38,20 @@ static std::vector<float> run1(const Model& m, ggml_context* c, T* out,
     ggml_cgraph* g = ggml_new_graph_custom(c, kGraphNodes, false);
     ggml_build_forward_expand(g, out);
     ggml_gallocr_t a = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m.backend));
-    if (!ggml_gallocr_alloc_graph(a, g)) throw std::runtime_error("shape_dec alloc");
+    if (!ggml_gallocr_alloc_graph(a, g)) {
+        // A failed reserve can still leave allocator/backend bookkeeping behind. Free it before
+        // the caller retries the same stage through the streamed ConvNeXt fallback.
+        ggml_gallocr_free(a);
+        throw std::runtime_error("shape_dec alloc");
+    }
     if (getenv("TRELLIS_DBG_ALLOC"))
         fprintf(stderr, "      [stage-alloc] nodes=%d  gallocr buffer = %.2f GB\n",
                 ggml_graph_n_nodes(g), ggml_gallocr_get_buffer_size(a, 0) / 1e9);
     for (auto& [t, d] : ins) ggml_backend_tensor_set(t, d, 0, ggml_nbytes(t));
-    if (ggml_backend_graph_compute(m.backend, g) != GGML_STATUS_SUCCESS) throw std::runtime_error("shape_dec compute");
+    if (ggml_backend_graph_compute(m.backend, g) != GGML_STATUS_SUCCESS) {
+        ggml_gallocr_free(a);
+        throw std::runtime_error("shape_dec compute");
+    }
     mem_probe("run1 computed (pre-readback)");
     std::vector<float> r = tensor_to_f32(out);
     ggml_gallocr_free(a);
@@ -94,15 +103,54 @@ static std::vector<float> decode_unet(const Model& m, const std::vector<float>&
         if (getenv("TRELLIS_DBG_MEM"))
             fprintf(stderr, "    == stage %d: C=%d nblk=%d N=%d | host h=%.2f GB nbr=%.2f GB\n",
                     si, st.C, st.nblk, N, (double)h.size() * 4 / 1e9, (double)nbr.size() * 4 / 1e9);
-        {   // ConvNeXt stage
-            ggml_context* c = mkctx();
-            T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, st.C, N); ggml_set_input(gh);
-            T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, N, 27);   ggml_set_input(gn);
-            T* x = gh;
-            for (int j = 0; j < st.nblk; ++j)
-                x = sparse_convnext(c, m, std::string("blocks.") + st.s + "." + std::to_string(j), x, gn, N);
-            h = run1(m, c, x, { {gh, h.data()}, {gn, nbr.data()} });
-            ggml_free(c);
+        {   // ConvNeXt stage. Once the full stage input itself reaches ~2 GiB, the old
+            // sparse_convnext concat chain multiplies that residency several-fold (the real
+            // stage-3 failure reserved 31.36 GB). Go straight to true chunk-local streaming
+            // for that case instead of deliberately provoking an OOM first. Smaller stages
+            // keep the fast whole-stage graph, with an allocation-failure fallback as a safety net.
+            size_t stream_at = 2ull * 1024 * 1024 * 1024;
+            if (const char* e = getenv("TRELLIS_CONVNEXT_STREAM_GB")) {
+                const double gb = atof(e);
+                if (gb > 0.0) stream_at = (size_t)(gb * 1024.0 * 1024.0 * 1024.0);
+            }
+            const size_t feature_bytes = (size_t)st.C * (size_t)N * sizeof(float);
+            const bool force_stream = getenv("TRELLIS_CONVNEXT_STREAM_ALL") != nullptr;
+            bool stream_stage = force_stream || feature_bytes >= stream_at;
+
+            if (!stream_stage) {
+                // The host vector `h` is only assigned after a successful run1(), so it remains
+                // the stage input if the fast-path allocation fails.
+                ggml_context* c = mkctx();
+                try {
+                    T* gh = ggml_new_tensor_2d(c, GGML_TYPE_F32, st.C, N); ggml_set_input(gh);
+                    T* gn = ggml_new_tensor_2d(c, GGML_TYPE_I32, N, 27);   ggml_set_input(gn);
+                    T* x = gh;
+                    for (int j = 0; j < st.nblk; ++j)
+                        x = sparse_convnext(c, m, std::string("blocks.") + st.s + "." + std::to_string(j), x, gn, N);
+                    std::vector<float> fast = run1(m, c, x, { {gh, h.data()}, {gn, nbr.data()} });
+                    ggml_free(c); c = nullptr;
+                    h = std::move(fast);
+                } catch (const std::runtime_error& e) {
+                    if (c) ggml_free(c);
+                    if (std::string(e.what()) != "shape_dec alloc") throw;
+                    stream_stage = true;
+                    fprintf(stderr,
+                            "      [shape-dec-stream] stage %d C=%d N=%d whole-stage alloc failed; "
+                            "retrying streamed\n", si, st.C, N);
+                }
+            } else if (getenv("TRELLIS_DBG_ALLOC")) {
+                fprintf(stderr,
+                        "      [shape-dec-stream] stage %d C=%d N=%d feature-table=%.2f GiB "
+                        ">= %.2f GiB; streaming\n",
+                        si, st.C, N, feature_bytes / 1073741824.0, stream_at / 1073741824.0);
+            }
+
+            if (stream_stage) {
+                for (int j = 0; j < st.nblk; ++j) {
+                    const std::string bp = std::string("blocks.") + st.s + "." + std::to_string(j);
+                    h = sparse_convnext_streamed(m, bp, h, st.C, nbr, N);
+                }
+            }
         }
         mem_probe("after ConvNeXt stage");
         const std::vector<uint8_t>* ext = guide_subs ? &(*guide_subs)[si] : nullptr;