From caca6d0f304dc660cce5a12c88cd58ecf6c09270 Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Wed, 16 Sep 2026 13:05:22 +0800 Subject: [PATCH 1/2] docs: guide model integration for FPM onboarding Signed-off-by: Yiming Liu --- python/aisimulate/docs/add_a_new_model.md | 4 + .../docs/fpm/end-to-end-workflow.md | 4 + .../aisimulate/docs/fpm/model-integration.md | 254 ++++++++++++++++++ 3 files changed, 262 insertions(+) create mode 100644 python/aisimulate/docs/fpm/model-integration.md diff --git a/python/aisimulate/docs/add_a_new_model.md b/python/aisimulate/docs/add_a_new_model.md index 34757afc2..5723ffebc 100644 --- a/python/aisimulate/docs/add_a_new_model.md +++ b/python/aisimulate/docs/add_a_new_model.md @@ -11,6 +11,10 @@ All paths below are relative to the repository root unless linked otherwise. The `aiconfigurator_core` source namespace is retained inside the `aisimulate` wheel; installing or rebuilding a separate AIConfigurator package is unnecessary. +For whole-forward FPM onboarding, use the [FPM model-integration procedure](fpm/model-integration.md) +alongside this guide. It covers pre-collection construction, memory, and analytical +checks; per-operation GPU timing collection is not required for that workflow. + ## Choose the smallest extension | What changed? | Start here | diff --git a/python/aisimulate/docs/fpm/end-to-end-workflow.md b/python/aisimulate/docs/fpm/end-to-end-workflow.md index 700ed19f4..277dc7c9b 100644 --- a/python/aisimulate/docs/fpm/end-to-end-workflow.md +++ b/python/aisimulate/docs/fpm/end-to-end-workflow.md @@ -11,6 +11,10 @@ parallelism, then consumes the newly collected data from an external directory. It assumes an existing, working GPU collection environment. Fill in the collection inputs from that environment before executing GPU steps. +For a new model, first complete the [model-integration procedure](model-integration.md): +reuse or add its analytical model class, verify memory and FPM SOL support on +CPU, and install the reviewed integration before planning GPU collection. + ```text Freeze plan -> GPU smoke -> GPU collection -> validate and publish pair | diff --git a/python/aisimulate/docs/fpm/model-integration.md b/python/aisimulate/docs/fpm/model-integration.md new file mode 100644 index 000000000..ba735b578 --- /dev/null +++ b/python/aisimulate/docs/fpm/model-integration.md @@ -0,0 +1,254 @@ + + +# Prepare a model for FPM collection and simulation + +FPM onboarding currently requires an AISimulate model class. Reuse a compatible +registered class or implement the missing model description before collecting +whole-forward timings. You do not need to collect per-operation GPU timings for +this workflow: the retained operations describe analytical work and memory. + +This guide covers the model-integration stage of self-service onboarding. Its +CPU examples use the bundled **Qwen3-0.6B configuration as a procedural example**; +they do not integrate Inkling or establish coverage for another model. After +reviewing and installing your integration, continue with the +[FPM collection-to-prediction workflow](end-to-end-workflow.md). + +## 1. Record the intended deployment + +Keep these inputs with the integration issue and eventual collection artifacts: + +| Input | What to inspect and record | +| --- | --- | +| Checkpoint | Canonical model ID, immutable revision, local `config.json` and optional quantization/processor metadata, and file hashes. AISimulate's config parser does not take a revision argument; use metadata from the pinned checkpoint in a local directory when an exact revision is needed. | +| Architecture | Layer types/counts, attention and KV heads, head dimension, FFN/expert dimensions, routing, shared experts, and any model-specific configuration fields. Compare these with the serving implementation. | +| Memory | Weight storage precision and replication/sharding; cache layout, quantization, sliding windows/compression, and fixed recurrent or decode state. | +| Runtime | GPU/system specification, backend and exact version, image digest, attention/MoE kernel choices, and effective GEMM/MoE/FMHA/KV/communication precision. Checkpoint weight precision alone does not specify all these values. | +| Topology and workload | TP, PP, attention DP, MoE TP/EP, CP, batch/concurrency range, input/output lengths, prefix lengths, and scheduler limits. Start with one topology and a bounded workload. | + +The current collection workflow targets vLLM with `PP=CP=1`. Start with ordinary +autoregressive text inference. Encoder/multimodal models and MTP have explicit +FPM construction restrictions; a model class existing for another mode is not +evidence that the intended FPM deployment is supported. + +## 2. Reuse or implement the model description + +Follow [How to add a new model](../add_a_new_model.md) for the registry and native +operation contracts. For this FPM use case, review these concrete responsibilities: + +| Responsibility | Source and required result | +| --- | --- | +| Parse the checkpoint | [`sdk/utils.py`](../../src/aiconfigurator_core/sdk/utils.py), especially `get_model_config_from_model_path()`. Confirm the parsed geometry and architecture-specific `extra_params`; add parsing only for fields the existing parser cannot represent correctly. | +| Select a family | [`sdk/common.py`](../../src/aiconfigurator_core/sdk/common.py), `ARCHITECTURE_TO_MODEL_FAMILY`. Reuse a family only when its operation pipeline, precision handling, and cache behavior match the model. A new architecture name alone does not require a dedicated class. | +| Construct a new family when needed | The [registry examples](../../src/aiconfigurator_core/sdk/models/README.md#adding-a-new-model) show `@register_model`, `create()`, and reusable blocks. Implement the actual `context_ops` and `generation_ops`; an empty placeholder class is insufficient. | +| Describe analytical work | Each operation needs correct dimensions, quantization, layer/repetition scaling, and per-rank parallelism. Preserve the context/generation attention and `logits_gemm` naming contracts. Inspect communication, overlap, and phase-specific behavior as well as GEMMs. | +| Describe memory | Operation `get_weights()` values provide the weight inventory. Model `get_kvcache_bytes_per_sequence()` and its inverse `get_kvcache_max_tokens()` must represent the real cache curve. Override both for non-linear/windowed/compressed layouts; do not extrapolate a one-token slope across a window boundary. | +| Support FPM analytical execution | Every retained operation must be supported by the Rust [`fpm_sol` evaluator](../../../../crates/core/src/perfmodel/operators/fpm_sol.rs) for the proposed deployment and shapes. Adding an ordinary native operation or a `SOL_FULL` diagnostic does not automatically add this support. | + +With `ModelConfig(forward_model="fpm")`, `get_model()` first constructs the +registered class, then centrally replaces each ordinary target +phase with an `FPMForwardOp`. The original phase operations remain inside +`sol_ops`, and their context weight inventory becomes `weight_bytes`. The model +class and cache methods remain available. Do not implement another phase +replacement or a Python latency/SOL formula in the new class. + +## 3. Run CPU checks before collecting timings + +From the repository root, use the [development environment](../../../../DEVELOPMENT.md): + +```bash +uv sync --project python/aisimulate --extra dev +source python/aisimulate/.venv/bin/activate +export FPM_CHECK_DIR="$(mktemp -d)" +``` + +The following example needs no model weights, GPU, or matching Qwen FPM pair. It +uses bundled model metadata and the H200 system/backend catalog. The memory +entry point still requires a resolvable catalog/database layout, even though +memory sizing does not query measured operation latencies. The example pins a +queryable backend version; select a version allowed by your installed +[`query_versions.yaml`](../../src/aiconfigurator_core/systems/query_versions.yaml) +when adapting it. + +```bash +python - <<'PY' +import hashlib +import json +import math +import os +from dataclasses import replace +from pathlib import Path + +from aisimulate_core.sdk import EngineHandle, ModelConfig, estimate_kv_cache +from aiconfigurator_core.sdk.engine import build_ops_json +from aiconfigurator_core.sdk.models import get_model +from aiconfigurator_core.sdk.utils import get_model_config_from_model_path + +out = Path(os.environ["FPM_CHECK_DIR"]) +model_path, system, backend, version = "Qwen/Qwen3-0.6B", "h200_sxm", "vllm", "0.24.0" +topology = dict(tp_size=1, pp_size=1, attention_dp_size=1, moe_tp_size=1, moe_ep_size=1) +options = ModelConfig(**topology) +info = get_model_config_from_model_path(model_path) +model = get_model(model_path, replace(options), backend) +fpm = get_model(model_path, replace(options, forward_model="fpm"), backend) + +phase_ops = {"prefill": model.context_ops, "decode": model.generation_ops} +original = {phase: json.loads(build_ops_json(ops)) for phase, ops in phase_ops.items()} +folded = { + "prefill": json.loads(build_ops_json(fpm.context_ops)), + "decode": json.loads(build_ops_json(fpm.generation_ops)), +} +weights = sum(op.get_weights() for op in model.context_ops) +assert weights > 0 +for phase in phase_ops: + assert original[phase] and len(folded[phase]) == 1 + whole = folded[phase][0]["FpmForward"] + assert whole["sol_ops"] == original[phase] + assert whole["weight_bytes"] == weights + +kv_bytes = {s: model.get_kvcache_bytes_per_sequence(s) for s in (1, 128, 8192)} +for s, budget in kv_bytes.items(): + assert budget > 0 and fpm.get_kvcache_bytes_per_sequence(s) == budget + capacity = model.get_kvcache_max_tokens(budget) + assert model.get_kvcache_bytes_per_sequence(capacity) <= budget + assert model.get_kvcache_bytes_per_sequence(capacity + 1) > budget + +memory = estimate_kv_cache( + model_path, system, backend, version, **topology, + max_num_tokens=2048, max_batch_size=8, + memory_fraction_kind="of_total", memory_fraction_value=0.9, + allow_naive_fallback=False, +) +assert memory["source"] == "native" and memory["total_kv_size_tokens"] > 0 + +# This is the per-operation SOL_FULL diagnostic, not FPM interpolation. +engine = EngineHandle.compile( + model_path, system, backend, backend_version=version, **topology, + database_mode="SOL", +) +sol = {} +for phase, ops in phase_ops.items(): + rows = engine.evaluate_ops_sol_json( + build_ops_json(ops), is_context=phase == "prefill", + batch_size=1, s=128 if phase == "prefill" else 129, + ) + assert rows and all(math.isfinite(v) and v >= 0 for row in rows for v in row[1:]) + assert sum(row[1] for row in rows) > 0 + sol[phase] = rows + +report = { + "model_path": model_path, "class": type(model).__name__, "family": model.model_family, + "system": system, "backend": backend, "backend_version": version, "topology": topology, + "geometry": {k: info[k] for k in ("architecture", "layers", "n", "n_kv", "d", "hidden_size")}, + "resolved_config_sha256": hashlib.sha256( + json.dumps(info["raw_config"], sort_keys=True).encode() + ).hexdigest(), + "quantization": {k: getattr(model.config, k).name for k in ( + "gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", "kvcache_quant_mode", "comm_quant_mode" + )}, + "op_counts": {phase: len(ops) for phase, ops in phase_ops.items()}, + "weights_bytes_per_rank": weights, "kv_bytes_per_sequence_per_rank": kv_bytes, + "memory": memory, "sol_full_ms": sol, +} +for name, value in (("integration.json", report), ("op-level.json", original), ("fpm.json", folded)): + (out / name).write_text(json.dumps(value, indent=2) + "\n") +print(json.dumps(report, indent=2)) +print(f"Saved integration evidence to {out}") +PY +``` + +**Expected for this example:** `LLAMAModel` / `LLAMA`, 28 layers, 16 attention +heads, 8 KV heads, head dimension 128, and 14 operation entries per phase. +The modeled weight inventory is 1,503,133,696 bytes per rank. BF16 KV costs +114,688 bytes per token; an 8,192-token sequence uses 939,524,096 bytes per rank. +The FPM wrapper preserves that inventory and curve. The diagnostic produces +finite nonnegative per-operation SOL values, including zero-cost communication +at TP=PP=1. These are analytical estimates, not measured latency or peak runtime +memory. `resolved_config_sha256` hashes parsed metadata, not checkpoint weights. + +For your integration, replace the example model and deployment inputs together. +Compare the emitted operation dimensions and scale factors with the serving +architecture, and add independent expected-value tests for those facts. Exercise +each supported topology/precision and the relevant cache/window boundaries; +three positive sequence lengths and a successful import are not sufficient +model-specific evidence. See the existing [model configuration and memory tests](../../tests/unit/sdk/models/test_model_config.py) +for patterns. Include an unsupported architecture/topology case that must fail +explicitly, rather than falling back to another family or memory estimator. + +## 4. Verify the FPM execution contract + +Run the existing CPU contract suites from the same checkout/environment: + +```bash +python -m pytest -p no:timeout -c python/aisimulate/pytest.ini \ + python/aisimulate/tests/unit/sdk/test_fpm_forward.py \ + python/aisimulate/tests/unit/sdk/test_opspec_coverage.py \ + python/aisimulate/tests/cross_package/test_single_oracle_contract.py +cargo test -p aisimulate-core --lib perfmodel::operators::fpm_sol +cargo test -p aisimulate-core --lib perfmodel::operators::fpm_forward +``` + +`-p no:timeout` follows the repository's macOS test guidance. These tests exercise +the Python-to-Rust operation contract, FPM routing, and native analytical/transfer +behavior using existing models and synthetic fixtures. They do not automatically +cover your new model's operations or demonstrate predictive accuracy. + +There is currently no public model-readiness call that evaluates an arbitrary +retained graph through the exact FPM `sol_total`/`fpm_sol` path without a pair. +`evaluate_ops_sol_json()` above uses a different diagnostic path and integer +workload coordinates. FPM transfer maps iteration totals to potentially +fractional per-request coordinates. Also, exact FPM hits and some linear +interpolations can succeed without ever evaluating `sol_ops`. + +Before declaring the new model ready for collection, add a model-specific CPU +regression that serializes its real FPM graph and queries deliberately synthetic +FPM cells at an off-site coordinate requiring SOL transfer. Cover both phases +and any relevant fractional coordinates. Use the pair/engine setup in +[`test_fpm_forward.py`](../../tests/unit/sdk/test_fpm_forward.py) and the native +tests `prefill_off_site_kv_resolves_through_the_dsa_sol_transfer`, +`decode_pairless_batch_resolves_through_the_dsa_sol_transfer`, and +`unsupported_sol_family_is_lazy` in +[`fpm_forward.rs`](../../../../crates/core/src/perfmodel/operators/fpm_forward.rs) +as contract examples. Assert an independently justified result and a clear +failure for an unsupported shape or identity; an exact-hit-only fixture misses +this dependency. Keep synthetic timings in test fixtures, never publish them as +collected data. If an operation is unsupported, implement and numerically test +its Rust FPM SOL path before relying on sparse collection. + +## 5. Review, install, and hand off to collection + +Submit the model/parser/operation changes with the focused model tests, CPU +reports, resolved deployment inputs, and any explicit unsupported cases. Follow +the repository [review contract](../../../../REVIEW.md). After review, install +the approved AISimulate revision in the environment that will plan collection +and run prediction. For a source checkout at that revision: + +```bash +git rev-parse HEAD +uv sync --project python/aisimulate --extra dev +source python/aisimulate/.venv/bin/activate +python -c 'import importlib.metadata; print(importlib.metadata.version("aisimulate"))' +aisimulate predict --help +aisimulate recommend --help +``` + +Record the commit as well as the package version; several source revisions may +share a version. Rerun the CPU checks in a fresh process against the installed +integration. Preserve the reports with the pinned config files and test output. + +The collection handoff consists of the canonical checkpoint identity/revision, +config files/hashes, reviewed AISimulate revision, exact runtime image/version, +hardware/topology/precision choices, workload limits, and the CPU evidence. +The metadata used for local construction must describe the same checkpoint +mounted in the collection runtime; a temporary local path is not a portable +replacement for the canonical identity in the published pair. + +Proceed to [freeze a bounded plan, smoke, collect, resume, and validate the pair](end-to-end-workflow.md#3-freeze-and-inspect-the-plan). +Only whole-forward silicon timings are collected for FPM. Then exercise the +ordinary standalone AISimulate/Mocker FPM `predict` and `recommend` paths on +covered candidates and save the matching inputs, data provenance, reports, and +reproduction commands. Model integration, collection completion, functional +prediction, and agreement with independent silicon measurements are separate +acceptance results; this guide establishes the procedure for the first stage. From d338ebc37f2c6e4abc496975bce9e05c33e8a55a Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Thu, 17 Sep 2026 04:57:41 +0800 Subject: [PATCH 2/2] docs: make analytical integration an optional FPM route Signed-off-by: Yiming Liu --- python/aisimulate/docs/add_a_new_model.md | 7 ++- .../docs/fpm/end-to-end-workflow.md | 8 ++- .../aisimulate/docs/fpm/model-integration.md | 55 ++++++++++++++----- 3 files changed, 50 insertions(+), 20 deletions(-) diff --git a/python/aisimulate/docs/add_a_new_model.md b/python/aisimulate/docs/add_a_new_model.md index 5723ffebc..cac6912d2 100644 --- a/python/aisimulate/docs/add_a_new_model.md +++ b/python/aisimulate/docs/add_a_new_model.md @@ -11,9 +11,10 @@ All paths below are relative to the repository root unless linked otherwise. The `aiconfigurator_core` source namespace is retained inside the `aisimulate` wheel; installing or rebuilding a separate AIConfigurator package is unnecessary. -For whole-forward FPM onboarding, use the [FPM model-integration procedure](fpm/model-integration.md) -alongside this guide. It covers pre-collection construction, memory, and analytical -checks; per-operation GPU timing collection is not required for that workflow. +For whole-forward FPM onboarding, first review the [FPM execution routes and metadata requirements](fpm/model-integration.md). +Use its registered-model/SOL procedure alongside this guide when choosing to +reuse or add an analytical model class. The planned class-independent route is +separate; neither FPM route requires per-operation GPU timing collection. ## Choose the smallest extension diff --git a/python/aisimulate/docs/fpm/end-to-end-workflow.md b/python/aisimulate/docs/fpm/end-to-end-workflow.md index 277dc7c9b..2dc804654 100644 --- a/python/aisimulate/docs/fpm/end-to-end-workflow.md +++ b/python/aisimulate/docs/fpm/end-to-end-workflow.md @@ -11,9 +11,11 @@ parallelism, then consumes the newly collected data from an external directory. It assumes an existing, working GPU collection environment. Fill in the collection inputs from that environment before executing GPU steps. -For a new model, first complete the [model-integration procedure](model-integration.md): -reuse or add its analytical model class, verify memory and FPM SOL support on -CPU, and install the reviewed integration before planning GPU collection. +For a new model, first record the [model metadata and intended FPM execution route](model-integration.md). +The prediction examples in this revision use a registered analytical class; +follow the optional registered-model/SOL procedure when choosing that route. +The planned class-independent route will use direct measured-time interpolation +and independent resource metadata, without adding an op-level model class. ```text Freeze plan -> GPU smoke -> GPU collection -> validate and publish pair diff --git a/python/aisimulate/docs/fpm/model-integration.md b/python/aisimulate/docs/fpm/model-integration.md index ba735b578..914759325 100644 --- a/python/aisimulate/docs/fpm/model-integration.md +++ b/python/aisimulate/docs/fpm/model-integration.md @@ -3,18 +3,45 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Prepare a model for FPM collection and simulation - -FPM onboarding currently requires an AISimulate model class. Reuse a compatible -registered class or implement the missing model description before collecting -whole-forward timings. You do not need to collect per-operation GPU timings for -this workflow: the retained operations describe analytical work and memory. - -This guide covers the model-integration stage of self-service onboarding. Its -CPU examples use the bundled **Qwen3-0.6B configuration as a procedural example**; -they do not integrate Inkling or establish coverage for another model. After -reviewing and installing your integration, continue with the -[FPM collection-to-prediction workflow](end-to-end-workflow.md). +# Prepare model metadata and choose an FPM execution route + +FPM onboarding needs pinned model and deployment metadata, correct memory and +KV-cache accounting, and matching whole-forward timing data. An op-level model +class is one way to supply analytical work and resource information; it is not +the intended prerequisite for every FPM onboarding. Per-operation GPU timing +collection is not required. + +In this revision, ordinary FPM `predict` and `recommend` still construct a +registered analytical class and can use its SOL estimates for timing transfer. +The separate decoupling change will add direct measured-time interpolation and +independent model/resource metadata for models without a registered class. It +will preserve the existing registered-model/SOL route. Collection bootstrap can +already resolve some unregistered configurations; this does not establish that +the current prediction path can construct those models. + +Record the common inputs in section 1 for either route. Sections 2–5 describe +the **optional registered-model/SOL integration procedure**. Their runnable CPU +examples use the bundled **Qwen3-0.6B configuration as a procedural example**; +they do not integrate Inkling or establish coverage for another model. The +[FPM collection-to-prediction workflow](end-to-end-workflow.md) covers the later +collection and simulation stages. + +## Planned class-independent route + +The self-service pilot will onboard Inkling on NVIDIA GB200 without adding an +op-level model class. It will cover TP, DEP, and TEP configurations with separately +pinned topology identities. The exact checkpoint, runtime, allocation, and +strategy degrees still need to be resolved; the current guided CLI foundation +generates pure-TP configurations only. + +The separate decoupling change must resolve model geometry, weight storage and +sharding, and cache layout from validated metadata rather than a retained +operation graph. Direct prefill interpolation will use measured prompt curves +and two-sided KV brackets at the same batch size, including wider brackets when +both neighboring curves cover the query. It must keep exact deployment identity, +reject unsupported metadata and uncovered queries explicitly, and avoid hidden +SOL fallback for unregistered models. 2D interpolation remains experimental. +No class-independent API or runnable example is introduced by this guide. ## 1. Record the intended deployment @@ -33,10 +60,10 @@ autoregressive text inference. Encoder/multimodal models and MTP have explicit FPM construction restrictions; a model class existing for another mode is not evidence that the intended FPM deployment is supported. -## 2. Reuse or implement the model description +## 2. Registered-model route: reuse or implement the model description Follow [How to add a new model](../add_a_new_model.md) for the registry and native -operation contracts. For this FPM use case, review these concrete responsibilities: +operation contracts. For this registered-model/SOL route, review these concrete responsibilities: | Responsibility | Source and required result | | --- | --- |