[0.9.1][BUGFIX] [mtp][pd] FIX mtp torchair bug (#2610)

### What this PR does / why we need it?
In the pd Disaggregation scenario, the first token of the inference
after the d node receives the kv follows the eager mode.

Fixes:
Running with MTP torchair graph mode with Prefilling Decoding
Disaggregation , if all requests processed by the D node are requests
just transmitted from the P node, it will break the torchair graph.

Reason: During PD Disaggregation , the P node only transmits the KV
cache and prompt to the D node, not the actual tokens inferred (neither
the main model tokens nor the MTP tokens are transmitted). Therefore,
the D node will treat this request as one without MTP tokens for
inference (seq_len=1).
The community does not have graph mode issues because the community's
attention has a seq_len=1 for each batch during the decode phase.
We have issues because the graph mode pads according to processing 2
tokens per request. When there are some seq_len=1 and some seq_len=2,
padding is done at the end. If all requests received by the D node are
seq_len=1, padding cannot be performed normally according to the
attention's fia operator constraints.

Solution:

The kv consumer uses extra torchair graph padding to avoid breaking FIA
graph constrains (The one this PR implemented).

The kv producer provides the correct tokens to the kv consumer, so that
our graph mode constraints are not broken, and all logic is the same as
the PD mixed deployment . Since we are using the community scheduler,
the modification requires patching the vllm scheduler, but
theoretically, performance should be better. (Maybe later )

Signed-off-by: xuyexiong <xuyexiong@huawei.com>

Author: JC-ut0

vllm_ascend/worker/model_runner_v1.py

@@ -363,6 +363,13 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.use_cached_kv_cache_bytes = ascend_config.torchair_graph_config.use_cached_kv_cache_bytes
         self.torchair_graph_batch_sizes = ascend_config.torchair_graph_config.graph_batch_sizes
 
+        # kv role
+        self.is_kv_producer = False
+        self.is_kv_consumer = False
+        if vllm_config.kv_transfer_config is not None:
+            self.is_kv_producer = vllm_config.kv_transfer_config.is_kv_producer
+            self.is_kv_consumer = vllm_config.kv_transfer_config.is_kv_consumer
+
         if ascend_config.torchair_graph_config.graph_batch_sizes_init:
             self.init_torchair_graph_batch_sizes()
 
@@ -394,13 +401,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         # NOTE: we need to use `in_profile_run` to determine whether `enable_force_load_balance` is True
         self.in_profile_run = False
 
-        # kv role
-        self.is_kv_producer = False
-        self.is_kv_consumer = False
-        if vllm_config.kv_transfer_config is not None:
-            self.is_kv_producer = vllm_config.kv_transfer_config.is_kv_producer
-            self.is_kv_consumer = vllm_config.kv_transfer_config.is_kv_consumer
-
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
         output.
@@ -2339,7 +2339,13 @@ def select_torchair_padded_batch_size(self, batch_size: int):
     def check_torchair_graph_batch_sizes(self):
         # return graph_batch_sizes according to the number of tokens
         # first pad according to the number of requests
-        if len(self.torchair_graph_batch_sizes) == 0:
+        if self.is_kv_consumer:
+            # pd disaggregation scenario may incorrectly calculate the batch, so we force set it to max_num_reqs
+            self.torchair_graph_batch_sizes = [self.max_num_reqs]
+            logger.warning(
+                "is kv_consumer, torch_graph_batch_sizes sets to [max_num_seqs]"
+            )
+        elif len(self.torchair_graph_batch_sizes) == 0:
             self.torchair_graph_batch_sizes = [1, self.max_num_reqs]
         else:
             self.torchair_graph_batch_sizes = sorted(
@@ -2355,10 +2361,23 @@ def check_torchair_graph_batch_sizes(self):
                 self.torchair_graph_batch_sizes.append(self.max_num_reqs)
 
         # we need to make sure that we can deal with max_num_req when `self.decode_token_per_req` is not 1
-        self.torchair_graph_batch_sizes = [
-            graph_batch_size * self.decode_token_per_req
-            for graph_batch_size in self.torchair_graph_batch_sizes
-        ]
+        if self.decode_token_per_req > 1:
+            # pd disaggregation scenario need redundant_batch_sizes to avoid each batch's seq_len exceed 16 tokens
+            if self.is_kv_consumer:
+                FIA_SEQ_LEN_LIMIT = 16
+                self.torchair_graph_batch_sizes = [
+                    (graph_batch_size +
+                     math.ceil(graph_batch_size / FIA_SEQ_LEN_LIMIT) +
+                     math.ceil(graph_batch_size * self.decode_token_per_req /
+                               FIA_SEQ_LEN_LIMIT / FIA_SEQ_LEN_LIMIT)) *
+                    self.decode_token_per_req
+                    for graph_batch_size in self.torchair_graph_batch_sizes
+                ]
+            else:
+                self.torchair_graph_batch_sizes = [
+                    graph_batch_size * self.decode_token_per_req
+                    for graph_batch_size in self.torchair_graph_batch_sizes
+                ]
 
         # NOTE: when enable_expert_parallel, we need to check if `graph_batch_size` is divisible by `tp_size`
         tp_size = self.parallel_config.tensor_parallel_size