vllm check

core/pytorch-cpu-requirements.txt

@@ -1,3 +1,9 @@
+<<<<<<< HEAD
 --pre
 torch==2.1.0
 mpmath==1.3.0
+=======
+--pre --index-url https://download.pytorch.org/whl/nightly/cpu
+torch
+torchvision
+>>>>>>> 3982839 (gemma wip)

core/requirements.txt

@@ -5,5 +5,4 @@
 -f https://openxla.github.io/iree/pip-release-links.html
 
 -r pytorch-cpu-requirements.txt
--r torchvision-requirements.txt
 -r iree-requirements.txt

core/shark_turbine/aot/builtins/jittable.py

@@ -244,6 +244,7 @@ def flat_wrapped_f(*args):
             hooks=_Hooks(proc_trace.module_builder),
             py_attr_tracker=proc_trace.module_builder.fx_py_attr_tracker,
         )
+        print(gm.graph)
         fx_importer.import_stateless_graph(gm.graph, func_name=self.function_name)
 
         # TODO: Real debugging options

core/shark_turbine/dynamo/passes.py

@@ -53,6 +53,10 @@
     torch.ops.aten._scaled_dot_product_flash_attention.default,
 ]
 
+if hasattr(torch.ops.aten, "_scaled_dot_product_flash_attention_for_cpu"):
+    DEFAULT_DECOMPOSITIONS.append(
+        torch.ops.aten._scaled_dot_product_flash_attention_for_cpu
+    )
 
 def apply_decompositions(
     gm: torch.fx.GraphModule,