nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

gptoss: use fa swa (#17522)

Browse Source
This commit is contained in:
wozeparrot
2026-08-13 10:39:22 -07:00
committed by GitHub
parent 0252cb8fa7
commit 849074f0db
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/models/gpt_oss.py
+4 -4
View File
@@ -182,12 +182,12 @@ class GPTOSS:
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
if sliding:
attn = self._sliding_attention(xq, xk, xv, sinks)
elif getenv("HK_FLASH_ATTENTION"):
if getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
elif sliding:
attn = self._sliding_attention(xq, xk, xv, sinks)
else:
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)