[inductor] Fix FlexAttention autotuning with inlined captured buffers (#194555)

## Human Note
Instead of manually building up the list of inputs just let the interpreting lowering populate teh fields as it lowers it wil contain the full list

## Agent Report
# Result

Fixed FlexAttention max-autotune argument construction for captured buffers whose producers are inlined into Triton templates.

# Root cause

FlexAttention's forward and decoding lowerings independently constructed the autotune input list from the logical inputs plus captured score/mask buffers. That list is not necessarily the generated kernel's ABI.

For the reported relative-position bias, the logical capture is one computed bias buffer, but subgraph codegen inlines its indirect-gather producer. The generated Triton kernel therefore accepts two captured source buffers: the integer index tensor and the bias table. The launcher declares 12 tensor arguments followed by `_grid_0`, `_grid_1`, and `_grid_2`, while autotuning supplied only 11 tensor/output arguments followed by the three grid values. Positional shifting consumed two grid values as tensor arguments and produced the reported missing `_grid_2` exception.

`TritonTemplate.generate` already records the generated kernel dependencies in each choice's `input_nodes`. The fix makes FlexAttention and flex-decoding use that authoritative list for autotuning instead of reconstructing a logical list that can disagree after producer inlining. Runtime code generation and producer fusion remain unchanged.

# Changed files

- `torch/_inductor/kernel/flex/flex_attention.py`
  - Use generated forward-template input nodes for autotune tensor construction.
- `torch/_inductor/kernel/flex/flex_decoding.py`
  - Use generated decoding-template input nodes for autotune tensor construction.
- `test/inductor/test_flex_attention.py`
  - Add a max-autotune regression with an internally gathered relative-position bias, parameterized across the general Triton and Triton decoding backends.

# Validation

Hardware and build: NVIDIA GB200 (SM100), PyTorch `2.15.0a0+git1a955c7`. The reporter used an RTX 5070 Ti with PyTorch 2.13.0+cu132.

Behavioral tests:

- Before the fix, both new parameterized cases failed with `launcher() missing 1 required positional argument: '_grid_2'`.
- After the fix and simplification (`atol=5e-3, rtol=0`):

```text
python test/inductor/test_flex_attention.py -k max_autotune_with_gathered_captured_buffer
Ran 2 tests in 21.077s
OK
```

- Existing nearby max-autotune forward/backward coverage passed:

```text
python test/inductor/test_flex_attention.py TestFlexAttentionCUDA.test_max_autotune_cuda
Ran 1 test in 32.949s
OK
```

- The exact issue-style reproduction also passed with `mode="max-autotune"` (CUDA graphs enabled) and with `mode="max-autotune-no-cudagraphs"`.
- A scratch matrix covering direct, pointwise-derived, and gather-derived captured biases under forced flex decoding passed all three cases.
- A max-autotune forward/backward training smoke test completed with finite query and bias-table gradients.

Prerequisite checks:

- `git diff --check`: passed.
- Python compilation of all three changed files: passed.
- `spin quicklint`: passed.
- An earlier full `spin lint` run found no changed-Python issues but exited nonzero on pre-existing `SHELLCHECK` findings in `.ci/pytorch/test.sh`, which was not modified.

# Upstream status and uncertainty

Fetched `origin/main` at `10583bc4225`. No related fix is present in the relevant lowering, template, benchmark request, or launcher code, so the older job base did not obscure an upstream resolution and no rebase was required for the diagnosis.

The CUDA validation was on GB200 rather than the reporter's RTX 5070 Ti. The failure reproduced identically before the fix, and the fix changes Python-side argument plumbing rather than architecture-specific kernel logic. The parameterized test also permits XPU, but the tightened `5e-3` absolute bound was validated only on CUDA in this job.

<details>
<summary>Agent Worklog</summary>

# Objective

Investigate PyTorch issue #194515: max-autotune crashes while benchmarking flex-decoding when `score_mod` captures a bias produced by an indirect gather.

# Starting state

- Worktree: `pytorch/`, detached at `1a955c70e48`.
- Working tree was clean at takeover.
- Job environment: `.venv/bin/python`.
- Relevant prior notes: FlexAttention gather performance analysis and B200 FlexAttention configuration tuning; neither documents this launcher-argument failure.
- `system_prompt.md` was absent; task context is in `prime.md`.

# Issue and upstream status

- Read the full issue and all current comments with `gh issue view 194515 --comments`; there are currently no comments.
- Fetched `origin/main` at `10583bc4225`. The relevant FlexAttention lowering and template/autotune plumbing contain no upstream fix, so rebasing is not required to establish or address the bug.

# Reproduction

- Reproduced on an NVIDIA GB200 (SM100) with the job build `2.15.0a0+git1a955c7`, versus the reporter's RTX 5070 Ti and PyTorch 2.13.0+cu132.
- Exact issue-style repro: `pytorch/agent_space/repro_194515.py`.
- Command used `gpu-run auto`, `PYTHONPATH=$PWD`, and a clean `TORCHINDUCTOR_CACHE_DIR`.
- Failure matches the report: `TypeError: launcher() missing 1 required positional argument: '_grid_2'` during in-process Triton template benchmarking.
- A forced general `TRITON` FlexAttention variant also fails identically, so the invariant is shared by flex-attention and flex-decoding rather than decode-specific.

# Root cause

- Flex lowering passed the logical captured bias buffer as one autotune input.
- Template subgraph codegen inlined the gather producer. The generated kernel therefore took the gather's two source buffers (`indices`, then `table`) instead of the one materialized bias buffer.
- The generated launcher declared 12 tensor arguments plus `_grid_0`, `_grid_1`, `_grid_2`, while the benchmark call supplied only 11 tensor/output arguments before the grid tuple. Positional shifting consumed `_grid_0` and `_grid_1` as ordinary arguments and left `_grid_2` missing.
- `TritonTemplate.generate` already records the generated choice's actual kernel dependencies in `choice.input_nodes` (explicit template inputs followed by captured producer inputs). The flex lowerings instead independently reconstructed a logical input list, violating that contract when one captured buffer expands to multiple producer dependencies.

# Changes

- Flex forward and decoding autotuning now use the generated choice's dependency-expanded input nodes.
- Added a focused max-autotune regression covering forced `TRITON` and `TRITON_DECODE` with an internally gathered relative-position bias.
- Confirmed both new parameterized cases fail before the fix with the reported missing `_grid_2` error.
- The initial strict numerical check exposed expected TF32-vs-eager differences (maximum absolute error about 0.0038); the simplified test uses a tighter `atol=5e-3, rtol=0` bound.

# Validation

- New regression after the fix: 2 tests passed on GB200.
- Existing `TestFlexAttentionCUDA.test_max_autotune_cuda`: passed, including its forward/backward autotuning coverage.
- Exact issue repro passed under both `max-autotune-no-cudagraphs` and `max-autotune`.
- Scratch direct/pointwise/gather capture matrix: all three passed under forced flex decoding.
- Separate training smoke repro: max-autotune forward and backward completed with finite query/table gradient norms.
- `git diff --check` and Python compilation of changed files passed.
- `spin quicklint`: passed after simplification.
- An earlier full `spin lint` run reached pre-existing ShellCheck SC2317/SC2086 findings in `.ci/pytorch/test.sh`; no changed Python file had an issue.

# Simplification review

- Removed the single-use tensor factory and compiled-callable local from the regression.
- Reduced each source explanation to one invariant-focused line at the affected call site.
- Tightened the numerical assertion from `atol=rtol=1e-2` to `atol=5e-3, rtol=0` and re-ran both cases successfully.
- Kept one parameterized test: forced `TRITON` and `TRITON_DECODE` independently exercise the two affected paths.
- Review confirmed all generated configurations for each backend have one identical input signature, while CuTeDSL/FLASH already uses separately realized capture inputs. Backward remains unchanged because its captures are AOT-autograd intermediates and existing backward max-autotune coverage passes.
- Two requested Fable reviewers were launched twice but could not start because the local provider's GCP token helper rejects its configured `gcp-direct` mode. Two Sol reviewers were launched; one completed the actionable simplification review, while the other returned no substantive result before its provider run ended.

</details>

---
*This PR was generated by [ptq](https://github.com/drisspg/pt_job_queue) with human review.*

Pull Request resolved: https://github.com/pytorch/pytorch/pull/194555
Approved by: https://github.com/pytorchgreenlight