fix(NeighborhoodAttention): correct axis bound in neighborhood clamp mask

[qiaozhijian]

Summary

NeighborhoodAttention._make_mask has a latent bug in the mask_mod neighborhood clamp: the per-axis upper bounds for the height and width axes are swapped.

split() returns (t, h_row, w_col) (row-major: rem // cw is the row/height index, rem % cw is the column/width index). But mask_mod clamped:

• the h_row axis (qx) with the width upper bound cw - 1 - kw2, and
• the w_col axis (qy) with the height upper bound ch - 1 - kh2.

Each axis's neighborhood boundary clamp must use its own axis size. The bounds were bound to the wrong axes.

This is latent because the model is trained/evaluated on square grids (ch == cw), where the two upper bounds are equal and the swap is elementwise-identical — so the bug never surfaces. On non-square grids (ch != cw, e.g. 16:9 → 18×32 latent), the width-direction columns get clamped to the height upper bound, corrupting the neighborhood mask on the wider half of the image.

Reproduction

Non-square grid ct,ch,cw = 1,18,32 (i.e. 288×512 / 16), kernel (1,7,7). Right-edge query q_idx = 31 (h_row=0, w_col=31), list the w_cols it attends to:

import torch
from rayder.model import NeighborhoodAttention

ct, ch, cw, kt, kh, kw = 1, 18, 32, 1, 7, 7
mask_mod = NeighborhoodAttention._make_mask(ct, ch, cw, kt, kh, kw)

def split(idx, hw=ch*cw, cw_=cw):
    rem = idx % hw
    return int(idx // hw), int(rem // cw_), int(rem % cw_)

q_idx = 31
w_cols = sorted(set(
    split(kv)[2] for kv in range(ch*cw)
    if bool(mask_mod(0, 0, torch.tensor(q_idx), torch.tensor(kv)).item())
))
print("attended w_cols:", w_cols)

• Before (buggy): attended w_cols: [11, 12, 13, 14, 15, 16, 17] — a right-edge query attends to the middle columns (clamped to the height bound 14). Wrong.
• After (fixed): attended w_cols: [25, 26, 27, 28, 29, 30, 31] — right-edge query attends the right-edge columns (clamped to cw-1-kw2 = 28). Correct.

Square-grid equivalence

For square grids the fix is provably a no-op. Verified by comparing the old vs new mask_mod output for all (q_idx, kv_idx) pairs on ch=cw ∈ {16, 18, 32}:

[square ch=cw=16] all 256x256 pairs identical old vs new -> PASS
[square ch=cw=18] all 324x324 pairs identical old vs new -> PASS
[square ch=cw=32] all 1024x1024 pairs identical old vs new -> PASS

So existing square-resolution checkpoints/behavior are unaffected.

Fix

Clamp each axis with its own bound; rename qx/qy → qh/qw (and kx/ky → kh_/kw_) for clarity.

         def mask_mod(b, h, q_idx, kv_idx) -> Bool[torch.Tensor, ""]:
-            qt, qx, qy = split(q_idx)
-            kt_, kx, ky = split(kv_idx)
+            qt, qh, qw = split(q_idx)
+            kt_, kh_, kw_ = split(kv_idx)
             return (
                 ((qt.clamp(kt2, ct - 1 - kt2) - kt_).abs() <= kt2)
-                & ((qx.clamp(kw2, cw - 1 - kw2) - kx).abs() <= kw2)
-                & ((qy.clamp(kh2, ch - 1 - kh2) - ky).abs() <= kh2)
+                & ((qh.clamp(kh2, ch - 1 - kh2) - kh_).abs() <= kh2)
+                & ((qw.clamp(kw2, cw - 1 - kw2) - kw_).abs() <= kw2)
             )

Only this one function is touched.

Made with Cursor