[NPU]grok2 model support (#17119)

Co-authored-by: cy <chenyang08056032@163.com>

.github/workflows/nightly-test-npu.yml

@@ -55,11 +55,12 @@ jobs:
           # copy required file from our daily cache
           cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
           # copy download through proxy
-          curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+          curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
 
       - name: Print Log Information
         run: |
           bash scripts/ci/npu/npu_log_print.sh
+
       - name: Run test
         timeout-minutes: 240
         env:
@@ -102,7 +103,7 @@ jobs:
           # copy required file from our daily cache
           cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
           # copy download through proxy
-          curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+          curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
 
       - name: Print Log Information
         run: |
@@ -148,7 +149,7 @@ jobs:
           # copy required file from our daily cache
           cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
           # copy download through proxy
-          curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+          curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
 
       - name: Print Log Information
         run: |

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -831,6 +831,7 @@ class AscendAttnBackend(AttentionBackend):
                     layer.qk_head_dim <= 128
                     and causal
                     and forward_batch.encoder_lens is None
+                    and layer.logit_cap == 0
                     and not getattr(self, "use_native_sdpa", False)
                 ):
                     if not self.use_alibi:
@@ -896,6 +897,8 @@ class AscendAttnBackend(AttentionBackend):
                         scaling=layer.scaling,
                         enable_gqa=use_gqa,
                         causal=causal,
+                        logit_cap=layer.logit_cap,
+                        logit_capping_method=layer.logit_capping_method,
                     )
                     attn_output = attn_output.view(
                         -1, layer.tp_q_head_num * layer.v_head_dim
@@ -1022,7 +1025,7 @@ class AscendAttnBackend(AttentionBackend):
                     layer.layer_id
                 )
                 kv_cache = torch.cat([k_cache, v_cache], dim=-1)
-                attn_output = self.native_attn._run_sdpa_forward_extend(
+                attn_output = self.native_attn.run_sdpa_forward_extend(
                     q,
                     attn_output,
                     kv_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
@@ -1525,7 +1528,7 @@ class AscendAttnBackend(AttentionBackend):
                 )
             # there are some accuracy issues in cross attention scene to use torch_npu._npu_flash_attention_qlens
             # forward_batch.encoder_lens is not None in cross attention scend, we add native attn to solve accuracy issues
-            elif forward_batch.encoder_lens is None:
+            elif forward_batch.encoder_lens is None and layer.logit_cap == 0:
                 query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
                 num_tokens = query.shape[0]
                 if not self.use_alibi:
@@ -1585,6 +1588,8 @@ class AscendAttnBackend(AttentionBackend):
                     scaling=layer.scaling,
                     enable_gqa=use_gqa,
                     causal=False,
+                    logit_cap=layer.logit_cap,
+                    logit_capping_method=layer.logit_capping_method,
                 )
             return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
         else:

python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import math
+
 import torch
 from torch.nn.functional import scaled_dot_product_attention
 
@@ -8,6 +10,49 @@ class AscendTorchNativeAttnBackend:
     def __init__(self):
         pass
 
+    def scaled_dot_product_attention_with_softcapping(
+        self,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        is_causal=False,
+        scale=None,
+        enable_gqa=False,
+        logit_cap=0.0,
+        logit_capping_method="tanh",
+    ) -> torch.Tensor:
+        L, S = query.size(-2), key.size(-2)
+        scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+        attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+        if is_causal:
+            assert attn_mask is None
+            temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(
+                diagonal=0
+            )
+            attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+            attn_bias.to(query.dtype)
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+            else:
+                attn_bias = attn_mask + attn_bias
+
+        if enable_gqa:
+            key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
+            value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
+
+        attn_weight = query @ key.transpose(-2, -1) * scale_factor
+
+        if logit_cap > 0:
+            if logit_capping_method == "tanh":
+                attn_weight = logit_cap * torch.tanh(attn_weight / logit_cap)
+
+        attn_weight += attn_bias
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+        return attn_weight @ value
+
     def run_sdpa_forward_extend(
         self,
         query: torch.Tensor,
@@ -24,6 +69,8 @@ class AscendTorchNativeAttnBackend:
         scaling=None,
         enable_gqa=False,
         causal=False,
+        logit_cap: float = 0.0,
+        logit_capping_method: str = "tanh",
     ):
         """Run the extend forward by using torch native sdpa op.
 
@@ -94,18 +141,34 @@ class AscendTorchNativeAttnBackend:
                 per_req_key = per_req_key.to(per_req_query.dtype)
                 per_req_value = per_req_value.to(per_req_query.dtype)
 
-            per_req_out_redudant = (
+            if logit_cap > 0:
-                scaled_dot_product_attention(
+                per_req_out_redudant = (
-                    per_req_query_redudant.unsqueeze(0),
+                    self.scaled_dot_product_attention_with_softcapping(
-                    per_req_key.unsqueeze(0),
+                        per_req_query_redudant.unsqueeze(0),
-                    per_req_value.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
-                    enable_gqa=enable_gqa,
+                        per_req_value.unsqueeze(0),
-                    scale=scaling,
+                        enable_gqa=enable_gqa,
-                    is_causal=causal,
+                        scale=scaling,
+                        is_causal=causal,
+                        logit_cap=logit_cap,
+                        logit_capping_method=logit_capping_method,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            else:
+                per_req_out_redudant = (
+                    scaled_dot_product_attention(
+                        per_req_query_redudant.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
                 )
-                .squeeze(0)
-                .movedim(query.dim() - 2, 0)
-            )
             output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :]
             start_q, start_kv = end_q, end_kv
         return output
@@ -124,6 +187,8 @@ class AscendTorchNativeAttnBackend:
         scaling=None,
         enable_gqa=False,
         causal=False,
+        logit_cap: float = 0.0,
+        logit_capping_method: str = "tanh",
     ):
         """Run the decode forward by using torch native sdpa op.
 
@@ -180,18 +245,34 @@ class AscendTorchNativeAttnBackend:
                 per_req_key = per_req_key.to(per_req_query.dtype)
                 per_req_value = per_req_value.to(per_req_query.dtype)
 
-            per_req_out = (
+            if logit_cap > 0:
-                scaled_dot_product_attention(
+                per_req_out = (
-                    per_req_query.unsqueeze(0),
+                    self.scaled_dot_product_attention_with_softcapping(
-                    per_req_key.unsqueeze(0),
+                        per_req_query.unsqueeze(0),
-                    per_req_value.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
-                    enable_gqa=enable_gqa,
+                        per_req_value.unsqueeze(0),
-                    scale=scaling,
+                        enable_gqa=enable_gqa,
-                    is_causal=causal,
+                        scale=scaling,
+                        is_causal=causal,
+                        logit_cap=logit_cap,
+                        logit_capping_method=logit_capping_method,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            else:
+                per_req_out = (
+                    scaled_dot_product_attention(
+                        per_req_query.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
                 )
-                .squeeze(0)
-                .movedim(query.dim() - 2, 0)
-            )
             output[start_q:end_q, :, :] = per_req_out
             start_q, start_kv = end_q, end_kv
 

python/sglang/srt/models/grok.py

@@ -18,6 +18,7 @@ import math
 from typing import Iterable, Optional, Tuple
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
@@ -59,7 +60,9 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.utils import add_prefix
+from sglang.srt.utils import add_prefix, is_npu
+
+_is_npu = is_npu()
 
 logger = logging.getLogger(__name__)
 
@@ -143,7 +146,7 @@ class Grok1MoE(nn.Module):
             top_k=top_k,
             renormalize=False,
             layer_id=layer_id,
-            custom_routing_function=custom_routing_function,
+            custom_routing_function=None if _is_npu else custom_routing_function,
         )
 
         self.experts = FusedMoE(
@@ -162,8 +165,21 @@ class Grok1MoE(nn.Module):
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        topk_output = self.topk(hidden_states, self.gate.weight)
+        if not _is_npu:
-        return self.experts(hidden_states, topk_output)
+            topk_output = self.topk(hidden_states, self.gate.weight)
+            return self.experts(hidden_states, topk_output)
+        else:
+            orig_shape = hidden_states.shape
+            hidden_states = hidden_states.view(-1, self.hidden_size)
+
+            router_logits, _ = self.gate(hidden_states)
+            router_logits = self.router_logit_softcapping * F.tanh(
+                router_logits / self.router_logit_softcapping
+            )
+            topk_output = self.topk(hidden_states, router_logits)
+
+            final_hidden_states = self.experts(hidden_states, topk_output)
+            return final_hidden_states.view(orig_shape)
 
 
 def _yarn_linear_ramp_mask(
@@ -228,6 +244,8 @@ class ScalingRotaryEmbedding(RotaryEmbedding):
         self.attn_factor = attn_factor
         self.beta_fast = beta_fast
         self.beta_slow = beta_slow
+        if _is_npu:
+            dtype = torch.float32
         # Get n-d magnitude scaling corrected for interpolation
         self.mscale = float(_yarn_get_mscale(self.scaling_factor) * attn_factor)
         super().__init__(
@@ -396,6 +414,7 @@ class Grok1Attention(nn.Module):
                 max_position=max_position,
                 base=int(self.rope_theta),
                 is_neox_style=True,
+                dtype=torch.float32 if _is_npu else None,
             )
             pos_encoding_mode = "NONE"
 
@@ -425,7 +444,12 @@ class Grok1Attention(nn.Module):
         qkv, _ = self.qkv_proj(hidden_states)
 
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
+        if not _is_npu:
+            q, k = self.rotary_emb(positions, q, k)
+        else:
+            odtype = q.dtype
+            q, k = self.rotary_emb(positions, q.to(torch.float32), k.to(torch.float32))
+            q, k = q.to(odtype), k.to(odtype)
 
         attn_output = self.attn(q, k, v, forward_batch)
 

test/registered/ascend/llm_models/test_ascend_grok_2.py

@@ -0,0 +1,29 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ci.ci_register import register_npu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_npu_ci(est_time=400, suite="nightly-16-npu-a3", nightly=True)
+
+
+class TestGrok2(GSM8KAscendMixin, CustomTestCase):
+    model = "/root/.cache/modelscope/hub/models/huihui-ai/grok-2"
+    accuracy = 0.91
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.8",
+        "--attention-backend",
+        "ascend",
+        "--disable-radix-cache",
+        "--disable-cuda-graph",
+        "--tokenizer-path",
+        "/root/.cache/modelscope/hub/models/huihui-ai/grok-2/tokenizer.tok.json",
+        "--tp-size",
+        "16",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()