feat: recycle K/V in cuda graph

src/kernl/implementations/cuda_graph.py

@@ -25,10 +25,29 @@ def cuda_graphs_wrapper(
     pool: (int, int) = torch.cuda.graph_pool_handle(),
 ):
     """
-    From torchdynamo
+    This function is a wrapper for the model to be used with cuda graphs.
+    It is used to create a cuda graph and to execute it.
+
+    @param model: model to be wrapped
+    @param inputs: original inputs of the model
+    @param copy_outputs: if True, the outputs will be cloned so it can be mutated, etc
+    @param pool: cuda graph pool handle, to share pool between graphs
+    @return: a callable to run the graph
     """
+
     assert isinstance(inputs, (list, tuple)), f"inputs is of type {type(inputs)} instead of list"
-    static_inputs = [torch.zeros_like(x) for x in inputs]
+    static_inputs = list()
+
+    for i in inputs:
+        if getattr(i, "reuse_counter", 0) > 0:
+            i.reuse_counter += 1
+            static_inputs.append(i)
+        else:
+            t = torch.empty_like(i)
+            i.reuse_counter = 1
+            t.reuse_counter = 0
+            static_inputs.append(t)
+
     # required warmup, not just for perf but for correctness
     torch.cuda.synchronize()
     stream = torch.cuda.Stream()
@@ -52,8 +71,13 @@ def cuda_graphs_wrapper(
     def run(*new_inputs):
         assert isinstance(new_inputs, (list, tuple)), f"inputs is of type {type(new_inputs)} instead of list"
         assert len(static_inputs) == len(new_inputs), f"{len(static_inputs)} == {len(new_inputs)}"
-        for dst, src in zip(static_inputs, new_inputs):
-            dst.copy_(src)  # cuda graph can only read data from the same address
+        # cuda graph can only read data from the same address
+        for src, dst in zip(new_inputs, static_inputs):
+            # some tensors are reused from call to call,
+            # if this is the second times or more, they already contain the right data
+            if dst.reuse_counter <= 1:
+                dst.copy_(src)
+
         graph.replay()
         if copy_outputs:
             return [x.clone() for x in static_outputs]

test/test_torchdynamo.py

@@ -132,3 +132,15 @@ def test_t5():
             do_sample=False,
         )
         assert "La maison est merveilleuse." in tokenizer.batch_decode(output_sequences, skip_special_tokens=True)[0]
+
+    task2 = "translate English to French: What are you doing here sir, may I help you or call someone?"
+    inputs = tokenizer(task2, return_tensors="pt", padding=True).to("cuda")
+    with torch.inference_mode(), torch.autocast(dtype=torch.float16, cache_enabled=True, device_type="cuda"):
+        output_sequences = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            min_length=1,
+            max_length=22,
+            do_sample=False,
+        )
+        assert "Que faites-vous ici, Monsieur" in tokenizer.batch_decode(output_sequences, skip_special_tokens=True)[0]