fix mtp torchair

Author: JC-ut0

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -30,18 +30,23 @@ def get_masked_input_and_mask(
         added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
     # torch.compile will fuse all of the pointwise ops below
     # into a single kernel, making it very fast
-    org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ <
-                                                          org_vocab_end_index)
-    added_vocab_mask = (input_ >= added_vocab_start_index) & (
-        input_ < added_vocab_end_index)
-    added_offset = added_vocab_start_index - (
-        org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding
-    valid_offset = (org_vocab_start_index *
-                    org_vocab_mask) + (added_offset * added_vocab_mask)
-    vocab_mask = org_vocab_mask | added_vocab_mask
+    org_vocab_mask = (input_ >= org_vocab_start_index) & (
+        input_ < org_vocab_end_index)
+    # Adapt: avoid create added_vocab_mask when added_vocab_start_index == added_vocab_end_index.
+    if added_vocab_start_index == added_vocab_end_index:
+        valid_offset = (org_vocab_start_index * org_vocab_mask)
+        vocab_mask = org_vocab_mask
+    else:
+        added_vocab_mask = (input_ >= added_vocab_start_index) & (
+            input_ < added_vocab_end_index)
+        added_offset = added_vocab_start_index - (
+            org_vocab_end_index -
+            org_vocab_start_index) - num_org_vocab_padding
+        valid_offset = (org_vocab_start_index *
+                        org_vocab_mask) + (added_offset * added_vocab_mask)
+        vocab_mask = org_vocab_mask | added_vocab_mask
+    # Adapt end.
     input_ = vocab_mask * (input_ - valid_offset)
-    #FIXME(xyx) refactor this
-    torch._dynamo.mark_static(vocab_mask)
     return input_, ~vocab_mask
 
 

vllm_ascend/worker/mtp_proposer_v1.py

@@ -100,9 +100,7 @@ def prepare_inputs(
         # [a, b, c] -> [a - n1, b - n2, c - n3]
         num_tokens_per_req = query_len_per_req - num_rejected_tokens
         if is_torchair_graph:
-            cu_num_tokens = torch.empty_like(cu_target_query_lens)
-            torch.cumsum(num_tokens_per_req, dim=0, out=cu_num_tokens[1:])
-            cu_num_tokens[0] = 0
+            cu_num_tokens = cu_target_query_lens
             relative_index = query_len_per_req - num_rejected_tokens - 1
             token_indices = cu_num_tokens[:-1] + relative_index
         elif force_one_token:
@@ -239,7 +237,8 @@ def propose(
                         input_ids=self.input_ids[:num_input_tokens],
                         positions=self.positions[:num_input_tokens],
                         previous_hidden_states=self.
-                        hidden_states[:num_input_tokens])
+                        hidden_states[:num_input_tokens],
+                        kv_caches=self.runner.kv_caches[-1:])
         sample_hidden_states = hidden_states[last_token_indices]
         logits = self.model.compute_logits(sample_hidden_states, None)
         draft_token_ids = logits.argmax(dim=-1)