miniworld-engine 2.5.3__tar.gz → 2.5.5__tar.gz

{miniworld_engine-2.5.3/src/miniworld_engine.egg-info → miniworld_engine-2.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: miniworld-engine
-Version: 2.5.3
+Version: 2.5.5
 Summary: H100 + A100 CUDA / Triton kernels for AlphaFold3-style modules, with a PyTorch reference backend per module, accuracy and speed measurers, and a version-based judge.
 License-Expression: MIT AND Apache-2.0 AND BSD-3-Clause
 Requires-Python: >=3.10
@@ -22,8 +22,8 @@ Requires-Dist: pydantic
 Requires-Dist: pyyaml
 Requires-Dist: cuda-bindings
 Requires-Dist: ninja
-Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
-Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.5; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.5; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
 Provides-Extra: baselines
 Requires-Dist: cuequivariance-torch<0.13,>=0.12; extra == "baselines"
 Provides-Extra: flash
@@ -212,7 +212,7 @@ A paper describing the engine is in preparation. Until then, please cite the rep
   author  = {Park, Sanggeun},
   title   = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
   year    = {2026},
-  version = {2.5.3},
+  version = {2.5.5},
   note    = {TODO: replace with the paper reference}
 }
 ```

{miniworld_engine-2.5.3 → miniworld_engine-2.5.5}/README.md

@@ -176,7 +176,7 @@ A paper describing the engine is in preparation. Until then, please cite the rep
   author  = {Park, Sanggeun},
   title   = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
   year    = {2026},
-  version = {2.5.3},
+  version = {2.5.5},
   note    = {TODO: replace with the paper reference}
 }
 ```

{miniworld_engine-2.5.3 → miniworld_engine-2.5.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.5.3"
+version = "2.5.5"
 requires-python = ">=3.10"
 readme = "README.md"
 description = "H100 + A100 CUDA / Triton kernels for AlphaFold3-style modules, with a PyTorch reference backend per module, accuracy and speed measurers, and a version-based judge."
@@ -24,8 +24,8 @@ dependencies = [
   # miniworld_engine/_prebuilt/ (content-keyed paths; the few arch-independent artifacts, multi-arch fatbins / host-only glue, share a
   # key and are shipped by both, either copy serves both GPUs); elsewhere a kernel
   # JIT-builds on first use (needs a CUDA toolkit; the CUTLASS / MathDx headers are fetched by kernels/_headers.py, see docs/install.md).
-  'miniworld-engine-kernels-sm90-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
-  'miniworld-engine-kernels-sm80-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
+  'miniworld-engine-kernels-sm90-cu129==2.5.5; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
+  'miniworld-engine-kernels-sm80-cu129==2.5.5; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
 ]
 
 [project.optional-dependencies]

{miniworld_engine-2.5.3 → miniworld_engine-2.5.5}/src/miniworld_engine/judge/__init__.py

@@ -345,7 +345,17 @@ def owed(
     module classes measured again although current (a measurer fix that changes only their results)."""
     tokens, given = load_tokens(now.gpu), answers()
     tally, reasons, questions = collections.Counter(), collections.Counter(), collections.Counter()
-    todo = collections.defaultdict(lambda: {"rows": set(), "lengths": set(), "backends": set(), "dtypes": set(), "modes": set(), "n": 0})
+    todo = collections.defaultdict(
+        lambda: {
+            "rows": set(),
+            "lengths": set(),
+            "backends": set(),
+            "dtypes": set(),
+            "modes": set(),
+            "n": 0,
+            "cells": collections.Counter(),
+        }
+    )
     seen = set()
     for kind in CONFIGS:
         for rec in required(kind):
@@ -373,6 +383,7 @@ def owed(
                 t["dtypes"].add(rec["dtype"])
                 t["modes"].add(rec["mode"])
                 t["n"] += 1
+                t["cells"][(rec["row"], rec["length"])] += 1
             elif state == "pending":
                 for part, old, new in why:
                     questions[(part, str(old), str(new))] += 1
@@ -395,26 +406,49 @@ def owed(
     return 0
 
 
+def _split_cells(cells: collections.Counter, parts: int) -> list[list[tuple[int, int]]]:
+    """The owed (row, length) cells of one (kind, module) in ``parts`` balanced chunks: by row when there are enough rows, else by
+    (row, length) cell -- so one large module runs on several GPUs at once instead of holding one group."""
+    rows = collections.Counter()
+    for (row, _), n in cells.items():
+        rows[row] += n
+    units = [[c for c in cells if c[0] == row] for row in rows] if len(rows) >= parts else [[c] for c in cells]
+    chunks, load = [[] for _ in range(parts)], [0] * parts
+    for unit in sorted(units, key=lambda u: -sum(cells[c] for c in u)):
+        i = load.index(min(load))
+        chunks[i].extend(unit)
+        load[i] += sum(cells[c] for c in unit)
+    return [c for c in chunks if c]
+
+
 def write_plan(plan_dir: str, todo: dict, n_groups: int = 4) -> None:
-    """The owed run: one config per (kind, module) selecting the owed rows / lengths / backends / dtypes / modes (a superset of the
-    owed cases when they do not form a product: the extra cases are measured again, which is harmless), in ``n_groups`` balanced groups
-    (one measurement job each, run in parallel)."""
+    """The owed run: configs selecting the owed rows / lengths / backends / dtypes / modes (a superset of the owed cases when they do
+    not form a product: the extra cases are measured again, which is harmless), in ``n_groups`` balanced groups (one measurement job
+    each, run in parallel). A (kind, module) owing more than its share of the total is split into several configs (by row, else by
+    (row, length)), each its own label ``<module>.p<i>`` (its own record file and done marker), so a lone large module still fills
+    every group."""
     d = Path(plan_dir)
     (d / "cfg").mkdir(parents=True, exist_ok=True)
+    total = sum(t["n"] for t in todo.values())
+    share = max(1, -(-total // max(1, n_groups)))
     jobs = []
     for (kind, module), t in sorted(todo.items()):
-        base = yaml.safe_load(CONFIGS[kind].read_text())
-        base.update(
-            modules=[module],
-            rows=sorted(t["rows"]),
-            lengths=sorted(t["lengths"]),
-            backends=sorted(t["backends"] | {"pytorch"}),
-            dtypes=sorted(t["dtypes"]),
-            modes=sorted(t["modes"]),
-        )
-        path = d / "cfg" / f"{kind}.{module}.yaml"
-        path.write_text(yaml.safe_dump(base, sort_keys=False))
-        jobs.append((t["n"], kind, module, str(path)))
+        parts = min(n_groups, -(-t["n"] // share), len(t["cells"]) or 1)
+        chunks = _split_cells(t["cells"], parts) if parts > 1 else [list(t["cells"])]
+        for i, cells in enumerate(chunks, 1):
+            label = module if len(chunks) == 1 else f"{module}.p{i}"
+            base = yaml.safe_load(CONFIGS[kind].read_text())
+            base.update(
+                modules=[module],
+                rows=sorted({r for r, _ in cells}) if cells else sorted(t["rows"]),
+                lengths=sorted({n for _, n in cells}) if cells else sorted(t["lengths"]),
+                backends=sorted(t["backends"] | {"pytorch"}),
+                dtypes=sorted(t["dtypes"]),
+                modes=sorted(t["modes"]),
+            )
+            path = d / "cfg" / f"{kind}.{label}.yaml"
+            path.write_text(yaml.safe_dump(base, sort_keys=False))
+            jobs.append((sum(t["cells"][c] for c in cells) if cells else t["n"], kind, label, str(path)))
     groups, load = [[] for _ in range(n_groups)], [0] * n_groups
     for n, kind, module, path in sorted(jobs, reverse=True):
         i = load.index(min(load))
@@ -422,7 +456,7 @@ def write_plan(plan_dir: str, todo: dict, n_groups: int = 4) -> None:
         load[i] += n
     for i, g in enumerate(groups, 1):
         (d / f"G{i}").write_text("".join(line + "\n" for line in g))
-    print(f"[judge] plan: {len(jobs)} (kind, module) runs in {plan_dir}, groups of {load} owed cases")
+    print(f"[judge] plan: {len(jobs)} runs in {plan_dir}, groups of {load} owed cases")
 
 
 # ------------------------------------------------------------------------------------------------------------------- sign, status

{miniworld_engine-2.5.3 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/_capture.py

@@ -189,6 +189,8 @@ def lookup_inputs(store: dict, key, build, limit: int = 64, valid=None, alive=No
         cid = _graph_capture_id()
         if cid == 0:
             return build()
+    if alive is not None:
+        _INPUT_STORES[id(store)] = (store, alive)
     sk = (cid, key)
     pinned = _pin_if_static_capture(store, sk)
     hit = store.get(sk)
@@ -206,6 +208,30 @@ def lookup_inputs(store: dict, key, build, limit: int = 64, valid=None, alive=No
     return out
 
 
+#: id(store) -> (the store, its ``alive(entry)``): every per-input cache ``lookup_inputs`` filled with an ``alive`` test, for
+#: ``release_dead_inputs``.
+_INPUT_STORES: dict[int, tuple[dict, object]] = {}
+
+
+def release_dead_inputs() -> int:
+    """Drop every per-input cache entry (``lookup_inputs`` / ``hoisted_table``: a hoisted pair bias, conditioning tables) whose input
+    tensor is gone; returns how many were dropped. ``lookup_inputs`` drops them only at its next call on the same store, so a table of a
+    finished item (a token DiT's hoisted pair bias: [heads x blocks, S, S], 11.3 GiB for OpenDDE's 3968 structural tokens) otherwise
+    outlives its item until the next item reaches the same module -- through the whole next trunk. A graph that read such an entry read
+    its input too, so it is no longer replayable either (as at ``lookup_inputs``'s own drop). Cheap (a weak-reference test per entry);
+    called by :class:`~miniworld_engine.modules.PairformerBlock` on every eager call, and by callers between items."""
+    dropped = 0
+    for store, alive in list(_INPUT_STORES.values()):
+        dead = [k for k, e in store.items() if not alive(e)]  # type: ignore[operator]
+        pinned = _PINNED.get(id(store))
+        for k in dead:
+            del store[k]
+            if pinned is not None:
+                pinned.discard(k)
+        dropped += len(dead)
+    return dropped
+
+
 #: (scope, (tag, ((pointer, version, dtype, shape), ...))) -> (weak references to the parameters, the forms). Also the Transition's ``Ws^T``
 #: (``transition/cuda/_weights.py``) and what FoldForge clears before it loads a model.
 _FORMS: dict = {}

miniworld_engine-2.5.5/src/miniworld_engine/kernels/_wide.py

@@ -0,0 +1,72 @@
+"""64-bit index arithmetic for Triton kernels, paid only where a tensor can pass 2**31 elements.
+
+Triton computes ``tl.program_id`` / ``tl.arange`` / loop counters and every integer argument below 2**31 in int32, and an int32 product
+wraps BEFORE it reaches the pointer. A kernel whose per-program base is int64 (``program_id(..).to(tl.int64)``) is still exposed through a
+separate in-tile term such as ``cols[None, :] * stride_col``: on a channel-major (m-major) view the channel stride is the row count
+M = B * L * L, and ``K * M`` passes 2**31 at the pair sizes the engine serves (d 512 at L 2048, OpenDDE's refiner pair d 384 at
+S 2944 / 3968). Casting that term on every call costs the hot loop (A100 TriMul output projection d 384: 4-5 %), so the launcher decides
+with ``needs_wide`` and passes ``WIDE`` as a constexpr: below the bound the kernel compiles to the 32-bit code it always was.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+INT32_LIMIT = 2**31
+
+
+@triton.jit
+def wide(x, WIDE: tl.constexpr):
+    """``x`` as int64 when the launcher set WIDE (an address product can pass 2**31), else unchanged."""
+    return x.to(tl.int64) if WIDE else x  # WIDE is a constexpr: the untaken branch is never traced
+
+
+def max_offset(t: torch.Tensor) -> int:
+    """The largest element offset a kernel forms when it walks ``t`` at its own strides (sum of (size - 1) * |stride|)."""
+    return sum((s - 1) * abs(st) for s, st in zip(t.shape, t.stride(), strict=True) if s > 0)
+
+
+def needs_wide(*items: torch.Tensor | int | None) -> bool:
+    """True when any tensor's largest offset (or any int extent product given directly) reaches 2**31."""
+    for it in items:
+        if it is None:
+            continue
+        n = max_offset(it) if isinstance(it, torch.Tensor) else int(it)
+        if n >= INT32_LIMIT:
+            return True
+    return False
+
+
+@triton.jit
+def row_tile(N_CTX, BLOCK: tl.constexpr, FLAT: tl.constexpr):
+    """``(row, tile)`` of the grid ``row_tile_grid`` launches over ``rows x cdiv(N_CTX, BLOCK)`` tiles; ``row`` is int64.
+
+    CUDA caps grid axes y and z at 65535, so a per-row axis there (attention over B * H * L rows: H 16 at L 4096 already passes it) fails
+    to launch; axis x takes 2**31 - 1. Up to 65535 rows the grid is (tiles, 1, rows) as it always was; past it (``FLAT``, set by the
+    launcher) it is one axis x of rows * tiles programs, row-major, so the tile index stays the fastest-varying one and every row's tiles
+    still run back to back. The flat decode (an integer division per program) cost ~1 % on the H100 attention backward, hence the switch."""
+    if FLAT:
+        n_tiles = tl.cdiv(N_CTX, BLOCK)
+        pid = tl.program_id(0)
+        row = (pid // n_tiles).to(tl.int64)
+        tile = pid % n_tiles
+    else:
+        row = tl.program_id(2).to(tl.int64)
+        tile = tl.program_id(0)
+    return row, tile
+
+
+GRID_YZ_MAX = 65535
+
+
+def rows_flat(rows: int) -> bool:
+    """The ``FLAT`` constexpr of ``row_tile``: True when the rows do not fit grid axis z."""
+    return rows > GRID_YZ_MAX
+
+
+def row_tile_grid(rows: int, n_ctx: int, block: int) -> list[int]:
+    """The grid ``row_tile`` decodes: ``(tiles, 1, rows)``, or ``rows * tiles`` programs on axis x past 65535 rows."""
+    tiles = (n_ctx + block - 1) // block
+    return [rows * tiles] if rows_flat(rows) else [tiles, 1, rows]

{miniworld_engine-2.5.3 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/adaln/triton/inference.py

@@ -23,6 +23,7 @@ eager/compiled path runs — one launch, cond_aff read once, TF32 policy for fp3
 from __future__ import annotations
 
 from miniworld_engine.kernels._launch_config import configured
+from miniworld_engine.kernels._wide import needs_wide, wide
 
 # The weighted-LayerNorm kernel that used to live here was fused3.py's `_ln_kernel` with
 # HAS_W=True -- bitwise equal on the output (.bench/direct.out). Imported now.
@@ -135,6 +136,7 @@ def _cond_affine(
         0,  # no second reduction
         HAS_W=True,
         shape_key=pack(shape_key, N=N),
+        WIDE=needs_wide(cond, aff),
     )
     return aff
 
@@ -194,6 +196,7 @@ def _cond_affine_stats(
         HAS_W=True,
         shape_key=pack(shape_key, N=N, N2=N2),
         HAS_STATS=True,
+        WIDE=needs_wide(cond, aff, x),
     )
     return aff, rstd, c1
 
@@ -204,7 +207,22 @@ def _cond_affine_stats(
 @configured("adaln_epilogue_triton", key=["shape_key"], default=dict(BLOCK_K=64, BLOCK_M1=64, num_warps=2, num_stages=2))
 @triton.jit
 def _adaln_epilogue_kernel(
-    X, SB, Y, M, N: tl.constexpr, eps, sx0, sx1, ss0, ss1, sy0, sy1, BLOCK_M1: tl.constexpr, BLOCK_K: tl.constexpr, shape_key
+    X,
+    SB,
+    Y,
+    M,
+    N: tl.constexpr,
+    eps,
+    sx0,
+    sx1,
+    ss0,
+    ss1,
+    sy0,
+    sy1,
+    BLOCK_M1: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    shape_key,
+    WIDE: tl.constexpr = False,  # 64-bit column offsets (launcher: needs_wide on the activations)
 ):
     # SB is (M, 2N): cols [0:N] = scale (incl. scale_b), [N:2N] = bias.
     row = tl.program_id(0).to(tl.int64)
@@ -224,23 +242,23 @@ def _adaln_epilogue_kernel(
     if BLOCK_K >= N:
         cols = tl.arange(0, BLOCK_K)
         mask = rmask[:, None] & (cols[None, :] < N)
-        x = tl.load(X + rm[:, None] * sx0 + cols[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
+        x = tl.load(X + rm[:, None] * sx0 + wide(cols, WIDE)[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
         mean = tl.sum(x, axis=1) / N
         xc = tl.where(mask, x - mean[:, None], 0.0)
         var = tl.sum(xc * xc, axis=1) / N
         rstd = 1.0 / tl.sqrt(var + eps)
         x_hat = xc * rstd[:, None]
-        scale = tl.load(SB + rm[:, None] * ss0 + cols[None, :] * ss1, mask=mask, other=0.0).to(tl.float32)
-        bias = tl.load(SB + rm[:, None] * ss0 + (cols[None, :] + N) * ss1, mask=mask, other=0.0).to(tl.float32)
+        scale = tl.load(SB + rm[:, None] * ss0 + wide(cols, WIDE)[None, :] * ss1, mask=mask, other=0.0).to(tl.float32)
+        bias = tl.load(SB + rm[:, None] * ss0 + (wide(cols, WIDE)[None, :] + N) * ss1, mask=mask, other=0.0).to(tl.float32)
         y = tl.sigmoid(scale) * x_hat + bias
-        tl.store(Y + rm[:, None] * sy0 + cols[None, :] * sy1, y.to(Y.dtype.element_ty), mask=mask)
+        tl.store(Y + rm[:, None] * sy0 + wide(cols, WIDE)[None, :] * sy1, y.to(Y.dtype.element_ty), mask=mask)
     else:
         s = tl.zeros([BLOCK_M1], dtype=tl.float32)
         ss = tl.zeros([BLOCK_M1], dtype=tl.float32)
         for n0 in range(0, N, BLOCK_K):
             cols = n0 + tl.arange(0, BLOCK_K)
             mask = rmask[:, None] & (cols[None, :] < N)
-            x = tl.load(X + rm[:, None] * sx0 + cols[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
+            x = tl.load(X + rm[:, None] * sx0 + wide(cols, WIDE)[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
             s += tl.sum(x, axis=1)
             ss += tl.sum(x * x, axis=1)
         mean = s / N
@@ -249,12 +267,12 @@ def _adaln_epilogue_kernel(
         for n0 in range(0, N, BLOCK_K):
             cols = n0 + tl.arange(0, BLOCK_K)
             mask = rmask[:, None] & (cols[None, :] < N)
-            x = tl.load(X + rm[:, None] * sx0 + cols[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
+            x = tl.load(X + rm[:, None] * sx0 + wide(cols, WIDE)[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
             x_hat = (x - mean[:, None]) * rstd[:, None]
-            scale = tl.load(SB + rm[:, None] * ss0 + cols[None, :] * ss1, mask=mask, other=0.0).to(tl.float32)
-            bias = tl.load(SB + rm[:, None] * ss0 + (cols[None, :] + N) * ss1, mask=mask, other=0.0).to(tl.float32)
+            scale = tl.load(SB + rm[:, None] * ss0 + wide(cols, WIDE)[None, :] * ss1, mask=mask, other=0.0).to(tl.float32)
+            bias = tl.load(SB + rm[:, None] * ss0 + (wide(cols, WIDE)[None, :] + N) * ss1, mask=mask, other=0.0).to(tl.float32)
             y = tl.sigmoid(scale) * x_hat + bias
-            tl.store(Y + rm[:, None] * sy0 + cols[None, :] * sy1, y.to(Y.dtype.element_ty), mask=mask)
+            tl.store(Y + rm[:, None] * sy0 + wide(cols, WIDE)[None, :] * sy1, y.to(Y.dtype.element_ty), mask=mask)
 
 
 def _adaln_epilogue_fake(x, sb, eps, shape_key=None):
@@ -294,6 +312,7 @@ def _adaln_epilogue(x: torch.Tensor, sb: torch.Tensor, eps: float, shape_key: in
         y.stride(0),
         y.stride(1),
         shape_key=pack(shape_key, N=N),
+        WIDE=needs_wide(x, sb, y),
     )
     return y
 
@@ -346,6 +365,7 @@ def _adaln_gemm_gate_kernel(
     GROUP_M: tl.constexpr,
     shape_key,
     PRE: tl.constexpr = False,
+    WIDE: tl.constexpr = False,  # 64-bit row offsets (launcher: needs_wide on the activations)
 ):
     """y = sigmoid(CondN @ SW + SB) * (x*rstd - c1) + CondN @ BW, one program per output tile.
 
@@ -364,7 +384,7 @@ def _adaln_gemm_gate_kernel(
     block looping over all of NX. That was worth 0.776 -> 0.580 ms on its own: a per-row program
     serialises NX/BLOCK_N iterations inside one CTA, which is parallelism left on the floor.
     """
-    pid_m, pid_n = tile_order(tl.program_id(0), tl.cdiv(M, BLOCK_M), tl.cdiv(NX, BLOCK_N), GROUP_M)
+    pid_m, pid_n = tile_order(wide(tl.program_id(0), WIDE), tl.cdiv(M, BLOCK_M), tl.cdiv(NX, BLOCK_N), GROUP_M)
     rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
     cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     rmask = rows < M
@@ -447,6 +467,7 @@ def _adaln_gemm_gate(
         y.stride(0),
         shape_key=pack(shape_key, NX=NX, NC=NC),
         PRE=pre,
+        WIDE=needs_wide(cond_n, x, y),
     )
     return y
 
@@ -500,6 +521,7 @@ def _cond_affine_stats_saved(
         MC=mean_c,
         RC=rstd_c,
         SAVE_C=True,
+        WIDE=needs_wide(cond, aff, x),
     )
     return aff, rstd, c1, mean_c, rstd_c
 
@@ -591,6 +613,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
     BLOCK_K_NC: tl.constexpr,
     shape_key,
     PRE: tl.constexpr = False,
+    WIDE: tl.constexpr = False,  # 64-bit column offsets (launcher: needs_wide on the activations)
 ):
     pid = tl.program_id(0).to(tl.int64)
     rows = pid * BLOCK_M1 + tl.arange(0, BLOCK_M1)
@@ -601,7 +624,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
     for cs in range(0, NC, BLOCK_K_NC):
         cc = cs + tl.arange(0, BLOCK_K_NC)
         cm = row_mask[:, None] & (cc[None, :] < NC)
-        co = rows[:, None] * scr + cc[None, :] * scc
+        co = rows[:, None] * scr + wide(cc, WIDE)[None, :] * scc
         v = tl.load(Cond + co, mask=cm, other=0.0).to(tl.float32)
         mean_c += tl.sum(tl.where(cm, v, 0.0), axis=1)
     mean_c /= NC
@@ -609,7 +632,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
     for cs in range(0, NC, BLOCK_K_NC):
         cc = cs + tl.arange(0, BLOCK_K_NC)
         cm = row_mask[:, None] & (cc[None, :] < NC)
-        co = rows[:, None] * scr + cc[None, :] * scc
+        co = rows[:, None] * scr + wide(cc, WIDE)[None, :] * scc
         v = tl.load(Cond + co, mask=cm, other=0.0).to(tl.float32)
         vc = tl.where(cm, v - mean_c[:, None], 0.0)
         var_c += tl.sum(vc * vc, axis=1)
@@ -620,7 +643,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
     for xs in range(0, NX, BLOCK_K_NX):
         xc = xs + tl.arange(0, BLOCK_K_NX)
         xm = row_mask[:, None] & (xc[None, :] < NX)
-        xo = rows[:, None] * sxr + xc[None, :] * sxc
+        xo = rows[:, None] * sxr + wide(xc, WIDE)[None, :] * sxc
         v = tl.load(X + xo, mask=xm, other=0.0).to(tl.float32)
         mean_x += tl.sum(tl.where(xm, v, 0.0), axis=1)
     mean_x /= NX
@@ -628,7 +651,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
     for xs in range(0, NX, BLOCK_K_NX):
         xc = xs + tl.arange(0, BLOCK_K_NX)
         xm = row_mask[:, None] & (xc[None, :] < NX)
-        xo = rows[:, None] * sxr + xc[None, :] * sxc
+        xo = rows[:, None] * sxr + wide(xc, WIDE)[None, :] * sxc
         v = tl.load(X + xo, mask=xm, other=0.0).to(tl.float32)
         vc = tl.where(xm, v - mean_x[:, None], 0.0)
         var_x += tl.sum(vc * vc, axis=1)
@@ -638,7 +661,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
     for xs in range(0, NX, BLOCK_K_NX):
         xc = xs + tl.arange(0, BLOCK_K_NX)
         xm = row_mask[:, None] & (xc[None, :] < NX)
-        xo = rows[:, None] * sxr + xc[None, :] * sxc
+        xo = rows[:, None] * sxr + wide(xc, WIDE)[None, :] * sxc
         xv = tl.load(X + xo, mask=xm, other=0.0).to(tl.float32)
         x_hat = (xv - mean_x[:, None]) * rstd_x[:, None]
 
@@ -648,7 +671,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
         for cs in range(0, NC, BLOCK_K_NC):
             cc = cs + tl.arange(0, BLOCK_K_NC)
             cm = row_mask[:, None] & (cc[None, :] < NC)
-            co = rows[:, None] * scr + cc[None, :] * scc
+            co = rows[:, None] * scr + wide(cc, WIDE)[None, :] * scc
             v = tl.load(Cond + co, mask=cm, other=0.0).to(tl.float32)
             cond_norm = (v - mean_c[:, None]) * rstd_c[:, None]
             lnw = tl.load(LnW + cc, mask=cc < NC, other=0.0).to(tl.float32)
@@ -667,7 +690,7 @@ def _adaln_fused_kernel(  # noqa: PLR0915
             bias += tf32_dot(cond_aff, bw, out_dtype=tl.float32, pre_b=PRE)
         scale += scale_b[None, :]
         y = tl.sigmoid(scale) * x_hat + bias
-        yo = rows[:, None] * syr + xc[None, :] * syc
+        yo = rows[:, None] * syr + wide(xc, WIDE)[None, :] * syc
         tl.store(Y + yo, y.to(Y.dtype.element_ty), mask=xm)
 
 
@@ -731,6 +754,7 @@ def adaln_inference_fused(
         PRE=pre,
         # L is x's pre-flatten shape[-2]; m = B*A is the row count, not the shape.
         shape_key=atom_key(length_of(orig_x_shape), NX=nx, NC=nc),
+        WIDE=needs_wide(x2d, cond2d, y),
     )
     return y.reshape(orig_x_shape)
 

{miniworld_engine-2.5.3 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/adaln/triton/ln_strided.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 from miniworld_engine.kernels._compile import opaque
 from miniworld_engine.kernels._launch_config import configured
+from miniworld_engine.kernels._wide import wide
 
 import torch
 import triton
@@ -75,6 +76,7 @@ def _ln_kernel(
     MC=None,
     RC=None,
     SAVE_C: tl.constexpr = False,
+    WIDE: tl.constexpr = False,  # 64-bit column offsets (launcher: needs_wide on the activations)
 ):
     row = tl.program_id(0).to(tl.int64)
     rm = row * BLOCK_M1 + tl.arange(0, BLOCK_M1)
@@ -95,7 +97,7 @@ def _ln_kernel(
         cols = tl.arange(0, BLOCK_K)
         cmask = cols < N
         mask = rmask[:, None] & cmask[None, :]
-        x = tl.load(X + rm[:, None] * sx0 + cols[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
+        x = tl.load(X + rm[:, None] * sx0 + wide(cols, WIDE)[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
         mean = tl.sum(x, axis=1) / N
         xc = tl.where(mask, x - mean[:, None], 0.0)
         var = tl.sum(xc * xc, axis=1) / N
@@ -107,14 +109,14 @@ def _ln_kernel(
         if HAS_W:
             w = tl.load(W + cols, mask=cmask, other=0.0).to(tl.float32)[None, :]
             xn = xn * w
-        tl.store(Y + rm[:, None] * sy0 + cols[None, :] * sy1, xn.to(Y.dtype.element_ty), mask=mask)
+        tl.store(Y + rm[:, None] * sy0 + wide(cols, WIDE)[None, :] * sy1, xn.to(Y.dtype.element_ty), mask=mask)
     else:
         s = tl.zeros([BLOCK_M1], dtype=tl.float32)
         ss = tl.zeros([BLOCK_M1], dtype=tl.float32)
         for n0 in range(0, N, BLOCK_K):
             cols = n0 + tl.arange(0, BLOCK_K)
             mask = rmask[:, None] & (cols[None, :] < N)
-            x = tl.load(X + rm[:, None] * sx0 + cols[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
+            x = tl.load(X + rm[:, None] * sx0 + wide(cols, WIDE)[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
             s += tl.sum(x, axis=1)
             ss += tl.sum(x * x, axis=1)
         mean = s / N
@@ -127,12 +129,12 @@ def _ln_kernel(
             cols = n0 + tl.arange(0, BLOCK_K)
             cmask = cols < N
             mask = rmask[:, None] & cmask[None, :]
-            x = tl.load(X + rm[:, None] * sx0 + cols[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
+            x = tl.load(X + rm[:, None] * sx0 + wide(cols, WIDE)[None, :] * sx1, mask=mask, other=0.0).to(tl.float32)
             xn = (x - mean[:, None]) * rstd[:, None]
             if HAS_W:
                 w = tl.load(W + cols, mask=cmask, other=0.0).to(tl.float32)[None, :]
                 xn = xn * w
-            tl.store(Y + rm[:, None] * sy0 + cols[None, :] * sy1, xn.to(Y.dtype.element_ty), mask=mask)
+            tl.store(Y + rm[:, None] * sy0 + wide(cols, WIDE)[None, :] * sy1, xn.to(Y.dtype.element_ty), mask=mask)
     if HAS_STATS:
         # a second row-wise reduction in the same launch: the LayerNorm statistics of X2 (M, N2), as (rstd, mean * rstd) -- what
         # `layernorm_stats_triton` returns (the same sums, the same formula), without its own launch