From 1819ef8dcf82f25445ceb66debd3c29e2bc8417c Mon Sep 17 00:00:00 2001 From: "zhijian.qiao" Date: Tue, 7 Jul 2026 18:53:01 +0800 Subject: [PATCH] fix(NeighborhoodAttention): correct axis bound in neighborhood clamp mask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit split() returns (t, h_row, w_col), but mask_mod clamped h_row with the width upper bound (cw-1-kw2) and w_col with the height upper bound (ch-1-kh2) — the axis bounds were swapped. Latent under square training (ch==cw makes the two bounds equal); on non-square grids (ch!=cw) the width-direction columns get clamped to the height upper bound, corrupting the neighborhood mask on the wider half of the image. Fix: clamp each axis with its own bound. Rename qx/qy -> qh/qw for clarity. Behavior unchanged for square grids. Co-authored-by: Cursor --- rayder/model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rayder/model.py b/rayder/model.py index 6ff5eb3..c8dcc7c 100644 --- a/rayder/model.py +++ b/rayder/model.py @@ -359,12 +359,12 @@ def split( return t, torch.div(rem, cw, rounding_mode="floor"), rem.remainder(cw) 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) ) return mask_mod