issue/297 - compile all paged batch sizes up to 64

Author: wooway777

csrc/engine/compiler/paged_compiler.cpp

@@ -18,10 +18,7 @@ inline void set_minus_one(infinicore::Tensor &tensor) {
 namespace infinilm::engine {
 PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
     : GraphCompiler(model, barrier) {
-    for (size_t b = 1; b < 32; b++) {
-        decode_batch_sizes_.push_back(b);
-    }
-    for (size_t b = 32; b < 64; b += 8) {
+    for (size_t b = 1; b < 64; ++b) {
         decode_batch_sizes_.push_back(b);
     }
     for (size_t b = 64; b < 128; b += 16) {

python/infinilm/llm/llm.py

@@ -334,21 +334,23 @@ def _check_request_finished(self, req: InferenceRequest, token_id: int) -> bool:
             req.finish_reason = FinishReason.LENGTH
             return True
 
-        # Check EOS token
-        eos_ids = req.eos_token_ids or self.eos_token_ids
-        if eos_ids and token_id in eos_ids:
-            req.finish_reason = FinishReason.EOS_TOKEN
-            return True
-
-        # Check stop strings
-        # Remove stop string from generated_text if STOP_STRING finish reason
-        stop_strings = req.sampling_params.stop or []
-        for stop_str in stop_strings:
-            if req.generated_text.endswith(stop_str):
-                req.generated_text = req.generated_text[: -len(stop_str)]
-                req.finish_reason = FinishReason.STOP_STRING
+        if not req.sampling_params.ignore_eos:
+            # Check EOS token - only stop if ignore_eos is False
+            eos_ids = req.eos_token_ids or self.eos_token_ids
+            if eos_ids and token_id in eos_ids:
+                req.finish_reason = FinishReason.EOS_TOKEN
                 return True
 
+            # While ignoring EOS, stop strings are also ignored to avoid requiring additional arguments for benchmarking.
+            # Check stop strings
+            # Remove stop string from generated_text if STOP_STRING is the finishing reason
+            stop_strings = req.sampling_params.stop or []
+            for stop_str in stop_strings:
+                if req.generated_text.endswith(stop_str):
+                    req.generated_text = req.generated_text[: -len(stop_str)]
+                    req.finish_reason = FinishReason.STOP_STRING
+                    return True
+
         return False
 
     def tokenize(self, text: str) -> List[int]:

python/infinilm/llm/sampling_params.py

@@ -15,7 +15,10 @@ class SamplingParams:
     top_k: int = 1
     max_tokens: Optional[int] = None
     stop: Optional[List[str]] = None
-    stop_token_ids: Optional[List[int]] = None  # Placeholder for future usage, not currently handled
+    stop_token_ids: Optional[List[int]] = (
+        None  # Placeholder for future usage, not currently handled
+    )
+    ignore_eos: bool = False
 
     def __post_init__(self):
         if self.stop is None:
@@ -32,4 +35,5 @@ def clone(self) -> "SamplingParams":
             max_tokens=self.max_tokens,
             stop=self.stop.copy() if self.stop else None,
             stop_token_ids=self.stop_token_ids.copy() if self.stop_token_ids else None,
+            ignore_eos=self.ignore_eos,
         )

python/infinilm/server/inference_server.py

@@ -109,6 +109,7 @@ def __init__(
         port: int = 8000,
         enable_graph: bool = False,
         attn_backend: str = "default",
+        ignore_eos: bool = False,
     ):
         """Initialize inference server.
 
@@ -150,6 +151,7 @@ def __init__(
         self.port = port
         self.enable_graph = enable_graph
         self.attn_backend = attn_backend
+        self.ignore_eos = ignore_eos
 
         self.engine: AsyncLLMEngine = None
 
@@ -331,6 +333,7 @@ def pick(key: str, default):
             top_k=int(pick("top_k", self.top_k)),
             max_tokens=int(max_tokens) if max_tokens is not None else None,
             stop=stop,
+            ignore_eos=self.ignore_eos,
         )
 
     async def _stream_chat(self, request_id: str, data: dict, http_request: Request):
@@ -382,7 +385,11 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request)
                 # Skip EOS token text for OpenAI API compatibility
                 # Check if this token is an EOS token by comparing token_id with eos_token_ids
                 eos_token_ids = self.engine.engine.eos_token_ids
-                is_eos_token = eos_token_ids and token_output.token_id in eos_token_ids
+                is_eos_token = (
+                    not sampling_params.ignore_eos
+                    and eos_token_ids
+                    and token_output.token_id in eos_token_ids
+                )
 
                 if not is_eos_token and token_output.token_text:
                     # Send token
@@ -631,6 +638,13 @@ def parse_args():
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Logging level",
     )
+    parser.add_argument(
+        "--ignore-eos",
+        action="store_true",
+        dest="ignore_eos",
+        default=False,
+        help="Ignore EOS token and continue generation",
+    )
 
     return parser.parse_args()
 
@@ -644,21 +658,22 @@ def main():
     server = InferenceServer(
         model_path=cfg.model,
         device=device,
-        dtype=cfg.dtype,
-        tensor_parallel_size=cfg.tp,
-        cache_type=cfg.cache_type,
-        max_tokens=cfg.max_tokens,
-        max_batch_size=cfg.max_batch_size,
-        num_blocks=cfg.num_blocks,
-        block_size=cfg.block_size,
-        max_cache_len=cfg.max_cache_len,
-        temperature=cfg.temperature,
-        top_p=cfg.top_p,
-        top_k=cfg.top_k,
-        host=cfg.host,
-        port=cfg.port,
-        enable_graph=cfg.enable_graph,
-        attn_backend=cfg.attn,
+        dtype=args.dtype,
+        tensor_parallel_size=args.tp,
+        cache_type=args.cache_type,
+        max_tokens=args.max_tokens,
+        max_batch_size=args.max_batch_size,
+        num_blocks=args.num_blocks,
+        block_size=args.block_size,
+        max_cache_len=args.max_cache_len,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        top_k=args.top_k,
+        host=args.host,
+        port=args.port,
+        enable_graph=args.enable_graph,
+        attn_backend=args.attn,
+        ignore_eos=args.ignore_eos,
     )
     server.start()