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

[thomas-schweich]

Summary

SpecializedCLM._Block passed 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 asserts 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 train with --no-compile or sdpa_math (MATH backend) on ROCm — i.e. it couldn't use the flash + compile + bf16 path the main model already uses.

The main PAWN model's Attention.forward already forces q/k/v contiguous for exactly this reason (pawn/model.py). This applies the same 3-line fix to the CLM block.

Change

• pawn/specialized_clm.py: .contiguous() on v (after transpose) and on q/k (after RoPE), mirroring model.py.

Effect

specialized_clm now trains with flash + torch.compile + bf16 on ROCm (verified locally on a 7900 XT: clean single-bucket compile, no flash-backward assert).

Provenance

Found while benchmarking v1 cotrain vs the v2 supernet — the CLM was the one path that still fell back off flash on ROCm. Targeting main so it lands pre-v2.

Test plan

• Local: specialized_clm trains to completion with default (flash + compile + bf16) on ROCm, no stride assert.
• CI green.

[claude]

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()

[claude]

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.

[claude]

Review

Fix is correct and minimal. The three .contiguous() calls on v (post-transpose), q, and k (post-RoPE) directly mirror pawn/model.py:113-116, which already carries the same fix for the main backbone. torch.stack(...).reshape(x.shape) in _apply_rope and .transpose(1,2) on v both produce non-contiguous layouts; the ROCm inductor flash-attention backward kernel requires contiguous strides and asserts otherwise.

Findings

• Correctness: fix is sound. Non-contiguous layout after transpose + _apply_rope → torch.stack is exactly why model.py added .contiguous() in the same positions. Same root cause, same fix.
• Comment placement (nit): the comment block describes "q/k/v" but v is already assigned on the line before it. Suggestion posted inline to hoist the comment above all three lines.
• Pre-existing gap — --sdpa-math silently ignored: model.py wraps F.scaled_dot_product_attention in sdpa_kernel(SDPA_BACKEND), so the --sdpa-math escape hatch works there. specialized_clm._Block.forward calls SDPA bare and never reads SDPA_BACKEND, so that flag has no effect on this strategy. Not introduced by this PR; flagged inline as a follow-up.
• Tests: existing tests exercise forward/backward shape and gradient flow on CPU, which is enough to catch a complete regression. The flash-backward stride assert is ROCm + compile-specific and not feasible to catch in CI — the author's note of local verification on a 7900 XT is the right signal here.

One nit and one pre-existing gap to track; nothing blocking.