Porting DeeSeek v2/r1 PRs (vllm-project#1756)

## Essential Elements of an Effective PR Description Checklist
- [ ] The purpose of the PR, such as "Fix some issue (link existing
issues this PR will resolve)".
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before
and after, or e2e results


## Porting List
1. HabanaAI#1402
2. HabanaAI#1504
3. HabanaAI#1404

<!--- pyml disable-next-line no-emphasis-as-heading -->

Author: Wei-Lin-Intel

vllm/attention/backends/mla/common.py

@@ -1171,9 +1171,9 @@ def get_scales(layer: LinearBase) -> torch.Tensor:
             [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
 
         # Convert from (L, N, V) to (N, L, V)
-        self.W_UV = W_UV.transpose(0, 1)
+        self.W_UV = W_UV.transpose(0, 1).contiguous()
         # Convert from (L, N, P) to (N, P, L)
-        self.W_UK_T = W_UK.permute(1, 2, 0)
+        self.W_UK_T = W_UK.permute(1, 2, 0).contiguous()
 
     def _compute_prefill_context(
         self,

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -944,20 +944,32 @@ def grouped_topk(
         # scores for expert selection but original scores for routing weights
         original_scores = scores
         scores = scores + e_score_correction_bias.unsqueeze(0)
-        group_scores = (scores.view(num_token, num_expert_group,
-                                    -1).topk(2, dim=-1)[0].sum(dim=-1))
+
+        scores_tmp = scores.clone().reshape(num_token, num_expert_group, -1)
+        top1_val, top1_idx = torch.max(scores_tmp, dim=-1)
+        scores_tmp.scatter_(-1, top1_idx.unsqueeze(-1), torch.finfo(scores.dtype).min)
+        group_scores, top2_idx = torch.max(scores_tmp, dim=-1)
+        group_scores.add_(top1_val)
     else:
         group_scores = scores.view(num_token, num_expert_group,
                                    -1).max(dim=-1).values  # [n, n_group]
-    group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
-                           sorted=False)[1]  # [n, top_k_group]
-    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
-    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
-    score_mask = group_mask.unsqueeze(-1).expand(
-        num_token, num_expert_group,
-        scores.shape[-1] // num_expert_group).reshape(num_token, -1)  # [n, e]
-    tmp_scores = scores.masked_fill(~score_mask.bool(),
-                                    float("-inf"))  # [n, e]
+
+    if num_token > 1024:
+        group_mask = torch.zeros_like(group_scores)
+        for i in range(topk_group):
+            _, group_idx = torch.max(group_scores, dim=-1)
+            group_mask.scatter_(1, group_idx.unsqueeze(-1), 1)
+            if i < topk_group - 1:
+                group_scores.scatter_(1, group_idx.unsqueeze(-1), torch.finfo(scores.dtype).min)
+    else:
+        group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
+                               sorted=False)[1]  # [n, top_k_group]
+        group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+        group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+
+    tmp_scores = scores.reshape(num_token, num_expert_group, -1) + \
+                     ((1 - group_mask) * torch.finfo(scores.dtype).min).unsqueeze(-1)
+    tmp_scores = tmp_scores.reshape(num_token, -1)
 
     if e_score_correction_bias is not None:
         topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
@@ -972,7 +984,7 @@ def grouped_topk(
     if renormalize:
         topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
 
-    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+    return topk_weights.to(torch.bfloat16), topk_ids.to(torch.int32)
 
 
 def get_config_dtype_str(

vllm/model_executor/layers/rotary_embedding.py

@@ -921,7 +921,7 @@ def _compute_cos_sin_cache(self) -> torch.Tensor:
         cache = torch.cat((cos, sin), dim=-1)
         return cache
 
-    def forward(
+    def forward_native(
         self,
         positions: torch.Tensor,
         query: torch.Tensor,
@@ -963,6 +963,48 @@ def forward(
             key = key_rot
         return query, key
 
+    def forward_hpu(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from habana_frameworks.torch.hpex.kernels import (
+            RotaryPosEmbeddingMode, apply_rotary_pos_emb)
+        """PyTorch-native implementation equivalent to forward()."""
+        query_rot = query[..., :self.rotary_dim]
+        key_rot = key[..., :self.rotary_dim]
+        if self.rotary_dim < self.head_size:
+            query_pass = query[..., self.rotary_dim:]
+            key_pass = key[..., self.rotary_dim:]
+
+        self.cos_sin_cache = self.cos_sin_cache.to(positions.device)
+        cos_sin = self.cos_sin_cache[torch.add(positions, offsets)
+                                     if offsets is not None else positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        rope_mode: RotaryPosEmbeddingMode
+        if self.is_neox_style:
+            rope_mode = RotaryPosEmbeddingMode.BLOCKWISE
+            # NOTE(woosuk): Here we assume that the positions tensor has the
+            # shape [batch_size, seq_len].
+            cos = cos.repeat(1, 1, 2).unsqueeze(-2)
+            sin = sin.repeat(1, 1, 2).unsqueeze(-2)
+        else:
+            rope_mode = RotaryPosEmbeddingMode.PAIRWISE
+            cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
+            sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
+        query_rot = apply_rotary_pos_emb(query_rot, cos, sin, None, 0, rope_mode)
+        key_rot = apply_rotary_pos_emb(key_rot, cos, sin, None, 0, rope_mode)
+
+        if self.rotary_dim < self.head_size:
+            query = torch.cat((query_rot, query_pass), dim=-1)
+            key = torch.cat((key_rot, key_pass), dim=-1)
+        else:
+            query = query_rot
+            key = key_rot
+        return query, key
+
 
 class Llama3RotaryEmbedding(RotaryEmbedding):
 

vllm/model_executor/models/deepseek_v2.py

@@ -156,8 +156,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         input_shape = hidden_states.shape
         hidden_dim = input_shape[-1]
         hidden_states = hidden_states.view(-1, hidden_dim)
-        if self.n_shared_experts is not None:
-            shared_output = self.shared_experts(hidden_states)
         # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.gate(hidden_states)
 
@@ -170,9 +168,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             # See DeepseekV2DecoderLayer for more details.
             final_hidden_states = self.experts(hidden_states=hidden_states,
                                                router_logits=router_logits)
+
+        if self.n_shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states)
+
         if shared_output is not None:
             if hidden_states.dtype != torch.float16:
-                final_hidden_states = final_hidden_states + shared_output
+                final_hidden_states.add_(shared_output)
             else:
                 # Fix FP16 overflow
                 # See DeepseekV2DecoderLayer for more details.