From 395bdb0ce12d9802e5f8327d2f4d4011346545ec Mon Sep 17 00:00:00 2001
From: alint77 <ali.nt1377@gmail.com>
Date: Thu, 2 Jul 2026 09:20:02 +0200
Subject: [PATCH 1/6] feat(dion2/nordion2): selection_scope="global_capped" --
 global selection at local comm cost

Third selection scope for the FSDP2 row-sharded path, between "global" (exact
selection, full-shard comm) and "local" (fraction-scaled comm, per-shard
approximate selection):

1. all-gather the slice L1 norms only (one fp32 per row, ~KBs vs MBs of rows);
2. every rank computes the same global top-(k*world_size) threshold from the
   same gathered tensor (rank-consistent by construction);
3. each rank sends its globally-selected rows through the same fixed k-slot
   alltoall as "local" (local_comm_size=k, wire format unchanged). A rank
   owning more winners than slots defers the overflow: only selected rows get
   error-feedback decay, so deferred winners keep accumulating momentum and
   win a slot on a later step. Spare slots fill with the rank's next-best rows.

Implementation: dion2_pre_accumulate_norms (compiled: M+=G, stacked norms) +
dion2_capped_priority_async (eager async all-gather -> capacity priority,
yields in megabatch style) + a `priority` kwarg on dion2_pre_orthogonalize
that swaps the topk key and skips re-accumulation. NorDion2 reuses all three.
Unsharded/DDP falls through to the existing no-op selection path.

Measured (NorDion2 5120-hidden/8-layer, 4xH100, fraction=0.25, gram-NS):
a2a 1311 MB/step (= local; global ships 5243 MB), norms all-gather 2.5 MB +
0.02 ms, opt.step 33.2 ms (local 31.0, global 35.5), end-to-end step time and
MFU equal across scopes at this size. 2-rank functional tests: selection ==
exact global top-k when winners fit per-rank capacity; on overflow the
deferred winners are selected on the following step (EF deferral); even and
uneven shard divisions both pass; existing selection_scope test suite 11/11.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 dion/dion2.py    | 148 ++++++++++++++++++++++++++++++++++++++++++++---
 dion/nordion2.py |  27 ++++++++-
 2 files changed, 164 insertions(+), 11 deletions(-)

diff --git a/dion/dion2.py b/dion/dion2.py
index 2a853eb..470d5fe 100644
--- a/dion/dion2.py
+++ b/dion/dion2.py
@@ -1,5 +1,6 @@
 import math
 import torch
+import torch.distributed as dist
 from collections import defaultdict
 from torch import Tensor
 from torch.distributed import ProcessGroup
@@ -48,7 +49,13 @@ class Dion2(DistributedOrthoBase):
             assembled whole matrix -- layout-invariant/reproducible and better-
             converging. "local": per-shard top-k (union) -- cheaper comm but a
             sharding-dependent approximation that converges slightly worse; opt
-            in when comm-bound at large scale. No-op off the row-sharded path.
+            in when comm-bound at large scale. "global_capped": global selection
+            at local comm cost -- all-gather the row norms only (~KBs), compute
+            the global top-k everywhere, and send each rank's globally-selected
+            rows through the fixed ``k``-slot pipes; a rank owning more winners
+            than slots defers the overflow (error feedback re-selects them next
+            step), one owning fewer fills spare slots with its next-best rows.
+            No-op off the row-sharded path.
 
     Dion2 optimizer by Ahn et al.: TBD
     """
@@ -78,9 +85,10 @@ def __init__(
             raise ValueError(f"Invalid learning rate: {lr}")
         if not (0.0 < fraction <= 1.0):
             raise ValueError(f"fraction must be in (0, 1], got {fraction}")
-        if selection_scope not in ("local", "global"):
+        if selection_scope not in ("local", "global", "global_capped"):
             raise ValueError(
-                f"selection_scope must be 'local' or 'global', got {selection_scope!r}"
+                f"selection_scope must be 'local', 'global', or 'global_capped', "
+                f"got {selection_scope!r}"
             )
         if ef_decay < 0.0:
             raise ValueError(f"Invalid ef_decay: {ef_decay}")
@@ -237,6 +245,14 @@ def dion2_update_megabatch_async(
       sharding-dependent approximation of the true top-k (world-size variant).
       Cheaper comm (the win grows with model size), but converges slightly
       worse; opt in when comm-bound at large scale.
+    - ``"global_capped"``: global selection at local comm cost. The slice L1
+      norms are all-gathered (one fp32 per slice, ~KBs), every rank computes the
+      same global top-``k*world_size`` threshold, and each rank sends its
+      globally-selected slices through the same fixed ``k``-slot pipes as
+      "local". Overflow winners (a rank owning more winners than slots) are
+      deferred -- error feedback keeps their momentum accumulating so they win a
+      slot on a later step; spare slots fill with the rank's next-best slices.
+      Near-global selection quality at "local" comm cost.
 
     Off the row-sharded path (per-head, single-GPU, batch-sharded) each rank
     already holds whole matrices, so local and global selection coincide and
@@ -353,6 +369,22 @@ def dion2_update_megabatch_async(
         k = None
     global_comm_dim_size = global_dim_size
 
+    # "global_capped": accumulate momentum + all-gather the slice norms, then
+    # rank slices by the capacity priority (global winners first) instead of the
+    # local norms. Same fixed k-slot comm as "local" below.
+    priority = None
+    if (
+        selection_scope == "global_capped"
+        and comm_dim is not None
+        and process_group is not None
+    ):
+        norms = dion2_pre_accumulate_norms(
+            G=to_local(G), M=to_local(M), select_dim=select_dim
+        )
+        priority = yield from dion2_capped_priority_async(
+            norms, padded_local, k, device_rank, world_size, process_group
+        )
+
     # Pre-orthogonalize: momentum update + submatrix selection
     U_selected, indices_list = dion2_pre_orthogonalize(
         G=to_local(G),
@@ -361,6 +393,7 @@ def dion2_update_megabatch_async(
         ef_decay=ef_decay,
         select_dim=select_dim,
         k_override=k,
+        priority=priority,
     )
 
     # Orthogonalize via shared megabatch communication
@@ -426,6 +459,90 @@ def dion2_update_megabatch_async(
 )
 
 
+@_inductor_workaround
+@torch.compile(fullgraph=True)
+def dion2_pre_accumulate_norms(
+    G: List[Tensor], M: List[Tensor], select_dim: int
+) -> Tensor:
+    """
+    Phase A of the "global_capped" scope: update momentum with the gradient and
+    return the stacked slice L1 norms ``[N, local_size]`` (fp32). The caller
+    all-gathers these tiny norms across ranks (an eager collective that cannot
+    live inside this compiled graph) and feeds the resulting selection priority
+    back into ``dion2_pre_orthogonalize`` via ``priority`` -- which then must
+    NOT re-accumulate the gradient.
+    """
+    dtype = M[0].dtype
+    G = [g.to(dtype=dtype) for g in G]
+    torch._foreach_add_(M, G)
+    norm_dim = -1 if select_dim == -2 else -2
+    if M[0].size(select_dim) == 0:
+        return torch.zeros((len(M), 0), dtype=torch.float32, device=M[0].device)
+    return torch.stack(M, dim=0).norm(p=1, dim=norm_dim).float()
+
+
+def dion2_capped_priority_async(
+    norms: Tensor,  # [N, local_size] fp32 from dion2_pre_accumulate_norms
+    padded_local: int,
+    k: int,
+    device_rank: int,
+    world_size: int,
+    process_group: Optional[ProcessGroup],
+) -> Generator[None, None, Tensor]:
+    """
+    "global_capped" selection priority: all-gather the slice L1 norms (one fp32
+    per slice -- ~KBs vs MBs for the rows), compute the exact global
+    top-``k*world_size`` threshold identically on every rank, and return a
+    priority tensor ``[N, local_size]`` such that a plain local
+    ``topk(priority, k)`` implements the capacity rule:
+
+      - slices at or above the global threshold ("winners") outrank ALL
+        non-winners (a per-matrix offset larger than any norm), so a rank sends
+        as many of its winners as fit in its ``k`` slots -- overflow winners are
+        deferred, and since only *selected* slices get error-feedback decay,
+        their momentum keeps accumulating and they win a slot on a later step;
+      - spare slots fill with the rank's next-best non-winners (ordinary norm
+        order), harmless extra orthogonalization exactly like "local" scope.
+
+    Rank consistency is by construction: every rank computes the threshold from
+    the same gathered tensor. Async generator in the megabatch style: yields at
+    the collective, returns the priority.
+    """
+    N, local_size = norms.shape
+    # Pad short/empty shards to the rank-uniform padded_local. L1 norms are
+    # nonnegative, so -1 padding can never enter the winner set unless
+    # k*world_size exceeds the real row count (tiny-matrix degenerate case,
+    # where "everything is a winner" is the correct answer anyway).
+    if local_size < padded_local:
+        norms_pad = torch.nn.functional.pad(
+            norms, (0, padded_local - local_size), value=-1.0
+        )
+    else:
+        norms_pad = norms
+    gathered = torch.empty(
+        world_size * N * padded_local, dtype=torch.float32, device=norms.device
+    )
+    work = dist.all_gather_into_tensor(
+        gathered, norms_pad.reshape(-1).contiguous(),
+        group=process_group, async_op=True,
+    )
+    yield
+    work.wait()
+
+    # [world, N, padded_local] -> per-matrix global threshold = k_total-th value
+    all_norms = gathered.view(world_size, N, padded_local)
+    k_total = min(k * world_size, world_size * padded_local)
+    vals, _ = torch.topk(
+        all_norms.permute(1, 0, 2).reshape(N, -1), k_total, dim=-1
+    )
+    thresh = vals[:, -1:]
+    own = norms  # this rank's real slices (unpadded view)
+    winner = own >= thresh
+    # Offset strictly larger than any norm: winners always outrank non-winners.
+    offset = vals[:, :1] + 1.0
+    return own + winner.to(own.dtype) * offset
+
+
 @_inductor_workaround
 @torch.compile(fullgraph=True)
 def dion2_pre_orthogonalize(
@@ -435,6 +552,7 @@ def dion2_pre_orthogonalize(
     ef_decay: Tensor,
     select_dim: int,
     k_override: Optional[int] = None,
+    priority: Optional[Tensor] = None,
 ) -> Tuple[List[Tensor], List[Tensor]]:
     """
     Update momentum with gradient and compute the input to orthogonalization.
@@ -454,6 +572,13 @@ def dion2_pre_orthogonalize(
     to the available count. The pad up to ``k_override`` is not done here -- the
     downstream megabatch pads U to ``local_comm_size=k_override`` for a uniform
     alltoall, and indices stay at the real selected count.
+
+    ``priority`` (the "global_capped" scope) replaces the locally-computed L1
+    norms as the top-k ranking key ``[N, local_size]``. Passing it also means
+    the momentum was ALREADY accumulated by ``dion2_pre_accumulate_norms`` (the
+    priority derives from post-accumulation norms), so the ``M += G`` here is
+    skipped. Everything downstream (gather, EF on selected, bf16 convert) is
+    identical.
     """
     dtype = M[0].dtype
 
@@ -476,9 +601,12 @@ def dion2_pre_orthogonalize(
         k = max(1, int(math.ceil(fraction * num_select)))
     k_topk = min(k, num_select)
 
-    # Update momentum: M = M + G
-    G = [g.to(dtype=dtype) for g in G]
-    torch._foreach_add_(M, G)
+    # Update momentum: M = M + G. Skipped when `priority` is given ("global_
+    # capped"): dion2_pre_accumulate_norms already accumulated so the gathered
+    # norms reflect the post-accumulation momentum.
+    if priority is None:
+        G = [g.to(dtype=dtype) for g in G]
+        torch._foreach_add_(M, G)
 
     # Empty local shard along select_dim: FSDP2 contiguous chunking leaves this
     # rank with a size-0 shard when the param's sharded dim is smaller than
@@ -497,8 +625,12 @@ def dion2_pre_orthogonalize(
 
     M_stacked = torch.stack(M, dim=0)
 
-    # Compute L1 norm along norm_dim (sum of absolute values)
-    slice_norms = M_stacked.norm(p=1, dim=norm_dim)
+    # Top-k ranking key: the capped-scope priority when given (winners outrank
+    # all non-winners), otherwise the local L1 norms along norm_dim.
+    if priority is not None:
+        slice_norms = priority
+    else:
+        slice_norms = M_stacked.norm(p=1, dim=norm_dim)
 
     # Batched topk: indices shape (batch_size, k_topk). k_topk <= num_select is
     # guaranteed, so this never raises even on a short remainder shard.
diff --git a/dion/nordion2.py b/dion/nordion2.py
index 80eba38..f83002a 100644
--- a/dion/nordion2.py
+++ b/dion/nordion2.py
@@ -18,6 +18,8 @@
     dion2_pre_orthogonalize,
     dion2_post_orthogonalize,
     dion2_pre_accumulate,
+    dion2_pre_accumulate_norms,
+    dion2_capped_priority_async,
     _make_select_and_orthogonalize,
 )
 from .normuon import normuon_normalization_stacked, _normuon_normalization_core
@@ -55,7 +57,10 @@ class NorDion2(DistributedOrthoBase):
             assembled whole matrix -- layout-invariant/reproducible and better-
             converging. "local": per-shard top-k (union) -- cheaper comm but a
             sharding-dependent approximation that converges slightly worse; opt
-            in when comm-bound at large scale. No-op off the row-sharded path.
+            in when comm-bound at large scale. "global_capped": global selection
+            at local comm cost via a norms-only all-gather; overflow winners are
+            deferred through error feedback (see Dion2). No-op off the
+            row-sharded path.
 
     NorDion2 optimizer applying Dion2 update to NorMuon
     """
@@ -85,9 +90,10 @@ def __init__(
             raise ValueError(f"Invalid learning rate: {lr}")
         if not (0.0 < fraction <= 1.0):
             raise ValueError(f"fraction must be in (0, 1], got {fraction}")
-        if selection_scope not in ("local", "global"):
+        if selection_scope not in ("local", "global", "global_capped"):
             raise ValueError(
-                f"selection_scope must be 'local' or 'global', got {selection_scope!r}"
+                f"selection_scope must be 'local', 'global', or 'global_capped', "
+                f"got {selection_scope!r}"
             )
         if mu < 0.0:
             raise ValueError(f"Invalid momentum factor (mu): {mu}")
@@ -340,6 +346,20 @@ def nordion2_update_megabatch_async(
         k = None
     global_comm_dim_size = global_dim_size
 
+    # "global_capped": norms-only all-gather + capacity priority; see dion2.
+    priority = None
+    if (
+        selection_scope == "global_capped"
+        and comm_dim is not None
+        and process_group is not None
+    ):
+        norms = dion2_pre_accumulate_norms(
+            G=to_local(G), M=to_local(M), select_dim=select_dim
+        )
+        priority = yield from dion2_capped_priority_async(
+            norms, padded_local, k, device_rank, world_size, process_group
+        )
+
     # Update momentum and compute the inputs for orthogonalization
     # Dion2 pre-orthogonalizes differs from NorMuon by applying damping before updating momentum
     U_selected, indices_list = dion2_pre_orthogonalize(
@@ -349,6 +369,7 @@ def nordion2_update_megabatch_async(
         ef_decay=momentum,
         select_dim=select_dim,
         k_override=k,
+        priority=priority,
     )
 
     # Orthogonalize via shared megabatch communication

From 6ce9ac5ca3b22e85f5dbac07dc1c2987cda7ba7a Mon Sep 17 00:00:00 2001
From: alint77 <ali.nt1377@gmail.com>
Date: Thu, 2 Jul 2026 13:01:32 +0200
Subject: [PATCH 2/6] fix(capped): winners-only updates + single-owner
 accumulation (PR#99 review)

The review's asks, plus one deeper fix found while implementing them:

The committed capped scope was trajectory-identical to "local". The global
threshold is a value cut, so a rank's winners are always its top-norm rows;
filling spare slots with next-best rows and giving fillers EF decay made the
selected set exactly the per-rank top-k (verified analytically and by 20k-trial
brute force). Fixed: spare slots now travel as zero rows (zeros pass through
Newton-Schulz as exact zeros, so fillers get no weight update) and fillers skip
error-feedback decay, so only global winners are applied and overflow deferral
is real.

Review items:
- Accumulation contract: pass the threshold (not a priority tensor) into
  dion2_pre_orthogonalize; dion2_pre_accumulate_norms is the single accumulation
  owner and the priority offset (a topk no-op) is gone.
- Document the budget as top-(k*world_size), not top-(ceil(fraction*global));
  they coincide when world_size divides the sharded dim.
- Drop the dead device_rank parameter; raise on negative pad instead of letting
  F.pad silently trim; comment the >=-tie inflation.
- Port tests into tests/test_dion2_selection_scope.py: capped == global when
  winners fit, overflow deferral across steps (the test that distinguishes
  capped from local), uneven shards, fraction=1.0 == global, empty shard,
  bit-identical threshold across ranks, negative-pad raise. 21/21 pass on
  2x/4x H100.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 dion/dion2.py                       | 162 ++++++++++++----------
 dion/nordion2.py                    |  23 ++--
 tests/test_dion2_selection_scope.py | 199 +++++++++++++++++++++++++++-
 3 files changed, 307 insertions(+), 77 deletions(-)

diff --git a/dion/dion2.py b/dion/dion2.py
index 470d5fe..bac9be6 100644
--- a/dion/dion2.py
+++ b/dion/dion2.py
@@ -51,10 +51,11 @@ class Dion2(DistributedOrthoBase):
             sharding-dependent approximation that converges slightly worse; opt
             in when comm-bound at large scale. "global_capped": global selection
             at local comm cost -- all-gather the row norms only (~KBs), compute
-            the global top-k everywhere, and send each rank's globally-selected
-            rows through the fixed ``k``-slot pipes; a rank owning more winners
-            than slots defers the overflow (error feedback re-selects them next
-            step), one owning fewer fills spare slots with its next-best rows.
+            the same global top-``k*world_size`` threshold everywhere, and send
+            each rank's globally-selected rows through the fixed ``k``-slot
+            pipes; a rank owning more winners than slots defers the overflow
+            (error feedback re-selects them on a later step), one owning fewer
+            sends zeros in the spare slots so only winners produce updates.
             No-op off the row-sharded path.
 
     Dion2 optimizer by Ahn et al.: TBD
@@ -248,11 +249,16 @@ def dion2_update_megabatch_async(
     - ``"global_capped"``: global selection at local comm cost. The slice L1
       norms are all-gathered (one fp32 per slice, ~KBs), every rank computes the
       same global top-``k*world_size`` threshold, and each rank sends its
-      globally-selected slices through the same fixed ``k``-slot pipes as
-      "local". Overflow winners (a rank owning more winners than slots) are
-      deferred -- error feedback keeps their momentum accumulating so they win a
-      slot on a later step; spare slots fill with the rank's next-best slices.
-      Near-global selection quality at "local" comm cost.
+      globally-selected slices ("winners") through the same fixed ``k``-slot
+      pipes as "local". Overflow winners (a rank owning more winners than slots)
+      are deferred -- they skip error-feedback decay, so their momentum keeps
+      accumulating and they win a slot on a later step. Spare slots travel as
+      zeros: non-winners are never applied and never decayed, which is what
+      distinguishes this scope from "local" (a per-rank top-``k`` fill would be
+      trajectory-identical to "local", since a rank's winners are always its
+      top-norm slices). Note the selection budget is top-``k*world_size``, not
+      top-``ceil(fraction*global)``: the two coincide when ``world_size``
+      divides the sharded dim and the budget is slightly larger otherwise.
 
     Off the row-sharded path (per-head, single-GPU, batch-sharded) each rank
     already holds whole matrices, so local and global selection coincide and
@@ -369,10 +375,11 @@ def dion2_update_megabatch_async(
         k = None
     global_comm_dim_size = global_dim_size
 
-    # "global_capped": accumulate momentum + all-gather the slice norms, then
-    # rank slices by the capacity priority (global winners first) instead of the
-    # local norms. Same fixed k-slot comm as "local" below.
-    priority = None
+    # "global_capped": accumulate momentum + all-gather the slice norms (one
+    # fp32 per slice, ~KBs), and hand the global top-(k*world_size) threshold to
+    # pre_orthogonalize, which applies the capacity rule (winners-only updates,
+    # zeros in spare slots). Same fixed k-slot comm as "local" below.
+    thresh = None
     if (
         selection_scope == "global_capped"
         and comm_dim is not None
@@ -381,8 +388,8 @@ def dion2_update_megabatch_async(
         norms = dion2_pre_accumulate_norms(
             G=to_local(G), M=to_local(M), select_dim=select_dim
         )
-        priority = yield from dion2_capped_priority_async(
-            norms, padded_local, k, device_rank, world_size, process_group
+        thresh = yield from dion2_capped_threshold_async(
+            norms, padded_local, k, world_size, process_group
         )
 
     # Pre-orthogonalize: momentum update + submatrix selection
@@ -393,7 +400,7 @@ def dion2_update_megabatch_async(
         ef_decay=ef_decay,
         select_dim=select_dim,
         k_override=k,
-        priority=priority,
+        thresh=thresh,
     )
 
     # Orthogonalize via shared megabatch communication
@@ -468,9 +475,11 @@ def dion2_pre_accumulate_norms(
     Phase A of the "global_capped" scope: update momentum with the gradient and
     return the stacked slice L1 norms ``[N, local_size]`` (fp32). The caller
     all-gathers these tiny norms across ranks (an eager collective that cannot
-    live inside this compiled graph) and feeds the resulting selection priority
-    back into ``dion2_pre_orthogonalize`` via ``priority`` -- which then must
-    NOT re-accumulate the gradient.
+    live inside this compiled graph) and feeds the resulting global threshold
+    back into ``dion2_pre_orthogonalize`` via ``thresh``. This function is the
+    single owner of momentum accumulation on the capped path: passing ``thresh``
+    tells ``dion2_pre_orthogonalize`` NOT to re-accumulate the gradient, so the
+    gathered norms always reflect the post-accumulation momentum.
     """
     dtype = M[0].dtype
     G = [g.to(dtype=dtype) for g in G]
@@ -481,38 +490,39 @@ def dion2_pre_accumulate_norms(
     return torch.stack(M, dim=0).norm(p=1, dim=norm_dim).float()
 
 
-def dion2_capped_priority_async(
+def dion2_capped_threshold_async(
     norms: Tensor,  # [N, local_size] fp32 from dion2_pre_accumulate_norms
     padded_local: int,
     k: int,
-    device_rank: int,
     world_size: int,
     process_group: Optional[ProcessGroup],
 ) -> Generator[None, None, Tensor]:
     """
-    "global_capped" selection priority: all-gather the slice L1 norms (one fp32
-    per slice -- ~KBs vs MBs for the rows), compute the exact global
-    top-``k*world_size`` threshold identically on every rank, and return a
-    priority tensor ``[N, local_size]`` such that a plain local
-    ``topk(priority, k)`` implements the capacity rule:
-
-      - slices at or above the global threshold ("winners") outrank ALL
-        non-winners (a per-matrix offset larger than any norm), so a rank sends
-        as many of its winners as fit in its ``k`` slots -- overflow winners are
-        deferred, and since only *selected* slices get error-feedback decay,
-        their momentum keeps accumulating and they win a slot on a later step;
-      - spare slots fill with the rank's next-best non-winners (ordinary norm
-        order), harmless extra orthogonalization exactly like "local" scope.
+    "global_capped" selection threshold: all-gather the slice L1 norms (one fp32
+    per slice -- ~KBs vs MBs for the rows) and return the per-matrix global
+    top-``k*world_size`` threshold ``[N, 1]``, computed identically on every
+    rank. ``dion2_pre_orthogonalize`` compares its local norms against this
+    value to apply the capacity rule (winners-only updates -- see its
+    docstring).
 
     Rank consistency is by construction: every rank computes the threshold from
-    the same gathered tensor. Async generator in the megabatch style: yields at
-    the collective, returns the priority.
+    the same bit-identical gathered tensor, and the threshold is a *value* (the
+    ``k_total``-th largest norm), so it does not depend on ``topk`` tie-breaking
+    order. Async generator in the megabatch style: yields at the collective,
+    returns the threshold.
     """
     N, local_size = norms.shape
     # Pad short/empty shards to the rank-uniform padded_local. L1 norms are
     # nonnegative, so -1 padding can never enter the winner set unless
     # k*world_size exceeds the real row count (tiny-matrix degenerate case,
     # where "everything is a winner" is the correct answer anyway).
+    if local_size > padded_local:
+        # F.pad with a negative amount would silently TRIM norms; fail loudly
+        # instead (mirrors the analogous guard in megabatch_orthogonalize_async).
+        raise ValueError(
+            f"Local norm count {local_size} exceeds padded size {padded_local}. "
+            "This should not happen with FSDP2 contiguous sharding."
+        )
     if local_size < padded_local:
         norms_pad = torch.nn.functional.pad(
             norms, (0, padded_local - local_size), value=-1.0
@@ -535,12 +545,7 @@ def dion2_capped_priority_async(
     vals, _ = torch.topk(
         all_norms.permute(1, 0, 2).reshape(N, -1), k_total, dim=-1
     )
-    thresh = vals[:, -1:]
-    own = norms  # this rank's real slices (unpadded view)
-    winner = own >= thresh
-    # Offset strictly larger than any norm: winners always outrank non-winners.
-    offset = vals[:, :1] + 1.0
-    return own + winner.to(own.dtype) * offset
+    return vals[:, -1:]
 
 
 @_inductor_workaround
@@ -552,7 +557,7 @@ def dion2_pre_orthogonalize(
     ef_decay: Tensor,
     select_dim: int,
     k_override: Optional[int] = None,
-    priority: Optional[Tensor] = None,
+    thresh: Optional[Tensor] = None,
 ) -> Tuple[List[Tensor], List[Tensor]]:
     """
     Update momentum with gradient and compute the input to orthogonalization.
@@ -573,12 +578,21 @@ def dion2_pre_orthogonalize(
     downstream megabatch pads U to ``local_comm_size=k_override`` for a uniform
     alltoall, and indices stay at the real selected count.
 
-    ``priority`` (the "global_capped" scope) replaces the locally-computed L1
-    norms as the top-k ranking key ``[N, local_size]``. Passing it also means
-    the momentum was ALREADY accumulated by ``dion2_pre_accumulate_norms`` (the
-    priority derives from post-accumulation norms), so the ``M += G`` here is
-    skipped. Everything downstream (gather, EF on selected, bf16 convert) is
-    identical.
+    ``thresh`` (the "global_capped" scope) is the per-matrix global
+    top-``k*world_size`` norm threshold ``[N, 1]`` from
+    ``dion2_capped_threshold_async``. Passing it means the momentum was ALREADY
+    accumulated by ``dion2_pre_accumulate_norms`` (the single accumulation
+    owner on the capped path -- the gathered norms must reflect the
+    post-accumulation momentum), so the ``M += G`` here is skipped. The top-k
+    gather itself is unchanged: a rank's winners (slices at/above ``thresh``)
+    are by construction its top-norm slices, so they already lead the plain
+    norm-order selection. What changes is what happens to the non-winner
+    "filler" slices occupying spare slots: they are masked to zero in the
+    communicated ``U_selected`` (zero slices pass through Newton-Schulz as
+    exact zeros, so they produce no weight update) and they skip error-feedback
+    decay, letting their momentum accumulate until they cross the global
+    threshold on a later step. Without this masking, capped selection would be
+    trajectory-identical to "local" scope.
     """
     dtype = M[0].dtype
 
@@ -601,10 +615,10 @@ def dion2_pre_orthogonalize(
         k = max(1, int(math.ceil(fraction * num_select)))
     k_topk = min(k, num_select)
 
-    # Update momentum: M = M + G. Skipped when `priority` is given ("global_
-    # capped"): dion2_pre_accumulate_norms already accumulated so the gathered
-    # norms reflect the post-accumulation momentum.
-    if priority is None:
+    # Update momentum: M = M + G. Skipped when `thresh` is given ("global_
+    # capped"): dion2_pre_accumulate_norms is the single accumulation owner on
+    # that path, so the gathered norms reflect the post-accumulation momentum.
+    if thresh is None:
         G = [g.to(dtype=dtype) for g in G]
         torch._foreach_add_(M, G)
 
@@ -624,13 +638,7 @@ def dion2_pre_orthogonalize(
         return U_selected, indices_list
 
     M_stacked = torch.stack(M, dim=0)
-
-    # Top-k ranking key: the capped-scope priority when given (winners outrank
-    # all non-winners), otherwise the local L1 norms along norm_dim.
-    if priority is not None:
-        slice_norms = priority
-    else:
-        slice_norms = M_stacked.norm(p=1, dim=norm_dim)
+    slice_norms = M_stacked.norm(p=1, dim=norm_dim)
 
     # Batched topk: indices shape (batch_size, k_topk). k_topk <= num_select is
     # guaranteed, so this never raises even on a short remainder shard.
@@ -653,20 +661,40 @@ def dion2_pre_orthogonalize(
         )
         selected_stacked = torch.gather(M_stacked, dim=-1, index=indices_expanded)
 
-    # Apply error feedback decay to selected slices in the original M tensors.
-    # We reuse the already-gathered slices and write them back (scaled) using
-    # scatter_, which places values into positions specified by the index tensor.
+    # "global_capped" capacity rule: only slices at/above the global threshold
+    # ("winners") are applied. Winners lead the norm-order selection above by
+    # construction, so `indices` needs no change; non-winner "fillers" in spare
+    # slots are masked to zero for communication (zero slices pass through
+    # Newton-Schulz as exact zeros -> no weight update) and keep their momentum
+    # un-decayed so it accumulates until they win. `>=` marks hard ties at the
+    # threshold as winners, which can inflate the winner set past k*world_size
+    # -- measure-zero for continuous fp32 norms and bounded by the per-rank k
+    # cap, so benign.
+    if thresh is not None:
+        sel_norms = torch.gather(slice_norms, dim=-1, index=indices)  # [N, k_topk]
+        winner = sel_norms >= thresh
+        winner_exp = winner.unsqueeze(-1) if select_dim == -2 else winner.unsqueeze(-2)
+        ef_factor = torch.where(winner_exp, ef_decay, torch.ones_like(ef_decay))
+        U_stacked = selected_stacked * winner_exp.to(selected_stacked.dtype)
+    else:
+        ef_factor = ef_decay
+        U_stacked = selected_stacked
+
+    # Apply error feedback decay to selected slices in the original M tensors
+    # (capped scope: winners only). We reuse the already-gathered slices and
+    # write them back (scaled) using scatter_, which places values into
+    # positions specified by the index tensor.
     indices_list = list(indices.unbind(dim=0))
-    selected_list = list(selected_stacked.unbind(dim=0))
-    for m, idx, selected in zip(M, indices_list, selected_list):
+    ef_src_list = list((selected_stacked * ef_factor).unbind(dim=0))
+    for m, idx, ef_src in zip(M, indices_list, ef_src_list):
         if select_dim == -2:
             idx_exp = idx.unsqueeze(-1).expand(*idx.shape, m.size(-1))
         else:
             idx_exp = idx.unsqueeze(-2).expand(*idx.shape[:-1], m.size(-2), idx.shape[-1])
-        m.scatter_(dim=select_dim, index=idx_exp, src=selected * ef_decay)
+        m.scatter_(dim=select_dim, index=idx_exp, src=ef_src)
 
     # Convert to bf16 and unstack for communication
-    U_selected = list(selected_stacked.to(dtype=torch.bfloat16).unbind(dim=0))
+    U_selected = list(U_stacked.to(dtype=torch.bfloat16).unbind(dim=0))
 
     return U_selected, indices_list
 
diff --git a/dion/nordion2.py b/dion/nordion2.py
index f83002a..92d503a 100644
--- a/dion/nordion2.py
+++ b/dion/nordion2.py
@@ -19,7 +19,7 @@
     dion2_post_orthogonalize,
     dion2_pre_accumulate,
     dion2_pre_accumulate_norms,
-    dion2_capped_priority_async,
+    dion2_capped_threshold_async,
     _make_select_and_orthogonalize,
 )
 from .normuon import normuon_normalization_stacked, _normuon_normalization_core
@@ -57,9 +57,10 @@ class NorDion2(DistributedOrthoBase):
             assembled whole matrix -- layout-invariant/reproducible and better-
             converging. "local": per-shard top-k (union) -- cheaper comm but a
             sharding-dependent approximation that converges slightly worse; opt
-            in when comm-bound at large scale. "global_capped": global selection
-            at local comm cost via a norms-only all-gather; overflow winners are
-            deferred through error feedback (see Dion2). No-op off the
+            in when comm-bound at large scale. "global_capped": global
+            top-``k*world_size`` selection at local comm cost via a norms-only
+            all-gather; winners-only updates (spare slots carry zeros), overflow
+            winners deferred through error feedback (see Dion2). No-op off the
             row-sharded path.
 
     NorDion2 optimizer applying Dion2 update to NorMuon
@@ -346,8 +347,12 @@ def nordion2_update_megabatch_async(
         k = None
     global_comm_dim_size = global_dim_size
 
-    # "global_capped": norms-only all-gather + capacity priority; see dion2.
-    priority = None
+    # "global_capped": norms-only all-gather + global threshold; winners-only
+    # updates, zeros in spare slots -- see dion2. (Filler slices come back from
+    # Newton-Schulz as exact zeros, so they get no weight update; their variance
+    # buffer rows still decay by muon_beta2 this step, a benign side effect that
+    # recovers once the row wins a slot.)
+    thresh = None
     if (
         selection_scope == "global_capped"
         and comm_dim is not None
@@ -356,8 +361,8 @@ def nordion2_update_megabatch_async(
         norms = dion2_pre_accumulate_norms(
             G=to_local(G), M=to_local(M), select_dim=select_dim
         )
-        priority = yield from dion2_capped_priority_async(
-            norms, padded_local, k, device_rank, world_size, process_group
+        thresh = yield from dion2_capped_threshold_async(
+            norms, padded_local, k, world_size, process_group
         )
 
     # Update momentum and compute the inputs for orthogonalization
@@ -369,7 +374,7 @@ def nordion2_update_megabatch_async(
         ef_decay=momentum,
         select_dim=select_dim,
         k_override=k,
-        priority=priority,
+        thresh=thresh,
     )
 
     # Orthogonalize via shared megabatch communication
diff --git a/tests/test_dion2_selection_scope.py b/tests/test_dion2_selection_scope.py
index ea7e475..178647d 100644
--- a/tests/test_dion2_selection_scope.py
+++ b/tests/test_dion2_selection_scope.py
@@ -32,7 +32,7 @@
 import torch.multiprocessing as mp
 from torch.distributed.tensor import DeviceMesh, distribute_tensor, Shard
 
-from dion.dion2 import Dion2
+from dion.dion2 import Dion2, dion2_capped_threshold_async
 from dion.nordion2 import NorDion2
 from dion.polar_express import polar_express
 
@@ -241,3 +241,200 @@ def test_fraction_one_local_equals_global(OptCls, kw, tmp_path):
     dl = _run_sharded(OptCls, "local", 1.0, kw, 2, 29630, tmp_path)
     dg = _run_sharded(OptCls, "global", 1.0, kw, 2, 29640, tmp_path)
     torch.testing.assert_close(dl["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL)
+
+
+# ---- selection_scope="global_capped" ----
+#
+# The capacity rule: all-gather slice norms only, every rank derives the same
+# global top-(k*world_size) threshold, and only slices at/above it ("winners")
+# are applied -- spare slots travel as zeros and non-winners skip error-feedback
+# decay so their momentum accumulates until they win. NOTE the budget is
+# top-(k*world_size), NOT top-(ceil(fraction*global)); they coincide when
+# world_size divides the sharded dim (test below) and the capped budget is
+# slightly larger otherwise (empty-shard test below exploits this: with
+# k*world >= real rows, everything wins and capped(f) == global(1.0)).
+
+
+def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
+                   fraction, boost, steps):
+    """Multi-step worker with a crafted gradient: small noise everywhere, plus
+    rows in ``boost`` set to a constant so their momentum L1 norm is exactly the
+    boost magnitude. Steps after the first use a zero gradient (isolates the
+    error-feedback dynamics). Saves the full W after every step."""
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    torch.cuda.set_device(rank)
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+    dev = torch.device(f"cuda:{rank}")
+    mesh = DeviceMesh("cuda", list(range(world_size)))
+
+    rows, cols = shape
+    g = torch.Generator(device=dev).manual_seed(1234)
+    W0 = torch.randn(rows, cols, generator=g, device=dev)
+    G = torch.randn(rows, cols, generator=g, device=dev) * 0.01
+    for r, mag in boost.items():
+        G[r] = mag / cols  # exact row L1 norm == mag
+
+    p = torch.nn.Parameter(distribute_tensor(W0, mesh, [Shard(0)]))
+    opt = OptCls(
+        [dict(params=[p])],
+        distributed_mesh=mesh,
+        lr=0.1,
+        fraction=fraction,
+        newton_schulz_func=_ns,
+        selection_scope=scope,
+        **kw,
+    )
+    Ws = []
+    for s in range(steps):
+        grad = G if s == 0 else torch.zeros_like(G)
+        p.grad = distribute_tensor(grad, mesh, [Shard(0)])
+        opt.step()
+        Ws.append(p.detach().full_tensor().cpu())
+    if rank == 0:
+        torch.save({"W0": W0.cpu(), "G": G.cpu(), "Ws": Ws}, out_path)
+    dist.destroy_process_group()
+
+
+def _run_capped(OptCls, scope, kw, port, tmp, *, shape=(16, 32), fraction=0.25,
+                boost=None, steps=1):
+    out = str(tmp / f"out_{port}.pt")
+    mp.spawn(
+        _capped_worker,
+        args=(2, port, OptCls, scope, kw, out, shape, fraction, boost or {}, steps),
+        nprocs=2,
+        join=True,
+    )
+    return torch.load(out)
+
+
+def _changed_rows(W_after, W_before):
+    """Row indices where anything changed at all (bitwise)."""
+    return set((W_after != W_before).any(dim=-1).nonzero().flatten().tolist())
+
+
+@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
+@pytest.mark.parametrize("OptCls,kw", [
+    (Dion2, {}),
+    (NorDion2, {"mu": 0.95, "muon_beta2": 0.95}),
+])
+def test_capped_scope_matches_global_when_winners_fit(OptCls, kw, tmp_path):
+    """When every rank owns at most k winners and world_size divides the rows
+    (so the top-(k*world) budget equals global's top-ceil(fraction*global)),
+    capped must equal global-scope exactly. 16 rows / 2 ranks / fraction 0.25:
+    k=2, budget 4; the 4 boosted rows are split 2 per rank."""
+    boost = {0: 3.0, 1: 2.9, 8: 2.8, 9: 2.7}
+    dc = _run_capped(OptCls, "global_capped", kw, 29680, tmp_path, boost=boost)
+    dg = _run_capped(OptCls, "global", kw, 29681, tmp_path, boost=boost)
+    torch.testing.assert_close(dc["Ws"][0], dg["Ws"][0],
+                               rtol=BF16_RTOL, atol=BF16_ATOL)
+
+
+@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
+@pytest.mark.parametrize("OptCls,kw", [
+    (Dion2, {"ef_decay": 0.5, "weight_decay": 0.0}),
+    (NorDion2, {"mu": 0.5, "muon_beta2": 0.95, "weight_decay": 0.0}),
+])
+def test_capped_scope_defers_overflow_winners(OptCls, kw, tmp_path):
+    """All 4 global winners live on rank 0 but it only has k=2 slots: step 1
+    applies exactly its top-2 {0,1} (rows 2,3 deferred WITHOUT error-feedback
+    decay; rank 1's spare slots carry zeros, so NO row of rank 1 changes --
+    this is what distinguishes capped from local, which would fill and apply
+    rank 1's top-2). Step 2 (zero grad): rows 0,1 decayed to 1.5/1.45 while the
+    deferred 2,3 kept 2.8/2.7, so exactly {2,3} are applied. weight_decay=0
+    makes non-applied rows bitwise unchanged."""
+    boost = {0: 3.0, 1: 2.9, 2: 2.8, 3: 2.7}
+    d = _run_capped(OptCls, "global_capped", kw, 29682, tmp_path,
+                    boost=boost, steps=2)
+    assert _changed_rows(d["Ws"][0], d["W0"]) == {0, 1}
+    assert _changed_rows(d["Ws"][1], d["Ws"][0]) == {2, 3}
+
+
+@pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs")
+@pytest.mark.parametrize("OptCls,kw", [
+    (Dion2, {}),
+    (NorDion2, {"mu": 0.95, "muon_beta2": 0.95}),
+])
+def test_capped_scope_uneven_shards_fraction_one_equals_global(OptCls, kw, tmp_path):
+    """Uneven division (17 rows over 4 ranks -> [5,5,5,2] + one short shard):
+    at fraction=1.0 the threshold is the -1 pad value, everything wins, and
+    capped must equal global (both apply all 17 rows; the megabatch zero-pads
+    the short shard's spare slots)."""
+    dc = _run_sharded(OptCls, "global_capped", 1.0, kw, 4, 29683, tmp_path,
+                      shape=(17, 32))
+    dg = _run_sharded(OptCls, "global", 1.0, kw, 4, 29684, tmp_path,
+                      shape=(17, 32))
+    torch.testing.assert_close(dc["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL)
+
+
+@pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs")
+def test_capped_scope_uneven_shards_runs(tmp_path):
+    """Partial-fraction sanity on uneven shards: finite and actually updates."""
+    d = _run_sharded(Dion2, "global_capped", 0.35, {}, 4, 29685, tmp_path,
+                     shape=(17, 32))
+    assert torch.isfinite(d["W"]).all()
+    assert not torch.allclose(d["W"], d["W0"], atol=1e-6)
+
+
+@pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs")
+def test_capped_scope_empty_shard(tmp_path):
+    """3 rows over 4 ranks chunk to [1,1,1,0]: rank 3's shard is empty. The
+    capped budget k*world = 4 exceeds the 3 real rows, so the threshold is the
+    -1 pad value and every real row wins -- which also makes capped(0.5) equal
+    global(1.0), a live demonstration of the top-(k*world) vs
+    top-(ceil(fraction*global)) budget distinction."""
+    dc = _run_sharded(Dion2, "global_capped", 0.5, {}, 4, 29686, tmp_path,
+                      shape=(3, 32))
+    assert torch.isfinite(dc["W"]).all()
+    dg = _run_sharded(Dion2, "global", 1.0, {}, 4, 29687, tmp_path,
+                      shape=(3, 32))
+    torch.testing.assert_close(dc["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL)
+
+
+def _thresh_consistency_worker(rank, world_size, port):
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    torch.cuda.set_device(rank)
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+    dev = torch.device(f"cuda:{rank}")
+
+    # Distinct norms everywhere (no ties), different on each rank; one short
+    # shard (rank 1 has 3 of padded_local=4 slices) to exercise the -1 pad.
+    N, padded_local, k = 2, 4, 2
+    local = padded_local - rank  # rank0: 4, rank1: 3
+    norms = (torch.arange(N * local, dtype=torch.float32, device=dev)
+             .reshape(N, local) * 1.7 + 3.0 + rank * 0.37)
+
+    gen = dion2_capped_threshold_async(norms, padded_local, k, world_size, None)
+    next(gen)
+    try:
+        next(gen)
+        raise AssertionError("generator should be exhausted after the collective")
+    except StopIteration as e:
+        thresh = e.value
+
+    assert thresh.shape == (N, 1)
+    # (f) rank consistency: every rank must hold the bit-identical threshold.
+    both = [torch.empty_like(thresh) for _ in range(world_size)]
+    dist.all_gather(both, thresh.contiguous())  # thresh is a topk slice view
+    assert torch.equal(both[0], both[1])
+    # With distinct norms and k_total <= real slices, exactly k*world_size
+    # slices per matrix sit at/above the threshold globally.
+    winners = (norms >= thresh).sum()
+    dist.all_reduce(winners)
+    assert winners.item() == N * k * world_size
+    dist.destroy_process_group()
+
+
+@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
+def test_capped_threshold_rank_consistent():
+    mp.spawn(_thresh_consistency_worker, args=(2, 29688), nprocs=2, join=True)
+
+
+def test_capped_threshold_negative_pad_raises():
+    """local_size > padded_local would make F.pad silently TRIM the norms; the
+    guard must raise instead. (Raise happens before the collective, so no
+    process group is needed.) Pure-CPU unit test."""
+    gen = dion2_capped_threshold_async(torch.zeros(1, 5), 3, 1, 2, None)
+    with pytest.raises(ValueError, match="exceeds padded size"):
+        next(gen)

From 987acaa0ea0919191685a050992acaac40966c44 Mon Sep 17 00:00:00 2001
From: alint77 <ali.nt1377@gmail.com>
Date: Fri, 3 Jul 2026 09:28:17 +0200
Subject: [PATCH 3/6] fix(capped): cast EF write-back to momentum dtype (bf16
 crash)

The winners-only refactor changed the EF factor from the 0-dim ef_decay
scalar to a dimensioned [N, k, 1] fp32 tensor. PyTorch promotion treats
these differently: bf16 * 0-dim fp32 stays bf16, but bf16 * dimensioned
fp32 promotes to fp32, and the in-place scatter_ into bf16 momentum then
raises "Expected self.dtype to be equal to src.dtype". Only fires with
bf16 optimizer state (real mixed-precision training); every existing test
used fp32 params, so the suite stayed green.

Fix per review: multiply in fp32, round once into M's dtype (a no-op on
all fp32 paths). Add a bf16-param smoke test over local/global_capped x
Dion2/NorDion2 so the promotion path can't regress. 25/25 on 2x/4x H100.

Reported-by: JohnLangford (1b NorDion2 bf16 run)
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 dion/dion2.py                       |  8 +++++--
 tests/test_dion2_selection_scope.py | 36 ++++++++++++++++++++++++-----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/dion/dion2.py b/dion/dion2.py
index bac9be6..6d380eb 100644
--- a/dion/dion2.py
+++ b/dion/dion2.py
@@ -683,9 +683,13 @@ def dion2_pre_orthogonalize(
     # Apply error feedback decay to selected slices in the original M tensors
     # (capped scope: winners only). We reuse the already-gathered slices and
     # write them back (scaled) using scatter_, which places values into
-    # positions specified by the index tensor.
+    # positions specified by the index tensor. The .to(dtype) matters for bf16
+    # momentum: the capped ef_factor is a *dimensioned* fp32 tensor, so
+    # bf16 * fp32 promotes to fp32 (unlike the 0-dim ef_decay scalar, which
+    # loses promotion to a dimensioned bf16), and scatter_ requires src dtype
+    # to match M exactly. Multiply in fp32, round once into M's dtype.
     indices_list = list(indices.unbind(dim=0))
-    ef_src_list = list((selected_stacked * ef_factor).unbind(dim=0))
+    ef_src_list = list((selected_stacked * ef_factor).to(dtype).unbind(dim=0))
     for m, idx, ef_src in zip(M, indices_list, ef_src_list):
         if select_dim == -2:
             idx_exp = idx.unsqueeze(-1).expand(*idx.shape, m.size(-1))
diff --git a/tests/test_dion2_selection_scope.py b/tests/test_dion2_selection_scope.py
index 178647d..3ce277d 100644
--- a/tests/test_dion2_selection_scope.py
+++ b/tests/test_dion2_selection_scope.py
@@ -256,11 +256,13 @@ def test_fraction_one_local_equals_global(OptCls, kw, tmp_path):
 
 
 def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
-                   fraction, boost, steps):
+                   fraction, boost, steps, dtype):
     """Multi-step worker with a crafted gradient: small noise everywhere, plus
     rows in ``boost`` set to a constant so their momentum L1 norm is exactly the
     boost magnitude. Steps after the first use a zero gradient (isolates the
-    error-feedback dynamics). Saves the full W after every step."""
+    error-feedback dynamics). Saves the full W after every step. ``dtype`` sets
+    the param (and therefore momentum) dtype -- bf16 exercises the
+    mixed-precision EF write-back path that fp32 params can never reach."""
     os.environ["MASTER_ADDR"] = "127.0.0.1"
     os.environ["MASTER_PORT"] = str(port)
     torch.cuda.set_device(rank)
@@ -270,8 +272,8 @@ def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
 
     rows, cols = shape
     g = torch.Generator(device=dev).manual_seed(1234)
-    W0 = torch.randn(rows, cols, generator=g, device=dev)
-    G = torch.randn(rows, cols, generator=g, device=dev) * 0.01
+    W0 = torch.randn(rows, cols, generator=g, device=dev).to(dtype)
+    G = (torch.randn(rows, cols, generator=g, device=dev) * 0.01).to(dtype)
     for r, mag in boost.items():
         G[r] = mag / cols  # exact row L1 norm == mag
 
@@ -297,11 +299,12 @@ def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
 
 
 def _run_capped(OptCls, scope, kw, port, tmp, *, shape=(16, 32), fraction=0.25,
-                boost=None, steps=1):
+                boost=None, steps=1, dtype=torch.float32):
     out = str(tmp / f"out_{port}.pt")
     mp.spawn(
         _capped_worker,
-        args=(2, port, OptCls, scope, kw, out, shape, fraction, boost or {}, steps),
+        args=(2, port, OptCls, scope, kw, out, shape, fraction, boost or {},
+              steps, dtype),
         nprocs=2,
         join=True,
     )
@@ -391,6 +394,27 @@ def test_capped_scope_empty_shard(tmp_path):
     torch.testing.assert_close(dc["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL)
 
 
+@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
+@pytest.mark.parametrize("scope", ["local", "global_capped"])
+@pytest.mark.parametrize("OptCls,kw", [
+    (Dion2, {}),
+    (NorDion2, {"mu": 0.95, "muon_beta2": 0.95}),
+])
+def test_scope_bf16_momentum(OptCls, kw, scope, tmp_path):
+    """bf16 params (=> bf16 momentum) must run through the selection scopes.
+    Guards the EF write-back dtype: the capped-scope ef_factor is a
+    *dimensioned* fp32 tensor, and ``bf16 * fp32[N,k,1]`` type-promotes to
+    fp32, which the in-place ``scatter_`` into bf16 momentum rejects with
+    "Expected self.dtype to be equal to src.dtype". A 0-dim fp32 scalar (the
+    pre-capped code, and the local branch) does NOT promote, so fp32-param
+    tests can never catch this -- it only fires in real bf16 training."""
+    port = 29690 + (scope == "global_capped") * 2 + (OptCls is NorDion2)
+    d = _run_capped(OptCls, scope, kw, port, tmp_path, dtype=torch.bfloat16)
+    W = d["Ws"][0]
+    assert torch.isfinite(W.float()).all()
+    assert not torch.equal(W, d["W0"])  # the step actually applied an update
+
+
 def _thresh_consistency_worker(rank, world_size, port):
     os.environ["MASTER_ADDR"] = "127.0.0.1"
     os.environ["MASTER_PORT"] = str(port)

From aaa2cf6408f4f6fd1e704fb013b8085ce12abdf4 Mon Sep 17 00:00:00 2001
From: alint77 <ali.nt1377@gmail.com>
Date: Fri, 3 Jul 2026 10:51:18 +0200
Subject: [PATCH 4/6] feat(capped): pooled packed transport -- near-dropless
 global selection, zero syncs

Replaces the per-matrix fixed-k capped transport with pooled packing ("T2",
see dion2_capped_pooled_plan): each (rank -> dest) a2a chunk has a fixed
budget of rows SHARED across the destination's per_rank matrices, packed
back-to-back with an int32 count header bit-cast into the chunk's first row.
One matrix's winner overflow uses another's spare room, so deferral drops
from per-matrix fluctuation to pooled fluctuation. The receiver parses the
sender's header instead of recomputing counts from the gathered norms --
no bit-identical-plan invariant across ranks, and a receiver-side clip
guards divergence. All collective sizes are construction-time constants;
data-dependence lives in GPU gathers/scatters (static shapes, no host syncs,
compile-safe).

Also in this change:
- Exact-count winner set (T0): stable-argsort tie-break, budget is exactly
  top-ceil(fraction*global) -- fixes the top-(k*world) overshoot and makes
  sum(winners) <= NS buffer a hard invariant the transport relies on.
- capacity_factor knob: None (default) = auto per group,
  1 + 2*sqrt((1-1/world)/(per_rank*k)); float >= 1.0 pins it. Groups whose
  budget would reach the full shard (or whose header cannot fit the first
  row) route to exact "global" scope.
- Cross-matrix deferral priority is scale-normalized (norm/thresh): a
  large-norm matrix cannot evict a small-norm matrix's winners from the
  shared chunk.
- The capped path now mirrors the global path's structure: returned shards
  are zero except at applied winner rows and feed the existing masked post
  (EF decay on applied rows only; deferred winners keep accumulating).
  dion2_pre_orthogonalize no longer has a capped mode.
- CAPPED_STATS module dict exposes deferred/winner row counts per step
  (GPU tensors, no sync) for deferral-rate logging.

Tests: pack/unpack identity roundtrip (pooled deferral, blackout matrix,
uneven counts); pooled cross-matrix skew with scale-normalized eviction
order; exact-count rank consistency; header bit-cast roundtrip; degeneracy
routing (empty shard, fraction=1.0); all prior scope + bf16 tests updated
for the multi-param harness. 28/28 on 2x/4x H100.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 dion/dion2.py                       | 559 +++++++++++++++++++++-------
 dion/nordion2.py                    | 114 ++++--
 tests/test_dion2_selection_scope.py | 238 ++++++++----
 3 files changed, 696 insertions(+), 215 deletions(-)

diff --git a/dion/dion2.py b/dion/dion2.py
index 6d380eb..b0573d0 100644
--- a/dion/dion2.py
+++ b/dion/dion2.py
@@ -11,6 +11,7 @@
 from .megabatch_base import (
     DistributedOrthoBase,
     megabatch_orthogonalize_async,
+    muon_update_newton_schulz,
     adjust_lr_spectral_norm,
     adjust_lr_rms_norm,
 )
@@ -49,14 +50,17 @@ class Dion2(DistributedOrthoBase):
             assembled whole matrix -- layout-invariant/reproducible and better-
             converging. "local": per-shard top-k (union) -- cheaper comm but a
             sharding-dependent approximation that converges slightly worse; opt
-            in when comm-bound at large scale. "global_capped": global selection
-            at local comm cost -- all-gather the row norms only (~KBs), compute
-            the same global top-``k*world_size`` threshold everywhere, and send
-            each rank's globally-selected rows through the fixed ``k``-slot
-            pipes; a rank owning more winners than slots defers the overflow
-            (error feedback re-selects them on a later step), one owning fewer
-            sends zeros in the spare slots so only winners produce updates.
-            No-op off the row-sharded path.
+            in when comm-bound at large scale. "global_capped": exact global
+            top-``ceil(fraction*global)`` selection at ~local comm cost --
+            all-gather the row norms only (~KBs), then pack each rank's winner
+            rows into fixed-size chunks POOLED across the megabatch's matrices
+            (an int32 count header rides in each chunk). A rank whose winners
+            overflow its pooled chunk defers the overflow via error feedback;
+            see ``capacity_factor``. Groups where the packing cannot save comm
+            route to "global". No-op off the row-sharded path.
+        capacity_factor: "global_capped" pooled-chunk slack. None (default) =
+            auto per group, ``1 + 2*sqrt((1-1/world)/(per_rank*k))`` (covers
+            ~2 sigma of winner-count fluctuation); a float >= 1.0 pins it.
 
     Dion2 optimizer by Ahn et al.: TBD
     """
@@ -80,6 +84,7 @@ def __init__(
         verbose: bool = False,
         triton_post_ortho: bool = False,
         selection_scope: str = "global",
+        capacity_factor: Optional[float] = None,
     ):
         # Validate hyperparameters
         if lr < 0.0:
@@ -91,6 +96,10 @@ def __init__(
                 f"selection_scope must be 'local', 'global', or 'global_capped', "
                 f"got {selection_scope!r}"
             )
+        if capacity_factor is not None and capacity_factor < 1.0:
+            raise ValueError(
+                f"capacity_factor must be None (auto) or >= 1.0, got {capacity_factor}"
+            )
         if ef_decay < 0.0:
             raise ValueError(f"Invalid ef_decay: {ef_decay}")
         if len(betas) != 2 or betas[0] < 0.0 or betas[1] < 0.0:
@@ -113,6 +122,7 @@ def __init__(
             algorithm="dion2",
             step=0,
             selection_scope=selection_scope,
+            capacity_factor=capacity_factor,
         )
         super().__init__(
             params, distributed_mesh, "dion2", defaults,
@@ -163,6 +173,7 @@ def _create_ortho_tasks(
                 verbose=self.verbose,
                 triton_post_ortho=self._triton_post_ortho,
                 selection_scope=group["selection_scope"],
+                capacity_factor=group["capacity_factor"],
             )
 
             shape_groups: dict[tuple, list] = defaultdict(list)
@@ -226,6 +237,7 @@ def dion2_update_megabatch_async(
     verbose: bool = False,
     triton_post_ortho: bool = False,
     selection_scope: str = "global",  # "global" (exact whole-matrix top-k, default) or "local" (per-shard top-k; cheaper comm, sharding-variant)
+    capacity_factor: Optional[float] = None,  # "global_capped" pooled-chunk slack; None = auto (z=2 model)
 ) -> Generator[None, None, None]:
     """
     Mega-batched Dion2 update: processes ALL same-shape parameters in one
@@ -246,19 +258,19 @@ def dion2_update_megabatch_async(
       sharding-dependent approximation of the true top-k (world-size variant).
       Cheaper comm (the win grows with model size), but converges slightly
       worse; opt in when comm-bound at large scale.
-    - ``"global_capped"``: global selection at local comm cost. The slice L1
-      norms are all-gathered (one fp32 per slice, ~KBs), every rank computes the
-      same global top-``k*world_size`` threshold, and each rank sends its
-      globally-selected slices ("winners") through the same fixed ``k``-slot
-      pipes as "local". Overflow winners (a rank owning more winners than slots)
-      are deferred -- they skip error-feedback decay, so their momentum keeps
-      accumulating and they win a slot on a later step. Spare slots travel as
-      zeros: non-winners are never applied and never decayed, which is what
-      distinguishes this scope from "local" (a per-rank top-``k`` fill would be
-      trajectory-identical to "local", since a rank's winners are always its
-      top-norm slices). Note the selection budget is top-``k*world_size``, not
-      top-``ceil(fraction*global)``: the two coincide when ``world_size``
-      divides the sharded dim and the budget is slightly larger otherwise.
+    - ``"global_capped"``: exact global top-``ceil(fraction*global)`` selection
+      at ~"local" comm cost, via packed transport. The slice L1 norms are
+      all-gathered (one fp32 per slice, ~KBs) and every rank derives the same
+      exact-count winner set (stable-sort tie-break). Each rank packs its
+      winner rows into a fixed ``budget``-row chunk per destination, POOLED
+      across the destination's ``per_rank`` matrices, with an int32 count
+      header bit-cast into the chunk's first row -- the receiver parses the
+      sender's plan instead of recomputing it. Winners that overflow the
+      pooled budget are deferred: they skip error-feedback decay (only applied
+      rows are decayed, by the masked post), so their momentum accumulates and
+      they win later. Non-winners are never applied and never decayed -- a
+      per-rank top-k fill would be trajectory-identical to "local". Groups
+      whose budget would reach the full shard route to "global" instead.
 
     Off the row-sharded path (per-head, single-GPU, batch-sharded) each rank
     already holds whole matrices, so local and global selection coincide and
@@ -291,12 +303,6 @@ def dion2_update_megabatch_async(
     # comm_dim for sharded communication: use select_dim (which equals normalized shard_dim)
     comm_dim = select_dim if is_sharded else None
 
-    # Decide whether selection happens before communication ("local", and only
-    # meaningful on the row-sharded path) or after the whole matrix is assembled
-    # ("global"). Off the sharded path each rank holds whole matrices, so the two
-    # are identical and we keep the cheaper pre-comm selection.
-    global_scope = selection_scope == "global" and comm_dim is not None
-
     # On the sharded path X[0] must still be a DTensor, so .shape[comm_dim]
     # is the unsharded global size. The megabatch fn uses this to compute
     # the rank-consistent pad size for its alltoall. Catch the case where a
@@ -312,6 +318,89 @@ def dion2_update_megabatch_async(
     else:
         global_dim_size = None
 
+    # --- "global_capped" packed-path decision (row-sharded 2D only) ---
+    # Pooled chunk budget: per (rank -> dest) pair, B = ceil(c * per_rank * k)
+    # rows shared across the dest's per_rank matrices. c defaults to the
+    # z=2 fluctuation model (see dion2_capped_pack); if the budget reaches the
+    # full shard, packing saves nothing and the group routes to exact "global".
+    capped_packed = (
+        selection_scope == "global_capped"
+        and comm_dim is not None
+        and process_group is not None
+        and select_dim == -2
+        and X[0].ndim == 2
+    )
+    if capped_packed:
+        padded_local = (global_dim_size + world_size - 1) // world_size
+        k = max(1, int(math.ceil(fraction * padded_local)))
+        per_rank = (N + world_size - 1) // world_size
+        if capacity_factor is None:
+            c = 1.0 + 2.0 * math.sqrt((1.0 - 1.0 / world_size) / (per_rank * k))
+        else:
+            c = float(capacity_factor)
+        budget = int(math.ceil(c * per_rank * k))
+        # Route to exact "global" when packing cannot pay for itself (slack
+        # reaches the full shard) or the int32 count header cannot fit in the
+        # chunk's first row (2 bf16 slots per matrix).
+        if budget >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]:
+            capped_packed = False
+
+    # Decide whether selection happens before communication ("local", and only
+    # meaningful on the row-sharded path) or after the whole matrix is assembled
+    # ("global"). Off the sharded path each rank holds whole matrices, so the two
+    # are identical and we keep the cheaper pre-comm selection. "global_capped"
+    # groups that fell out of the packed path (degenerate budget, col-sharded,
+    # non-2D) take the exact global path.
+    global_scope = comm_dim is not None and (
+        selection_scope == "global"
+        or (selection_scope == "global_capped" and not capped_packed)
+    )
+
+    if capped_packed:
+        # --- "global_capped": winner-only packed transport ---
+        # Accumulate momentum (single owner on this path), all-gather the slice
+        # norms, pack the global winners into pooled fixed-size chunks, and
+        # come back with full-size shards that are zero except at the applied
+        # winner rows -- exactly the global path's masked-post format.
+        norms = dion2_pre_accumulate_norms(
+            G=to_local(G), M=to_local(M), select_dim=select_dim
+        )
+        U_ortho = yield from dion2_capped_packed_async(
+            M_local=to_local(M),
+            norms=norms,
+            padded_local=padded_local,
+            fraction=fraction,
+            global_size=global_dim_size,
+            device_rank=device_rank,
+            world_size=world_size,
+            process_group=process_group,
+            per_rank=per_rank,
+            budget=budget,
+            kw=k * world_size,
+            newton_schulz_func=newton_schulz_func,
+            flatten=flatten,
+            epsilon=epsilon,
+        )
+        if adjust_lr is None:
+            adjusted_lr = lr
+        elif adjust_lr == "spectral_norm":
+            adjusted_lr = adjust_lr_spectral_norm(lr, X[0].shape, flatten=flatten)
+        elif adjust_lr == "rms_norm":
+            adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten)
+        else:
+            raise ValueError(f"Unknown adjust_lr: {adjust_lr}")
+        dion2_post_orthogonalize_masked(
+            X=to_local(X),
+            M=to_local(M),
+            U=U_ortho,
+            base_lr=lr,
+            adjusted_lr=adjusted_lr,
+            weight_decay=weight_decay,
+            ef_decay=ef_decay,
+            select_dim=select_dim,
+        )
+        return
+
     if global_scope:
         # --- Global selection: send the full shard, select after assembly ---
         # No pre-comm selection; momentum gets the gradient and the whole shard
@@ -375,23 +464,6 @@ def dion2_update_megabatch_async(
         k = None
     global_comm_dim_size = global_dim_size
 
-    # "global_capped": accumulate momentum + all-gather the slice norms (one
-    # fp32 per slice, ~KBs), and hand the global top-(k*world_size) threshold to
-    # pre_orthogonalize, which applies the capacity rule (winners-only updates,
-    # zeros in spare slots). Same fixed k-slot comm as "local" below.
-    thresh = None
-    if (
-        selection_scope == "global_capped"
-        and comm_dim is not None
-        and process_group is not None
-    ):
-        norms = dion2_pre_accumulate_norms(
-            G=to_local(G), M=to_local(M), select_dim=select_dim
-        )
-        thresh = yield from dion2_capped_threshold_async(
-            norms, padded_local, k, world_size, process_group
-        )
-
     # Pre-orthogonalize: momentum update + submatrix selection
     U_selected, indices_list = dion2_pre_orthogonalize(
         G=to_local(G),
@@ -400,7 +472,6 @@ def dion2_update_megabatch_async(
         ef_decay=ef_decay,
         select_dim=select_dim,
         k_override=k,
-        thresh=thresh,
     )
 
     # Orthogonalize via shared megabatch communication
@@ -475,11 +546,10 @@ def dion2_pre_accumulate_norms(
     Phase A of the "global_capped" scope: update momentum with the gradient and
     return the stacked slice L1 norms ``[N, local_size]`` (fp32). The caller
     all-gathers these tiny norms across ranks (an eager collective that cannot
-    live inside this compiled graph) and feeds the resulting global threshold
-    back into ``dion2_pre_orthogonalize`` via ``thresh``. This function is the
-    single owner of momentum accumulation on the capped path: passing ``thresh``
-    tells ``dion2_pre_orthogonalize`` NOT to re-accumulate the gradient, so the
-    gathered norms always reflect the post-accumulation momentum.
+    live inside this compiled graph; see ``dion2_capped_select_async``). This
+    is the ONLY momentum accumulation on the capped path -- the packed
+    transport never calls ``dion2_pre_orthogonalize``, so the gathered norms
+    always reflect the post-accumulation momentum.
     """
     dtype = M[0].dtype
     G = [g.to(dtype=dtype) for g in G]
@@ -490,32 +560,41 @@ def dion2_pre_accumulate_norms(
     return torch.stack(M, dim=0).norm(p=1, dim=norm_dim).float()
 
 
-def dion2_capped_threshold_async(
+# Deferral instrumentation for the "global_capped" packed path. Updated per
+# megabatch step with GPU tensors (no host sync); external code may read and
+# log e.g. CAPPED_STATS["deferred_rows"] / CAPPED_STATS["winner_rows"].
+CAPPED_STATS: dict = {}
+
+
+def dion2_capped_select_async(
     norms: Tensor,  # [N, local_size] fp32 from dion2_pre_accumulate_norms
     padded_local: int,
-    k: int,
+    fraction: float,
+    global_size: int,
+    device_rank: int,
     world_size: int,
     process_group: Optional[ProcessGroup],
-) -> Generator[None, None, Tensor]:
+) -> Generator[None, None, Tuple[Tensor, Tensor]]:
     """
-    "global_capped" selection threshold: all-gather the slice L1 norms (one fp32
-    per slice -- ~KBs vs MBs for the rows) and return the per-matrix global
-    top-``k*world_size`` threshold ``[N, 1]``, computed identically on every
-    rank. ``dion2_pre_orthogonalize`` compares its local norms against this
-    value to apply the capacity rule (winners-only updates -- see its
-    docstring).
-
-    Rank consistency is by construction: every rank computes the threshold from
-    the same bit-identical gathered tensor, and the threshold is a *value* (the
-    ``k_total``-th largest norm), so it does not depend on ``topk`` tie-breaking
-    order. Async generator in the megabatch style: yields at the collective,
-    returns the threshold.
+    "global_capped" winner selection: all-gather the slice L1 norms (one fp32
+    per slice -- ~KBs vs MBs for the rows) and return, for this rank,
+
+      - ``winner`` bool ``[N, local_size]``: membership in the exact global
+        top-``k_total`` set, ``k_total = ceil(fraction * global_size)``. The
+        set is EXACT-COUNT: a stable descending argsort over the gathered
+        norms breaks ties deterministically by flat (rank, row) position, so
+        ties cannot inflate the winner set past the budget (the receiver-side
+        buffer bound in the packed transport relies on this).
+      - ``thresh`` ``[N, 1]``: the k_total-th largest norm per matrix, used to
+        scale-normalize cross-matrix deferral priorities in the pack step.
+
+    ``k_total <= global_size`` always, so the -1 shard padding can never win.
+    Rank consistency: every rank sorts the same bit-identical gathered tensor;
+    the receiver clip in ``dion2_capped_assemble`` guards the (pathological)
+    divergent case. Async generator in the megabatch style: yields at the
+    collective, returns the pair.
     """
     N, local_size = norms.shape
-    # Pad short/empty shards to the rank-uniform padded_local. L1 norms are
-    # nonnegative, so -1 padding can never enter the winner set unless
-    # k*world_size exceeds the real row count (tiny-matrix degenerate case,
-    # where "everything is a winner" is the correct answer anyway).
     if local_size > padded_local:
         # F.pad with a negative amount would silently TRIM norms; fail loudly
         # instead (mirrors the analogous guard in megabatch_orthogonalize_async).
@@ -539,13 +618,278 @@ def dion2_capped_threshold_async(
     yield
     work.wait()
 
-    # [world, N, padded_local] -> per-matrix global threshold = k_total-th value
-    all_norms = gathered.view(world_size, N, padded_local)
-    k_total = min(k * world_size, world_size * padded_local)
-    vals, _ = torch.topk(
-        all_norms.permute(1, 0, 2).reshape(N, -1), k_total, dim=-1
+    # [world, N, padded_local] -> flat per-matrix [N, world*padded_local] with
+    # rank-major flat position (r * padded_local + row).
+    flat = gathered.view(world_size, N, padded_local).permute(1, 0, 2).reshape(N, -1)
+    k_total = min(max(1, int(math.ceil(fraction * global_size))), global_size)
+    order = torch.argsort(flat, dim=-1, descending=True, stable=True)
+    winner_flat = torch.zeros_like(flat, dtype=torch.bool)
+    winner_flat.scatter_(1, order[:, :k_total], True)
+    thresh = flat.gather(1, order[:, k_total - 1 : k_total])
+    own = winner_flat.view(N, world_size, padded_local)[:, device_rank, :local_size]
+    return own.contiguous(), thresh
+
+
+@_inductor_workaround
+@torch.compile(fullgraph=True)
+def dion2_capped_pack(
+    M: List[Tensor],  # N tensors [local, cols] (momentum, already accumulated)
+    norms: Tensor,  # [N, local] fp32
+    winner: Tensor,  # [N, local] bool, from dion2_capped_select_async
+    thresh: Tensor,  # [N, 1] fp32
+    world_size: int,
+    per_rank: int,
+    budget: int,  # pooled row budget B per (rank -> dest) chunk
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """
+    Sender side of the packed "global_capped" transport. Matrices are assigned
+    to destination ranks in blocks of ``per_rank`` (megabatch order); for each
+    destination this rank packs its winner rows of that destination's matrices
+    back-to-back into a fixed ``budget``-row chunk -- slots are POOLED across
+    the per_rank matrices, so one matrix's overflow uses another's spare room.
+
+    If a destination's total winners exceed ``budget``, the lowest-priority
+    rows are deferred (error feedback: they keep full momentum and win later).
+    Priority is the SCALE-NORMALIZED ``norm / thresh`` -- raw norms are not
+    comparable across matrices, and a large-scale matrix must not evict a
+    small-scale matrix's winners from the shared chunk.
+
+    Returns ``(payload [world, budget, cols] bf16, counts [world, per_rank]
+    int32, src_index [world, budget] long, deferred_rows scalar)``. Empty
+    slots gather a dedicated zero row (never read by the receiver, which
+    parses only ``counts`` rows per segment). ``src_index`` maps each slot to
+    its flat local row (sentinel = the zero row) and is kept by the sender to
+    scatter the orthogonalized rows home; the packing is stable and
+    position-preserving -- the reverse path relies on it.
+    """
+    N = len(M)
+    local = M[0].size(-2)
+    cols = M[0].size(-1)
+    n_pad = world_size * per_rank
+    device = M[0].device
+
+    M_stacked = torch.stack(M, dim=0)  # [N, local, cols]
+    if n_pad > N:
+        # Virtual matrices padding the group to world_size * per_rank: no
+        # winners, zero rows.
+        M_stacked = torch.cat(
+            [M_stacked, torch.zeros(n_pad - N, local, cols, dtype=M_stacked.dtype, device=device)]
+        )
+        winner = torch.cat(
+            [winner, torch.zeros(n_pad - N, local, dtype=torch.bool, device=device)]
+        )
+        norms = torch.cat(
+            [norms, torch.zeros(n_pad - N, local, dtype=norms.dtype, device=device)]
+        )
+        thresh = torch.cat(
+            [thresh, torch.ones(n_pad - N, 1, dtype=thresh.dtype, device=device)]
+        )
+
+    prio = norms / thresh.clamp_min(1e-12)
+    score = torch.where(winner, prio, torch.full_like(prio, float("-inf")))
+    score_d = score.view(world_size, per_rank * local)
+    b_eff = min(budget, per_rank * local)
+    vals, idxs = torch.topk(score_d, b_eff, dim=-1)
+    kept_flat = torch.zeros_like(score_d, dtype=torch.bool)
+    kept_flat.scatter_(1, idxs, vals.isfinite())  # winners only; spare slots stay empty
+    kept = kept_flat.view(world_size, per_rank, local)
+
+    counts = kept.sum(dim=-1)  # [world, per_rank]
+    seg_border = counts.cumsum(dim=1) - counts  # exclusive, within chunk
+    # Slot of each kept row: matrix segment border + within-matrix ordinal
+    # (row order preserved -- stable packing).
+    ordinal = kept.cumsum(dim=-1) - 1
+    slot = seg_border.unsqueeze(-1) + ordinal
+    slot_safe = torch.where(kept, slot, torch.full_like(slot, budget))  # trash col
+
+    fid = torch.arange(n_pad * local, device=device).view(world_size, per_rank, local)
+    sentinel = n_pad * local  # index of the appended zero row
+    src_index = torch.full(
+        (world_size, budget + 1), sentinel, dtype=torch.long, device=device
+    )
+    src_index.scatter_(
+        1, slot_safe.reshape(world_size, -1), fid.reshape(world_size, -1)
+    )
+    src_index = src_index[:, :budget]
+
+    M_flat = torch.cat(
+        [M_stacked.reshape(n_pad * local, cols),
+         torch.zeros(1, cols, dtype=M_stacked.dtype, device=device)]
+    )
+    payload = M_flat[src_index].to(torch.bfloat16)  # [world, budget, cols]
+    deferred = winner.sum() - kept.sum()
+    return payload, counts.to(torch.int32), src_index, deferred
+
+
+@torch.compile(fullgraph=True)
+def dion2_capped_assemble(
+    payload: Tensor,  # [world, budget, cols] bf16 (received chunks, no header)
+    counts: Tensor,  # [world, per_rank] int32 (parsed headers)
+    kw: int,  # static NS row count per matrix (k * world_size)
+) -> Tuple[Tensor, Tensor]:
+    """
+    Receiver side: assemble each of this rank's ``per_rank`` matrices from the
+    variable-length segments of all senders' chunks into a static, zero-padded
+    ``[per_rank, kw, cols]`` Newton-Schulz input. Rows land in (sender rank,
+    within-segment) order; NS is permutation-equivariant and the reverse path
+    inverts the same mapping, so the order is only required to be stable.
+
+    Arrivals beyond ``kw`` for a matrix (impossible with exact-count winners;
+    possible only under a divergent sender) are clipped to a trash slot rather
+    than corrupting the buffer. Returns ``(ns_input, ns_index)`` where
+    ``ns_index [per_rank, kw]`` maps NS rows back to flat payload positions
+    (sentinel = zero row) for ``dion2_capped_repack``.
+    """
+    world_size, budget, cols = payload.shape
+    per_rank = counts.size(1)
+    device = payload.device
+    cnt = counts.to(torch.long)
+
+    seg_border = cnt.cumsum(dim=1) - cnt  # [world, per_rank] within sender chunk
+    dst_border = cnt.cumsum(dim=0) - cnt  # [world, per_rank] within NS rows
+    max_seg = min(budget, kw)
+    j = torch.arange(max_seg, device=device)
+    valid = j < cnt.unsqueeze(-1)  # [world, per_rank, max_seg]
+    src_pos = (
+        torch.arange(world_size, device=device).view(-1, 1, 1) * budget
+        + seg_border.unsqueeze(-1)
+        + j
+    )
+    src_pos = torch.where(valid, src_pos, torch.full_like(src_pos, world_size * budget))
+    dst_pos = dst_border.unsqueeze(-1) + j
+    dst_pos = torch.where(
+        valid & (dst_pos < kw), dst_pos, torch.full_like(dst_pos, kw)  # receiver clip
+    )
+
+    ns_index = torch.full(
+        (per_rank, kw + 1), world_size * budget, dtype=torch.long, device=device
+    )
+    ns_index.scatter_(
+        1,
+        dst_pos.permute(1, 0, 2).reshape(per_rank, -1),
+        src_pos.permute(1, 0, 2).reshape(per_rank, -1),
+    )
+    ns_index = ns_index[:, :kw]
+
+    pay_flat = torch.cat(
+        [payload.reshape(world_size * budget, cols),
+         torch.zeros(1, cols, dtype=payload.dtype, device=device)]
+    )
+    ns_input = pay_flat[ns_index]  # [per_rank, kw, cols], zero-padded tail
+    return ns_input, ns_index
+
+
+@torch.compile(fullgraph=True)
+def dion2_capped_repack(
+    ns_out: Tensor,  # [per_rank, kw, cols] orthogonalized
+    ns_index: Tensor,  # [per_rank, kw] from dion2_capped_assemble
+    world_size: int,
+    budget: int,
+) -> Tensor:
+    """Inverse of assemble: scatter NS rows back to their payload positions.
+    Padding/sentinel rows collapse onto a trash row; payload slots that carried
+    no real row stay exactly zero (so the sender scatters zeros = no update)."""
+    per_rank, kw, cols = ns_out.shape
+    out = torch.zeros(
+        world_size * budget + 1, cols, dtype=ns_out.dtype, device=ns_out.device
+    )
+    out.scatter_(0, ns_index.reshape(-1, 1).expand(-1, cols), ns_out.reshape(-1, cols))
+    return out[: world_size * budget].view(world_size, budget, cols)
+
+
+@torch.compile(fullgraph=True)
+def dion2_capped_unpack(
+    recv: Tensor,  # [world, budget, cols] orthogonalized rows, sender's layout
+    src_index: Tensor,  # [world, budget] from dion2_capped_pack
+    n_params: int,
+    local: int,
+    total_rows: int,  # n_pad * local (sentinel index)
+) -> List[Tensor]:
+    """Sender-side finish: scatter the returned rows to their home (matrix,
+    row) positions. Output is a list of full-size ``[local, cols]`` shards that
+    are exactly zero except at this rank's applied winner rows -- the format
+    ``dion2_post_orthogonalize_masked`` consumes (same as the global scope)."""
+    world_size, budget, cols = recv.shape
+    U_flat = torch.zeros(
+        total_rows + 1, cols, dtype=recv.dtype, device=recv.device
+    )
+    U_flat.scatter_(0, src_index.reshape(-1, 1).expand(-1, cols), recv.reshape(-1, cols))
+    return list(U_flat[: n_params * local].view(n_params, local, cols).unbind(0))
+
+
+def dion2_capped_packed_async(
+    M_local: List[Tensor],
+    norms: Tensor,
+    padded_local: int,
+    fraction: float,
+    global_size: int,
+    device_rank: int,
+    world_size: int,
+    process_group: Optional[ProcessGroup],
+    per_rank: int,
+    budget: int,
+    kw: int,
+    newton_schulz_func: Callable,
+    flatten: bool,
+    epsilon: Tensor,
+) -> Generator[None, None, List[Tensor]]:
+    """
+    "global_capped" packed transport: winner selection (norms-only all-gather)
+    -> pooled pack -> fixed-size a2a with an int32 count header bit-cast into
+    each chunk's first row -> assemble + Newton-Schulz -> repack -> reverse a2a
+    -> scatter home. All collective sizes are construction-time constants; the
+    data-dependence lives entirely in GPU-side gathers/scatters (no host
+    syncs). The header makes the sender's packing plan the single source of
+    truth -- the receiver never recomputes winner counts.
+    """
+    winner, thresh = yield from dion2_capped_select_async(
+        norms, padded_local, fraction, global_size,
+        device_rank, world_size, process_group,
+    )
+    payload, counts, src_index, deferred = dion2_capped_pack(
+        M_local, norms, winner, thresh, world_size, per_rank, budget
+    )
+    CAPPED_STATS["deferred_rows"] = deferred
+    CAPPED_STATS["winner_rows"] = winner.sum()
+
+    # Header row: int32 counts bit-cast into bf16 (2 bf16 per int32). Counts
+    # are NOT representable as bf16 *values* (exact only to 256); the bit-cast
+    # round-trips exactly.
+    cols = payload.size(-1)
+    header = torch.zeros(
+        world_size, 1, cols, dtype=torch.bfloat16, device=payload.device
+    )
+    header[:, 0, : per_rank * 2] = counts.contiguous().view(torch.bfloat16)
+    chunks = torch.cat([header, payload], dim=1)  # [world, 1 + budget, cols]
+
+    send = [c.contiguous() for c in chunks.unbind(0)]
+    recv = [torch.empty_like(c) for c in send]
+    work = dist.all_to_all(recv, send, group=process_group, async_op=True)
+    yield
+    work.wait()
+
+    recv_t = torch.stack(recv)
+    counts_recv = recv_t[:, 0, : per_rank * 2].contiguous().view(torch.int32)
+    ns_input, ns_index = dion2_capped_assemble(
+        recv_t[:, 1:].contiguous(), counts_recv, kw
+    )
+    ns_out = muon_update_newton_schulz(
+        ns_input, newton_schulz_func=newton_schulz_func,
+        flatten=flatten, epsilon=epsilon,
+    )
+    back = dion2_capped_repack(ns_out, ns_index, world_size, budget)
+
+    send2 = [c.contiguous() for c in back.unbind(0)]
+    recv2 = [torch.empty_like(c) for c in send2]
+    work = dist.all_to_all(recv2, send2, group=process_group, async_op=True)
+    yield
+    work.wait()
+
+    local = M_local[0].size(-2)
+    n_pad = world_size * per_rank
+    return dion2_capped_unpack(
+        torch.stack(recv2), src_index, len(M_local), local, n_pad * local
     )
-    return vals[:, -1:]
 
 
 @_inductor_workaround
@@ -557,7 +901,6 @@ def dion2_pre_orthogonalize(
     ef_decay: Tensor,
     select_dim: int,
     k_override: Optional[int] = None,
-    thresh: Optional[Tensor] = None,
 ) -> Tuple[List[Tensor], List[Tensor]]:
     """
     Update momentum with gradient and compute the input to orthogonalization.
@@ -578,21 +921,8 @@ def dion2_pre_orthogonalize(
     downstream megabatch pads U to ``local_comm_size=k_override`` for a uniform
     alltoall, and indices stay at the real selected count.
 
-    ``thresh`` (the "global_capped" scope) is the per-matrix global
-    top-``k*world_size`` norm threshold ``[N, 1]`` from
-    ``dion2_capped_threshold_async``. Passing it means the momentum was ALREADY
-    accumulated by ``dion2_pre_accumulate_norms`` (the single accumulation
-    owner on the capped path -- the gathered norms must reflect the
-    post-accumulation momentum), so the ``M += G`` here is skipped. The top-k
-    gather itself is unchanged: a rank's winners (slices at/above ``thresh``)
-    are by construction its top-norm slices, so they already lead the plain
-    norm-order selection. What changes is what happens to the non-winner
-    "filler" slices occupying spare slots: they are masked to zero in the
-    communicated ``U_selected`` (zero slices pass through Newton-Schulz as
-    exact zeros, so they produce no weight update) and they skip error-feedback
-    decay, letting their momentum accumulate until they cross the global
-    threshold on a later step. Without this masking, capped selection would be
-    trajectory-identical to "local" scope.
+    (The "global_capped" scope does not pass through here at all -- it uses
+    the packed transport, ``dion2_capped_packed_async``, and the masked post.)
     """
     dtype = M[0].dtype
 
@@ -615,12 +945,9 @@ def dion2_pre_orthogonalize(
         k = max(1, int(math.ceil(fraction * num_select)))
     k_topk = min(k, num_select)
 
-    # Update momentum: M = M + G. Skipped when `thresh` is given ("global_
-    # capped"): dion2_pre_accumulate_norms is the single accumulation owner on
-    # that path, so the gathered norms reflect the post-accumulation momentum.
-    if thresh is None:
-        G = [g.to(dtype=dtype) for g in G]
-        torch._foreach_add_(M, G)
+    # Update momentum: M = M + G
+    G = [g.to(dtype=dtype) for g in G]
+    torch._foreach_add_(M, G)
 
     # Empty local shard along select_dim: FSDP2 contiguous chunking leaves this
     # rank with a size-0 shard when the param's sharded dim is smaller than
@@ -661,35 +988,15 @@ def dion2_pre_orthogonalize(
         )
         selected_stacked = torch.gather(M_stacked, dim=-1, index=indices_expanded)
 
-    # "global_capped" capacity rule: only slices at/above the global threshold
-    # ("winners") are applied. Winners lead the norm-order selection above by
-    # construction, so `indices` needs no change; non-winner "fillers" in spare
-    # slots are masked to zero for communication (zero slices pass through
-    # Newton-Schulz as exact zeros -> no weight update) and keep their momentum
-    # un-decayed so it accumulates until they win. `>=` marks hard ties at the
-    # threshold as winners, which can inflate the winner set past k*world_size
-    # -- measure-zero for continuous fp32 norms and bounded by the per-rank k
-    # cap, so benign.
-    if thresh is not None:
-        sel_norms = torch.gather(slice_norms, dim=-1, index=indices)  # [N, k_topk]
-        winner = sel_norms >= thresh
-        winner_exp = winner.unsqueeze(-1) if select_dim == -2 else winner.unsqueeze(-2)
-        ef_factor = torch.where(winner_exp, ef_decay, torch.ones_like(ef_decay))
-        U_stacked = selected_stacked * winner_exp.to(selected_stacked.dtype)
-    else:
-        ef_factor = ef_decay
-        U_stacked = selected_stacked
-
-    # Apply error feedback decay to selected slices in the original M tensors
-    # (capped scope: winners only). We reuse the already-gathered slices and
-    # write them back (scaled) using scatter_, which places values into
-    # positions specified by the index tensor. The .to(dtype) matters for bf16
-    # momentum: the capped ef_factor is a *dimensioned* fp32 tensor, so
-    # bf16 * fp32 promotes to fp32 (unlike the 0-dim ef_decay scalar, which
-    # loses promotion to a dimensioned bf16), and scatter_ requires src dtype
-    # to match M exactly. Multiply in fp32, round once into M's dtype.
+    # Apply error feedback decay to selected slices in the original M tensors.
+    # We reuse the already-gathered slices and write them back (scaled) using
+    # scatter_, which places values into positions specified by the index
+    # tensor. The .to(dtype) guards bf16 momentum: scatter_ requires src dtype
+    # to match M exactly, and it keeps a single fp32-multiply-then-round if
+    # ef_decay ever becomes a dimensioned tensor (a 0-dim fp32 scalar loses
+    # type promotion to bf16, but a dimensioned one would win it).
     indices_list = list(indices.unbind(dim=0))
-    ef_src_list = list((selected_stacked * ef_factor).to(dtype).unbind(dim=0))
+    ef_src_list = list((selected_stacked * ef_decay).to(dtype).unbind(dim=0))
     for m, idx, ef_src in zip(M, indices_list, ef_src_list):
         if select_dim == -2:
             idx_exp = idx.unsqueeze(-1).expand(*idx.shape, m.size(-1))
@@ -698,7 +1005,7 @@ def dion2_pre_orthogonalize(
         m.scatter_(dim=select_dim, index=idx_exp, src=ef_src)
 
     # Convert to bf16 and unstack for communication
-    U_selected = list(U_stacked.to(dtype=torch.bfloat16).unbind(dim=0))
+    U_selected = list(selected_stacked.to(dtype=torch.bfloat16).unbind(dim=0))
 
     return U_selected, indices_list
 
diff --git a/dion/nordion2.py b/dion/nordion2.py
index 92d503a..76b65c1 100644
--- a/dion/nordion2.py
+++ b/dion/nordion2.py
@@ -19,7 +19,7 @@
     dion2_post_orthogonalize,
     dion2_pre_accumulate,
     dion2_pre_accumulate_norms,
-    dion2_capped_threshold_async,
+    dion2_capped_packed_async,
     _make_select_and_orthogonalize,
 )
 from .normuon import normuon_normalization_stacked, _normuon_normalization_core
@@ -57,11 +57,11 @@ class NorDion2(DistributedOrthoBase):
             assembled whole matrix -- layout-invariant/reproducible and better-
             converging. "local": per-shard top-k (union) -- cheaper comm but a
             sharding-dependent approximation that converges slightly worse; opt
-            in when comm-bound at large scale. "global_capped": global
-            top-``k*world_size`` selection at local comm cost via a norms-only
-            all-gather; winners-only updates (spare slots carry zeros), overflow
-            winners deferred through error feedback (see Dion2). No-op off the
-            row-sharded path.
+            in when comm-bound at large scale. "global_capped": exact global
+            top-``ceil(fraction*global)`` selection at ~local comm cost via a
+            norms-only all-gather + pooled packed transport; overflow winners
+            deferred through error feedback (see Dion2, incl.
+            ``capacity_factor``). No-op off the row-sharded path.
 
     NorDion2 optimizer applying Dion2 update to NorMuon
     """
@@ -85,6 +85,7 @@ def __init__(
         newton_schulz_func: Optional[Callable] = None,
         triton_post_ortho: bool = False,
         selection_scope: str = "global",
+        capacity_factor: Optional[float] = None,
     ):
         # Validate hyperparameters
         if lr < 0.0:
@@ -96,6 +97,10 @@ def __init__(
                 f"selection_scope must be 'local', 'global', or 'global_capped', "
                 f"got {selection_scope!r}"
             )
+        if capacity_factor is not None and capacity_factor < 1.0:
+            raise ValueError(
+                f"capacity_factor must be None (auto) or >= 1.0, got {capacity_factor}"
+            )
         if mu < 0.0:
             raise ValueError(f"Invalid momentum factor (mu): {mu}")
         if muon_beta2 < 0.0:
@@ -121,6 +126,7 @@ def __init__(
             algorithm="nordion2",
             step=0,
             selection_scope=selection_scope,
+            capacity_factor=capacity_factor,
         )
         super().__init__(
             params, distributed_mesh, "nordion2", defaults,
@@ -187,6 +193,7 @@ def _create_ortho_tasks(
                 newton_schulz_func=self._newton_schulz_func,
                 triton_post_ortho=self._triton_post_ortho,
                 selection_scope=group["selection_scope"],
+                capacity_factor=group["capacity_factor"],
             )
 
             shape_groups: dict[tuple, list] = defaultdict(list)
@@ -253,6 +260,7 @@ def nordion2_update_megabatch_async(
     newton_schulz_func: Optional[Callable] = None,
     triton_post_ortho: bool = False,
     selection_scope: str = "global",
+    capacity_factor: Optional[float] = None,
 ) -> Generator[None, None, None]:
     """
     Mega-batched NorDion2 update: processes ALL same-shape parameters in one
@@ -273,7 +281,6 @@ def nordion2_update_megabatch_async(
 
     # comm_dim for sharded communication: use select_dim
     comm_dim = select_dim if is_sharded else None
-    global_scope = selection_scope == "global" and comm_dim is not None
 
     # On the sharded path X[0] must still be a DTensor, so .shape[comm_dim]
     # is the unsharded global size. The megabatch fn uses this to compute
@@ -290,6 +297,80 @@ def nordion2_update_megabatch_async(
     else:
         global_dim_size = None
 
+    # --- "global_capped" packed-path decision; mirrors dion2 (see there) ---
+    capped_packed = (
+        selection_scope == "global_capped"
+        and comm_dim is not None
+        and process_group is not None
+        and X[0].ndim == 2
+    )
+    if capped_packed:
+        padded_local = (global_dim_size + world_size - 1) // world_size
+        k = max(1, int(math.ceil(fraction * padded_local)))
+        per_rank = (N + world_size - 1) // world_size
+        if capacity_factor is None:
+            c = 1.0 + 2.0 * math.sqrt((1.0 - 1.0 / world_size) / (per_rank * k))
+        else:
+            c = float(capacity_factor)
+        budget = int(math.ceil(c * per_rank * k))
+        # Route to exact "global" when packing cannot pay for itself (slack
+        # reaches the full shard) or the int32 count header cannot fit in the
+        # chunk's first row (2 bf16 slots per matrix).
+        if budget >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]:
+            capped_packed = False
+
+    global_scope = comm_dim is not None and (
+        selection_scope == "global"
+        or (selection_scope == "global_capped" and not capped_packed)
+    )
+
+    if capped_packed:
+        # --- "global_capped": winner-only packed transport (see dion2) ---
+        # Returns full-size shards, zero except at applied winner rows -- the
+        # masked post's format. Deferred winners skip both EF decay and the
+        # muon_beta2 variance update (the masked post only touches nonzero
+        # rows), so their momentum accumulates until they win a slot.
+        norms = dion2_pre_accumulate_norms(
+            G=to_local(G), M=to_local(M), select_dim=select_dim
+        )
+        U_ortho = yield from dion2_capped_packed_async(
+            M_local=to_local(M),
+            norms=norms,
+            padded_local=padded_local,
+            fraction=fraction,
+            global_size=global_dim_size,
+            device_rank=device_rank,
+            world_size=world_size,
+            process_group=process_group,
+            per_rank=per_rank,
+            budget=budget,
+            kw=k * world_size,
+            newton_schulz_func=newton_schulz_func,
+            flatten=flatten,
+            epsilon=epsilon,
+        )
+        if adjust_lr is None:
+            adjusted_lr = lr
+        elif adjust_lr == "spectral_norm":
+            adjusted_lr = adjust_lr_spectral_norm(lr, X[0].shape, flatten=flatten)
+        elif adjust_lr == "rms_norm":
+            adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten)
+        else:
+            raise ValueError(f"Unknown adjust_lr: {adjust_lr}")
+        nordion2_post_orthogonalize_masked(
+            X=to_local(X),
+            M=to_local(M),
+            V=to_local(V),
+            U=U_ortho,
+            base_lr=lr,
+            adjusted_lr=adjusted_lr,
+            weight_decay=weight_decay,
+            ef_decay=momentum,
+            muon_beta2=muon_beta2,
+            select_dim=select_dim,
+        )
+        return
+
     if global_scope:
         # --- Global selection: send the full shard, select after assembly ---
         U_full = dion2_pre_accumulate(G=to_local(G), M=to_local(M))
@@ -347,24 +428,6 @@ def nordion2_update_megabatch_async(
         k = None
     global_comm_dim_size = global_dim_size
 
-    # "global_capped": norms-only all-gather + global threshold; winners-only
-    # updates, zeros in spare slots -- see dion2. (Filler slices come back from
-    # Newton-Schulz as exact zeros, so they get no weight update; their variance
-    # buffer rows still decay by muon_beta2 this step, a benign side effect that
-    # recovers once the row wins a slot.)
-    thresh = None
-    if (
-        selection_scope == "global_capped"
-        and comm_dim is not None
-        and process_group is not None
-    ):
-        norms = dion2_pre_accumulate_norms(
-            G=to_local(G), M=to_local(M), select_dim=select_dim
-        )
-        thresh = yield from dion2_capped_threshold_async(
-            norms, padded_local, k, world_size, process_group
-        )
-
     # Update momentum and compute the inputs for orthogonalization
     # Dion2 pre-orthogonalizes differs from NorMuon by applying damping before updating momentum
     U_selected, indices_list = dion2_pre_orthogonalize(
@@ -374,7 +437,6 @@ def nordion2_update_megabatch_async(
         ef_decay=momentum,
         select_dim=select_dim,
         k_override=k,
-        thresh=thresh,
     )
 
     # Orthogonalize via shared megabatch communication
diff --git a/tests/test_dion2_selection_scope.py b/tests/test_dion2_selection_scope.py
index 3ce277d..fcfc4e8 100644
--- a/tests/test_dion2_selection_scope.py
+++ b/tests/test_dion2_selection_scope.py
@@ -32,7 +32,7 @@
 import torch.multiprocessing as mp
 from torch.distributed.tensor import DeviceMesh, distribute_tensor, Shard
 
-from dion.dion2 import Dion2, dion2_capped_threshold_async
+from dion.dion2 import Dion2, dion2_capped_select_async
 from dion.nordion2 import NorDion2
 from dion.polar_express import polar_express
 
@@ -243,26 +243,28 @@ def test_fraction_one_local_equals_global(OptCls, kw, tmp_path):
     torch.testing.assert_close(dl["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL)
 
 
-# ---- selection_scope="global_capped" ----
+# ---- selection_scope="global_capped" (pooled packed transport) ----
 #
-# The capacity rule: all-gather slice norms only, every rank derives the same
-# global top-(k*world_size) threshold, and only slices at/above it ("winners")
-# are applied -- spare slots travel as zeros and non-winners skip error-feedback
-# decay so their momentum accumulates until they win. NOTE the budget is
-# top-(k*world_size), NOT top-(ceil(fraction*global)); they coincide when
-# world_size divides the sharded dim (test below) and the capped budget is
-# slightly larger otherwise (empty-shard test below exploits this: with
-# k*world >= real rows, everything wins and capped(f) == global(1.0)).
+# All-gather slice norms only; every rank derives the same EXACT-COUNT global
+# top-ceil(fraction*global) winner set (stable-sort tie-break). Each rank
+# packs its winner rows into fixed-size chunks pooled across the megabatch's
+# matrices (int32 count header bit-cast into each chunk's first row); winners
+# overflowing the pooled budget defer via error feedback (no EF decay ->
+# momentum accumulates -> selected later). Only applied rows are decayed, by
+# the same masked post the global scope uses. Groups where packing cannot
+# save comm route to the exact "global" scope.
 
 
 def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
-                   fraction, boost, steps, dtype):
-    """Multi-step worker with a crafted gradient: small noise everywhere, plus
-    rows in ``boost`` set to a constant so their momentum L1 norm is exactly the
-    boost magnitude. Steps after the first use a zero gradient (isolates the
-    error-feedback dynamics). Saves the full W after every step. ``dtype`` sets
-    the param (and therefore momentum) dtype -- bf16 exercises the
-    mixed-precision EF write-back path that fp32 params can never reach."""
+                   fraction, boosts, steps, dtype):
+    """Multi-step, multi-param worker with crafted gradients: small noise
+    everywhere, plus rows in ``boosts[i]`` (one dict per param) set to a
+    constant so their momentum L1 norm is exactly the boost magnitude. Steps
+    after the first use a zero gradient (isolates the error-feedback
+    dynamics). Same-shape params land in ONE megabatch group, so multiple
+    params exercise the pooled packed transport. Saves the full W of every
+    param after every step. ``dtype`` sets the param (and therefore momentum)
+    dtype -- bf16 exercises the mixed-precision paths fp32 can never reach."""
     os.environ["MASTER_ADDR"] = "127.0.0.1"
     os.environ["MASTER_PORT"] = str(port)
     torch.cuda.set_device(rank)
@@ -272,14 +274,18 @@ def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
 
     rows, cols = shape
     g = torch.Generator(device=dev).manual_seed(1234)
-    W0 = torch.randn(rows, cols, generator=g, device=dev).to(dtype)
-    G = (torch.randn(rows, cols, generator=g, device=dev) * 0.01).to(dtype)
-    for r, mag in boost.items():
-        G[r] = mag / cols  # exact row L1 norm == mag
+    W0s, Gs, params = [], [], []
+    for boost in boosts:
+        W0 = torch.randn(rows, cols, generator=g, device=dev).to(dtype)
+        G = (torch.randn(rows, cols, generator=g, device=dev) * 0.01).to(dtype)
+        for r, mag in boost.items():
+            G[r] = mag / cols  # exact row L1 norm == mag
+        W0s.append(W0)
+        Gs.append(G)
+        params.append(torch.nn.Parameter(distribute_tensor(W0, mesh, [Shard(0)])))
 
-    p = torch.nn.Parameter(distribute_tensor(W0, mesh, [Shard(0)]))
     opt = OptCls(
-        [dict(params=[p])],
+        [dict(params=params)],
         distributed_mesh=mesh,
         lr=0.1,
         fraction=fraction,
@@ -289,21 +295,23 @@ def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape,
     )
     Ws = []
     for s in range(steps):
-        grad = G if s == 0 else torch.zeros_like(G)
-        p.grad = distribute_tensor(grad, mesh, [Shard(0)])
+        for p, G in zip(params, Gs):
+            grad = G if s == 0 else torch.zeros_like(G)
+            p.grad = distribute_tensor(grad, mesh, [Shard(0)])
         opt.step()
-        Ws.append(p.detach().full_tensor().cpu())
+        Ws.append([p.detach().full_tensor().cpu() for p in params])
     if rank == 0:
-        torch.save({"W0": W0.cpu(), "G": G.cpu(), "Ws": Ws}, out_path)
+        torch.save({"W0": [w.cpu() for w in W0s], "Ws": Ws}, out_path)
     dist.destroy_process_group()
 
 
 def _run_capped(OptCls, scope, kw, port, tmp, *, shape=(16, 32), fraction=0.25,
                 boost=None, steps=1, dtype=torch.float32):
     out = str(tmp / f"out_{port}.pt")
+    boosts = boost if isinstance(boost, list) else [boost or {}]
     mp.spawn(
         _capped_worker,
-        args=(2, port, OptCls, scope, kw, out, shape, fraction, boost or {},
+        args=(2, port, OptCls, scope, kw, out, shape, fraction, boosts,
               steps, dtype),
         nprocs=2,
         join=True,
@@ -329,28 +337,30 @@ def test_capped_scope_matches_global_when_winners_fit(OptCls, kw, tmp_path):
     boost = {0: 3.0, 1: 2.9, 8: 2.8, 9: 2.7}
     dc = _run_capped(OptCls, "global_capped", kw, 29680, tmp_path, boost=boost)
     dg = _run_capped(OptCls, "global", kw, 29681, tmp_path, boost=boost)
-    torch.testing.assert_close(dc["Ws"][0], dg["Ws"][0],
+    torch.testing.assert_close(dc["Ws"][0][0], dg["Ws"][0][0],
                                rtol=BF16_RTOL, atol=BF16_ATOL)
 
 
 @pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
 @pytest.mark.parametrize("OptCls,kw", [
-    (Dion2, {"ef_decay": 0.5, "weight_decay": 0.0}),
-    (NorDion2, {"mu": 0.5, "muon_beta2": 0.95, "weight_decay": 0.0}),
+    (Dion2, {"ef_decay": 0.5, "weight_decay": 0.0, "capacity_factor": 1.0}),
+    (NorDion2, {"mu": 0.5, "muon_beta2": 0.95, "weight_decay": 0.0,
+                "capacity_factor": 1.0}),
 ])
 def test_capped_scope_defers_overflow_winners(OptCls, kw, tmp_path):
-    """All 4 global winners live on rank 0 but it only has k=2 slots: step 1
-    applies exactly its top-2 {0,1} (rows 2,3 deferred WITHOUT error-feedback
-    decay; rank 1's spare slots carry zeros, so NO row of rank 1 changes --
+    """All 4 global winners live on rank 0 but its chunk budget is pinned to
+    k=2 rows (capacity_factor=1.0, single matrix): step 1 sends exactly its
+    top-2 {0,1} by normalized priority (rows 2,3 deferred WITHOUT
+    error-feedback decay; rank 1 sends nothing, so NO row of rank 1 changes --
     this is what distinguishes capped from local, which would fill and apply
-    rank 1's top-2). Step 2 (zero grad): rows 0,1 decayed to 1.5/1.45 while the
-    deferred 2,3 kept 2.8/2.7, so exactly {2,3} are applied. weight_decay=0
-    makes non-applied rows bitwise unchanged."""
+    rank 1's top-2). Step 2 (zero grad): rows 0,1 decayed to 1.5/1.45 while
+    the deferred 2,3 kept 2.8/2.7, so exactly {2,3} are applied.
+    weight_decay=0 makes non-applied rows bitwise unchanged."""
     boost = {0: 3.0, 1: 2.9, 2: 2.8, 3: 2.7}
     d = _run_capped(OptCls, "global_capped", kw, 29682, tmp_path,
                     boost=boost, steps=2)
-    assert _changed_rows(d["Ws"][0], d["W0"]) == {0, 1}
-    assert _changed_rows(d["Ws"][1], d["Ws"][0]) == {2, 3}
+    assert _changed_rows(d["Ws"][0][0], d["W0"][0]) == {0, 1}
+    assert _changed_rows(d["Ws"][1][0], d["Ws"][0][0]) == {2, 3}
 
 
 @pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs")
@@ -372,24 +382,25 @@ def test_capped_scope_uneven_shards_fraction_one_equals_global(OptCls, kw, tmp_p
 
 @pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs")
 def test_capped_scope_uneven_shards_runs(tmp_path):
-    """Partial-fraction sanity on uneven shards: finite and actually updates."""
-    d = _run_sharded(Dion2, "global_capped", 0.35, {}, 4, 29685, tmp_path,
-                     shape=(17, 32))
+    """Partial-fraction sanity on uneven shards (one short shard) through the
+    real packed path: capacity_factor is pinned to 1.0 because the auto slack
+    on this tiny k would reach the full shard and route the group to global."""
+    d = _run_sharded(Dion2, "global_capped", 0.35, {"capacity_factor": 1.0},
+                     4, 29685, tmp_path, shape=(17, 32))
     assert torch.isfinite(d["W"]).all()
     assert not torch.allclose(d["W"], d["W0"], atol=1e-6)
 
 
 @pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs")
 def test_capped_scope_empty_shard(tmp_path):
-    """3 rows over 4 ranks chunk to [1,1,1,0]: rank 3's shard is empty. The
-    capped budget k*world = 4 exceeds the 3 real rows, so the threshold is the
-    -1 pad value and every real row wins -- which also makes capped(0.5) equal
-    global(1.0), a live demonstration of the top-(k*world) vs
-    top-(ceil(fraction*global)) budget distinction."""
+    """3 rows over 4 ranks chunk to [1,1,1,0]: rank 3's shard is empty. For a
+    matrix this tiny the pooled budget reaches the full shard, so the group
+    must route to the exact global scope -- capped(f) == global(f), with the
+    exact top-ceil(fraction*global) budget (2 rows at f=0.5, not k*world)."""
     dc = _run_sharded(Dion2, "global_capped", 0.5, {}, 4, 29686, tmp_path,
                       shape=(3, 32))
     assert torch.isfinite(dc["W"]).all()
-    dg = _run_sharded(Dion2, "global", 1.0, {}, 4, 29687, tmp_path,
+    dg = _run_sharded(Dion2, "global", 0.5, {}, 4, 29687, tmp_path,
                       shape=(3, 32))
     torch.testing.assert_close(dc["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL)
 
@@ -410,12 +421,12 @@ def test_scope_bf16_momentum(OptCls, kw, scope, tmp_path):
     tests can never catch this -- it only fires in real bf16 training."""
     port = 29690 + (scope == "global_capped") * 2 + (OptCls is NorDion2)
     d = _run_capped(OptCls, scope, kw, port, tmp_path, dtype=torch.bfloat16)
-    W = d["Ws"][0]
+    W = d["Ws"][0][0]
     assert torch.isfinite(W.float()).all()
-    assert not torch.equal(W, d["W0"])  # the step actually applied an update
+    assert not torch.equal(W, d["W0"][0])  # the step actually applied an update
 
 
-def _thresh_consistency_worker(rank, world_size, port):
+def _select_consistency_worker(rank, world_size, port):
     os.environ["MASTER_ADDR"] = "127.0.0.1"
     os.environ["MASTER_PORT"] = str(port)
     torch.cuda.set_device(rank)
@@ -424,41 +435,142 @@ def _thresh_consistency_worker(rank, world_size, port):
 
     # Distinct norms everywhere (no ties), different on each rank; one short
     # shard (rank 1 has 3 of padded_local=4 slices) to exercise the -1 pad.
-    N, padded_local, k = 2, 4, 2
+    N, padded_local = 2, 4
+    global_size = 7  # 4 + 3 real rows
+    fraction = 0.5  # k_total = ceil(0.5 * 7) = 4 per matrix
     local = padded_local - rank  # rank0: 4, rank1: 3
     norms = (torch.arange(N * local, dtype=torch.float32, device=dev)
              .reshape(N, local) * 1.7 + 3.0 + rank * 0.37)
 
-    gen = dion2_capped_threshold_async(norms, padded_local, k, world_size, None)
+    gen = dion2_capped_select_async(
+        norms, padded_local, fraction, global_size, rank, world_size, None
+    )
     next(gen)
     try:
         next(gen)
         raise AssertionError("generator should be exhausted after the collective")
     except StopIteration as e:
-        thresh = e.value
+        winner, thresh = e.value
 
+    assert winner.shape == (N, local) and winner.dtype == torch.bool
     assert thresh.shape == (N, 1)
     # (f) rank consistency: every rank must hold the bit-identical threshold.
     both = [torch.empty_like(thresh) for _ in range(world_size)]
-    dist.all_gather(both, thresh.contiguous())  # thresh is a topk slice view
+    dist.all_gather(both, thresh.contiguous())
     assert torch.equal(both[0], both[1])
-    # With distinct norms and k_total <= real slices, exactly k*world_size
-    # slices per matrix sit at/above the threshold globally.
-    winners = (norms >= thresh).sum()
+    # Exact-count winner set: exactly ceil(fraction*global) winners per matrix
+    # across all ranks -- the receiver buffer bound relies on this.
+    winners = winner.sum()
     dist.all_reduce(winners)
-    assert winners.item() == N * k * world_size
+    assert winners.item() == N * 4
     dist.destroy_process_group()
 
 
 @pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
-def test_capped_threshold_rank_consistent():
-    mp.spawn(_thresh_consistency_worker, args=(2, 29688), nprocs=2, join=True)
+def test_capped_select_rank_consistent():
+    mp.spawn(_select_consistency_worker, args=(2, 29688), nprocs=2, join=True)
 
 
-def test_capped_threshold_negative_pad_raises():
+def test_capped_select_negative_pad_raises():
     """local_size > padded_local would make F.pad silently TRIM the norms; the
     guard must raise instead. (Raise happens before the collective, so no
     process group is needed.) Pure-CPU unit test."""
-    gen = dion2_capped_threshold_async(torch.zeros(1, 5), 3, 1, 2, None)
+    gen = dion2_capped_select_async(torch.zeros(1, 5), 3, 0.5, 8, 0, 2, None)
     with pytest.raises(ValueError, match="exceeds padded size"):
         next(gen)
+
+
+def test_capped_header_bitcast_roundtrip():
+    """The packed transport ships int32 per-matrix counts bit-cast into the
+    bf16 chunk's first row (2 bf16 slots per int32). The bit-cast must
+    round-trip exactly for counts far beyond 256 (where bf16 *values* lose
+    integer exactness -- the failure mode this encoding avoids). Pure CPU."""
+    counts = torch.tensor([[320, 4097, 0], [70000, 1, 2 ** 20]], dtype=torch.int32)
+    world, per_rank = counts.shape
+    cols = 32
+    header = torch.zeros(world, cols, dtype=torch.bfloat16)
+    header[:, : per_rank * 2] = counts.contiguous().view(torch.bfloat16)
+    back = header[:, : per_rank * 2].contiguous().view(torch.int32)
+    assert torch.equal(back, counts)
+
+
+def test_capped_pack_roundtrip_identity():
+    """pack -> (simulated a2a) -> assemble -> identity-NS -> repack ->
+    (simulated reverse a2a) -> unpack must return each rank's kept winner rows
+    bit-exactly (bf16) at their home positions and exact zeros elsewhere, with
+    sent rows always a subset of the winner set. Covers pooled deferral (4
+    winners vs budget 3), a matrix with zero winners anywhere ("blackout"),
+    and uneven per-(rank, matrix) counts. Pure CPU, no process group."""
+    from dion.dion2 import (
+        dion2_capped_pack, dion2_capped_assemble,
+        dion2_capped_repack, dion2_capped_unpack,
+    )
+    torch.manual_seed(0)
+    world, per_rank, N, local, cols, budget, kw = 2, 2, 4, 4, 8, 3, 4
+
+    Ms, normss = [], []
+    for _ in range(world):
+        M = [torch.randn(local, cols) + 5 for _ in range(N)]
+        Ms.append(M)
+        normss.append(torch.stack([m.abs().sum(-1) for m in M]))
+    w0 = torch.zeros(N, local, dtype=torch.bool)
+    w1 = torch.zeros(N, local, dtype=torch.bool)
+    w0[0, :3] = True; w1[0, 0] = True  # m0: rank0 pooled 3+1 > budget -> defer
+    w0[1, 1] = True; w1[1, 2] = True   # m1: 1 + 1
+    w0[2, :2] = True                   # m2: 2 + 0; m3: blackout
+    winners = [w0, w1]
+    threshs = [torch.ones(N, 1)] * world  # uniform scale: prio == raw norm
+
+    packs = [
+        dion2_capped_pack(Ms[r], normss[r], winners[r], threshs[r],
+                          world, per_rank, budget)
+        for r in range(world)
+    ]
+    backs = {}
+    for d in range(world):
+        recv = torch.stack([packs[r][0][d] for r in range(world)])
+        cnt = torch.stack([packs[r][1][d] for r in range(world)])
+        ns_in, ns_ix = dion2_capped_assemble(recv, cnt, kw)
+        backs[d] = dion2_capped_repack(ns_in, ns_ix, world, budget)
+
+    assert int(packs[0][3]) == 1 and int(packs[1][3]) == 0  # deferred counts
+    for r in range(world):
+        recv2 = torch.stack([backs[d][r] for d in range(world)])
+        U = dion2_capped_unpack(recv2, packs[r][2], N, local,
+                                world * per_rank * local)
+        for i in range(N):
+            sent = U[i].abs().sum(-1) > 0
+            assert bool((sent <= winners[r][i]).all())  # winners only
+            assert torch.equal(U[i][sent], Ms[r][i].to(torch.bfloat16)[sent])
+            assert U[i][~sent].abs().sum() == 0
+
+
+@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs")
+def test_capped_pooled_slots_absorb_cross_matrix_skew(tmp_path):
+    """The T2 value proposition. 4 same-shape params form ONE megabatch group;
+    world=2 gives per_rank=2, so matrices {0,1} assemble on rank 0 and {2,3}
+    on rank 1, and rank 0's chunk to dest 0 POOLS matrices 0 and 1 with a
+    budget of per_rank*k = 4 rows (capacity_factor pinned to 1.0).
+
+    Matrix 0 has ALL 4 of its global winners on rank 0 (needs 4 > k=2 rows of
+    "its own" slots); matrix 1 contributes 1 rank-0 winner at a 100x smaller
+    norm scale but the HIGHEST scale-normalized priority (3.0/2.7 = 1.111).
+    Rank 0's pooled winners = 5 > budget 4, so exactly one row defers -- and
+    it must be matrix 0's LOWEST normalized priority (row 3, 275/275 = 1.0),
+    NOT matrix 1's small-raw-norm winner, which a raw-norm packing order
+    would evict. Per-matrix walls (the pre-pooling design) could never have
+    sent matrix 0's third winner at all. weight_decay=0 makes unapplied rows
+    bitwise unchanged."""
+    boosts = [
+        {0: 300.0, 1: 290.0, 2: 280.0, 3: 275.0},  # m0: prios 1.09/1.05/1.02/1.0
+        {0: 3.0, 8: 2.9, 9: 2.8, 10: 2.7},          # m1: prio(row0) = 1.111
+        {},
+        {},
+    ]
+    kw = {"capacity_factor": 1.0, "weight_decay": 0.0}
+    d = _run_capped(Dion2, "global_capped", kw, 29694, tmp_path, boost=boosts)
+    # Matrix 0: rank 0 keeps its top-3 by normalized priority; row 3 deferred.
+    assert _changed_rows(d["Ws"][0][0], d["W0"][0]) == {0, 1, 2}
+    # Matrix 1: the small-scale winner survives the shared chunk; rank 1's
+    # winners {8,9,10} ride its own uncontended chunk.
+    assert _changed_rows(d["Ws"][0][1], d["W0"][1]) == {0, 8, 9, 10}

From 569b395b2d33a8252bd2619d6d7024d7df344667 Mon Sep 17 00:00:00 2001
From: alint77 <ali.nt1377@gmail.com>
Date: Fri, 3 Jul 2026 13:16:01 +0200
Subject: [PATCH 5/6] fix(capped): review fixes -- stats semantics, stable
 overflow order, routing break-even

- CAPPED_STATS is now a reliable rank-local per-step reading: cleared once
  per optimizer step (top of _create_ortho_tasks, executed on the first task
  pull) and accumulated across the step's packed megabatches, instead of
  last-writer-wins across interleaved shape-group tasks. Empty dict means no
  packed group ran (no stale reuse); semantics documented at the definition.
- Overflow deferral uses stable argsort instead of topk: tied normalized
  priorities break by flat (matrix, row) order, so WHICH row defers is
  reproducible across runs/devices -- matching the winner selection's stable
  tie contract. (Sender-local either way; this is a determinism fix, not a
  cross-rank one.)
- Degeneracy routing accounts for the header row: the packed chunk ships
  1 + budget rows, so break-even against the full shard is budget + 1.

28/28 on 2x/4x H100.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 dion/dion2.py    | 35 +++++++++++++++++++++++++----------
 dion/nordion2.py | 14 ++++++++++----
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/dion/dion2.py b/dion/dion2.py
index b0573d0..bc76bb3 100644
--- a/dion/dion2.py
+++ b/dion/dion2.py
@@ -148,6 +148,9 @@ def _create_ortho_tasks(
         Mega-batched Dion2 task creation: groups ALL same-shape parameters
         into a single task to minimize communication rounds and kernel launches.
         """
+        # New optimizer step: reset the rank-local capped-deferral counters so
+        # they accumulate exactly this step's packed megabatches.
+        CAPPED_STATS.clear()
         for group in param_groups:
             assert group["algorithm"] == self._algo_name
             assert all(
@@ -339,10 +342,12 @@ def dion2_update_megabatch_async(
         else:
             c = float(capacity_factor)
         budget = int(math.ceil(c * per_rank * k))
-        # Route to exact "global" when packing cannot pay for itself (slack
-        # reaches the full shard) or the int32 count header cannot fit in the
-        # chunk's first row (2 bf16 slots per matrix).
-        if budget >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]:
+        # Route to exact "global" when packing cannot pay for itself -- the
+        # packed chunk sends 1 + budget rows (count header included), so
+        # break-even against sending the full shard is budget + 1 -- or when
+        # the int32 count header cannot fit in the chunk's first row (2 bf16
+        # slots per matrix).
+        if budget + 1 >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]:
             capped_packed = False
 
     # Decide whether selection happens before communication ("local", and only
@@ -560,9 +565,13 @@ def dion2_pre_accumulate_norms(
     return torch.stack(M, dim=0).norm(p=1, dim=norm_dim).float()
 
 
-# Deferral instrumentation for the "global_capped" packed path. Updated per
-# megabatch step with GPU tensors (no host sync); external code may read and
-# log e.g. CAPPED_STATS["deferred_rows"] / CAPPED_STATS["winner_rows"].
+# Deferral instrumentation for the "global_capped" packed path. RANK-LOCAL
+# semantics: cleared once per optimizer step (at ortho-task creation) and
+# ACCUMULATED across that step's packed megabatches (shape groups), as GPU
+# scalar tensors -- reading them costs no sync until the consumer calls
+# .item(). Empty dict => no packed group ran this step (do not reuse stale
+# values). "winner_rows" counts THIS rank's winners and can legitimately be
+# zero; aggregate across ranks before forming a deferral ratio.
 CAPPED_STATS: dict = {}
 
 
@@ -689,7 +698,13 @@ def dion2_capped_pack(
     score = torch.where(winner, prio, torch.full_like(prio, float("-inf")))
     score_d = score.view(world_size, per_rank * local)
     b_eff = min(budget, per_rank * local)
-    vals, idxs = torch.topk(score_d, b_eff, dim=-1)
+    # Stable argsort (not topk): on tied normalized priorities the flat
+    # (matrix, row) order is the deterministic secondary key, so WHICH row
+    # defers on overflow is reproducible across runs/devices -- matching the
+    # stable tie contract of the winner selection itself.
+    order = torch.argsort(score_d, dim=-1, descending=True, stable=True)
+    idxs = order[:, :b_eff]
+    vals = torch.gather(score_d, 1, idxs)
     kept_flat = torch.zeros_like(score_d, dtype=torch.bool)
     kept_flat.scatter_(1, idxs, vals.isfinite())  # winners only; spare slots stay empty
     kept = kept_flat.view(world_size, per_rank, local)
@@ -849,8 +864,8 @@ def dion2_capped_packed_async(
     payload, counts, src_index, deferred = dion2_capped_pack(
         M_local, norms, winner, thresh, world_size, per_rank, budget
     )
-    CAPPED_STATS["deferred_rows"] = deferred
-    CAPPED_STATS["winner_rows"] = winner.sum()
+    CAPPED_STATS["deferred_rows"] = CAPPED_STATS.get("deferred_rows", 0) + deferred
+    CAPPED_STATS["winner_rows"] = CAPPED_STATS.get("winner_rows", 0) + winner.sum()
 
     # Header row: int32 counts bit-cast into bf16 (2 bf16 per int32). Counts
     # are NOT representable as bf16 *values* (exact only to 256); the bit-cast
diff --git a/dion/nordion2.py b/dion/nordion2.py
index 76b65c1..233c2e4 100644
--- a/dion/nordion2.py
+++ b/dion/nordion2.py
@@ -15,6 +15,7 @@
 )
 from .opt_utils import AsyncTask, to_local
 from .dion2 import (
+    CAPPED_STATS,
     dion2_pre_orthogonalize,
     dion2_post_orthogonalize,
     dion2_pre_accumulate,
@@ -168,6 +169,9 @@ def _create_ortho_tasks(
         Mega-batched NorDion2 task creation: groups ALL same-shape parameters
         into a single task to minimize communication rounds and kernel launches.
         """
+        # New optimizer step: reset the rank-local capped-deferral counters so
+        # they accumulate exactly this step's packed megabatches.
+        CAPPED_STATS.clear()
         for group in param_groups:
             assert group["algorithm"] == self._algo_name
             assert all(
@@ -313,10 +317,12 @@ def nordion2_update_megabatch_async(
         else:
             c = float(capacity_factor)
         budget = int(math.ceil(c * per_rank * k))
-        # Route to exact "global" when packing cannot pay for itself (slack
-        # reaches the full shard) or the int32 count header cannot fit in the
-        # chunk's first row (2 bf16 slots per matrix).
-        if budget >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]:
+        # Route to exact "global" when packing cannot pay for itself -- the
+        # packed chunk sends 1 + budget rows (count header included), so
+        # break-even against sending the full shard is budget + 1 -- or when
+        # the int32 count header cannot fit in the chunk's first row (2 bf16
+        # slots per matrix).
+        if budget + 1 >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]:
             capped_packed = False
 
     global_scope = comm_dim is not None and (

From e58126703ee6ba6b56c90fd8fd02bbcc3d580c2a Mon Sep 17 00:00:00 2001
From: alint77 <ali.nt1377@gmail.com>
Date: Fri, 3 Jul 2026 20:38:13 +0200
Subject: [PATCH 6/6] feat(capped): fused Triton masked post-orthogonalize for
 global/global_capped

The masked posts (selection scopes "global" and "global_capped") previously
had no triton_post_ortho variant: only the indexed local-scope post got the
fused kernel from #81, so local enjoyed a single fused pass while the masked
scopes ran the generic compiled path (separate full-matrix sweeps for weight
decay, mask reduction, EF decay, NorMuon normalization, and weight update --
X crossing HBM twice, U twice, ~12 kernels per matrix).

Adds row-programmed masked kernels that derive selection from the nonzero
rows of U (no index map, preserving the masked transport's no-host-sync
property -- hyperparameter scalars are 0-dim CPU tensors, so .item() is
sync-free):

- _dion2_post_ortho_masked_kernel: single pass fusing weight decay, masked
  update (a*x + neg_b*u, fp32, one rounding) and EF decay on M. Unselected
  rows skip M entirely (bitwise-identical to multiply-by-one).
- _nordion2_post_ortho_masked_{stats,apply}_kernel: two phases forced by the
  NorMuon Frobenius rescale (matrix-level reduction). Phase A computes row
  sumsq, updates V, EF-decays M, and emits per-row partials; a batched
  deterministic torch reduction (no fp32 atomics) produces the per-matrix
  rescale; phase B applies the weight update.

All four masked call sites (dion2/nordion2 x global/capped) dispatch on the
existing triton_post_ortho flag, with fallback to the compiled masked posts
for CPU/no-triton, column selection, tensor subclasses, and mixed shapes.

Launch count per shape group drops from ~12*n_mat to 2*n_mat + 4, and memory
traffic to one pass over each of W/M/V/U. Tests: function-level parity vs the
compiled masked posts (bitwise on unselected rows, dtype-aware tolerances on
selected; edge cases: none/all selected, single row, non-pow2 wide cols, 3D
batch, empty shards) plus a 2-GPU NorDion2 global_capped end-to-end A/B.
28+17 relevant tests pass on 4xH100. (The 8 pre-existing
test_post_ortho_triton_falls_back_for_wrapper_subclass failures also fail at
the clean PR head -- inductor BackendCompilerFailed compiling the fused
fallback against _MockWrapperTensor on this torch build -- unrelated.)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dion/dion2.py                                |  10 +-
 dion/dion2_triton.py                         | 366 +++++++++++++++++++
 dion/nordion2.py                             |  10 +-
 tests/test_dion2_post_ortho_masked_triton.py | 301 +++++++++++++++
 4 files changed, 683 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_dion2_post_ortho_masked_triton.py

diff --git a/dion/dion2.py b/dion/dion2.py
index bc76bb3..bee8835 100644
--- a/dion/dion2.py
+++ b/dion/dion2.py
@@ -394,7 +394,10 @@ def dion2_update_megabatch_async(
             adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten)
         else:
             raise ValueError(f"Unknown adjust_lr: {adjust_lr}")
-        dion2_post_orthogonalize_masked(
+        _masked_post = dion2_post_orthogonalize_masked
+        if triton_post_ortho:
+            from .dion2_triton import dion2_post_orthogonalize_masked_triton as _masked_post
+        _masked_post(
             X=to_local(X),
             M=to_local(M),
             U=U_ortho,
@@ -439,7 +442,10 @@ def dion2_update_megabatch_async(
         # U_ortho rows are exactly zero except at the globally-selected positions
         # this rank owns. Apply error-feedback decay to those rows of M and the
         # masked weight update, both keyed off the nonzero mask (no indices).
-        dion2_post_orthogonalize_masked(
+        _masked_post = dion2_post_orthogonalize_masked
+        if triton_post_ortho:
+            from .dion2_triton import dion2_post_orthogonalize_masked_triton as _masked_post
+        _masked_post(
             X=to_local(X),
             M=to_local(M),
             U=U_ortho,
diff --git a/dion/dion2_triton.py b/dion/dion2_triton.py
index cf6a59a..a62163d 100644
--- a/dion/dion2_triton.py
+++ b/dion/dion2_triton.py
@@ -231,3 +231,369 @@ def dion2_post_orthogonalize_triton(
             index_map.stride(0),
             SELECT_ROWS=SELECT_ROWS,
         )
+
+
+# ---------------------------------------------------------------------------
+# Masked posts (selection_scope="global"/"global_capped")
+#
+# The masked post-orthogonalize has no index list: U arrives full-shard-shaped
+# and exactly zero outside the selected rows, so selection membership is
+# derived from the nonzero rows of U itself. These kernels fuse the entire
+# masked post into one (Dion2) or two (NorDion2) row-programmed passes,
+# replacing the compiled path's separate full-matrix sweeps (weight-decay
+# foreach, mask reduction, EF decay, NorMuon normalization, weight update).
+# Like the indexed kernel above, updates compute in fp32 with a single
+# rounding on store, and in-place writes rule out @triton.autotune.
+#
+# Scalars (lr, weight decay, ...) arrive as 0-dim CPU tensors, so .item() is
+# free (no GPU sync) and values are passed to the kernels by value.
+# ---------------------------------------------------------------------------
+
+
+@triton.heuristics(
+    {
+        "BLOCK_N": lambda args: max(128, min(1024, triton.next_power_of_2(args["N"]))),
+        "num_warps": lambda args: 8 if args["N"] >= 2048 else 4,
+    }
+)
+@triton.jit
+def _dion2_post_ortho_masked_kernel(
+    X_ptr,
+    M_ptr,
+    U_ptr,
+    a,
+    neg_b,
+    ef,
+    Mrows,
+    N,
+    x_stride_b,
+    x_stride_m,
+    x_stride_n,
+    m_stride_b,
+    m_stride_m,
+    m_stride_n,
+    u_stride_b,
+    u_stride_m,
+    u_stride_n,
+    BLOCK_N: tl.constexpr,
+):
+    """Fused masked post for Dion2: one program per (batch, row).
+
+    Pass 1 reduces |u| over the row to derive the selected flag (U is exactly
+    zero on non-selected rows). Pass 2 applies, in fp32 with one rounding:
+      selected:   x = a*x + neg_b*u   and   m = ef*m
+      unselected: x = a*x             and   m untouched (ef would be *1)
+    Skipping M entirely on unselected rows is bitwise identical to the
+    compiled path's multiply-by-one and saves the M read+write there.
+    """
+    pid = tl.program_id(0)
+    b = pid // Mrows
+    row = pid % Mrows
+    X_ptr += b * x_stride_b + row * x_stride_m
+    M_ptr += b * m_stride_b + row * m_stride_m
+    U_ptr += b * u_stride_b + row * u_stride_m
+
+    absacc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        msk = idx < N
+        u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32)
+        absacc += tl.abs(u)
+    sel = tl.sum(absacc) > 0.0
+
+    if sel:
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            msk = idx < N
+            x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32)
+            u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32)
+            m = tl.load(M_ptr + idx * m_stride_n, mask=msk, other=0.0).to(tl.float32)
+            tl.store(X_ptr + idx * x_stride_n, a * x + neg_b * u, mask=msk)
+            tl.store(M_ptr + idx * m_stride_n, ef * m, mask=msk)
+    else:
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            msk = idx < N
+            x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32)
+            tl.store(X_ptr + idx * x_stride_n, a * x, mask=msk)
+
+
+def dion2_post_orthogonalize_masked_triton(
+    X: List[Tensor],
+    M: List[Tensor],
+    U: List[Tensor],
+    base_lr: Tensor,
+    adjusted_lr: Tensor,
+    weight_decay: Tensor,
+    ef_decay: Tensor,
+    select_dim: int,
+):
+    """Triton-fused version of ``dion2_post_orthogonalize_masked``.
+
+    Falls back to the compiled masked post when the fast path does not apply
+    (no triton/CUDA, column selection, or tensor-subclass params). Only row
+    selection (``select_dim == -2``) is fused: that is the only mode the
+    global/global_capped scopes produce today.
+    """
+    use_fallback = (
+        not TRITON_AVAILABLE
+        or select_dim != -2
+        or not X
+        or not X[0].is_cuda
+        or any(is_traceable_wrapper_subclass(x) for x in X)
+    )
+    if use_fallback:
+        from .dion2 import dion2_post_orthogonalize_masked
+
+        dion2_post_orthogonalize_masked(
+            X=X, M=M, U=U, base_lr=base_lr, adjusted_lr=adjusted_lr,
+            weight_decay=weight_decay, ef_decay=ef_decay, select_dim=select_dim,
+        )
+        return
+
+    a = (1 - base_lr * weight_decay).item()
+    neg_b = (-adjusted_lr).item()
+    # Match the compiled path exactly: it multiplies M by ef_decay rounded to
+    # M's dtype (``ef_decay.to(M[0].dtype)``), so round before upcasting.
+    ef = ef_decay.to(M[0].dtype).item()
+
+    for x, m, u in zip(X, M, U):
+        if x.numel() == 0:
+            continue
+        if not (x.is_contiguous() and m.is_contiguous() and u.is_contiguous()):
+            raise ValueError(
+                "dion2_post_orthogonalize_masked_triton requires contiguous tensors"
+            )
+        orig_shape = x.shape
+        Mrows, N = orig_shape[-2], orig_shape[-1]
+        B = x.numel() // (Mrows * N)
+        x_flat = x.reshape(B, Mrows, N)
+        m_flat = m.reshape(B, Mrows, N)
+        u_flat = u.reshape(B, Mrows, N)
+
+        _dion2_post_ortho_masked_kernel[(B * Mrows,)](
+            x_flat, m_flat, u_flat,
+            a, neg_b, ef,
+            Mrows, N,
+            x_flat.stride(0), x_flat.stride(1), x_flat.stride(2),
+            m_flat.stride(0), m_flat.stride(1), m_flat.stride(2),
+            u_flat.stride(0), u_flat.stride(1), u_flat.stride(2),
+        )
+
+
+@triton.heuristics(
+    {
+        "BLOCK_N": lambda args: max(128, min(1024, triton.next_power_of_2(args["N"]))),
+        "num_warps": lambda args: 8 if args["N"] >= 2048 else 4,
+    }
+)
+@triton.jit
+def _nordion2_post_ortho_masked_stats_kernel(
+    U_ptr,
+    V_ptr,
+    M_ptr,
+    P_ptr,
+    ef,
+    lerp_w,
+    N,
+    u_stride_m,
+    u_stride_n,
+    v_stride_m,
+    m_stride_m,
+    m_stride_n,
+    p_stride_r,
+    BLOCK_N: tl.constexpr,
+):
+    """NorDion2 masked post, phase A: one program per row.
+
+    Computes the row's sum of squares of U (fp32), derives the selected flag,
+    updates the per-neuron variance V (lerp toward mean-square, selected rows
+    only -- unselected rows are stored back unchanged), applies error-feedback
+    decay to selected rows of momentum M, and emits per-row partials into P:
+      P[0, row] = sum_sq(u_row)                (-> Frobenius norm of U)
+      P[1, row] = sum_sq(u_row) / denom^2      (-> Frobenius norm of U/denom)
+      P[2, row] = denom = sqrt(v_new) + 1e-8   (fp32, pre-rounding V store)
+    The matrix-level Frobenius rescale needs all rows, hence the two-phase
+    split; the epilogue reduces P and phase B applies the weight update.
+    """
+    row = tl.program_id(0)
+    U_ptr += row * u_stride_m
+    M_ptr += row * m_stride_m
+    V_ptr += row * v_stride_m
+
+    ssq = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for off in range(0, N, BLOCK_N):
+        idx = off + tl.arange(0, BLOCK_N)
+        msk = idx < N
+        u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32)
+        ssq += u * u
+    sumsq = tl.sum(ssq)
+    sel = sumsq > 0.0
+
+    v_old = tl.load(V_ptr).to(tl.float32)
+    neuron = sumsq / N.to(tl.float32)
+    v_lerped = v_old + lerp_w * (neuron - v_old)
+    v_new = tl.where(sel, v_lerped, v_old)
+    denom = tl.sqrt(v_new) + 1e-8
+    tl.store(V_ptr, v_new)
+
+    tl.store(P_ptr + 0 * p_stride_r + row, sumsq)
+    tl.store(P_ptr + 1 * p_stride_r + row, sumsq / (denom * denom))
+    tl.store(P_ptr + 2 * p_stride_r + row, denom)
+
+    if sel:
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            msk = idx < N
+            m = tl.load(M_ptr + idx * m_stride_n, mask=msk, other=0.0).to(tl.float32)
+            tl.store(M_ptr + idx * m_stride_n, ef * m, mask=msk)
+
+
+@triton.heuristics(
+    {
+        "BLOCK_N": lambda args: max(128, min(1024, triton.next_power_of_2(args["N"]))),
+        "num_warps": lambda args: 8 if args["N"] >= 2048 else 4,
+    }
+)
+@triton.jit
+def _nordion2_post_ortho_masked_apply_kernel(
+    X_ptr,
+    U_ptr,
+    P_ptr,
+    scale_ptr,
+    mat_idx,
+    a,
+    neg_b,
+    N,
+    x_stride_m,
+    x_stride_n,
+    u_stride_m,
+    u_stride_n,
+    p_stride_r,
+    BLOCK_N: tl.constexpr,
+):
+    """NorDion2 masked post, phase B: one program per row.
+
+    Applies, in fp32 with a single rounding on store:
+      selected:   x = a*x + neg_b * scale * u / denom
+      unselected: x = a*x   (weight decay only; u is exactly zero anyway)
+    where scale = ||U||_F / max(||U/denom||_F, 1e-8) comes from the epilogue
+    reduction over phase A's partials.
+    """
+    row = tl.program_id(0)
+    X_ptr += row * x_stride_m
+    U_ptr += row * u_stride_m
+
+    sumsq = tl.load(P_ptr + 0 * p_stride_r + row)
+    sel = sumsq > 0.0
+
+    if sel:
+        denom = tl.load(P_ptr + 2 * p_stride_r + row)
+        scale = tl.load(scale_ptr + mat_idx).to(tl.float32)
+        coef = neg_b * scale / denom
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            msk = idx < N
+            x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32)
+            u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32)
+            tl.store(X_ptr + idx * x_stride_n, a * x + coef * u, mask=msk)
+    else:
+        for off in range(0, N, BLOCK_N):
+            idx = off + tl.arange(0, BLOCK_N)
+            msk = idx < N
+            x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32)
+            tl.store(X_ptr + idx * x_stride_n, a * x, mask=msk)
+
+
+def nordion2_post_orthogonalize_masked_triton(
+    X: List[Tensor],
+    M: List[Tensor],
+    V: List[Tensor],
+    U: List[Tensor],
+    base_lr: Tensor,
+    adjusted_lr: Tensor,
+    weight_decay: Tensor,
+    ef_decay: Tensor,
+    muon_beta2: Tensor,
+    select_dim: int,
+):
+    """Triton-fused version of ``nordion2_post_orthogonalize_masked``.
+
+    Two row-programmed kernels per matrix plus one batched epilogue reduction
+    per megabatch group (the NorMuon Frobenius rescale is a matrix-level
+    reduction, so a grid-wide dependency forces the two-phase split; the
+    epilogue is a deterministic torch reduction rather than fp32 atomics).
+    Falls back to the compiled masked post when the fast path does not apply.
+    All matrices in a megabatch shape group share one shape; mixed shapes fall
+    back too.
+    """
+    shapes = {tuple(x.shape) for x in X}
+    use_fallback = (
+        not TRITON_AVAILABLE
+        or select_dim != -2
+        or not X
+        or not X[0].is_cuda
+        or X[0].ndim != 2
+        or len(shapes) != 1
+        or any(is_traceable_wrapper_subclass(x) for x in X)
+    )
+    if use_fallback:
+        from .nordion2 import nordion2_post_orthogonalize_masked
+
+        nordion2_post_orthogonalize_masked(
+            X=X, M=M, V=V, U=U, base_lr=base_lr, adjusted_lr=adjusted_lr,
+            weight_decay=weight_decay, ef_decay=ef_decay,
+            muon_beta2=muon_beta2, select_dim=select_dim,
+        )
+        return
+
+    Mrows, N = X[0].shape
+    if Mrows == 0 or N == 0:
+        return
+
+    a = (1 - base_lr * weight_decay).item()
+    neg_b = (-adjusted_lr).item()
+    ef = ef_decay.to(M[0].dtype).item()
+    lerp_w = (1 - muon_beta2).item()
+
+    for x, m, v, u in zip(X, M, V, U):
+        if not (
+            x.is_contiguous() and m.is_contiguous() and u.is_contiguous()
+            and v.is_contiguous()
+        ):
+            raise ValueError(
+                "nordion2_post_orthogonalize_masked_triton requires contiguous tensors"
+            )
+
+    n_mat = len(X)
+    device = X[0].device
+    # Per-row partials: [n_mat, 3, rows] fp32 (sumsq, sumsq/denom^2, denom).
+    P = torch.empty((n_mat, 3, Mrows), dtype=torch.float32, device=device)
+
+    for i, (m, v, u) in enumerate(zip(M, V, U)):
+        _nordion2_post_ortho_masked_stats_kernel[(Mrows,)](
+            u, v, m, P[i],
+            ef, lerp_w,
+            N,
+            u.stride(0), u.stride(1),
+            v.stride(0),
+            m.stride(0), m.stride(1),
+            P.stride(1),
+        )
+
+    # Epilogue: matrix-level Frobenius rescale factors, one deterministic
+    # reduction for the whole group. Matches the compiled path's
+    # norm_U / clamp(norm_U_new, 1e-8).
+    sums = P[:, :2, :].sum(dim=-1)  # [n_mat, 2]
+    scale = sums[:, 0].sqrt() / sums[:, 1].sqrt().clamp_min(1e-8)
+    scale = scale.contiguous()
+
+    for i, (x, u) in enumerate(zip(X, U)):
+        _nordion2_post_ortho_masked_apply_kernel[(Mrows,)](
+            x, u, P[i], scale, i,
+            a, neg_b,
+            N,
+            x.stride(0), x.stride(1),
+            u.stride(0), u.stride(1),
+            P.stride(1),
+        )
diff --git a/dion/nordion2.py b/dion/nordion2.py
index 233c2e4..1e9709f 100644
--- a/dion/nordion2.py
+++ b/dion/nordion2.py
@@ -363,7 +363,10 @@ def nordion2_update_megabatch_async(
             adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten)
         else:
             raise ValueError(f"Unknown adjust_lr: {adjust_lr}")
-        nordion2_post_orthogonalize_masked(
+        _masked_post = nordion2_post_orthogonalize_masked
+        if triton_post_ortho:
+            from .dion2_triton import nordion2_post_orthogonalize_masked_triton as _masked_post
+        _masked_post(
             X=to_local(X),
             M=to_local(M),
             V=to_local(V),
@@ -407,7 +410,10 @@ def nordion2_update_megabatch_async(
         # rows of M, runs NorMuon per-neuron normalization on the selected rows
         # (updating the local variance buffer V in place), and applies the
         # masked weight update -- all keyed off the nonzero mask, no indices.
-        nordion2_post_orthogonalize_masked(
+        _masked_post = nordion2_post_orthogonalize_masked
+        if triton_post_ortho:
+            from .dion2_triton import nordion2_post_orthogonalize_masked_triton as _masked_post
+        _masked_post(
             X=to_local(X),
             M=to_local(M),
             V=to_local(V),
diff --git a/tests/test_dion2_post_ortho_masked_triton.py b/tests/test_dion2_post_ortho_masked_triton.py
new file mode 100644
index 0000000..b9e3121
--- /dev/null
+++ b/tests/test_dion2_post_ortho_masked_triton.py
@@ -0,0 +1,301 @@
+"""Tests for the Triton masked post-orthogonalize kernels.
+
+``dion2_post_orthogonalize_masked_triton`` and
+``nordion2_post_orthogonalize_masked_triton`` fuse the masked posts used by
+the "global" and "global_capped" selection scopes. Parity is checked against
+the compiled masked posts on identical inputs:
+
+- Unselected rows must match bitwise for X (both compute a*x with one
+  rounding) and be untouched for M/V.
+- Selected rows differ at FP-rounding level: the compiled Dion2 path rounds
+  ``neg_lr * u`` to U's dtype before the add (two/three roundings), the
+  Triton kernel computes ``a*x + neg_b*u`` in fp32 with a single rounding.
+
+The end-to-end test runs NorDion2 with selection_scope="global_capped" on
+2 GPUs, triton_post_ortho True vs False, and compares final params.
+"""
+
+import multiprocessing as mp
+import os
+import pytest
+import torch
+
+from dion.dion2 import dion2_post_orthogonalize_masked
+from dion.nordion2 import nordion2_post_orthogonalize_masked
+from dion.dion2_triton import (
+    TRITON_AVAILABLE,
+    dion2_post_orthogonalize_masked_triton,
+    nordion2_post_orthogonalize_masked_triton,
+)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+CUDA_AVAILABLE = torch.cuda.is_available()
+TRITON_AND_CUDA = CUDA_AVAILABLE and TRITON_AVAILABLE
+
+torch._dynamo.config.cache_size_limit = 64
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_masked_data(
+    n_mat, rows, cols, sel_rows_per_mat, seed=42, x_dtype=torch.float32,
+    with_v=False,
+):
+    """Lists of X, M, (V,) U with U exactly zero outside selected rows.
+
+    ``sel_rows_per_mat`` is a list (len n_mat) of per-matrix selected-row
+    counts; 0 means no rows selected on that matrix.
+    """
+    torch.manual_seed(seed)
+    X, M, V, U, sels = [], [], [], [], []
+    for i in range(n_mat):
+        X.append(torch.randn(rows, cols, device=DEVICE, dtype=x_dtype))
+        M.append(torch.randn(rows, cols, device=DEVICE, dtype=x_dtype))
+        if with_v:
+            V.append(
+                torch.rand(rows, 1, device=DEVICE, dtype=x_dtype) + 0.01
+            )
+        u = torch.zeros(rows, cols, device=DEVICE, dtype=torch.bfloat16)
+        k = sel_rows_per_mat[i]
+        sel = torch.randperm(rows, device=DEVICE)[:k]
+        if k:
+            u[sel] = torch.randn(k, cols, device=DEVICE, dtype=torch.bfloat16)
+        U.append(u)
+        mask = torch.zeros(rows, dtype=torch.bool, device=DEVICE)
+        mask[sel] = True
+        sels.append(mask)
+    scalars = dict(
+        base_lr=torch.tensor(0.01),
+        adjusted_lr=torch.tensor(0.02),
+        weight_decay=torch.tensor(0.1),
+        ef_decay=torch.tensor(0.95),
+    )
+    if with_v:
+        scalars["muon_beta2"] = torch.tensor(0.95)
+        return X, M, V, U, sels, scalars
+    return X, M, U, sels, scalars
+
+
+def _tols(dtype):
+    if dtype == torch.bfloat16:
+        return dict(atol=1e-2, rtol=1e-2)
+    return dict(atol=1e-6, rtol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# Function-level parity: Dion2 masked
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton")
+@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(
+    "rows,cols,sel_counts",
+    [
+        (64, 128, [16, 0, 64]),  # partial / none / all selected
+        (1, 32, [1, 0, 1]),      # single row
+        (128, 4099, [32, 7, 0]), # cols > BLOCK_N, non-pow2
+    ],
+)
+def test_dion2_masked_parity(x_dtype, rows, cols, sel_counts):
+    X, M, U, sels, sc = _make_masked_data(
+        len(sel_counts), rows, cols, sel_counts, x_dtype=x_dtype
+    )
+    X_ref = [x.clone() for x in X]
+    M_ref = [m.clone() for m in M]
+
+    dion2_post_orthogonalize_masked(
+        X=X_ref, M=M_ref, U=U, select_dim=-2,
+        base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"],
+        weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"],
+    )
+    dion2_post_orthogonalize_masked_triton(
+        X=X, M=M, U=U, select_dim=-2,
+        base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"],
+        weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"],
+    )
+
+    for x, x_ref, m, m_ref, sel in zip(X, X_ref, M, M_ref, sels):
+        # Unselected rows: both paths compute a*x with one rounding -> bitwise.
+        torch.testing.assert_close(x[~sel], x_ref[~sel], rtol=0, atol=0)
+        # M untouched on unselected rows (compiled multiplies by exactly 1).
+        torch.testing.assert_close(m[~sel], m_ref[~sel], rtol=0, atol=0)
+        torch.testing.assert_close(x, x_ref, **_tols(x_dtype))
+        torch.testing.assert_close(m, m_ref, **_tols(x_dtype))
+
+
+@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton")
+def test_dion2_masked_batched_3d():
+    """Leading batch dim: (B, rows, cols) with per-(b,row) selection."""
+    torch.manual_seed(0)
+    B, rows, cols = 3, 32, 64
+    x = torch.randn(B, rows, cols, device=DEVICE)
+    m = torch.randn(B, rows, cols, device=DEVICE)
+    u = torch.zeros(B, rows, cols, device=DEVICE, dtype=torch.bfloat16)
+    for b in range(B):
+        sel = torch.randperm(rows)[: 4 * (b + 1)]
+        u[b, sel] = torch.randn(len(sel), cols, device=DEVICE, dtype=torch.bfloat16)
+    args = dict(
+        base_lr=torch.tensor(0.01), adjusted_lr=torch.tensor(0.02),
+        weight_decay=torch.tensor(0.1), ef_decay=torch.tensor(0.95),
+        select_dim=-2,
+    )
+    x_ref, m_ref = x.clone(), m.clone()
+    dion2_post_orthogonalize_masked(X=[x_ref], M=[m_ref], U=[u], **args)
+    dion2_post_orthogonalize_masked_triton(X=[x], M=[m], U=[u], **args)
+    torch.testing.assert_close(x, x_ref, **_tols(torch.float32))
+    torch.testing.assert_close(m, m_ref, **_tols(torch.float32))
+
+
+@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton")
+def test_dion2_masked_empty_shard_skipped():
+    args = dict(
+        base_lr=torch.tensor(0.01), adjusted_lr=torch.tensor(0.02),
+        weight_decay=torch.tensor(0.1), ef_decay=torch.tensor(0.95),
+        select_dim=-2,
+    )
+    x = torch.empty(0, 64, device=DEVICE)
+    m = torch.empty(0, 64, device=DEVICE)
+    u = torch.empty(0, 64, device=DEVICE, dtype=torch.bfloat16)
+    # Must not raise.
+    dion2_post_orthogonalize_masked_triton(X=[x], M=[m], U=[u], **args)
+
+
+# ---------------------------------------------------------------------------
+# Function-level parity: NorDion2 masked
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton")
+@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(
+    "rows,cols,sel_counts",
+    [
+        (64, 128, [16, 0, 64]),
+        (128, 4099, [32, 7, 0]),
+        (16, 20480, [4, 16, 0]),  # production-wide rows
+    ],
+)
+def test_nordion2_masked_parity(x_dtype, rows, cols, sel_counts):
+    X, M, V, U, sels, sc = _make_masked_data(
+        len(sel_counts), rows, cols, sel_counts, x_dtype=x_dtype, with_v=True
+    )
+    X_ref = [x.clone() for x in X]
+    M_ref = [m.clone() for m in M]
+    V_ref = [v.clone() for v in V]
+
+    kw = dict(
+        base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"],
+        weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"],
+        muon_beta2=sc["muon_beta2"], select_dim=-2,
+    )
+    nordion2_post_orthogonalize_masked(X=X_ref, M=M_ref, V=V_ref, U=U, **kw)
+    nordion2_post_orthogonalize_masked_triton(X=X, M=M, V=V, U=U, **kw)
+
+    for x, x_ref, m, m_ref, v, v_ref, sel in zip(
+        X, X_ref, M, M_ref, V, V_ref, sels
+    ):
+        torch.testing.assert_close(x[~sel], x_ref[~sel], rtol=0, atol=0)
+        torch.testing.assert_close(m[~sel], m_ref[~sel], rtol=0, atol=0)
+        torch.testing.assert_close(v[~sel], v_ref[~sel], rtol=0, atol=0)
+        torch.testing.assert_close(x, x_ref, **_tols(x_dtype))
+        torch.testing.assert_close(m, m_ref, **_tols(x_dtype))
+        torch.testing.assert_close(v, v_ref, **_tols(x_dtype))
+
+
+@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton")
+def test_nordion2_masked_no_rows_selected_anywhere():
+    """All-zero U: pure weight decay, V and M untouched, no NaN from the
+    Frobenius rescale (norm 0 / clamp)."""
+    X, M, V, U, sels, sc = _make_masked_data(
+        2, 32, 64, [0, 0], with_v=True
+    )
+    kw = dict(
+        base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"],
+        weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"],
+        muon_beta2=sc["muon_beta2"], select_dim=-2,
+    )
+    X_ref = [x.clone() for x in X]
+    M_ref = [m.clone() for m in M]
+    V_ref = [v.clone() for v in V]
+    nordion2_post_orthogonalize_masked(X=X_ref, M=M_ref, V=V_ref, U=U, **kw)
+    nordion2_post_orthogonalize_masked_triton(X=X, M=M, V=V, U=U, **kw)
+    for x, x_ref in zip(X, X_ref):
+        assert torch.isfinite(x).all()
+        torch.testing.assert_close(x, x_ref, rtol=0, atol=0)
+    for m, m_ref in zip(M, M_ref):
+        torch.testing.assert_close(m, m_ref, rtol=0, atol=0)
+    for v, v_ref in zip(V, V_ref):
+        torch.testing.assert_close(v, v_ref, rtol=0, atol=0)
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: NorDion2 global_capped, triton vs compiled post (2 GPUs)
+# ---------------------------------------------------------------------------
+
+def _e2e_worker(rank, world_size, port, triton_post, out_path):
+    import torch.distributed as dist
+    from dion import NorDion2
+
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+    device = torch.device("cuda", rank)
+
+    torch.manual_seed(1234)  # same init on all ranks (replicated DDP-style)
+    full = [torch.randn(64, 96) for _ in range(3)]
+    params = [
+        torch.nn.Parameter(
+            torch.tensor_split(w, world_size, dim=0)[rank].clone().to(device)
+        )
+        for w in full
+    ]
+    opt = NorDion2(
+        [dict(params=params)],
+        distributed_mesh=dist.group.WORLD,
+        lr=0.01,
+        fraction=0.25,
+        selection_scope="global_capped",
+        capacity_factor=1.5,
+        use_triton=False,
+        triton_post_ortho=triton_post,
+    )
+    torch.manual_seed(99)
+    for _ in range(3):
+        grads = [torch.randn(64, 96) for _ in range(3)]
+        for p, g in zip(params, grads):
+            p.grad = torch.tensor_split(g, world_size, dim=0)[rank].clone().to(device)
+        opt.step()
+        opt.zero_grad()
+
+    if rank == 0:
+        torch.save([p.detach().cpu() for p in params], out_path)
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    not TRITON_AND_CUDA or torch.cuda.device_count() < 2,
+    reason="requires >=2 GPUs + triton",
+)
+def test_nordion2_capped_e2e_triton_vs_compiled(tmp_path):
+    ctx = mp.get_context("spawn")
+    outs = {}
+    for triton_post, port in ((False, 29711), (True, 29713)):
+        out = tmp_path / f"params_{triton_post}.pt"
+        procs = [
+            ctx.Process(
+                target=_e2e_worker, args=(r, 2, port, triton_post, str(out))
+            )
+            for r in range(2)
+        ]
+        for p in procs:
+            p.start()
+        for p in procs:
+            p.join(300)
+            assert p.exitcode == 0
+        outs[triton_post] = torch.load(out)
+
+    for p_ref, p_tri in zip(outs[False], outs[True]):
+        torch.testing.assert_close(p_tri, p_ref, atol=1e-5, rtol=1e-4)