[PD] Improve optimistic prefill (#30951)

test/registered/disaggregation/test_disaggregation_optimistic_prefill.py

@@ -29,7 +29,7 @@ def rid_that_forces_retry(prefix: str) -> str:
         req = SimpleNamespace(
             rid=rid,
             is_retracted=False,
-            time_stats=SimpleNamespace(prefill_retry_count=0),
+            prefill_attempt_count=0,
         )
         if should_force_retry(req):
             return rid
@@ -72,7 +72,7 @@ class TestOptimisticPrefill(
         envs.SGLANG_TEST_FORCE_OPTIMISTIC_PREFILL_RETRY_PROB.set(FORCE_RETRY_PROB)
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.extra_prefill_args = [
-            "--optimistic-prefill-retries",
+            "--optimistic-prefill-attempts",
             "3",
             "--chunked-prefill-size",
             "128",
@@ -129,7 +129,11 @@ class TestOptimisticPrefill(
 
         self.assertGreater(j["meta_info"]["prompt_tokens"], 512)
         assert len(output_logprobs) == completion_tokens
-        assert len(input_logprobs) > 0
+        # Input logprobs must be complete: retried or pending chunks must not
+        # drop their logprobs.
+        self.assertGreaterEqual(
+            len(input_logprobs), j["meta_info"]["prompt_tokens"] - 1
+        )
 
 
 class TestOptimisticPrefillFailure(PDDisaggregationServerBase):
@@ -150,7 +154,7 @@ class TestOptimisticPrefillFailure(PDDisaggregationServerBase):
 
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.extra_prefill_args = [
-            "--optimistic-prefill-retries",
+            "--optimistic-prefill-attempts",
             "3",
             "--chunked-prefill-size",
             "128",