[Feature] Add support for custom DeepSeek modeling in ACL Graph mode

Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

Author: yiz-liu

vllm_ascend/attention/mla_v1.py

@@ -9,6 +9,8 @@
                                               MLAAttentionImpl)
 from vllm.attention.backends.utils import PAD_SLOT_ID
 from vllm.config import get_current_vllm_config
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.utils import direct_register_custom_op
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                LinearBase, RowParallelLinear,
                                                UnquantizedLinearMethod)
@@ -669,130 +671,180 @@ def forward(
         kv_cache: torch.Tensor,
         attn_metadata: M,
         output: Optional[torch.Tensor] = None,
+        trace_flag: bool = True,
     ) -> torch.Tensor:
         assert output is not None, "Output tensor must be provided."
-        if attn_metadata is None:
-            # Profiling run.
-            return output
-        self.running_in_graph = self.enable_graph_mode and attn_metadata.attn_state == AscendAttentionState.DecodeOnly
-        num_actual_toks = attn_metadata.num_actual_tokens
-        if k_pe is None and not self.running_in_graph:
-            kv_c, k_pe = self.kv_a_proj_with_mqa(
-                hidden_states_or_kv_c_normed)[0].split(
-                    [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
-            kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
+        if trace_flag:
+            torch.ops.vllm.unified_ascend_mla_attention_with_output(
+                query=hidden_states_or_q_c,
+                key=hidden_states_or_kv_c_normed,
+                value=k_pe,
+                output=output,
+                layer_name=layer.layer_name)
         else:
-            kv_c_normed = hidden_states_or_kv_c_normed
-        assert attn_metadata.num_decodes is not None and \
-        attn_metadata.num_prefills is not None and \
-        attn_metadata.num_decode_tokens is not None
-        has_decode = attn_metadata.num_decodes > 0
-        has_prefill = attn_metadata.num_prefills > 0
-        num_decode_tokens = attn_metadata.num_decode_tokens
-        if not self.running_in_graph:
-            # Inputs and outputs may be padded for CUDA graphs
-            output_padded = output
-            output = output[:num_actual_toks, ...]
-            kv_c_normed = kv_c_normed[:num_actual_toks, ...]
-            prefill_k_c_normed = kv_c_normed[num_decode_tokens:]
-        if not self.running_in_graph:
-            hidden_states_or_q_c = hidden_states_or_q_c[:num_actual_toks, ...]
-            decode_hs_or_q_c = hidden_states_or_q_c[:num_decode_tokens]
-            prefill_hs_or_q_c = hidden_states_or_q_c[num_decode_tokens:]
-            k_pe = k_pe[:num_actual_toks, ...]
-            k_pe = k_pe.unsqueeze(1)
-            decode_k_pe = k_pe[:num_decode_tokens]
-            prefill_k_pe = k_pe[num_decode_tokens:]
-        else:
-            decode_hs_or_q_c = hidden_states_or_q_c
-        if has_decode:
-            decode_k_nope = None
-            assert attn_metadata.decode is not None
-            decode_ql_nope, decode_q_pe = \
-                self._q_proj_and_k_up_proj(decode_hs_or_q_c)
-            if self.running_in_graph:
-                seq_len = self.rotary_emb.max_position_embeddings
-                cos = self.rotary_emb.cos_cached[:seq_len].to(
-                    dtype=decode_q_pe.dtype)
-                sin = self.rotary_emb.sin_cached[:seq_len].to(
-                    dtype=decode_q_pe.dtype)
-                cos = cos[attn_metadata.decode.input_positions]
-                sin = sin[attn_metadata.decode.input_positions]
-                cos = cos[:, None, None, :]
-                sin = sin[:, None, None, :]
-                decode_q_pe = self.rope_single(decode_q_pe, cos, sin)
-                decode_k_pe, decode_k_nope = self.exec_kv(
-                    hidden_states_or_kv_c_normed, cos, sin, kv_cache,
-                    attn_metadata.slot_mapping)
+            if attn_metadata is None:
+                # Profiling run.
+                return output
+            self.running_in_graph = self.enable_graph_mode and attn_metadata.attn_state == AscendAttentionState.DecodeOnly
+            num_actual_toks = attn_metadata.num_actual_tokens
+            if k_pe is None and not self.running_in_graph:
+                kv_c, k_pe = self.kv_a_proj_with_mqa(
+                    hidden_states_or_kv_c_normed)[0].split(
+                        [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
+            else:
+                kv_c_normed = hidden_states_or_kv_c_normed
+            assert attn_metadata.num_decodes is not None and \
+            attn_metadata.num_prefills is not None and \
+            attn_metadata.num_decode_tokens is not None
+            has_decode = attn_metadata.num_decodes > 0
+            has_prefill = attn_metadata.num_prefills > 0
+            num_decode_tokens = attn_metadata.num_decode_tokens
+            if not self.running_in_graph:
+                # Inputs and outputs may be padded for CUDA graphs
+                output_padded = output
+                output = output[:num_actual_toks, ...]
+                kv_c_normed = kv_c_normed[:num_actual_toks, ...]
+                prefill_k_c_normed = kv_c_normed[num_decode_tokens:]
+            if not self.running_in_graph:
+                hidden_states_or_q_c = hidden_states_or_q_c[:num_actual_toks, ...]
+                decode_hs_or_q_c = hidden_states_or_q_c[:num_decode_tokens]
+                prefill_hs_or_q_c = hidden_states_or_q_c[num_decode_tokens:]
+                k_pe = k_pe[:num_actual_toks, ...]
+                k_pe = k_pe.unsqueeze(1)
+                decode_k_pe = k_pe[:num_decode_tokens]
+                prefill_k_pe = k_pe[num_decode_tokens:]
             else:
-                decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
-                    attn_metadata.decode.input_positions,
-                    decode_q_pe.contiguous(),
-                    decode_k_pe,
-                    max_seq_len=attn_metadata.decode.max_seq_lens)
-        if has_prefill:
-            assert attn_metadata.prefill is not None
-            prefill_q = self.q_proj(prefill_hs_or_q_c)[0]\
-                .view(-1, self.num_heads, self.qk_head_dim)
-            prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
-            prefill_q_nope = prefill_q[..., :self.qk_nope_head_dim]
+                decode_hs_or_q_c = hidden_states_or_q_c
+            if has_decode:
+                decode_k_nope = None
+                assert attn_metadata.decode is not None
+                decode_ql_nope, decode_q_pe = \
+                    self._q_proj_and_k_up_proj(decode_hs_or_q_c)
+                if self.running_in_graph:
+                    seq_len = self.rotary_emb.max_position_embeddings
+                    cos = self.rotary_emb.cos_cached[:seq_len].to(
+                        dtype=decode_q_pe.dtype)
+                    sin = self.rotary_emb.sin_cached[:seq_len].to(
+                        dtype=decode_q_pe.dtype)
+                    cos = cos[attn_metadata.decode.input_positions]
+                    sin = sin[attn_metadata.decode.input_positions]
+                    cos = cos[:, None, None, :]
+                    sin = sin[:, None, None, :]
+                    decode_q_pe = self.rope_single(decode_q_pe, cos, sin)
+                    decode_k_pe, decode_k_nope = self.exec_kv(
+                        hidden_states_or_kv_c_normed, cos, sin, kv_cache,
+                        attn_metadata.slot_mapping)
+                else:
+                    decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
+                        attn_metadata.decode.input_positions,
+                        decode_q_pe.contiguous(),
+                        decode_k_pe,
+                        max_seq_len=attn_metadata.decode.max_seq_lens)
+            if has_prefill:
+                assert attn_metadata.prefill is not None
+                prefill_q = self.q_proj(prefill_hs_or_q_c)[0]\
+                    .view(-1, self.num_heads, self.qk_head_dim)
+                prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
+                prefill_q_nope = prefill_q[..., :self.qk_nope_head_dim]
+                if self.enable_graph_mode:
+                    num_tokens = prefill_hs_or_q_c.shape[0]
+                    prefill_k_pe = prefill_k_pe.view(num_tokens, self.num_kv_heads,
+                                                    -1)
+                    if self.rotary_emb.__class__.__name__ == 'RotaryEmbedding':
+                        # NOTE: When scaling not specified
+                        ori_q_pe_shape, ori_k_pe_shape = prefill_q_pe.shape, prefill_k_pe.shape
+                        prefill_q_pe = prefill_q_pe.reshape(num_tokens, -1)
+                        prefill_k_pe = prefill_k_pe.reshape(num_tokens, -1)
+                        prefill_q_pe, prefill_k_pe = self.rotary_emb(
+                            attn_metadata.prefill.input_positions, prefill_q_pe,
+                            prefill_k_pe)
+                        prefill_q_pe = prefill_q_pe.view(ori_q_pe_shape)
+                        prefill_k_pe = prefill_k_pe.view(ori_k_pe_shape)
+                    else:
+                        prefill_q_pe, prefill_k_pe = self.rotary_emb(
+                            attn_metadata.prefill.input_positions, prefill_q_pe,
+                            prefill_k_pe)
+                    prefill_q = torch.cat([prefill_q_nope, prefill_q_pe], dim=-1)
+                else:
+                    prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
+                        attn_metadata.prefill.input_positions,
+                        prefill_q_pe.contiguous(),
+                        prefill_k_pe,
+                        max_seq_len=attn_metadata.prefill.max_seq_lens)
             if self.enable_graph_mode:
-                num_tokens = prefill_hs_or_q_c.shape[0]
-                prefill_k_pe = prefill_k_pe.view(num_tokens, self.num_kv_heads,
-                                                 -1)
-                if self.rotary_emb.__class__.__name__ == 'RotaryEmbedding':
-                    # NOTE: When scaling not specified
-                    ori_q_pe_shape, ori_k_pe_shape = prefill_q_pe.shape, prefill_k_pe.shape
-                    prefill_q_pe = prefill_q_pe.reshape(num_tokens, -1)
-                    prefill_k_pe = prefill_k_pe.reshape(num_tokens, -1)
-                    prefill_q_pe, prefill_k_pe = self.rotary_emb(
-                        attn_metadata.prefill.input_positions, prefill_q_pe,
-                        prefill_k_pe)
-                    prefill_q_pe = prefill_q_pe.view(ori_q_pe_shape)
-                    prefill_k_pe = prefill_k_pe.view(ori_k_pe_shape)
+                if len(kv_cache) > 0 and kv_cache[0].numel(
+                ) > 0 and attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
+                    slots = attn_metadata.slot_mapping
+                    # NOTE: Separate the kv cache in advance to avoid OOM or other issues
+                    torch_npu._npu_reshape_and_cache(key=kv_c_normed.view(
+                        num_tokens, self.num_kv_heads, -1),
+                                                    value=prefill_k_pe,
+                                                    key_cache=kv_cache[0],
+                                                    value_cache=kv_cache[1],
+                                                    slot_indices=slots)
+            elif kv_cache.numel() > 0:
+                key = torch.cat([
+                    kv_c_normed.view([num_actual_toks, self.num_kv_heads, -1]),
+                    k_pe
+                ],
+                                dim=2)
+                torch_npu._npu_reshape_and_cache_siso(
+                    key=key,
+                    key_cache=kv_cache,
+                    slot_indices=attn_metadata.slot_mapping.flatten())
+            if has_prefill:
+                output[num_decode_tokens:] = self._forward_prefill(
+                    prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
+                    attn_metadata)
+            if has_decode:
+                if self.running_in_graph:
+                    return self._forward_decode(decode_ql_nope, decode_q_pe,
+                                                decode_k_nope, decode_k_pe,
+                                                kv_cache, attn_metadata)
                 else:
-                    prefill_q_pe, prefill_k_pe = self.rotary_emb(
-                        attn_metadata.prefill.input_positions, prefill_q_pe,
-                        prefill_k_pe)
-                prefill_q = torch.cat([prefill_q_nope, prefill_q_pe], dim=-1)
-            else:
-                prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
-                    attn_metadata.prefill.input_positions,
-                    prefill_q_pe.contiguous(),
-                    prefill_k_pe,
-                    max_seq_len=attn_metadata.prefill.max_seq_lens)
-        if self.enable_graph_mode:
-            if len(kv_cache) > 0 and kv_cache[0].numel(
-            ) > 0 and attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
-                slots = attn_metadata.slot_mapping
-                # NOTE: Separate the kv cache in advance to avoid OOM or other issues
-                torch_npu._npu_reshape_and_cache(key=kv_c_normed.view(
-                    num_tokens, self.num_kv_heads, -1),
-                                                 value=prefill_k_pe,
-                                                 key_cache=kv_cache[0],
-                                                 value_cache=kv_cache[1],
-                                                 slot_indices=slots)
-        elif kv_cache.numel() > 0:
-            key = torch.cat([
-                kv_c_normed.view([num_actual_toks, self.num_kv_heads, -1]),
-                k_pe
-            ],
-                            dim=2)
-            torch_npu._npu_reshape_and_cache_siso(
-                key=key,
-                key_cache=kv_cache,
-                slot_indices=attn_metadata.slot_mapping.flatten())
-        if has_prefill:
-            output[num_decode_tokens:] = self._forward_prefill(
-                prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
-                attn_metadata)
-        if has_decode:
-            if self.running_in_graph:
-                return self._forward_decode(decode_ql_nope, decode_q_pe,
-                                            decode_k_nope, decode_k_pe,
-                                            kv_cache, attn_metadata)
-            else:
-                output[:num_decode_tokens] = self._forward_decode(
-                    decode_ql_nope, decode_q_pe, decode_k_nope, decode_k_pe,
-                    kv_cache, attn_metadata)
-        return output_padded
+                    output[:num_decode_tokens] = self._forward_decode(
+                        decode_ql_nope, decode_q_pe, decode_k_nope, decode_k_pe,
+                        kv_cache, attn_metadata)
+            return output_padded
+
+
+def unified_ascend_mla_attention_with_output(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> None:
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    self = forward_context.no_compile_layers[layer_name]
+    kv_cache = self.kv_cache[forward_context.virtual_engine]
+    self.impl.forward(self,
+                      query,
+                      key,
+                      value,
+                      kv_cache,
+                      attn_metadata,
+                      output,
+                      trace_flag=False)
+    return
+
+
+def unified_mla_attention_with_output_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="unified_ascend_mla_attention_with_output",
+    op_func=unified_ascend_mla_attention_with_output,
+    mutates_args=["output"],
+    fake_impl=unified_mla_attention_with_output_fake,
+    dispatch_key="PrivateUse1",
+)

vllm_ascend/models/deepseek_v2.py

@@ -35,6 +35,7 @@
 from torch import nn
 from transformers import PretrainedConfig
 from vllm.attention import Attention, AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
                          get_current_vllm_config)
 from vllm.distributed import (get_dp_group, get_pp_group,
@@ -207,13 +208,47 @@ def __init__(
         self.dp_size = get_dp_group().world_size
         batch_size = vllm_config.scheduler_config.max_num_seqs
         self.enable_mc2 = int(os.environ.get("VLLM_ENABLE_MC2", '0')) == 1
+        additional_config = vllm_config.additional_config
+        self.enable_graph_mode = False
+        if additional_config:
+            self.enable_graph_mode = additional_config.get(
+                "enable_graph_mode", False)
 
         params_dtype = torch.get_default_dtype()
         self.final_hidden_states = torch.zeros(
             [batch_size, config.hidden_size], dtype=params_dtype, device="npu")
         self.tp_group = get_tp_group().device_group
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.enable_graph_mode:
+            return self._forward(hidden_states)
+        else:
+            return self._forward_eager(hidden_states)
+    
+    def _forward_eager(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        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)
+        final_hidden_states = self.experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits) * self.routed_scaling_factor
+
+        # NOTE(Yizhou): Quite strange that the order of these two operations
+        # is reversed in the original vLLM code
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+        if shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+    def _forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         attn_metadata = get_forward_context().attn_metadata
         if attn_metadata is None:
             # for profile run
@@ -538,6 +573,7 @@ def forward(
         return hidden_states, residual
 
 
+@support_torch_compile
 class CustomDeepseekV2Model(nn.Module):
 
     fall_back_to_pt_during_load = False