fix

Author: niushengxiao

test/benchmark/static_inference/static_benchmark.py

@@ -34,12 +34,9 @@
 from lightllm.models.mistral_mtp.model import MistralMTPModel
 from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel
 from lightllm.server.api_cli import make_argument_parser
-from lightllm.server.router.model_infer.mode_backend.mtp_pre_process import (
-    prepare_mtp_prefill_inputs,
-)
 from lightllm.utils.config_utils import auto_set_fused_shared_experts, get_dtype, get_vocab_size
 from lightllm.utils.dist_utils import init_distributed_env
-from lightllm.utils.envs_utils import set_env_start_args
+from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args
 
 
 DEFAULT_BATCH_SIZES = [2, 8, 16, 32, 64, 128]
@@ -130,18 +127,6 @@ def empty_multimodal_params(batch_size: int) -> List[Dict]:
     return [{"images": [], "audios": []} for _ in range(batch_size)]
 
 
-def mtp_prefill_chunk_size(batch_size: int, prompt_len: int, batch_max_tokens: int) -> int:
-    """Bound each MTP setup prefill step by the production token budget."""
-    if batch_size <= 0:
-        raise ValueError(f"batch_size must be positive, got {batch_size}")
-    if batch_max_tokens < batch_size:
-        raise ValueError(
-            "MTP prefill requires at least one token per request in a step: "
-            f"batch_size={batch_size}, batch_max_tokens={batch_max_tokens}"
-        )
-    return max(1, min(int(prompt_len), int(batch_max_tokens) // int(batch_size)))
-
-
 class StaticBenchmarkExecutor:
     def __init__(
         self,
@@ -199,115 +184,53 @@ def _run_prefill_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResul
 
     def _run_decode_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult:
         mtp_enabled = self._mtp_enabled()
-        token_rows = self.token_source.batch(case.batch_size, case.context_len) if mtp_enabled else None
         measured_tokens = case.batch_size * case.output_len
         elapsed = 0.0
-        ttft_elapsed = 0.0
-        decode_step_count = 0
         iters = self._case_iters(warmup)
 
         for _ in range(iters):
             self._reset_model_cache()
+            candidate_width = self._mtp_step_width() if mtp_enabled else 1
+            req_idx, seq_len, next_ids = self._materialize_context_for_decode(case, candidate_width)
             if mtp_enabled:
-                torch.cuda.synchronize()
-                ttft_start = time.perf_counter()
-                req_idx, seq_len, next_ids = self._prefill_for_decode(case, token_rows, mtp_enabled)
-                torch.cuda.synchronize()
-                ttft_elapsed += time.perf_counter() - ttft_start
-                step_elapsed, step_count = self._run_mtp_decode_steps(
+                elapsed += self._run_mtp_decode_steps(
                     case=case,
                     req_idx=req_idx,
                     seq_len=seq_len,
                     next_ids=next_ids,
                 )
-                elapsed += step_elapsed
-                decode_step_count += step_count
             else:
-                req_idx, seq_len, next_ids = self._materialize_context_for_decode(case)
                 elapsed += self._run_plain_decode_steps(
                     case=case,
                     req_idx=req_idx,
                     seq_len=seq_len,
                     next_ids=next_ids,
                 )
-                decode_step_count += case.output_len
 
         self._reset_model_cache()
-        inter_token_latency_ms = elapsed * 1000.0 / max(1, decode_step_count) if iters > 0 else None
+        # An MTP verification round may commit multiple tokens; ITL is per
+        # logical output token so it remains comparable with plain decode.
+        logical_output_tokens = case.output_len * iters
+        inter_token_latency_ms = elapsed * 1000.0 / max(1, logical_output_tokens) if iters > 0 else None
         return self._make_result(
             case,
             elapsed,
             measured_tokens,
             warmup,
-            ttft_elapsed_s=ttft_elapsed if mtp_enabled else None,
             inter_token_latency_ms=inter_token_latency_ms,
         )
 
-    def _materialize_context_for_decode(self, case: BenchmarkCase):
-        """Allocate historical KV slots so decode can be measured without prefill."""
+    def _materialize_context_for_decode(self, case: BenchmarkCase, candidate_width: int = 1):
+        """Allocate synthetic historical KV so decode throughput excludes prefill."""
         req_idx = self._alloc_req_indexes(case.batch_size)
         self._materialize_cached_prefix(req_idx, case.context_len)
         seq_len = cpu_i32_full((case.batch_size,), case.context_len)
-        next_ids = torch.from_numpy(np.ascontiguousarray(self.token_source.batch(case.batch_size, 1).reshape(-1))).to(
-            torch.int64
-        )
-        return req_idx, seq_len, next_ids
-
-    def _prefill_for_decode(self, case: BenchmarkCase, token_rows: np.ndarray, mtp_enabled: bool):
-        req_idx = self._alloc_req_indexes(case.batch_size)
-        chunk_size = (
-            mtp_prefill_chunk_size(case.batch_size, case.context_len, self.args.batch_max_tokens)
-            if mtp_enabled
-            else None
-        )
-        inputs = self._build_prefill_inputs(
-            token_rows=token_rows,
-            req_idx=req_idx,
-            prompt_len=case.context_len,
-            chunk_size=chunk_size,
-        )
-        output = None
-        next_ids = None
-        for model_input in inputs:
-            output = self._forward_prefill_input(model_input, allow_overlap=not mtp_enabled)
-            self._touch_output(output)
-            next_ids = self._argmax_ids(output.logits)
-            if mtp_enabled:
-                # Main and draft KV must advance together for every chunk.
-                next_ids = self._fill_mtp_prefill_kv(case, model_input, output, next_ids)
-        assert output is not None
-        assert next_ids is not None
-
-        seq_len = cpu_i32_full((case.batch_size,), case.context_len)
+        tokens = np.ascontiguousarray(self.token_source.batch(case.batch_size, candidate_width))
+        next_ids = torch.from_numpy(tokens).to(torch.int64)
+        if candidate_width == 1:
+            next_ids = next_ids.reshape(-1)
         return req_idx, seq_len, next_ids
 
-    def _fill_mtp_prefill_kv(
-        self,
-        case: BenchmarkCase,
-        main_prefill_input: ModelInput,
-        main_output: ModelOutput,
-        first_next_ids: torch.Tensor,
-    ):
-        draft_input = main_prefill_input
-        draft_output = main_output
-        current_next_ids = first_next_ids.cuda(non_blocking=True)
-        mtp_candidates = [current_next_ids.detach().cpu()]
-        for draft_index in range(self._num_mtp_modules()):
-            draft_input = prepare_mtp_prefill_inputs(
-                model_input=draft_input,
-                b_next_token_ids=current_next_ids,
-                mtp_draft_input_hiddens=draft_output.mtp_main_output_hiddens,
-            )
-            draft_output = self.draft_models[draft_index].forward(draft_input)
-            current_next_ids = self._argmax_ids(draft_output.logits).cuda(non_blocking=True)
-            mtp_candidates.append(current_next_ids.detach().cpu())
-
-        step_width = self._mtp_step_width()
-        while len(mtp_candidates) < step_width:
-            mtp_candidates.append(mtp_candidates[-1])
-        next_ids = torch.stack(mtp_candidates[:step_width], dim=1)
-        return next_ids
-
     def _run_plain_decode_steps(
         self,
         case: BenchmarkCase,
@@ -343,13 +266,12 @@ def _run_mtp_decode_steps(
         req_idx: torch.Tensor,
         seq_len: torch.Tensor,
         next_ids: torch.Tensor,
-    ) -> tuple:
+    ) -> float:
         elapsed = 0.0
-        step_count = 0
         generated_len = 0
         step_width = self._mtp_step_width()
         base_req_idx, b_mtp_index = self._build_mtp_decode_index_tensors(req_idx, step_width)
-        current_candidates = next_ids
+        current_candidates = next_ids.cuda(non_blocking=True)
 
         while generated_len < case.output_len:
             accepted_width = self._sample_mtp_accept_width(step_width, case.output_len - generated_len)
@@ -396,13 +318,11 @@ def _run_mtp_decode_steps(
                     accepted_width=accepted_width,
                 )
                 .detach()
-                .cpu()
             )
             seq_len += accepted_width
             generated_len += accepted_width
-            step_count += 1
 
-        return elapsed, step_count
+        return elapsed
 
     def _run_mtp_draft_decode(
         self,
@@ -413,7 +333,7 @@ def _run_mtp_draft_decode(
     ):
         draft_input = model_input
         draft_output = model_output
-        draft_next_ids = self._argmax_ids(model_output.logits).cuda(non_blocking=True)
+        draft_next_ids = torch.argmax(model_output.logits, dim=-1).detach().to(torch.int64)
         generated = [draft_next_ids.detach()]
 
         temporary_mem = None
@@ -429,7 +349,7 @@ def _run_mtp_draft_decode(
             draft_input.mtp_draft_input_hiddens = draft_output.mtp_main_output_hiddens
             draft_model = self.draft_models[step % self._num_mtp_modules()]
             draft_output = draft_model.forward(draft_input)
-            draft_next_ids = self._argmax_ids(draft_output.logits).cuda(non_blocking=True)
+            draft_next_ids = torch.argmax(draft_output.logits, dim=-1).detach().to(torch.int64)
             generated.append(draft_next_ids.detach())
 
             if self.args.mtp_mode.startswith("eagle") and step + 1 < self.args.mtp_step:
@@ -567,7 +487,7 @@ def _make_decode_input(
             total_token_num=int(seq_len.sum().item()),
             max_q_seq_len=1,
             max_kv_seq_len=max_kv_seq_len,
-            input_ids=input_ids.to(torch.int64).cpu(),
+            input_ids=input_ids.to(torch.int64),
             b_req_idx=req_idx,
             b_mtp_index=mtp_index,
             b_seq_len=seq_len,
@@ -1212,6 +1132,15 @@ def init_deferred_cudagraph(args: SimpleNamespace, cases: Sequence[BenchmarkCase
     model_kvargs["graph_max_batch_size"] = graph_batch_size
     model_kvargs["disable_cudagraph"] = False
 
+    decode_batch_sizes = {int(case.batch_size) for case in cases if case.stage == "decode"}
+    if decode_batch_sizes == {graph_batch_size}:
+        # A single profile case only replays its exact MTP-expanded batch. Avoid
+        # spending startup time and graph memory on smaller batches it never uses.
+        args.graph_split_batch_size = 0
+        args.graph_grow_step_size = graph_batch_size
+        set_env_start_args(args)
+        get_env_start_args.cache_clear()
+
     if args.enable_decode_microbatch_overlap:
         graph_batch_size //= 2
     model.graph_max_batch_size = graph_batch_size * (int(args.mtp_step) + 1)
@@ -1369,13 +1298,6 @@ def _is_dp_group_leader(args: SimpleNamespace, rank_id: int) -> bool:
     return rank_id % _dp_world_size(args) == 0
 
 
-def _raw_decode_step_count(result: Dict) -> int:
-    inter_token_latency_ms = result.get("inter_token_latency_ms")
-    if inter_token_latency_ms is None or inter_token_latency_ms <= 0:
-        return 0
-    return max(1, int(round(float(result["elapsed_ms"]) / float(inter_token_latency_ms))))
-
-
 def aggregate_rank_results(args: SimpleNamespace, messages: Sequence[Dict]) -> List[Dict]:
     """Aggregate rank-local measurements into one global result per case."""
     by_case: Dict[int, List[Dict]] = {}
@@ -1416,10 +1338,11 @@ def aggregate_rank_results(args: SimpleNamespace, messages: Sequence[Dict]) -> L
             logical_tokens = batch_size * int(first["context_len"]) * iters
             first["logical_tps"] = logical_tokens / elapsed_s if elapsed_s > 0 else 0.0
 
-        slowest_item = max(rank_items, key=lambda item: float(item["result"]["elapsed_ms"]))
-        decode_step_count = _raw_decode_step_count(slowest_item["result"])
-        if decode_step_count > 0:
-            first["inter_token_latency_ms"] = elapsed_ms / decode_step_count
+        if first["stage"] == "decode" and iters > 0:
+            # DP latency follows the slowest rank but keeps the per-token
+            # denominator; MTP verification rounds are not output tokens.
+            logical_output_tokens = int(first["output_len"]) * iters
+            first["inter_token_latency_ms"] = elapsed_ms / max(1, logical_output_tokens)
 
         ttft_values = [
             float(item["result"]["ttft_ms"]) for item in rank_items if item["result"].get("ttft_ms") is not None