[CPU] Add GPT-OSS model optimization for CPU (#16775)

Co-authored-by: mingfeima <mingfei.ma@intel.com> Co-authored-by: jianan-gu <jianan.gu@intel.com>
2026-05-29 16:05:26 +08:00
parent 5601b7139d
commit 3ecf2c76ad
35 changed files with 2000 additions and 530 deletions
@@ -9,6 +9,91 @@ torch.manual_seed(1234)


 class TestDecodeAttention(CustomTestCase):
+    def _scaled_dot_product_attention(self, Q, K, V, S, scaling, sliding_window):
+        # sliding_window <= 0 means no sliding window
+        # Q: [n_tokens_q, n_heads, q_mult, d_head]
+        # K: [n_tokens_kv, n_heads, d_head]
+        # V: [n_tokens_kv, n_heads, d_head]
+        n_tokens_q, n_heads, q_mult, d_head = Q.shape
+        n_tokens_kv = K.shape[0]
+        assert K.shape == (n_tokens_kv, n_heads, d_head)
+        assert V.shape == (n_tokens_kv, n_heads, d_head)
+        K = K[:, :, None, :].expand(-1, -1, q_mult, -1)
+        V = V[:, :, None, :].expand(-1, -1, q_mult, -1)
+        S = S.reshape(n_heads, q_mult, 1, 1).expand(-1, -1, n_tokens_q, -1)
+        if n_tokens_q == n_tokens_kv:  # Prefill
+            mask = torch.triu(
+                Q.new_full((n_tokens_q, n_tokens_kv), -float("inf")), diagonal=1
+            )
+        else:  # Decode
+            mask = Q.new_zeros((n_tokens_q, n_tokens_kv))
+        if sliding_window is not None and sliding_window > 0:
+            mask += torch.tril(
+                mask.new_full((n_tokens_q, n_tokens_kv), -float("inf")),
+                diagonal=n_tokens_kv - n_tokens_q - sliding_window,
+            )
+        QK = torch.einsum("qhmd,khmd->hmqk", Q, K)
+        QK *= scaling
+        QK += mask[None, None, :, :]
+        QK = torch.cat([QK, S], dim=-1)
+        W = torch.softmax(QK, dim=-1)
+        W = W[..., :-1]
+        attn = torch.einsum("hmqk,khmd->qhmd", W, V)
+        return attn.reshape(n_tokens_q, -1)
+
+    def _run_sdpa_forward_decode_sink(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        num_kv_heads: int,
+        q_mult: int,
+        scaling=None,
+        sliding_window=None,
+        attention_sinks=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # TODO: this loop process a sequence per iter, this is inefficient.
+            # Need optimize the performance later.
+            seq_len_q = 1
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + seq_len_q
+            end_kv = start_kv + seq_len_kv
+            per_req_query = query[:, start_q:end_q, :]
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+
+            per_req_query = per_req_query.permute(1, 0, 2).reshape(
+                seq_len_q, num_kv_heads, q_mult, per_req_query.shape[-1]
+            )
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_key = per_req_key.permute(1, 0, 2)
+            per_req_value = per_req_value.permute(1, 0, 2)
+
+            per_req_out = self._scaled_dot_product_attention(
+                per_req_query,
+                per_req_key,
+                per_req_value,
+                attention_sinks,
+                scaling=scaling,
+                sliding_window=sliding_window,
+            ).reshape(seq_len_q, -1, per_req_value.shape[-1])
+            output[start_q:end_q, :, :] = per_req_out
+            start_q, start_kv = end_q, end_kv
+        return output
+
    def _run_sdpa_forward_decode(
        self,
        query: torch.Tensor,
@@ -68,11 +153,11 @@ class TestDecodeAttention(CustomTestCase):
        return output

    def _test_grouped_decode_attention_once(
-        self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
+        self, B, H_Q, H_KV, D, D_V, sliding_window, sink, is_cross_attn, dtype, device
    ):
        # This represents the number of tokens already in the sequence
        seq_len = 1024
-        encoder_len = 10
+        encoder_len = 0 if sink else 10
        total_tokens = B * (seq_len + encoder_len)
        sm_scale = 1.0 / (D**0.5)
        logit_cap = 0.0
@@ -81,6 +166,7 @@ class TestDecodeAttention(CustomTestCase):

        # q represents the new token being generated, one per batch
        q = torch.randn(B, H_Q, D, dtype=dtype, device=device)
+        sinks = torch.rand(H_Q, dtype=dtype, device=device) * 10

        # k_buffer and v_buffer represent all previous tokens
        k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
@@ -134,22 +220,41 @@ class TestDecodeAttention(CustomTestCase):
            sm_scale,
            logit_cap,
            is_cross_attn,
+            sliding_window if sliding_window is not None else 0,
            encoder_lens,
+            sinks if sink else None,
        )

-        self._run_sdpa_forward_decode(
-            q,
-            o_grouped,
-            k_buffer,
-            v_buffer,
-            req_to_token,
-            b_req_idx,
-            b_seq_len,
-            scaling=sm_scale,
-            enable_gqa=enable_gqa,
-            encoder_lens=encoder_lens,
-            is_cross_attn=is_cross_attn,
-        )
+        if sink:
+            self._run_sdpa_forward_decode_sink(
+                q,
+                o_grouped,
+                k_buffer,
+                v_buffer,
+                req_to_token,
+                b_req_idx,
+                b_seq_len,
+                num_kv_heads=H_KV,
+                q_mult=H_Q // H_KV if enable_gqa else 1,
+                scaling=sm_scale,
+                sliding_window=sliding_window if sliding_window is not None else None,
+                attention_sinks=sinks,
+                enable_gqa=enable_gqa,
+            )
+        else:
+            self._run_sdpa_forward_decode(
+                q,
+                o_grouped,
+                k_buffer,
+                v_buffer,
+                req_to_token,
+                b_req_idx,
+                b_seq_len,
+                scaling=sm_scale,
+                enable_gqa=enable_gqa,
+                encoder_lens=encoder_lens,
+                is_cross_attn=is_cross_attn,
+            )
        cos_sim = torch.nn.functional.cosine_similarity(
            o.flatten(), o_grouped.flatten(), dim=0
        )
@@ -173,11 +278,26 @@ class TestDecodeAttention(CustomTestCase):

        for B, H_Q, H_KV, D, D_V in configs:
            for dtype in [torch.bfloat16, torch.float16]:
+                for sink in [True, False]:
+                    if D != D_V and sink:
+                        continue
+                    for sliding_window in [None, 10]:
+                        if sliding_window is not None and not sink:
+                            continue
+                        self._test_grouped_decode_attention_once(
+                            B,
+                            H_Q,
+                            H_KV,
+                            D,
+                            D_V,
+                            sliding_window,
+                            sink,
+                            False,
+                            dtype=dtype,
+                            device=device,
+                        )
                self._test_grouped_decode_attention_once(
-                    B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
-                )
-                self._test_grouped_decode_attention_once(
-                    B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
+                    B, H_Q, H_KV, D, D_V, None, False, True, dtype=dtype, device=device
                )

    def test_grouped_decode_attention(self):
@@ -9,6 +9,37 @@ torch.manual_seed(1234)


 class TestExtendAttention(CustomTestCase):
+    def _scaled_dot_product_attention(self, Q, K, V, S, scaling, sliding_window):
+        # sliding_window <= 0 means no sliding window
+        # Q: [n_tokens_q, n_heads, q_mult, d_head]
+        # K: [n_tokens_kv, n_heads, d_head]
+        # V: [n_tokens_kv, n_heads, d_head]
+        n_tokens_q, n_heads, q_mult, d_head = Q.shape
+        n_tokens_kv = K.shape[0]
+        assert K.shape == (n_tokens_kv, n_heads, d_head)
+        assert V.shape == (n_tokens_kv, n_heads, d_head)
+        K = K[:, :, None, :].expand(-1, -1, q_mult, -1)
+        V = V[:, :, None, :].expand(-1, -1, q_mult, -1)
+        S = S.reshape(n_heads, q_mult, 1, 1).expand(-1, -1, n_tokens_q, -1)
+        if n_tokens_q == n_tokens_kv:  # Prefill
+            mask = torch.triu(
+                Q.new_full((n_tokens_q, n_tokens_kv), -float("inf")), diagonal=1
+            )
+        else:  # Decode
+            mask = Q.new_zeros((n_tokens_q, n_tokens_kv))
+        if sliding_window is not None and sliding_window > 0:
+            mask += torch.tril(
+                mask.new_full((n_tokens_q, n_tokens_kv), -float("inf")),
+                diagonal=n_tokens_kv - n_tokens_q - sliding_window,
+            )
+        QK = torch.einsum("qhmd,khmd->hmqk", Q, K)
+        QK *= scaling
+        QK += mask[None, None, :, :]
+        QK = torch.cat([QK, S], dim=-1)
+        W = torch.softmax(QK, dim=-1)
+        W = W[..., :-1]
+        attn = torch.einsum("hmqk,khmd->qhmd", W, V)
+        return attn.reshape(n_tokens_q, -1)

    def _run_sdpa_forward_extend(
        self,
@@ -83,6 +114,70 @@ class TestExtendAttention(CustomTestCase):
            start_q, start_kv = end_q, end_kv
        return output

+    def _run_sdpa_forward_extend_sink(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_prefix_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        num_kv_heads: int,
+        q_mult: int,
+        scaling=None,
+        sliding_window=None,
+        attention_sinks=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+        assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
+        assert seq_lens.shape[0] == extend_seq_lens.shape[0]
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # TODO: this loop process a sequence per iter, this is inefficient.
+            # Need optimize the performance later.
+            extend_seq_len_q = extend_seq_lens[seq_idx]
+            prefill_seq_len_q = extend_prefix_lens[seq_idx]
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + extend_seq_len_q
+            end_kv = start_kv + seq_len_kv
+            per_req_query = query[:, start_q:end_q, :]
+            per_req_query_redudant = torch.empty(
+                (per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]),
+                dtype=per_req_query.dtype,
+                device=per_req_query.device,
+            )
+            per_req_query_redudant[:, prefill_seq_len_q:, :] = per_req_query
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            per_req_query_redudant = per_req_query_redudant.permute(1, 0, 2).reshape(
+                seq_len_kv, num_kv_heads, q_mult, per_req_query_redudant.shape[-1]
+            )
+            per_req_key = per_req_key.permute(1, 0, 2)
+            per_req_value = per_req_value.permute(1, 0, 2)
+            per_req_out_redudant = self._scaled_dot_product_attention(
+                per_req_query_redudant,
+                per_req_key,
+                per_req_value,
+                attention_sinks,
+                scaling=scaling,
+                sliding_window=sliding_window,
+            ).reshape(seq_len_kv, -1, per_req_value.shape[-1])
+            output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :]
+            start_q, start_kv = end_q, end_kv
+        return output
+
    def _test_extend_attention_once(
        self,
        B,
@@ -91,6 +186,8 @@ class TestExtendAttention(CustomTestCase):
        H_KV,
        D,
        DV,
+        sliding_window=None,
+        has_sink=False,
        mla=False,
        is_cross_attn=False,
        *,
@@ -105,9 +202,13 @@ class TestExtendAttention(CustomTestCase):
            b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)

        encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
+        scale = 20
        if mla:
            b_seq_len_prefix.zero_()
            encoder_lens.zero_()
+        if has_sink:
+            encoder_lens.zero_()
+            scale = 1

        if b_seq_len_extend is None:
            b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
@@ -140,6 +241,7 @@ class TestExtendAttention(CustomTestCase):
        k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype)
        v_extend = torch.empty((extend_token_num, H_KV, DV), dtype=dtype)
        q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
+        sinks = torch.rand(H_Q, dtype=dtype)

        for i in range(B):
            extend_start_in_buffer = (
@@ -155,7 +257,7 @@ class TestExtendAttention(CustomTestCase):
                extend_start_in_buffer:extend_end_in_buffer
            ]
            q_extend[extend_start:extend_end] = (
-                torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * 20
+                torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * scale
            )

        # q_extend, k_extend, v_extend, k_buffer and v_buffer supports non-contiguous tensors
@@ -179,22 +281,40 @@ class TestExtendAttention(CustomTestCase):

        enable_gqa = H_Q != H_KV
        o_ref = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
-        self._run_sdpa_forward_extend(
-            q_extend,
-            o_ref,
-            k_buffer,
-            v_buffer,
-            req_to_tokens,
-            b_req_idx,
-            b_seq_len,
-            b_seq_len_prefix,
-            b_seq_len_extend,
-            scaling=sm_scale,
-            enable_gqa=enable_gqa,
-            causal=not is_cross_attn,
-            is_cross_attn=is_cross_attn,
-            encoder_lens=encoder_lens,
-        )
+        if has_sink:
+            self._run_sdpa_forward_extend_sink(
+                q_extend,
+                o_ref,
+                k_buffer,
+                v_buffer,
+                req_to_tokens,
+                b_req_idx,
+                b_seq_len,
+                b_seq_len_prefix,
+                b_seq_len_extend,
+                H_KV,
+                H_Q // H_KV if enable_gqa else 1,
+                scaling=sm_scale,
+                sliding_window=sliding_window,
+                attention_sinks=sinks,
+            )
+        else:
+            self._run_sdpa_forward_extend(
+                q_extend,
+                o_ref,
+                k_buffer,
+                v_buffer,
+                req_to_tokens,
+                b_req_idx,
+                b_seq_len,
+                b_seq_len_prefix,
+                b_seq_len_extend,
+                scaling=sm_scale,
+                enable_gqa=enable_gqa,
+                causal=not is_cross_attn,
+                is_cross_attn=is_cross_attn,
+                encoder_lens=encoder_lens,
+            )

        o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
        torch.ops.sgl_kernel.extend_attention_cpu(
@@ -213,7 +333,9 @@ class TestExtendAttention(CustomTestCase):
            sm_scale,
            logit_cap,
            is_cross_attn,
+            sliding_window if sliding_window is not None else 0,
            encoder_lens,
+            sinks if has_sink else None,
        )

        torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
@@ -224,16 +346,29 @@ class TestExtendAttention(CustomTestCase):
                if is_mla and is_cross_attn:
                    continue
                self._test_extend_attention_once(
-                    1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
+                    1, 123, 1, 1, 128, 96, None, False, is_mla, is_cross_attn
                )
                self._test_extend_attention_once(
-                    1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
+                    1, 123, 16, 1, 128, 96, None, False, is_mla, is_cross_attn
                )
                self._test_extend_attention_once(
-                    4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
+                    4, 1230, 16, 4, 128, 96, None, False, is_mla, is_cross_attn
                )
                self._test_extend_attention_once(
-                    1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
+                    1, 9000, 16, 1, 32, 32, None, False, is_mla, is_cross_attn
+                )
+        for has_sink in [True, False]:
+            for sliding_window in [None, 10, 128]:
+                if not has_sink and sliding_window is not None:
+                    continue
+                self._test_extend_attention_once(
+                    1, 123, 16, 4, 64, 64, sliding_window, has_sink, False, False
+                )
+                self._test_extend_attention_once(
+                    1, 20, 16, 1, 64, 64, sliding_window, has_sink, False, False
+                )
+                self._test_extend_attention_once(
+                    1, 20, 1, 1, 64, 64, sliding_window, has_sink, False, False
                )

    def test_extend_attention_large_seq_causal_mask(self):
@@ -116,6 +116,8 @@ class TestMLA(CustomTestCase):
            sm_scale,
            logit_cap,
            False,
+            0,
+            None,
            None,
        )