feat(trace): embed runnable init() in every TraceTemplate (#3221)

## Summary

- Adds an `init=` field to `TraceTemplate` that carries a Python
function
  building a ready-to-call dict of inputs for the corresponding API.
- Init functions take the template's `Var` axes as keyword-only args
  (e.g. `rmsnorm_trace.init(batch_size=32)`); `Const` axes (hidden_size,
  num_heads, ...) are baked into the function body.
- Distributions are copied from each API's unit test under `tests/`,
  with shared patterns (paged-KV indices, ragged indptrs, RoPE cos/sin
  caches, sampling probs) factored into
  `flashinfer/trace/templates/_init_helpers.py`.
- The init function source is embedded in every dumped trace JSON
  under the `"init"` key (alongside `"reference"`), with helpers
  inlined ahead of the per-template body so the embedded snippet is
  **self-contained** — a downstream consumer can `exec()` it without
  flashinfer installed and call e.g. `_rmsnorm_init(batch_size=32)`
  directly.
- Symmetric to `func.fi_trace(...)`, every `@flashinfer_api(trace=...)`
  decorated function gets a `func.fi_init(...)` attribute.

## Coverage

| Category | Templates with init | Total | Coverage |
| --- | --- | --- | --- |
| activation, cascade, comm, gdn, gemm, mamba, norm, page, quantize,
rope, sampling | 74 | 74 | 100% |
| attention (core: gqa_paged_decode/prefill, gqa_ragged, mla_paged_*,
single_*, concat_mla_k) | 9 | 22 | 41% |
| moe (all fp8 + fp4 block_scale routing variants) | 12 | 25 | 48% |
| **Total** | **95** | **121** | — |

JSON regeneration: 86 auto-dumped JSONs under
`tests/trace/fi_trace_out/`,
**82 (95%) now contain the embedded init source**. The 4 without
(segment_gemm_run,
trtllm_batch_decode_mla, xqa_batch_decode, xqa_batch_decode_mla) are
advanced wrapper APIs deferred for follow-up; adding init for each is a
~30-line copy from the corresponding test file.

## Architecture

### Schema (`flashinfer/trace/template.py`)

```python
TraceTemplate(
    op_type="rmsnorm",
    name_prefix="rmsnorm",
    axes={"batch_size": Var(), "hidden_size": Const(abbrev="h")},
    inputs={...},
    outputs={...},
    reference=_rmsnorm_reference,
    init=_rmsnorm_init,           # NEW
)
```

### Init function contract

```python
def _<name>_init(
    *,
    <var_axis_1>: int,            # required
    <var_axis_2>: int = <default>,
    ...,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    seed: int = 0,
) -> dict[str, Any]:
    """Returns a dict whose keys are the API's Python parameter names.
    For wrapper-class APIs (paged decode/prefill, MLA), returns
    `{"plan": {...}, "run": {...}}` so callers can drive plan() and
    run() separately."""
```

### JSON embedding (`_render_init_source`)

```python
import math
import torch

# ----- shared init helpers -----
def make_paged_kv_indices(batch_size, num_pages_per_seq, page_size, ...):
    ...
def make_probs(batch_size, vocab_size, ...):
    ...
# ... 8 helpers total ...

# ----- init -----
def _rmsnorm_init(*, batch_size, hidden_size=4096, ...):
    torch.manual_seed(seed)
    return {
        "input": torch.randn(batch_size, hidden_size, ...),
        "weight": torch.randn(hidden_size, ...),
    }
```

## Tests

A new `tests/trace/test_template_init.py` auto-discovers every template
with init via `_TRACE_REGISTRY` and runs four assertions per template:
1. Signature is keyword-only with all `Var` axes accepted as kwargs.
2. Smoke-runs on CPU for canonical Var values without raising.
3. KV-cache invariants (`kv_indptr` monotonic, `kv_indptr[-1] ==
   kv_indices.numel()`, `kv_last_page_len > 0`) when those arrays are
   present.
4. `fi_trace(**init(...))` round-trips and embeds the init source.

```
$ pytest tests/trace/test_template_init.py \
         tests/trace/test_fi_trace.py \
         tests/trace/test_fi_trace_template_consistency.py
674 passed, 88 skipped
```

## Test plan

- [x] CPU-only test suite (`tests/trace/`) passes.
- [x] `tests/trace/example.py` regeneration produces 86 JSONs, 82 with
      `"init":` field; remaining 4 match the deferred wrapper list.
- [x] Pre-commit (ruff check + ruff format + mypy) passes on all
      changed files.
- [ ] Reviewer to confirm the JSON `"init"` snippet format matches the
      flashinfer-bench schema expectations.
- [ ] CI to pick up the new test file under `tests/trace/`.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Trace templates now support optional init callables to auto-generate
deterministic example inputs; many templates across attention, GEMM/BMM,
RoPE, norm, sampling, MoE, paging and quantization now expose init
hooks.
* A shared set of reusable tensor- and FP8/FP4 quantization helpers
added for consistent test inputs; init sources are embedded in many
trace fixtures.

* **Tests**
* New test module auto-discovers templates with init hooks and validates
init signatures, CPU smoke runs, KV-cache invariants, and trace
round‑trip consistency.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Avery Huang <averyh@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>