Error running 06-fused-attention.py

Hi, I just compiled the newest master branch and tried to run 06-fused-attention.py in tutorial by uncommenting the last line

bench_flash_attention.run(save_path='.', print_data=True)

However, I got these errors when running it:

Traceback (most recent call last):
  File "06-fused-attention.py", line 359, in <module>
    bench_flash_attention.run(save_path='.', print_data=True)
  File "/work_1a/nyn/pkgs/triton/python/triton/testing.py", line 308, in run
    self._run(bench, save_path, show_plots, print_data)
  File "/work_1a/nyn/pkgs/triton/python/triton/testing.py", line 263, in _run
    ret = self.fn(**x_args, **{bench.line_arg: y}, **bench.args)
  File "06-fused-attention.py", line 343, in bench_flash_attention
    ms = triton.testing.do_bench(fn, percentiles=None, warmup=warmup, rep=rep)
  File "/work_1a/nyn/pkgs/triton/python/triton/testing.py", line 137, in do_bench
    fn()
  File "06-fused-attention.py", line 342, in <lambda>
    fn = lambda: o.backward(do, retain_graph=True)
  File "/data/home/nyn/miniconda3/envs/efm/lib/python3.8/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/data/home/nyn/miniconda3/envs/efm/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/data/home/nyn/miniconda3/envs/efm/lib/python3.8/site-packages/torch/autograd/function.py", line 253, in apply
    return user_fn(self, *args)
  File "06-fused-attention.py", line 253, in backward
    _bwd_kernel[(ctx.grid[1],)](
  File "/work_1a/nyn/pkgs/triton/python/triton/code_gen.py", line 999, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/work_1a/nyn/pkgs/triton/python/triton/code_gen.py", line 988, in __call__
    return _triton.runtime.launch(wargs, self.fn.do_not_specialize, cache_key, self.fn.arg_names,
  File "/work_1a/nyn/pkgs/triton/python/triton/code_gen.py", line 956, in add_to_cache
    return self.fn._warmup(key, arg_types=arg_types, device=device_idx, attributes=attributes, constants=constants, num_warps=num_warps, num_stages=num_stages,
  File "/work_1a/nyn/pkgs/triton/python/triton/code_gen.py", line 1285, in _warmup
    binary = self._compile(**compile)
  File "/work_1a/nyn/pkgs/triton/python/triton/code_gen.py", line 1320, in _compile
    name, asm, shared_mem = _triton.code_gen.compile_ttir(backend, generator.module, device, num_warps, num_stages, extern_libs)
RuntimeError: Internal Triton PTX codegen error: 
ptxas /tmp/fileOTTBHd, line 4004; fatal   : Parsing error near ':': syntax error
ptxas fatal   : Ptx assembly aborted due to errors

Looks to me that something is wrong when doing the backward pass. Any idea what is the issue?