[Bugfix][Nixl] Fix Preemption Bug

[robertgshaw2-redhat]

SUMMARY:

• Prior to this change, if the request is recved but there is not enough blocks to schedule the request on this step, we would lose the num_prealloc_computed_tokens state which was ephemeral per step
• This would cause the server to crash as it violates an invariant which we assert
• Update to use num_computed_tokens instead, which is persistent across steps so we properly handle this case

This also opens up a question in my mind of what to do when we run out of memory after a recv. We should handle this in a separate PR / design doc.

The following justfile (running eval will fail on main but works here):

# Setting this allows creating a symlink to Justfile from another dir
set working-directory := "/home/rshaw/vllm/pd_examples/"

# Needed for the proxy server
vllm-directory := "/home/rshaw/vllm/" 

MODEL := "Qwen/Qwen3-0.6B"
TP_SIZE := "1"
PREFILL_GPUS := "0"
DECODE_GPUS := "1"

port PORT: 
  @python port_allocator.py {{PORT}}


prefill:
    VLLM_SERVER_DEV_MODE=1 \
    VLLM_NIXL_SIDE_CHANNEL_PORT=$(just port 5557) \
    CUDA_VISIBLE_DEVICES={{PREFILL_GPUS}} \
    vllm serve {{MODEL}} \
      --port $(just port 8100) \
      --enforce-eager \
      --disable-log-requests \
      --tensor-parallel-size {{TP_SIZE}} \
      --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'

decode:
    VLLM_SERVER_DEV_MODE=1 \
    VLLM_NIXL_SIDE_CHANNEL_PORT=$(just port 5558) \
    CUDA_VISIBLE_DEVICES={{DECODE_GPUS}} \
    vllm serve {{MODEL}} \
      --port $(just port 8200) \
      --enforce-eager \
      --disable-log-requests \
      --tensor-parallel-size {{TP_SIZE}} \
      --num_gpu_blocks_override 1000 \
      --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'

proxy:
    python "{{vllm-directory}}tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
      --port $(just port 8192) \
      --prefiller-port $(just port 8100) \
      --decoder-port $(just port 8200)


send_request:
  curl -X POST http://localhost:$(just port 8192)/v1/completions \
    -H "Content-Type: application/json" \
    -d '{ \
      "model": "{{MODEL}}", \
      "prompt": "Red Hat is the best open source company by far across Linux, K8s, and AI, and vLLM has the greatest community in open source AI software infrastructure. I love vLLM because", \
      "max_tokens": 150, \
      "temperature": 0.7 \
    }'

benchmark:
  python {{vllm-directory}}/benchmarks/benchmark_serving.py --port $(just port 8192) --model {{MODEL}} --dataset-name random --random-input-len 1000 --random-output-len 100

benchmark_one_concurrent:
  python {{vllm-directory}}/benchmarks/benchmark_one_concurrent_req.py \
    --model {{MODEL}} \
    --input-len 1000 \
    --output-len 100 \
    --num-requests 10 \
    --seed 12147 \
    --port $(just port 8192)

reset_prefix_cache:
  curl -X POST http://localhost:$(just port 8100)/reset_prefix_cache && \
  curl -X POST http://localhost:$(just port 8200)/reset_prefix_cache
  

eval:
  lm_eval --model local-completions --tasks gsm8k \
    --model_args model={{MODEL}},base_url=http://127.0.0.1:$(just port 8192)/v1/completions,num_concurrent=1000,max_retries=3,tokenized_requests=False \
    --limit 1000

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mgoin]

num_new_[local/external]_computed_tokens helps a lot with legibility, nice refactor

[njhill]

Thanks @robertgshaw2-redhat. Nice fix and I agree this is much clearer.

I think this may be the correct fix for #18485? cc @NickLucche

[njhill]

I know this came up in a prior PR review but are you sure we want to call this when num_external_computed_tokens == 0?

IMO it's clearer to not do so since the connector has declared that it is not providing any tokens in this case, so should not have any interest in the allocation. In the MultiConnector where we assume that we'll load from at most one connector, we only invoke update_state_after_alloc on the first connector that returned nonzero from get_num_new_matched_tokens.

[robertgshaw2-redhat]

I think that the connector should decide what to do in the case that num_external_tokens is 0.

For instance, in the case of the NIXLConnector, we currently leak memory on the P worker is the D worker has a full local prefix cache hit since the D worker does not have an opportunity to call send_notif() to alert the P worker it can free its blocks.

I will remove the change from this PR and make another one to make this fix more explicit.

[robertgshaw2-redhat]

Thanks @robertgshaw2-redhat. Nice fix and I agree this is much clearer.

I think this may be the correct fix for #18485? cc @NickLucche

No, this is a separate issue that I will address.

[mgoin]

Thanks for the clear logic and unit test