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 4, 12:30 UTC
[Bugfix][CPU] Backport PT cpp codegen indirect_assert scalar-mask fix (#40973)
Signed-off-by: Lalithnarayan C <Lalithnarayan.C@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lalithnarayan C committed
Apr 29, 2026 at 19:51 UTC
5560cac7e25b1a2c3c15506c885af4911c5611d9
1 file changed
+117
vllm/env_override.py
+117
@@ -634,3 +634,120 @@ def _patch_fxgraphcache_pickle_if_needed():
634
635
636
_patch_fxgraphcache_pickle_if_needed()
637
+
638
+# ===================================================
639
+# torch 2.11 Inductor cpp codegen indirect_assert scalar-mask fix
640
+# ===================================================
641
+# CppVecKernel.indirect_assert wraps a scalar mask with
642
+# `VecMask<...>(scalar)`, which is not a valid constructor and triggers a
643
+# C++ compile error during torch.compile of any model that does indirect
644
+# indexing inside a tail-vectorized loop (e.g. Qwen3-VL).
645
+# Failure looks like:
646
+# no matching function for call to 'VecMask<int64_t,2>::VecMask(int&)'
647
+# Upstream fix in PyTorch mainline replaces the call with
648
+# `VecMask<...>::from(scalar)`, see pytorch/pytorch#178148 (lands in 2.12).
649
+# This is a thin backport for torch >= 2.11 and < 2.12; remove once the
650
+# minimum supported torch is 2.12.
651
+
652
+
653
+def _apply_cpp_indirect_assert_patch():
654
+ """Replace CppVecKernel.indirect_assert with a fixed copy that uses
655
+ `VecMask<...>::from(scalar)` for scalar masks.
656
+
657
+ Idempotent: marks the class with `_vllm_indirect_assert_patched` after
658
+ the first apply.
659
+ """
660
+ from torch._inductor.codegen.cpp import CppVecKernel
661
+
662
+ if getattr(CppVecKernel, "_vllm_indirect_assert_patched", False):
663
+ return
664
+
665
+ from torch._inductor.codegen.cpp import CppCSEVariable, cexpr_index
666
+
667
+ def patched_indirect_assert(self, var, lower, upper, mask=None):
668
+ assert isinstance(var, CppCSEVariable)
669
+ assert var.dtype is not None
670
+ if not var.is_vec:
671
+ if isinstance(mask, CppCSEVariable) and mask.is_vec:
672
+ mask = f"({mask}).all_masked()"
673
+ return super(CppVecKernel, self).indirect_assert(var, lower, upper, mask)
674
+ lower_scalar = lower
675
+ upper_scalar = upper
676
+ if lower:
677
+ lower = f"{self._get_vec_type(var.dtype)}({lower})"
678
+ if upper:
679
+ upper = f"{self._get_vec_type(var.dtype)}({upper})"
680
+ if lower and upper:
681
+ cond = f"({lower} <= {var}) & ({var} < {upper})"
682
+ cond_print = f"{lower_scalar} <= {var} < {upper_scalar}"
683
+ elif lower:
684
+ cond = f"{lower} <= {var}"
685
+ cond_print = f"{lower_scalar} <= {var}"
686
+ else:
687
+ assert upper
688
+ cond = f"{var} < {upper}"
689
+ cond_print = f"{var} < {upper_scalar}"
690
+ cond = f"{self._get_mask_type(var.dtype)}({cond})"
691
+ if mask:
692
+ if not mask.is_vec:
693
+ # Backport of pytorch/pytorch#178148 -- use ::from for
694
+ # scalar masks so g++ picks the correct overload.
695
+ mask = f"{self._get_mask_type(var.dtype)}::from({mask})"
696
+ cond = f"({cond}) | ~({mask})"
697
+ if self.tail_size:
698
+ cond = (
699
+ f"{self._get_mask_type(var.dtype)}::set("
700
+ f"{self._get_mask_type(var.dtype)}::from(1)"
701
+ f", ({cond}), {cexpr_index(self.tail_size)})"
702
+ )
703
+ cond = f"({cond}).all_masked()"
704
+ return f'{self.assert_function}({cond}, "index out of bounds: {cond_print}")'
705
+
706
+ CppVecKernel.indirect_assert = patched_indirect_assert
707
+ CppVecKernel._vllm_indirect_assert_patched = True # type: ignore[attr-defined]
708
+
709
+
710
+def _patch_cpp_indirect_assert_if_needed():
711
+ """Apply cpp codegen indirect_assert backport when on torch 2.11.x.
712
+
713
+ Defers application until torch._inductor.codegen.cpp is naturally
714
+ imported by Inductor. Importing it eagerly during vllm.__init__ pulls
715
+ in torch._inductor.scheduler, whose top-level
716
+ `import torch._inductor.async_compile` can fail with
717
+ `ModuleNotFoundError: import of torch._inductor.async_compile halted;
718
+ None in sys.modules` depending on the import order on the runner
719
+ (observed in vLLM CPU CI).
720
+ """
721
+ if not is_torch_equal_or_newer("2.11.0") or is_torch_equal_or_newer("2.12.0.dev"):
722
+ return
723
+
724
+ import sys
725
+
726
+ target_name = "torch._inductor.codegen.cpp"
727
+ if target_name in sys.modules:
728
+ _apply_cpp_indirect_assert_patch()
729
+ return
730
+
731
+ import importlib.abc
732
+
733
+ class _CppCodegenPatchFinder(importlib.abc.MetaPathFinder):
734
+ def find_spec(self, fullname, path, target=None):
735
+ if fullname != target_name:
736
+ return None
737
+ sys.meta_path.remove(self)
738
+ spec = importlib.util.find_spec(fullname)
739
+ if spec is None or spec.loader is None:
740
+ return None
741
+ original_exec = spec.loader.exec_module
742
+
743
+ def _exec_then_patch(module):
744
+ original_exec(module)
745
+ _apply_cpp_indirect_assert_patch()
746
+
747
+ spec.loader.exec_module = _exec_then_patch # type: ignore[method-assign]
748
+ return spec
749
+
750
+ sys.meta_path.insert(0, _CppCodegenPatchFinder())
751
+
752
+
753
+_patch_cpp_indirect_assert_if_needed()