feat(logging,trace): cuda-graph-compatible level-5/10 logging + fi_trace template additions/fixes (#3172)

## Summary

Two related changes to `@flashinfer_api`:

1. **CUDA-graph compatibility for `FLASHINFER_LOGLEVEL=5` and `=10`.**
Both levels previously had host-side paths that were unsafe inside
`torch.cuda.graph(...)` capture (`tensor.min().item()`, `tensor.cpu()`,
tensor repr on nested CUDA tensors, etc.). Level 5 skipped stats under
capture; Level 10 could invalidate graph capture or only preserve the
latest replay state.
2. **`fi_trace` template additions/fixes** uncovered while validating
(1) end-to-end against sglang DSR1 FP8 TP=8 with `--attention-backend
trtllm_mla`. Several flashinfer entry points exercised by that workload
had bare `@flashinfer_api` (no trace template) or had templates whose
declared shapes did not match the actual runtime layout.

## Part A: CUDA-graph-compatible logging

### Level 5: stats under CUDA graph capture

Replaces the `[statistics skipped: CUDA graph capture in progress]` path
for supported dtypes with a small captured CUDA kernel that computes
min/max/mean/nan/inf and emits one line via device-side `printf`. The
launch is captured into the graph, so the printf fires on every
`g.replay()`. The host log records a correlation marker so the
kernel-emitted line can be matched back to the API call/argument.

Supported dtypes: `float32`, `float16`, `bfloat16`, `int32`, `int64`,
`uint8`. Other dtypes (for example fp8/fp4) fall back to the legacy skip
message.

Files: `csrc/api_log_stats.cu`,
`csrc/flashinfer_api_log_stats_binding.cu`,
`flashinfer/jit/api_log_stats.py`, `flashinfer/api_logging.py`,
`flashinfer/aot.py`.

Example output inside `torch.cuda.graph(...)`:

```text
FlashInfer API Call: my_op
  arg[0]:
    Tensor(
      shape=(64, 64) ...
      [stats deferred to GPU kernel: id=1; look for '[flashinfer stats] id=1 ...' in graph replay output]
    )
```

On `g.replay()`:

```text
[flashinfer stats] id=1 numel=4096 min=-3.42 max=3.59 mean=0.01 nan=0 inf=0
```

### Level 10: tensor dumps under CUDA graph capture

Current behavior under CUDA graph capture:

- Capture does **not** call `.cpu()`, allocate pinned host memory, or
insert D2H copy nodes into the captured graph.
- Captured level-10 calls record tensor references plus metadata and
defer disk writes.
- `torch.cuda.CUDAGraph.capture_begin` / `capture_end` are wrapped to
tag deferred dumps with the owning graph id.
- `torch.cuda.CUDAGraph.replay()` is wrapped so FlashInfer automatically
flushes graph dumps after every replay for that graph. No sglang code
injection is needed.
- Each replay flush synchronizes, materializes current tensor values to
CPU, writes root `inputs.pt` / `outputs.pt` compatibility files, and
also writes immutable `graph_flushes/flush_XXXX/` snapshots.
- Nested tensors inside `list` / `tuple` / `dict` inputs are recursively
extracted and dumped with stable keys such as `arg_2__0`, `arg_2__1`,
while metadata records `tensor_key` links so `replay_from_dump()` can
reconstruct containers.
- A process-exit / SIGTERM flush remains as a fallback, but normal
PyTorch graph replay preserves every replay automatically.

This means root `inputs.pt` / `outputs.pt` contain the latest flushed
replay, while `graph_flushes/flush_XXXX/` preserves per-replay
snapshots.

Caveats:

- Level 10 graph dumping is intentionally expensive. It adds
synchronization and disk writes to every graph replay.
- Dumping nested KV-cache tensors can create large dumps quickly. Use
`FLASHINFER_DUMP_INCLUDE`, `FLASHINFER_DUMP_MAX_COUNT`, and
`FLASHINFER_DUMP_MAX_SIZE_GB` for short targeted debug runs.
- Deferred graph tensor references are retained until
`clear_graph_dumps()` or process exit.
- `kill -9` can still lose pending deferred writes.

Files: `flashinfer/api_logging.py`, `docs/logging.rst`,
`tests/utils/test_logging.py`.

## Part B: `fi_trace` template additions / fixes

End-to-end validation under sglang DSR1 + `trtllm_mla` showed these
gaps:

1. **`gemm_fp8_nt_groupwise` had no `trace=` template.** The op fires
heavily in DSR1 + `flashinfer_trtllm` MoE but produced no trace JSON.
2. **`trtllm_batch_decode_mla_trace` and `xqa_batch_decode_mla_trace`**
declared `kv_cache` as rank-3 `[num_pages, page_size, head_dim_qk]`, but
the kernel accepts and sglang passes the rank-4 `[num_pages, 1,
page_size, head_dim_qk]` form. They were also missing
`skip_softmax_threshold_scale_factor`, and workspace dtype was corrected
from `int8` to `uint8`.
3. **`mla_rope_quantize_fp8_trace`** inherited the rank-3 GQA rope quant
axes with `num_k_heads`, but MLA passes rank-2 K tensors with
`num_k_heads=1` collapsed.
4. **`gemm_fp8_nt_groupwise_trace.b_scale`** axes were corrected to
`[N_div_block, K_div_block]` to match the trtllm path and sglang runtime
layout.

Files: `flashinfer/gemm/gemm_base.py`,
`flashinfer/trace/templates/{gemm,attention,rope}.py`,
`tests/trace/example.py`, regenerated JSONs in
`tests/trace/fi_trace_out/`.

## Test plan

### Unit / formatting

- [x] `pre-commit run --files docs/logging.rst flashinfer/api_logging.py
tests/utils/test_logging.py` passed.
- [x] `CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/home/averyh/flashinfer-pr3172
pytest -q tests/utils/test_logging.py` passed: `20 passed, 2 warnings`.
- [x] `python -m compileall -q flashinfer/api_logging.py
tests/utils/test_logging.py` passed.
- [x] `git diff --check` passed.
- [x] `pytest tests/trace/` passed previously: 440 passed, 8 skipped.
- [x] `tests/utils/test_logging_replay.py` partial local run: 14 passed,
2 failed due local environment/JIT setup before replay validation:
- `test_bmm_fp8_replay`: cuDNN reported multiple CUDA runtime libraries,
`libcudart.so.12` and `libcudart.so.13`.
- `test_mm_fp4_replay`: local JIT build failed because
`cutlass/arch/barrier.h` was missing.

### Manual CUDA graph smoke tests

- [x] Manual eager dispatch over fp32 (with NaN + Inf), bf16, int32:
values match expected stats.
- [x] Manual capture/replay at level 5: `g.replay()` after mutating the
input shows updated stats; multiple replays work.
- [x] Manual capture/replay at level 10: captured graph plus repeated
replays with different inputs; replay hook writes root files and
immutable `graph_flushes/flush_XXXX/` snapshots.
- [x] Manual level-10 nested tensor capture test: tuple tensors are
dumped as separate tensor files and reconstructed by
`replay_from_dump()`.

### SGLang single-GPU validation, May 13 2026

Environment:

- GPU: NVIDIA B200
- PyTorch: `2.9.1+cu128`
- sglang: `0.5.10.post1`
- FlashInfer import path:
`/home/averyh/flashinfer-pr3172/flashinfer/__init__.py`
- Model: `meta-llama/Llama-3.2-3B-Instruct`
- Command shape: `python -m sglang.launch_server --attention-backend
flashinfer --sampling-backend flashinfer --cuda-graph-bs 1
--cuda-graph-max-bs 1 --context-length 512 --max-total-tokens 1024
--mem-fraction-static 0.55 --dtype bfloat16`

Level 5 result:

- [x] Server arguments showed `attention_backend='flashinfer'`,
`sampling_backend='flashinfer'`, `disable_cuda_graph=False`,
`cuda_graph_bs=[1]`.
- [x] SGLang completed `Capture cuda graph bs [1]` and `Capture cuda
graph end`.
- [x] Completion request returned successfully: `" Paris. The capital"`.
- [x] SGLang request log showed `cuda graph: True`.
- [x] FlashInfer level-5 log had 322726 lines, including 283 `[stats
deferred to GPU kernel: id=...]` markers.
- [x] Server stdout contained device-side `[flashinfer stats] id=...`
lines.
- [x] No `[statistics skipped: CUDA graph capture in progress]` lines
were observed for this run.

Level 10 result:

- [x] Server arguments showed `attention_backend='flashinfer'`,
`sampling_backend='flashinfer'`, `disable_cuda_graph=False`,
`cuda_graph_bs=[1]`.
- [x] SGLang completed CUDA graph capture and served the same completion
request successfully.
- [x] Final recursive-dump validation used
`FLASHINFER_DUMP_INCLUDE='BatchDecodeWithPagedKVCacheWrapper.run'` and
`FLASHINFER_DUMP_MAX_COUNT=64`.
- [x] FlashInfer log showed 16 graph-deferred input/output sections and
12 automatic `flush_graph_dumps (CUDAGraph.replay): wrote ...` replay
flushes.
- [x] Dump tree contained 64 dump dirs, 8 dirs with `graph_flushes`, and
96 immutable replay snapshot dirs.
- [x] Example snapshot input keys were `arg_1`, `arg_2__0`, `arg_2__1`,
proving both the query tensor and nested K/V cache tuple tensors were
dumped.
- [x] Example snapshot output key was `result`.
- [x] Consecutive `flush_0001` to `flush_0004` snapshots had different
input/output sums, confirming replay snapshots are not just the last
buffer state.
- [x] Dump size for this targeted run was about 647 MB, which confirms
correctness and also the expected cost of recursive K/V tensor dumping.
- [x] No `Capture cuda graph failed`,
`cudaErrorStreamCaptureInvalidated`, or `BatchDecodeWithPagedKVCache
failed` errors in the final level-10 run.

### Larger integration validation

- [x] End-to-end sglang DSR1 FP8 TP=8 on 8x B200 at level 3, full bench:
40/40 requests in 99 s, output throughput 416 tok/s, mean TPOT 9.3
ms/token, no CUDA error, every `(api, axes)` tuple produced a trace
JSON.
- [x] End-to-end sglang at level 5: deferred-stats kernel JIT-builds
without racing graph capture, captures into graph, fires on replay. No
`cudaErrorStreamCaptureInvalidated`. Runtime is dominated by unthrottled
device-side printf rate.

## PR state

- Head commit after latest update: `008a2836`
- Base refreshed against upstream main commit: `103fcf86`
- Branch pushed to `yyihuang/flashinfer:cuda-graph-api-logging`

---------

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