ai-dynamo/aisimulate/aic-auto-collect

AIC Auto Collect

Start with the project playbook

Read `docs/perf-model/collector/upgrade.md` before changing collector code: it is the canonical upgrade/bring-up workflow. This skill is the run pipeline; the repository-owned policy in .claude/rules/collector/ is authoritative over anything restated here.

Goal

You are already on a GPU node with one or more GPUs. Your job is to collect AIC perf data for one target (GPU system, framework/backend, framework version), fix collector failures, recollect until the perf files are good enough, verify AIC can consume the new data, and maintain a draft PR to ai-dynamo/aiconfigurator as the checkpoint.

This is a long-running workflow. Be persistent. Iterate carefully. Do not hide failures by shrinking coverage or deleting hard cases unless the configuration is genuinely unsupported and documented.

Operating Mode

Assume the user wants the complete bring-up unless they explicitly ask for a smoke test only.

• Run collectors inside the target backend runtime containers. Do not claim data was collected for a backend/version unless the collector code ran in that container on the target GPU.
• Full collection means every supported registry op for the requested backend/version has either finished with real rows or is explicitly unsupported by that runtime/GPU with evidence.
• Smoke and --limit runs are only gates before the full run; never present them as final data.
• Keep working through multi-hour or multi-day runs. Use stable resume checkpoints, inspect errors, patch, and retry.
• The unit of work is one draft PR per (system, backend, backend_version). Use that PR as the checkpoint, not a private local branch.
• Make small sign-off commits as progress checkpoints after each collect_xx.py/op-family collection finishes. If a single collection runs longer than about one hour, checkpoint safe collector fixes, current perf outputs, and the failure/status summary before continuing.
• If multiple frameworks are requested, finish and checkpoint one framework/version PR first, then create a new draft PR for the next framework/version.
• Do not synthesize missing rows. An unsupported runtime or kernel path gets a collector/skips.py entry or fails and is recorded; it is never filled with fake numbers.
• If the PR body says customer support is ready, verify the default AIC workflow can actually choose a valid configuration for at least the expected representative model/backend/system path.

Required Inputs

Establish these before collecting:

• Target AIC repo and branch.
• Target backend and version: sglang, vllm, or trtllm. Treat "all three" as an ordered queue of separate draft PRs.
• Target runtime container image for each backend, such as:
• nvcr.io/nvidia/ai-dynamo/tensorrtllm-runtime:<tag>
• nvcr.io/nvidia/ai-dynamo/sglang-runtime:<tag>
• nvcr.io/nvidia/ai-dynamo/vllm-runtime:<tag>
• Target system name, such as rtx_pro_6000_server.
• Target ops: start narrow, then expand to all relevant registry ops.
• Whether to collect power columns.
• Whether the GPU type is already supported by AIC.

Discover the node:

nvidia-smi
nvidia-smi --query-gpu=name,memory.total,power.limit --format=csv
python3 - <<'PY'
import torch
print("cuda available:", torch.cuda.is_available())
print("device count:", torch.cuda.device_count())
for i in range(torch.cuda.device_count()):
    print(i, torch.cuda.get_device_name(i), torch.cuda.get_device_capability(i))
PY

Protect the host and other workloads

Run framework inspection and collection inside the pinned container. Treat the host driver, CUDA toolkit, kernel, container runtime, and system packages as immutable.

Before every GPU run, inspect:

docker ps
nvidia-smi
nvidia-smi --query-compute-apps=pid,process_name,used_gpu_memory --format=csv
df -h
df -h /dev/shm

Use task-unique container names and persistent output/checkpoint/log roots. Clean only task-owned workers and containers. Never kill or restart another container, restart Docker, or run a global prune. If another workload blocks the run, report it.

When the runtime root is read-only, mount every framework JIT/cache location that the selected path may create, not only the generic XDG cache. This can include DeepGEMM (~/.deep_gemm or DG_JIT_CACHE_DIR), TensorRT-LLM, TileLang, Triton/TorchInductor, CUDA, and FlashInfer caches. A family-wide EROFS failure while creating one of these directories is runner plumbing, not evidence that its shapes or kernel are unsupported. Preserve the rejected attempt, fix only the task runner, and rerun in a fresh namespace.

Treat an OOM as unclassified until the same case fails on a clean GPU. Check for stale workers, retained weights, descriptor/JIT caches, and oversized dummy allocations before adding a capacity rule.

Phase 1: Prepare Repo

1. Clone or update ai-dynamo/aiconfigurator.
2. Create a branch named like data/<system>-<backend>-<version> or codex/<system>-<backend>-<version>.
3. Immediately create a draft PR for this (system, backend, version) once the branch exists and the scope is known. Use it as the remote checkpoint even before all data is ready.
4. Install only lightweight local dev dependencies needed for tests. Do not install the target backend locally unless the collector is not running in a backend container.
5. Set:

export PYTHONPATH="$PWD"
export COLLECTOR_LOG_DIR="$PWD/collector_logs"
export COLLECTOR_CHECKPOINT_DIR="$PWD/collector_checkpoints"
mkdir -p "$COLLECTOR_LOG_DIR"
mkdir -p "$COLLECTOR_CHECKPOINT_DIR"

Local validation often works with the repo's pinned environment:

uv run --frozen ruff check <files>
uv run --frozen ruff format --check <files>
uv run --frozen pytest <tests> -q

If the local uv environment lacks framework packages such as torch, run framework-dependent collector smoke tests in the runtime container instead of trying to mutate the host environment.

Container pattern:

docker run --rm -it --gpus all --ipc=host --network host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v "$PWD:/workspace" -w /workspace \
  <runtime-image> bash
 
export PYTHONPATH=/workspace
export COLLECTOR_LOG_DIR=/workspace/collector_logs/<backend>/<version>/<op>
mkdir -p "$COLLECTOR_LOG_DIR"
python collector/collect.py --backend <backend> --ops <op> --smoke
python collector/collect.py --backend <backend> --ops <op> \
  --checkpoint-dir "$COLLECTOR_CHECKPOINT_DIR/<backend>/<version>/<op>" --resume

Inside the container, record the real framework version:

python - <<'PY'
import importlib.metadata as m
for pkg in ("tensorrt_llm", "sglang", "vllm"):
    try:
        print(pkg, m.version(pkg))
    except Exception:
        pass
PY

Update the draft PR body whenever scope changes or a meaningful checkpoint lands. Include current status, what is still partial, and links or paths to failure summaries.

Phase 2: Add Or Validate System Metadata

If the system is new or incomplete:

1. Add src/aisimulate_core/systems/<system>.yaml.
2. Add <system> to SupportedSystems in

src/aisimulate_core/sdk/common.py.

1. Create src/aisimulate_core/systems/data/<system>/.
2. Populate YAML with conservative, documented values:

• gpu.mem_bw
• gpu.mem_capacity
• tensor core FLOPS fields relevant to the architecture
• gpu.power
• gpu.sm_version
• node.num_gpus_per_node
• inter-node, intra-node, PCIe bandwidth
• misc.nccl_version
• misc.other_mem

Prefer verified values from nvidia-smi, node docs, or nearby existing YAML files. If a value is an estimate, leave a comment.

Run:

pytest tests/unit/sdk/test_common.py -q

Phase 3: Framework Version And Kernel Preflight

Before collecting a new framework version or a new GPU SM target, check whether collector code matches the runtime. This avoids spending hours on known-bad grids.

1. Record the exact runtime package version from inside the container. Do not trust only the image tag.
2. Inspect collector/<backend>/registry.py for version routes and the concrete collect_xx.py files registered for this version.
3. Read the op's collector/skips.py entries and the top of each relevant collect_xx.py file: docstring, __compat__ and deviation notes.
4. Compare against the framework source for the exact version/tag when there is any sign of API or kernel churn:

• SGLang: attention backends, DeepGEMM, FlashInfer, DeepSeek/DSV module code, MoE runner paths.
• vLLM: attention/MLA modules, quantized GEMM paths, GDN/Mamba modules, custom all-reduce.
• TRT-LLM: attention plugins, MLA/DSA modules, MoE, compute-scale, GDN/Mamba modules.

1. For a new SM such as SM120, check kernel architecture support before full collection: tiny smoke cases, import probes, and direct framework source guards.
2. When the runtime cannot run a whole block (an op, dtype or model on this SM), add a collector/skips.py entry with the reason before full collection, rather than discovering the same failure across thousands of tasks. Scattered failures still run and are recorded.
3. If a special framework image exists for a model family or kernel family, try that image before declaring the op unsupported. Record its exact package version and image digest.

Commit and push preflight collector fixes before starting long full runs.

Phase 4: Choose Collection Plan

Read the backend registry:

python3 - <<'PY'
from collector.sglang.registry import REGISTRY as SGLANG
from collector.vllm.registry import REGISTRY as VLLM
from collector.trtllm.registry import REGISTRY as TRTLLM
for name, reg in [("sglang", SGLANG), ("vllm", VLLM), ("trtllm", TRTLLM)]:
    print(name, [e.op for e in reg])
PY

Recommended order:

1. gemm
2. moe
3. attention / MLA ops
4. module-level ops such as DSA, GDN, MHC, WideEP
5. communication ops separately, because they may require multi-GPU or full-node ownership

Collect only what the question needs (.claude/rules/collector/overview.md, "Collection cost"). An upgrade or a fix runs the affected ops and cases on one system per SM generation. A full plan is for bringing up a new system, or for a family whose kernels changed at a pin bump; the other families declare reuse. To run exact cases, pass each case's full string as its own --case-filter, for example --case-filter "['fp8_block', 64, 32, 4096]". --plan-only --sm <sm> --plan-table DIR lists every planned case of an op.

For new GPUs, prioritize ops used by the target models and backend first, but do not claim full support until all expected ops for that backend/version are collected or explicitly marked unsupported.

For an all-backend bring-up, run this as separate draft PRs:

1. Finish one GPU/framework/framework_version PR.
2. Push final data and collector fixes for that framework.
3. Create the next draft PR for the next framework/version.

For each backend/version, enumerate registry ops and create a tracking table with: op, collector file, perf filename, smoke status, full status, row count, duplicate-key status, AIC sanity status, and unsupported notes. Keep this table in a local scratch file and summarize it in the PR body.

Phase 5: Progressive Collection Loop

For each op:

python3 collector/collect.py --backend <backend> --ops <op> --smoke
python3 collector/collect.py --backend <backend> --ops <op> --shuffle --limit 20
python3 collector/collect.py --backend <backend> --ops <op> --shuffle --limit 100
python3 collector/collect.py --backend <backend> --ops <op> \
  --checkpoint-dir "$COLLECTOR_CHECKPOINT_DIR/<backend>/<version>/<op>" --resume

Campaign mode (full plan, several GPUs) — use the checked-in launcher, do not hand-roll one