A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 11:45 UTC
[torch.compile] Add patch for fullgraph compilation (#42686)
Signed-off-by: Luka Govedič <luka.govedic@gmail.com>
Luka Govedič committed
May 17, 2026 at 15:49 UTC
966903eb93a053a908fbf8b931fcebfb28c4741a
1 file changed
+62
vllm/env_override.py
+62
@@ -758,3 +758,65 @@ def _patch_cpp_indirect_assert_if_needed():
758
759
760
_patch_cpp_indirect_assert_if_needed()
761
+
762
+
763
+def _patch_should_realize_on_reuse():
764
+ """
765
+ This patches the materialization heuristic in Inductor compilation.
766
+ Patched in torch 2.12: https://github.com/pytorch/pytorch/pull/176994
767
+ vLLM issue: https://github.com/vllm-project/vllm/issues/27828
768
+
769
+ Without this patch, Inductor inlined a chain of computations when compiling
770
+ the whole model graph at once. For example, the residual in fused_add_rms_norm
771
+ would get recomputed every time, and this would cascade through the model.
772
+
773
+ Note that pytorch/pytorch#176994 updated multiple callsites,
774
+ this is a simplified version of the fix.
775
+ """
776
+
777
+ # Only apply for 2.11, fix included in 2.12
778
+ if not is_torch_equal("2.11.0"):
779
+ return
780
+
781
+ def should_realize_on_reuse_patched(self, users: int) -> bool:
782
+ """
783
+ A heuristic to decide if we should realize a tensor
784
+ that is used multiple times.
785
+ """
786
+ from torch._inductor import config
787
+ from torch._inductor.ir import Pointwise, Reduction, is_cpu
788
+ from torch._inductor.virtualized import V
789
+
790
+ if users > 1 and isinstance(self.data, (Pointwise, Reduction)):
791
+ if is_cpu(self.data):
792
+ # Heuristic for realizing reused result of heavy ops on cpu
793
+ opcount = self.data.inner_fn_opcount()
794
+ heavy_ops = ["exp", "sigmoid"] # a list of heavy ops
795
+ if any(x in opcount.used_ops for x in heavy_ops):
796
+ return True
797
+ if self.has_large_inner_fn():
798
+ return True
799
+ # Size-aware cost model comparing total memory traffic:
800
+ # Inline: total_read_bytes * users
801
+ # Materialize: total_read_bytes + output_bytes * (1 + users)
802
+ # This naturally handles broadcast reads (small buffers are
803
+ # cheap to re-read) and dtype promotions (fp32 outputs cost
804
+ # more to write than bf16 inputs cost to read).
805
+ total_read_bytes = sum(
806
+ V.graph.get_dep_size_hint(dep) for dep in self.get_reads()
807
+ )
808
+ output_bytes = (
809
+ V.graph.sizevars.optimization_hint(self.data.get_numel(), fallback=0)
810
+ * self.data.dtype.itemsize
811
+ )
812
+ if total_read_bytes > 0 and output_bytes > 0:
813
+ return total_read_bytes * (users - 1) >= output_bytes * (1 + users)
814
+ return self.num_reads() > config.realize_reads_threshold
815
+ return False
816
+
817
+ from torch._inductor.ir import StorageBox
818
+
819
+ StorageBox.should_realize_on_reuse = should_realize_on_reuse_patched
820
+
821
+
822
+_patch_should_realize_on_reuse()