[profiler] support PyTorch profiler enablement

[mcalman]

This PR implements the required methods to support profiling vLLM with PyTorch Profiler.

Using PyTorch Profiler with vLLM

Offline profiling

Enable torch profiler, can also be set on cmd line:

os.environ["VLLM_TORCH_PROFILER_DIR"] = "./vllm_profile"

Start and stop profiling:

llm.start_profile()
outputs = llm.generate(prompts, sampling_params)
llm.stop_profile()

Online profiling

Start server with profiling enabled:

VLLM_RPC_TIMEOUT=1800000 VLLM_TORCH_PROFILER_DIR=./vllm_profile python3 -m vllm.entrypoints.openai.api_server --model /models/llama-7b-chat --max-model-len=2048 --block-size=128

Start and stop profiler on the client side:

requests.post("http://0.0.0.0:8000/start_profile")
client.completions.create(...)
requests.post("http://0.0.0.0:8000/stop_profile")

[github-actions]

👋 Hi! Thank you for contributing to vLLM support on Spyre.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, first install the linting requirements, then run format.sh and commit the changes. This can be done with uv directly:

uv pip install --group lint

Or this can be done with pip:

uv pip compile --group lint > requirements-lint.txt
pip install -r requirements-lint.txt
bash format.sh

Now you are good to go 🚀

[marceloamaral]

Looks great @mcalman!

Just want to mention that the Spyre events will only be collected if PyTorch is using the Kineto that supports Spyre events.

[marceloamaral]

@dpatel-ops could you please have a look on this PR?

[joerunde]

Not gonna lie I'm not an expert in torch profiling.

With the internal package installed I spun this up, profiled it and loaded up the profile in perfetto:

Then in tensorboard I see the gpu listed as AIU 🌶️🌶️🌶️

Reverting the profiler install and running a new profile I see only cpu devices listed:

so, LGTM!

[joerunde]

🚀