Skip to content

fix(specialized_clm): contiguous q/k/v for flash backward under torch.compile on ROCm #116

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
thomas-schweich wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(specialized_clm): contiguous q/k/v for flash backward under torch.compile on ROCm #116

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions pawn/specialized_clm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Comment on lines +63 to +69

@claude claude Bot May 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: the comment says "Force contiguous q/k/v" but v is already on the line before it, so the explanation arrives after the first call it covers. Moving the comment block above all three makes the intent clearer and matches how you'd normally read it:

Suggested change
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()
# 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.
v = self.wv(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2).contiguous()
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)

@claude claude Bot May 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pre-existing gap (not introduced by this PR, but visible here): model.py's Attention.forward wraps this call with sdpa_kernel(SDPA_BACKEND) when SDPA_BACKEND is not None, so --sdpa-math takes effect there. specialized_clm never checks SDPA_BACKEND and always calls SDPA bare, so --sdpa-math is silently ignored for the specialized_clm strategy. Not a blocker for this fix, but worth a follow-up if --sdpa-math ever needs to work here as a debugging escape hatch.

attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, -1)
x = x + self.wo(attn_out)
Expand Down
Loading