From 849074f0dbdc9e97095aef09335f85cbea848089 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Fri, 14 Aug 2026 01:39:22 +0800 Subject: [PATCH] gptoss: use fa swa (#17522) --- examples/mlperf/models/gpt_oss.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/mlperf/models/gpt_oss.py b/examples/mlperf/models/gpt_oss.py index 2fb03ad603..d58d615ef0 100644 --- a/examples/mlperf/models/gpt_oss.py +++ b/examples/mlperf/models/gpt_oss.py @@ -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)