Update on "Changes needed for transformer inference (includes Elias's cudagraph"

Chillee · Chillee · commit 70edd0a2a6bd · 2023-08-16T14:47:35.000-07:00
mutation for params)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov anijain2305

[ghstack-poisoned]
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -149,6 +149,7 @@ def convolution_backward(
     )
     return (grad_inp, grad_weight, grad_bias)
 
+
 @register_decomposition(aten.slice_scatter)
 def slice_scatter(self, src, dim=0, start=None, end=None, step=1):
     if start is None or end is None:
@@ -157,14 +158,16 @@ def slice_scatter(self, src, dim=0, start=None, end=None, step=1):
         return src
     return NotImplemented
 
+
 @register_decomposition(aten.slice)
 def slice(self, dim=0, start=None, end=None, step=1):
     if start is None or end is None:
         return NotImplemented
-    if start == 0 and end >= 2**63 -1 and step == 1:
+    if start == 0 and end >= 2**63 - 1 and step == 1:
         return self
     return NotImplemented
 
+
 @register_decomposition([aten.log2])
 def log2(x):
     return torch.log(x) * (1.0 / math.log(2.0))
@@ -210,8 +213,10 @@ def bmm(self, batch2):
 
 @register_decomposition([aten.mm])
 def mm(self, input2):
-    if config.decompose_mm_to_mv:
-        if self.shape[0] == 1 and input2.shape[1] == 1:
+    # Our matrix vector multiplies only achieve peak bandwidth with coordinate descent tuning.
+    # todo: Look into why and fix it (hopefully)
+    if config.coordinate_descent_tuning:
+        if self.shape[0] == 1 or input2.shape[1] == 1:
             return (self.unsqueeze(2) * input2.unsqueeze(0)).sum(dim=1)
     if self.device == "cpu":
         if (
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
@@ -2,10 +2,10 @@
 import itertools
 import logging
 import operator
+from collections import namedtuple
 from typing import List, Optional, Union
 
 from sympy import Expr
-from collections import namedtuple
 
 import torch
 import torch._inductor as inductor
@@ -554,7 +554,6 @@ def is_valid_splitwithsizes_cat(match):
     return True
 
 
-
 def reinplace_scatters(graph):
     """
     Reinplaces scatter operations in easy cases where the node being mutated
@@ -569,7 +568,7 @@ def reinplace_scatters(graph):
         if node.target == aten.copy_.default:
             copy_nodes[(node.args[0], node.args[1])] = node
             mutated_inputs.add(node.args[0])
-        elif node.op == 'output':
+        elif node.op == "output":
             pass
         else:
             break
@@ -580,9 +579,9 @@ def reinplace_scatters(graph):
         aten.index_put.default: InplaceableOp(aten.index_put_.default, 0),
     }
     for node in graph.nodes:
-        if (inplaceable_op := inplaceable_ops.get(node.target, False)):
+        if inplaceable_op := inplaceable_ops.get(node.target, False):
             mutated_arg = node.args[inplaceable_op.mutated_arg]
-            if mutated_arg.op == 'placeholder':
+            if mutated_arg.op == "placeholder":
                 if not (copy_node := copy_nodes.get((mutated_arg, node), False)):
                     continue
 
