# Questions on piecewise torch compile design **URL:** \ **Category:** torch.compile integration\ **Created:** [May 7, 2025, 12:52pm UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589 "2025-05-07T12:52:39Z")\ **Posts on this page:** 20\ **Page:** 2
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 22, 2025, 9:17am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/21 "2025-05-22T09:17:44Z")
why lm\_head can not be piecewise compiled ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 22, 2025, 9:18am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/22 "2025-05-22T09:18:03Z")
It seems likely that `lm_head` cannot be piecewise compiled because it is not a custom op or a natural graph boundary like attention, and thus is not included in the list of splitting ops for piecewise compilation. Piecewise compilation in vLLM is designed to split the graph at specific operations (e.g., attention) that benefit from eager execution or have special requirements, while `lm_head` is typically a simple linear layer that does not require such treatment ([design doc](https://docs.vllm.ai/en/latest/design/v1/torch_compile.html#cudagraph-capture)). Would you like more detail on how splitting ops are chosen or how to customize them? Sources: - [design doc](https://docs.vllm.ai/en/latest/design/v1/torch_compile.html#cudagraph-capture) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 22, 2025, 9:22am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/23 "2025-05-22T09:22:57Z")
so will cuda graph capture lm\_head region? ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 22, 2025, 9:23am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/24 "2025-05-22T09:23:16Z")
No, CUDA graph capture in vLLM does not include the `lm_head` region. The computation graph for CUDA graph capture only covers the model from input IDs/embeddings up to the final hidden states; the `lm_head` projection and sampling operations are explicitly excluded from the captured graph ([design doc](https://docs.vllm.ai/en/latest/design/v1/torch_compile.html#computation-graph-processing)). Would you like more detail on what is included in the captured graph? Sources: - [design doc](https://docs.vllm.ai/en/latest/design/v1/torch_compile.html#computation-graph-processing) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 22, 2025, 9:32am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/27 "2025-05-22T09:32:26Z")
and v0 engine will capture lm\_head and logitsprocessor? ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 22, 2025, 9:32am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/28 "2025-05-22T09:32:29Z")
In vLLM v0, CUDA graph capture includes the `lm_head` region but does not capture the logits processor. The computation graph in v0 typically covers the model up to and including the `lm_head` projection, but sampling and logits processing (such as temperature scaling or penalties) are performed outside the captured graph in Python ([design doc](https://docs.vllm.ai/en/latest/design/v1/torch_compile.html#computation-graph-processing)). Would you like more detail on the v0 graph capture process? Sources: - [https://github.com/vllm-project/vllm/pull/9715](https://github.com/vllm-project/vllm/pull/9715) - [https://github.com/vllm-project/vllm/pull/16072](https://github.com/vllm-project/vllm/pull/16072) ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 22, 2025, 9:32am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/29 "2025-05-22T09:32:44Z")
In vLLM v0, CUDA graph capture includes the `lm_head` region but does not capture the `logitsprocessor`. The `lm_head` (final projection layer) is part of the model’s forward pass and is included in the captured CUDA graph, while the `logitsprocessor` (which applies post-processing to logits) is executed outside the CUDA graph in Python for flexibility ([issue 1019](https://github.com/vllm-project/vllm/issues/1019), [custom\_all\_reduce.py](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/custom_all_reduce.py)). Would you like more detail on the execution flow or graph capture boundaries? Sources: - [issue 1019](https://github.com/vllm-project/vllm/issues/1019) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 22, 2025, 9:41am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/30 "2025-05-22T09:41:28Z")
how about the benefit to capture lm\_head in cuda graph in further development ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 22, 2025, 9:42am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/31 "2025-05-22T09:42:02Z")
Capturing `lm_head` in the CUDA graph can reduce kernel launch overhead and improve end-to-end latency, especially for small batch sizes or low-latency scenarios. This is because `lm_head` is a simple linear projection and including it in the graph allows for more efficient, fused execution with the rest of the model ([issue 1019](https://github.com/vllm-project/vllm/issues/1019)). However, the benefit is modest for large batch sizes or throughput-bound workloads, since most compute time is spent in the main transformer layers. The main advantage is in latency-sensitive use cases. Would you like more detail on when this optimization is most impactful? Sources: - [https://github.com/vllm-project/vllm/pull/16072](https://github.com/vllm-project/vllm/pull/16072) - [https://docs.vllm.ai/en/stable/design/v1/torch\_compile.html#cudagraph-capture](https://docs.vllm.ai/en/stable/design/v1/torch_compile.html#cudagraph-capture) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 23, 2025, 3:23am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/32 "2025-05-23T03:23:28Z")
why piecewise compile disable custom op? ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 23, 2025, 3:23am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/33 "2025-05-23T03:23:51Z")
Piecewise compilation disables most custom ops (except attention) because custom ops can interfere with graph splitting and optimization in `torch.compile`. Inductor and CUDA graph capture require clear graph boundaries, and custom ops may introduce side effects or in-place mutations that break assumptions needed for correct piecewise execution ([PR 9058](https://github.com/vllm-project/vllm/pull/9058), [custom\_op.py](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/custom_op.py)). Would you like more detail on how to control which custom ops are enabled? Sources: - [PR 9058](https://github.com/vllm-project/vllm/pull/9058) - [custom\_op.py](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/custom_op.py) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 27, 2025, 8:00am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/34 "2025-05-27T08:00:36Z")
when compile with enable\_sequence\_parallelism, why the computation\_graph seems no changed except moving rmsnorm’s implementation to cuda kernel? allreduce didn’t split to reduce scatter and all gather on Qwen3 model. ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 27, 2025, 8:01am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/35 "2025-05-27T08:01:03Z")
When compiling with `enable_sequence_parallelism`, the computation graph may appear unchanged except for moving RMSNorm to a CUDA kernel because the sequence parallelism pass only rewrites certain patterns (e.g., matmul → allreduce → rmsnorm) if the model and kernel patterns match exactly. For Qwen3, if the model’s graph does not match the expected pattern or does not use the custom RMSNorm op, the allreduce will not be split into reduce-scatter and all-gather ([PR 16155](https://github.com/vllm-project/vllm/pull/16155), [sequence\_parallelism.py](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/sequence_parallelism.py)). This means sequence parallelism is only applied if the model uses the custom RMSNorm op and matches the pass’s pattern. Otherwise, the graph remains unchanged. Would you like more detail on how to enable or debug sequence parallelism for your model? Sources: - [PR 16155](https://github.com/vllm-project/vllm/pull/16155) - [sequence\_parallelism.py](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/sequence_parallelism.py) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 28, 2025, 2:59am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/36 "2025-05-28T02:59:53Z")
how to understand sequence parallelism is not enabled for because symbolic shape can’t be split ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 28, 2025, 3:00am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/37 "2025-05-28T03:00:17Z")
Sequence parallelism is not enabled if the symbolic shape (e.g., batch size or token count) cannot be evenly divided by the tensor parallel size. The sequence parallelism pass checks if the shape is divisible; if not, it skips rewriting the graph, so reduce-scatter/all-gather are not inserted ([sequence\_parallelism.py](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/sequence_parallelism.py)). Would you like more detail on how this check is performed or how to debug it? Sources: - [sequence\_parallelism.py](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/sequence_parallelism.py) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [May 29, 2025, 12:53pm UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/38 "2025-05-29T12:53:56Z")
how fused\_add\_rms\_norm pattern work for the weight with 2-dim tensor [4,4], because the cuda kernel interface only accept the rms\_norm)weight parameter for only 1 dim? ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [May 29, 2025, 12:54pm UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/39 "2025-05-29T12:54:16Z")
It seems likely that the `fused_add_rms_norm` CUDA kernel in vLLM expects the `weight` parameter to be a 1-dimensional tensor matching the hidden size, not a 2D tensor. If you pass a 2D weight tensor like `[4, 4]`, it will raise a shape mismatch error, since the kernel and the Python wrapper both check that `weight`’s shape matches the last dimension of the input ([layernorm\_kernels.cu](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu), [layernorm.py](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/layernorm.py)). Would you like more detail on how to adapt your weight tensor or the expected input shapes? Sources: - [layernorm\_kernels.cu](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu) - [layernorm.py](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/layernorm.py) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [June 18, 2025, 8:01am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/40 "2025-06-18T08:01:18Z")
how to force prefill stage without cuda graph? ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [June 18, 2025, 8:01am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/41 "2025-06-18T08:01:40Z")
To force the prefill stage to run without CUDA graph in vLLM, set `enforce_eager=True` when creating the LLM instance, or use the `--enforce-eager` flag in the CLI. This disables CUDA graph compilation for all stages, including prefill ([optimization.md](https://github.com/vllm-project/vllm/blob/main/docs/configuration/optimization.md), [issue 3137](https://github.com/vllm-project/vllm/issues/3137)). Example: ```python from vllm import LLM llm = LLM(model="your-model", enforce_eager=True) ``` Would you like more detail on the effects or alternatives? Sources: - [optimization.md](https://github.com/vllm-project/vllm/blob/main/docs/configuration/optimization.md) - [issue 3137](https://github.com/vllm-project/vllm/issues/3137) ---
**Author:** ![Juelianqvq](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/juelianqvq/32/280_2.png) [@Juelianqvq](https://discuss.vllm.ai/u/Juelianqvq)\ **Post date:** [June 30, 2025, 1:31am UTC](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589/42 "2025-06-30T01:31:10Z")
why we need \_set\_default\_torch\_dtype when loading weights? [Previous page](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589.md?page=1) [Next page](https://discuss.vllm.ai/t/questions-on-piecewise-torch-compile-design/589.md?page=3)