Activate the LTX allocator gates and pinned staging

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.1.0.10
+Version: 0.1.0.11
 Authors@R: c(
     person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
            comment = c(ORCID = "0009-0005-4248-604X")),

NEWS.md

@@ -1,3 +1,20 @@
+# diffuseR 0.1.0.11 (development)
+
+* `ltx23_tune_gc()` now takes effect: it used to set the allocator gate
+  options one call after `torch::cuda_is_available()` had started torch
+  (which reads them exactly once at init), so every option it set was
+  inert. Options now land before the first torch call, the three CUDA
+  gates are pushed into the live allocator, and `.onLoad` defaults
+  `torch.threshold_call_gc` (which has no live setter). Measured: ~4-5%
+  off LTX render walls (63.7/62.2/61.9 s vs 67.9/64.2/65.6 s at
+  768x512x49 NF4).
+* `diffuseR.pin_staging` now defaults to TRUE for LTX phase offload:
+  page-locked host copies make onload a DMA transfer and offload a
+  pointer swap, saving ~7 s per render for +9 s of one-time
+  page-locking at pipeline load (break-even on the second render).
+  Opt out with `options(diffuseR.pin_staging = FALSE)` under host
+  memory pressure.
+
 # diffuseR 0.1.0.10 (development)
 
 * The generators accept a three-level `verbose`: "silent", "progress",

R/memory_ltx23.R

@@ -106,8 +106,17 @@ ltx23_memory_profile <- function(vram_gb = NULL) {
 #' actual footprint is safe here because the LTX hot loops compute into
 #' persistent scratch buffers (near-zero per-step garbage). Also raises
 #' the host-allocation GC threshold and defaults
-#' \code{PYTORCH_CUDA_ALLOC_CONF} to expandable segments. Must run before
-#' the first CUDA op; user-set options win.
+#' \code{PYTORCH_CUDA_ALLOC_CONF} to expandable segments. User-set
+#' options win.
+#'
+#' \code{start_torch()} reads the gate options exactly once, so setting
+#' them after torch has started is inert on its own. The three CUDA
+#' gates are therefore also pushed into the live allocator here (the
+#' \code{.flux_gc_gates} pattern), which makes this function effective
+#' whenever it runs. The host-side \code{torch.threshold_call_gc} has
+#' no live setter; the package defaults it in \code{.onLoad} so torch
+#' reads it at init in any session that loads diffuseR before running
+#' torch ops.
 #'
 #' @param footprint_gb Numeric. Expected resident GPU footprint in GB
 #'   (NF4 transformer: ~12).
@@ -120,15 +129,10 @@ ltx23_tune_gc <- function(footprint_gb = 12, total_gb = NULL) {
     if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) {
         Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:True")
     }
-    if (!torch::cuda_is_available()) {
-        return(invisible(NULL))
-    }
-    if (is.null(total_gb)) {
-        total_gb <- .detect_vram(use_free = FALSE)
-        if (!isTRUE(total_gb > 0)) {
-            return(invisible(NULL))
-        }
-    }
+    # Options are set before the first torch:: call in this function:
+    # torch::cuda_is_available() below starts torch, and start_torch
+    # reads these once. (The old order set them one call after init
+    # read them - inert by construction.)
     if (is.null(getOption("torch.threshold_call_gc"))) {
         options(torch.threshold_call_gc = 16000)
     }
@@ -144,8 +148,27 @@ ltx23_tune_gc <- function(footprint_gb = 12, total_gb = NULL) {
     }
     rate <- NULL
     if (is.null(getOption("torch.cuda_allocator_reserved_rate"))) {
-        rate <- min(0.92, max(0.20, footprint_gb / total_gb))
-        options(torch.cuda_allocator_reserved_rate = rate)
+        if (is.null(total_gb)) {
+            total_gb <- .detect_vram(use_free = FALSE)
+        }
+        if (isTRUE(total_gb > 0)) {
+            rate <- min(0.92, max(0.20, footprint_gb / total_gb))
+            options(torch.cuda_allocator_reserved_rate = rate)
+        }
+    }
+    if (!torch::cuda_is_available()) {
+        return(invisible(rate))
+    }
+    # Torch is usually long started by the time a loader calls this:
+    # push the CUDA gates into the live allocator directly.
+    push <- get0("cpp_set_cuda_allocator_allocator_thresholds",
+                 envir = asNamespace("torch"))
+    if (is.function(push)) {
+        try(push(
+                 getOption("torch.cuda_allocator_reserved_rate", 0.2),
+                 getOption("torch.cuda_allocator_allocated_rate", 0.8),
+                 getOption("torch.cuda_allocator_allocated_reserved_rate", 0.8)
+            ), silent = TRUE)
     }
     invisible(rate)
 }

R/staging_ltx23.R

@@ -10,16 +10,17 @@
 #' re-points the tensors at the still-valid pinned copies — weights
 #' are immutable during inference, so offload moves no bytes at all.
 #'
-#' Costs: the pinned copies are non-swappable host RAM for the life
-#' of the pipeline (about the model's CPU footprint), and page-locking
-#' ~20GB adds ~30s to pipeline load. Measured per render the win is
-#' small (~2.7s: warm pinned onload 0.54s vs pageable 0.99s, offload
-#' 0.06s vs 2.36s) because the real phase-transition cost was
-#' allocator pool regrowth, fixed separately by the loader's pool
-#' pre-warm and by not emptying the CUDA cache between phases. Off by
-#' default; enable for long multi-render sessions with
-#' \code{options(diffuseR.pin_staging = TRUE)} before
-#' \code{ltx23_load_pipeline} (breaks even after ~11 renders).
+#' Costs: the model's host copies become non-swappable for the life
+#' of the pipeline (no extra RAM - set_data repoints the same
+#' tensors), and page-locking adds ~9s to pipeline load. Measured
+#' post byte-LUT (768x512x49, NF4, RTX 5060 Ti): ~7s saved per render
+#' (warm renders 64-66s pageable vs 57-59s pinned; denoise and decode
+#' identical, the delta is pure transfer), so pinning breaks even on
+#' the second render and costs a single-render session ~2s net. On by
+#' default; page-locking failure falls back silently per component,
+#' and \code{options(diffuseR.pin_staging = FALSE)} before
+#' \code{ltx23_load_pipeline} opts out (e.g. under host memory
+#' pressure, where unswappable pages turn thrashing into OOM).
 #'
 #' @name staging_ltx23
 NULL

R/txt2vid_ltx23.R

@@ -349,7 +349,7 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda",
     }
 
     if (phase_offload && device == "cuda" &&
-        isTRUE(getOption("diffuseR.pin_staging", FALSE))) {
+        isTRUE(getOption("diffuseR.pin_staging", TRUE))) {
         # Page-lock every phase-offloaded component once so the
         # per-render CPU<->GPU moves run at full PCIe rate (offload
         # becomes a pointer swap; see staging_ltx23.R). Falls back

R/zzz.R

@@ -3,3 +3,16 @@ NULL
 
 # torch indexing uses `..` as a Python-style ellipsis
 utils::globalVariables("..")
+
+.onLoad <- function(libname, pkgname) {
+    # start_torch() reads torch.threshold_call_gc exactly once at torch
+    # init and there is no live setter (unlike the CUDA gates, which
+    # ltx23_tune_gc pushes via cpp), so this is the only place the
+    # option can land in time. The 4000 MB default fires an R gc for
+    # every few GB of host allocation; raising it measurably cuts gc
+    # counts, though an LTX load-time A/B showed no wall-clock change
+    # (those gcs were cheap). A user-set option wins.
+    if (is.null(getOption("torch.threshold_call_gc"))) {
+        options(torch.threshold_call_gc = 16000)
+    }
+}

inst/tinytest/test_memory_ltx23.R

@@ -0,0 +1,49 @@
+# ltx23_tune_gc option semantics: defaults land when unset, user-set
+# values win, and the reserved rate clamps to [0.20, 0.92]. The cpp
+# gate push itself is CUDA-only and exercised implicitly wherever a
+# GPU is present.
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+  exit_file("torch not fully installed")
+}
+
+library(diffuseR)
+
+gate_opts <- c("torch.threshold_call_gc",
+               "torch.cuda_allocator_allocated_rate",
+               "torch.cuda_allocator_allocated_reserved_rate",
+               "torch.cuda_allocator_reserved_rate")
+old <- options()[gate_opts]
+names(old) <- gate_opts
+clear <- function() {
+  for (o in gate_opts) {
+    opt <- list(NULL)
+    names(opt) <- o
+    options(opt)
+  }
+}
+
+clear()
+rate <- ltx23_tune_gc(footprint_gb = 12, total_gb = 16)
+expect_equal(rate, 0.75)
+expect_equal(getOption("torch.cuda_allocator_reserved_rate"), 0.75)
+expect_equal(getOption("torch.threshold_call_gc"), 16000)
+expect_equal(getOption("torch.cuda_allocator_allocated_rate"), 0.95)
+expect_equal(getOption("torch.cuda_allocator_allocated_reserved_rate"), 0.95)
+
+# User-set values win; the reserved rate is then not recomputed
+clear()
+options(torch.cuda_allocator_reserved_rate = 0.5,
+        torch.threshold_call_gc = 4000)
+rate <- ltx23_tune_gc(footprint_gb = 12, total_gb = 16)
+expect_null(rate)
+expect_equal(getOption("torch.cuda_allocator_reserved_rate"), 0.5)
+expect_equal(getOption("torch.threshold_call_gc"), 4000)
+
+# Clamps: floor 0.20 for small footprints, ceiling 0.92 for tight fits
+clear()
+expect_equal(ltx23_tune_gc(footprint_gb = 12, total_gb = 100), 0.20)
+clear()
+expect_equal(ltx23_tune_gc(footprint_gb = 12, total_gb = 12), 0.92)
+
+options(old)

man/ltx23_tune_gc.Rd

@@ -23,6 +23,17 @@ weights that fires on nearly every allocation. Raising the rate to the
 actual footprint is safe here because the LTX hot loops compute into
 persistent scratch buffers (near-zero per-step garbage). Also raises
 the host-allocation GC threshold and defaults
-\code{PYTORCH_CUDA_ALLOC_CONF} to expandable segments. Must run before
-the first CUDA op; user-set options win.
+\code{PYTORCH_CUDA_ALLOC_CONF} to expandable segments. User-set
+options win.
+}
+\details{
+\code{start_torch()} reads the gate options exactly once, so setting
+them after torch has started is inert on its own. The three CUDA
+gates are therefore also pushed into the live allocator here (the
+\code{.flux_gc_gates} pattern), which makes this function effective
+whenever it runs. The host-side \code{torch.threshold_call_gc} has
+no live setter; the package defaults it in \code{.onLoad} so torch
+reads it at init in any session that loads diffuseR before running
+torch ops.
+
 }

man/staging_ltx23.Rd

@@ -14,15 +14,16 @@ re-points the tensors at the still-valid pinned copies — weights
 are immutable during inference, so offload moves no bytes at all.
 }
 \details{
-Costs: the pinned copies are non-swappable host RAM for the life
-of the pipeline (about the model's CPU footprint), and page-locking
-~20GB adds ~30s to pipeline load. Measured per render the win is
-small (~2.7s: warm pinned onload 0.54s vs pageable 0.99s, offload
-0.06s vs 2.36s) because the real phase-transition cost was
-allocator pool regrowth, fixed separately by the loader's pool
-pre-warm and by not emptying the CUDA cache between phases. Off by
-default; enable for long multi-render sessions with
-\code{options(diffuseR.pin_staging = TRUE)} before
-\code{ltx23_load_pipeline} (breaks even after ~11 renders).
+Costs: the model's host copies become non-swappable for the life
+of the pipeline (no extra RAM - set_data repoints the same
+tensors), and page-locking adds ~9s to pipeline load. Measured
+post byte-LUT (768x512x49, NF4, RTX 5060 Ti): ~7s saved per render
+(warm renders 64-66s pageable vs 57-59s pinned; denoise and decode
+identical, the delta is pure transfer), so pinning breaks even on
+the second render and costs a single-render session ~2s net. On by
+default; page-locking failure falls back silently per component,
+and \code{options(diffuseR.pin_staging = FALSE)} before
+\code{ltx23_load_pipeline} opts out (e.g. under host memory
+pressure, where unswappable pages turn thrashing into OOM).
 
 }