Lowrank sum params

src/pyjuice/__init__.py

@@ -13,7 +13,7 @@
 from pyjuice.model import compile, TensorCircuit
 
 # TensorCircuit layers
-from pyjuice.layer import InputLayer, ProdLayer, SumLayer
+from pyjuice.layer import InputLayer, ProdLayer, SumLayer, ExternalParamsSumLayer
 
 # Construction methods
 from pyjuice.nodes import multiply, summate, inputs, set_block_size, structural_properties

src/pyjuice/io/serialization.py

@@ -5,7 +5,7 @@
 from functools import partial
 from typing import Sequence
 
-from pyjuice.nodes import CircuitNodes, InputNodes, ProdNodes, SumNodes, inputs, multiply, summate
+from pyjuice.nodes import CircuitNodes, InputNodes, ProdNodes, SumNodes, ExternalParamsSumNodes, inputs, multiply, summate
 
 
 def serialize_nodes(root_ns: CircuitNodes):
@@ -30,6 +30,9 @@ def serialize_nodes(root_ns: CircuitNodes):
         if ns.is_prod() or ns.is_sum():
             ns_info["edge_ids"] = ns.edge_ids.detach().cpu().numpy().copy()
 
+        if isinstance(ns, ExternalParamsSumNodes):
+            ns_info["external_params"] = pickle.dumps(ns.external_params)
+
         if hasattr(ns, "_params") and ns._params is not None:
             ns_info["params"] = ns._params.detach().cpu().numpy().copy()
 
@@ -102,7 +105,12 @@ def deserialize_nodes(nodes_list: Sequence):
             else:
                 params = None
 
-            ns = summate(*chs, edge_ids = edge_ids, params = params, block_size = block_size)
+            external_params = ns_info.get("external_params", None)
+            if external_params is not None:
+                external_params = pickle.loads(external_params)
+
+            ns = summate(*chs, edge_ids = edge_ids, params = params, block_size = block_size,
+                         external_params = external_params)
 
             if "zero_param_mask" in ns_info:
                 zero_param_mask = torch.from_numpy(ns_info["zero_param_mask"])

src/pyjuice/layer/__init__.py

@@ -2,4 +2,7 @@
 from .input_layer import InputLayer
 from .prod_layer import ProdLayer
 from .sum_layer import SumLayer
+from .external_sum_layer import ExternalParamsSumLayer, ExternalNodeInfo, StagedExternalParams, \
+                               EXTERNAL_PARAMS_BUFFER_KWARG, EXTERNAL_PARAMS_GRAD_BUFFER_KWARG, \
+                               EXTERNAL_PARAMS_KWARG, EXTERNAL_PARAMS_GRAD_KWARG
 from .layer_group import LayerGroup

src/pyjuice/layer/input_layer.py

@@ -854,6 +854,11 @@ def add_missing_flows(self, node_flows: torch.Tensor, node_mars: Optional[torch.
                 TILE_SIZE_K = 1
 
             if _pre_accum_nflows:
+                # This launch walks the SOURCE nodes -- it indexes `source_nids`, which has one entry
+                # per source node, not per node -- so its bound is `layer_num_source_nodes`. Bounding
+                # it by `layer_num_nodes` instead reads `source_nids` past its end whenever the layer
+                # has tied nodes, and the garbage node ids it returns are then used to gather
+                # `node_flows` / `s_pids` / `s_pfids` and to scatter into `param_flows`.
                 layer_num_source_nodes = self.source_nids.size(0)
 
                 self._accum_flows_to_tied_nodes(node_flows, logspace_flows)
@@ -871,7 +876,7 @@ def add_missing_flows(self, node_flows: torch.Tensor, node_mars: Optional[torch.
                     source_nids_ptr = self.source_nids,
                     scale = scale,
                     node_offset = node_offset,
-                    layer_num_nodes = layer_num_nodes,
+                    num_active_nodes = layer_num_source_nodes,
                     logspace_flows = logspace_flows,
                     are_source_nodes = True,
                     BLOCK_SIZE = BLOCK_SIZE,
@@ -892,7 +897,7 @@ def add_missing_flows(self, node_flows: torch.Tensor, node_mars: Optional[torch.
                     source_nids_ptr = self.source_nids,
                     scale = scale,
                     node_offset = node_offset,
-                    layer_num_nodes = layer_num_nodes,
+                    num_active_nodes = layer_num_nodes,
                     logspace_flows = logspace_flows,
                     are_source_nodes = False,
                     BLOCK_SIZE = BLOCK_SIZE,
@@ -1221,13 +1226,13 @@ def _flows_kernel_template(flow_fn, params_ptr, param_flows_ptr, node_flows_ptr,
 
     @staticmethod
     def _missing_flows_kernel_template(node_flows_ptr, params_ptr, param_flows_ptr, s_pids_ptr, s_pfids_ptr,
-                                       metadata_ptr, s_mids_ptr, source_nids_ptr, scale, node_offset: tl.constexpr, layer_num_nodes: tl.constexpr, 
+                                       metadata_ptr, s_mids_ptr, source_nids_ptr, scale, node_offset: tl.constexpr, num_active_nodes: tl.constexpr,
                                        logspace_flows: tl.constexpr, are_source_nodes: tl.constexpr, BLOCK_SIZE: tl.constexpr, TILE_SIZE_K: tl.constexpr):
         pid = tl.program_id(axis = 0)
         block_start = pid * BLOCK_SIZE
 
         offsets = block_start + tl.arange(0, BLOCK_SIZE)
-        mask = offsets < layer_num_nodes
+        mask = offsets < num_active_nodes
 
         # Get the local node ids
         if are_source_nodes:

src/pyjuice/layer/kernels/c/ele_backward_sum.cu

@@ -26,6 +26,7 @@
 #include <cuda_fp16.h>
 #include <cuda.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <cute/tensor.hpp>
 #include <cute/atom/mma_atom.hpp>
 #include <cute/atom/copy_atom.hpp>
@@ -200,6 +201,7 @@ void ele_backward_sum(torch::Tensor element_flows, torch::Tensor element_mars, t
             element_flows.data_ptr<float>(), element_mars.data_ptr<float>(), params.data_ptr<float>(),
             chids.data_ptr<long>(), ebase.data_ptr<long>(), pbase.data_ptr<long>(),
             (int)batch, (int)BSK, (int)BSM, (int)knt, off, g_descNf, g_descNm);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
     }
 }
 

src/pyjuice/layer/kernels/c/input_layers/__init__.py

@@ -25,7 +25,10 @@ def _jit_plain(name: str, source_file: str):
     if not ENABLE_CUDA_KERNELS or not torch.cuda.is_available():
         return None
     cc = torch.cuda.get_device_capability()
-    cuda_cflags = ["-O3", f"-arch=sm_{cc[0]}{cc[1]}a", "--use_fast_math", "-DNDEBUG"]
+    # NO `a` suffix. These kernels use nothing arch-specific, and the suffix only exists for the
+    # accelerated targets (90a, 100a, 120a, ...) -- there is no `sm_80a` or `sm_86a`, so asking for one
+    # fails the compile on every pre-Hopper GPU and silently drops the layer onto the Triton path.
+    cuda_cflags = ["-O3", f"-arch=sm_{cc[0]}{cc[1]}", "--use_fast_math", "-DNDEBUG"]
     from pyjuice.utils.cuda_ext import jit_load
     try:
         return jit_load(name, sources=[os.path.join(_THIS_DIR, source_file)],
@@ -48,10 +51,14 @@ def cat_backward_is_available() -> bool:
 
 def cat_backward(param_flows: torch.Tensor, node_flows: torch.Tensor, data: torch.Tensor,
                  vids: torch.Tensor, s_pfids: torch.Tensor, layer_num_nodes: int, batch_size: int,
-                 node_offset: int, num_cats: int, logspace: bool) -> None:
+                 node_offset: int, num_cats: int, logspace: bool) -> bool:
     """Categorical input-layer backward: param_flows[node, cat] += sum_b (exp if logspace) node_flows
     where data[vid[node], b] == cat. The flush is NON-ATOMIC, so the caller must guarantee distinct
     (untied), 4-aligned ``s_pfids`` and num_cats/batch divisible by 4 (see the gate in
-    ``input_layer.backward``)."""
-    _cat_module.cat_backward(param_flows, node_flows, data, vids, s_pfids, int(layer_num_nodes),
-                             int(batch_size), int(node_offset), int(num_cats), int(logspace))
+    ``input_layer.backward``).
+
+    Returns False WITHOUT running if the ``num_cats``-float histogram exceeds the shared memory this
+    device grants a block; the caller must then use the Triton path."""
+    return bool(_cat_module.cat_backward(param_flows, node_flows, data, vids, s_pfids,
+                                         int(layer_num_nodes), int(batch_size), int(node_offset),
+                                         int(num_cats), int(logspace)))

src/pyjuice/layer/kernels/c/input_layers/cat_backward.cu

@@ -13,6 +13,7 @@
 // bases (s_pfids multiples of 4, node_offset & batch_size multiples of 4) — enforced by the input_layer gate.
 #include <torch/extension.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <cuda_runtime.h>
 
 #ifndef NTH
@@ -51,17 +52,43 @@ __global__ void cat_backward_kernel(float* __restrict__ param_flows, const float
     }
 }
 
-void cat_backward(torch::Tensor param_flows, torch::Tensor node_flows, torch::Tensor data,
+// The histogram is `num_cats` floats, which for a large vocabulary is well past the 48 KB a block gets
+// by default. That ceiling is not a fixed property of the machine -- a block can opt in to most of the
+// SM's shared memory, and how much that is varies by part (64 KB, 100 KB, 227 KB) -- so it is queried,
+// and the opt-in is requested, rather than either assumed or given up on.
+//
+// Returns false when even the opt-in ceiling is too small; the caller then falls back to Triton. It is
+// a return value and not a `TORCH_CHECK` because this is an optional fast path, and a silent one before:
+// the launch was issued unopted-in and unchecked, so past ~12288 categories it failed, `param_flows`
+// was never accumulated, and the caller reported success anyway.
+bool cat_backward_fits(int num_cats) {
+    int dev = 0, smax = 0;
+    cudaGetDevice(&dev);
+    cudaDeviceGetAttribute(&smax, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+    return (size_t)num_cats * sizeof(float) <= (size_t)smax;
+}
+
+
+bool cat_backward(torch::Tensor param_flows, torch::Tensor node_flows, torch::Tensor data,
                   torch::Tensor vids, torch::Tensor s_pfids, int layer_num_nodes, int batch_size,
                   long node_offset, int num_cats, int logspace) {
     int smem = num_cats * sizeof(float);
+    if (!cat_backward_fits(num_cats)) return false;
     auto stream = c10::cuda::getCurrentCUDAStream();
     float* pf = param_flows.data_ptr<float>();
     const float* nf = node_flows.data_ptr<float>();
     const long* vp = vids.data_ptr<long>();
     const long* sp = s_pfids.data_ptr<long>();
-#define LAUNCH(T) cat_backward_kernel<T><<<layer_num_nodes, NTH, smem, stream>>>( \
-        pf, nf, data.data_ptr<T>(), vp, sp, layer_num_nodes, batch_size, node_offset, num_cats, logspace)
+#define LAUNCH(T)                                                                                      \
+    do {                                                                                               \
+        if (smem > 48 * 1024)                                                                          \
+            TORCH_CHECK(cudaFuncSetAttribute(cat_backward_kernel<T>,                                   \
+                            cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess,         \
+                        "cat_backward: device refused ", smem, " B of shared memory");                 \
+        cat_backward_kernel<T><<<layer_num_nodes, NTH, smem, stream>>>(                                \
+            pf, nf, data.data_ptr<T>(), vp, sp, layer_num_nodes, batch_size, node_offset, num_cats,    \
+            logspace);                                                                                 \
+    } while (0)
     switch (data.scalar_type()) {
         case torch::kByte:  LAUNCH(uint8_t);  break;
         case torch::kChar:  LAUNCH(int8_t);   break;
@@ -70,9 +97,14 @@ void cat_backward(torch::Tensor param_flows, torch::Tensor node_flows, torch::Te
         case torch::kLong:  LAUNCH(int64_t);  break;
         default: TORCH_CHECK(false, "cat_backward: unsupported data dtype ", data.scalar_type());
     }
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
 #undef LAUNCH
+    return true;
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("cat_backward", &cat_backward, "Categorical input-layer backward (smem-histogram)");
+    m.def("cat_backward", &cat_backward,
+          "Categorical input-layer backward (smem-histogram); false if the histogram does not fit");
+    m.def("cat_backward_fits", &cat_backward_fits,
+          "Whether this device grants enough shared memory for a num_cats histogram");
 }

src/pyjuice/layer/kernels/c/par_backward_sum.cu

@@ -31,6 +31,7 @@
 #include <cuda_fp16.h>
 #include <cuda.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <cute/tensor.hpp>
 #include <cute/atom/mma_atom.hpp>
 #include <cute/atom/copy_atom.hpp>
@@ -242,7 +243,8 @@ void par_backward_sum(torch::Tensor param_flows, torch::Tensor node_flows, torch
         dim3 grid(gx, chunk);
         par_kernel<<<grid, NTH, smem, c10::cuda::getCurrentCUDAStream()>>>(param_flows.data_ptr<float>(), params.data_ptr<float>(),
             nbase.data_ptr<long>(), cbase.data_ptr<long>(), pbase.data_ptr<long>(), fbase.data_ptr<long>(),
-            (int)batch, (int)block_size, (int)num_edges, bnt, (int)use_atomic, off, g_dNf, g_dNm, g_dEm); }
+            (int)batch, (int)block_size, (int)num_edges, bnt, (int)use_atomic, off, g_dNf, g_dNm, g_dEm);
+        C10_CUDA_KERNEL_LAUNCH_CHECK(); }
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

src/pyjuice/layer/kernels/c/smallbatch_ele_backward.cu

@@ -29,6 +29,7 @@
 #include <torch/extension.h>
 #include <cuda_runtime.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <vector>
 
 // Combine two online-logsumexp accumulators (max `m`, linear-sum `l`) with a second point/accumulator
@@ -105,6 +106,7 @@ static void launch_ele(torch::Tensor element_flows, torch::Tensor element_mars,
         element_flows.data_ptr<float>(), element_mars.data_ptr<float>(), node_flows.data_ptr<float>(),
         node_mars.data_ptr<float>(), params.data_ptr<float>(), chids.data_ptr<long>(),
         ebase.data_ptr<long>(), pbase.data_ptr<long>(), batch, block_size, cs_block_size, num_edges);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
 
 // Dispatch on (WARPS from `cfg`) x (BMAX = next pow2 of batch, for register-tight state).

src/pyjuice/layer/kernels/c/smallbatch_forward_sum.cu

@@ -23,6 +23,7 @@
 #include <torch/extension.h>
 #include <cuda_runtime.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <vector>
 
 template <int SPLIT>
@@ -81,6 +82,7 @@ static void launch_sb(torch::Tensor node_mars, torch::Tensor element_mars, torch
         node_mars.data_ptr<float>(), element_mars.data_ptr<float>(), params.data_ptr<float>(),
         nids.data_ptr<long>(), ebase.data_ptr<long>(), pbase.data_ptr<long>(),
         batch, block_size, num_edges);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
 
 // cfg id -> SPLIT (edge-warps). SPLIT=32 was the sweep winner; the autotuner picks per layer/batch.

src/pyjuice/layer/kernels/c/smallbatch_par_backward.cu

@@ -24,6 +24,7 @@
 #include <torch/extension.h>
 #include <cuda_runtime.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <vector>
 
 template <int EY, int BMAX>
@@ -80,6 +81,7 @@ static void launch_par(torch::Tensor param_flows, torch::Tensor node_flows, torc
         param_flows.data_ptr<float>(), node_flows.data_ptr<float>(), node_mars.data_ptr<float>(),
         element_mars.data_ptr<float>(), params.data_ptr<float>(), nids.data_ptr<long>(),
         cids.data_ptr<long>(), pids.data_ptr<long>(), pfids.data_ptr<long>(), batch, block_size, num_edges);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
 
 // Dispatch on (EY from `cfg`) x (BMAX = next pow2 of batch, for register-tight state).

src/pyjuice/layer/kernels/c/tlmm_forward_sum.cu

@@ -24,6 +24,7 @@
 #include <cuda_bf16.h>
 #include <cuda.h>
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAException.h>
 #include <vector>
 #include <cute/tensor.hpp>
 #include <cute/atom/mma_atom.hpp>
@@ -228,6 +229,7 @@ static void launch_cfg(torch::Tensor node_mars, torch::Tensor element_mars, torc
         tlmm_kernel<BM, BN, WM, WN><<<grid, NTH, smem, c10::cuda::getCurrentCUDAStream()>>>(
             node_mars.data_ptr<float>(), params.data_ptr<float>(), nids.data_ptr<long>(),
             ebase.data_ptr<long>(), pbase.data_ptr<long>(), batch, block_size, knt, off, desc);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
     }
 }
 

src/pyjuice/layer/layer_group.py

@@ -17,8 +17,13 @@ def __init__(self, layers: Sequence[Layer]):
 
         assert len(layers) >= 1, "A `LayerGroup` must contains at least 1 layer."
 
+        # Layers of one group are evaluated interchangeably, so they must be of one kind -- but sum
+        # layers may legitimately be a mix of `SumLayer` and its external-parameter subclasses, which
+        # share its forward/backward interface and land at the same depth.
         for i in range(1, len(layers)):
-            assert type(layers[i]) == type(layers[0])
+            assert type(layers[i]) == type(layers[0]) or \
+                (isinstance(layers[i], SumLayer) and isinstance(layers[0], SumLayer)), \
+                f"A `LayerGroup` must contain layers of one kind, got {type(layers[0])} and {type(layers[i])}."
 
         if isinstance(layers[0], InputLayer):
             self.layer_type = "input"