From c4ddde42109aef30122a3558af416040a8e78431 Mon Sep 17 00:00:00 2001 From: Thomas Schweich Date: Thu, 28 May 2026 21:55:58 -0700 Subject: [PATCH] fix(specialized_clm): contiguous q/k/v so flash backward works under torch.compile on ROCm SpecializedCLM._Block fed non-contiguous q/k/v (transpose + RoPE, no .contiguous()) into F.scaled_dot_product_attention. Under torch.compile on ROCm the inductor flash-attention *backward* kernel then hit a stride mismatch: AssertionError: ... torch.ops.aten._scaled_dot_product_flash_attention_backward expected size 4==4, stride 16384==64 at dim=1 ... so the dedicated CLM could only run with --no-compile or sdpa_math. The main PAWN model's Attention.forward already forces q/k/v contiguous for exactly this reason (pawn/model.py); apply the same to the CLM block. Lets specialized_clm train with flash + torch.compile + bf16 on ROCm. Found while benchmarking v1 cotrain vs the v2 supernet. --- pawn/specialized_clm.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pawn/specialized_clm.py b/pawn/specialized_clm.py index 3f895cf..f0737a6 100644 --- a/pawn/specialized_clm.py +++ b/pawn/specialized_clm.py @@ -60,9 +60,13 @@ def forward( h = self.attn_norm(x) q = self.wq(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.wk(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) - v = self.wv(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) - q = _apply_rope(q, rope_cos, rope_sin) - k = _apply_rope(k, rope_cos, rope_sin) + v = self.wv(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2).contiguous() + # Force contiguous q/k/v before SDPA so flash attention's backward + # gets the strides it expects under torch.compile on ROCm — same + # fix as the main model's Attention.forward (pawn/model.py). Without + # it the inductor flash-backward kernel hits a stride mismatch. + q = _apply_rope(q, rope_cos, rope_sin).contiguous() + k = _apply_rope(k, rope_cos, rope_sin).contiguous() attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True) attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, -1) x = x + self.wo(attn_out)