Gemma3 NF4 on GPU: quantized encoder, batch encoding, byte-LUT eager dequant

DESCRIPTION

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

NAMESPACE

@@ -66,6 +66,8 @@ export(flux2_unpack_latents_with_ids)
 export(flux2_unpatchify_latents)
 export(flux2_vae_decoder)
 export(gemma3_config_ltx2)
+export(gemma3_encode_batch)
+export(gemma3_quantize_nf4)
 export(gemma3_text_model)
 export(gemma3_tokenizer)
 export(img2img)
@@ -74,6 +76,7 @@ export(latents_to_video)
 export(load_decoder_safetensors)
 export(load_decoder_weights)
 export(load_flux2_vae_decoder)
+export(load_gemma3_nf4)
 export(load_gemma3_text_encoder)
 export(load_model_component)
 export(load_pipeline)

NEWS.md

@@ -1,3 +1,25 @@
+# diffuseR 0.1.0.13 (development)
+
+* `gemma3_quantize_nf4()` + `load_gemma3_nf4()` put the Gemma3 12B text
+  encoder on the GPU: the 336 projection weights quantize to a ~8 GB
+  NF4 artifact (one-time 52 s; vision tower dropped), which loads in
+  ~12 s and encodes in ~7 s on CUDA vs ~30 s load + ~24 s encode for
+  the fp32 CPU path, at 8 GB host RAM instead of 45. Pinned phase swap
+  costs 0.3 s on / 0 s off after a one-time 3.3 s page-lock.
+  `load_gemma3_text_encoder()` dispatches to the artifact
+  automatically. Embedding drift (cosine 0.989 vs fp32) renders
+  scene-equivalent same-seed videos.
+* The eager NF4 dequant now shares the byte-LUT with the jit block
+  stack (one index per packed byte, gather straight into the output),
+  speeding every eager NF4 forward (Gemma3, FLUX.1, LTX fallbacks).
+* Fixed `txt2vid_ltx2(decode_audio = FALSE, filename = )`: partial
+  matching handed the raw audio latents to the WAV writer.
+* `gemma3_encode_batch()` encodes a prompt vector in sub-batches with
+  optional per-prompt disk caching (resumable; returns cache paths):
+  encode a whole episode list once, then swap the encoder off and
+  render from the cached embeddings. Works around a torch_save bug
+  where storage-offset views serialize from the base storage.
+
 # diffuseR 0.1.0.12 (development)
 
 * LTX video decode runs untiled when the estimated activation cost

R/gemma3_text_encoder.R

@@ -550,9 +550,33 @@ gemma3_config_ltx2 <- function() {
 # Weight Loading
 # -----------------------------------------------------------------------------
 
+# Model config with HF defaults, shared by the checkpoint and NF4
+# artifact loaders (the raw text_config travels in the NF4 manifest)
+.gemma3_build_config <- function(config_raw) {
+    list(
+         vocab_size = config_raw$vocab_size %||% 262208L,
+         hidden_size = config_raw$hidden_size %||% 3840L,
+         intermediate_size = config_raw$intermediate_size %||% 15360L,
+         num_hidden_layers = config_raw$num_hidden_layers %||% 48L,
+         num_attention_heads = config_raw$num_attention_heads %||% 16L,
+         num_key_value_heads = config_raw$num_key_value_heads %||% 8L,
+         head_dim = config_raw$head_dim %||% 256L,
+         max_position_embeddings = config_raw$max_position_embeddings %||% 131072L,
+         rms_norm_eps = config_raw$rms_norm_eps %||% 1e-6,
+         rope_theta = config_raw$rope_theta %||% 1000000.0, # Gemma3 uses 1M
+         rope_scaling = list(factor = config_raw$rope_scaling$factor %||% 8.0),
+         sliding_window = config_raw$sliding_window %||% 1024L,
+         sliding_window_pattern = config_raw$sliding_window_pattern %||% 6L,
+         attn_logit_softcapping = config_raw$attn_logit_softcapping, # NULL = no softcapping (HF default)
+         query_pre_attn_scalar = config_raw$query_pre_attn_scalar %||% 256L
+    )
+}
+
 #' Load Gemma3 Text Model from safetensors
 #'
 #' Loads pre-trained Gemma3 weights from HuggingFace safetensors files.
+#' An NF4 artifact directory (from \code{\link{gemma3_quantize_nf4}})
+#' dispatches to \code{\link{load_gemma3_nf4}}.
 #'
 #' @param model_path Character. Path to directory containing model files.
 #' @param device Character. Device to load model to.
@@ -562,6 +586,15 @@ gemma3_config_ltx2 <- function() {
 #' @export
 load_gemma3_text_encoder <- function(model_path, device = "cpu",
                                      dtype = "float16", verbose = TRUE) {
+    # An NF4 artifact directory (gemma3_quantize_nf4) loads through the
+    # quantized path; bf16 compute unless the caller says otherwise
+    if (file.exists(file.path(model_path, "manifest.json"))) {
+        if (identical(dtype, "float16")) {
+            dtype <- "bfloat16"
+        }
+        return(load_gemma3_nf4(model_path, device = device, dtype = dtype,
+                               verbose = verbose))
+    }
     # Load config
     config_path <- file.path(model_path, "config.json")
     if (!file.exists(config_path)) {
@@ -575,23 +608,7 @@ load_gemma3_text_encoder <- function(model_path, device = "cpu",
         config_raw <- config_raw$text_config
     }
 
-    config <- list(
-                   vocab_size = config_raw$vocab_size %||% 262208L,
-                   hidden_size = config_raw$hidden_size %||% 3840L,
-                   intermediate_size = config_raw$intermediate_size %||% 15360L,
-                   num_hidden_layers = config_raw$num_hidden_layers %||% 48L,
-                   num_attention_heads = config_raw$num_attention_heads %||% 16L,
-                   num_key_value_heads = config_raw$num_key_value_heads %||% 8L,
-                   head_dim = config_raw$head_dim %||% 256L,
-                   max_position_embeddings = config_raw$max_position_embeddings %||% 131072L,
-                   rms_norm_eps = config_raw$rms_norm_eps %||% 1e-6,
-                   rope_theta = config_raw$rope_theta %||% 1000000.0, # Gemma3 uses 1M
-                   rope_scaling = list(factor = config_raw$rope_scaling$factor %||% 8.0),
-                   sliding_window = config_raw$sliding_window %||% 1024L,
-                   sliding_window_pattern = config_raw$sliding_window_pattern %||% 6L,
-                   attn_logit_softcapping = config_raw$attn_logit_softcapping, # NULL = no softcapping (HF default)
-                   query_pre_attn_scalar = config_raw$query_pre_attn_scalar %||% 256L
-    )
+    config <- .gemma3_build_config(config_raw)
 
     if (verbose) {
         message(sprintf("Creating Gemma3 model: %d layers, hidden_size=%d",
@@ -866,3 +883,106 @@ if (!exists("%||%", mode = "function")) {
         x
     }
 }
+
+#' Batch-encode prompts with Gemma3, cached to disk
+#'
+#' Encodes a vector of prompts in sub-batches (bounding activation
+#' VRAM) and optionally caches each prompt's result under
+#' \code{cache_dir}, keyed by the prompt text and sequence length.
+#' Already-cached prompts are skipped, so an interrupted batch resumes
+#' where it stopped. Embeddings land on the CPU either way; the
+#' renderer moves them per phase.
+#'
+#' Budget note: each prompt's result is the full hidden-state stack the
+#' LTX connectors consume - roughly 0.4 GB at 1024 tokens - so caching
+#' N prompts needs ~0.4N GB of disk.
+#'
+#' @param prompts Character vector.
+#' @param model,tokenizer As in \code{\link{encode_with_gemma3}}.
+#' @param batch_size Integer. Prompts per forward pass (default 4;
+#'   raise on cards with headroom, lower if the encode OOMs).
+#' @param cache_dir Optional directory. When given, each prompt is
+#'   written to \code{gemma3-<md5>.pt} and the return value is the
+#'   character vector of those paths (load one with
+#'   \code{torch::torch_load}); when NULL, the return value is a list
+#'   of \code{list(prompt_embeds, prompt_attention_mask)} in prompt
+#'   order.
+#' @param max_sequence_length,device,verbose As in
+#'   \code{\link{encode_with_gemma3}}.
+#'
+#' @return Character vector of cache paths (with \code{cache_dir}) or
+#'   a list of per-prompt embedding results (without).
+#'
+#' @export
+gemma3_encode_batch <- function(prompts, model = NULL, tokenizer = NULL,
+                                batch_size = 4L, cache_dir = NULL,
+                                max_sequence_length = 1024L, device = "cuda",
+                                verbose = TRUE) {
+    prompts <- as.character(prompts)
+    n <- length(prompts)
+    if (!n) {
+        stop("No prompts given")
+    }
+
+    cache_paths <- NULL
+    if (!is.null(cache_dir)) {
+        dir.create(cache_dir, recursive = TRUE, showWarnings = FALSE)
+        cache_paths <- vapply(prompts, function(p) {
+            tf <- tempfile()
+            writeLines(c(p, as.character(max_sequence_length)), tf)
+            key <- unname(tools::md5sum(tf))
+            unlink(tf)
+            file.path(cache_dir, paste0("gemma3-", key, ".pt"))
+        }, character(1), USE.NAMES = FALSE)
+        todo <- which(!file.exists(cache_paths))
+        if (verbose && length(todo) < n) {
+            message(sprintf("%d/%d prompts already cached", n - length(todo),
+                            n))
+        }
+    } else {
+        todo <- seq_len(n)
+    }
+
+    if (is.null(cache_dir)) {
+        results <- vector("list", n)
+    } else {
+        results <- NULL
+    }
+    done <- 0L
+    for (chunk in split(todo, ceiling(seq_along(todo) / batch_size))) {
+        enc <- encode_with_gemma3(prompts[chunk], model = model,
+                                  tokenizer = tokenizer,
+                                  max_sequence_length = max_sequence_length,
+                                  device = device, verbose = FALSE)
+        embeds <- enc$prompt_embeds$cpu()
+        mask <- enc$prompt_attention_mask$cpu()
+        for (j in seq_along(chunk)) {
+            # clone(), not contiguous(): a narrow of a contiguous
+            # tensor is already contiguous, so contiguous() keeps the
+            # storage-offset view - and torch_save serializes offset
+            # views from the base storage (every row would save as
+            # row 1). clone() forces fresh storage.
+            one <- list(
+                        prompt_embeds = embeds$narrow(1L, j, 1L)$clone(),
+                        prompt_attention_mask = mask$narrow(1L, j, 1L)$clone()
+            )
+            if (is.null(cache_dir)) {
+                results[[chunk[j]]] <- one
+            } else {
+                torch::torch_save(one, cache_paths[chunk[j]])
+            }
+        }
+        rm(enc, embeds, mask)
+        gc(verbose = FALSE)
+        done <- done + length(chunk)
+        if (verbose) {
+            message(sprintf("  encoded %d/%d prompts", done, length(todo)))
+        }
+    }
+
+    if (is.null(cache_dir)) {
+        results
+    } else {
+        cache_paths
+    }
+}

R/nf4_ltx23.R

@@ -82,42 +82,32 @@ ltx23_nf4_dequantize <- function(packed, absmax, shape,
     if (is.null(out)) {
         out <- torch::torch_empty(shape, dtype = dtype, device = packed$device)
     }
-    out_flat <- out$view(-1L)
+    out_pairs <- out$view(c(-1L, 2L))
+    out_blocks <- out$view(c(-1L, block))
 
     n_bytes <- packed$shape[1]
     bytes_per_chunk <- max(chunk_elements %/% 2L, block)
     scratch <- .ltx23_get_dequant_scratch(min(bytes_per_chunk, n_bytes),
         packed$device)
+    # Byte LUT [256, 2] in the target dtype (shared with the jit block
+    # stack): one 1-based index per PACKED byte and a single gather of
+    # (hi, lo) pairs written straight into out - no nibble unpack, no
+    # per-value int64 index, no float32 staging copy
+    tbl <- .ltx23_jit_table(packed$device, dtype)
+    am <- absmax$to(dtype = dtype)
 
     start <- 1L
     torch::with_no_grad({
         while (start <= n_bytes) {
             len <- min(bytes_per_chunk, n_bytes - start + 1L)
-            chunk <- packed$narrow(1L, start, len)
-
-            # Fully in-place nibble unpack into persistent scratch:
-            # hi = byte %/% 16, lo = byte - 16 * hi
-            hi <- scratch$hi$narrow(1L, 1L, len)
-            lo <- scratch$lo$narrow(1L, 1L, len)
-            hi$copy_(chunk)$div_(16L, rounding_mode = "floor")
-            lo$copy_(hi)$mul_(-16L)$add_(chunk)
-
-            # Interleave into the (1-based) int64 index scratch
-            idx <- scratch$idx$narrow(1L, 1L, len * 2L)
-            idx_pairs <- idx$view(c(-1L, 2L))
-            idx_pairs$narrow(2L, 1L, 1L)$squeeze(2L)$copy_(hi)
-            idx_pairs$narrow(2L, 2L, 1L)$squeeze(2L)$copy_(lo)
-            idx$add_(1L)
-
-            vals <- scratch$vals$narrow(1L, 1L, len * 2L)
-            .ltx23_index_select_into(vals, scratch$table, idx)
+            idx <- scratch$idx$narrow(1L, 1L, len)
+            idx$copy_(packed$narrow(1L, start, len))$add_(1L)
+            .ltx23_index_select_into(out_pairs$narrow(1L, start, len), tbl, idx)
 
             block_start <- ((start - 1L) * 2L) %/% block + 1L
             n_blocks <- (len * 2L) %/% block
-            scales <- absmax$narrow(1L, block_start, n_blocks)
-            vals$view(c(-1L, block))$mul_(scales$unsqueeze(2L))
-
-            out_flat$narrow(1L, (start - 1L) * 2L + 1L, len * 2L)$copy_(vals)
+            out_blocks$narrow(1L, block_start, n_blocks)$
+            mul_(am$narrow(1L, block_start, n_blocks)$unsqueeze(2L))
             start <- start + len
         }
     })
@@ -159,18 +149,7 @@ ltx23_nf4_dequantize <- function(packed, absmax, shape,
     if (is.null(scratch) || scratch$n_bytes < n_bytes) {
         scratch <- list(
                         n_bytes = n_bytes,
-                        hi = torch::torch_empty(n_bytes, dtype = torch::torch_uint8(),
-                device = device),
-                        lo = torch::torch_empty(n_bytes, dtype = torch::torch_uint8(),
-                device = device),
-                        idx = torch::torch_empty(n_bytes * 2L,
-                dtype = torch::torch_long(),
-                device = device),
-                        vals = torch::torch_empty(n_bytes * 2L,
-                dtype = torch::torch_float32(),
-                device = device),
-                        table = torch::torch_tensor(.ltx23_nf4_table,
-                dtype = torch::torch_float32(),
+                        idx = torch::torch_empty(n_bytes, dtype = torch::torch_long(),
                 device = device)
         )
         .ltx23_dequant_scratch[[key]] <- scratch