DISABLED test_max_autotune_with_captured_cuda (__main__.TestFlexDecodingCUDA)

Platforms: linux, slow

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/flakytest?name=test_max_autotune_with_captured_cuda&suite=TestFlexDecodingCUDA&limit=100) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/88730879525).

Over the past 6 hours, it has been determined flaky in 8 workflow(s) with 16 failures and 8 successes.

**Debugging instructions (after clicking on the recent samples link):**
DO NOT ASSUME THINGS ARE OKAY IF THE CI IS GREEN. We now shield flaky tests from developers so CI will thus be green but it will be harder to parse the logs.
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Click on the Test step of the job so that it is expanded. Otherwise, the grepping will not work.
3. Grep for `test_max_autotune_with_captured_cuda`
4. There should be several instances run (as flaky tests are rerun in CI) from which you can study the logs.

Sample error message

```
Traceback (most recent call last):
File "/__w/pytorch/pytorch/test/inductor/test_flex_decoding.py", line 1889, in test_max_autotune_with_captured
self.run_test(bias_mod, device=device)
File "/__w/pytorch/pytorch/test/inductor/test_flex_decoding.py", line 468, in run_test
compiled_out, compiled_lse = compiled_sdpa(q, k, v, return_lse=True)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1205, in compile_wrapper
raise e.remove_dynamo_frames() from None # see TORCHDYNAMO_VERBOSE=1
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1186, in _compile_fx_inner
raise InductorError(e, currentframe()).with_traceback(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1163, in _compile_fx_inner
mb_compiled_graph = fx_codegen_and_compile(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1983, in fx_codegen_and_compile
return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1642, in codegen_and_compile
graph.run(*example_inputs)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1132, in run
return super().run(*args)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/fx/interpreter.py", line 197, in run
self.env[node] = self.run_node(node)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/graph.py", line 2059, in run_node
result = super().run_node(n)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/fx/interpreter.py", line 294, in run_node
return getattr(self, n.op)(n.target, args, kwargs)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1564, in call_function
raise LoweringException(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1541, in call_function
out = lowerings[target](*args, **kwargs)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/lowering.py", line 524, in wrapped
out = decomp_fn(*args, **kwargs)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/kernel/flex/flex_attention.py", line 293, in flex_attention
return create_flex_decoding_kernel(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/kernel/flex/flex_decoding.py", line 464, in create_flex_decoding_kernel
buf_ACC, _ = autotune_select_algorithm(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 6107, in autotune_select_algorithm
return cache(*args, **kwargs)
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 4083, in __call__
precompile_fn = self.make_precompile_fn(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 4674, in make_precompile_fn
timings = self.lookup(
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 431, in lookup
precision = torch.get_float32_matmul_precision()
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/__init__.py", line 2073, in get_float32_matmul_precision
return _C._get_float32_matmul_precision()
torch._inductor.exc.InductorError: LoweringException: RuntimeError: PyTorch is checking the matmul precision without a specific backend name,Current status indicate that you have used mix of the legacy and new APIs to set the matmul precision. We suggest only using the new API for matmul precision. See also: https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices
target: flex_attention
args[0]: TensorBox(StorageBox(
InputBuffer(name='primals_1', layout=FixedLayout('cuda:0', torch.float16, size=[4, 16, 1, 64], stride=[1024, 64, 64, 1]))
))
args[1]: TensorBox(StorageBox(
InputBuffer(name='primals_2', layout=FixedLayout('cuda:0', torch.float16, size=[4, 8, 2048, 64], stride=[1048576, 131072, 64, 1]))
))
args[2]: TensorBox(StorageBox(
InputBuffer(name='primals_3', layout=FixedLayout('cuda:0', torch.float16, size=[4, 8, 2048, 64], stride=[1048576, 131072, 64, 1]))
))
args[3]: Subgraph(name='sdpa_score0', graph_module=(), inductor_config_patches=None, graph=None)
args[4]: (1, 1, TensorBox(StorageBox(
ComputedBuffer(name='buf0', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _ = index
tmp0 = ops.constant(1, torch.int32)
return tmp0
,
ranges=[1, 1, 1],
origin_node=full,
origins=OrderedSet([full]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2098, in _create_empty_block_mask,
kv_num_blocks=torch.ones([1, 1, 1], dtype=torch.int32, device=device),,
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), TensorBox(StorageBox(
ComputedBuffer(name='buf1', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _, _ = index
tmp0 = ops.constant(0, torch.int32)
return tmp0
,
ranges=[1, 1, 1, 1],
origin_node=full_default,
origins=OrderedSet([full_default]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2099, in _create_empty_block_mask,
kv_indices=torch.zeros([1, 1, 1, 1], dtype=torch.int32, device=device),,
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), None, None, TensorBox(StorageBox(
ComputedBuffer(name='buf10', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _ = index
tmp0 = ops.load(buf5, 0)
tmp1 = ops.to_dtype(tmp0, torch.int64, src_dtype=torch.int32, use_compute_types=True)
tmp2 = ops.to_dtype(tmp1, torch.int32, src_dtype=torch.int64, use_compute_types=True)
return tmp2
,
ranges=[1, 1, 1],
origin_node=convert_element_type_1,
origins=OrderedSet([convert_element_type_1, sum_1]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2097, in _create_empty_block_mask,
return BlockMask.from_kv_blocks(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 1118, in from_kv_blocks,
q_num_blocks, q_indices = _transpose_ordered(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 503, in _transpose_ordered,
return _dense_to_ordered(dense.transpose(-2, -1)),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 490, in _dense_to_ordered,
num_blocks_in_row.to(torch.int32, memory_format=torch.contiguous_format),,
,
},
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2097, in _create_empty_block_mask,
return BlockMask.from_kv_blocks(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 1118, in from_kv_blocks,
q_num_blocks, q_indices = _transpose_ordered(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 503, in _transpose_ordered,
return _dense_to_ordered(dense.transpose(-2, -1)),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 486, in _dense_to_ordered,
num_blocks_in_row = dense_mask.sum(dim=-1),
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), TensorBox(StorageBox(
ComputedBuffer(name='buf9', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _, _ = index
tmp0 = ops.load(buf8, 0)
tmp1 = ops.to_dtype(tmp0, torch.int32, src_dtype=torch.int64, use_compute_types=True)
return tmp1
,
ranges=[1, 1, 1, 1],
origin_node=convert_element_type,
origins=OrderedSet([convert_element_type]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2097, in _create_empty_block_mask,
return BlockMask.from_kv_blocks(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 1118, in from_kv_blocks,
q_num_blocks, q_indices = _transpose_ordered(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 503, in _transpose_ordered,
return _dense_to_ordered(dense.transpose(-2, -1)),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 491, in _dense_to_ordered,
col_indices.to(torch.int32, memory_format=torch.contiguous_format),,
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), None, None, None, None, None, None, 1073741824, 1073741824, Subgraph(name='sdpa_mask0', graph_module=(), inductor_config_patches=None, graph=None))
args[5]: 0.125
args[6]: {'BACKEND': 'AUTO', 'PRESCALE_QK': False, 'ROWS_GUARANTEED_SAFE': False, 'BLOCKS_ARE_CONTIGUOUS': False, 'WRITE_DQ': True, 'OUTPUT_LOGSUMEXP': True, 'OUTPUT_MAX': False}
args[7]: (TensorBox(StorageBox(
InputBuffer(name='primals_7', layout=FixedLayout('cuda:0', torch.float32, size=[2048], stride=[1]))
)), TensorBox(StorageBox(
InputBuffer(name='primals_6', layout=FixedLayout('cuda:0', torch.float32, size=[1], stride=[1]))
)), TensorBox(StorageBox(
InputBuffer(name='primals_4', layout=FixedLayout('cuda:0', torch.float32, size=[4], stride=[1]))
)), TensorBox(StorageBox(
InputBuffer(name='primals_5', layout=FixedLayout('cuda:0', torch.float32, size=[16], stride=[1]))
)))
args[8]: ()RuntimeError: PyTorch is checking the matmul precision without a specific backend name,Current status indicate that you have used mix of the legacy and new APIs to set the matmul precision. We suggest only using the new API for matmul precision. See also: https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices
target: flex_attention
args[0]: TensorBox(StorageBox(
InputBuffer(name='primals_1', layout=FixedLayout('cuda:0', torch.float16, size=[4, 16, 1, 64], stride=[1024, 64, 64, 1]))
))
args[1]: TensorBox(StorageBox(
InputBuffer(name='primals_2', layout=FixedLayout('cuda:0', torch.float16, size=[4, 8, 2048, 64], stride=[1048576, 131072, 64, 1]))
))
args[2]: TensorBox(StorageBox(
InputBuffer(name='primals_3', layout=FixedLayout('cuda:0', torch.float16, size=[4, 8, 2048, 64], stride=[1048576, 131072, 64, 1]))
))
args[3]: Subgraph(name='sdpa_score0', graph_module=(), inductor_config_patches=None, graph=None)
args[4]: (1, 1, TensorBox(StorageBox(
ComputedBuffer(name='buf0', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _ = index
tmp0 = ops.constant(1, torch.int32)
return tmp0
,
ranges=[1, 1, 1],
origin_node=full,
origins=OrderedSet([full]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2098, in _create_empty_block_mask,
kv_num_blocks=torch.ones([1, 1, 1], dtype=torch.int32, device=device),,
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), TensorBox(StorageBox(
ComputedBuffer(name='buf1', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _, _ = index
tmp0 = ops.constant(0, torch.int32)
return tmp0
,
ranges=[1, 1, 1, 1],
origin_node=full_default,
origins=OrderedSet([full_default]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2099, in _create_empty_block_mask,
kv_indices=torch.zeros([1, 1, 1, 1], dtype=torch.int32, device=device),,
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), None, None, TensorBox(StorageBox(
ComputedBuffer(name='buf10', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _ = index
tmp0 = ops.load(buf5, 0)
tmp1 = ops.to_dtype(tmp0, torch.int64, src_dtype=torch.int32, use_compute_types=True)
tmp2 = ops.to_dtype(tmp1, torch.int32, src_dtype=torch.int64, use_compute_types=True)
return tmp2
,
ranges=[1, 1, 1],
origin_node=convert_element_type_1,
origins=OrderedSet([convert_element_type_1, sum_1]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2097, in _create_empty_block_mask,
return BlockMask.from_kv_blocks(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 1118, in from_kv_blocks,
q_num_blocks, q_indices = _transpose_ordered(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 503, in _transpose_ordered,
return _dense_to_ordered(dense.transpose(-2, -1)),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 490, in _dense_to_ordered,
num_blocks_in_row.to(torch.int32, memory_format=torch.contiguous_format),,
,
},
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2097, in _create_empty_block_mask,
return BlockMask.from_kv_blocks(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 1118, in from_kv_blocks,
q_num_blocks, q_indices = _transpose_ordered(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 503, in _transpose_ordered,
return _dense_to_ordered(dense.transpose(-2, -1)),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 486, in _dense_to_ordered,
num_blocks_in_row = dense_mask.sum(dim=-1),
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), TensorBox(StorageBox(
ComputedBuffer(name='buf9', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]), data=Pointwise(
'cuda',
torch.int32,
def inner_fn(index):
_, _, _, _ = index
tmp0 = ops.load(buf8, 0)
tmp1 = ops.to_dtype(tmp0, torch.int32, src_dtype=torch.int64, use_compute_types=True)
return tmp1
,
ranges=[1, 1, 1, 1],
origin_node=convert_element_type,
origins=OrderedSet([convert_element_type]),
stack_traces = {,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2484, in flex_attention,
block_mask = _create_empty_block_mask(query, key),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2097, in _create_empty_block_mask,
return BlockMask.from_kv_blocks(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 1118, in from_kv_blocks,
q_num_blocks, q_indices = _transpose_ordered(,
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 503, in _transpose_ordered,
return _dense_to_ordered(dense.transpose(-2, -1)),
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 491, in _dense_to_ordered,
col_indices.to(torch.int32, memory_format=torch.contiguous_format),,
,
}
), _split_size=None, _original_inner_fn=None, _original_ranges=None, _original_reduction_ranges=None)
)), None, None, None, None, None, None, 1073741824, 1073741824, Subgraph(name='sdpa_mask0', graph_module=(), inductor_config_patches=None, graph=None))
args[5]: 0.125
args[6]: {'BACKEND': 'AUTO', 'PRESCALE_QK': False, 'ROWS_GUARANTEED_SAFE': False, 'BLOCKS_ARE_CONTIGUOUS': False, 'WRITE_DQ': True, 'OUTPUT_LOGSUMEXP': True, 'OUTPUT_MAX': False}
args[7]: (TensorBox(StorageBox(
InputBuffer(name='primals_7', layout=FixedLayout('cuda:0', torch.float32, size=[2048], stride=[1]))
)), TensorBox(StorageBox(
InputBuffer(name='primals_6', layout=FixedLayout('cuda:0', torch.float32, size=[1], stride=[1]))
)), TensorBox(StorageBox(
InputBuffer(name='primals_4', layout=FixedLayout('cuda:0', torch.float32, size=[4], stride=[1]))
)), TensorBox(StorageBox(
InputBuffer(name='primals_5', layout=FixedLayout('cuda:0', torch.float32, size=[16], stride=[1]))
)))
args[8]: ()
Found from :
File "/opt/python-3.10-venv/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2577, in flex_attention
out, lse, max_scores = flex_attention_hop(

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

To execute this test, run the following from the base repo dir:
python test/inductor/test_flex_decoding.py TestFlexDecodingCUDA.test_max_autotune_with_captured_cuda

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0
```

Test file path: `inductor/test_flex_decoding.py`

For all disabled tests (by GitHub issue), see https://hud.pytorch.org/disabled.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

Start with test/inductor/test_flex_decoding.py, especially test_max_autotune_with_captured at line 1889 and run_test at line 468. Review the linked flaky-test samples and workflow logs, expanding the Test step and grepping for test_max_autotune_with_captured_cuda. Done means the failure is understood and the disabled test can run reliably in CI.

Written by the indexing model from the issue text.