[Bugfix][TPU][V1] Fix recompilation

[NickLucche]

It appears recompilation was due to a misconfiguration when pre-compiling the num_reqs.
Added tests to CI to confirm this upstream.

[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.

🚀

[robertgshaw2-redhat]

nice job nicolo!

[mgoin]

MIN_NUM_SEQS is 8, should this be rounding up to the nearest number divisible by MIN_NUM_SEQS?

[NickLucche]

not necessarily see in capture_model, max_num_reqs still gets compiled anyway due to how padding with upper limit is implemented.

[NickLucche]

btw I'm ok with forcing everything to be nicely divisible by MIN_NUM_SEQS; just I remember max_num_seqs used to be padded and then it was changed for reasons

[yarongmu-google]

tests/tpu/test_compilation.py failed it seems w/ this error:

Processed prompts: 100%|████████████████████████████████████████████████████| 1/1 [00:45<00:00, 46.00s/it, est. speed input: 109.60 toks/s, output: 0.11 toks/s]

  | WARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5
  | FAILEDWARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5

Source: https://buildkite.com/vllm/fastcheck/builds/18374#0195d341-8f23-4eb8-bd34-998e6055b2ff

[bvrockwell]

tests/tpu/test_compilation.py failed it seems w/ this error:

Processed prompts: 100%|████████████████████████████████████████████████████| 1/1 [00:45<00:00, 46.00s/it, est. speed input: 109.60 toks/s, output: 0.11 toks/s]

  | WARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5   | FAILEDWARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5

Source: https://buildkite.com/vllm/fastcheck/builds/18374#0195d341-8f23-4eb8-bd34-998e6055b2ff

@NickLucche would we be able to set --enforce-eager=False on all tests except test_compilation.py please to capture whether recompilation is resolved fully?

[robertgshaw2-redhat]

tests/tpu/test_compilation.py failed it seems w/ this error:

Processed prompts: 100%|████████████████████████████████████████████████████| 1/1 [00:45<00:00, 46.00s/it, est. speed input: 109.60 toks/s, output: 0.11 toks/s]

  | WARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5   | FAILEDWARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5
Source: https://buildkite.com/vllm/fastcheck/builds/18374#0195d341-8f23-4eb8-bd34-998e6055b2ff

@NickLucche would we be able to set --enforce-eager=False on all tests except test_compilation.py please to capture whether recompilation is resolved fully?

It’s false by default.

[bvrockwell]

tests/tpu/test_compilation.py failed it seems w/ this error:

Processed prompts: 100%|████████████████████████████████████████████████████| 1/1 [00:45<00:00, 46.00s/it, est. speed input: 109.60 toks/s, output: 0.11 toks/s]

  | WARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5   | FAILEDWARNING 03-26 16:58:08 [parallel_state.py:1093] torch._C._host_emptyCache() only available in Pytorch >=2.5
Source: https://buildkite.com/vllm/fastcheck/builds/18374#0195d341-8f23-4eb8-bd34-998e6055b2ff

@NickLucche would we be able to set --enforce-eager=False on all tests except test_compilation.py please to capture whether recompilation is resolved fully?

It’s false by default.

https://github.com/NickLucche/vllm/blob/tpu-fix-recompilation/tests/v1/tpu/test_basic.py#L34

Am I reading this wrong, where should I be checking this otherwise?

[lsy323]

@NickLucche thank you for fixing this! qq: is this change actually fix the recompilation issue?

[yaochengji]

I also have such confusion. IMO, the code fix the recompilation issue when max_num_reqs is not power of 2. But in our tests, it's already power of 2.

[bvrockwell]

Does it make sense to also enable the following in test_basic.py? enforce_eager=False

[yaochengji]

I think so. We should use the default value of enforce_eager (which is False) in most cases.

[bvrockwell]

indeed, it is set to True in test_basic.py

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @NickLucche.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yaochengji]

@NickLucche Thanks for fixing this. Left a few comments.

[yaochengji]

Why don't we use torch.compile for this method?

Usually it's easier to know the boundary of TPU computation and avoid recompilation if we wrap it inside a torch.compile

[NickLucche]

once main is stable we'll turn that on. I'd like to do that in a separate PR. Last time around compilation got slower, just wanted to be cautious.

[yaochengji]

I did an experiment base on you branch, I cleaned the xla cache each time before the execution.
without torch.compile: sampler pre-compilation time is 35.79 [secs]
with torch.compile: sampler pre-compilation time is 35.51 [secs]

The compilation time difference is negligible. In the meantime, torch.compile can speed up the execution because the guard check of torch.compile is usually faster than torch/xla's graph trace.

[NickLucche]

happy to add it back then! Just would like to merge this PR first if you don't mind, it still fixes the pre-compiliation when max_num_reqs is not a power of 2

[yaochengji]

Could you add it here in this PR? Just one line code change.

[yaochengji]

I think so. We should use the default value of enforce_eager (which is False) in most cases.

[yaochengji]

we'd better add xm.mark_step() before this line unless we use torch.compile.

[NickLucche]

isn't it redundant? A sync to CPU will still cause the graph to be flushed and executed

[yaochengji]

TL;DR, xm.mark_step computes all the pending tensor, but B.cpu() only computes the exact tensor B.

E.g. say we have the code

A = ...
B = op(A)

The graph output generated by xm.mark_step() and B.cpu() are different.

For xm.mark_step(), we will get both A and B as outputs.
For B.cpu(), only B is the output.

Then if we have another xm.mark_step() later, as A's result is not returned in the previous computation, we have to compute A again.

[NickLucche]

Thanks for explaining this is really helpful.
I don't think this is the case as we only need the sampled output tokens from the sampler step, but I could also add it for completeness.

[yaochengji]

Although we don't intend to use other tensors, a later xm.mark_step will still try to get the results of them. Then we have redundant computation.

BTW, we have an implicit xm.mark_step when using torch.compile. What's one reason I recommend using torch.compile when possible.

[yaochengji]

And when we use torch.compile, we don't need so many xm.mark_step.

torch.compile is much more easier for pytorch developer to understand.

[NickLucche]

@NickLucche thank you for fixing this! qq: is this change actually fix the recompilation issue?

So I've added the test you mentioned in the previous PR, test_sampler.py with enforce_eager=False and you can run it with VLLM_XLA_CHECK_RECOMPILATION=1 .
If you did the same thing prior to this PR it would show a graph compilation was detected at runtime.

Mentioned it in this thread too #15309.

[NickLucche]

Also another test worth running that I had mentioned on slack but not here is a benchmark with the check recompilation on.

VLLM_XLA_CHECK_RECOMPILATION=1 VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct \
 --disable-log-requests \
 --port 8004 \
 --gpu-memory-utilization 0.95 \
 --max-num-seqs 512 \
 --max-num-batched-tokens 512 \
 --tensor-parallel-size 1 \
 --max-model-len 2048 > "$VLLM_LOG" 2>&1 &

Feel free to ping me if you spot any case where the check fails though!

[DarkLight1337]

Can you push a dummy commit to retry the doc build?

[NickLucche]

thanks for the ping @DarkLight1337

[bvrockwell]

Thanks for looking into this @NickLucche ! @yaochengji recommended this yesterday and it's a good idea: could we please set enforce-eager for all tests (like test_basic.py, where this is still set to True currently) except for test_compilation.py. This ensures all tests are checking for recompilation and so nothing slips inadvertently in the future.

[NickLucche]

Sure, it's a bit out of scope for this PR I'd suggest we open another one.
Here I only meant to fix the issues related to max_num_reqs when compiling on main + adding a test (test 7) for it.

[yaochengji]

Thanks for your fix again. Could you address my comments in this PR?

[yaochengji]

Although we don't intend to use other tensors, a later xm.mark_step will still try to get the results of them. Then we have redundant computation.

BTW, we have an implicit xm.mark_step when using torch.compile. What's one reason I recommend using torch.compile when possible.

[yaochengji]

And when we use torch.compile, we don't need so many xm.mark_step.

torch.compile is much more easier for pytorch developer to understand.

[yaochengji]

Could you add it here in this PR? Just one line code change.