fix(kernel): unbreak the nightly kernel-cache wheel build

python/freetoken/kernel/aot_models.py

@@ -63,6 +63,13 @@ class AotModel:
     arch_aliases: tuple[str, ...] = ()
 
 
+def fp8_block_scale_pad(rows: int, cols: int) -> int:
+    """Trailing scale-bank dim padded so per-expert row bytes are 16B-aligned (fused copy)."""
+    while (rows * cols * 2) % 16:
+        cols += 1
+    return cols
+
+
 def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int) -> dict[str, int]:
     """Per-expert row bytes for each offload bank a format registers.
 
@@ -77,8 +84,6 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
     if fmt == "fp8_block":
         # qwen3_5_moe/weight.py _build_fp8_expert_banks: fp8 weights + bf16 128x128 block
         # scales, trailing scale dim 16B-padded (same helper as the loader)
-        from freetoken.moe.offload_cache import fp8_block_scale_pad
-
         B = 128
         return {
             "gate_up": 2 * I * H,
@@ -417,7 +422,8 @@ def aggregate_fast_index_copy_feature_sizes() -> tuple[int, ...]:
     sizes: set[int] = set(TEST_FEATURE_SIZES)
     for model in SUPPORTED_MODELS:
         sizes.update(fast_index_copy_feature_sizes(model))
-    return tuple(sorted(sizes))
+    # the per-bank kernel copies rows in fixed 128-byte steps; other sizes cannot compile
+    return tuple(sorted(size for size in sizes if size % 128 == 0))
 
 
 __all__ = [

python/freetoken/moe/offload_cache.py

@@ -77,11 +77,8 @@
     "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"),
 }
 
-def fp8_block_scale_pad(rows: int, cols: int) -> int:
-    """Trailing scale-bank dim padded so per-expert row bytes are 16B-aligned (fused copy)."""
-    while (rows * cols * 2) % 16:
-        cols += 1
-    return cols
+# lives in kernel/aot_models.py: the AOT row table shares it and must stay importable in the torch-only kernel-cache build env, which cannot import freetoken.moe
+from freetoken.kernel.aot_models import fp8_block_scale_pad
 
 
 # bytes per (expert, layer) as f(hidden, moe_intermediate), from the bank shapes above; keep in sync with _BANK_SCHEMAS
@@ -350,6 +347,19 @@ def set_bank_sources(
             self._init_prefill_overlap_buffers()
 
     def _build_copy_plan(self) -> None:
+        self._build_fused_copy_plan()
+        if self._copy_fused_ok or self.device.type != "cuda" or not self.banks:
+            return
+        for name in self.bank_schema:
+            cache = self.bank_caches[name]
+            feat = math.prod(cache.shape[1:]) * cache.element_size()
+            if feat % 128:
+                raise RuntimeError(
+                    f"MoE bank {name!r} rows are {feat} bytes (not a multiple of 128): "
+                    f"only the fused multi-bank copy can move them, but it is disabled"
+                )
+
+    def _build_fused_copy_plan(self) -> None:
         """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes).
 
         Built once here (and on :meth:`rebuild`, which reallocates the slot caches);