[Benchmark nvFusion Symbol]: kernel time number is different between online fd execution and the saved script execution

🐛 Bug

The perf numbers measured in the following 2 ways are different:

1. Gets the nvfusion symbol from the trace and measures

fd=nvfusion_symbols._call_ctx["nvFusion0"].last_used
kernel_time, _= kernel_time_with_torch_profiler(fd.execute, inputs)

2. gets the repro code from fd and runs the script

To Reproduce

The script of the 1st way:

import torch

from litgpt import Config, GPT
from functools import partial
from torch.testing import make_tensor

from torch.profiler import profile, ProfilerActivity
from time import sleep
def kernel_time_with_torch_profiler(callable, inputs):
    # Simple warm-up
    for i in range(5):
        callable(inputs)

    time = 0.0
    with profile(activities=[ProfilerActivity.CUDA]) as prof:
        callable(inputs)
        # sleep(5)

    for evt in prof.events():
        time += evt.device_time
    # Also returning profile object
    # if anyone wants to investigate, (not required)
    return time, prof


cfg = Config.from_name("Llama-3-8B")

# Uses a reduced configuration
cfg.n_layer = 3
cfg.block_size = 1024
batch_dim = 4

torch.set_default_dtype(torch.bfloat16)
make = partial(make_tensor, low=0, high=255, device='cuda', dtype=torch.int64)

with torch.device('cuda'):
    model = GPT(cfg)
    shape = (batch_dim, cfg.block_size)
    x = make(shape)


from thunder.dynamo import thunderfx
cm = thunderfx(model)
cm(x)

trace = cm.last_traces[0]
bsyms = [bsym for bsym in trace.bound_symbols if "nvFusion" in bsym.sym.name]

nvfuser_callable = bsyms[0]._call_ctx[bsyms[0].sym.name]
fd = nvfuser_callable.last_used
# print(fd)
inputs = [
    torch.testing.make_tensor((4, 1024, 4096), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((4096,), dtype=torch.bfloat16, device='cuda:0'),
]

torch.cuda.synchronize()

kernel_time, prof = kernel_time_with_torch_profiler(fd.execute, inputs)
print(kernel_time)

It outputs:

root@0fc4a86a2d36:/wayan/lightning-thunder# python repro/timingtest/nvfusion_kerneltime.py
26.367999999999995

The script of the second way

# CUDA devices:
#  0: NVIDIA RTX 6000 Ada Generation
#  1: NVIDIA RTX 6000 Ada Generation
# torch version: 2.7.0a0+ecf3bae40a.nvInternal
# cuda version: 12.8
# nvfuser version: 0.2.26+git37bc4fa
import torch
from nvfuser import FusionDefinition, DataType

def nvfuser_fusion_id0(fd : FusionDefinition) -> None :
    T0 = fd.define_tensor(shape=[4, 1024, 4096], contiguity=[True, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T1 = fd.define_tensor(shape=[4096], contiguity=[True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[0])
    T2 = fd.ops.cast(T0, dtype=DataType.Float)
    T3 = fd.ops.cast(T1, dtype=DataType.Float)
    T8 = fd.ops.broadcast_in_dim(T3, shape=[4, 1024, 4096], broadcast_dims=[2])
    T9 = fd.ops.mul(T2, T2)
    T10 = fd.ops.sum(T9, dims=[2], keepdim=False, dtype=DataType.Null)
    T15 = fd.ops.broadcast_in_dim(T10, shape=[4, 1024, 1], broadcast_dims=[0, 1])
    S16 = fd.define_scalar(4096.00, dtype=DataType.Double)
    S17 = fd.ops.reciprocal(S16)
    T18 = fd.ops.mul(T15, S17)
    S19 = fd.define_scalar(1.00000e-05, dtype=DataType.Double)
    T20 = fd.ops.add(T18, S19)
    T21 = fd.ops.rsqrt(T20)
    T26 = fd.ops.broadcast_in_dim(T21, shape=[4, 1024, 4096], broadcast_dims=[0, 1, 2])
    T27 = fd.ops.mul(T2, T26)
    T28 = fd.ops.mul(T27, T8)
    T29 = fd.ops.cast(T28, dtype=DataType.BFloat16)
    fd.add_output(T21)
    fd.add_output(T29)

with FusionDefinition() as fd:
    nvfuser_fusion_id0(fd)

inputs = [
    torch.testing.make_tensor((4, 1024, 4096), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((4096,), dtype=torch.bfloat16, device='cuda:0'),
]


from torch.profiler import profile, ProfilerActivity
from time import sleep
def kernel_time_with_torch_profiler(callable, inputs):
    # Simple warm-up
    for i in range(5):
        callable(inputs)

    time = 0.0
    with profile(activities=[ProfilerActivity.CUDA]) as prof:
        callable(inputs)
        # sleep(5)

    for evt in prof.events():
        time += evt.device_time
    # Also returning profile object
    # if anyone wants to investigate, (not required)
    return time, prof

kernel_time, prof = kernel_time_with_torch_profiler(fd.execute, inputs)
print(kernel_time)

it outputs:

root@0fc4a86a2d36:/wayan/lightning-thunder# python repro/timingtest/nvfuserfd_kernel.py
17.183999999999287

And by adding the sleep(5), the timing number doesn't change much

ENV:

CUDA Version: 12.8
NVIDIA RTX 6000 Ada Generation

Expected behavior

The timing numbers are similar

cc @crcrpar @tfogal

Start with the two repro scripts, repro/timingtest/nvfusion_kerneltime.py and repro/timingtest/nvfuserfd_kernel.py, and compare the profiler events produced by fd.execute in each path. Investigate the difference between the traced nvFusion symbol's last_used execution and the reconstructed FusionDefinition execution. Done means explaining the discrepancy and making the measured kernel times comparable under the stated CUDA environment.