feat: support transformers benchmark

python/infinilm/base_config.py

@@ -311,8 +311,8 @@ def _add_common_args(self):
             "--backend",
             type=str,
             default="cpp",
-            choices=["python", "cpp", "torch", "vllm"],
-            help="backend type",
+            choices=["python", "cpp", "infinilm", "torch", "transformers", "vllm"],
+            help="backend type (cpp/infinilm, transformers/torch, or vllm)",
         )
 
         self.parser.add_argument(

test/bench/backends/__init__.py

@@ -0,0 +1,5 @@
+from .infinilm import InfiniLMBenchmark
+from .transformers import TransformersBenchmark
+from .vllm import VLLMBenchmark
+
+__all__ = ["InfiniLMBenchmark", "TransformersBenchmark", "VLLMBenchmark"]

test/bench/backends/_vllm_hip_compat/sitecustomize.py

@@ -0,0 +1,15 @@
+import os
+
+if os.getenv("INFINILM_VLLM_FORCE_ROCM") == "1":
+    try:
+        import vllm.third_party.pynvml as pynvml
+
+        class HygonNVMLError(RuntimeError):
+            pass
+
+        def disabled_nvml_init():
+            raise HygonNVMLError("NVML detection disabled for a HIP-native build")
+
+        pynvml.nvmlInit = disabled_nvml_init
+    except ImportError:
+        pass

test/bench/backends/base.py

@@ -0,0 +1,91 @@
+import time
+from abc import ABC, abstractmethod
+
+
+def render_ceval(tokenizer, conversation):
+    return (
+        tokenizer.apply_chat_template(
+            conversation=conversation,
+            add_generation_prompt=True,
+            tokenize=False,
+        )
+        + "正确答案是"
+    )
+
+
+def render_mmlu(tokenizer, question, choices):
+    choices_text = "\n".join(
+        [f"{chr(65 + i)}. {choice}" for i, choice in enumerate(choices)]
+    )
+    instruction = (
+        "You are a multiple-choice question solver. "
+        "Select the correct option and respond with only the letter A, B, C, or D."
+    )
+    prompt = f"{instruction}\n\nQuestion: {question}\n{choices_text}\nAnswer:"
+
+    if hasattr(tokenizer, "apply_chat_template"):
+        conversation = [
+            {"role": "system", "content": instruction},
+            {"role": "user", "content": f"{question}\n{choices_text}\n"},
+        ]
+        try:
+            return (
+                tokenizer.apply_chat_template(
+                    conversation=conversation,
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+                + "The answer is: "
+            )
+        except Exception:
+            pass
+    return prompt
+
+
+class BaseBenchmark(ABC):
+    def __init__(self, benchmark):
+        self.benchmark = benchmark
+        self.total_tokens = 0
+        self.total_time = 0.0
+
+    def encode_text(self, text):
+        return self.tokenizer.encode(text)
+
+    def decode_token(self, token_id):
+        return self.tokenizer.decode(token_id)
+
+    def max_context_len(self):
+        return self.config_dict.get("max_position_embeddings", 2048)
+
+    def render_input_content(self, *args, **kwargs):
+        renderer = getattr(self, "processor", self.tokenizer)
+        if self.benchmark == "ceval":
+            return render_ceval(renderer, *args, **kwargs)
+        if self.benchmark == "mmlu":
+            return render_mmlu(renderer, *args, **kwargs)
+        raise ValueError(f"Unknown benchmark: {self.benchmark}")
+
+    def record_generation(self, output_text, input_tokens, new_tokens, start_time):
+        elapsed = time.perf_counter() - start_time
+        total_tokens = input_tokens + new_tokens
+        throughput = total_tokens / elapsed if elapsed > 0 else 0.0
+
+        print(output_text)
+        print()
+        print(f"Total time: {elapsed * 1000:.2f} ms")
+        print(f"Input tokens: {input_tokens}")
+        print(f"New tokens: {new_tokens}")
+        print(f"Total tokens processed: {total_tokens}")
+        print(f"Throughput: {throughput:.2f} tok/s")
+
+        self.total_tokens += total_tokens
+        self.total_time += elapsed
+        return output_text
+
+    @abstractmethod
+    def generate(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def destroy_model_instance(self):
+        pass

test/bench/backends/infinilm.py

@@ -0,0 +1,96 @@
+import json
+import os
+import time
+
+from .base import BaseBenchmark
+
+
+class InfiniLMBenchmark(BaseBenchmark):
+    """InfiniLM backend using the scheduler-backed high-level API."""
+
+    def __init__(
+        self,
+        model_dir_path,
+        device_type_str="cpu",
+        tensor_parallel_size=1,
+        benchmark="ceval",
+        enable_paged_attn=False,
+        enable_graph=False,
+        attn_backend="default",
+    ):
+        from infinilm import LLM
+
+        super().__init__(benchmark)
+
+        device_map = {
+            "cpu": "cpu",
+            "nvidia": "cuda",
+            "cambricon": "mlu",
+            "ascend": "npu",
+            "metax": "cuda",
+            "moore": "musa",
+            "iluvatar": "cuda",
+            "kunlun": "cuda",
+            "hygon": "cuda",
+            "ali": "cuda",
+            "cuda": "cuda",
+            "mlu": "mlu",
+            "musa": "musa",
+            "npu": "npu",
+        }
+        device_name = device_map.get(device_type_str.lower(), "cpu")
+
+        with open(os.path.join(model_dir_path, "config.json"), "r") as f:
+            self.config_dict = json.load(f)
+
+        if enable_paged_attn and attn_backend == "default":
+            attn_backend = "paged-attn"
+
+        print("Loading model with InfiniLM backend...")
+        print(f"Graph compilation: {'enabled' if enable_graph else 'disabled'}")
+        print(f"Attention backend: {attn_backend}")
+
+        self.model = LLM(
+            model_path=model_dir_path,
+            device=device_name,
+            tensor_parallel_size=tensor_parallel_size,
+            cache_type="paged" if enable_paged_attn else "static",
+            max_batch_size=1,
+            num_blocks=128,
+            block_size=256,
+            enable_graph=enable_graph,
+            attn_backend=attn_backend,
+        )
+        self.processor = self.model.engine.processor
+        self.tokenizer = self.processor.get_tokenizer()
+        print("InfiniLM model loaded successfully")
+
+    def generate(self, *args, max_steps=500, topp_=1.0, topk_=1, temperature_=1.0):
+        from infinilm import SamplingParams
+
+        prompt = self.render_input_content(*args)
+        print(prompt, end="", flush=True)
+
+        start_time = time.perf_counter()
+        request_output = self.model.generate(
+            prompts=prompt,
+            sampling_params=SamplingParams(
+                max_tokens=max_steps,
+                temperature=temperature_,
+                top_k=topk_,
+                top_p=topp_,
+            ),
+            use_tqdm=False,
+        )[0]
+
+        completion = request_output.outputs[0]
+        return self.record_generation(
+            completion.text,
+            len(request_output.prompt_token_ids or []),
+            len(completion.token_ids),
+            start_time,
+        )
+
+    def destroy_model_instance(self):
+        del self.model
+        print("InfiniLM model destroyed")

test/bench/backends/transformers.py

@@ -0,0 +1,90 @@
+import json
+import os
+import time
+
+from .base import BaseBenchmark
+
+
+class TransformersBenchmark(BaseBenchmark):
+    """Hugging Face Transformers backend."""
+
+    def __init__(self, model_dir_path, device_type_str="cpu", benchmark="ceval"):
+        import torch
+        import transformers
+
+        super().__init__(benchmark)
+
+        supported_devices = {"cpu", "cuda", "mlu", "musa", "npu"}
+        if device_type_str not in supported_devices:
+            raise ValueError(
+                f"Transformers backend unsupported device type: {device_type_str}"
+            )
+        self.device = torch.device(device_type_str)
+
+        with open(os.path.join(model_dir_path, "config.json"), "r") as f:
+            self.config_dict = json.load(f)
+
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
+            model_dir_path, trust_remote_code=True
+        )
+
+        print("Loading model with Transformers backend...")
+        self.model = transformers.AutoModelForCausalLM.from_pretrained(
+            model_dir_path,
+            dtype="auto",
+            trust_remote_code=True,
+        ).to(self.device)
+        self.model.eval()
+        print("Transformers model loaded successfully")
+
+        eos_token_id = self.config_dict.get("eos_token_id")
+        self.eos_token_id = (
+            [eos_token_id] if isinstance(eos_token_id, int) else eos_token_id
+        )
+
+    def _synchronize(self):
+        import torch
+
+        device_module = getattr(torch, self.device.type, None)
+        synchronize = getattr(device_module, "synchronize", None)
+        if synchronize is not None:
+            synchronize()
+
+    def generate(self, *args, max_steps=500, topp_=1.0, topk_=1, temperature_=1.0):
+        import torch
+
+        prompt = self.render_input_content(*args)
+        print(prompt, end="", flush=True)
+        tokens = self.encode_text(prompt)
+        input_ids = torch.tensor([tokens], device=self.device)
+
+        self._synchronize()
+        start_time = time.perf_counter()
+        outputs = self.model.generate(
+            input_ids=input_ids,
+            max_new_tokens=max_steps,
+            do_sample=temperature_ > 0,
+            temperature=temperature_,
+            top_k=topk_,
+            top_p=topp_,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=(
+                self.tokenizer.pad_token_id
+                if self.tokenizer.pad_token_id is not None
+                else self.tokenizer.eos_token_id
+            ),
+        )
+        self._synchronize()
+
+        generated_ids = outputs[0][len(tokens) :]
+        output_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+        return self.record_generation(
+            output_text,
+            len(tokens),
+            generated_ids.numel(),
+            start_time,
+        )
+
+    def destroy_model_instance(self):
+        del self.model
+        print("Transformers model destroyed")