Byte-LUT NF4 dequant in the JIT block stack

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.1.0.8
+Version: 0.1.0.9
 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,11 @@
+# diffuseR 0.1.0.9 (development)
+
+* The LTX-2.3 JIT block stack dequantizes NF4 weights through a
+  precomputed [256, 2] byte lookup table (one embedding gather in the
+  compute dtype) instead of the int64 shift/stack/gather chain, cutting
+  per-step dequant memory traffic ~6x (isolated benchmark: 5.7x; the
+  dequant was the measured per-step wall at ~4.4 s of a 6.5 s step).
+
 # diffuseR 0.1.0.8 (development)
 
 * The Qwen3 encoder builds its additive attention mask in the query

R/dit_ltx23_modules.R

@@ -175,7 +175,7 @@ ltx23_ada_layer_norm_single <- torch::nn_module(
 # adaptively against a memory budget and all large temporaries live in
 # reusable scratch buffers.
 .ltx23_sdpa_manual <- function(query, key, value, attention_mask = NULL,
-    chunk_size = NULL) {
+                               chunk_size = NULL) {
     head_dim <- query$shape[length(query$shape)]
     scale <- 1.0 / sqrt(head_dim)
     key_t <- key$transpose(-2L, -1L)
@@ -241,11 +241,10 @@ ltx23_ada_layer_norm_single <- torch::nn_module(
     }
 
     rows <- min(n_q, chunk_size)
-    attn_buf <- .ltx23_get_attn_buffer(
-        "scores", c(b, heads, rows, n_k), query$dtype, query$device
-    )
+    attn_buf <- .ltx23_get_attn_buffer("scores", c(b, heads, rows, n_k),
+                                       query$dtype, query$device)
     out_buf <- .ltx23_get_attn_buffer(
-        "output", c(b, heads, n_q, d_v), query$dtype, query$device
+                                      "output", c(b, heads, n_q, d_v), query$dtype, query$device
     )
 
     if (n_q <= chunk_size) {
@@ -278,10 +277,9 @@ ltx23_ada_layer_norm_single <- torch::nn_module(
     function() {
         if (!resolved) {
             if ("torch_scaled_dot_product_attention" %in%
-                    getNamespaceExports("torch")) {
-                fn <<- getExportedValue(
-                    "torch", "torch_scaled_dot_product_attention"
-                )
+                              getNamespaceExports("torch")) {
+                fn <<- getExportedValue("torch",
+                                        "torch_scaled_dot_product_attention")
             }
             resolved <<- TRUE
         }
@@ -299,7 +297,7 @@ ltx23_ada_layer_norm_single <- torch::nn_module(
         fn(query, key, value, dropout_p = 0, is_causal = FALSE)
     } else {
         fn(query, key, value, attn_mask = attention_mask,
-           dropout_p = 0, is_causal = FALSE)
+            dropout_p = 0, is_causal = FALSE)
     }
 }
 
@@ -308,8 +306,8 @@ ltx23_ada_layer_norm_single <- torch::nn_module(
 .ltx23_sdpa <- function(query, key, value, attention_mask = NULL,
                         chunk_size = NULL) {
     use_fused <- is.null(chunk_size) &&
-        isTRUE(getOption("diffuseR.ltx23_fused_sdpa", TRUE)) &&
-        !is.null(.ltx23_fused_sdpa_fn())
+    isTRUE(getOption("diffuseR.ltx23_fused_sdpa", TRUE)) &&
+    !is.null(.ltx23_fused_sdpa_fn())
     if (use_fused) {
         return(.ltx23_fused_sdpa(query, key, value, attention_mask))
     }

R/jit_ltx23.R

@@ -37,24 +37,17 @@ NULL
 
 .ltx23_jit_source <- function() {
     "
-def nf4_lin(x: Tensor, packed: Tensor, absmax: Tensor, bias: Tensor, table: Tensor) -> Tensor:
+def nf4_lin(x: Tensor, packed: Tensor, absmax: Tensor, bias: Tensor, table2: Tensor) -> Tensor:
+    # table2 is the [256, 2] byte LUT (hi nibble, lo nibble) in the
+    # compute dtype: one embedding gather replaces the shift/and/stack
+    # int64-index chain, whose intermediates cost ~70 bytes of traffic
+    # per half-byte weight and made dequant the per-step wall. absmax
+    # is cast down before the scale so the product never promotes to
+    # float32 (which would materialize a full-precision weight copy).
     rows = bias.size(0)
-    n = packed.size(0)
-    step = 4194304
-    outs: List[Tensor] = []
-    i = 0
-    while i < n:
-        j = min(i + step, n)
-        chunk = packed.narrow(0, i, j - i)
-        hi = torch.bitwise_right_shift(chunk, 4).long()
-        lo = torch.bitwise_and(chunk, 15).long()
-        idx = torch.stack([hi, lo], -1).flatten()
-        vals = torch.index_select(table, 0, idx)
-        sc = absmax.narrow(0, i * 2 // 64, (j - i) * 2 // 64)
-        outs.append((vals.reshape(-1, 64) * sc.unsqueeze(1)).flatten().type_as(x))
-        i = j
-    w = torch.cat(outs, 0).reshape([rows, n * 2 // rows])
-    return torch.linear(x, w, bias)
+    vals = torch.embedding(table2, packed.int())
+    w = vals.reshape([-1, 64]) * absmax.type_as(table2).unsqueeze(1)
+    return torch.linear(x, w.reshape([rows, -1]).type_as(x), bias)
 
 def rmsn(x: Tensor) -> Tensor:
     v = x.float().pow(2).mean(-1, keepdim=True)
@@ -97,14 +90,14 @@ def modtok(vada: Tensor, j: int, idx: Optional[Tensor]) -> Tensor:
         return v
     return v.index_select(1, idx)
 
-def attn_nf4(x: Tensor, ctx: Tensor, ws: List[Tensor], base: int, table: Tensor, heads: int,
+def attn_nf4(x: Tensor, ctx: Tensor, ws: List[Tensor], base: int, table2: Tensor, heads: int,
              q_cos: Optional[Tensor], q_sin: Optional[Tensor],
              k_cos: Optional[Tensor], k_sin: Optional[Tensor],
              mask: Optional[Tensor]) -> Tensor:
     gl = torch.linear(x, ws[base], ws[base + 1])
-    q = rmsn_w(nf4_lin(x, ws[base + 2], ws[base + 3], ws[base + 4], table), ws[base + 14])
-    k = rmsn_w(nf4_lin(ctx, ws[base + 5], ws[base + 6], ws[base + 7], table), ws[base + 15])
-    v = nf4_lin(ctx, ws[base + 8], ws[base + 9], ws[base + 10], table)
+    q = rmsn_w(nf4_lin(x, ws[base + 2], ws[base + 3], ws[base + 4], table2), ws[base + 14])
+    k = rmsn_w(nf4_lin(ctx, ws[base + 5], ws[base + 6], ws[base + 7], table2), ws[base + 15])
+    v = nf4_lin(ctx, ws[base + 8], ws[base + 9], ws[base + 10], table2)
     if q_cos is not None and q_sin is not None:
         q = rope(q, q_cos, q_sin)
     if k_cos is not None and k_sin is not None:
@@ -116,11 +109,11 @@ def attn_nf4(x: Tensor, ctx: Tensor, ws: List[Tensor], base: int, table: Tensor,
     o = o.transpose(1, 2).flatten(2).type_as(x)
     gates = torch.sigmoid(gl) * 2.0
     o = (o.unflatten(-1, [heads, -1]) * gates.unsqueeze(-1)).flatten(2)
-    return nf4_lin(o, ws[base + 11], ws[base + 12], ws[base + 13], table)
+    return nf4_lin(o, ws[base + 11], ws[base + 12], ws[base + 13], table2)
 
-def ff_nf4(x: Tensor, ws: List[Tensor], base: int, table: Tensor) -> Tensor:
-    h = torch.gelu(nf4_lin(x, ws[base], ws[base + 1], ws[base + 2], table), approximate=\"tanh\")
-    return nf4_lin(h, ws[base + 3], ws[base + 4], ws[base + 5], table)
+def ff_nf4(x: Tensor, ws: List[Tensor], base: int, table2: Tensor) -> Tensor:
+    h = torch.gelu(nf4_lin(x, ws[base], ws[base + 1], ws[base + 2], table2), approximate=\"tanh\")
+    return nf4_lin(h, ws[base + 3], ws[base + 4], ws[base + 5], table2)
 
 def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor,
               temb: Tensor, temb_a: Tensor,
@@ -130,30 +123,30 @@ def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor,
               cav_cos: Tensor, cav_sin: Tensor, caa_cos: Tensor, caa_sin: Tensor,
               enc_mask: Optional[Tensor], aenc_mask: Optional[Tensor],
               cond_idx: Optional[Tensor],
-              ws: List[Tensor], base: int, table: Tensor,
+              ws: List[Tensor], base: int, table2: Tensor,
               heads: int, aheads: int) -> Tuple[Tensor, Tensor]:
     vada = mods(ws[base + 108], temb, 9)
     aada = mods(ws[base + 109], temb_a, 9)
 
     nh = rmsn(h) * (modtok(vada, 1, cond_idx) + 1.0) + modtok(vada, 0, cond_idx)
-    ax = attn_nf4(nh, nh, ws, base, table, heads, v_cos, v_sin, v_cos, v_sin, None)
+    ax = attn_nf4(nh, nh, ws, base, table2, heads, v_cos, v_sin, v_cos, v_sin, None)
     h = h + ax * modtok(vada, 2, cond_idx)
 
     nah = rmsn(ah) * (aada.select(2, 1) + 1.0) + aada.select(2, 0)
-    aax = attn_nf4(nah, nah, ws, base + 16, table, aheads, a_cos, a_sin, a_cos, a_sin, None)
+    aax = attn_nf4(nah, nah, ws, base + 16, table2, aheads, a_cos, a_sin, a_cos, a_sin, None)
     ah = ah + aax * aada.select(2, 2)
 
     pada = mods(ws[base + 110], tp, 2)
     apada = mods(ws[base + 111], tpa, 2)
 
     nh = rmsn(h) * (modtok(vada, 7, cond_idx) + 1.0) + modtok(vada, 6, cond_idx)
     encm = enc * (pada.select(2, 1) + 1.0) + pada.select(2, 0)
-    ax = attn_nf4(nh, encm, ws, base + 32, table, heads, None, None, None, None, enc_mask)
+    ax = attn_nf4(nh, encm, ws, base + 32, table2, heads, None, None, None, None, enc_mask)
     h = h + ax * modtok(vada, 8, cond_idx)
 
     nah = rmsn(ah) * (aada.select(2, 7) + 1.0) + aada.select(2, 6)
     aencm = aenc * (apada.select(2, 1) + 1.0) + apada.select(2, 0)
-    aax = attn_nf4(nah, aencm, ws, base + 48, table, aheads, None, None, None, None, aenc_mask)
+    aax = attn_nf4(nah, aencm, ws, base + 48, table2, aheads, None, None, None, None, aenc_mask)
     ah = ah + aax * aada.select(2, 8)
 
     nh = rmsn(h)
@@ -165,19 +158,19 @@ def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor,
 
     mnh = nh * (vca.select(2, 0) + 1.0) + vca.select(2, 1)
     mna = nah * (aca.select(2, 0) + 1.0) + aca.select(2, 1)
-    a2v = attn_nf4(mnh, mna, ws, base + 64, table, aheads, cav_cos, cav_sin, caa_cos, caa_sin, None)
+    a2v = attn_nf4(mnh, mna, ws, base + 64, table2, aheads, cav_cos, cav_sin, caa_cos, caa_sin, None)
     h = h + vcg.select(2, 0) * a2v
 
     mnh = nh * (vca.select(2, 2) + 1.0) + vca.select(2, 3)
     mna = nah * (aca.select(2, 2) + 1.0) + aca.select(2, 3)
-    v2a = attn_nf4(mna, mnh, ws, base + 80, table, aheads, caa_cos, caa_sin, cav_cos, cav_sin, None)
+    v2a = attn_nf4(mna, mnh, ws, base + 80, table2, aheads, caa_cos, caa_sin, cav_cos, cav_sin, None)
     ah = ah + acg.select(2, 0) * v2a
 
     nh = rmsn(h) * (modtok(vada, 4, cond_idx) + 1.0) + modtok(vada, 3, cond_idx)
-    h = h + ff_nf4(nh, ws, base + 96, table) * modtok(vada, 5, cond_idx)
+    h = h + ff_nf4(nh, ws, base + 96, table2) * modtok(vada, 5, cond_idx)
 
     nah = rmsn(ah) * (aada.select(2, 4) + 1.0) + aada.select(2, 3)
-    ah = ah + ff_nf4(nah, ws, base + 102, table) * aada.select(2, 5)
+    ah = ah + ff_nf4(nah, ws, base + 102, table2) * aada.select(2, 5)
 
     return (h, ah)
 
@@ -189,14 +182,14 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor,
               cav_cos: Tensor, cav_sin: Tensor, caa_cos: Tensor, caa_sin: Tensor,
               enc_mask: Optional[Tensor], aenc_mask: Optional[Tensor],
               cond_idx: Optional[Tensor],
-              ws: List[Tensor], table: Tensor,
+              ws: List[Tensor], table2: Tensor,
               n_blocks: int, heads: int, aheads: int) -> Tuple[Tensor, Tensor]:
     i = 0
     while i < n_blocks:
         h, ah = block_nf4(h, ah, enc, aenc, temb, temb_a, tcss, tcass, tcg, tcag,
                           tp, tpa, v_cos, v_sin, a_cos, a_sin,
                           cav_cos, cav_sin, caa_cos, caa_sin,
-                          enc_mask, aenc_mask, cond_idx, ws, i * 114, table, heads, aheads)
+                          enc_mask, aenc_mask, cond_idx, ws, i * 114, table2, heads, aheads)
         i += 1
     return (h, ah)
 "
@@ -269,14 +262,19 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor,
     isTRUE(block$audio_cross_attn_adaln)
 }
 
-# NF4 level table on the right device (cached per device)
-.ltx23_jit_table <- function(device) {
-    key <- paste(device$type, device$index %||% 0L, sep = "|")
+# Byte-level NF4 dequant LUT [256, 2] (hi nibble, lo nibble), cached
+# per device+dtype. Built in the compute dtype so the in-graph scale
+# multiply never promotes: float32 (CPU, tests) reproduces the old
+# fp32 dequant exactly; bf16 (GPU) accepts ~0.4% weight rounding (far
+# below NF4 quantization noise) for half the dequant traffic.
+.ltx23_jit_table <- function(device, dtype) {
+    key <- paste(device$type, device$index %||% 0L, dtype$.type(), sep = "|")
     tbl <- .ltx23_jit_env[[key]]
     if (is.null(tbl)) {
-        tbl <- torch::torch_tensor(.ltx23_nf4_table,
-                                   dtype = torch::torch_float32(),
-                                   device = device)
+        byte <- 0:255
+        pairs <- cbind(.ltx23_nf4_table[byte %/% 16L + 1L],
+                       .ltx23_nf4_table[byte %% 16L + 1L])
+        tbl <- torch::torch_tensor(pairs, dtype = dtype, device = device)
         .ltx23_jit_env[[key]] <- tbl
     }
     tbl
@@ -308,7 +306,7 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor,
     # unname: an nn_module_list yields named children, and a named R
     # list marshals to TorchScript as Dict[str, Tensor], not List[Tensor]
     ws <- unname(do.call(c, lapply(blocks, .ltx23_jit_pack_block)))
-    table <- .ltx23_jit_table(hidden_states$device)
+    table <- .ltx23_jit_table(hidden_states$device, hidden_states$dtype)
     heads <- blocks[[1]]$attn1$heads
     aheads <- blocks[[1]]$audio_attn1$heads
 

R/reshard.R

@@ -41,7 +41,7 @@ reshard_safetensors <- function(input, output_dir,
     shard <- list()
     shard_size <- 0
     shard_idx <- 0L
-    key_to_shard <- integer(0)  # 1-based shard index per key (named)
+    key_to_shard <- integer(0) # 1-based shard index per key (named)
     total_size <- 0
 
     flush <- function() {
@@ -65,10 +65,10 @@ reshard_safetensors <- function(input, output_dir,
         for (key in keys) {
             t <- handle$get_tensor(key)
             bytes <- prod(t$shape) *
-                as.integer(switch(as.character(t$dtype),
-                                  Float = 4L, Double = 8L, Half = 2L,
-                                  BFloat16 = 2L, Byte = 1L, Char = 1L,
-                                  Long = 8L, Int = 4L, 4L))
+            as.integer(switch(as.character(t$dtype),
+                              Float = 4L, Double = 8L, Half = 2L,
+                              BFloat16 = 2L, Byte = 1L, Char = 1L,
+                              Long = 8L, Int = 4L, 4L))
             if (shard_size > 0 && shard_size + bytes > shard_bytes) {
                 flush()
             }

inst/tinytest/test_jit_ltx23.R

@@ -239,3 +239,14 @@ torch::with_no_grad({
 })
 expect_true(max_abs_diff(out_ci[[1]], ref_ci[[1]]) < 1e-4)
 expect_true(max_abs_diff(out_ci[[2]], ref_ci[[2]]) < 1e-4)
+
+# The byte LUT decodes every byte to its (hi nibble, lo nibble)
+# codebook pair in order
+lut <- diffuseR:::.ltx23_jit_table(torch::torch_device("cpu"),
+  torch::torch_float32())
+expect_equal(as.integer(lut$shape), c(256L, 2L))
+nf4_levels <- diffuseR:::.ltx23_nf4_table
+for (b in c(0L, 1L, 15L, 16L, 137L, 240L, 255L)) {
+  expect_equal(as.numeric(lut[b + 1L, ]),
+    c(nf4_levels[b %/% 16L + 1L], nf4_levels[b %% 16L + 1L]))
+}