[BugFix] fix speculative decoding memory leak when speculation is disabled

[noyoshi]

There is a memory leak associated with self.previous_hidden_states on the spec_decode_worker when using --speculative-disable-by-batch-size.

It comes from the fact that we do not clean up the tensors except inside of the speculative decoding codepath, but when it is disabled via batch size, we go through the _run_no_spec function which never frees up the hidden states.

You can verify this by checking the output of

llm = vllm.LLM(
    model="model",
    speculative_model="spec-model",
    quantization="fp8", 
    enable_prefix_caching=False,
    enable_chunked_prefill=True,
    num_speculative_tokens=3,
)
sampling_params = llm.get_default_sampling_params()
sampling_params.max_tokens = 100
sampling_params.temperature = 0.0
for i in range(1000):  # Simulate multiple inference runs
    prompt = f"Test prompt {i}"
    outputs = llm.generate([prompt] * 2, sampling_params)
    
    if i % 10 == 0:
        print(f"Allocated: {torch.cuda.memory_allocated() / 1024**2} MB")
        print(f"Cached: {torch.cuda.memory_reserved() / 1024**2} MB")

You can hard code no_spec=True in execute_model in the speculative decoding worker to simulate the behavior of a model that exceeds its max speculative batch size, and remains there for a long period of time.

I think this was missed because it was probably tested with the following scenario:

start model with speculation enabled, disable by batch size = X
batch size hits X, speculation disabled 
--> CUDA memory starts growing without stopping
batch size goes below X, speculation re-enabled
--> CUDA memory reset / fixed via self.previous_hidden_state = None

It won't show up in this example, but in the real world often times once the model goes above X (the max batch size to run speculation), it will stay under heavy load for a long period of time - say for the entire work day (8 hours), at which point you will likely run into a CUDA oom and increasingly bad performance.

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

🚀

[NickLucche]

Hey thanks a lot for the contribution!
This looks correct, but I think to reproduce it more easily you need to set speculative_disable_by_batch_size and also use the async interface to stack up reqs.
Otherwise memory should be constant.

Would you mind turning the example into a unit test that we can run to verify everything is working as intended?

[NickLucche]

Nice work with the test!

What do you think about using the online API in the test with

    with RemoteOpenAIServer(model_name, server_args) as remote_server:
        client: openai.AsyncOpenAI = remote_server.get_async_client()

I think it serves our purpose a bit better. Can you give it a try?
We should still be able to track the CUDA memory usage from the test process.

Other than that we're good to go here.

[noyoshi]

@NickLucche I think its actually a lot harder to test that way - because you need to simulate a situation where you never have smaller batch sizes (which re-enable speculation and clear the variable).

So something like

Low QPS initially (optional)
High constant QPS

• Check cuda memory usage
  Maintain High QPS (but same as earlier, so it should be held constant)
• Check cuda memory usage

Not quite sure how to instrument that out with the standard open AI client

[NickLucche]

We have utils in benchmark_serving.py to send at a request rate.
Otherwise it should still show a mem increase if you just keep sending async requests in a loop as you're bound to cross the batch size limit that disables spec decoding.

Anyways my suggestion is non-blocking, we can still merge this.

[njhill]

Thanks @noyoshi!

FYI @LiuXiaoxuanPKU

[njhill]

@noyoshi could you merge in latest main branch? Not sure what the issue is with the tests.

[noyoshi]

@njhill Hows it looking? Can someone run the CI :)

[DarkLight1337]

PTAL at the failing test

[noyoshi]

Test failures look completely unrelated to my single line change, is there a way to re-run them? Should I merge off master again?

[DarkLight1337]

Yeah, can you try merging from main? Sorry for the delay

[ProExpertProg]

Btw - this broke spec_decoding tests on main because the EngineArgs interface changed. Pushing a fix.
Fixed in #17754