Check memory allocation with torchdynamo cudagraph about kernl HOT 1 CLOSED

Experiment:

Without any warmup, we launch inference with the following shapes : (1, 64), (8, 256), (8, 256), (1, 64), (8, 256), (8, 64)
We want to see how the model works with these shape changes. For that we compute the elapsed time (in ms) for each shape.
First with torch:

baseline time (1, 64): 690.441873
baseline time (8, 256): 7.029635
baseline time (8, 256): 6.178398
baseline time (1, 64): 6.019235999999999
baseline time (8, 256): 6.248811
baseline time (8, 64): 6.215109

We see that there is a warmup phase and then every run run similarly.

With torchdynamo :

dynamo time (1, 64): 5574.487258
dynamo time (8, 256): 5407.364699
dynamo time (8, 256): 5.446753999999999
dynamo time (1, 64): 5.2307109999999994
dynamo time (8, 256): 5.044296999999999
dynamo time (8, 64): 5459.834921

For every new shape there is a 5s computation overhead, but for existing shape the running time is a bit faster than torch

With torchdynamo and cudagraph:

dynamo_cuda_graphs time (1, 64): 5566.876034999999
dynamo_cuda_graphs time (8, 256): 5604.906061
dynamo_cuda_graphs time (8, 256): 0.20713399999999998
dynamo_cuda_graphs time (1, 64): 0.265165
dynamo_cuda_graphs time (8, 256): 0.21359399999999998
dynamo_cuda_graphs time (8, 64): 5579.558674

We have the same behavior than without cuda graph but with very faster runtime for existing shapes

Torchdynamo with our triton kernel fusion:

dynamo_optimized time (1, 64): 33224.978761
dynamo_optimized time (8, 256): 25587.014703999997
dynamo_optimized time (8, 256): 12.889158
dynamo_optimized time (1, 64): 12.483979999999999
dynamo_optimized time (8, 256): 13.551530999999999
dynamo_optimized time (8, 64): 25261.002185999998

With the kernel optimization, startup time is very long ~33s, and new shapes takes ~26s, but already encountered shapes are faster, but still slower than pytorch and son on

Torchdynamo with our triton kernel fusion and cudagraph:

dynamo_optimized_cuda_graphs time (1, 64): 6147.142248
dynamo_optimized_cuda_graphs time (8, 256): 6309.444265
dynamo_optimized_cuda_graphs time (8, 256): 0.15262299999999998
dynamo_optimized_cuda_graphs time (1, 64): 0.18226299999999998
dynamo_optimized_cuda_graphs time (8, 256): 0.174434
dynamo_optimized_cuda_graphs time (8, 64): 6389.431353999999

There is a 6s overhead with new shapes, but it's the fastest on already-encountered shape.

Actually, there is more than 6s overhead, because it uses the previous graph from the previous run. If we swap in order dynamo_optimized and dynamo_optimized_cuda_graphs, the time with new shapes exchange.

Overhead with cudagraph doesn't seem too important compared to vanilla torchdynamo, but our model with our fused kernel seems to be very long to compute (~30s), for every new shape on input.

from kernl.