[gaudi] Deepseek v2 mla and add ep to unquantized moe (#3287)

sywangyi · web-flow · commit ebb26f0ccd8e · 2025-07-07T11:29:39.000+02:00
Signed-off-by: Wang, Yi A &lt;yi.a.wang@intel.com&gt;
diff --git a/Dockerfile_gaudi b/Dockerfile_gaudi
@@ -118,9 +118,9 @@ ENTRYPOINT ["./entrypoint.sh"]
 # Final image
 FROM base
 
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
-ENV HABANA_VISIBLE_DEVICES all
-ENV OMPI_MCA_btl_vader_single_copy_mechanism NONE
+ENV HF_HUB_ENABLE_HF_TRANSFER=1
+ENV HABANA_VISIBLE_DEVICES=all
+ENV OMPI_MCA_btl_vader_single_copy_mechanism=NONE
 
 COPY backends/gaudi/tgi-entrypoint.sh /tgi-entrypoint.sh
 RUN chmod +x /tgi-entrypoint.sh
diff --git a/backends/gaudi/server/text_generation_server/layers/moe/fp8.py b/backends/gaudi/server/text_generation_server/layers/moe/fp8.py
@@ -51,10 +51,12 @@ def __init__(
         self.rank = weights.process_group.rank()
         self.ep_rank = self.rank
         self.use_ep = os.getenv("USE_EXPERT_PARALLEL", "true").lower() == "true"
-
+        if (n_experts + self.world_size - 1) // self.world_size < 4:
+            self.use_ep = False
         if self.use_ep:
-            n_experts = (n_experts + self.world_size - 1) // self.world_size
-            self.ep_offset = self.ep_rank * n_experts
+            n_experts_per_rank = (n_experts + self.world_size - 1) // self.world_size
+            self.ep_offset = self.ep_rank * n_experts_per_rank
+            n_experts = min(n_experts_per_rank, n_experts - self.ep_offset)
         else:
             self.ep_offset = 0
 
diff --git a/backends/gaudi/server/text_generation_server/layers/moe/unquantized.py b/backends/gaudi/server/text_generation_server/layers/moe/unquantized.py
@@ -7,6 +7,7 @@
 from vllm_hpu_extension.ops import VllmMixtureOfExpertsOp
 import habana_frameworks.torch as htorch
 import torch.nn.functional as F
+import os
 
 
 class UnquantizedSparseMoELayer(nn.Module):
@@ -39,23 +40,42 @@ def __init__(
         self.weight_block_size = weights.weights_loader.weight_block_size
         self.scoring_func = scoring_func
         self.e_score_correction_bias = e_score_correction_bias
+        self.rank = weights.process_group.rank()
+        self.world_size = weights.process_group.size()
+        self.use_ep = os.getenv("USE_EXPERT_PARALLEL", "true").lower() == "true"
+        if (n_experts + self.world_size - 1) // self.world_size < 4:
+            self.use_ep = False
+        if self.use_ep:
+            n_experts_per_rank = (n_experts + self.world_size - 1) // self.world_size
+            self.ep_offset = self.rank * n_experts_per_rank
+            n_experts = min(n_experts_per_rank, n_experts - self.ep_offset)
+            experts_min = self.ep_offset
+            experts_max = self.ep_offset + n_experts - 1
+        else:
+            self.ep_offset = 0
+            experts_min = 0
+            experts_max = n_experts - 1
 
         self.gate_up_proj = _load_expert_multi_weights_col(
             prefix=prefix,
             n_experts=n_experts,
             gate_proj_name=gate_proj_name,
             up_proj_name=up_proj_name,
             weights=weights,
+            use_ep=self.use_ep,
+            ep_offset=self.ep_offset,
         )
 
         self.down_proj = _load_expert_weights_row(
             prefix=prefix,
             n_experts=n_experts,
             name=down_proj_name,
             weights=weights,
+            use_ep=self.use_ep,
+            ep_offset=self.ep_offset,
         )
 
-        self.MoeOp = VllmMixtureOfExpertsOp(n_experts, 0, n_experts - 1)
+        self.MoeOp = VllmMixtureOfExpertsOp(n_experts, experts_min, experts_max)
         for i in range(n_experts):
             self.MoeOp.w13_list[i].set_weight(self.gate_up_proj[i])
             self.MoeOp.w2_list[i].set_weight(self.down_proj[i])
@@ -87,12 +107,23 @@ def _load_expert_multi_weights_col(
     gate_proj_name: str,
     up_proj_name: str,
     weights: Weights,
+    use_ep: bool = False,
+    ep_offset: int = 0,
 ) -> torch.Tensor:
     all_weight = None
     for i in range(n_experts):
-        weight = weights.get_multi_weights_col(
-            [f"{prefix}.{i}.{gate_proj_name}", f"{prefix}.{i}.{up_proj_name}"], 0
-        )
+        if not use_ep:
+            weight = weights.get_multi_weights_col(
+                [f"{prefix}.{i}.{gate_proj_name}", f"{prefix}.{i}.{up_proj_name}"], 0
+            )
+        else:
+            weight = weights.get_multi_weights(
+                [
+                    f"{prefix}.{i+ep_offset}.{gate_proj_name}",
+                    f"{prefix}.{i+ep_offset}.{up_proj_name}",
+                ],
+                0,
+            )
 
         assert isinstance(weight, UnquantizedWeight)
 
@@ -116,12 +147,19 @@ def _load_expert_weights_row(
     n_experts: int,
     name: str,
     weights: Weights,
+    use_ep: bool = False,
+    ep_offset: int = 0,
 ) -> torch.Tensor:
     all_weight = None
     for i in range(n_experts):
-        weight = weights.get_weights_row(
-            f"{prefix}.{i}.{name}",
-        )
+        if not use_ep:
+            weight = weights.get_weights_row(
+                f"{prefix}.{i}.{name}",
+            )
+        else:
+            weight = weights.get_weights(
+                f"{prefix}.{i+ep_offset}.{name}",
+            )
 
         assert isinstance(weight, UnquantizedWeight)
 
diff --git a/backends/gaudi/server/text_generation_server/models/custom_modeling/flash_deepseek_v2_modeling.py b/backends/gaudi/server/text_generation_server/models/custom_modeling/flash_deepseek_v2_modeling.py
@@ -28,11 +28,12 @@
     TensorParallelEmbedding,
     TensorParallelRowLinear,
     get_linear,
+    Fp8Linear,
 )
 from text_generation_server.layers.attention import (
     Seqlen,
     attention,
-    paged_attention,
+    paged_attention_mla,
     set_block_mapping,
     HPUPagedAttentionMetadata,
 )
@@ -44,6 +45,18 @@
 import habana_frameworks.torch as htorch
 
 
+def get_and_maybe_dequant_weights(layer: torch.nn.Module) -> torch.Tensor:
+    if isinstance(layer, Fp8Linear):
+        eye = torch.eye(
+            layer.qweight.shape[-1], dtype=torch.bfloat16, device=layer.qweight.device
+        )
+        dequant_weights = layer(eye)
+        del eye
+        # standardize to (output, input)
+        return dequant_weights.T
+    return layer.weight
+
+
 class DeepseekV2Config(PretrainedConfig):
     def __init__(
         self,
@@ -246,6 +259,45 @@ def __init__(
             0, self.num_key_value_heads, dtype=torch.int32, device=weights.device
         ).repeat_interleave(self.num_groups)
 
+        kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj.linear).T
+        kv_b_proj_weight = kv_b_proj_weight.view(
+            self.kv_lora_rank,
+            self.num_heads,
+            self.qk_nope_head_dim + self.value_head_size,
+        )
+
+        W_UK, W_UV = kv_b_proj_weight.split(
+            [self.qk_nope_head_dim, self.value_head_size], dim=-1
+        )
+        # Convert from (L, N, V) to (N, L, V)
+        self.W_UV = W_UV.transpose(0, 1)
+        # Convert from (L, N, P) to (N, P, L)
+        self.W_UK_T = W_UK.permute(1, 2, 0)
+
+    def _q_proj_and_k_up_proj(self, x):
+        q_proj = self.q_proj if self.q_lora_rank is None else self.q_b_proj
+        q_nope, q_pe = (
+            q_proj(x)
+            .view(-1, self.num_heads, self.head_size)
+            .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        )
+
+        # Convert from (B, N, P) to (N, B, P)
+        q_nope = q_nope.transpose(0, 1)
+        # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
+        ql_nope = torch.bmm(q_nope, self.W_UK_T)
+        # Convert from (N, B, L) to (B, N, L)
+        return ql_nope.transpose(0, 1), q_pe
+
+    def _v_up_proj_and_o_proj(self, x):
+        # Convert from (B, N, L) to (N, B, L)
+        x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
+        # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
+        x = torch.bmm(x, self.W_UV)
+        # Convert from (N, B, V) to (B, N * V)
+        x = x.transpose(0, 1).reshape(-1, self.num_heads * self.value_head_size)
+        return self.o_proj(x)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -258,28 +310,28 @@ def forward(
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ):
         if self.q_lora_rank is None:
-            query = self.q_proj(hidden_states)
+            hidden_states_or_q_c = hidden_states
         else:
-            query = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))[0])
-        query = query.view(-1, self.num_heads, self.head_size)
-
-        _, query_pe = torch.split(
-            query, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
-        )
+            hidden_states_or_q_c = self.q_a_layernorm(self.q_a_proj(hidden_states))[0]
 
         compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
         compressed_kv, key_pe = torch.split(
             compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
         )
 
         key_pe = key_pe.view(-1, 1, self.qk_rope_head_dim)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv.contiguous())[0]).view(
-            -1, self.num_key_value_heads, self.qk_nope_head_dim + self.value_head_size
-        )
+        kv_c_normed = self.kv_a_layernorm(compressed_kv.contiguous())[0]
 
-        key_nope, value = torch.split(
-            kv, [self.qk_nope_head_dim, self.value_head_size], dim=-1
-        )
+        # Prefill
+        if cu_seqlen_prefill is not None:
+            q_proj = self.q_proj if self.q_lora_rank is None else self.q_b_proj
+            query = q_proj(hidden_states_or_q_c)
+            query = query.view(-1, self.num_heads, self.head_size)
+            query_nope, query_pe = torch.split(
+                query, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+            )
+        else:
+            query_nope, query_pe = self._q_proj_and_k_up_proj(hidden_states_or_q_c)
 
         batch_size, heads, head_dim = query_pe.shape
         query_pe = (
@@ -294,33 +346,47 @@ def forward(
             .reshape(batch_size, heads, head_dim)
         )
         self.rotary_emb(query_pe, key_pe, cos, sin)
-
-        query[..., self.qk_nope_head_dim :] = query_pe
-        key = torch.empty_like(query)
-        key[..., : self.qk_nope_head_dim] = key_nope
-        key[..., self.qk_nope_head_dim :] = key_pe
-
-        # We need to pad the heads because Flash Attention does not support
-        # qk and v with different head sizes.
-        query = torch.nn.functional.pad(
-            query, (0, self.head_pad_size - self.head_size), value=0
-        )
-        key = torch.nn.functional.pad(
-            key, (0, self.head_pad_size - self.head_size), value=0
-        )
-        value = torch.nn.functional.pad(
-            value, (0, self.head_pad_size - self.value_head_size), value=0
+        latent_vec_k = torch.concat(
+            (kv_c_normed, key_pe.view(-1, self.qk_rope_head_dim)), dim=-1
         )
+        latent_vec_k = latent_vec_k.view(-1, self.qk_rope_head_dim + self.kv_lora_rank)
+
+        latent_vec_k = latent_vec_k.unflatten(0, (slots.size(0), -1))
 
         kv_cache.store(
-            key=key,
-            value=value,
+            key=latent_vec_k,
+            value=None,
             slots=slots,
             kv_scales=self.kv_scales,
         )
 
-        # Prefill
         if cu_seqlen_prefill is not None:
+            kv = self.kv_b_proj(kv_c_normed).view(
+                -1,
+                self.num_key_value_heads,
+                self.qk_nope_head_dim + self.value_head_size,
+            )
+
+            key_nope, value = torch.split(
+                kv, [self.qk_nope_head_dim, self.value_head_size], dim=-1
+            )
+            query[..., self.qk_nope_head_dim :] = query_pe
+            key = torch.empty_like(query)
+            key[..., : self.qk_nope_head_dim] = key_nope
+            key[..., self.qk_nope_head_dim :] = key_pe
+
+            # We need to pad the heads because Flash Attention does not support
+            # qk and v with different head sizes.
+            query = torch.nn.functional.pad(
+                query, (0, self.head_pad_size - self.head_size), value=0
+            )
+            key = torch.nn.functional.pad(
+                key, (0, self.head_pad_size - self.head_size), value=0
+            )
+            value = torch.nn.functional.pad(
+                value, (0, self.head_pad_size - self.value_head_size), value=0
+            )
+
             # flash attention
             attn_output = attention(
                 query=query,
@@ -331,24 +397,26 @@ def forward(
                 seqlen=seqlen,
                 softmax_scale=self.softmax_scale,
             )
-        # Decode
+            attn_output = attn_output[..., : self.value_head_size]
+
+            return self.o_proj(
+                attn_output.reshape(-1, self.num_heads * self.value_head_size)
+            )
         else:
-            attn_output = paged_attention(
+            # Decode
+            query = torch.cat([query_nope, query_pe], dim=-1)
+            attn_output = paged_attention_mla(
                 query,
                 kv_cache,
                 self.kv_head_mapping,
                 self.softmax_scale,
                 seqlen,
                 kv_scales=self.kv_scales,
                 hpu_attention_meta=hpu_attention_meta,
+                kv_lora_rank=self.kv_lora_rank,
             )
-
-        # Remove padding.
-        attn_output = attn_output[..., : self.value_head_size]
-
-        return self.o_proj(
-            attn_output.reshape(-1, self.num_heads * self.value_head_size)
-        )
+            attn_output = self._v_up_proj_and_o_proj(attn_output)
+            return attn_output
 
 
 class DeepseekV2MLP(nn.Module):
diff --git a/backends/gaudi/server/text_generation_server/models/custom_modeling/flash_qwen3_moe_modeling.py b/backends/gaudi/server/text_generation_server/models/custom_modeling/flash_qwen3_moe_modeling.py
@@ -21,6 +21,7 @@
 from text_generation_server.layers.attention import (
     attention,
     paged_attention,
+    set_block_mapping,
     Seqlen,
     HPUPagedAttentionMetadata,
 )
@@ -466,6 +467,10 @@ def forward(
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
     ) -> torch.Tensor:
+        if hpu_attention_meta is not None:
+            hpu_attention_meta = set_block_mapping(
+                hpu_attention_meta, inputs_embeds.shape[0]
+            )
 
         hidden_states = inputs_embeds
 
diff --git a/backends/gaudi/server/text_generation_server/models/flash_causal_lm.py b/backends/gaudi/server/text_generation_server/models/flash_causal_lm.py
@@ -1606,7 +1606,7 @@ def init_kv_cache(
     ):
         self.kv_cache = []
         empty_cache()
-        if self.config.model_type == "deepseek_v3":
+        if self.config.model_type in ["deepseek_v3", "deepseek_v2"]:
             self.kv_cache = [
                 KVCompressCache(
                     num_blocks=num_blocks,
@@ -1646,7 +1646,7 @@ def warmup(
         # Inspired by the original implementation in [vllm](https://github.com/vllm-project/vllm)
         # Calculate the number of blocks that can be allocated with the free memory
         dtype_size = torch.tensor([], dtype=self.kv_cache_dtype).element_size()
-        if self.config.model_type == "deepseek_v3":
+        if self.config.model_type in ["deepseek_v3", "deepseek_v2"]:
             cache_block_size = BLOCK_SIZE * (
                 self.config.kv_lora_rank + self.config.qk_rope_head_dim
             )
diff --git a/backends/gaudi/server/text_generation_server/tgi_service.py b/backends/gaudi/server/text_generation_server/tgi_service.py
diff --git a/docs/source/backends/gaudi.mdx b/docs/source/backends/gaudi.mdx
