From 3eae51edd74cd072546a1502bcc029ff67220e53 Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Mon, 28 Sep 2026 16:15:45 +0800 Subject: [PATCH 1/7] fix: count only uncached tokens against the legacy agg prefill budget The legacy agg estimator ignored prefix-cache hits when scheduling: - run_agg (base_backend.py) sized the mixed (prefill-bearing) steps as ceil(isl * b / ctx_tokens) with the FULL isl, and derived the balance score, the mixed-step decode tokens, the TTFT chunk count and the prefilling-request count from it. --prefix only reached the per-step cost, so a batch whose requests were 90% cached still ran as many mixed steps as a cold batch. - The Rust mixed step (runtime.rs mixed_step_breakdown_with) credited the cached prefix in pass 1 as prefix * floor(ctx / isl), which is 0 whenever ctx < isl (chunked prefill), and pass 2 packed ceil(ctx / isl) requests of the full isl. ctx_tokens now means what the engines' knobs cap: a per-step budget of UNCACHED (new) prefill tokens (SGLang --chunked-prefill-size, vLLM max_num_batched_tokens, the TRT-LLM scheduler's max_num_tokens). - run_agg packs requests by isl_new = isl - prefix: mixed-step count, balance score, mixed-step decode tokens, TTFT chunk count ceil(isl_new / ctx) and prefilling-request count all use it. The schedule budget is capped at b * isl_new for every batch size, and when it covers the whole batch (ctx_tokens >= b * isl_new) every request prefills in the one mixed step with no decode request priced alongside. Packing by isl_new makes ceil(ctx / isl_new) > b reachable with ordinary inputs (a large prefix), which would publish negative decode-request counts; the old full-isl schedule already did so for an explicit ctx_tokens > b * isl, and priced phantom requests at b == 1. prefix >= isl is rejected up front. - Pass 1 prices ctx + decode new tokens. The cached prefix of the floor(ctx / isl_new) complete requests (at least the one being chunked; none without prefill) is passed as KV context, so an op that folds attention into a module outside the `context_attention` name (the DeepSeek/Kimi `context_mla_block`) still sees it; as before, those requests are priced as one sequence over their summed prefixes. Token-major ops ignore it. With ctx == isl_new this is exactly the previous (ctx + decode - prefix, prefix) query. - Pass 2 packs floor(ctx / isl_new) complete requests of isl_new new tokens over prefix cached ones, plus the partial request of a non-multiple budget weighted by its fill fraction, divided by ceil(isl_new / ctx). prefix == 0 keeps the legacy ceil packing. - The --ctx-tokens default is max(isl - prefix, 1), visual tokens included (CLI, Task agg), and the agg sweep grid and its guards are built on isl_new, so high-hit sweeps no longer come back empty. - Docs, CLI help and SDK docstrings (MixedStepInput, run_mixed) describe the uncached budget. Behaviour changes: - prefix == 0: unchanged, except when ctx_tokens >= b * isl (the whole batch prefills in one step), where the old schedule priced phantom decode or prefill requests. The agg sweep never reaches this regime. - prefix > 0, default budget: the new default max(isl - prefix, 1) prices every prefix parity case exactly as the old default (isl) did (all engine-step and compile-engine prefix records match at rtol 1e-12). - prefix > 0, any explicit budget other than isl - prefix is re-priced. Below it, a request's uncached prefill needs ceil(isl_new / ctx) chunks instead of ceil(isl / ctx). Above it, a step carries the uncached tokens of more requests: e.g. DeepSeek-V3 b200 vLLM TP8/EP8, isl 2048, osl 16, prefix 1024, bs 2 (the deepseek-v3-b200-vllm-shape-prefix-heavy parity shape) with an explicit --ctx-tokens 2048 (the old default) now prefills both requests in one step (mixed step 55.78 -> 89.76 ms, TTFT 117.65 -> 155.88 ms, TPOT 16.05 -> 15.33 ms); its default budget 1024 keeps 55.78 / 117.65 / 16.05. - The FPM mixed branch already packed ceil(ctx / (isl - prefix)); with the new default budget it prices one request per step at prefix > 0 instead of ceil(isl / (isl - prefix)). Tests: Rust oracle tests for the mixed step (prefix 0 bit-identical to the legacy composition; chunked, full and partial budgets; exact pass-1 KV context for a module-attention op at ctx < isl_new, = isl_new, > 2 isl_new and ctx 0; DSv4.1 branch); Python unit tests for the schedule, the caps at every batch size (including ctx_tokens == b * isl_new, osl 1) with no decode request priced when the whole batch prefills, the TTFT chunk count, the default budget and the sweep grid; e2e tests for explicit ISL-sized budgets and the DeepSeek-V3 MLA module path (warm mixed step > cold at equal new tokens). The parity harness and pin_goldens feed the default budget max(isl - prefix, 1) to the prefix cases through one shared helper. Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- .../parity_tests/perfmodel/pin_goldens.py | 6 +- .../perfmodel/test_compile_engine_parity.py | 13 +- .../perfmodel/test_engine_step_parity.py | 12 +- crates/core/src/perfmodel/engine/runtime.rs | 626 +++++++++++++++--- crates/core/src/perfmodel/py.rs | 3 +- docs/cli/legacy-aic-user-guide.md | 6 +- .../src/aisimulate/legacy_cli/api.py | 13 +- .../src/aisimulate/legacy_cli/main.py | 10 +- .../src/aisimulate/sdk/inference_session.py | 7 +- .../aisimulate/src/aisimulate/sdk/predict.py | 2 +- python/aisimulate/src/aisimulate/sdk/sweep.py | 28 +- .../aisimulate/src/aisimulate/sdk/task_v2.py | 13 +- .../sdk/backends/base_backend.py | 91 ++- .../src/aisimulate_core/sdk/engine.py | 3 +- .../aisimulate_core/sdk/rust_engine_step.py | 4 +- .../src/aisimulate_core/sdk/step_estimate.py | 3 + .../tests/e2e/cli/test_cli_estimate_static.py | 79 +++ .../aisimulate/tests/unit/cli/test_cli_api.py | 75 +++ .../unit/sdk/backends/test_base_backend.py | 329 ++++++++- .../speculation/test_consumer_equivalence.py | 34 +- .../tests/unit/sdk/sweep/test_sweep.py | 71 ++ .../unit/sdk/task_v2/test_task_config.py | 40 ++ 22 files changed, 1313 insertions(+), 155 deletions(-) diff --git a/crates/core/parity_tests/perfmodel/pin_goldens.py b/crates/core/parity_tests/perfmodel/pin_goldens.py index 43cc0bf49..7bfceb029 100644 --- a/crates/core/parity_tests/perfmodel/pin_goldens.py +++ b/crates/core/parity_tests/perfmodel/pin_goldens.py @@ -177,7 +177,11 @@ def _compile_case_references(case) -> dict[str, float]: return { "static_ctx": float(ctx), "static_gen": float(gen), - "mixed_step": float(handle.mixed_step_latency(case.isl, case.batch_size, case.isl, osl, case.prefix)), + "mixed_step": float( + handle.mixed_step_latency( + compile_parity._default_ctx_tokens(case), case.batch_size, case.isl, osl, case.prefix + ) + ), "decode_step": float(handle.decode_step_latency(case.batch_size, case.isl, osl)), } diff --git a/crates/core/parity_tests/perfmodel/test_compile_engine_parity.py b/crates/core/parity_tests/perfmodel/test_compile_engine_parity.py index 0ea8fb495..c4cfa2cde 100644 --- a/crates/core/parity_tests/perfmodel/test_compile_engine_parity.py +++ b/crates/core/parity_tests/perfmodel/test_compile_engine_parity.py @@ -306,9 +306,16 @@ def _assert_within(name: str, python_value: float, new_value: float, *, backend: (512, 4, 4096, 128, 0), # chunked prefill: ctx_tokens < isl (512, 4, 4096, 128, 256), # chunked + cached prefix (300, 7, 1000, 64, 100), # ragged chunk + prefix + decode overlap + (4096, 4, 4096, 128, 256), # isl-sized budget + cached prefix: one request plus a partial one ] +def _default_ctx_tokens(case) -> int: + """The agg default budget: one request's uncached prefill. Shared by the + mixed-step test and the golden pin path (pin_goldens.py).""" + return max(case.isl - case.prefix, 1) + + def _chunked_prefill_key(ctx_tokens: int, gen_tokens: int, isl: int, osl: int, prefix: int) -> str: return f"chunked_prefill::ctx{ctx_tokens}_gen{gen_tokens}_isl{isl}_osl{osl}_prefix{prefix}::mixed_step" @@ -333,7 +340,11 @@ class TestCompileEngineMixedStepParity: @pytest.mark.parametrize("case", _SUBSET_CASES) def test_mixed_step(self, case: EngineStepParityCase) -> None: handle = _compile_handle(case) - new_val = handle.mixed_step_latency(case.isl, case.batch_size, case.isl, max(case.osl, 2), case.prefix) + # One request's uncached prefill as the budget (the agg default), as in + # test_engine_step_parity._mix_step_shape. + new_val = handle.mixed_step_latency( + _default_ctx_tokens(case), case.batch_size, case.isl, max(case.osl, 2), case.prefix + ) py_val = _case_reference(case, "mixed_step") _assert_within("mixed_step", py_val, new_val, backend=case.backend_name) diff --git a/crates/core/parity_tests/perfmodel/test_engine_step_parity.py b/crates/core/parity_tests/perfmodel/test_engine_step_parity.py index 7842f1706..bbc508efc 100644 --- a/crates/core/parity_tests/perfmodel/test_engine_step_parity.py +++ b/crates/core/parity_tests/perfmodel/test_engine_step_parity.py @@ -758,7 +758,8 @@ def call(): backend_name=case.backend_name, backend_version=case.backend_version, batch_size=case.agg_batch_size, - ctx_tokens=case.agg_ctx_tokens or case.isl, + # The CLI default budget: one request's uncached prefill. + ctx_tokens=case.agg_ctx_tokens or max(case.isl - case.prefix, 1), isl=case.isl, osl=case.osl, prefix=case.prefix, @@ -896,12 +897,13 @@ def call(): def _mix_step_shape(case: EngineStepParityCase) -> dict: """Mix-step (chunked-prefill + decode) shape for a smoke case. - Treats the case's single prefill request as one chunk with `case.isl` - isl-equivalent tokens (matching Python's agg orchestration). Decode - batch is `case.batch_size` (matches the FPM constructor below). + Treats the case's single prefill request as one chunk: the budget is its + `case.isl - case.prefix` uncached tokens (the agg default), over + `case.prefix` cached ones. Decode batch is `case.batch_size` (matches the + FPM constructor below). """ return { - "ctx_tokens": case.isl, + "ctx_tokens": max(case.isl - case.prefix, 1), "gen_tokens": case.batch_size, "isl": case.isl, "osl": max(case.osl, 2), diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index ca67b169c..800ac077c 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -771,36 +771,43 @@ impl Engine { decode.decode_kv_ceiling(&self.db) } - /// One mixed (chunked-prefill + decode) step latency. LITERAL mirror of - /// Python `_get_mix_step_latency` / `run_mixed`, which composes three - /// filtered phase passes (`_run_context_phase` / `_run_generation_phase` - /// with `op_filter`) that query ONLY the ops each pass consumes — the - /// same name-keyed sets the `ContextOpFilter` / - /// `only_generation_attention` walks below visit (issue #1498 - /// follow-through: Python used to run the full lists and discard, so a - /// raise in a discarded query was a one-sided error surface): + /// One mixed (chunked-prefill + decode) step latency: the three-pass + /// composition of the legacy Python `_get_mix_step_latency` / + /// `run_mixed`, each pass querying ONLY the ops it consumes (the + /// name-keyed sets the `ContextOpFilter` / `only_generation_attention` + /// walks below visit; issue #1498 follow-through). `ctx_tokens` budgets + /// UNCACHED (new) prefill tokens — what SGLang `--chunked-prefill-size`, + /// vLLM `max_num_batched_tokens` and the TRT-LLM scheduler's + /// `max_num_tokens` cap — so with `isl_new = isl - prefix`: /// /// ```text - /// // Pass 1 — combined non-attention work: - /// // run_static(batch=1, isl=ctx+gen, osl=1, - /// // prefix=prefix*floor(ctx/isl), mode=static_ctx) - /// // sum every op EXCEPT "context_attention" + /// // Pass 1 — combined non-attention work (every budget token is new): + /// // run_static(batch=1, isl=ctx+gen*(nextn+1), osl=1, + /// // prefix=(ctx > 0 && prefix > 0) + /// // ? prefix * max(floor(ctx/isl_new), 1) : 0, + /// // mode=static_ctx) + /// // sum every op EXCEPT "context_attention"; the prefix is KV + /// // context for ops that fold attention into a module /// // Pass 2 — context attention at the prefill shape: - /// // run_static(batch=ceil(ctx/isl), isl=isl, osl=1, prefix=prefix) - /// // take ONLY "context_attention", divide by ceil(isl/ctx) + /// // prefix == 0 or ctx < isl_new: + /// // run_static(batch=ceil(ctx/isl_new), isl=isl, osl=1, prefix) + /// // prefix > 0 and ctx >= isl_new (fill = (ctx % isl_new) / isl_new): + /// // (1 - fill) * run_static(batch=floor(ctx/isl_new), isl, prefix) + /// // + fill * run_static(batch=floor(ctx/isl_new) + 1, isl, prefix) + /// // i.e. the last, partial request weighs its fill fraction + /// // take ONLY "context_attention", divide by ceil(isl_new/ctx) /// // Pass 3 — decode attention (only when gen_tokens > 0): /// // run_static(batch=gen, isl=isl+osl//2, osl=2, mode=static_gen) /// // -> one step at s = isl + osl//2 + 1 with the (nextn+1) batch /// // take ONLY "generation_attention" /// ``` /// - /// Note the Python conventions this deliberately preserves (they differed - /// from the pre-rewrite FPM packing): pass 1 uses - /// `ctx + gen * (nextn + 1)` tokens (the speculative-progress model — - /// every decode request verifies one target plus all drafts in the - /// combined pass, mirroring Python `run_mixed`'s `decode_query_tokens`), - /// the cached prefix multiplier is `floor(ctx/isl)` (not ceil), and the - /// pass-3 kv position carries `_run_generation_phase`'s `+1`. + /// Pass 1 uses `ctx + gen * (nextn + 1)` tokens (the speculative-progress + /// model — every decode request verifies one target plus all drafts in + /// the combined pass, mirroring `run_mixed`'s `decode_query_tokens`) and + /// the pass-3 kv position carries `_run_generation_phase`'s `+1`. The + /// prefix-free pass-2 packing is the legacy `ceil(ctx/isl)` and stays + /// bit-for-bit (frozen goldens); see [`Self::context_attention_groups`]. /// /// The imbalance-correction scales mirror the `RuntimeConfig` fields /// Python threads into each pass (`base_backend.py:950-1043`). @@ -829,9 +836,14 @@ impl Engine { /// decode_attention]`` for one mixed engine iteration — the three passes /// of the `_get_mix_step_latency` composition reported separately: pass 1 /// is the shared non-attention work, pass 2 the context-attention slice - /// (already divided by `ceil(isl/ctx)`), pass 3 the decode-attention - /// slice. [`Engine::mixed_step_latency`] is their sum; the agg - /// speculative scheduler consumes the components. + /// (already divided by `ceil((isl - prefix)/ctx)`), pass 3 the + /// decode-attention slice. [`Engine::mixed_step_latency`] is their sum; + /// the agg speculative scheduler consumes the components. + /// + /// `ctx_tokens` budgets UNCACHED (new) prefill tokens — what SGLang + /// `--chunked-prefill-size`, vLLM `max_num_batched_tokens` and TRT-LLM + /// `max_num_tokens` cap — so a request with `prefix` cached tokens + /// contributes `isl - prefix` tokens to it; the prefix is KV context only. pub fn mixed_step_breakdown( &self, ctx_tokens: u32, @@ -854,10 +866,41 @@ impl Engine { ) } + /// `(requests, weight)` batches the context-attention pass prices for a + /// `ctx_tokens` budget of uncached tokens when every request carries + /// `isl_new` new tokens over its cached prefix; the weighted results sum + /// to the pass. + /// + /// With a cached prefix the budget is filled the way the schedulers fill + /// it: `floor(ctx/isl_new)` complete requests plus ONE partial request + /// of the remaining `ctx % isl_new` tokens. The partial request weighs + /// its fill fraction of one more batched request — the pass is the + /// convex combination of the floor-packed and the ceil-packed batch, so + /// it is exact at full fills and tends to the floor-packed batch as the + /// remainder vanishes. The perf tables are per-launch batch measurements: + /// pricing the remainder as a standalone query would charge a second + /// kernel floor and forfeit the batch efficiency a varlen prefill launch + /// actually has (a 256-token standalone query costs more than adding a + /// whole 1792-token request to the batch), while `ceil(ctx/isl_new)` + /// complete requests overstate a non-multiple budget by up to one whole + /// request. At `prefix == 0` the legacy `ceil(ctx/isl)` packing is kept + /// so the frozen prefix-free goldens hold bit-for-bit; with `ctx < + /// isl_new` one whole request is priced and the caller divides by the + /// chunk count `ceil(isl_new/ctx)`. + fn context_attention_groups(ctx_tokens: u32, isl_new: u32, prefix: u32) -> Vec<(u32, f64)> { + let complete = ctx_tokens / isl_new; + let partial = ctx_tokens % isl_new; + if prefix == 0 || complete == 0 || partial == 0 { + return vec![(ctx_tokens.div_ceil(isl_new), 1.0)]; + } + let fill = f64::from(partial) / f64::from(isl_new); + vec![(complete, 1.0 - fill), (complete + 1, fill)] + } + /// [`Self::mixed_step_breakdown`] with a per-op sink. The sink observes /// `(pass, op, result)` for every queried op with RAW (undivided) pass-2 - /// values; the per-op wrapper applies the `ceil(isl/ctx)` division to the - /// FOLDED entries (fold-then-divide, matching Python and the scalar + /// values; the per-op wrapper applies the `ceil((isl - prefix)/ctx)` + /// division to the FOLDED entries (fold-then-divide, matching the scalar /// bucket bit-for-bit). #[allow(clippy::too_many_arguments)] fn mixed_step_breakdown_with( @@ -912,14 +955,16 @@ impl Engine { "V4.1 prefill requires isl > prefix".into(), )); } - // The SDK's packed context count includes prefix for complete - // requests. A remainder describes this iteration's partial extend. + // `ctx_tokens` budgets UNCACHED (new) tokens, so complete requests + // pack by `isl - prefix` new tokens each. A remainder describes + // this iteration's partial extend. + let isl_new = isl.saturating_sub(prefix).max(1); let mut prefills = Vec::with_capacity(2); - if ctx_tokens / isl > 0 { - prefills.push((ctx_tokens / isl, isl - prefix, prefix)); + if ctx_tokens / isl_new > 0 { + prefills.push((ctx_tokens / isl_new, isl_new, prefix)); } - if ctx_tokens % isl > 0 { - prefills.push((1, ctx_tokens % isl, prefix)); + if ctx_tokens % isl_new > 0 { + prefills.push((1, ctx_tokens % isl_new, prefix)); } return self.dsv41_mixed_workload( &prefills, @@ -930,26 +975,45 @@ impl Engine { on_op, ); } - // Python divides by `isl` (`floor(ctx/isl)`, `ceil(ctx/isl)`) without - // a guard — callers always pass isl >= 1. Clamp to avoid a Rust - // div-by-zero panic on degenerate input Python would crash on. + // Callers always pass isl >= 1; clamp to avoid a div-by-zero panic + // on degenerate input. let isl = isl.max(1); + // `ctx_tokens` budgets UNCACHED (new) prefill tokens, so a request + // contributes `isl_new = isl - prefix` tokens to it and the cached + // prefix is KV context for the attention pass only. The scheduling + // layer (`run_agg`) packs requests by the same `isl_new`. + if ctx_tokens > 0 && prefix >= isl { + return Err(AicError::InvalidEngineConfig(format!( + "isl must be greater than 0 after removing prefix, but got {}", + isl as i64 - prefix as i64 + ))); + } + let isl_new = isl.saturating_sub(prefix).max(1); // ---- Pass 1: combined non-attention work ---- // Speculative progress model: every decode request verifies one // target token plus all scheduled drafts, so the combined pass sees // `gen * (nextn + 1)` decode tokens (mirrors Python `run_mixed`'s // `decode_query_tokens`). Acceptance does not reduce this - // current-iteration work. + // current-iteration work. Every context token in the budget is a + // new token, so no prefix is subtracted from the token count. let decode_query_tokens = gen_tokens.saturating_mul(self.nextn.saturating_add(1)); let combined = ctx_tokens + decode_query_tokens; - let prefix1 = prefix * (ctx_tokens / isl); // prefix * floor(ctx/isl) - if prefix1 >= combined { - return Err(AicError::InvalidEngineConfig(format!( - "isl must be greater than 0 after removing prefix, but got {}", - combined as i64 - prefix1 as i64 - ))); - } + // The cached prefix still reaches this pass as KV context: an op that + // folds attention into a module outside the `context_attention` + // name (DeepSeek/Kimi `context_mla_block`) reads it through `prefix`, + // while token-major ops (GEMM, MoE, comm, norms) only read the token + // count. DSA/DSV4/MSA modules are named `context_attention` and get + // the prefix in pass 2. The budget holds `floor(ctx / isl_new)` + // complete requests, at least the one whose chunk it carries; as in + // the legacy query, they are priced as one sequence over their + // summed prefixes. With `ctx == isl_new` this is exactly the legacy + // `(combined - prefix, prefix)` query, and prefix 0 is unchanged. + let prefix1 = if ctx_tokens == 0 || prefix == 0 { + 0 + } else { + prefix.saturating_mul((ctx_tokens / isl_new).max(1)) + }; let mut shared_non_attention = 0.0; for op in &self.context_ops { // Only target operations share a forward across prefill and @@ -961,7 +1025,7 @@ impl Engine { op, &self.db, 1, - combined - prefix1, + combined, prefix1, seq_imbalance_correction_scale, None, @@ -971,19 +1035,16 @@ impl Engine { } // ---- Pass 2: context attention at the prefill shape ---- - // Python: batch = ceil(ctx/isl), effective_isl = isl - prefix, then - // latency["context_attention"] / ceil(isl/ctx). With ctx_tokens == 0 - // Python's `np.ceil(isl/0)` is +inf and the division yields 0 — skip. + // The weighted request batches filling the budget + // (`context_attention_groups`), each of `isl_new` new tokens over + // `prefix` cached ones; the fold is then divided by ceil(isl_new/ctx), + // the chunk count when one request's uncached prefill spans several + // steps. With ctx_tokens == 0 that division would be by +inf and + // yield 0 — skip. let mut context_attention = 0.0_f64; if ctx_tokens > 0 { - if prefix >= isl { - return Err(AicError::InvalidEngineConfig(format!( - "isl must be greater than 0 after removing prefix, but got {}", - isl as i64 - prefix as i64 - ))); - } - let batch2 = ctx_tokens.div_ceil(isl); - let scale2 = isl.div_ceil(ctx_tokens) as f64; + let groups = Self::context_attention_groups(ctx_tokens, isl_new, prefix); + let scale2 = isl_new.div_ceil(ctx_tokens) as f64; let mut attn = 0.0; for op in &self.context_ops { if !op.is_context_attention() && !op.name().starts_with("draft_") { @@ -992,22 +1053,23 @@ impl Engine { // Draft prefill uses the same whole-prefill amortization // as target attention, independently of decode work. It // never sees the combined target verification token count. - let result = query_context_op( - op, - &self.db, - batch2, - isl - prefix, - prefix, - seq_imbalance_correction_scale, - None, - )?; - attn += result.latency_ms; - // RAW results to the sink; the per-op wrapper divides the - // FOLDED values by scale2 with one true division per name - // (Python folds `context_attention` into one key, then - // `latency_dict["context_attention"] / scale_factor` — - // fold-then-divide, `base_backend.py:1244-1246`). - on_op(MixedPass::ContextAttention, op, result); + for &(batch2, weight) in &groups { + let result = query_context_op( + op, + &self.db, + batch2, + isl_new, + prefix, + seq_imbalance_correction_scale, + None, + )? + .scaled(weight); + attn += result.latency_ms; + // RAW results to the sink; the per-op wrapper divides + // the FOLDED values by scale2 with one true division per + // name (fold-then-divide, matching this scalar bucket). + on_op(MixedPass::ContextAttention, op, result); + } } context_attention = attn / scale2; } @@ -1575,7 +1637,7 @@ impl Engine { /// [`Self::mixed_step_breakdown`] with the per-op values kept: /// `(shared_non_attention, context_attention, decode_attention)` lists of /// `(name, latency_ms, energy_wms, source)`. Context-attention entries - /// arrive already divided by the `ceil(isl/ctx)` scale. + /// arrive already divided by the `ceil((isl - prefix)/ctx)` scale. #[allow(clippy::too_many_arguments)] pub fn mixed_step_breakdown_per_op( &self, @@ -1696,10 +1758,12 @@ impl Engine { )?; let mut ctx_attn = ctx_attn.into_values(); if ctx_tokens > 0 && !self.has_dsv41_stages() { - // Mirror the scalar bucket and Python's fold-then-single-true- - // division (`base_backend.py:1244-1246`): one `/ scale2` per - // folded name, never a per-entry reciprocal multiply. - let scale2 = isl.max(1).div_ceil(ctx_tokens) as f64; + // Mirror the scalar bucket's fold-then-single-true-division: one + // `/ scale2` per folded name, never a per-entry reciprocal + // multiply. `scale2` is the chunk count of the UNCACHED prefill + // (`isl - prefix`), exactly as in `mixed_step_breakdown_with`. + let isl_new = isl.max(1).saturating_sub(prefix).max(1); + let scale2 = isl_new.div_ceil(ctx_tokens) as f64; for entry in &mut ctx_attn { entry.1 /= scale2; entry.2 /= scale2; @@ -2372,6 +2436,30 @@ mod tests { Engine::build(spec, Arc::new(db)).unwrap() } + /// `build_engine` with the context attention renamed so it rides in + /// pass 1, as production `context_mla_block` does: an op outside the + /// `context_attention` name whose cost reads the cached prefix. + fn build_engine_with_module_attention() -> Engine { + let db = PerfDatabase::load_resolved( + &systems_root(), + "b200_sxm", + "vllm", + "0.24.0", + true, + false, + false, + ) + .unwrap(); + let mut ops = context_ops(); + for op in &mut ops { + if let Op::ContextAttention(attention) = op { + attention.name = "context_attention_module".into(); + } + } + let spec = EngineSpec::new(fixture_engine_config(None), ops, generation_ops()); + Engine::build(spec, Arc::new(db)).unwrap() + } + fn runtime(batch_size: u32, isl: u32, osl: u32) -> RuntimeConfig { RuntimeConfig { batch_size, @@ -3021,6 +3109,354 @@ mod tests { assert_eq!(ms, breakdown[0]); } + /// Legacy (pre-uncached-budget) three-pass composition, hand-rolled from + /// the fixture ops: pass 1 priced `ctx + decode - prefix * floor(ctx/isl)` + /// new tokens over that prefix, pass 2 priced `ceil(ctx/isl)` requests of + /// `isl - prefix` new tokens divided by `ceil(isl/ctx)`. At `prefix == 0` + /// it coincides with the current semantics, so it doubles as the + /// bit-for-bit oracle for the prefix-free path. + fn legacy_mixed_reference( + engine: &Engine, + ctx: u32, + decode: u32, + isl: u32, + osl: u32, + prefix: u32, + ) -> [f64; 3] { + let prefix1 = prefix * (ctx / isl); + let mut shared = 0.0; + let mut attn = 0.0; + for op in &engine.context_ops { + if op.is_context_attention() { + attn += query_context_op( + op, + &engine.db, + ctx.div_ceil(isl), + isl - prefix, + prefix, + 1.0, + None, + ) + .unwrap() + .latency_ms; + } else { + shared += query_context_op( + op, + &engine.db, + 1, + ctx + decode - prefix1, + prefix1, + 1.0, + None, + ) + .unwrap() + .latency_ms; + } + } + [ + shared, + attn / isl.div_ceil(ctx) as f64, + decode_attention_reference(engine, decode, isl, osl), + ] + } + + /// Pass 3 is independent of the budget semantics: decode attention for + /// `decode` requests at `s = isl + osl/2 + 1` (fixture nextn = 0). + fn decode_attention_reference( + engine: &Engine, + decode_requests: u32, + isl: u32, + osl: u32, + ) -> f64 { + let mut decode = 0.0; + for op in &engine.generation_ops { + if op.is_generation_attention() { + decode += query_generation_op( + op, + &engine.db, + decode_requests, + 1, + isl + osl / 2 + 1, + 1.0, + 0, + None, + ) + .unwrap() + .latency_ms; + } + } + decode + } + + /// Static prefill context attention of `batch` requests that each extend + /// `prefix` cached tokens by `new_tokens` — the per-request primitive the + /// mixed-step compositions below are hand-assembled from. + fn static_context_attention(engine: &Engine, batch: u32, new_tokens: u32, prefix: u32) -> f64 { + engine + .context_ops + .iter() + .filter(|op| op.is_context_attention()) + .map(|op| { + query_context_op(op, &engine.db, batch, new_tokens, prefix, 1.0, None) + .unwrap() + .latency_ms + }) + .sum() + } + + /// Pass 1 as one forward of `tokens` new tokens (budget plus decode) + /// that reads `kv_prefix` cached tokens as context: every non-attention + /// context op at that shape. + fn shared_pass_reference(engine: &Engine, tokens: u32, kv_prefix: u32) -> f64 { + engine + .context_ops + .iter() + .filter(|op| !op.is_context_attention()) + .map(|op| { + query_context_op(op, &engine.db, 1, tokens, kv_prefix, 1.0, None) + .unwrap() + .latency_ms + }) + .sum() + } + + #[test] + fn mixed_step_prefix_zero_is_bit_for_bit_the_legacy_composition() { + // ctx < isl (chunked), ctx == isl, ctx > isl at a non-multiple, and a + // decode-heavy shape: at prefix == 0 the uncached-budget semantics + // must reproduce the legacy composition exactly. + let engine = build_engine(None); + for (ctx, decode, isl, osl) in [ + (512_u32, 4_u32, 4096_u32, 128_u32), + (1024, 2, 1024, 8), + (3000, 5, 2048, 64), + (256, 64, 4096, 512), + ] { + let got = engine + .mixed_step_breakdown(ctx, decode, isl, osl, 0, 1.0, 1.0) + .unwrap(); + let legacy = legacy_mixed_reference(&engine, ctx, decode, isl, osl, 0); + assert_eq!(got[1], legacy[0], "pass 1 ctx={ctx} isl={isl}"); + assert_eq!(got[2], legacy[1], "pass 2 ctx={ctx} isl={isl}"); + assert_eq!(got[3], legacy[2], "pass 3 ctx={ctx} isl={isl}"); + assert_eq!(got[0], got[1] + got[2] + got[3]); + } + } + + #[test] + fn mixed_step_chunked_prefill_with_prefix_follows_uncached_isl() { + // ctx < isl with a cached prefix: the chunk count is + // ceil(isl_new/ctx) = ceil(900/300) = 3, not the legacy + // ceil(isl/ctx) = 4. Pass 1 prices the 300 + 7 new tokens over the + // chunked request's 100 cached tokens (the legacy pass dropped them: + // floor(300/1000) == 0). + let engine = build_engine(None); + let (ctx, decode, isl, osl, prefix) = (300_u32, 7_u32, 1000_u32, 64_u32, 100_u32); + let got = engine + .mixed_step_breakdown(ctx, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + let legacy = legacy_mixed_reference(&engine, ctx, decode, isl, osl, prefix); + assert_eq!(got[1], shared_pass_reference(&engine, 307, 100)); + assert!(got[2] > 0.0); + // One request of 900 new tokens over 100 cached, amortized over its + // 3 chunks of 300 (legacy: over 4). + assert_eq!(got[2], static_context_attention(&engine, 1, 900, 100) / 3.0); + assert!((got[2] - legacy[1] * 4.0 / 3.0).abs() < 1e-9 * got[2]); + assert_eq!(got[3], legacy[2]); + } + + #[test] + fn mixed_step_full_prefill_with_prefix_prices_every_budget_token_as_new() { + // ctx >= isl with a cached prefix: the legacy pass 1 subtracted + // prefix * floor(ctx/isl) from the budget (double-crediting the cache + // once the step count already uses isl_new); now every one of the + // 2048 ctx tokens is new, and the budget's floor(2048/768) = 2 + // complete requests contribute their 2 * 256 cached tokens as KV + // context only. Pass 2 fills the 2048-token budget with requests + // of 768 new tokens over 256 cached: 2 complete ones (1536) plus a + // partial request of the remaining 512, weighing 512/768 = 2/3 of a + // third batched request — not ceil(2048/768) = 3 complete requests. + let engine = build_engine(None); + let (ctx, decode, isl, osl, prefix) = (2048_u32, 2_u32, 1024_u32, 8_u32, 256_u32); + let got = engine + .mixed_step_breakdown(ctx, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + let legacy = legacy_mixed_reference(&engine, ctx, decode, isl, osl, prefix); + assert_eq!(got[1], shared_pass_reference(&engine, 2050, 512)); + assert_ne!( + got[1], legacy[0], + "pass 1 must no longer subtract the prefix from the token count" + ); + let two = static_context_attention(&engine, 2, 768, 256); + let three = static_context_attention(&engine, 3, 768, 256); + let expected_attn = two / 3.0 + three * 2.0 / 3.0; + assert!((got[2] - expected_attn).abs() < 1e-12 * expected_attn); + assert!(two < got[2] && got[2] < three); + // The legacy composition packed ceil(2048/1024) = 2 requests on the + // full isl and never saw the third one the budget partly fills. + assert_eq!(legacy[1], two); + assert_eq!(got[3], legacy[2]); + // The per-op surface folds to the same three buckets. + let (shared, ctx_attn, dec_attn) = engine + .mixed_step_breakdown_per_op(ctx, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + let sum = |rows: &[PerOpValue]| rows.iter().map(|r| r.1).sum::(); + assert!((sum(&shared) - got[1]).abs() < 1e-12); + assert!((sum(&ctx_attn) - got[2]).abs() < 1e-12); + assert!((sum(&dec_attn) - got[3]).abs() < 1e-12); + } + + #[test] + fn mixed_step_pass_one_reads_the_cached_prefix_of_the_complete_requests() { + // A module-attention op priced in pass 1 must see the cached prefix + // of the floor(ctx/isl_new) complete requests the budget holds (at + // least the one being chunked) and nothing when no prefill runs. + // isl 1000, prefix 100 -> isl_new 900; decode 7. + let engine = build_engine_with_module_attention(); + let (decode, isl, osl, prefix) = (7_u32, 1000_u32, 64_u32, 100_u32); + for (ctx, kv_prefix) in [ + (300_u32, 100_u32), // one chunk of one request + (900, 100), // exactly one request + (2250, 200), // two complete requests plus a partial one + ] { + let got = engine + .mixed_step_breakdown(ctx, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + let expected = shared_pass_reference(&engine, ctx + decode, kv_prefix); + assert_eq!(got[1], expected, "pass 1 ctx={ctx}"); + assert!( + got[1] > shared_pass_reference(&engine, ctx + decode, 0), + "the module op must read the prefix (ctx={ctx})" + ); + // Context attention proper is absent from this fixture. + assert_eq!(got[2], 0.0); + } + // No prefill scheduled: the prefix is inert in pass 1 as well. + let decode_only = engine + .mixed_step_breakdown(0, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + assert_eq!(decode_only[1], shared_pass_reference(&engine, decode, 0)); + } + + #[test] + fn mixed_step_prefix_at_or_beyond_isl_rejects_prefill_only() { + let engine = build_engine(None); + assert!( + engine + .mixed_step_breakdown(64, 2, 512, 8, 512, 1.0, 1.0) + .is_err() + ); + assert!( + engine + .mixed_step_breakdown(64, 2, 512, 8, 600, 1.0, 1.0) + .is_err() + ); + // A decode-only iteration schedules no prefill; the prefix is inert. + let decode_only = engine + .mixed_step_breakdown(0, 2, 512, 8, 512, 1.0, 1.0) + .unwrap(); + assert_eq!( + decode_only, + engine + .mixed_step_breakdown(0, 2, 512, 8, 0, 1.0, 1.0) + .unwrap() + ); + assert_eq!(decode_only[2], 0.0); + assert!(decode_only[3] > 0.0); + } + + #[test] + fn mixed_step_partial_request_prices_an_isl_sized_budget_between_floor_and_ceil_packing() { + // The regime every explicit `--ctx-tokens ` run lands in: the + // 2048-token budget holds one complete request of 1920 new tokens + // plus a 128-token partial request, both over 128 cached tokens. The + // partial request weighs 128/1920 = 1/15 of a second batched + // request; ceil(2048/1920) = 2 complete requests would charge 3840 + // new tokens of attention. (The fixture database is launch-bound at + // these toy shapes -- one 128-token request costs a third of a + // 2048-token one -- so the warm step is not compared with the cold + // one here; the real-engine CLI regression covers that direction.) + let engine = build_engine(None); + let (ctx, decode, isl, osl) = (2048_u32, 4_u32, 2048_u32, 512_u32); + let cold = engine + .mixed_step_breakdown(ctx, decode, isl, osl, 0, 1.0, 1.0) + .unwrap(); + let warm = engine + .mixed_step_breakdown(ctx, decode, isl, osl, 128, 1.0, 1.0) + .unwrap(); + // Pass 1 sees 2048 + 4 new tokens either way; warm, the one complete + // request's 128 cached tokens ride along as KV context. + assert_eq!(cold[1], shared_pass_reference(&engine, 2052, 0)); + assert_eq!(warm[1], shared_pass_reference(&engine, 2052, 128)); + assert_eq!(warm[3], cold[3]); + let one = static_context_attention(&engine, 1, 1920, 128); + let two = static_context_attention(&engine, 2, 1920, 128); + let expected_attn = one * 14.0 / 15.0 + two / 15.0; + assert!((warm[2] - expected_attn).abs() < 1e-12 * expected_attn); + assert!(one < warm[2] && warm[2] < two); + // The per-op surface folds both requests under one name and matches. + let (_, ctx_attn, _) = engine + .mixed_step_breakdown_per_op(ctx, decode, isl, osl, 128, 1.0, 1.0) + .unwrap(); + assert!((ctx_attn.iter().map(|r| r.1).sum::() - warm[2]).abs() < 1e-12); + } + + #[test] + fn mixed_step_prefix_free_non_multiple_budget_keeps_legacy_ceil_packing() { + // prefix == 0 keeps ceil(3000/2048) = 2 complete requests (frozen + // goldens); the partial-request pricing is gated on a cached prefix. + let engine = build_engine(None); + let got = engine + .mixed_step_breakdown(3000, 5, 2048, 64, 0, 1.0, 1.0) + .unwrap(); + assert_eq!(got[2], static_context_attention(&engine, 2, 2048, 0)); + assert_eq!( + Engine::context_attention_groups(3000, 2048, 0), + vec![(2, 1.0)] + ); + // 3000 = 1 * 2047 + 953: the partial request weighs 953/2047. + let fill = 953.0 / 2047.0; + assert_eq!( + Engine::context_attention_groups(3000, 2047, 1), + vec![(1, 1.0 - fill), (2, fill)] + ); + assert_eq!( + Engine::context_attention_groups(4094, 2047, 1), + vec![(2, 1.0)] + ); + assert_eq!( + Engine::context_attention_groups(300, 900, 100), + vec![(1, 1.0)] + ); + } + + #[test] + fn dsv41_complete_extends_with_prefix_pack_by_uncached_tokens() { + // 4096-token requests with 3584 cached: the 1024-token budget holds + // two complete 512-token extends; 1100 holds those plus a 76-token + // partial extend. Encoder stages see every new token, the bounded + // decoder stage min(extend, 128) per request, decode adds 3 requests + // per stage (linear memory probes, one unit per token). + let engine = dsv41_probe_engine(true); + let unit = dsv41_probe_token_ms(&engine); + let parts = engine + .mixed_step_breakdown(1024, 3, 4096, 32, 3584, 1.0, 1.0) + .unwrap(); + assert!((parts[2] / unit - (1024.0 + 2.0 * 128.0)).abs() < 1e-9); + assert!((parts[1] / unit - (1024.0 + 2.0 * 128.0 + 6.0)).abs() < 1e-9); + let parts = engine + .mixed_step_breakdown(1100, 3, 4096, 32, 3584, 1.0, 1.0) + .unwrap(); + assert!((parts[2] / unit - (1100.0 + 2.0 * 128.0 + 76.0)).abs() < 1e-9); + assert!((parts[1] / unit - (1100.0 + 2.0 * 128.0 + 76.0 + 6.0)).abs() < 1e-9); + let (shared, context, decode) = engine + .mixed_step_breakdown_per_op(1100, 3, 4096, 32, 3584, 1.0, 1.0) + .unwrap(); + assert!((shared.iter().map(|v| v.1).sum::() - parts[1]).abs() < 1e-12); + assert!((context.iter().map(|v| v.1).sum::() - parts[2]).abs() < 1e-12); + assert!((decode.iter().map(|v| v.1).sum::() - parts[3]).abs() < 1e-12); + } + #[test] fn mixed_draft_phases_preserve_native_results_and_target_composition() { use crate::operators::op::TokenScaleOp; @@ -3070,26 +3506,30 @@ mod tests { gen_ops.extend(gen_draft.clone()); let engine = Engine::build(EngineSpec::new(config, ctx_ops, gen_ops), db).unwrap(); - for (ctx, generation, prefix) in - [(0_u32, 7, 0), (128, 0, 64), (128, 7, 64), (8192, 7, 64)] - { + // Requests pack by their UNCACHED prefill length (4096 - 64 = + // 4032): a 128-token budget is one chunk of one request, an + // 8192-token budget holds two complete requests plus a 128-token + // partial one weighing 128/4032 of a third batched request. + let fill = 128.0 / 4032.0; + for (ctx, generation, prefix, batches) in [ + (0_u32, 7, 0, &[][..]), + (128, 0, 64, &[(1_u32, 1.0_f64)][..]), + (128, 7, 64, &[(1, 1.0)][..]), + (8192, 7, 64, &[(2, 1.0 - fill), (3, fill)][..]), + ] { + let isl_new = 4096 - prefix; let mut ctx_expected = PerformanceResult::zero(); if ctx > 0 { for op in &ctx_draft { - ctx_expected = ctx_expected.plus( - query_context_op( - op, - &engine.db, - ctx.div_ceil(4096), - 4096 - prefix, - prefix, - 1.0, - None, - ) - .unwrap(), - ); + for &(batch, weight) in batches { + ctx_expected = ctx_expected.plus( + query_context_op(op, &engine.db, batch, isl_new, prefix, 1.0, None) + .unwrap() + .scaled(weight), + ); + } } - ctx_expected = ctx_expected.scaled(1.0 / 4096_u32.div_ceil(ctx) as f64); + ctx_expected = ctx_expected.scaled(1.0 / isl_new.div_ceil(ctx) as f64); } let mut gen_expected = PerformanceResult::zero(); if generation > 0 { @@ -3143,7 +3583,7 @@ mod tests { ) .unwrap(); if ctx > 0 { - observed_ctx = observed_ctx.scaled(1.0 / 4096_u32.div_ceil(ctx) as f64); + observed_ctx = observed_ctx.scaled(1.0 / isl_new.div_ceil(ctx) as f64); } // Includes energy, source, SOL components and fallback records. assert_eq!(observed_ctx, ctx_expected); diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 060b898e3..42e417f99 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -565,7 +565,8 @@ impl AicEngine { /// ``(shared_non_attention, context_attention, decode_attention)`` lists /// of ``(name, latency_ms, energy_wms, source)`` tuples. The /// context-attention entries arrive already divided by the - /// ``ceil(isl/ctx)`` scale, so each list sums to its breakdown bucket. + /// ``ceil((isl - prefix)/ctx)`` scale, so each list sums to its breakdown + /// bucket. ``ctx_tokens`` budgets UNCACHED prefill tokens. #[pyo3(signature = (ctx_tokens, gen_tokens, isl, osl, prefix=0, seq_imbalance_correction_scale=1.0, gen_seq_imbalance_correction_scale=1.0))] diff --git a/docs/cli/legacy-aic-user-guide.md b/docs/cli/legacy-aic-user-guide.md index 6370aa29d..e5a7a8527 100644 --- a/docs/cli/legacy-aic-user-guide.md +++ b/docs/cli/legacy-aic-user-guide.md @@ -138,7 +138,7 @@ aiconfigurator cli estimate --model-path Qwen/Qwen3-32B --system h200_sxm --tp-s - `--isl`: Input sequence length. Default: `1024` - `--osl`: Output sequence length. Default: `1024` - `--batch-size` (alias `--bs`): Batch size (max concurrent requests, used for agg/static). Default: `128` -- `--ctx-tokens`: Context tokens budget for IFB scheduling (agg only). Default: same as ISL +- `--ctx-tokens`: The scheduler's per-step budget of uncached (new) context tokens for IFB scheduling (agg only): SGLang `--chunked-prefill-size`, vLLM `max_num_batched_tokens`, the TRT-LLM scheduler's `max_num_tokens` (TRT-LLM's build-time `max_num_tokens` for activation memory is a separate setting). Default: ISL (including any visual tokens) minus `--prefix` (one request's full uncached prefill per step) - `--tp-size` (alias `--tp`): Tensor parallelism size. Default: `1` - `--pp-size` (alias `--pp`): Pipeline parallelism size. Default: `1` - `--attention-dp-size` (alias `--dp`): Attention data parallelism size. Default: `1` @@ -149,7 +149,7 @@ aiconfigurator cli estimate --model-path Qwen/Qwen3-32B --system h200_sxm --tp-s - `--fmha-quant-mode`: FMHA quantization mode (auto-inferred if omitted). DeepSeek-V3 / Kimi-K2.5 context attention (MLA prefill) has no fp8 FMHA data — auto-inferred fp8 is downgraded to `bfloat16` for context-touching estimates (agg, prefill, static, static_ctx, AFD prefill), while decode/generation keeps fp8. Explicitly passing `fp8` for a context estimate is rejected with a clear error. - `--moe-quant-mode`: MoE quantization mode (auto-inferred if omitted) - `--comm-quant-mode`: Communication quantization mode (auto-inferred; default `half`) -- `--prefix`: Prefix cache length (subset of ISL already cached per request). Default: `0` +- `--prefix`: Prefix cache length (subset of ISL already cached per request). Default: `0`. Cached tokens are KV context only: the mixed-step schedule packs requests by their uncached `ISL - prefix` tokens, so at a fixed `--ctx-tokens` a larger prefix means fewer mix steps and higher throughput. TTFT drops when the uncached prefill needs fewer chunks (`ceil((ISL - prefix) / ctx_tokens)`) or queues behind fewer steps; a step that already holds one whole request costs about the same as without the prefix (its new tokens attend over more cached KV), so TTFT is roughly flat in that regime. A budget that is not a multiple of `ISL - prefix` is filled with one partial request, never rounded up to a whole one; the per-step request count is capped at `--batch-size` - `--nextn`: MTP draft length, or `auto` to use the checkpoint's `num_nextn_predict_layers` (absent/0 keeps MTP disabled). Default: `0`; MTP is never enabled implicitly — pass `--nextn` explicitly to model it - `--nextn-accepted`: Average accepted draft tokens per decode step (`0 <= nextn_accepted <= nextn`). Required when the draft depth is > 0 (including via `--nextn auto`); use a measured value from your deployment - `--stride`: (static modes only) OSL-sweep stride used by `run_static`; ignored for `agg`/`disagg`. Default: `32` @@ -451,7 +451,7 @@ Beyond `--ttft`, `--tpot`, `--isl`, `--osl`, and `--prefix`, `default` mode acce - `--backend-version`: Backend database version — a queryable slot version or the aliases `current` / `previous` / `next`. Default: `current`. - `--free-gpu-memory-fraction`: Fraction of free GPU memory TRT-LLM allocates for KV cache (default: `1.0`). Filters batch sizes that would exceed KV cache capacity. - `--max-seq-len`: TRT-LLM `--max_seq_len` (default: `isl + osl`). Controls how many KV blocks are pre-allocated per sequence; set to match your deployment for accurate KV-capacity filtering. -- `--enable-chunked-prefill`: Enable chunked prefill for a finer-grained context-token sweep. When off (default), the context-token stride is aligned to ISL for faster sweeping. +- `--enable-chunked-prefill`: Enable chunked prefill for a finer-grained context-token sweep. When off (default), the context-token stride is aligned to the uncached ISL (ISL minus `--prefix`) for faster sweeping. - `--enable-wideep`: **Deprecated and ignored for large-EP modeling** (accepted with a one-time warning). On SGLang, it still narrows the default `moe_tp` candidates to `[1]`; explicit `*_moe_tp_candidates` values take precedence. Large-EP (wideEP) is explored automatically — see the note below. - `--attention-backend`: Attention kernel backend — one of `fa3`, `triton`, `trtllm_mha`, `flashinfer`, `fla`, or `default`. It applies to every model graph with standard dense `ContextAttention`/`GenerationAttention` operations and to supported DeepSeek MLA/WideEP paths. Support is backend-, performance-table-, and version-specific; unsupported named values fail closed. For modeling, unset or `default` uses the mapped framework default when available and otherwise the safe `default` fallback. SGLang WideEP maps unset/`default` to `flashinfer` and also supports `fa3`. The deployment generator emits supported named SGLang values, omits unset/`default`, and rejects `fla` for SGLang 0.5.14. - `--moe-backend`: Explicit SGLang MoE backend. `megamoe` is a real kernel selection (use it to model DeepSeek-V4 MegaMoE on Blackwell); `deepep_moe` is deprecated: it is ignored for modeling (large-EP is explored automatically from data coverage), but on SGLang it still narrows the default `moe_tp` candidates to `[1]` — explicit `*_moe_tp_candidates` always win. diff --git a/python/aisimulate/src/aisimulate/legacy_cli/api.py b/python/aisimulate/src/aisimulate/legacy_cli/api.py index 76d27ecc4..8e0c11c93 100644 --- a/python/aisimulate/src/aisimulate/legacy_cli/api.py +++ b/python/aisimulate/src/aisimulate/legacy_cli/api.py @@ -1154,8 +1154,13 @@ def cli_estimate( vLLM mm_encoder_tp_mode="data" / SGLang --mm-enable-dp-encoder semantics). False models the legacy TP-sharded encoder. batch_size: Batch size (max concurrent requests, used for agg mode). Default is 128. - ctx_tokens: Context tokens budget for IFB scheduling (agg mode only). - Default is None, which uses ``isl`` as the budget. + ctx_tokens: The scheduler's per-step budget of uncached (new) context + tokens for IFB scheduling (agg mode only): SGLang + ``--chunked-prefill-size``, vLLM ``max_num_batched_tokens``, the + TRT-LLM scheduler's ``max_num_tokens`` (TRT-LLM's build-time + ``max_num_tokens`` for activation memory is a separate kwarg). + Defaults to the effective isl minus ``prefix`` (one request's full + uncached prefill per step) when None. tp_size: Tensor parallelism size. Default is 1. Also serves as fallback for prefill/decode TP when their specific args are omitted. pp_size: Pipeline parallelism size. Default is 1. @@ -1760,7 +1765,9 @@ def _run_agg_estimate( model = get_model(model_path, model_config, backend_name) if ctx_tokens is None: - ctx_tokens = isl + BaseBackend._visual_context_tokens(model, runtime_config) + # The budget counts UNCACHED prefill tokens: "one full request per + # mixed step" is the effective isl minus the cached prefix. + ctx_tokens = max(isl + BaseBackend._visual_context_tokens(model, runtime_config) - int(prefix or 0), 1) database = load_database(system_name) backend = get_backend(backend_name) session = InferenceSession(model, database, backend) diff --git a/python/aisimulate/src/aisimulate/legacy_cli/main.py b/python/aisimulate/src/aisimulate/legacy_cli/main.py index e8de82df4..9440c2a67 100644 --- a/python/aisimulate/src/aisimulate/legacy_cli/main.py +++ b/python/aisimulate/src/aisimulate/legacy_cli/main.py @@ -602,7 +602,8 @@ def _add_default_mode_arguments(parser): action="store_true", default=False, help="Enable chunked prefill for finer-grained context token sweep during optimization. " - "When off (default), context token stride is aligned to ISL for faster sweeping.", + "When off (default), context token stride is aligned to the uncached ISL (ISL minus --prefix) " + "for faster sweeping.", ) parser.add_argument( "--free-gpu-memory-fraction", @@ -1054,7 +1055,12 @@ def _add_estimate_mode_arguments(parser): "--ctx-tokens", type=int, default=None, - help="Context tokens budget for IFB scheduling (agg only). Default: same as ISL.", + help=( + "The scheduler's per-step budget of uncached (new) context tokens for IFB scheduling (agg only): " + "SGLang --chunked-prefill-size / vLLM max_num_batched_tokens / TRT-LLM scheduler max_num_tokens " + "(TRT-LLM's build-time max_num_tokens for activation memory is separate). " + "Default: ISL (including any visual tokens) minus --prefix." + ), ) # Shared parallelism defaults (also used as fallback for prefill/decode-specific args) diff --git a/python/aisimulate/src/aisimulate/sdk/inference_session.py b/python/aisimulate/src/aisimulate/sdk/inference_session.py index 68a25f6e4..9721aad5d 100644 --- a/python/aisimulate/src/aisimulate/sdk/inference_session.py +++ b/python/aisimulate/src/aisimulate/sdk/inference_session.py @@ -127,7 +127,12 @@ def run_mixed( runtime_config: config.RuntimeConfig, step: MixedStepInput, ) -> StepEstimate: - """Estimate one mixed prefill/decode engine iteration.""" + """Estimate one mixed prefill/decode engine iteration. + + ``step.context_tokens`` is a budget of UNCACHED prefill tokens: with + ``runtime_config.prefix`` cached tokens, requests contribute their + ``isl - prefix`` new tokens to it (see ``MixedStepInput``). + """ return self._backend.run_mixed(self._model, self._database, runtime_config, step) def run_agg(self, runtime_config: config.RuntimeConfig, **kwargs) -> InferenceSummary: diff --git a/python/aisimulate/src/aisimulate/sdk/predict.py b/python/aisimulate/src/aisimulate/sdk/predict.py index edd0416e2..9a204cdac 100644 --- a/python/aisimulate/src/aisimulate/sdk/predict.py +++ b/python/aisimulate/src/aisimulate/sdk/predict.py @@ -125,7 +125,7 @@ def predict_agg_worker( ctx_tokens: Per-step context-token budget. Callers enumerate candidate values; chunked-prefill semantics are implicit in which values are tried (chunked off → ``ctx_tokens`` - constrained to multiples of isl). + constrained to multiples of the uncached ``isl - prefix``). predictor: Optional Predictor strategy. Defaults to the analytic zero-queue predictor (zero behavior change from direct backend.run_agg calls). Future Mocker / dynamic predictors diff --git a/python/aisimulate/src/aisimulate/sdk/sweep.py b/python/aisimulate/src/aisimulate/sdk/sweep.py index aa2945994..8dcabd08c 100644 --- a/python/aisimulate/src/aisimulate/sdk/sweep.py +++ b/python/aisimulate/src/aisimulate/sdk/sweep.py @@ -335,12 +335,19 @@ def _sweep_one_parallel_agg( # the legacy find_best_agg_result_under_constraints (its isl_eff) and as # run_agg's internal accounting. isl = runtime_config.isl + BaseBackend._visual_context_tokens(model, runtime_config) + # ctx_tokens budgets UNCACHED prefill tokens (run_agg's isl_new), so the + # feasibility guards and the capped-gen dedup pack requests by it. + isl_new = max(isl - int(runtime_config.prefix or 0), 1) osl = runtime_config.osl ttft_target = runtime_config.ttft tpot_target = runtime_config.tpot b_list = [b for b in _DEFAULT_AGG_BATCH_SCHEDULE if b <= max_batch_size] - ctx_tokens_list = _agg_ctx_tokens_list(isl, ctx_stride, enable_chunked_prefill) + # The grid is built on isl_new too, so its points are multiples of the + # uncached prefill and the guards below admit the same (b, requests) + # shapes as at prefix 0 (a grid on the full isl would be guarded out + # almost entirely for a heavily cached workload). + ctx_tokens_list = _agg_ctx_tokens_list(isl_new, ctx_stride, enable_chunked_prefill) # The capped-gen dedup below assumes the non-speculative schedule # (decode_iterations == osl). With speculative progress the boundary is @@ -357,18 +364,20 @@ def _sweep_one_parallel_agg( saw_model_fit = False saw_memory_fit = False perf_misses = 0 + guarded_in = 0 for b in b_list: for ctx_tokens in ctx_tokens_list: # batch / ctx_tokens balance guards (legacy semantics) - if b - np.ceil(ctx_tokens / isl) < 0: + if b - np.ceil(ctx_tokens / isl_new) < 0: break - if b > 1 and (b - np.ceil(ctx_tokens / isl) < 1): + if b > 1 and (b - np.ceil(ctx_tokens / isl_new) < 1): break + guarded_in += 1 # Skip equivalent gen_tokens slices to avoid recomputing the same point. if dedup_gen_slices: - balance_score = isl * b / ctx_tokens / osl + balance_score = isl_new * b / ctx_tokens / osl if balance_score > 1: gen_tokens = b // balance_score if gen_tokens > 1 and gen_tokens in capped_b: @@ -417,6 +426,15 @@ def _sweep_one_parallel_agg( results_per_ops_source.append(summary.get_per_ops_source()) results_moe_comm_fallbacks.append(merge_moe_comm_fallbacks(summary.get_moe_comm_fallbacks())) + if guarded_in == 0: + logger.warning( + "agg sweep: no (batch_size, ctx_tokens) point passed the batch/ctx guards " + "(isl=%s, prefix=%s, batch sizes %s, ctx grid %s)", + isl, + runtime_config.prefix, + b_list, + ctx_tokens_list, + ) if not results_dict_list: return pd.DataFrame(columns=common.ColumnsAgg), saw_model_fit, saw_memory_fit, perf_misses @@ -516,7 +534,7 @@ def sweep_agg( top_k: Per-(parallel, tpot) top-K rows to keep before concat. max_batch_size: Upper bound on batch size sweep. ctx_stride: Stride for ctx_tokens sweep. - enable_chunked_prefill: When False, ctx_tokens snaps to multiples of isl. + enable_chunked_prefill: When False, ctx_tokens snaps to multiples of the uncached isl (isl - prefix). free_gpu_memory_fraction: TRT-LLM-only KV cache fraction. max_seq_len: TRT-LLM-only per-slot KV cache budget. diff --git a/python/aisimulate/src/aisimulate/sdk/task_v2.py b/python/aisimulate/src/aisimulate/sdk/task_v2.py index a24afc41a..fcd78ea6f 100644 --- a/python/aisimulate/src/aisimulate/sdk/task_v2.py +++ b/python/aisimulate/src/aisimulate/sdk/task_v2.py @@ -3118,9 +3118,10 @@ def run_single_agg( Args: tp / pp / dp / moe_tp / moe_ep: parallelism for this single point. batch_size: concurrency (max in-flight requests). - ctx_tokens: per-step context-token budget for the IFB - scheduler. Defaults to ``self.isl`` (full prefill in - one step) -- matching ``cli_estimate`` semantics. + ctx_tokens: per-step budget of UNCACHED context tokens for the + IFB scheduler. Defaults to the effective isl minus the + cached prefix (one request's full uncached prefill per + step) -- matching ``cli_estimate`` semantics. Returns: Row dict in ``common.ColumnsAgg`` schema, equivalent to one @@ -3166,7 +3167,11 @@ def run_single_agg( backend=backend, database=database, runtime_config=runtime_config, - ctx_tokens=ctx_tokens if ctx_tokens is not None else self._prefill_effective_isl(runtime_config), + ctx_tokens=( + ctx_tokens + if ctx_tokens is not None + else max(self._prefill_effective_isl(runtime_config) - int(runtime_config.prefix or 0), 1) + ), predictor=self.predictor, speculative_profile=self.build_speculative_profile(), **backend_kwargs, diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py index 6eff2a9b0..4b3128d23 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py @@ -150,6 +150,10 @@ def _mix_step_gen_tokens(self, b: int, ctx_tokens: int, isl: int, decode_iterati generating). This method encodes the engine's scheduling policy for how many decode-phase requests participate alongside the prefilling request(s). + ``isl`` is the UNCACHED prefill length of one request (effective isl + minus the cached prefix): ``ctx_tokens`` budgets new tokens, so that is + the length requests pack by. + Subclasses should override to match their engine's scheduling behaviour. """ steps_to_finish_ctx = np.ceil(isl * b / ctx_tokens) @@ -1462,6 +1466,12 @@ def run_mixed( ``run_static`` / ``run_agg`` do, so direct callers (e.g. ``InferenceSession.run_mixed``) get the effective sequence length without pre-adjusting the config. + + ``step.context_tokens`` budgets UNCACHED prefill tokens (SGLang + ``--chunked-prefill-size``, vLLM ``max_num_batched_tokens``, the + TRT-LLM scheduler's ``max_num_tokens``). The full ``isl`` and + ``prefix`` go to the engine, which packs requests by ``isl - prefix`` + and keeps the prefix as KV context for the attention pass only. """ isl = int(runtime_config.isl or 0) osl = int(runtime_config.osl or 0) @@ -1498,8 +1508,11 @@ def run_mixed( latency_ms = float(components["latency_ms"]) energy_wms = float(components["energy_wms"]) if self._has_visual_context_work(model, runtime_config): - visual_batch = int(np.ceil(step.context_tokens / isl)) - visual_scale = float(np.ceil(isl / step.context_tokens)) + # Requests per mixed step follow the uncached prefill length, as + # in the engine's context-attention pass. + isl_new = max(isl - prefix, 1) + visual_batch = int(np.ceil(step.context_tokens / isl_new)) + visual_scale = float(np.ceil(isl_new / step.context_tokens)) visual_latency, visual_energy, visual_source = self._run_visual_context_phase( model, database, @@ -1658,12 +1671,33 @@ def run_agg( """Run the agg (continuous-batching) inference for a single (b, ctx_tokens) point.""" text_isl = runtime_config.isl osl = runtime_config.osl - prefix = runtime_config.prefix + prefix = int(runtime_config.prefix or 0) b = runtime_config.batch_size img_ctx_tokens = self._visual_context_tokens(model, runtime_config) isl = text_isl + img_ctx_tokens ctx_tokens = kwargs.get("ctx_tokens") assert ctx_tokens is not None, "ctx_tokens is required" + if prefix >= isl: + # Same contract as run_mixed / the engine: a request must carry at + # least one uncached token. + raise ValueError(f"prefix ({prefix}) must be smaller than the effective isl ({isl}) for an agg run") + # ``ctx_tokens`` budgets UNCACHED (new) prefill tokens per mixed step — + # what SGLang --chunked-prefill-size, vLLM max_num_batched_tokens and + # the TRT-LLM scheduler's max_num_tokens actually cap — so the schedule + # below packs requests by ``isl_new``. The cached prefix still reaches + # the per-step cost (run_mixed, as KV context) and the KV footprint + # (_get_memory_usage); both receive the full ``isl`` and ``prefix``. + isl_new = max(isl - prefix, 1) + # A mixed step cannot hold more uncached tokens than the batch owns, so + # the schedule runs on the budget capped at b * isl_new: the packed + # request count then never exceeds b. Packing by isl_new makes + # ceil(ctx_tokens / isl_new) > b reachable with ordinary inputs (a + # large prefix), which would publish negative decode-request counts; + # the old full-isl schedule already did so for an explicit + # ctx_tokens > b * isl, and for b == 1 priced phantom requests. + # ``ctx_tokens`` stays the published engine knob. With prefix 0 this + # only differs from the old schedule when ctx_tokens >= b * isl. + ctx_budget = min(ctx_tokens, b * isl_new) # None (or an omitted kwarg) means the caller did not model speculative # progress here; the summary then stays eligible for the upper-layer # post-hoc projection (SpeculativeDecodingProfile.project_summary). @@ -1680,7 +1714,7 @@ def run_agg( f"decode_tokens_per_iteration must be finite and within [1, nextn + 1={max_decode_progress:g}]" ) decode_iterations = 1.0 + max(osl - 1, 0) / decode_tokens_per_iteration - balance_score = isl * b / ctx_tokens / decode_iterations + balance_score = isl_new * b / ctx_budget / decode_iterations # Backend-specific kwargs (TRT-LLM: max_seq_len / max_num_tokens / # free_gpu_memory_fraction; vLLM / SGLang: free_gpu_memory_fraction). @@ -1723,16 +1757,22 @@ def run_agg( encoder_memory_total = encoder_memory.get("total", 0.0) # Compute the mean-field number of engine iterations needed to consume - # all context and commit the requested output tokens. - steps_to_finish_ctx = np.ceil(isl * b / ctx_tokens) + # all UNCACHED context and commit the requested output tokens. + steps_to_finish_ctx = np.ceil(isl_new * b / ctx_budget) num_mix_steps = num_genonly_steps = 0 num_mix_steps_for_tpot_calc = 0 # correction for tpot calc only if b > 1: - num_mix_gen_tokens = self._mix_step_gen_tokens(b, ctx_tokens, isl, decode_iterations) - assert num_mix_gen_tokens >= 1, ( - f"num_mix_gen_tokens: {num_mix_gen_tokens}, b: {b}, ctx_tokens: {ctx_tokens}, isl: {isl}" - ) - num_mix_ctx_tokens = ctx_tokens + if ctx_budget >= b * isl_new: + # The whole batch prefills in the single mixed step, so no + # decode request rides along with it. + num_mix_gen_tokens = 0 + else: + num_mix_gen_tokens = self._mix_step_gen_tokens(b, ctx_budget, isl_new, decode_iterations) + assert num_mix_gen_tokens >= 1, ( + f"num_mix_gen_tokens: {num_mix_gen_tokens}, b: {b}, ctx_tokens: {ctx_tokens}, " + f"ctx_budget: {ctx_budget}, isl: {isl}, prefix: {prefix}" + ) + num_mix_ctx_tokens = ctx_budget if steps_to_finish_ctx >= decode_iterations: num_mix_steps = steps_to_finish_ctx num_genonly_steps = 0 @@ -1746,7 +1786,7 @@ def run_agg( elif b == 1: # special case for b=1 num_mix_steps = 1 - num_mix_ctx_tokens = ctx_tokens + num_mix_ctx_tokens = ctx_budget num_mix_gen_tokens = 0 num_genonly_steps = max(decode_iterations - 1.0, 0.0) num_genonly_tokens = 1 @@ -1797,9 +1837,10 @@ def run_agg( # TTFT: per-request prefill time * queuing factor, plus encoder latency. # _mix_step_efficiency reduces mix_step_latency_ms based on the fraction of # decode tokens in the step. For TTFT we need the pure prefill cost (no decode - # tokens alongside), so we undo that efficiency reduction first. + # tokens alongside), so we undo that efficiency reduction first. A request + # needs ceil(isl_new / ctx_budget) chunks for its uncached prefill. _prefill_step_ms = mix_step_latency_ms / mix_efficiency if mix_efficiency > 0 else mix_step_latency_ms - _ttft_per_request = _prefill_step_ms * np.ceil(isl / ctx_tokens) + self._prefill_dispatch_overhead_ms(model) + _ttft_per_request = _prefill_step_ms * np.ceil(isl_new / ctx_budget) + self._prefill_dispatch_overhead_ms(model) ttft = encoder_latency_ms + _ttft_per_request * self._ttft_queuing_factor(b, steps_to_finish_ctx) logger.debug( f"ttft: prefill_step={_prefill_step_ms:.2f}ms qf={self._ttft_queuing_factor(b, steps_to_finish_ctx):.2f}" @@ -1842,7 +1883,9 @@ def run_agg( agg_power_avg_w = total_energy_wms / total_latency_ms if total_latency_ms > 0 else 0.0 logger.debug(f"Aggregated power: {agg_power_avg_w}W (from {total_energy_wms} W·ms / {total_latency_ms} ms)") - num_ctx_requests = np.ceil(ctx_tokens / isl) + # Requests prefilling in a mixed step: complete ones plus the partial + # request that fills a non-multiple budget, never more than the batch. + num_ctx_requests = min(np.ceil(ctx_budget / isl_new), float(b)) num_gen_requests = b - num_ctx_requests if b == 1: num_ctx_requests = 1 @@ -1857,7 +1900,7 @@ def run_agg( if b > 1: # will not be corrected by balance score when it's larger than 1.0 # in order to indicate what's happening - num_tokens = num_gen_requests + ctx_tokens + num_tokens = num_gen_requests + ctx_budget # Only the decode requests' tokens verify nextn+1 under speculative # decoding; the context share is processed once (see the MTP # correction in _get_memory_usage). @@ -1865,7 +1908,7 @@ def run_agg( else: # b == 1 starts with a context-only step; the later decode peak is # compared below when the workload schedules output tokens. - num_tokens = ctx_tokens + num_tokens = ctx_budget mtp_scaled_tokens = 0 memory = self._get_memory_usage( @@ -2089,6 +2132,9 @@ def find_best_agg_result_under_constraints( """ isl = runtime_config.isl isl_eff = isl + self._visual_context_tokens(model, runtime_config) + # ctx_tokens budgets UNCACHED prefill tokens (run_agg's isl_new), so the + # feasibility guards and the capped-gen dedup pack requests by it. + isl_new = max(isl_eff - int(runtime_config.prefix or 0), 1) osl = runtime_config.osl ttft = runtime_config.ttft tpot = runtime_config.tpot @@ -2122,7 +2168,10 @@ def find_best_agg_result_under_constraints( # during the loop, as b, ctx_tokens and system memory are monotonic, we can break the # inner loop when the system is oom. b_list = [b for b in b_list_default if b <= max_batch_size] - ctx_tokens_list = self._get_ctx_tokens_list_for_agg_sweep(isl_eff, ctx_stride, enable_chunked_prefill) + # The grid is built on isl_new too, so its points are multiples of the + # uncached prefill and the guards below admit the same (b, requests) + # shapes as at prefix 0. + ctx_tokens_list = self._get_ctx_tokens_list_for_agg_sweep(isl_new, ctx_stride, enable_chunked_prefill) results_df = pd.DataFrame(columns=common.ColumnsAgg) results_dict_list: list[dict] = [] @@ -2131,16 +2180,16 @@ def find_best_agg_result_under_constraints( all_oom = True for b in b_list: for ctx_tokens in ctx_tokens_list: - if b - np.ceil(ctx_tokens / isl_eff) < 0: # allow b==1 + if b - np.ceil(ctx_tokens / isl_new) < 0: # allow b==1 break if b > 1 and ( - b - np.ceil(ctx_tokens / isl_eff) < 1 + b - np.ceil(ctx_tokens / isl_new) < 1 ): # general case, to ensure there's at least one gen req break # filter out repeated records for balance score correction - balance_score = isl_eff * b / ctx_tokens / osl + balance_score = isl_new * b / ctx_tokens / osl if balance_score > 1: gen_tokens = b // balance_score if gen_tokens > 1 and gen_tokens in capped_b: diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index 16b7f5e3b..105959086 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -1093,7 +1093,8 @@ def mixed_step_breakdown_per_op( ]: """``mixed_step_breakdown`` with the per-op values kept: ``(shared_non_attention, context_attention, decode_attention)`` lists; - context-attention entries arrive already divided by ``ceil(isl/ctx)``.""" + context-attention entries arrive already divided by + ``ceil((isl - prefix)/ctx)``; ``ctx_tokens`` budgets uncached tokens.""" return self._engine.mixed_step_breakdown_per_op( int(ctx_tokens), int(gen_tokens), diff --git a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py index 21046816e..3f0ffd9ba 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -698,7 +698,7 @@ def estimate_mixed_step_latency_with_rust( Delegates to ``EngineHandle.mixed_step_latency``. The Rust ``Engine::mixed_step_latency`` is a literal mirror of Python's ``_get_mix_step_latency`` three-pass composition (combined non-attention, - context attention / ceil(isl/ctx), decode attention with the ``(nextn+1)`` + context attention / ceil((isl - prefix)/ctx), decode attention with the ``(nextn+1)`` batch), so the raw step args plus the runtime imbalance scales pass straight through with no Python-side pre-math. """ @@ -738,7 +738,7 @@ def estimate_mixed_step_breakdown_with_rust( ``run_mixed`` retains native operation names, including draft work in either phase, while preserving the legacy ``"context_attention (scaled)"`` and ``"generation_attention"`` keys. Context values are already divided - by ``ceil(isl/ctx)``. FPM retains its prefill/decode component split and separate draft rows, with + by ``ceil((isl - prefix)/ctx)``. FPM retains its prefill/decode component split and separate draft rows, with the target decode operation reported under ``"generation_attention"``. """ handle = _cached_engine_handle(model, database) diff --git a/python/aisimulate/src/aisimulate_core/sdk/step_estimate.py b/python/aisimulate/src/aisimulate_core/sdk/step_estimate.py index 2c0a450ea..40fe75005 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/step_estimate.py +++ b/python/aisimulate/src/aisimulate_core/sdk/step_estimate.py @@ -14,6 +14,9 @@ class MixedStepInput: """Scheduled work for one mixed prefill/decode engine iteration. + ``context_tokens`` is the step's budget of UNCACHED (new) prefill tokens, + what the engines' chunked-prefill knobs cap: requests contribute their + ``isl - prefix`` new tokens to it, and the cached prefix is KV context. ``num_decode_requests`` is the number of logical decode sequences. It is deliberately distinct from the target-model query-token count, which is derived from the model's speculative draft depth. diff --git a/python/aisimulate/tests/e2e/cli/test_cli_estimate_static.py b/python/aisimulate/tests/e2e/cli/test_cli_estimate_static.py index 7386eabd9..33d7db6ec 100644 --- a/python/aisimulate/tests/e2e/cli/test_cli_estimate_static.py +++ b/python/aisimulate/tests/e2e/cli/test_cli_estimate_static.py @@ -397,6 +397,85 @@ def test_agg_estimate_responds_to_common_prefix(): ) +def test_agg_estimate_explicit_isl_budget_prices_partial_requests_and_caps_at_the_batch(): + """Explicit ``--ctx-tokens `` (the pre-fix default) with a cached prefix. + + The 2048-token budget of UNCACHED tokens holds one complete request of + ``2048 - prefix`` new tokens plus a partial request of ``prefix`` tokens. + The mixed step must price that partial request as such -- its context + attention lies strictly between one request's and two requests' static + context attention (two was the old ceil packing) -- the request counts + must never exceed the batch, and a half-cached batch finishes its prefill + in half the mixed steps: shorter TTFT and higher throughput than cold. + """ + kwargs = _common_kwargs() + isl, batch_size = kwargs["isl"], kwargs["batch_size"] + + def mixed_context_attention(result) -> float: + return result.summary.get_step_estimates()["mixed"].component_latency_ms["context_attention"] + + def static_context_attention(prefill_batch_size: int, prefix: int) -> float: + static = cli_estimate(mode="static_ctx", prefix=prefix, **{**kwargs, "batch_size": prefill_batch_size}) + return static.summary.get_context_latency_dict()["context_attention"] + + cold = cli_estimate(mode="agg", prefix=0, ctx_tokens=isl, **kwargs) + assert (cold.raw["num_ctx_reqs"], cold.raw["num_gen_reqs"]) == (1, batch_size - 1) + assert mixed_context_attention(cold) == pytest.approx(static_context_attention(1, 0)) + # Without a prefix the default budget is the isl: the explicit run is the default run. + assert cold.ttft == cli_estimate(mode="agg", prefix=0, **kwargs).ttft + + warm = cli_estimate(mode="agg", prefix=128, ctx_tokens=isl, **kwargs) + assert (warm.raw["num_ctx_reqs"], warm.raw["num_gen_reqs"]) == (2, batch_size - 2) + assert static_context_attention(1, 128) < mixed_context_attention(warm) < static_context_attention(2, 128) + + half = cli_estimate(mode="agg", prefix=isl // 2, ctx_tokens=isl, **kwargs) + assert (half.raw["num_ctx_reqs"], half.raw["num_gen_reqs"]) == (2, batch_size - 2) + # Exactly two complete requests: no partial request to price. + assert mixed_context_attention(half) == pytest.approx(static_context_attention(2, isl // 2)) + assert half.ttft < cold.ttft + assert half.raw["tokens/s"] > cold.raw["tokens/s"] + + nearly = cli_estimate(mode="agg", prefix=isl - 48, ctx_tokens=isl, **kwargs) + # 48 uncached tokens per request: the batch owns 4 * 48 = 192 of the + # 2048-token budget, so every request prefills at once and none decodes. + assert (nearly.raw["num_ctx_reqs"], nearly.raw["num_gen_reqs"]) == (batch_size, 0) + assert nearly.raw["num_tokens"] == batch_size * 48 + assert nearly.raw["ctx_tokens"] == isl + assert nearly.ttft < half.ttft + assert nearly.raw["tokens/s"] > half.raw["tokens/s"] + + +def test_agg_mixed_step_prices_the_cached_prefix_for_mla_module_models(): + """DeepSeek-V3 builds MLA as one context module (no separate + ``context_attention`` op), so the cached prefix reaches the mixed step + only as KV context of the shared pass. At equal new tokens a warm step + must cost more than a cold one, and the default budget resolves to one + request's uncached prefill. (That the default prices exactly what the + old ISL-sized default did is pinned by the engine-step golden + ``deepseek-v3-b200-vllm-shape-prefix-heavy::mixed``.)""" + kwargs = dict( + model_path="deepseek-ai/DeepSeek-V3", + system_name="b200_sxm", + backend_name="vllm", + batch_size=4, + tp_size=8, + pp_size=1, + moe_tp_size=1, + moe_ep_size=8, + osl=64, + ) + + def shared(result) -> float: + return result.summary.get_step_estimates()["mixed"].component_latency_ms["shared_non_attention"] + + cold = cli_estimate(mode="agg", isl=1024, prefix=0, ctx_tokens=1024, **kwargs) + warm = cli_estimate(mode="agg", isl=16384, prefix=15360, ctx_tokens=1024, **kwargs) + assert shared(warm) > shared(cold) + default = cli_estimate(mode="agg", isl=16384, prefix=15360, **kwargs) + assert default.raw["ctx_tokens"] == 1024 + assert shared(default) == shared(warm) + + def test_agg_estimate_responds_to_common_nextn(): """Regression: ``--nextn`` is now a common param. Toggling MTP on must visibly change the agg estimate; previously the kwarg was silently diff --git a/python/aisimulate/tests/unit/cli/test_cli_api.py b/python/aisimulate/tests/unit/cli/test_cli_api.py index 373a3f341..ef90ef9ad 100644 --- a/python/aisimulate/tests/unit/cli/test_cli_api.py +++ b/python/aisimulate/tests/unit/cli/test_cli_api.py @@ -377,6 +377,81 @@ def _wrapping_build_model_config(*args, **kwargs): f"attention_backend not set in ModelConfig; got: {captured_configs[0].attention_backend}" ) + @pytest.mark.parametrize( + ("prefix", "explicit_ctx_tokens", "expected_ctx_tokens"), + [ + (0, None, 1024), # no cache: one full request per step + (512, None, 512), # isl - prefix + (1023, None, 1), # never below one uncached token + (512, 777, 777), # an explicit budget passes through unchanged + ], + ) + def test_agg_estimate_ctx_tokens_defaults_to_the_uncached_isl( + self, monkeypatch, prefix, explicit_ctx_tokens, expected_ctx_tokens + ): + """With --ctx-tokens omitted the agg budget is the effective isl minus + the cached prefix (one request's full uncached prefill per step); an + explicit value is forwarded as given.""" + import aisimulate.legacy_cli.api as api + import aisimulate.sdk.inference_session as inference_session + + captured = {} + + class _CaptureCompleteError(Exception): + pass + + class FakeSession: + def __init__(self, model, loaded_database, backend): + pass + + def run_agg(self, runtime_config, **kwargs): + captured["ctx_tokens"] = kwargs["ctx_tokens"] + captured["prefix"] = runtime_config.prefix + raise _CaptureCompleteError + + monkeypatch.setattr(api, "_resolve_moe_parallelism", lambda *args, **kwargs: (1, 1)) + monkeypatch.setattr(api, "resolve_context_fmha_by_data", lambda *args, **kwargs: None) + monkeypatch.setattr(api, "resolve_dsv4_moe_arch", lambda *args, **kwargs: None) + monkeypatch.setattr(api, "resolve_nvfp4_for_system", lambda *args, **kwargs: None) + monkeypatch.setattr(inference_session, "InferenceSession", FakeSession) + + with pytest.raises(_CaptureCompleteError): + api._run_agg_estimate( + model_path="Qwen/Qwen3-32B", + system_name="h200_sxm", + backend_name="trtllm", + resolved_version="test", + isl=1024, + osl=512, + image_height=0, + image_width=0, + num_images=1, + video_height=0, + video_width=0, + video_frames=0, + num_videos=0, + num_video_tokens=0, + enable_encoder_dp=True, + batch_size=8, + ctx_tokens=explicit_ctx_tokens, + tp_size=1, + pp_size=1, + attention_dp_size=1, + moe_tp_size=1, + moe_ep_size=1, + gemm_quant_mode=None, + kvcache_quant_mode=None, + fmha_quant_mode=None, + moe_quant_mode=None, + comm_quant_mode=None, + load_database=lambda _: object(), + get_backend=lambda _: object(), + get_model=lambda *_: object(), + prefix=prefix, + ) + + assert captured == {"ctx_tokens": expected_ctx_tokens, "prefix": prefix} + class TestCLIDefaultNextn: """cli_default exposes MTP control with the same semantics as the CLI flags.""" diff --git a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py index d3fcc12ec..4d3938c57 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import math from types import SimpleNamespace from unittest.mock import MagicMock @@ -459,8 +460,10 @@ def _record(*args, **kwargs): ctx_tokens=8, ) + # ctx_tokens budgets uncached tokens: 8 new tokens pack ceil(8 / (8 - 2)) = 2 + # prefilling requests, leaving b - 2 = 2 decode requests. assert len(memory_calls) == 1 - assert memory_calls[0].get("mtp_scaled_tokens") == 3 + assert memory_calls[0].get("mtp_scaled_tokens") == 2 def test_trtllm_budget_path_ignores_the_decode_share() -> None: @@ -1258,3 +1261,327 @@ def test_run_agg_zero_decode_steps_do_not_change_energy_coverage( assert summary.get_aggregate_energy_breakdown()["genonly_step"].latency_ms == 0 assert summary.get_power_data_coverage() == pytest.approx(expected) assert (_apply_power_coverage_gate(summary, summary.get_result_dict())["power_w"] is not None) is published + + +# --------------------------------------------------------------------------- +# Prefix-aware scheduling: ctx_tokens budgets UNCACHED prefill tokens +# --------------------------------------------------------------------------- + +# A long-prompt shape one token past a multiple of the 16384-token per-step +# budget (SGLang --chunked-prefill-size), so every ceil in the schedule rounds +# up: ISL 32769 / OSL 512 / bs 16. +_ISL, _OSL, _BS, _CTX = 32769, 512, 16, 16384 + + +def _fixed_step_costs(monkeypatch, backend: BaseBackend, mix_latency_ms: float = 10.0) -> list[RuntimeConfig]: + """Pin the per-step costs so run_agg exposes only its scheduling; return + the runtime configs handed to run_mixed.""" + seen: list[RuntimeConfig] = [] + + def _run_mixed(model_arg, database_arg, runtime_config_arg, step): + seen.append(runtime_config_arg) + return StepEstimate(latency_ms=mix_latency_ms, energy_wms=1.0) + + monkeypatch.setattr(backend, "run_mixed", _run_mixed) + monkeypatch.setattr( + backend, + "_get_genonly_step_estimate", + lambda *args, **kwargs: _decode_step(1.0, 1.0, {"decode": 1.0}, {"decode": "silicon"}, ()), + ) + return seen + + +def _prefix_sweep_config(prefix: int) -> RuntimeConfig: + return RuntimeConfig(batch_size=_BS, beam_width=1, isl=_ISL, osl=_OSL, prefix=prefix, engine_step_backend="rust") + + +@pytest.mark.parametrize( + ("prefix", "expected_mix_steps"), + [ + # ceil((isl - prefix) * b / ctx_tokens), hand-derived: + (0, 33), # ceil(32769 * 16 / 16384) = ceil(32.001) + (8192, 25), # ceil(24577 * 16 / 16384) = ceil(24.001) + (16384, 17), # ceil(16385 * 16 / 16384) = ceil(16.001) + (24576, 9), # ceil(8193 * 16 / 16384) = ceil(8.001) + (29491, 4), # ceil(3278 * 16 / 16384) = ceil(3.201) + ], +) +def test_run_agg_mix_step_count_follows_uncached_isl( + monkeypatch, backend: BaseBackend, model, database, prefix: int, expected_mix_steps: int +) -> None: + """The ctx_tokens budget counts uncached tokens (what the engines' + chunked-prefill knobs cap), so the mixed steps needed to drain a batch's + prefill shrink with the cached prefix. The engine still receives the full + isl and prefix through the unmodified runtime config, and the scheduling + hook receives the uncached length.""" + seen = _fixed_step_costs(monkeypatch, backend) + hook_isl: list[int] = [] + base_hook = backend._mix_step_gen_tokens + + def _hook(b, ctx_tokens, isl, decode_iterations): + hook_isl.append(isl) + return base_hook(b, ctx_tokens, isl, decode_iterations) + + monkeypatch.setattr(backend, "_mix_step_gen_tokens", _hook) + + summary = backend.run_agg(model, database, _prefix_sweep_config(prefix), ctx_tokens=_CTX) + + scheduling = summary.get_step_estimates()["scheduling"] + assert scheduling["num_mix_steps"] == expected_mix_steps + assert scheduling["num_genonly_steps"] == _OSL - expected_mix_steps + result = summary.get_result_dict() + assert result["num_ctx_reqs"] == math.ceil(_CTX / (_ISL - prefix)) + assert result["num_gen_reqs"] == _BS - result["num_ctx_reqs"] + assert hook_isl == [_ISL - prefix] + assert seen[0].isl == _ISL and seen[0].prefix == prefix + + +def test_run_agg_prefix_zero_schedule_is_unchanged(monkeypatch, backend: BaseBackend, model, database) -> None: + """prefix=0 keeps the legacy schedule: ceil(isl*b/ctx) mixed steps, + ceil(ctx/isl) prefilling requests, ceil(isl/ctx) TTFT chunks and the + isl*b/ctx/osl balance score.""" + _fixed_step_costs(monkeypatch, backend, mix_latency_ms=10.0) + summary = backend.run_agg(model, database, _prefix_sweep_config(0), ctx_tokens=_CTX) + scheduling = summary.get_step_estimates()["scheduling"] + result = summary.get_result_dict() + assert scheduling["num_mix_steps"] == 33 + assert scheduling["num_genonly_steps"] == 479 + assert result["num_ctx_reqs"] == 1 + assert result["num_gen_reqs"] == 15 + assert result["balance_score"] == pytest.approx(_ISL * _BS / _CTX / _OSL) + # 3 chunks of 16384 for a 32769-token prefill, times the queuing factor. + assert result["ttft"] == pytest.approx(10.0 * 3 * backend._ttft_queuing_factor(_BS, 33)) + + +def test_run_agg_ttft_chunk_count_follows_uncached_isl(monkeypatch, model, database) -> None: + """A request's TTFT multiplies the prefill step by ceil((isl - prefix) / + ctx_tokens): a cold 32769-token request takes 3 chunks of 16384, one with + 16384 cached tokens takes 2 and one with 24576 cached takes 1.""" + ttfts = [] + for prefix, chunks, steps in ((0, 3, 33), (16384, 2, 17), (24576, 1, 9)): + # A fresh backend per prefix: the run_agg result cache is keyed on + # the prefix by a separate fix, not by this change. + backend = _TestBackend() + _fixed_step_costs(monkeypatch, backend, mix_latency_ms=10.0) + result = backend.run_agg(model, database, _prefix_sweep_config(prefix), ctx_tokens=_CTX).get_result_dict() + assert result["ttft"] / backend._ttft_queuing_factor(_BS, steps) == pytest.approx(10.0 * chunks) + ttfts.append(result["ttft"]) + assert ttfts[0] > ttfts[1] > ttfts[2] + + +@pytest.mark.parametrize("prefix", [8, 9]) +def test_run_agg_rejects_prefix_at_or_beyond_isl(monkeypatch, backend: BaseBackend, model, database, prefix) -> None: + """The engine rejects a mixed step whose requests carry no uncached + token; run_agg applies the same contract before scheduling.""" + _fixed_step_costs(monkeypatch, backend) + with pytest.raises(ValueError, match="prefix"): + backend.run_agg( + model, + database, + RuntimeConfig(batch_size=2, beam_width=1, isl=8, osl=5, prefix=prefix, engine_step_backend="rust"), + ctx_tokens=8, + ) + + +@pytest.mark.parametrize( + ("b", "prefix", "isl_new"), + [ + # ceil(16384 / 3278) = 5 prefilling requests would exceed the batch of 4. + (4, 29491, 3278), + # Every request carries a single uncached token: 16384 of them would + # be 16384 requests. + (16, _ISL - 1, 1), + ], +) +def test_run_agg_caps_prefilling_requests_at_the_batch( + monkeypatch, backend: BaseBackend, model, database, b: int, prefix: int, isl_new: int +) -> None: + """A mixed step cannot hold more uncached tokens than the batch owns: the + schedule runs on min(ctx_tokens, b * isl_new), so the request counts and + the activation token count handed to the memory model stay non-negative + and the step is priced at the tokens the batch can actually supply. The + published ctx_tokens column keeps the requested engine knob.""" + steps: list[MixedStepInput] = [] + + def _run_mixed(model_arg, database_arg, runtime_config_arg, step): + steps.append(step) + return StepEstimate(latency_ms=10.0, energy_wms=1.0) + + memory_calls: list[dict] = [] + + def _memory(model_arg, database_arg, batch_size, beam_width, isl, osl, **kwargs): + memory_calls.append(kwargs) + return {"total": 1.0} + + monkeypatch.setattr(backend, "run_mixed", _run_mixed) + monkeypatch.setattr( + backend, + "_get_genonly_step_estimate", + lambda *args, **kwargs: _decode_step(1.0, 1.0, {"decode": 1.0}, {"decode": "silicon"}, ()), + ) + monkeypatch.setattr(backend, "_get_memory_usage", _memory) + + summary = backend.run_agg( + model, + database, + RuntimeConfig(batch_size=b, beam_width=1, isl=_ISL, osl=_OSL, prefix=prefix, engine_step_backend="rust"), + ctx_tokens=_CTX, + ) + + result = summary.get_result_dict() + assert result["num_ctx_reqs"] == b + assert result["num_gen_reqs"] == 0 + assert result["gen_tokens"] == 0 + assert result["ctx_tokens"] == _CTX + assert result["num_tokens"] == b * isl_new + assert steps[0].context_tokens == b * isl_new + # Every request prefills in this step: no decode request rides along. + assert steps[0].num_decode_requests == 0 + assert summary.get_step_estimates()["scheduling"]["num_mix_steps"] == 1 + assert memory_calls[0]["num_tokens"] == b * isl_new + assert memory_calls[0]["mtp_scaled_tokens"] == 0 + # The whole batch prefills in one step of one chunk each. + assert result["ttft"] == pytest.approx(10.0 * backend._ttft_queuing_factor(b, 1)) + + +@pytest.mark.parametrize( + ("b", "prefix", "ctx_tokens", "capped", "osl"), + [ + # b == 1 with an explicit budget above its 16385 uncached tokens (e.g. + # the pre-fix default --ctx-tokens ): no phantom second request. + (1, 16384, _ISL, _ISL - 16384, _OSL), + # prefix 0, b == 1, a budget above the isl: one request's prefill. + (1, 0, 2 * _ISL, _ISL, _OSL), + # prefix 0, b == 2, a budget above b * isl: the old schedule packed + # ceil(ctx/isl) = 3 prefilling requests out of 2 (negative decode). + (2, 0, 2 * _ISL + 100, 2 * _ISL, _OSL), + # Exactly the batch's uncached tokens: the whole batch prefills in the + # one step, at prefix 0 and with a prefix, for osl > 1 and osl == 1. + (2, 0, 2 * _ISL, 2 * _ISL, _OSL), + (2, 16384, 2 * (_ISL - 16384), 2 * (_ISL - 16384), _OSL), + (2, 16384, 2 * (_ISL - 16384), 2 * (_ISL - 16384), 1), + ], +) +def test_run_agg_caps_the_budget_for_every_batch_size( + monkeypatch, backend: BaseBackend, model, database, b: int, prefix: int, ctx_tokens: int, capped: int, osl: int +) -> None: + """The b * isl_new cap applies to every batch size and prefix: the mixed + step and the activation token count are priced at the tokens the batch + actually owns, while the published ctx_tokens column keeps the knob. At + prefix 0 the schedule only changes in this regime (ctx > b * isl), where + the old one was not physical.""" + steps: list[MixedStepInput] = [] + + def _run_mixed(model_arg, database_arg, runtime_config_arg, step): + steps.append(step) + return StepEstimate(latency_ms=10.0, energy_wms=1.0) + + memory_calls: list[dict] = [] + + def _memory(model_arg, database_arg, batch_size, beam_width, isl, osl, **kwargs): + memory_calls.append(kwargs) + return {"total": 1.0} + + monkeypatch.setattr(backend, "run_mixed", _run_mixed) + monkeypatch.setattr( + backend, + "_get_genonly_step_estimate", + lambda *args, **kwargs: _decode_step(1.0, 1.0, {"decode": 1.0}, {"decode": "silicon"}, ()), + ) + monkeypatch.setattr(backend, "_get_memory_usage", _memory) + + summary = backend.run_agg( + model, + database, + RuntimeConfig(batch_size=b, beam_width=1, isl=_ISL, osl=osl, prefix=prefix, engine_step_backend="rust"), + ctx_tokens=ctx_tokens, + ) + + result = summary.get_result_dict() + assert steps[0].context_tokens == capped + assert steps[0].num_decode_requests == 0 + assert result["ctx_tokens"] == ctx_tokens + assert result["num_tokens"] == capped + assert memory_calls[0]["num_tokens"] == capped + assert result["num_gen_reqs"] >= 0 + assert result["num_ctx_reqs"] == (1 if b == 1 else b) + if osl == 1: + return + # One chunk per request: the whole uncached prefill fits the budget. + assert result["ttft"] / backend._ttft_queuing_factor( + b, summary.get_step_estimates()["scheduling"]["num_mix_steps"] + ) == pytest.approx(10.0) + + +def test_run_agg_budget_cap_is_inert_when_the_batch_owns_more_tokens( + monkeypatch, backend: BaseBackend, model, database +) -> None: + """bs 16 with 3278 uncached tokens each owns 52448 tokens, more than the + 16384 budget: nothing is capped and the schedule is the plain isl_new + packing (5 prefilling + 11 decoding requests, 4 mixed steps).""" + seen = _fixed_step_costs(monkeypatch, backend) + summary = backend.run_agg(model, database, _prefix_sweep_config(29491), ctx_tokens=_CTX) + result = summary.get_result_dict() + assert result["num_ctx_reqs"] == 5 + assert result["num_gen_reqs"] == 11 + assert result["num_tokens"] == 11 + _CTX + assert summary.get_step_estimates()["scheduling"]["num_mix_steps"] == 4 + assert seen[0].prefix == 29491 + + +@pytest.mark.parametrize( + ("ctx_tokens", "expected_batch", "expected_scale"), + [ + # 60 uncached tokens hold ceil(60 / 20) = 3 requests in one chunk. + (60, 3, 1.0), + # 10 uncached tokens hold one request over ceil(20 / 10) = 2 chunks. + (10, 1, 2.0), + ], +) +def test_run_mixed_visual_context_batch_follows_uncached_isl( + monkeypatch, backend: BaseBackend, model, database, ctx_tokens: int, expected_batch: int, expected_scale: float +) -> None: + """Text isl 8 + 16 visual tokens = 24 effective tokens, 4 of them cached: + the visual context phase runs for ceil(ctx_tokens / 20) requests and its + cost is amortized over ceil(20 / ctx_tokens) chunks, like the engine's + context-attention pass.""" + from aisimulate.sdk.backends import base_backend as base_backend_module + + model.encoder_config = _vision_encoder_config() + model.visual_context_ops = [_StaticOp("visual_attention", latency_ms=1.0, energy_wms=1.0)] + + monkeypatch.setattr( + base_backend_module, + "estimate_mixed_step_breakdown_with_rust", + lambda *args, **kwargs: { + "latency_ms": 1.0, + "energy_wms": 1.0, + "component_latency_ms": {"shared_non_attention": 1.0}, + "component_energy_wms": {"shared_non_attention": 1.0}, + "per_op_latency_ms": {}, + "per_op_source": {}, + "moe_comm_fallbacks": (), + }, + ) + batches: list[int] = [] + + def _visual(model_arg, database_arg, runtime_config_arg, batch_size, *, include_energy=True): + batches.append(batch_size) + return {"visual_attention": 8.0}, {"visual_attention": 4.0}, {"visual_attention": "silicon"} + + monkeypatch.setattr(backend, "_run_visual_context_phase", _visual) + + step = backend.run_mixed( + model, + database, + RuntimeConfig( + isl=8, osl=6, prefix=4, num_images_per_request=1, num_image_tokens=16, engine_step_backend="rust" + ), + MixedStepInput(context_tokens=ctx_tokens, num_decode_requests=1), + ) + + assert batches == [expected_batch] + assert step.per_op_latency_ms["visual_attention"] == pytest.approx(8.0 / expected_scale) + assert step.component_latency_ms["context_attention"] == pytest.approx(8.0 / expected_scale) + assert step.latency_ms == pytest.approx(1.0 + 8.0 / expected_scale) + assert step.energy_wms == pytest.approx(1.0 + 4.0 / expected_scale) diff --git a/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py b/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py index 78a5c04e2..23d2e18f4 100644 --- a/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py +++ b/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py @@ -622,8 +622,6 @@ def run(model): @pytest.mark.integration @pytest.mark.parametrize("ctx_tokens,gen_requests,prefix", [(0, 7, 0), (128, 0, 64), (128, 7, 64), (8000, 7, 64)]) def test_mixed_draft_native_phases_match_independent_queries(real_database, ctx_tokens, gen_requests, prefix): - import math - from aisimulate_core.sdk.models import get_model from aisimulate_core.sdk.rust_engine_step import _cached_engine_handle @@ -639,18 +637,34 @@ def test_mixed_draft_native_phases_match_independent_queries(real_database, ctx_ assert not any(row[0].startswith("draft_") for row in shared) ctx_indices = [i for i, op in enumerate(model.context_ops) if op._name.startswith("draft_")] gen_indices = [i for i, op in enumerate(model.generation_ops) if op._name.startswith("draft_")] - expected_context = ( - handle.evaluate_context_ops( - ctx_indices, batch_size=math.ceil(ctx_tokens / 4000), s=4000 - prefix, prefix=prefix - ) - if ctx_tokens - else [] - ) + # ctx_tokens budgets UNCACHED tokens: requests of isl - prefix = 3936 new + # tokens over 64 cached ones fill it. Hand-derived per case (not the + # production formula): 128 tokens are one chunk of one request, amortized + # over its ceil(3936/128) = 31 chunks; 8000 tokens hold 2 complete requests + # (7872) plus a 128-token partial one weighing 128/3936 of a third. + isl_new = 4000 - prefix + expected_groups = { + (128, 64): ([(1, 1.0)], 31), + (8000, 64): ([(2, 1.0 - 128 / 3936), (3, 128 / 3936)], 1), + } + expected_context = [] + context_divisor = 1 + if ctx_tokens: + groups, context_divisor = expected_groups[(ctx_tokens, prefix)] + per_group = [ + handle.evaluate_context_ops(ctx_indices, batch_size=batch, s=isl_new, prefix=prefix) for batch, _ in groups + ] + for rows in zip(*per_group, strict=True): + assert len({row[0] for row in rows}) == 1 + latency = sum(weight * row[1] for (_, weight), row in zip(groups, rows, strict=True)) + energy = sum(weight * row[2] for (_, weight), row in zip(groups, rows, strict=True)) + sources = {row[3] for row in rows} + expected_context.append((rows[0][0], latency, energy, sources.pop() if len(sources) == 1 else "mixed")) expected_generation = ( handle.evaluate_generation_ops(gen_indices, batch_size=gen_requests * 4, s=4033) if gen_requests else [] ) for actual, expected, divisor in ( - (context, expected_context, math.ceil(4000 / ctx_tokens) if ctx_tokens else 1), + (context, expected_context, context_divisor), (generation, expected_generation, 1), ): drafts = {r[0]: r for r in actual if r[0].startswith("draft_")} diff --git a/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py b/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py index 5a9585a9f..9f6cb1842 100644 --- a/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py +++ b/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py @@ -8,6 +8,7 @@ the unit coverage here targets local control flow and terminal classification. """ +import math from unittest.mock import MagicMock import pandas as pd @@ -434,6 +435,76 @@ def _record(*, runtime_config, ctx_tokens, **_kwargs): assert points == [(1, 1040)] # 256 text + 784 post-merge video tokens +def _recording_agg_worker(points: list[tuple[int, int]]): + def _record(*, runtime_config, ctx_tokens, **_kwargs): + points.append((runtime_config.batch_size, ctx_tokens)) + summary = MagicMock() + summary.check_oom.return_value = False + summary.check_kv_cache_oom.return_value = False + summary.get_result_dict.return_value = {"ttft": 1.0, "tpot": 1.0, "seq/s": 1.0} + summary.get_per_ops_source.return_value = {} + return summary + + return _record + + +@pytest.mark.parametrize("prefix", [0, 1536, 2000]) +def test_sweep_agg_prefix_grid_covers_every_batch_and_mirrors_legacy(monkeypatch, prefix): + """With a cached prefix the ctx grid and the batch/ctx guards run on the + uncached isl, so every batch size keeps sweep points (a grid on the full + isl was guarded out almost entirely once prefix >= ~75% of isl), and the + legacy BaseBackend sweep visits exactly the same (b, ctx_tokens) points.""" + from aisimulate.sdk.backends.factory import get_backend + + isl, osl, max_batch_size = 2048, 64, 8 + isl_new = isl - prefix + runtime_config = config.RuntimeConfig(isl=isl, osl=osl, prefix=prefix, ttft=1e9, tpot=1e9) + + points: list[tuple[int, int]] = [] + monkeypatch.setattr(sweep, "predict_agg_worker", _recording_agg_worker(points)) + sweep._sweep_one_parallel_agg( + model=MagicMock(), + backend=MagicMock(), + database=MagicMock(), + runtime_config=runtime_config, + top_k=0, + max_batch_size=max_batch_size, + ctx_stride=512, + enable_chunked_prefill=False, + free_gpu_memory_fraction=None, + max_seq_len=None, + ) + + assert {b for b, _ in points} == set(range(1, max_batch_size + 1)) + for b, ctx_tokens in points: + assert ctx_tokens % isl_new == 0 + # legacy guards: at least one decoding request when b > 1 + assert b - math.ceil(ctx_tokens / isl_new) >= (1 if b > 1 else 0) + # b=1 sweeps exactly one request's uncached prefill. + assert [ctx for b, ctx in points if b == 1] == [isl_new] + + legacy_points: list[tuple[int, int]] = [] + legacy_record = _recording_agg_worker(legacy_points) + legacy_backend = get_backend("sglang") + monkeypatch.setattr( + legacy_backend, + "run_agg", + lambda model, database, runtime_config, **kwargs: legacy_record( + runtime_config=runtime_config, ctx_tokens=kwargs["ctx_tokens"] + ), + ) + legacy_backend.find_best_agg_result_under_constraints( + MagicMock(), + MagicMock(), + runtime_config, + top_k=0, + max_batch_size=max_batch_size, + ctx_stride=512, + enable_chunked_prefill=False, + ) + assert legacy_points == points + + # --------------------------------------------------------------------------- # sweep_disagg validation # --------------------------------------------------------------------------- diff --git a/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py b/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py index f5521ac86..8e94a33c5 100644 --- a/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py +++ b/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py @@ -1900,6 +1900,46 @@ def fake_predict_agg_worker(**kwargs): assert captured["predict_kwargs"]["ctx_tokens"] == t.isl +@pytest.mark.parametrize( + ("prefix", "explicit_ctx_tokens", "expected_offset"), + [ + (0, None, 0), # one full request per step + (512, None, 512), # isl - prefix + (512, 128, None), # an explicit budget passes through unchanged + ], +) +def test_run_single_agg_ctx_tokens_defaults_to_the_uncached_isl( + monkeypatch, prefix, explicit_ctx_tokens, expected_offset +): + """Task.run_single_agg mirrors cli_estimate: the default budget is the + effective isl minus the cached prefix, an explicit ctx_tokens is used as + given, and the floor is one uncached token.""" + from unittest.mock import MagicMock + + from aisimulate.sdk import predict + + captured: dict = {} + + def fake_predict_agg_worker(**kwargs): + captured["ctx_tokens"] = kwargs["ctx_tokens"] + captured["prefix"] = kwargs["runtime_config"].prefix + return _build_fake_summary() + + monkeypatch.setattr("aisimulate.sdk.perf_database.get_database_view", lambda *a, **kw: "db") + monkeypatch.setattr("aisimulate.sdk.backends.factory.get_backend", lambda name: MagicMock(name=f"backend-{name}")) + monkeypatch.setattr("aisimulate.sdk.models.get_model", lambda *a: MagicMock(name="model")) + monkeypatch.setattr(predict, "predict_agg_worker", fake_predict_agg_worker) + + t = Task(serving_mode="agg", model_path="deepseek-ai/DeepSeek-V3", system_name="h200_sxm", prefix=prefix) + t.run_single_agg(tp=4, pp=1, dp=1, moe_tp=1, moe_ep=1, batch_size=64, ctx_tokens=explicit_ctx_tokens) + expected = explicit_ctx_tokens if explicit_ctx_tokens is not None else t.isl - expected_offset + assert captured == {"ctx_tokens": expected, "prefix": prefix} + + near = Task(serving_mode="agg", model_path="deepseek-ai/DeepSeek-V3", system_name="h200_sxm", prefix=t.isl - 1) + near.run_single_agg(tp=4, batch_size=64) + assert captured["ctx_tokens"] == 1 + + def test_run_single_agg_rejects_disagg_task(): t = Task( serving_mode="disagg", From ccc48f003288024fc326644ea3af6d5a5d4e2fe6 Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Tue, 29 Sep 2026 08:48:49 +0800 Subject: [PATCH 2/7] test(parity): re-pin prefix chunked-prefill goldens for the uncached budget Deliberate modeling change carried by 3eae51ed ("fix: count only uncached tokens against the legacy agg prefill budget"). All default-budget prefix records (engine-step minimax-m25-b200-vllm-sampled-prefix and deepseek-v3-b200-vllm-shape-prefix-heavy mixed/agg, compile-engine minimax-m25-b200-vllm-sampled-prefix::mixed_step) are unchanged: they pass against the existing goldens at rtol 1e-12. compile_engine.json, chunked-prefill shapes on minimax-m25-b200-vllm: chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step 16.081171352808852 -> 16.414675748089216 (+2.07%) The request's 900 uncached tokens are amortized over ceil(900/300) = 3 chunks instead of ceil(1000/300) = 4; that is the whole delta (this model's pass-1 ops do not read the prefix). chunked_prefill::ctx4096_gen4_isl4096_osl128_prefix256::mixed_step new record: 57.89342286133688 (upstream prices this shape 55.30729111158813) An explicit ISL-sized budget with a cached prefix: 4096 uncached tokens now hold one request of 3840 new tokens plus a 256-token partial request (fill 1/15). The new default budget 3840 prices exactly the upstream value 55.30729111158813. chunked_prefill::ctx512_gen4_isl4096_osl128_prefix256 is unchanged because ceil(3840/512) = ceil(4096/512) = 8. Pinned on the clean tree at 3eae51ed with python/aisimulate/.venv/bin/python crates/core/parity_tests/perfmodel/pin_goldens.py \ --refresh chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step (the ctx4096 record is appended by the same run, as a newly declared shape). Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- .../parity_tests/perfmodel/goldens/compile_engine.json | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/core/parity_tests/perfmodel/goldens/compile_engine.json b/crates/core/parity_tests/perfmodel/goldens/compile_engine.json index fa3a1e8d7..a71984d6c 100644 --- a/crates/core/parity_tests/perfmodel/goldens/compile_engine.json +++ b/crates/core/parity_tests/perfmodel/goldens/compile_engine.json @@ -16,7 +16,11 @@ "post_freeze_pins": { "chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step": { "engine": "rust", - "git_head": "0a51476b6ab90bc0e475bd41d4b1c7abbef07b95" + "git_head": "3eae51edd74cd072546a1502bcc029ff67220e53" + }, + "chunked_prefill::ctx4096_gen4_isl4096_osl128_prefix256::mixed_step": { + "engine": "rust", + "git_head": "3eae51edd74cd072546a1502bcc029ff67220e53" }, "chunked_prefill::ctx512_gen4_isl4096_osl128_prefix0::mixed_step": { "engine": "rust", @@ -228,7 +232,8 @@ } }, "references": { - "chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step": 16.081171352808852, + "chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step": 16.414675748089216, + "chunked_prefill::ctx4096_gen4_isl4096_osl128_prefix256::mixed_step": 57.89342286133688, "chunked_prefill::ctx512_gen4_isl4096_osl128_prefix0::mixed_step": 17.664446149246828, "chunked_prefill::ctx512_gen4_isl4096_osl128_prefix256::mixed_step": 17.662495086777085, "gpt-oss-20b-b200-trtllm-isl1024-osl2::decode_step": 2.213361158220323, From 0153c8976c1a958e69946674a7557f1353a89ddd Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Tue, 29 Sep 2026 08:46:09 +0800 Subject: [PATCH 3/7] fix(sdk): key the run_agg result cache on prefix and correction scales BaseBackend.run_agg memoizes summaries in _agg_cache, but the key only carried (isl, osl, b, ctx_tokens, backend kwargs), the visual fields and the speculative progress. RuntimeConfig.prefix, seq_imbalance_correction_scale and gen_seq_imbalance_correction_scale were left out although all three change the answer: prefix now drives the mixed-step schedule (uncached budget), the attention cost, the activation footprint and the echoed result_dict["prefix"]; the two scales drive the mixed and decode step latencies. An SDK caller that reuses one backend or InferenceSession across prefixes therefore got the first request's summary back, echoing the stale prefix. The CLI is unaffected because cli_estimate builds a fresh backend per call. Add a runtime_cache_key (prefix, seq_imbalance_correction_scale, gen_seq_imbalance_correction_scale) to the composite key. Extending the composite key rather than _make_agg_cache_key keeps the TRT-LLM override intact. prefix is normalized as int(prefix or 0), as run_mixed does, so prefix=None and prefix=0 share one entry. A comment inventories every RuntimeConfig field run_agg reads and which key part carries it, and the run_agg budget comment now states that the KV footprint is sized from the full isl and does not depend on the prefix. No estimate math changes; identical requests still hit the cache. Tests: prefix separation, echo and cache hit on repeat; prefix=None sharing the prefix=0 entry; separation on each correction scale. Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- .../sdk/backends/base_backend.py | 28 ++++++- .../unit/sdk/backends/test_base_backend.py | 75 +++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py index 4b3128d23..9e7896888 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py @@ -1671,6 +1671,8 @@ def run_agg( """Run the agg (continuous-batching) inference for a single (b, ctx_tokens) point.""" text_isl = runtime_config.isl osl = runtime_config.osl + # Same normalization as run_mixed so None and 0 share one cache entry + # and the echoed result_dict["prefix"] is always an int. prefix = int(runtime_config.prefix or 0) b = runtime_config.batch_size img_ctx_tokens = self._visual_context_tokens(model, runtime_config) @@ -1685,8 +1687,9 @@ def run_agg( # what SGLang --chunked-prefill-size, vLLM max_num_batched_tokens and # the TRT-LLM scheduler's max_num_tokens actually cap — so the schedule # below packs requests by ``isl_new``. The cached prefix still reaches - # the per-step cost (run_mixed, as KV context) and the KV footprint - # (_get_memory_usage); both receive the full ``isl`` and ``prefix``. + # the per-step cost (run_mixed, as KV context), which receives the + # full ``isl`` and ``prefix``; the KV footprint is sized from the full + # ``isl`` and does not depend on the prefix. isl_new = max(isl - prefix, 1) # A mixed step cannot hold more uncached tokens than the batch owns, so # the schedule runs on the budget capped at b * isl_new: the packed @@ -1731,9 +1734,30 @@ def run_agg( runtime_config.num_videos_per_request, runtime_config.num_video_tokens, ) + # Cache identity: every RuntimeConfig field run_agg (or run_mixed / + # _get_genonly_step_estimate, which it delegates to) reads must appear + # in exactly one of the key parts below, or a reused backend serves + # another request's summary. Current inventory: + # _make_agg_cache_key : isl, osl, batch_size (+ ctx_tokens / agg kwargs) + # visual_cache_key : image_height, image_width, num_images_per_request, + # num_image_tokens, video_height, video_width, + # video_frames, num_videos_per_request, num_video_tokens + # runtime_cache_key : prefix (mixed-step schedule and attention cost, + # activation footprint via ctx_budget / + # num_gen_requests, echoed result_dict), + # seq_imbalance_correction_scale, + # gen_seq_imbalance_correction_scale (step latencies) + # engine_step_backend is validated live below and intentionally excluded. + # A new RuntimeConfig field that run_agg reads goes into runtime_cache_key. + runtime_cache_key = ( + prefix, + runtime_config.seq_imbalance_correction_scale, + runtime_config.gen_seq_imbalance_correction_scale, + ) cache_key = ( self._make_agg_cache_key(isl, osl, b, ctx_tokens, agg_extra), visual_cache_key, + runtime_cache_key, # Explicit progress and an omitted kwarg schedule identically at # 1.0 but record different scheduling metadata, so they must not # share a cache entry. diff --git a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py index 4d3938c57..36701dfa1 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py @@ -926,6 +926,81 @@ def _run_mixed(*args, **kwargs): assert len(backend._agg_cache) == 2 +def test_run_agg_cache_separates_prefix_on_a_reused_backend( + monkeypatch, + backend: BaseBackend, + model, + database, +) -> None: + """Two calls differing only in prefix must not share a cached summary.""" + mixed_calls: list[int] = [] + + def _run_mixed(*args, **kwargs): + prefix = args[2].prefix + mixed_calls.append(prefix) + # Make the answer depend on prefix so a stale hit is observable. + return StepEstimate(latency_ms=10.0 - prefix, energy_wms=100.0) + + monkeypatch.setattr(backend, "run_mixed", _run_mixed) + monkeypatch.setattr( + backend, + "_get_genonly_step_estimate", + lambda *args, **kwargs: _decode_step(1.0, 1.0, {}, {}, ()), + ) + + common_kwargs = dict(batch_size=2, beam_width=1, isl=8, osl=5, engine_step_backend="rust") + no_prefix = backend.run_agg(model, database, RuntimeConfig(**common_kwargs, prefix=0), ctx_tokens=8) + with_prefix = backend.run_agg(model, database, RuntimeConfig(**common_kwargs, prefix=4), ctx_tokens=8) + + assert mixed_calls == [0, 4] + assert len(backend._agg_cache) == 2 + assert no_prefix is not with_prefix + assert no_prefix.get_result_dict()["prefix"] == 0 + assert with_prefix.get_result_dict()["prefix"] == 4 + assert no_prefix.get_result_dict()["ttft"] != with_prefix.get_result_dict()["ttft"] + + # An identical request still hits the cache without re-estimating. + repeat = backend.run_agg(model, database, RuntimeConfig(**common_kwargs, prefix=4), ctx_tokens=8) + assert repeat is with_prefix + assert mixed_calls == [0, 4] + assert len(backend._agg_cache) == 2 + + # prefix=None normalizes like run_mixed (int(prefix or 0)): it shares the + # prefix=0 entry and echoes 0, not None. + assert backend.run_agg(model, database, RuntimeConfig(**common_kwargs, prefix=None), ctx_tokens=8) is no_prefix + assert mixed_calls == [0, 4] + assert len(backend._agg_cache) == 2 + + +@pytest.mark.parametrize("field", ["seq_imbalance_correction_scale", "gen_seq_imbalance_correction_scale"]) +def test_run_agg_cache_separates_seq_imbalance_correction_scales( + monkeypatch, + backend: BaseBackend, + model, + database, + field: str, +) -> None: + seen: list[float] = [] + + def _run_mixed(*args, **kwargs): + seen.append(getattr(args[2], field)) + return StepEstimate(latency_ms=1.0, energy_wms=1.0) + + monkeypatch.setattr(backend, "run_mixed", _run_mixed) + monkeypatch.setattr( + backend, + "_get_genonly_step_estimate", + lambda *args, **kwargs: _decode_step(1.0, 1.0, {}, {}, ()), + ) + + common_kwargs = dict(batch_size=2, beam_width=1, isl=8, osl=5, engine_step_backend="rust") + backend.run_agg(model, database, RuntimeConfig(**common_kwargs, **{field: 1.0}), ctx_tokens=8) + backend.run_agg(model, database, RuntimeConfig(**common_kwargs, **{field: 1.5}), ctx_tokens=8) + + assert seen == [1.0, 1.5] + assert len(backend._agg_cache) == 2 + + @pytest.mark.parametrize( ("engine_step_backend", "error", "match"), [ From a47d8de69493bef4ef640c1ebe7422ba0cf330c5 Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Tue, 29 Sep 2026 08:46:10 +0800 Subject: [PATCH 4/7] fix: weight the partial request's cached prefix in the mixed-step pass 1 Pass 1 passed the cached prefix of the floor(ctx / isl_new) complete requests as KV context, but ignored the partial request that a non-multiple budget also schedules. Pass 2 already prices that request by its fill fraction (context_attention_groups), so a module-attention op (DeepSeek/Kimi context_mla_block) saw less cached context in pass 1 than the step contains. Pass 1 now reuses the pass-2 request groups: a group of n requests reads n * prefix cached tokens, and the (complete, 1 - fill) and (complete + 1, fill) groups are weighted like pass 2. Ops whose cost does not depend on the prefix (GEMM, MoE, comm, norms) return equal values for both groups and keep one unweighted, exact value. The token count, prefix 0, no-prefill steps, chunked steps and exact-multiple budgets are unchanged, so no golden moves (all parity records pass; the prefix cases at rtol 1e-12). DeepSeek-V3 b200 vLLM TP8/EP8, bs 8, ISL 16384, prefix 15360: at --ctx-tokens 1536 (one request plus half of another) context_mla_block is 20.17 ms, between the one-request (1024: 11.08 ms) and two-request (2048: 31.81 ms) steps. Tests: the module-attention fixture pins the weighted pass-1 value at a 2.5-request budget and its per-op fold; the plain fixture pins that prefix-free ops stay exact for a partial request. Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- crates/core/src/perfmodel/engine/runtime.rs | 116 ++++++++++++++------ 1 file changed, 85 insertions(+), 31 deletions(-) diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 800ac077c..a8d8238f0 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -783,11 +783,11 @@ impl Engine { /// ```text /// // Pass 1 — combined non-attention work (every budget token is new): /// // run_static(batch=1, isl=ctx+gen*(nextn+1), osl=1, - /// // prefix=(ctx > 0 && prefix > 0) - /// // ? prefix * max(floor(ctx/isl_new), 1) : 0, - /// // mode=static_ctx) - /// // sum every op EXCEPT "context_attention"; the prefix is KV - /// // context for ops that fold attention into a module + /// // prefix=n*prefix, mode=static_ctx) + /// // for the same weighted request groups n as pass 2 (prefix 0 or + /// // ctx 0: n = 0); sum every op EXCEPT "context_attention". The + /// // prefix is KV context for ops that fold attention into a module; + /// // an op that does not read it keeps one unweighted value /// // Pass 2 — context attention at the prefill shape: /// // prefix == 0 or ctx < isl_new: /// // run_static(batch=ceil(ctx/isl_new), isl=isl, osl=1, prefix) @@ -1004,15 +1004,21 @@ impl Engine { // name (DeepSeek/Kimi `context_mla_block`) reads it through `prefix`, // while token-major ops (GEMM, MoE, comm, norms) only read the token // count. DSA/DSV4/MSA modules are named `context_attention` and get - // the prefix in pass 2. The budget holds `floor(ctx / isl_new)` - // complete requests, at least the one whose chunk it carries; as in - // the legacy query, they are priced as one sequence over their - // summed prefixes. With `ctx == isl_new` this is exactly the legacy - // `(combined - prefix, prefix)` query, and prefix 0 is unchanged. - let prefix1 = if ctx_tokens == 0 || prefix == 0 { - 0 + // the prefix in pass 2. The budget is filled exactly as in pass 2 + // (`context_attention_groups`): `floor(ctx / isl_new)` complete + // requests plus a partial request weighing its fill fraction, or the + // one request being chunked. As in the legacy query, the requests of + // a group are priced as one sequence over their summed prefixes, so a + // group of `n` requests reads `n * prefix` cached tokens. With + // `ctx == isl_new` this is exactly the legacy `(combined - prefix, + // prefix)` query; prefix 0 or no prefill reads none. + let kv_groups: Vec<(u32, f64)> = if ctx_tokens == 0 || prefix == 0 { + vec![(0, 1.0)] } else { - prefix.saturating_mul((ctx_tokens / isl_new).max(1)) + Self::context_attention_groups(ctx_tokens, isl_new, prefix) + .into_iter() + .map(|(requests, weight)| (prefix.saturating_mul(requests), weight)) + .collect() }; let mut shared_non_attention = 0.0; for op in &self.context_ops { @@ -1021,17 +1027,35 @@ impl Engine { if op.is_context_attention() || op.name().starts_with("draft_") { continue; } - let result = query_context_op( - op, - &self.db, - 1, - combined, - prefix1, - seq_imbalance_correction_scale, - None, - )?; - shared_non_attention += result.latency_ms; - on_op(MixedPass::SharedNonAttention, op, result); + let query = |kv_prefix: u32| { + query_context_op( + op, + &self.db, + 1, + combined, + kv_prefix, + seq_imbalance_correction_scale, + None, + ) + }; + let first = query(kv_groups[0].0)?; + if kv_groups.len() == 1 { + shared_non_attention += first.latency_ms; + on_op(MixedPass::SharedNonAttention, op, first); + continue; + } + let second = query(kv_groups[1].0)?; + if second.latency_ms == first.latency_ms && second.energy_wms == first.energy_wms { + // Prefix-free op: one unweighted row keeps the value exact. + shared_non_attention += first.latency_ms; + on_op(MixedPass::SharedNonAttention, op, first); + continue; + } + for (result, &(_, weight)) in [first, second].into_iter().zip(&kv_groups) { + let result = result.scaled(weight); + shared_non_attention += result.latency_ms; + on_op(MixedPass::SharedNonAttention, op, result); + } } // ---- Pass 2: context attention at the prefill shape ---- @@ -3307,23 +3331,26 @@ mod tests { } #[test] - fn mixed_step_pass_one_reads_the_cached_prefix_of_the_complete_requests() { + fn mixed_step_pass_one_reads_the_cached_prefix_of_the_packed_requests() { // A module-attention op priced in pass 1 must see the cached prefix - // of the floor(ctx/isl_new) complete requests the budget holds (at - // least the one being chunked) and nothing when no prefill runs. - // isl 1000, prefix 100 -> isl_new 900; decode 7. + // of the requests the budget holds, weighted like pass 2: the one + // request being chunked, exactly one request, or two complete + // requests plus a partial one weighing its fill fraction; and nothing + // when no prefill runs. isl 1000, prefix 100 -> isl_new 900; decode 7. let engine = build_engine_with_module_attention(); let (decode, isl, osl, prefix) = (7_u32, 1000_u32, 64_u32, 100_u32); for (ctx, kv_prefix) in [ (300_u32, 100_u32), // one chunk of one request (900, 100), // exactly one request - (2250, 200), // two complete requests plus a partial one ] { let got = engine .mixed_step_breakdown(ctx, decode, isl, osl, prefix, 1.0, 1.0) .unwrap(); - let expected = shared_pass_reference(&engine, ctx + decode, kv_prefix); - assert_eq!(got[1], expected, "pass 1 ctx={ctx}"); + assert_eq!( + got[1], + shared_pass_reference(&engine, ctx + decode, kv_prefix), + "pass 1 ctx={ctx}" + ); assert!( got[1] > shared_pass_reference(&engine, ctx + decode, 0), "the module op must read the prefix (ctx={ctx})" @@ -3331,6 +3358,21 @@ mod tests { // Context attention proper is absent from this fixture. assert_eq!(got[2], 0.0); } + // 2250 = 2 * 900 + 450: two complete requests (200 cached tokens) + // plus a partial one of fill 450/900 = 1/2 (300 cached tokens). + let got = engine + .mixed_step_breakdown(2250, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + let two = shared_pass_reference(&engine, 2257, 200); + let three = shared_pass_reference(&engine, 2257, 300); + let expected = two / 2.0 + three / 2.0; + assert!((got[1] - expected).abs() < 1e-12 * expected); + assert!(two < got[1] && got[1] < three); + // The per-op surface folds both weighted rows of the module op. + let (shared, _, _) = engine + .mixed_step_breakdown_per_op(2250, decode, isl, osl, prefix, 1.0, 1.0) + .unwrap(); + assert!((shared.iter().map(|r| r.1).sum::() - got[1]).abs() < 1e-12 * got[1]); // No prefill scheduled: the prefix is inert in pass 1 as well. let decode_only = engine .mixed_step_breakdown(0, decode, isl, osl, prefix, 1.0, 1.0) @@ -3338,6 +3380,18 @@ mod tests { assert_eq!(decode_only[1], shared_pass_reference(&engine, decode, 0)); } + #[test] + fn mixed_step_pass_one_keeps_prefix_free_ops_exact_for_a_partial_request() { + // In the plain fixture no pass-1 op reads the prefix: a non-multiple + // budget with a prefix must price pass 1 exactly as one unweighted + // query, not as a rounded convex combination of equal values. + let engine = build_engine(None); + let got = engine + .mixed_step_breakdown(2250, 7, 1000, 64, 100, 1.0, 1.0) + .unwrap(); + assert_eq!(got[1], shared_pass_reference(&engine, 2257, 0)); + } + #[test] fn mixed_step_prefix_at_or_beyond_isl_rejects_prefill_only() { let engine = build_engine(None); From df7fc8bebd9448dc6b671ae941d33c10cfe23315 Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Tue, 29 Sep 2026 08:46:10 +0800 Subject: [PATCH 5/7] docs: qualify the --prefix mix-step, throughput and TTFT claims A larger prefix at a fixed --ctx-tokens can reduce the number of mix steps and can raise throughput, but neither is guaranteed: the step count only drops once the batch's uncached tokens need fewer budget-sized steps, and each step's cost grows with the cached context. The TTFT sentence is qualified the same way: a step that already holds one whole request costs about the same as without the prefix only when attention is priced per request. For module-attention models (DeepSeek / Kimi MLA) the packed requests are priced as one sequence over their summed prefixes, so step cost and TTFT can grow with the prefix (DeepSeek-V3 b200 vLLM TP8/EP8, bs 8, ISL 16384, --ctx-tokens 16384: TTFT 871 ms at prefix 0, 1379 ms at prefix 12288). Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- docs/cli/legacy-aic-user-guide.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cli/legacy-aic-user-guide.md b/docs/cli/legacy-aic-user-guide.md index e5a7a8527..713501efe 100644 --- a/docs/cli/legacy-aic-user-guide.md +++ b/docs/cli/legacy-aic-user-guide.md @@ -149,7 +149,7 @@ aiconfigurator cli estimate --model-path Qwen/Qwen3-32B --system h200_sxm --tp-s - `--fmha-quant-mode`: FMHA quantization mode (auto-inferred if omitted). DeepSeek-V3 / Kimi-K2.5 context attention (MLA prefill) has no fp8 FMHA data — auto-inferred fp8 is downgraded to `bfloat16` for context-touching estimates (agg, prefill, static, static_ctx, AFD prefill), while decode/generation keeps fp8. Explicitly passing `fp8` for a context estimate is rejected with a clear error. - `--moe-quant-mode`: MoE quantization mode (auto-inferred if omitted) - `--comm-quant-mode`: Communication quantization mode (auto-inferred; default `half`) -- `--prefix`: Prefix cache length (subset of ISL already cached per request). Default: `0`. Cached tokens are KV context only: the mixed-step schedule packs requests by their uncached `ISL - prefix` tokens, so at a fixed `--ctx-tokens` a larger prefix means fewer mix steps and higher throughput. TTFT drops when the uncached prefill needs fewer chunks (`ceil((ISL - prefix) / ctx_tokens)`) or queues behind fewer steps; a step that already holds one whole request costs about the same as without the prefix (its new tokens attend over more cached KV), so TTFT is roughly flat in that regime. A budget that is not a multiple of `ISL - prefix` is filled with one partial request, never rounded up to a whole one; the per-step request count is capped at `--batch-size` +- `--prefix`: Prefix cache length (subset of ISL already cached per request). Default: `0`. Cached tokens are KV context only: the mixed-step schedule packs requests by their uncached `ISL - prefix` tokens, so at a fixed `--ctx-tokens` a larger prefix can reduce the number of mix steps (it does once the uncached tokens of the batch no longer need as many budget-sized steps) and can raise throughput, depending on how much each step's cost grows with the cached context. TTFT drops when the uncached prefill needs fewer chunks (`ceil((ISL - prefix) / ctx_tokens)`) or queues behind fewer steps; a step that already holds one whole request costs about the same as without the prefix when attention is priced per request (its new tokens attend over more cached KV), so TTFT is roughly flat in that regime. For models whose attention is a module op (DeepSeek/Kimi MLA), the requests packed into one step are priced as one sequence over their summed prefixes, so the step cost and TTFT can grow with the prefix. A budget that is not a multiple of `ISL - prefix` is filled with one partial request, never rounded up to a whole one; the per-step request count is capped at `--batch-size` - `--nextn`: MTP draft length, or `auto` to use the checkpoint's `num_nextn_predict_layers` (absent/0 keeps MTP disabled). Default: `0`; MTP is never enabled implicitly — pass `--nextn` explicitly to model it - `--nextn-accepted`: Average accepted draft tokens per decode step (`0 <= nextn_accepted <= nextn`). Required when the draft depth is > 0 (including via `--nextn auto`); use a measured value from your deployment - `--stride`: (static modes only) OSL-sweep stride used by `run_static`; ignored for `agg`/`disagg`. Default: `32` From d802e9ae5cb49d0a2956d096b6bb45daabc6c1bd Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Tue, 29 Sep 2026 09:37:10 +0800 Subject: [PATCH 6/7] fix(sdk): reject prefix >= isl before building the agg sweep grid run_agg rejects a request that carries no uncached token (prefix >= the effective isl). Both agg sweeps (sdk/sweep.py _sweep_one_parallel_agg and the legacy BaseBackend.find_best_agg_result_under_constraints) clamped isl_new to 1 instead, built a grid starting at ctx_tokens 1, and only failed at their first run_agg call, mid-sweep. They now raise the same ValueError up front ("prefix (P) must be smaller than the effective isl (I) for an agg sweep"), before any grid point is estimated. Test: both sweeps reject prefix == isl and prefix > isl without visiting a grid point. Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- python/aisimulate/src/aisimulate/sdk/sweep.py | 7 ++- .../sdk/backends/base_backend.py | 7 ++- .../tests/unit/sdk/sweep/test_sweep.py | 46 +++++++++++++++++++ 3 files changed, 58 insertions(+), 2 deletions(-) diff --git a/python/aisimulate/src/aisimulate/sdk/sweep.py b/python/aisimulate/src/aisimulate/sdk/sweep.py index 8dcabd08c..36227132d 100644 --- a/python/aisimulate/src/aisimulate/sdk/sweep.py +++ b/python/aisimulate/src/aisimulate/sdk/sweep.py @@ -337,7 +337,12 @@ def _sweep_one_parallel_agg( isl = runtime_config.isl + BaseBackend._visual_context_tokens(model, runtime_config) # ctx_tokens budgets UNCACHED prefill tokens (run_agg's isl_new), so the # feasibility guards and the capped-gen dedup pack requests by it. - isl_new = max(isl - int(runtime_config.prefix or 0), 1) + # run_agg rejects a request with no uncached token; fail before the grid + # is built instead of at its first point. + prefix = int(runtime_config.prefix or 0) + if prefix >= isl: + raise ValueError(f"prefix ({prefix}) must be smaller than the effective isl ({isl}) for an agg sweep") + isl_new = isl - prefix osl = runtime_config.osl ttft_target = runtime_config.ttft tpot_target = runtime_config.tpot diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py index 9e7896888..06b1b19b0 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py @@ -2158,7 +2158,12 @@ def find_best_agg_result_under_constraints( isl_eff = isl + self._visual_context_tokens(model, runtime_config) # ctx_tokens budgets UNCACHED prefill tokens (run_agg's isl_new), so the # feasibility guards and the capped-gen dedup pack requests by it. - isl_new = max(isl_eff - int(runtime_config.prefix or 0), 1) + # run_agg rejects a request with no uncached token; fail before the + # grid is built instead of at its first point. + prefix = int(runtime_config.prefix or 0) + if prefix >= isl_eff: + raise ValueError(f"prefix ({prefix}) must be smaller than the effective isl ({isl_eff}) for an agg sweep") + isl_new = isl_eff - prefix osl = runtime_config.osl ttft = runtime_config.ttft tpot = runtime_config.tpot diff --git a/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py b/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py index 9f6cb1842..9bacbc67b 100644 --- a/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py +++ b/python/aisimulate/tests/unit/sdk/sweep/test_sweep.py @@ -505,6 +505,52 @@ def test_sweep_agg_prefix_grid_covers_every_batch_and_mirrors_legacy(monkeypatch assert legacy_points == points +@pytest.mark.parametrize("prefix", [2048, 2049]) +def test_sweep_agg_rejects_prefix_at_or_beyond_isl_before_the_grid(monkeypatch, prefix): + """run_agg rejects a request with no uncached token, so both agg sweeps + must fail up front, before any grid point is estimated, rather than at + their first run_agg call.""" + from aisimulate.sdk.backends.factory import get_backend + + runtime_config = config.RuntimeConfig(isl=2048, osl=64, prefix=prefix, ttft=1e9, tpot=1e9) + + points: list[tuple[int, int]] = [] + monkeypatch.setattr(sweep, "predict_agg_worker", _recording_agg_worker(points)) + with pytest.raises(ValueError, match=r"prefix \(\d+\) must be smaller than the effective isl \(2048\)"): + sweep._sweep_one_parallel_agg( + model=MagicMock(), + backend=MagicMock(), + database=MagicMock(), + runtime_config=runtime_config, + top_k=0, + max_batch_size=8, + ctx_stride=512, + enable_chunked_prefill=False, + free_gpu_memory_fraction=None, + max_seq_len=None, + ) + assert points == [] + + legacy_backend = get_backend("sglang") + legacy_calls: list[int] = [] + monkeypatch.setattr( + legacy_backend, + "run_agg", + lambda model, database, runtime_config, **kwargs: legacy_calls.append(kwargs["ctx_tokens"]), + ) + with pytest.raises(ValueError, match=r"prefix \(\d+\) must be smaller than the effective isl \(2048\)"): + legacy_backend.find_best_agg_result_under_constraints( + MagicMock(), + MagicMock(), + runtime_config, + top_k=0, + max_batch_size=8, + ctx_stride=512, + enable_chunked_prefill=False, + ) + assert legacy_calls == [] + + # --------------------------------------------------------------------------- # sweep_disagg validation # --------------------------------------------------------------------------- From 1f4c130b877c2bf18bcf7e5be1274a766fc881a7 Mon Sep 17 00:00:00 2001 From: Kang Zhang Date: Tue, 29 Sep 2026 10:38:21 +0800 Subject: [PATCH 7/7] fix(sdk): price whole-batch prefill steps without decode requests or TPOT weight When a mixed step's prefilling requests (complete ones plus a partial last one, ceil(ctx_budget / isl_new)) cover the whole batch, no request is decoding during the first mixed step. run_agg only recognized this when the budget held every request's full uncached prefill (ctx >= b * isl_new): - A budget reaching into the last request (e.g. bs 2, ISL 2048, prefix 128, --ctx-tokens 2048: one complete request plus a 128-token partial one) still asked _mix_step_gen_tokens for decode requests, which returns at least one. The step was priced with b + 1 requests while the result reported num_gen_reqs = 0. - A mixed step without decode requests still entered the TPOT average, although it produces no output token and belongs to TTFT. Key both on the same condition, ceil(ctx_budget / isl_new) >= b. The mixed step prices no decode request. When gen-only steps follow, the first (decode-free) mixed step leaves the TPOT average: with a single mixed step (ctx >= b * isl_new) TPOT is the decode step; with the second mixed step of a partial last request, which is shared with the b - 1 requests already decoding, that step still counts (_tpot_mix_steps(num_mix_steps - 1)). SGLang and TRT-LLM TPOT is therefore continuous across ctx = (b - 1) * isl_new; vLLM, which counts every mixed step, drops by the one decode-free step there. Where decode_iterations <= steps_to_finish_ctx (osl <= 1 + decode_tokens_per_iteration) no gen-only step exists and the TPOT formula is unchanged. Qwen3-32B-FP8, h200 SGLang 0.5.14, TP2, bs 2, ISL 2048, OSL 128, prefix 128, --ctx-tokens 2048: mixed step 87.90 -> 76.91 ms, TTFT 171.40 -> 149.98 ms, TPOT 10.705 -> 10.618 ms. DeepSeek-V3 b200 vLLM TP8/EP8, bs 2, ISL 2048, OSL 16, prefix 1024, --ctx-tokens 2048 (a single whole-batch mixed step): TPOT 15.334 -> 10.372 ms, and the vLLM throughput cap follows (77.74 -> 96.32 tokens/s). At prefix 0 this applies only to explicit budgets above (b - 1) * ISL; the agg sweeps never reach that regime (their guard keeps a decoding request) and default budgets are unaffected. Tests: the partial-last-request case and its one-request contrast with hand-derived TPOT; TPOT equals the decode step on the single-step whole-batch cap cases. Signed-off-by: Kang Zhang Co-Authored-By: Claude Opus 5.5 --- .../sdk/backends/base_backend.py | 23 +++++-- .../unit/sdk/backends/test_base_backend.py | 62 ++++++++++++++++++- 2 files changed, 78 insertions(+), 7 deletions(-) diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py index 06b1b19b0..68dd0705b 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py @@ -1698,8 +1698,9 @@ def run_agg( # large prefix), which would publish negative decode-request counts; # the old full-isl schedule already did so for an explicit # ctx_tokens > b * isl, and for b == 1 priced phantom requests. - # ``ctx_tokens`` stays the published engine knob. With prefix 0 this - # only differs from the old schedule when ctx_tokens >= b * isl. + # ``ctx_tokens`` stays the published engine knob. With prefix 0 the + # schedule only differs from the old one when ctx_tokens > (b - 1) * isl, + # where the step's prefilling requests cover the whole batch. ctx_budget = min(ctx_tokens, b * isl_new) # None (or an omitted kwarg) means the caller did not model speculative # progress here; the summary then stays eligible for the upper-layer @@ -1786,9 +1787,11 @@ def run_agg( num_mix_steps = num_genonly_steps = 0 num_mix_steps_for_tpot_calc = 0 # correction for tpot calc only if b > 1: - if ctx_budget >= b * isl_new: - # The whole batch prefills in the single mixed step, so no - # decode request rides along with it. + # The step's prefilling requests (complete ones plus a partial last + # one) already cover the whole batch: no decode request rides + # along with them. + whole_batch_prefills = np.ceil(ctx_budget / isl_new) >= b + if whole_batch_prefills: num_mix_gen_tokens = 0 else: num_mix_gen_tokens = self._mix_step_gen_tokens(b, ctx_budget, isl_new, decode_iterations) @@ -1806,7 +1809,15 @@ def run_agg( num_mix_steps = steps_to_finish_ctx num_genonly_steps = decode_iterations - num_mix_steps num_genonly_tokens = b - num_mix_steps_for_tpot_calc = self._tpot_mix_steps(num_mix_steps) + # When the whole batch prefills, the first mixed step has no + # decoding request and produces no output token, so it leaves + # the TPOT average. A second mixed step (a partial last request + # finishing its prefill) is shared with the b - 1 requests + # already decoding, so it stays. + if whole_batch_prefills: + num_mix_steps_for_tpot_calc = self._tpot_mix_steps(num_mix_steps - 1) if num_mix_steps > 1 else 0 + else: + num_mix_steps_for_tpot_calc = self._tpot_mix_steps(num_mix_steps) elif b == 1: # special case for b=1 num_mix_steps = 1 diff --git a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py index 36701dfa1..c21ba2804 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py @@ -1510,8 +1510,10 @@ def _memory(model_arg, database_arg, batch_size, beam_width, isl, osl, **kwargs) assert result["ctx_tokens"] == _CTX assert result["num_tokens"] == b * isl_new assert steps[0].context_tokens == b * isl_new - # Every request prefills in this step: no decode request rides along. + # Every request prefills in this step: no decode request rides along, + # and the step does not enter the TPOT average (decode steps cost 1 ms). assert steps[0].num_decode_requests == 0 + assert result["tpot"] == pytest.approx(1.0) assert summary.get_step_estimates()["scheduling"]["num_mix_steps"] == 1 assert memory_calls[0]["num_tokens"] == b * isl_new assert memory_calls[0]["mtp_scaled_tokens"] == 0 @@ -1575,6 +1577,9 @@ def _memory(model_arg, database_arg, batch_size, beam_width, isl, osl, **kwargs) result = summary.get_result_dict() assert steps[0].context_tokens == capped assert steps[0].num_decode_requests == 0 + if b > 1 and osl > 2: + # The prefill-only step belongs to TTFT; TPOT is the decode step (1 ms). + assert result["tpot"] == pytest.approx(1.0) assert result["ctx_tokens"] == ctx_tokens assert result["num_tokens"] == capped assert memory_calls[0]["num_tokens"] == capped @@ -1588,6 +1593,61 @@ def _memory(model_arg, database_arg, batch_size, beam_width, isl, osl, **kwargs) ) == pytest.approx(10.0) +@pytest.mark.parametrize( + ("ctx_tokens", "prefilling", "decoding", "expected_tpot"), + [ + # One complete request (1920 new tokens) plus a 128-token partial one: + # both requests are prefilling in the first of ceil(3840/2048) = 2 + # mixed steps, which leaves TPOT; the second one (the partial request + # finishing while the other decodes) stays: (10 + 62 * 1) / 63. + (2048, 2, 0, (10.0 + 62 * 1.0) / 63), + # Exactly one request: the other one decodes alongside it, and both of + # the ceil(3840/1920) = 2 mixed steps count: (2 * 10 + 62 * 1) / 64. + (1920, 1, 1, (2 * 10.0 + 62 * 1.0) / 64), + ], +) +def test_run_agg_partial_last_request_prefills_the_whole_batch( + monkeypatch, + backend: BaseBackend, + model, + database, + ctx_tokens: int, + prefilling: int, + decoding: int, + expected_tpot: float, +) -> None: + """bs 2, isl 2048, prefix 128: once the budget reaches into the last + request, every request of the batch is prefilling, so the mixed step + prices no decode request (never b + 1 requests), and only the decode-free + first mixed step stays out of the TPOT average (mixed step 10 ms, decode + step 1 ms, 64 decode iterations).""" + steps: list[MixedStepInput] = [] + + def _run_mixed(model_arg, database_arg, runtime_config_arg, step): + steps.append(step) + return StepEstimate(latency_ms=10.0, energy_wms=1.0) + + monkeypatch.setattr(backend, "run_mixed", _run_mixed) + monkeypatch.setattr( + backend, + "_get_genonly_step_estimate", + lambda *args, **kwargs: _decode_step(1.0, 1.0, {"decode": 1.0}, {"decode": "silicon"}, ()), + ) + + summary = backend.run_agg( + model, + database, + RuntimeConfig(batch_size=2, beam_width=1, isl=2048, osl=64, prefix=128, engine_step_backend="rust"), + ctx_tokens=ctx_tokens, + ) + + result = summary.get_result_dict() + assert (result["num_ctx_reqs"], result["num_gen_reqs"]) == (prefilling, decoding) + assert steps[0].num_decode_requests == decoding + assert steps[0].context_tokens == ctx_tokens + assert result["tpot"] == pytest.approx(expected_tpot) + + def test_run_agg_budget_cap_is_inert_when_the_batch_owns_more_tokens( monkeypatch, backend: BaseBackend, model, database ) -> None: