[Bugfix] Add max_num_batched_tokens to InputBatch to make main CI pass

[Potabk]

What this PR does / why we need it?

Fix V1 error found by nightly_ci, broken by [v1] Pass BlockTable and KVCacheSpec to AttentionMetadataBuilders #17483, make InputBatch parameter consistent with vllm.

Does this PR introduce any user-facing change?

How was this patch tested?

CI passed

[Potabk]

@wangxiyuan

[Yikun]

diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index 76d3ea4..11355f8 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -55,6 +55,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 from vllm_ascend.attention.attention import AttentionMaskBuilder
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.platform import NPUPlatform
+from vllm_ascend.utils import vllm_version_is

 if TYPE_CHECKING:
     import xgrammar as xgr  # type: ignore[import-untyped]
@@ -186,14 +187,26 @@ class NPUModelRunner:
         # Request states.
         self.requests: Dict[str, CachedRequestState] = {}
         # Persistent batch.
-        self.input_batch = InputBatch(
-            max_num_reqs=self.max_num_reqs,
-            max_model_len=self.model_config.max_model_len,
-            max_num_blocks_per_req=self.max_num_blocks_per_req,
-            device=self.device,
-            pin_memory=True,
-            vocab_size=self.model_config.get_vocab_size(),
-        )
+        # Remove this when drop 0.8.5 suport
+        if vllm_version_is("0.8.5") or vllm_version_is("0.8.5.post1"):
+            self.input_batch = InputBatch(
+                max_num_reqs=self.max_num_reqs,
+                max_model_len=self.model_config.max_model_len,
+                max_num_blocks_per_req=self.max_num_blocks_per_req,
+                device=self.device,
+                pin_memory=True,
+                vocab_size=self.model_config.get_vocab_size(),
+            )
+        else:
+            self.input_batch = InputBatch(
+                max_num_reqs=self.max_num_reqs,
+                max_model_len=self.model_config.max_model_len,
+                max_num_blocks_per_req=self.max_num_blocks_per_req,
+                max_num_batched_tokens=self.max_num_tokens,
+                device=self.device,
+                pin_memory=True,
+                vocab_size=self.model_config.get_vocab_size(),
+            )

I suggest to keep it simple and will drop useless branch when we don't support v0.8.5

[Yikun]

@Potabk Thanks for fixing it, will merge this after CI passed

[Potabk]

@Yikun Some errors occurred in the CI, I think it's a problem with the python version. The python 3.10 we are using does not import TypedDict from typing , in the python 3.10, the method should import as from typing_extensions import TypedDict, NotRequired

[Potabk]

@Yikun Some errors occurred in the CI, I think it's a problem with the python version. The python 3.10 we are using does not import TypedDict from typing , in the python 3.10, the method should import as from typing_extensions import TypedDict, NotRequired

Found it! this commit broken, this issue should fixed after #17962 merged

[Yikun]

@Potabk Thanks for investigation, let's skipt it and recover main CI first

[Yikun]

CI passed, merge it to recover main CI