fix(inference): AUT-2539 handle string prompts and rejections in KV hand

megatron/core/inference/disaggregation/inference_state_handoff.py

@@ -654,7 +654,7 @@ def _release_pinned_handoff_ssm_slot(self, ssm_slot: int | None) -> None:
     def add_request_with_kv_handoff(
         self,
         request_id: int,
-        prompt: list,
+        prompt: str | list[int],
         sampling_params: "SamplingParams",
         kv_meta: dict,
         src_block_ids: list,
@@ -663,6 +663,11 @@ def add_request_with_kv_handoff(
 
         Side: decode engine; pull and push transport paths. A pull backend starts
         the read here, while a push backend posts the matching receive.
+
+        A string prompt is tokenized once here, as ``add_request`` tokenizes one, so
+        block hashing, the transfer size check and admission all use the same token
+        IDs. Validation errors raise before any cache state is reserved; the engine's
+        SUBMIT_REQUEST_WITH_KV handler turns them into a failed reply.
         """
         allocator = self.context.kv_block_allocator
         if not allocator.enable_prefix_caching:
@@ -710,6 +715,13 @@ def add_request_with_kv_handoff(
                 "Decode-only handoff does not yet transfer prompt or first-token log probabilities"
             )
 
+        if isinstance(prompt, str):
+            # Normalize before hashing so the deferred and pending import records,
+            # and the finalizer's add_request(), carry token IDs. add_request() then
+            # takes its token-list path and does not add a second BOS.
+            prompt = self.controller.tokenize_prompt(
+                self.controller.tokenizer, prompt, sampling_params.add_BOS
+            )
         prompt_tensor = torch.tensor(prompt, dtype=torch.int64)
         hashes = compute_block_hashes_batched(prompt_tensor, self.context.block_size_tokens)
         num_blocks = transfer_block_count(kv_meta, src_block_ids)

megatron/core/inference/engines/dynamic_engine.py

@@ -1830,6 +1830,9 @@ def _fail_submission(
         add_request raises. Registering a placeholder record with
         Status.FAILED lets ``_handle_failed_request`` publish the ENGINE_REPLY
         without leaving the client hanging or killing the engine loop.
+
+        The SUBMIT_REQUEST_WITH_KV handler calls it too, when
+        ``add_request_with_kv_handoff`` rejects a handoff before reserving cache state.
         """
         if self.rank == 0:
             warnings.warn(
@@ -4594,9 +4597,15 @@ def schedule_requests(self) -> int:
                 src_block_ids = msgpack.unpackb(message[2], raw=False)
                 sampling_params = SamplingParams.deserialize(sampling_params)
                 nvtx_range_push("add_request_with_kv_handoff")
-                self.add_request_with_kv_handoff(
-                    request_id, prompt, sampling_params, kv_meta, src_block_ids
-                )
+                # As on the plain path, a rejected handoff must still reply. Every MP
+                # rank receives the same message and raises the same validation error,
+                # so failing it here keeps the ranks' request state mirrored.
+                try:
+                    self.add_request_with_kv_handoff(
+                        request_id, prompt, sampling_params, kv_meta, src_block_ids
+                    )
+                except Exception as error:  # pylint: disable=broad-except
+                    self._fail_submission(request_id, sampling_params, error)
                 nvtx_range_pop("add_request_with_kv_handoff")
             elif header == Headers.RELEASE_KV:
                 # Coordinator-broadcast release. Unknown request ids are no-ops.

tests/unit_tests/inference/engines/test_dynamic_engine.py

@@ -2698,6 +2698,54 @@ def test_schedule_requests_drops_malformed_submit_request(malformed, caplog):
     assert msgpack.unpackb(broadcast[1], raw=False) == [len(malformed), len(good)]
 
 
+def test_schedule_requests_fails_a_rejected_kv_handoff_and_continues():
+    """A SUBMIT_REQUEST_WITH_KV that the handoff rejects fails its request with a reply.
+
+    Every MP rank receives the same broadcast message and raises the same
+    validation error, so failing the submission keeps their request state
+    mirrored instead of taking the engine loop down, and the batch continues.
+    """
+    params = SamplingParams(num_tokens_to_generate=1, termination_id=-1)
+    kv_meta = {"resume_tokens": [9]}
+    handoff = [
+        msgpack.packb(
+            [Headers.SUBMIT_REQUEST_WITH_KV.value, 41, params.serialize(), kv_meta],
+            use_bin_type=True,
+        ),
+        msgpack.packb("hello", use_bin_type=True),
+        msgpack.packb([7, 8], use_bin_type=True),
+    ]
+    good = _submit_request_message(42, params.serialize(), [3, 4], None)
+    engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine)
+    engine.prompt_preparer = None
+    engine.rank, engine.use_coordinator, engine.is_mp_coordinator = 1, True, True
+    engine.requests, engine.failed_request_ids = {}, []
+    rejection = RuntimeError("KV handoff received without a transfer backend")
+    engine.add_request_with_kv_handoff = mock.Mock(side_effect=rejection)
+    engine.add_request = mock.Mock()
+    engine._fail_submission = mock.Mock()
+    socket = engine.socket_for_receiving_requests = mock.Mock()
+    socket.recv_multipart.side_effect = [handoff, good, dynamic_engine.zmq.Again]
+    engine.model_parallel_publisher_socket, engine._pending_signals = mock.Mock(), deque()
+    engine.local_metadata_ledger_enabled = False
+    engine._drain_handoff_completion_notifications = mock.Mock(return_value=[])
+    engine._collect_failed_requests = mock.Mock(return_value=[])
+
+    assert engine.schedule_requests() == 2
+
+    engine.add_request_with_kv_handoff.assert_called_once()
+    args = engine.add_request_with_kv_handoff.call_args.args
+    assert args[0] == 41 and args[1] == "hello"
+    assert args[3] == kv_meta and args[4] == [7, 8]
+    engine._fail_submission.assert_called_once()
+    failed_id, failed_params, error = engine._fail_submission.call_args.args
+    assert failed_id == 41 and error is rejection
+    assert failed_params.num_tokens_to_generate == 1
+    # The next message in the batch is still admitted.
+    engine.add_request.assert_called_once()
+    assert engine.add_request.call_args.args[:2] == (42, [3, 4])
+
+
 def test_handle_failed_request_releases_vlm_request_data():
     """Media registered before admission is dropped when the request fails."""
     request = DynamicInferenceRequest(

tests/unit_tests/inference/test_disagg_handoff_lifecycle.py

@@ -1051,6 +1051,48 @@ def test_decode_handoff_defers_until_kv_capacity_is_available(handoff_loop):
     assert engine._kv_transfer_agent.calls == [(kv_meta, [100, 101], [10, 11])]
 
 
+def test_string_handoff_prompt_is_tokenized_once_before_import(handoff_loop):
+    engine = _HandoffHarness(handoff_loop)
+    engine.context.kv_block_allocator.capacity_available = False
+    prompt_tokens = [7, 1, 2, 3, 4, 5, 6, 8]
+    tokenizer = object()
+    engine.controller = SimpleNamespace(
+        tokenizer=tokenizer, tokenize_prompt=mock.Mock(return_value=prompt_tokens)
+    )
+    sampling_params = SamplingParams(num_tokens_to_generate=2, add_BOS=True)
+    kv_meta = {"request_id": 6, "resume_tokens": [99]}
+    hashes = compute_block_hashes_batched(
+        torch.tensor(prompt_tokens), engine.context.block_size_tokens
+    )
+
+    future = engine.add_request_with_kv_handoff(
+        6, "hello world", sampling_params, kv_meta, [100, 101]
+    )
+
+    # The block-count check and the capacity-deferred record use the token IDs.
+    engine.controller.tokenize_prompt.assert_called_once_with(tokenizer, "hello world", True)
+    (deferred,) = engine._deferred_kv_handoffs
+    assert deferred.prompt == prompt_tokens
+    assert deferred.hashes == hashes
+    assert not future.done()
+
+    engine.context.kv_block_allocator.capacity_available = True
+    engine._poll_pending_kv_imports()
+    _drain_loop(handoff_loop)
+    (pending,) = engine._pending_kv_imports
+    assert pending.prompt == prompt_tokens
+    assert engine._kv_transfer_agent.calls == [(kv_meta, [100, 101], [10, 11])]
+
+    with mock.patch(
+        "megatron.core.inference.disaggregation.inference_state_handoff.admit_prefilled_decode"
+    ):
+        engine._finalize_kv_handoff_import(pending)
+
+    # Admission receives the same token IDs, and nothing tokenizes the prompt again.
+    engine.controller.tokenize_prompt.assert_called_once()
+    assert engine.get_request(6).prompt_tokens.tolist() == prompt_tokens
+
+
 def test_handoff_submission_failure_is_reported_to_model_parallel_coordinator(handoff_loop):
     engine = _HandoffHarness(handoff_loop)
     engine._kv_transfer_agent.begin_pull_blocks = mock.Mock(