[NIXL][BUG FIX] Fix a bug for PD with host_buffer after merging 29665 by xuechendi · Pull Request #30420 · vllm-project/vllm

xuechendi · 2025-12-10T19:08:30Z

Purpose

After #29665 merged, the PR added remote_request_id=kv_transfer_params["remote_request_id"], to add_new_req

However, in prefill side, there is no remote_request_id in kv_transfer_params

=> This issue only occurs when running PD with host_buffer

Test Plan

PREFILL_BLOCK_SIZE=16 DECODE_BLOCK_SIZE=16 bash tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh --kv_buffer_device cpu

without the PR, test will fail immediately
With this PR, there will still be an accuracy issue, which fix is proposed at #30419

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

chatgpt-codex-connector · 2025-12-10T19:08:36Z

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

gemini-code-assist

Code Review

This pull request addresses a KeyError that occurs when using NIXL with a host buffer for prefill/decode operations. The issue stems from accessing remote_request_id in kv_transfer_params which may not always be present on the prefill side. The fix correctly uses dict.get() for safe access, preventing a crash. This change is sound and effectively resolves the bug. The removal of a blank line is a minor stylistic improvement. The changes are approved.

markmc

Thank you

This is only an issue with --kv_buffer_device cpu which I guess we don't test in CI

I don't think remote_request_id should be particularly special here - a prefill request should be able to just contain kv_transfer_params = {'do_remote_decode': True} IMO

markmc · 2025-12-11T08:59:54Z

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

            remote_block_ids=kv_transfer_params["remote_block_ids"],
            remote_engine_id=kv_transfer_params["remote_engine_id"],
-            remote_request_id=kv_transfer_params["remote_request_id"],
+            remote_request_id=kv_transfer_params.get("remote_request_id", ""),


I don't think we should require remote_engine_id, remote_block_ids, remote_request_id, remote_host, or remote_port for a prefill request - it just happens that existing proxy servers set those fields, but I don't think there's any reason to require them

Let's use .get(..., None) for all of them - AFAICT all the other fields are None in a prefill request currently

Ok, will do.
remote_request_id=kv_transfer_params.get("remote_request_id", None) => Failed pre-commit mypy, so I used remote_request_id=kv_transfer_params.get("remote_request_id", "") instead

Done, @markmc , please help to review again.

Sorry to be a pain, but I'm pretty sure the proxies are sending None for these values currently - if that's true (please double-check I'm not mistaken!) then we should update the ReqMeta type hints to allow these to be None

Yes, proxies are sending None! I tested with current default("" or 0), it also works, but I can update the datatype in ReqMeta to make it as None as well

@markmc
By adding Optional of None for ReqMeta, the mypy complains more now
Wondering if I can ease the pain by simply providing a default value("" or 0 or [])?

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1129: error: Argument 2 to "submit" of "Executor" has incompatible type "str | None"; expected "str" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1130: error: Argument 3 to "submit" of "Executor" has incompatible type "int | None"; expected "int" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1636: error: Argument 2 to "map" has incompatible type "list[int] | None"; expected "Iterable[int]" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1771: error: Need type annotation for "block_ids_to_permute" (hint: "block_ids_to_permute: list[<type>] = ...") [var-annotated] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1780: error: Argument 1 to "__iadd__" of "list" has incompatible type "list[int] | None"; expected "Iterable[Any]" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1920: error: Argument 1 to "_logical_to_kernel_block_ids" of "NixlConnectorWorker" has incompatible type "list[int] | None"; expected "list[int]" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1937: error: Argument 2 to "_background_nixl_handshake" of "NixlConnectorWorker" has incompatible type "str | None"; expected "str" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1975: error: Argument "dst_engine_id" to "_read_blocks" of "NixlConnectorWorker" has incompatible type "str | None"; expected "str" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1976: error: Argument "remote_request_id" to "_read_blocks" of "NixlConnectorWorker" has incompatible type "str | None"; expected "str" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1977: error: Argument "local_block_ids" to "_read_blocks" of "NixlConnectorWorker" has incompatible type "list[int] | None"; expected "list[int]" [arg-type] vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py:1978: error: Argument "remote_block_ids" to "_read_blocks" of "NixlConnectorWorker" has incompatible type "list[int] | None"; expected "list[int]" [arg-type]

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

xuechendi · 2025-12-11T16:07:52Z

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+            remote_engine_id=kv_transfer_params.get("remote_engine_id"),  # type: ignore
+            remote_request_id=kv_transfer_params.get("remote_request_id"),  # type: ignore
+            remote_host=kv_transfer_params.get("remote_host"),  # type: ignore
+            remote_port=kv_transfer_params.get("remote_port"),  # type: ignore


@markmc , do you think this works?
I tried updating ReqMeta with optional, then the mypy check exploded. => Since we only want to align with toy_proxy, so I am hoping to use # type ignore , in that case, we can still use None and make mypy happy

On the decode side, in reqs_to_save, we do not expect or need any of these remote_ fields - they are for the prefill side only. Make this more clear by putting these fields in their own dataclass which is only present on requests in reqs_to_recv. Signed-off-by: Mark McLoughlin <markmc@redhat.com>

markmc · 2025-12-12T16:36:32Z

The mypy errors are a sign that we can do better here - the current code basically disables type checking

   def _add_new_req(
         ...
         kv_transfer_params: dict[str, Any],
         _req = ReqMeta(
             ...
             remote_engine_id=kv_transfer_params["remote_engine_id"]

Because the value returned by the dict lookup is Any, mypy does no type checking - even though it could be None (and we know it definitely is None on the decode side runtime), mypy doesn't complain that we are assigning it to a str field. And what's worse, it's a really subtle way of disabling type checking ... it's not at all obvious to the reader, so someone reading this code might be very surprised to find these fields are None at runtime

I've sent you xuechendi#15 which will enable proper type checking. Please do check it over and make sure it makes sense to you 👍

[NIXL] Refactor prefill-only ReqMeta into RemoteMeta

xuechendi · 2025-12-12T16:53:07Z

The mypy errors are a sign that we can do better here - the current code basically disables type checking
   def _add_new_req(
         ...
         kv_transfer_params: dict[str, Any],
         _req = ReqMeta(
             ...
             remote_engine_id=kv_transfer_params["remote_engine_id"]
Because the value returned by the dict lookup is Any, mypy does no type checking - even though it could be None (and we know it definitely is None on the decode side runtime), mypy doesn't complain that we are assigning it to a str field. And what's worse, it's a really subtle way of disabling type checking ... it's not at all obvious to the reader, so someone reading this code might be very surprised to find these fields are None at runtime

I've sent you xuechendi#15 which will enable proper type checking. Please do check it over and make sure it makes sense to you 👍

Oh, that's quite big change...., I was not expecting to refactoring. Looks good to me, thanks
I have merged the change now @markmc

markmc

@NickLucche PTAL, since I've done a bunch of refactoring in the latest version

njhill

Thanks @xuechendi @markmc! The RemoteMeta refactoring seems much cleaner!

njhill · 2025-12-13T17:45:03Z

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

    return compat_hash


+@dataclass


nit: perf benefit

Suggested change

@dataclass

@dataclass(slots=True)

can add it to other dataclasses here too!

Edit: can do this in a follow-on

…vllm-project#30420) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

…vllm-project#30420) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com> Signed-off-by: Joachim Studnia <joachim@mistral.ai>

…vllm-project#30420) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

…vllm-project#30420) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com> Signed-off-by: Ubuntu <mjtaheri68@gmail.com>

…vllm-project#30420) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com> Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>

…vllm-project#30420) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

Fix a bug when creating meta at prefill side for save to host

3817016

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

xuechendi requested review from ApostaC and NickLucche as code owners December 10, 2025 19:08

mergify bot added the kv-connector label Dec 10, 2025

gemini-code-assist bot reviewed Dec 10, 2025

View reviewed changes

xuechendi mentioned this pull request Dec 11, 2025

[NIXL] Heterogeneous KV Layout and block_size - prefill NHD and nP > nD support #30448

Open

5 tasks

markmc requested changes Dec 11, 2025

View reviewed changes

markmc mentioned this pull request Dec 11, 2025

[NIXL][BUG FIX] Fix both failing issue and accuracy issue with nixl + host_buffer on CUDA #30419

Merged

5 tasks

update add_new_req by setting remote with default value

74aba66

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

xuechendi requested a review from markmc December 11, 2025 15:41

Use None as default and ignore mypy

1916c0b

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

xuechendi commented Dec 11, 2025

View reviewed changes

markmc mentioned this pull request Dec 12, 2025

[NIXL] Refactor prefill-only ReqMeta into RemoteMeta xuechendi/vllm#15

Merged

Merge pull request #15 from markmc/fix_pd_host_buffer

492cc40

[NIXL] Refactor prefill-only ReqMeta into RemoteMeta

mergify bot added the v1 label Dec 12, 2025

markmc approved these changes Dec 12, 2025

View reviewed changes

njhill approved these changes Dec 13, 2025

View reviewed changes

njhill added the ready ONLY add when PR is ready to merge/full CI is needed label Dec 13, 2025

njhill enabled auto-merge (squash) December 13, 2025 17:50

Merge branch 'main' into dev/fix_pd_host_buffer

a708261

njhill merged commit ae2e503 into vllm-project:main Dec 14, 2025
52 checks passed

Uh oh!

Conversation

xuechendi commented Dec 10, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

chatgpt-codex-connector bot commented Dec 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

markmc left a comment

Choose a reason for hiding this comment

Uh oh!

markmc Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

xuechendi Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

xuechendi Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

markmc Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

xuechendi Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xuechendi Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

xuechendi Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

markmc commented Dec 12, 2025

Uh oh!

xuechendi commented Dec 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

markmc left a comment

Choose a reason for hiding this comment

Uh oh!

njhill left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

njhill Dec 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

xuechendi commented Dec 10, 2025 •

edited by github-actions bot

Loading

xuechendi Dec 11, 2025 •

edited

Loading

xuechendi Dec 11, 2025 •

edited

Loading

xuechendi commented Dec 12, 2025 •

edited

Loading

njhill left a comment •

edited

Loading

njhill Dec 13, 2025 •

edited

Loading