diff --git a/infera/projection/agents/tuning_agent/inference_tuning.py b/infera/projection/agents/tuning_agent/inference_tuning.py index 519c7610c..a8b0f9eec 100644 --- a/infera/projection/agents/tuning_agent/inference_tuning.py +++ b/infera/projection/agents/tuning_agent/inference_tuning.py @@ -635,9 +635,7 @@ def validate_inference( if pool_dp < 1: return False, f"{pool}_attention_dp={pool_dp} must be >= 1" if pool_dp > 1 and pool_tp % pool_dp: - return False, ( - f"{pool}_attention_dp={pool_dp} must divide {pool}_tp={pool_tp}" - ) + return False, (f"{pool}_attention_dp={pool_dp} must divide {pool}_tp={pool_tp}") for pool, pool_tp, pool_ep in ( ("prefill", p_tp, cfg.prefill_ep), ("decode", d_tp, cfg.decode_ep), @@ -648,8 +646,7 @@ def validate_inference( return False, f"{pool}_ep={pool_ep} not in legal EP set {legality.ep}" if pool_ep > 1 and (pool_tp * cfg.pp) % pool_ep: return False, ( - f"{pool}_ep={pool_ep} must divide the {pool} pool's " - f"{pool_tp * cfg.pp} ranks" + f"{pool}_ep={pool_ep} must divide the {pool} pool's {pool_tp * cfg.pp} ranks" ) prefill_gpus = p_tp * cfg.pp * cfg.prefill_replicas decode_gpus = d_tp * cfg.pp * cfg.decode_replicas diff --git a/infera/projection/cli.py b/infera/projection/cli.py index b150128a3..ee9bd40ed 100644 --- a/infera/projection/cli.py +++ b/infera/projection/cli.py @@ -528,6 +528,30 @@ def _add_inference_args(parser): help="Host<->device bandwidth for the KV offload tier in GB/s. " "PCIe 5 x16 is ~64; a cache-coherent host link is ~900. Default: 64.", ) + parser.add_argument( + "--workload-resident-tokens", + type=int, + default=None, + help="Mean KV tokens one request of the target workload holds while " + "resident, which sets how many requests the pool can admit at once and " + "so where TTFT stops being a service time. For a spread of lengths pass " + "E[L^2]/E[L], not the mean: a request holds the pool for a time " + "proportional to its length, so the resident set is length-biased. Do " + "not discount it by the prefix-cache hit rate -- a cache hit skips " + "prefill compute, it does not free the blocks. Default: the configured " + "context (input + output/2), exact for a fixed-length workpoint.", + ) + parser.add_argument( + "--kv-pool-tokens", + type=int, + default=None, + help="Size of the KV pool the engine allocated, in tokens, as it " + 'reports at startup (vLLM\'s "GPU KV cache size", SGLang\'s "KV Cache ' + 'is allocated. #tokens"). The admission bound is evaluated against ' + "this instead of against the memory model's estimate, which is what " + "you want when validating against a deployment that is already " + "running. Default: unset, and the pool is predicted.", + ) # ---- Feature B: custom collective ops ---- coll = parser.add_argument_group("inference collectives (feature B)") coll.add_argument( @@ -616,6 +640,17 @@ def _add_inference_args(parser): "parallelism replicates their compressed KV latent instead of sharding " "it, so only splitting by request shrinks the cache a rank holds.", ) + par.add_argument( + "--decode-context-parallel-size", + type=int, + default=None, + help="Split each sequence's context across this many ranks during " + "decode (vLLM --decode-context-parallel-size, SGLang --dcp-size). The " + "other way to shrink an MLA cache: a rank stores context/N of every " + "request rather than all of it, so the replica holds N times as many " + "tokens. Unlike --attention-dp-size it shrinks a single request's " + "footprint, so it raises the concurrency ceiling even at batch 1.", + ) # ---- Feature A: prefill/decode disaggregation ---- dis = parser.add_argument_group("inference disaggregation (feature A)") dis.add_argument( @@ -790,6 +825,20 @@ def _add_inference_args(parser): "only (not throughput). A property of the serving stack, so it has no default " "worth guessing. Default 0.", ) + serv.add_argument( + "--uncached-prompt-latency-us", + type=float, + default=None, + help="TTFT latency per prompt token the prefix cache did not serve (us/token). " + "Added to TTFT and end-to-end latency only (not throughput). Default 0.", + ) + serv.add_argument( + "--uncached-prompt-latency-max-tokens", + type=int, + default=None, + help="Uncached tokens past which --uncached-prompt-latency-us stops growing. " + "Default 0 (no cap).", + ) serv.add_argument( "--prefill-rate-us-per-token", type=float, @@ -853,6 +902,14 @@ def _add_inference_args(parser): "prefill-chunk + concurrent-decode tokens per step; oversized steps split, " "raising TPOT. Default: 0 (unlimited).", ) + serv.add_argument( + "--max-num-seqs", + type=int, + default=None, + help="Engine cap on running sequences (vLLM --max-num-seqs, SGLang " + "--max-running-requests). Clients beyond it queue at the server. " + "Default: 0 (bounded by client concurrency only).", + ) # ---- Offered load / request rate (open-loop arrivals) ---- serv.add_argument( "--request-rate", @@ -939,6 +996,75 @@ def _add_inference_args(parser): "fixed-concurrency run is only a few waves long and the opening burst " "carries the highest TTFT of any request in it. Default: 0.1.", ) + serv.add_argument( + "--des-client-think-ms", + type=float, + default=0.0, + help="DES: how long a closed-loop client waits after its last token " + "before issuing its next request. Zero re-issues immediately, which " + "is a load generator saturating the server and not what an agentic " + "harness does: a turn ends, the agent runs a tool, and the next turn " + "is sent when that returns. It is a property of the workload, like " + "the prompt lengths and the prefix reuse, and leaving it out does not " + "leave the replay neutral -- it holds every client permanently " + "resident. On the AgentX ladders the hardware's clients are idle for " + "50-74% of each turn's cycle below saturation and ~0% above it, so " + "omitting it overstates throughput several-fold at the low rungs and " + "not at all at the high ones, which bends the curve rather than " + "shifting it. A Mooncake trace's per-request ``think_ms`` overrides " + "it request by request.", + ) + serv.add_argument( + "--des-client-idle-cap-ms", + type=float, + default=0.0, + help="DES: longest the engine may sit with no request in flight before " + "pending client timers are pulled forward (0 = no cap). Mirrors a " + "replay harness's system-idle guard (AIPerf AgentX: 10 s), which " + "shortens a single client's idle time far below the trace's gaps.", + ) + serv.add_argument( + "--des-admit-backlog-only", + action="store_true", + help="DES: measure the longest-prefix admission window against the " + "backlog rather than the client count, so that below what the KV pool " + "holds -- where nothing is queued and there is no choice for a policy " + "to make -- admission is the order the lanes arrived in. This is the " + "correct account of what a scheduler can reorder, and it is the band " + "where measured reuse collapses (kimik3 at C=14 thrashed to 0.67 " + "against 0.93 modelled), but on the AgentX corpus it improves ITL " + "ordering by one model-engine pair and costs one on throughput, so it " + "is off by default.", + ) + serv.add_argument( + "--des-warmup-requests", + type=int, + default=0, + help="DES: exclude this many requests, in the order the clients issued " + "them, from the reported latency distribution. Prefer this to " + "--des-warmup-frac for a closed loop: all C clients fire at once, so " + "the first C requests queue against each other and wait far longer " + "than anything after them, and dropping the earliest *completions* " + "keeps every one of them -- a request that waited 40s for a slot is " + "among the last to finish. The harness excludes the same transient by " + "advancing each lane AIPERF_WARMUP_REQUESTS_PER_LANE requests and " + "draining before it starts profiling, so the matching value is that " + "many per client. Default: 0 (use --des-warmup-frac).", + ) + serv.add_argument( + "--des-duration-s", + type=float, + default=0.0, + help="DES: stop the run after this many seconds of simulated time and " + "report over whatever completed, instead of running to " + "--des-num-requests. This is how a fixed-concurrency harness is " + "actually bounded, and under a closed loop the two are not " + "interchangeable: a request budget divided among C lanes fixes how " + "far each lane walks into its conversation, and an agentic turn's " + "prompt grows with its position, so a budget-bounded run offers " + "different prompt lengths at each concurrency than the measured run " + "it is compared against. Default: 0 (run to the request count).", + ) serv.add_argument( "--des-seed", type=int, @@ -1064,6 +1190,34 @@ def _add_inference_args(parser): help="DES: per-instance KV block-cache capacity (blocks kept resident, " "LRU-evicted under pressure). 0 = unbounded within the run. Default: 0.", ) + serv.add_argument( + "--des-whole-context-residency", + action="store_true", + help="DES: charge a running request the whole context it attends to, " + "not just the part of it that was not already cached. The discount is " + "an account of one physical copy, and it is only that where the " + "sharers are running at the same time. Under a closed loop on " + "conversational traffic they are not: each client has one turn in " + "flight, reuse is a conversation hitting its own previous turn, and " + "the turns running together belong to different conversations whose " + "histories do not overlap. Discounted anyway, a pool holding five of " + "these conversations reports room for hundreds, never fills, and so " + "never shows reuse falling away or the queue that follows it.", + ) + serv.add_argument( + "--des-cache-shares-pool", + action="store_true", + help="DES: size the prefix cache to what the running requests leave " + "free, instead of to the whole pool. A request is already not charged " + "for the blocks it hit on, because something else is holding them; " + "that something is this cache, and giving it the pool as well books " + "the same memory twice. Where the pool is roomy the correction is " + "nothing, and where it is not it is the whole behaviour: a pool with " + "space for five of these conversations cannot also hold their history, " + "so reuse falls away as concurrency climbs and TTFT leaves the scale " + "the low rungs were on. Double-booked, the replay reports reuse near " + "its ceiling and a flat TTFT straight through that.", + ) serv.add_argument( "--des-mooncake-trace", type=str, @@ -1084,6 +1238,17 @@ def _add_inference_args(parser): help="Attention kernel library (ROCm). Representative compute multiplier " "vs the Triton baseline. Default: engine default (1.0).", ) + kern.add_argument( + "--serving-engine", + type=str, + default=None, + help="Engine this projection is about (vllm, sglang, atom, or a build " + "such as mori-sglang). Simulate mode is unaffected -- it is analytical " + "and returns the same number whichever engine is named. This is a " + "regime axis, so it decides which measured anchor benchmark mode may " + "calibrate from: an anchor harvested on another engine is refused " + "rather than substituted. Default: unset (matches as before).", + ) kern.add_argument( "--sparse-attention-topk", type=int, @@ -1092,6 +1257,18 @@ def _add_inference_args(parser): "query. Attention scales toward topk/context for long contexts. " "Default: 0 (dense).", ) + kern.add_argument( + "--sparse-indexer-cost-scale", + type=float, + default=None, + help="What the sparse indexer's top-k selection costs on this stack " + "relative to a fused kernel. The arithmetic is the model's, the kernel " + "is the serving stack's: the same selection runs fused through aiter " + "on gfx950 and unfused through Torch on gfx942. 1.0 prices the fused " + "path; raise it to charge a stack serving without one. Only reaches " + "models priced from an indexer geometry (GLM-5.2, MiniMax-M3), not " + "ones carrying a per-layer compression schedule. Default: 1.0.", + ) kern.add_argument( "--sliding-window", type=int, diff --git a/infera/projection/configs/models/megatron/glm5.yaml b/infera/projection/configs/models/megatron/glm5.yaml index 50c8bfb1a..7fa424831 100644 --- a/infera/projection/configs/models/megatron/glm5.yaml +++ b/infera/projection/configs/models/megatron/glm5.yaml @@ -12,6 +12,27 @@ extends: # declared below so tooling does not have to infer it from this comment. sparse_attention_topk: 2048 +# IndexShare is a memory property as much as a compute one. The indexer stores +# a 128-wide key per token (``index_head_dim: 128``), but only on the layers +# that compute one: ``indexer_types`` marks 21 of the 78 layers ``full`` and +# the other 57 ``shared``, and a shared layer reuses the keys of the full +# layer above it rather than storing its own. GLM keeps them fp8, unlike +# MiniMax-M3's bf16, so here the index cache is a 6% correction rather than a +# third of the footprint: against vLLM at TP8 this predicts 47,616 bytes per +# token with an fp8 KV cache where the engine allocates 47,805, and 92,544 +# with a bf16 cache where it allocates 92,815. +sparse_index_head_dim: 128 +sparse_index_layers: 21 +sparse_index_dtype: fp8 + +# The indexer's query heads (HF ``index_n_heads: 32``). Storage is unaffected +# -- the 128-wide key above is shared across them -- but the arithmetic is +# not, because each head scores the whole context. Without it GLM has no +# indexer term at all and prices back to a floored top-k, which on this +# corpus is a constant 0.15 of dense from 136k to 365k: an over-read of the +# selection it names and a free pass on the indexing it does not. +sparse_index_n_heads: 32 + tokenizer_type: HuggingFaceTokenizer ##### TODO update to GLM-5 tokenizer tokenizer_model: THUDM/glm-4-9b # GLM-5 tokenizer requires newer transformers; use GLM-4 for mock_data testing diff --git a/infera/projection/configs/models/megatron/mamba_370M.yaml b/infera/projection/configs/models/megatron/mamba_370M.yaml index 6665da3af..208c5818e 100644 --- a/infera/projection/configs/models/megatron/mamba_370M.yaml +++ b/infera/projection/configs/models/megatron/mamba_370M.yaml @@ -1,4 +1,4 @@ -bases: +extends: - mamba_base.yaml # Mamba 370M configuration diff --git a/infera/projection/configs/models/megatron/mamba_base.yaml b/infera/projection/configs/models/megatron/mamba_base.yaml index d52fe6db2..75fd64aae 100644 --- a/infera/projection/configs/models/megatron/mamba_base.yaml +++ b/infera/projection/configs/models/megatron/mamba_base.yaml @@ -1,4 +1,4 @@ -bases: +extends: - language_model.yaml # Mamba-specific configuration diff --git a/infera/projection/configs/models/megatron/minimax_m2.5.yaml b/infera/projection/configs/models/megatron/minimax_m2.5.yaml index 717ece8af..622e72ab9 100644 --- a/infera/projection/configs/models/megatron/minimax_m2.5.yaml +++ b/infera/projection/configs/models/megatron/minimax_m2.5.yaml @@ -1,4 +1,4 @@ -bases: +extends: - language_model.yaml # MiniMax-M2.5 Model (GQA + MoE) diff --git a/infera/projection/configs/models/megatron/minimax_m3.yaml b/infera/projection/configs/models/megatron/minimax_m3.yaml index 5c37a7d5a..f636389f2 100644 --- a/infera/projection/configs/models/megatron/minimax_m3.yaml +++ b/infera/projection/configs/models/megatron/minimax_m3.yaml @@ -1,4 +1,4 @@ -bases: +extends: - language_model.yaml # MiniMax-M3 (GQA + sparse attention + MoE) @@ -16,6 +16,25 @@ bases: # as a serving default rather than as an architectural dimension. sparse_attention_topk: 2048 +# Choosing those 16 blocks has a memory cost, and on this model it is a third +# of the footprint. The indexer stores a 128-wide key per token +# (``sparse_index_dim: 128``) on each of the 57 layers where +# ``sparse_attention_freq`` is set -- the first three are dense and have no +# indexer -- and keeps them in bf16, so unlike the KV cache they do not shrink +# when it is quantized. Against vLLM at TP8 this predicts 29,952 bytes per +# token with an fp8 cache where the engine allocates 30,023, and 45,312 with a +# bf16 cache where it allocates 45,398. +sparse_index_head_dim: 128 +sparse_index_layers: 57 +sparse_index_dtype: bf16 + +# The indexer's query heads (HF ``sparse_num_index_heads: 4``). Eight times +# narrower than GLM-5.2's 32 against attention layers that are themselves +# four times cheaper -- GQA over 64 heads of 128 rather than absorbed MLA -- +# so M3's indexer lands at a comparable fraction of its own attention from +# very different parts. Both are invisible without this field. +sparse_index_n_heads: 4 + tokenizer_type: HuggingFaceTokenizer tokenizer_model: MiniMaxAI/MiniMax-M3 diff --git a/infera/projection/configs/models/megatron/qwen3_5_397B_A17B.yaml b/infera/projection/configs/models/megatron/qwen3_5_397B_A17B.yaml new file mode 100644 index 000000000..f5f9902df --- /dev/null +++ b/infera/projection/configs/models/megatron/qwen3_5_397B_A17B.yaml @@ -0,0 +1,91 @@ +extends: + - qwen2.5_base.yaml + +# Qwen3.5-397B-A17B — text stack of `Qwen3_5MoeForConditionalGeneration` +# (HF `text_config`). 397B total, ~17B active. +# https://huggingface.co/Qwen/Qwen3.5-397B-A17B +# +# The vision tower in the published checkpoint is not described here: every +# AgentX run behind this model is text, so nothing drives it. +# +# What makes this checkpoint serve the way it does is the attention layout. +# `full_attention_interval: 4` means only every fourth layer keeps a per-token +# KV cache -- 15 of 60 -- and the other 45 run a gated-delta linear-attention +# layer whose state is a fixed d x d matrix per head. So a token costs a +# quarter of the KV a 60-layer GQA model would store, and three quarters of +# the attention layers cost the same at 200k context as at 1k. At the context +# these runs use (AgentX averages well over 100k) that is not a correction, +# it is most of the model's serving behaviour, and it is why this model holds +# roughly four times the concurrency its parameter count suggests. +# +# Both halves are charged. `linear_attention_freq` derives the 15/45 split, +# and `linear_attention_head_dim` is the state extent (`linear_key_head_dim` +# in HF), which sets both the recurrent state's bytes and the FLOP-equivalent +# KV length of those layers. Sizing the memory without the arithmetic would +# let a search fit batches whose attention was still priced as 60-layer GQA. + +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: Qwen/Qwen3.5-397B-A17B + +# model architecture +num_layers: 60 +hidden_size: 4096 +num_attention_heads: 32 + +# Every layer is MoE (`mlp_only_layers` is empty), so no dense FFN width is +# ever used. The key is required by the schema; it is set to the expert width +# rather than a made-up number so that a misread never inflates a layer. +ffn_hidden_size: 1024 + +# MoE — 512 experts, 10 routed + 1 shared, expert intermediate 1024. +# 60 x 512 x 3 x 4096 x 1024 = 386.5B, which with attention and the two +# untied embedding matrices reproduces the published 397B. +moe_layer_freq: "1" +num_experts: 512 +moe_router_topk: 10 +moe_ffn_hidden_size: 1024 +moe_shared_expert_intermediate_size: 1024 +moe_shared_expert_gate: true +moe_router_load_balancing_type: global_aux_loss +moe_aux_loss_coeff: 0.001 +moe_router_dtype: fp32 + +# GQA on the full-attention layers: 32 Q heads, 2 KV heads, head_dim 256. +# The head dim is twice the usual 128, so the per-token cache of one of these +# layers is what a 4-KV-head 128-dim layer would store -- the saving is the +# layer count, not the head size. +group_query_attention: true +num_query_groups: 2 +qk_layernorm: true +kv_channels: 256 +attention_output_gate: true # attn_output_gate + +# Hybrid attention: full attention every 4th layer (15 of 60), the rest are +# gated-delta linear layers keeping a 128 x 128 state per head. +experimental_attention_variant: gated_delta_net +linear_attention_freq: 4 +linear_attention_head_dim: 128 # linear_key_head_dim +linear_attention_conv_kernel: 4 # linear_conv_kernel_dim +linear_conv_kernel_dim: 4 +linear_key_head_dim: 128 +linear_value_head_dim: 128 +linear_num_key_heads: 16 +linear_num_value_heads: 64 + +# One next-token-prediction head ships with the checkpoint. +mtp_num_layers: 1 + +# partial rotary embedding (partial_rotary_factor 0.25) on the full layers +rotary_base: 10000000 +rotary_percent: 0.25 + +# other settings +vocab_size: 248320 +make_vocab_size_divisible_by: 64 +max_position_embeddings: 262144 +norm_epsilon: 1.0e-6 +init_method_std: 0.02 +untie_embeddings_and_output_weights: true # tie_word_embeddings: false +apply_layernorm_1p: true +apply_wd_to_qk_layernorm: true +bias_swiglu_fusion: false diff --git a/infera/projection/configs/models/megatron/zebra_llama_1B.yaml b/infera/projection/configs/models/megatron/zebra_llama_1B.yaml index d1afa9531..996c49724 100644 --- a/infera/projection/configs/models/megatron/zebra_llama_1B.yaml +++ b/infera/projection/configs/models/megatron/zebra_llama_1B.yaml @@ -1,4 +1,4 @@ -bases: +extends: - mamba_base.yaml # Zebra Llama 8B configuration diff --git a/infera/projection/configs/models/megatron/zebra_llama_3B.yaml b/infera/projection/configs/models/megatron/zebra_llama_3B.yaml index 23090841f..190dc2bbf 100644 --- a/infera/projection/configs/models/megatron/zebra_llama_3B.yaml +++ b/infera/projection/configs/models/megatron/zebra_llama_3B.yaml @@ -1,4 +1,4 @@ -bases: +extends: - mamba_base.yaml # Zebra Llama 8B configuration diff --git a/infera/projection/configs/models/megatron/zebra_llama_8B.yaml b/infera/projection/configs/models/megatron/zebra_llama_8B.yaml index 0237a652d..0387f0115 100644 --- a/infera/projection/configs/models/megatron/zebra_llama_8B.yaml +++ b/infera/projection/configs/models/megatron/zebra_llama_8B.yaml @@ -1,4 +1,4 @@ -bases: +extends: - mamba_base.yaml # Zebra Llama 8B configuration diff --git a/infera/projection/core/projection/inference_projection/benchmark_serving.py b/infera/projection/core/projection/inference_projection/benchmark_serving.py index aa8b9213a..85ecb3bb3 100644 --- a/infera/projection/core/projection/inference_projection/benchmark_serving.py +++ b/infera/projection/core/projection/inference_projection/benchmark_serving.py @@ -49,6 +49,7 @@ import contextlib import json import os +import random import re import shlex import shutil @@ -112,6 +113,99 @@ def _capture_engine_output(path: str): os.close(saved[1]) +def _speculative_argv(args) -> list[str]: + """The flags that turn the draft head on, in each engine's spelling. + + Speculation was previously the one regime axis the served harvest could not + reach: the flags were accepted on the command line, recorded in the + artifact's regime signature, and never passed to the engine. The anchor + therefore described a non-speculative machine while claiming to describe a + speculative one. Since every AgentX deployment serves with MTP, that is the + difference between an anchor the store accepts and one it refuses. + + The spellings are not translations of each other. DeepSeek's MTP head is + loaded through SGLang's EAGLE path at topk 1 -- a single draft chain -- + rather than NEXTN, whose V3/R1 loader does not accept the V4 architecture. + ATOM names the method directly, and vLLM takes one JSON object. + + Acceptance is pinned rather than earned, for the reason given on + ``--speculative-acceptance-length``: these harvests run on dummy weights. + """ + method = getattr(args, "speculative_method", None) + k = getattr(args, "speculative_num_tokens", None) + if not method or not k: + return [] + al = getattr(args, "speculative_acceptance_length", None) + draft = getattr(args, "speculative_draft_model", None) + # The method is part of the regime, not a hint. Collapsing every method to + # MTP made the EAGLE3 and DSpark harvests look for draft layers inside + # checkpoints that do not have them: MiniMax-M3-MXFP4 under vLLM failed + # "Failed to load MTP layer 0 weights" with the EAGLE3 draft it was handed + # never passed on. + separate_draft = str(method).lower() in ("eagle3", "eagle", "draft_model", "dflash") + + if args.serving_backend == "sglang": + if separate_draft: + argv = [ + "--speculative-algorithm", + "EAGLE3" if str(method).lower() == "eagle3" else "EAGLE", + "--speculative-num-steps", + str(k), + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + str(k + 1), + ] + if draft: + argv += ["--speculative-draft-model-path", str(draft)] + if al: + os.environ.setdefault("SGLANG_SIMULATE_ACC_LEN", str(al)) + os.environ.setdefault("SGLANG_SIMULATE_ACC_METHOD", "match-expected") + os.environ.setdefault("SGLANG_SIMULATE_ACC_TOKEN_MODE", "real-draft-token") + return argv + # num-steps is the chain depth; the draft-token count includes the + # verified token, hence k + 1. + argv = [ + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + str(k), + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + str(k + 1), + ] + # SGLang has no acceptance flag; it reads the simulation from the + # environment. Set here rather than asked of the caller so the anchor + # and the deployment agree without a second place to keep in step. + if al: + os.environ.setdefault("SGLANG_SIMULATE_ACC_LEN", str(al)) + os.environ.setdefault("SGLANG_SIMULATE_ACC_METHOD", "match-expected") + os.environ.setdefault("SGLANG_SIMULATE_ACC_TOKEN_MODE", "real-draft-token") + return argv + + if args.serving_backend == "atom": + # ATOM has one MTP method for every model's in-checkpoint head, and + # refuses a draft checkpoint with it; EAGLE3 requires one. + atom_method = "mtp" if str(method).lower().endswith("mtp") else str(method) + argv = ["--method", atom_method, "--num-speculative-tokens", str(k)] + if draft and atom_method != "mtp": + argv += ["--draft-model", str(draft)] + if al: + argv += ["--spec-decode-acceptance-length", str(al)] + return argv + + # vLLM resolves the generic "mtp" to the model's own head, and takes every + # other method by name. + spec: dict = {"method": str(method), "num_speculative_tokens": int(k)} + if draft: + spec["model"] = str(draft) + if al: + spec["rejection_sample_method"] = "synthetic" + spec["synthetic_acceptance_length"] = al + return ["--speculative-config", json.dumps(spec)] + + def _engine_argv(args, port: int, tp: int) -> list[str]: """The flags one engine wants for the intent every engine shares. @@ -175,6 +269,7 @@ def _engine_argv(args, port: int, tp: int) -> list[str]: argv += ["--enable-expert-parallel"] if args.enforce_eager: argv += ["--enforce-eager"] + argv += _speculative_argv(args) if args.quantization: argv += ["--quantization", args.quantization] if args.kv_cache_dtype: @@ -216,6 +311,18 @@ def _engine_argv(args, port: int, tp: int) -> list[str]: return argv + shlex.split(args.server_args or "") +def _trust_remote_code(args) -> bool: + """Whether the checkpoint's own code is trusted, however it was asked for. + + The client loads the same tokenizer the server does, so a checkpoint that + needs remote code (Kimi-K3) fails in the client too when the flag reached + the server only through --server-args. + """ + return bool(getattr(args, "trust_remote_code", False)) or ( + "--trust-remote-code" in shlex.split(getattr(args, "server_args", None) or "") + ) + + def _build_engine(args, argv: list[str], port: int, tp: int): """The engine to measure, launched the way the platform launches it. @@ -288,10 +395,59 @@ def _run_client( "--ignore-eos", "--percentile-metrics", "ttft,tpot,itl,e2el", + # Same seed-0 default as ATOM's vendored copy below, with the same + # consequence: repeated probes hit the prefix cache. + "--seed", + str(random.SystemRandom().randint(1, 2**31 - 1)), + "--save-result", + "--result-filename", + result, + ] + if _trust_remote_code(args): + cmd.append("--trust-remote-code") + elif kind == "atom": + # vLLM's legacy benchmark_serving.py, vendored into ATOM: same flags, + # same result keys. It uses relative imports, so it only runs as a + # module. Its random dataset draws exact lengths by default. + cmd = [ + sys.executable, + "-m", + "atom.benchmarks.benchmark_serving", + "--backend", + "vllm", + "--model", + args.model, + "--host", + "127.0.0.1", + "--port", + str(port), + "--endpoint", + "/v1/completions", + "--dataset-name", + "random", + "--random-input-len", + str(input_len), + "--random-output-len", + str(output_len), + "--num-prompts", + str(num_prompts), + "--max-concurrency", + str(batch), + "--ignore-eos", + "--percentile-metrics", + "ttft,tpot,itl,e2el", + # Its prompts are a pure function of the seed, which defaults to + # 0, so every repeat of a probe resends the same prompts and the + # prefix cache serves them: the prefill probe then prices a few + # hundred new tokens per request instead of the whole prompt. + "--seed", + str(random.SystemRandom().randint(1, 2**31 - 1)), "--save-result", "--result-filename", result, ] + if _trust_remote_code(args): + cmd.append("--trust-remote-code") else: cmd = [ sys.executable, @@ -325,6 +481,11 @@ def _run_client( str(num_prompts), "--max-concurrency", str(batch), + # Seeded like the other two clients (its default is 1): one seed + # draws the same leading ids at every prompt length, so each longer + # probe prompt is a cache hit on the shorter one before it. + "--seed", + str(random.SystemRandom().randint(1, 2**31 - 1)), "--output-file", result, ] @@ -356,6 +517,14 @@ def client_kind(args) -> str: """ if shutil.which("vllm"): return "vllm" + # The ATOM image carries neither vLLM nor SGLang; its load generator is + # its own vendored copy of vLLM's legacy client. Named rather than probed: + # which client an engine ships is a fact about the engine, and importing to + # check it answers a different question -- whether the module resolves in + # *this* process, which for a harvest is the submitting host and not the + # container the server runs in. + if args.serving_backend == "atom": + return "atom" if args.serving_backend == "sglang": return "sglang" raise RuntimeError( @@ -368,7 +537,7 @@ def client_kind(args) -> str: def _client_result(path: str, kind: str) -> dict: """One run's metrics, however its client chose to write them down.""" with open(path) as fh: - if kind == "vllm": + if kind in ("vllm", "atom"): return json.load(fh) # SGLang writes JSON Lines, one object per run appended to the file. lines = [line for line in fh.read().splitlines() if line.strip()] @@ -438,6 +607,19 @@ def _measure_concurrency(port: int, batch: int, args, out_dir: str) -> float: f"models. Expect TPOT to be over-predicted at high concurrency." ) return float(doc["mean_tpot_ms"]) + # The median gap can only sit above the mean one when the stream delivers + # tokens in bursts: SGLang serving Llama-3.1-8B coalesces its ~4 ms steps, + # so half the gaps read ~0 and the rest two steps, and the median read + # 8.91 ms at batch 1 and 14.58 at 16 where mean TPOT read 4.45 and 7.94, + # matching the served engine. Mean TPOT is then the cleaner of the two. + tpot = doc.get("mean_tpot_ms") + if tpot and float(itl) > 1.2 * float(tpot): + print( + f"[inferasim:Inference:Serving] median ITL {float(itl):.2f} ms exceeds " + f"mean TPOT {float(tpot):.2f} ms at concurrency {batch}: the stream " + f"is coalescing tokens, so the decode step is read from mean TPOT." + ) + return float(tpot) return float(itl) @@ -1035,6 +1217,15 @@ def run_serving_benchmark(args) -> dict: "kv_cache_dtype": args.kv_cache_dtype, "enforce_eager": args.enforce_eager, "use_aiter": os.environ.get("VLLM_ROCM_USE_AITER", "0") == "1", + # Recorded explicitly, including the "off" case. The store rebuilds + # each anchor's regime from this meta and discards the signature the + # harvest computed, so an axis that is absent here is absent from + # the index -- and a missing speculative key does not read as "off", + # it reads as "measured before this was tracked", which is refused + # for a speculating target rather than assumed compatible. + "speculative_method": getattr(args, "speculative_method", None), + "speculative_num_tokens": getattr(args, "speculative_num_tokens", None), + "speculative_acceptance_length": getattr(args, "speculative_acceptance_length", None), "server_args": args.server_args or None, # Recorded explicitly, not left to be re-derived from the flag # string by every reader. A prefill measured against a warm prefix diff --git a/infera/projection/core/projection/inference_projection/benchmark_vllm.py b/infera/projection/core/projection/inference_projection/benchmark_vllm.py index 466b72cb2..b1b3d42fe 100644 --- a/infera/projection/core/projection/inference_projection/benchmark_vllm.py +++ b/infera/projection/core/projection/inference_projection/benchmark_vllm.py @@ -214,8 +214,9 @@ def warmup_gpu_count(target_tp, target_ep=1, target_pp=1): an anchor for the degrees it did *not* measure, a four-GPU anchor was the best single choice at 6.6% -- better than the full-node eight-GPU anchor at 7.6% -- because it interpolates in both directions where an end rung has to - extrapolate. It also asks for at most half a node, so a warmup does not wait - on a full one. + extrapolate. The cap is a GPU count, not a fraction of a node, because node + sizes differ; a small fixed allocation also keeps a warmup from waiting on a + large one. """ tp = max(1, int(target_tp or 1)) return min(tp, WARMUP_GPU_CAP) @@ -1371,6 +1372,19 @@ def main(argv=None): help="draft model path for methods that need a separate " "checkpoint (unused for 'deepseek_mtp')", ) + ap.add_argument( + "--speculative-acceptance-length", + type=float, + default=None, + help="pin the mean accepted length per step instead of letting the " + "draft head earn it. An anchor is harvested under --load-format " + "dummy, and acceptance is a property of the weights: random ones " + "accept at chance, so an unpinned speculative harvest measures a " + "draft head that never hits and reports the non-speculative decode " + "rate under a speculative label. The deployments pin the same " + "committed golden acceptance for their throughput runs, so passing " + "it here measures the regime they serve in", + ) ap.add_argument( "--no-aiter", action="store_true", diff --git a/infera/projection/core/projection/inference_projection/collectives.py b/infera/projection/core/projection/inference_projection/collectives.py index 261a22e9f..1be214d62 100644 --- a/infera/projection/core/projection/inference_projection/collectives.py +++ b/infera/projection/core/projection/inference_projection/collectives.py @@ -24,6 +24,7 @@ from __future__ import annotations +import json import os from dataclasses import dataclass @@ -298,6 +299,11 @@ def __init__( ) nn = num_nodes if num_nodes else int(os.environ.get("NNODES", "1")) hw = dict(coll_config.hardware_config or {}) + # A calibrated system profile (JSON object, or a path to one), e.g. the + # one-parameter xGMI mesh fit {"node_bw": 594.5, "node_topology": "mesh"}. + prof = os.environ.get("INFERASIM_COLL_HW", "").strip() + if prof: + hw.update(json.load(open(prof)) if os.path.isfile(prof) else json.loads(prof)) for key in ("node_bw", "pod_bw"): if key in ic: hw.setdefault(key, ic[key]) @@ -319,6 +325,15 @@ def __init__( # on-node dispatch at pod bandwidth. ``1`` makes the test read # ``ep <= node_size``, the question actually being asked; where ``ep`` # genuinely exceeds a node, the crossing is still charged. + self._args_intra = get_default_args( + num_nodes=1, + gpus_per_node=gpn, + tp=self.tp, + pp=self.pp, + ep=self.ep, + cp=self.cp, + hardware_config=hw, + ) self._a2a_args = get_default_args( num_nodes=nn, gpus_per_node=gpn, @@ -487,12 +502,43 @@ def _floor_small_msg_overhead(self, us: float, msg: int, gpus: int, *, is_a2a: b # -- Pipeline P2P ---------------------------------------------------------- def pp_p2p_ms(self, batch: int, tokens: int) -> float: - """Activation send/recv across (pp-1) stage boundaries per forward.""" + """Activation send/recv across (pp-1) stage boundaries per forward, and + the last stage's sampled tokens back to the first, which cannot + schedule that micro-batch again until they land.""" if self.pp <= 1 or not self.cc.include_pp_p2p: return 0.0 msg = max(1, batch * tokens * self.hidden * 2 // self.cp) - us = cm.sendrecv(self._args, msg) - return (self.pp - 1) * (us / 1000.0) + ov = float(getattr(self.cc, "pp_stage_overhead_us", 0.0) or 0.0) + tok = max(1, batch * 8) + nodes = self.pp_nodes() + if nodes <= 1: + fwd = (self.pp - 1) * (cm.sendrecv(self._args, msg) + ov) + back = cm.sendrecv(self._args, tok) + ov + return (fwd + back) / 1000.0 + # Stages are placed contiguously, so a pipeline over n nodes has n-1 + # forward boundaries on the NIC, and the token return from the last + # stage to the first crosses it too; the rest stay on xGMI. + cross = nodes - 1 + intra = cm.sendrecv(self._args_intra, msg) + ov + fwd = (self.pp - 1 - cross) * intra + cross * (self._nic_sendrecv_us(msg) + ov) + back = self._nic_sendrecv_us(tok) + ov + return (fwd + back) / 1000.0 + + def pp_nodes(self) -> int: + """Nodes the pipeline spans: ``INFERASIM_PP_NODES`` when the placement + is given, else as many as ``tp * pp`` GPUs fill.""" + env = os.getenv("INFERASIM_PP_NODES") + if env: + return max(1, min(self.pp, int(env))) + gpn = max(1, int(self._args.node_size)) + return max(1, min(self.pp, -(-self.tp * self.pp // gpn))) + + def _nic_sendrecv_us(self, msg: float) -> float: + """One point-to-point transfer over a single NIC (us).""" + a = self._args + raw_pod = getattr(a, "_raw_pod_bw", a.pod_bw / a.bw_eff) + bw = raw_pod * getattr(a, "p2p_bw_eff", 0.80) + return (msg / bw) * 1.0e-3 + a.pod_lat + a.kernel_launch_latency # -- KV-cache transfer (disaggregation) ------------------------------------ diff --git a/infera/projection/core/projection/inference_projection/des.py b/infera/projection/core/projection/inference_projection/des.py index 37d3ac055..1c336615d 100644 --- a/infera/projection/core/projection/inference_projection/des.py +++ b/infera/projection/core/projection/inference_projection/des.py @@ -60,11 +60,11 @@ from __future__ import annotations import csv +import heapq import json import math import os import random -from collections import OrderedDict from dataclasses import dataclass, field from infera.projection.core.projection.training_config import InferenceConfig @@ -128,11 +128,60 @@ class _Req: first_token_ms: float = -1.0 finish_ms: float = -1.0 itls: list[float] = field(default_factory=list) + # When the request last emitted a token; the client's inter-token gap runs + # from here, including steps it sat out (an exclusive prefill batch). + last_emit_ms: float = -1.0 + # Whether the credit below applies. It is a property of *who the sharers + # are*, not of the engine: see ``reserved_kv``. + shared_prefix_credit: bool = True + # Closed loop only: the client's idle time before issuing this request, + # from the trace. ``None`` falls back to the run's single think time. + think_ms: float | None = None + # Host-side TTFT latency; see ``_set_host_ms``. + host_ms: float = 0.0 + # Session replay (closed loop): which client session this request belongs + # to, in the order a shared sampler hands sessions out, and the requests + # it waits on -- the previous turn of its own stream (``dep``) and any + # subagents it joins. ``sess < 0`` is the plain one-request-per-slot loop. + sess: int = -1 + dep: int = -1 + joins: tuple = () @property def reserved_kv(self) -> int: - # Full-ISL reservation: the whole sequence is reserved up front. - return self.prompt_len + self.output_len + # Full-ISL reservation: the whole sequence is reserved up front, less + # the leading blocks that are already resident. A prefix-cache hit + # means those blocks are *in* the pool -- some earlier request put them + # there and this one attends to the same physical KV -- so charging + # them again per request double-counts the one copy the engine keeps. + # On an agentic replay that is not a small correction: at the 96-98% + # block reuse these corpora carry it inflates a request's footprint + # 25-50x, so a measured 7.67M-token pool looks full at ~19 concurrent + # and the replay queues from there, while the MI355X ladder reports its + # pool at 0.1-0.8% utilization at every concurrency it was run at and + # never binds at all. + # + # The credit is not refcounted: when the request that warmed a block + # retires, its reservation is released even though the block stays + # resident for reuse. Both directions are approximations of a + # refcounted pool, and undercharging shared blocks lands far nearer the + # measured occupancy than charging every sharer in full. A cold run is + # untouched -- ``cached_prefix`` is zero without a cache to hit, which + # is every fixed-sequence workpoint. + # + # The credit accounts for one physical copy only where the sharers are + # resident *at the same time*. On the agentic corpora they are not: + # reuse there is a conversation hitting its own previous turn, and a + # closed loop gives each client one turn in flight, so the requests + # running together are always different conversations whose contexts + # are disjoint past a short shared head. Credited anyway, C concurrent + # 200k-token histories cost a few thousand tokens each and the pool + # never binds at any concurrency these ladders reach -- which is why + # reuse holds at its ceiling in the replay and degrades on every one + # of these systems as C approaches what the pool holds. + if not self.shared_prefix_credit: + return self.prompt_len + self.output_len + return self.prompt_len - self.cached_prefix + self.output_len @property def in_prefill(self) -> bool: @@ -171,6 +220,10 @@ class DESResult: # optional raw per-request latency samples (only when return_samples=True); # used to pool distributions across instances in the multi-instance driver. samples: dict[str, list[float]] | None = None + # Prompt plus generated tokens/s over the same requests and span as + # system_throughput_tps -- InferenceX's total_tput_tps, which counts every + # prompt token, cached or not. + system_total_throughput_tps: float = 0.0 class _CostKernel: @@ -181,6 +234,43 @@ def __init__(self, projector: InferencePerformanceProjector, q_len: int): self._q = q_len self._decode: dict[tuple, float] = {} self._mixed: dict[tuple, float] = {} + # The simulated scheduler is one of this many that step in lockstep + # (SGLang attention-DP ranks meet at every MoE layer), so a step costs + # what the combined batch costs. + self._sync = max(1, int(os.getenv("INFERASIM_DES_LOCKSTEP", "1") or 1)) + # Opt-in: under attention-DP the requests one prefill step carries + # belong to different ranks, which attend over their own share in + # parallel; only the MoE sees the whole step's tokens. + self._dp = 1 + if os.getenv("INFERASIM_DES_DPA_SPREAD") == "1": + from .kv_cache import attention_dp_size + + self._dp = max(1, int(attention_dp_size(projector.cfg))) + # Pipeline stages hand this many micro-batches round. Once two or more + # are in flight no stage waits on its neighbour, so each pays its + # scheduler host time (pp_host_us + pp_host_per_req_us per request) in + # series with its forward. + self._pp_n = max(1, int(os.getenv("INFERASIM_PP_INFLIGHT", "1") or 1)) + self._pp_stages = 1 + self._pp_host = (0.0, 0.0) + if self._pp_n > 1: + try: + cc = projector.cfg.collective_config + self._pp_stages = max( + 1, projector.cfg.model_parallel_config.pipeline_model_parallel_size + ) + self._pp_host = ( + float(getattr(cc, "pp_host_us", 0.0) or 0.0) / 1000.0, + float(getattr(cc, "pp_host_per_req_us", 0.0) or 0.0) / 1000.0, + ) + except Exception: + self._pp_stages = 1 + + def _pp_sync(self, v: float, batch: int) -> float: + if self._pp_n <= 1: + return v + h0, h1 = self._pp_host + return v + self._pp_stages * (h0 + h1 * max(1, batch)) @staticmethod def _bucket(ctx: int) -> int: @@ -194,17 +284,41 @@ def decode_step_ms(self, batch: int, ctx: int) -> float: key = (batch, self._bucket(ctx)) v = self._decode.get(key) if v is None: - v = self._p.decode_step_latency_ms(batch, key[1], self._q) + v = self._pp_sync( + self._p.decode_step_latency_ms(batch * self._sync, key[1], self._q), batch + ) self._decode[key] = v return v def mixed_step_ms( - self, num_decode: int, prefill_tokens: int, ctx: int, prefill_kv: int + self, + num_decode: int, + prefill_tokens: int, + ctx: int, + prefill_kv: int, + prefill_reqs: int = 1, ) -> float: - key = (num_decode, self._tok(prefill_tokens), self._bucket(ctx), self._bucket(prefill_kv)) + spread = min(self._dp, max(1, int(prefill_reqs))) + key = ( + num_decode, + self._tok(prefill_tokens), + self._bucket(ctx), + self._bucket(prefill_kv), + spread, + ) v = self._mixed.get(key) if v is None: - v = self._p.mixed_step_latency_ms(num_decode, key[1], key[2], key[3], self._q) + # One chunk per rank: under attention-DP each rank attends over its + # own chunk with every head, and only the MoE sees all of them. + kw = {"prefill_batch": self._sync} if self._sync > 1 else {} + tok = key[1] + if spread > 1: + kw = {"prefill_batch": spread} + tok = max(1, key[1] // spread) + v = self._p.mixed_step_latency_ms( + num_decode * self._sync, tok, key[2], key[3], self._q, **kw + ) + v = self._pp_sync(v, key[0] + 1) self._mixed[key] = v if os.getenv("INFERASIM_DEBUG_DES_STEPS"): print( @@ -214,6 +328,124 @@ def mixed_step_ms( return v +def _set_host_ms(reqs: list[_Req], req_cfg) -> None: + """Per-request host latency: ``request_overhead_ms``, plus + ``tokenize_overhead_us`` per prompt token (the prompt arrives as text and is + tokenized server-side after the TTFT clock starts), plus + ``uncached_prompt_latency_us`` for each prompt token the cache did not serve, + up to ``uncached_prompt_latency_max_tokens`` of them. + + Charged to TTFT and end-to-end latency only: the scheduler, and a closed + loop's issue times, are unchanged. + """ + fixed = max(0.0, float(getattr(req_cfg, "request_overhead_ms", 0.0) or 0.0)) + tok = max(0.0, float(getattr(req_cfg, "tokenize_overhead_us", 0.0) or 0.0)) / 1000.0 + unc = max(0.0, float(getattr(req_cfg, "uncached_prompt_latency_us", 0.0) or 0.0)) / 1000.0 + cap = int(getattr(req_cfg, "uncached_prompt_latency_max_tokens", 0) or 0) + for r in reqs: + u = max(0, r.prompt_len - r.cached_prefix) + r.host_ms = fixed + tok * r.prompt_len + unc * (min(u, cap) if cap > 0 else u) + + +def _resident_cap( + reqs: list[_Req], kv_cache_tokens: int, max_running: int, enabled: bool = False +) -> int: + """How many of these requests the KV pool holds at once. + + Length-biased, not the arithmetic mean: a request occupies the pool for a + time proportional to its own length, so the set resident at any instant is + sampled in proportion to length and the plain mean overstates how many + fit. Zero when there is no pool to divide, which leaves the caller's + ordering policy unconstrained. + + Off unless ``enabled``, because measuring the reordering window against + this is a modelling choice that did not pay for itself on the AgentX + corpus. It is the right account of what a scheduler can reorder, and it + does improve ITL ordering (8 of 11 model-engine pairs at 0.90 or better, + against 9), but it costs more than that in throughput (8 pairs to 7) -- + reuse at the low rungs falls further than the hardware's did. Kept + available because the effect it describes is real and the band it applies + to is narrow; a caller that wants it has to ask. + """ + if not enabled: + return 0 + if kv_cache_tokens <= 0 or not reqs: + return 0 + ctx = [max(1, r.prompt_len + r.output_len) for r in reqs] + biased = sum(c * c for c in ctx) / max(1, sum(ctx)) + fits = int(kv_cache_tokens / max(1.0, biased)) + return max(1, min(fits, max_running if max_running > 0 else fits)) + + +def _free_cache_blocks( + kv_cache_tokens: int, live_tokens: float, block_size: int, cap_blocks: int +) -> int: + """Block-cache capacity left over once the running set has its KV. + + There is one pool. ``_Req.reserved_kv`` already declines to charge a + request for the leading blocks it hit on, on the grounds that an earlier + request is holding them -- which is right, and which only balances if + whatever *is* holding them is charged instead. The block cache is that + holder, and it was being handed the whole pool at the same time as the + running set, so both structures booked the same memory and neither ever + saw it run out. + + The error is invisible wherever the pool is large against the working set + and total wherever it is not. On the AgentX ladders DeepSeek-V4 has room + for 44-132 conversations and reuse holds near its ceiling at every + concurrency run; GLM-5.2 on MI325X has room for 5.5, and the hardware's + reuse falls 0.79 -> 0.30 -> 0.12 across C=4,5,6 as the live contexts crowd + the cached prefixes out, taking TTFT from 3s to 198s. Double-booked, the + replay kept reporting a 0.96 hit rate and a flat TTFT through the whole + collapse. + + Residency does not depend on reuse -- a hit skips prefill compute, it does + not free blocks -- so the live figure measured with the cache at one + capacity is still the live figure at another, and a single corrective pass + lands on the answer rather than approaching it. + + Never returns zero: ``_BlockStore`` reads a zero capacity as unbounded, and + a pool with no room left is the opposite of that. + """ + if kv_cache_tokens <= 0 or block_size <= 0: + return cap_blocks + free = max(0, int(kv_cache_tokens) - int(max(0.0, live_tokens))) + blocks = free // int(block_size) + if cap_blocks > 0: + blocks = min(blocks, cap_blocks) + return max(1, int(blocks)) + + +def _scored_sample(done: list[_Req], warmup_frac: float, warmup_requests: int) -> list[_Req]: + """The requests whose latencies get reported. + + ``warmup_requests`` drops the opening transient by *issue* order, which is + the only order that removes it. A closed loop starts every one of its C + clients at once, so the first C requests queue against each other and wait + far longer than anything that follows -- a mean of 17 s against a run + median of 0 on glm5.2 at C=8, and 16% of all the queue wait in the run. + Dropping a fraction of the *earliest completions* instead, as + ``warmup_frac`` does, keeps every one of them: a request that waited 40 s + for a slot is among the last to finish, not the first. + + The harness has the same transient and excludes it the same way -- it + advances each lane by ``AIPERF_WARMUP_REQUESTS_PER_LANE`` requests, waits + for them to drain, and only then starts profiling -- so matching it is a + matter of dropping the same requests rather than a comparable number of + them. At least half the run is always kept, so a short trace still reports + over something. + """ + keep = done + if warmup_requests > 0: + cut = min(int(warmup_requests), len(done) // 2) + keep = [r for r in done if r.idx >= cut] + if len(keep) >= 8: + return keep + keep = done + drop = int(len(keep) * max(0.0, min(0.9, warmup_frac))) + return keep[drop:] if len(keep) - drop >= 8 else keep + + def _generate_arrivals( n: int, rate_per_s: float, model: str, rng: random.Random, burstiness: float = 1.0 ) -> list[float]: @@ -324,6 +556,7 @@ def simulate_once( arrival_model: str = "poisson", num_requests: int = 400, warmup_frac: float = 0.1, + warmup_requests: int = 0, seed: int = 0, arrivals: list[float] | None = None, burstiness: float = 1.0, @@ -334,8 +567,12 @@ def simulate_once( prebuilt: list[_Req] | None = None, return_samples: bool = False, closed_loop_clients: int = 0, + closed_loop_think_ms: float = 0.0, + whole_context_residency: bool = False, prefill_exclusive: bool = False, new_seqs_per_step: int = 0, + duration_ms: float = 0.0, + closed_loop_idle_cap_ms: float = 0.0, ) -> DESResult: """Run one single-engine DES at a fixed offered load. @@ -351,6 +588,25 @@ def simulate_once( which is what the harness does and what makes its *mean* TTFT carry an opening-burst transient; ``warmup_frac`` decides whether that transient is scored, so match it to whatever the harness reports over. + + ``duration_ms`` stops the run on the clock instead of on a request count, + which is how a fixed-concurrency harness is actually bounded: it runs for + a set wall time and reports whatever completed. The distinction is not + cosmetic under a closed loop. A request budget split across ``C`` lanes + fixes how far *each lane* walks into its conversation -- 400 turns deep at + C=1 against 40 at C=16 for the same budget -- and an agentic turn's prompt + grows with its position in the conversation, so a budget-bounded replay + offers systematically different prompt lengths at each concurrency than + the run it is being compared against. A clock-bounded one lets the lane + depth fall out of how fast the engine actually is, which is the same thing + that decided it on the hardware. + + A request's own ``think_ms`` (from the trace) overrides + ``closed_loop_think_ms``. ``closed_loop_idle_cap_ms > 0`` bounds how long + the engine can sit with nothing in flight: pending client timers are pulled + forward so the next request lands at most that long after the engine went + idle. That is the replay harness's system-idle guard, and it is what makes + a single client's idle time much shorter than the trace's own gaps. """ req = inference_config.request_config input_len = max(1, req.input_seq_len) @@ -367,6 +623,7 @@ def simulate_once( rng = random.Random(seed) clients = max(0, int(closed_loop_clients)) closed_loop = clients > 0 + think_ms = max(0.0, float(closed_loop_think_ms)) # ---- workload (arrivals + per-request lengths) ---- if closed_loop: @@ -391,8 +648,9 @@ def simulate_once( pending = _build_workload( n_total, [0.0] * n_total, input_len, output_len, range_ratio, rng ) + session_mode = any(r.sess >= 0 for r in pending) for i, r in enumerate(pending): - r.arrival_ms = 0.0 if i < clients else math.inf + r.arrival_ms = 0.0 if (i < clients and not session_mode) else math.inf elif prebuilt is not None: # Caller supplied a fully-formed request list (multi-instance router: # arrivals, lengths, prefix ids + seeded ``num_computed`` for hits). @@ -425,7 +683,24 @@ def simulate_once( if cached > 0: r.cached_prefix = cached r.num_computed = cached - pending.sort(key=lambda r: (r.arrival_ms, r.idx)) + if whole_context_residency: + for r in pending: + r.shared_prefix_credit = False + _set_host_ms(pending, req) + session_mode = closed_loop and any(r.sess >= 0 for r in pending) + if session_mode: + pending.sort(key=lambda r: r.idx) + # Concurrency here counts sessions, not requests: a session's subagents + # run at once, so C=1 still puts several requests in flight. The server + # admits them up to its own sequence limit (KV admission still applies); + # capping the batch at C instead queued a subagent behind its sibling's + # whole decode, 42 s at p99 on DeepSeek-V4 at C=1. + max_running = max(max_running, _SERVER_MAX_RUNNING) + else: + pending.sort(key=lambda r: (r.arrival_ms, r.idx)) + cap_seqs = int(getattr(req, "max_num_seqs", 0) or 0) + if cap_seqs > 0: + max_running = min(max_running, cap_seqs) n = len(pending) kernel = _CostKernel(projector, q_len) @@ -434,6 +709,63 @@ def simulate_once( # index below it carries a finite arrival, which is what lets the ingest # pointer above stay a simple in-order scan. next_unissued = clients if closed_loop else n + # Closed loop: issued-but-not-arrived requests, keyed on arrival. Per-request + # think times make arrivals non-monotone in issue order, so the in-order + # ingest scan cannot be used for them. + idle_cap_ms = max(0.0, float(closed_loop_idle_cap_ms or 0.0)) + arrivals_heap: list[tuple[float, int, _Req]] = [] + # Opt-in serial front end (us per prompt token): one process tokenizes + # every prompt in send order before the scheduler can see it. TTFT is still + # timed from the send, so a prompt queued behind others pays their tokenize. + fe_us = float(os.environ.get("INFERASIM_DES_SERIAL_FRONTEND_US", "0") or 0.0) + fe_heap: list[tuple[float, int, _Req]] = [] + fe_free = 0.0 + horizon = duration_ms if duration_ms and duration_ms > 0 else math.inf + if session_mode: + # A request is released when everything it waits on has finished, plus + # its own recorded gap; a session starts when a client slot frees, + # which is when the session ``clients`` places ahead of it drains. + # That is the shared-sampler client the AgentX harness runs: whichever + # lane frees first takes the next trace, and a lane holds its slot + # until its whole tree -- root and every subagent -- has drained. + by_idx = {r.idx: r for r in pending} + need: dict[int, int] = {} + ready_at: dict[int, float] = {} + waiters: dict[int, list] = {} + sess_reqs: dict[int, list] = {} + for r in pending: + deps = [d for d in ((r.dep,) + tuple(r.joins)) if d >= 0 and d in by_idx] + need[r.idx] = len(deps) + ready_at[r.idx] = 0.0 + for d in deps: + waiters.setdefault(d, []).append(r) + sess_reqs.setdefault(r.sess, []).append(r) + sess_left = {k: len(v) for k, v in sess_reqs.items()} + sess_order = sorted(sess_reqs) + next_sess = [0] + + def _release(r: _Req, t: float) -> None: + r.arrival_ms = t + (r.think_ms or 0.0) + if r.arrival_ms < horizon: + heapq.heappush(arrivals_heap, (r.arrival_ms, r.idx, r)) + + def _start_session(t: float) -> None: + if next_sess[0] >= len(sess_order): + return + k = sess_order[next_sess[0]] + next_sess[0] += 1 + for r in sess_reqs[k]: + if need[r.idx] == 0: + _release(r, t) + + for _ in range(clients): + _start_session(0.0) + next_arrival = n + next_unissued = n + elif closed_loop: + arrivals_heap = [(pending[i].arrival_ms, i, pending[i]) for i in range(clients)] + heapq.heapify(arrivals_heap) + next_arrival = n waiting: list[_Req] = [] running: list[_Req] = [] done: list[_Req] = [] @@ -448,6 +780,7 @@ def simulate_once( pk_qtokens = 0 pk_prefill_steps = 0 pk_kv_peak = 0 + pk_kv_sum = 0 # Exact worst-case iteration bound (batch=1: every token its own step). Both # the per-request chunk cap and the shared token budget can split a prefill, @@ -463,15 +796,33 @@ def simulate_once( max_steps = total_work + n + 16 steps = 0 - while len(done) < n and steps < max_steps: + # A session replay stops issuing at the horizon and lets what is already + # in flight finish, which is what the harness reports over. + while ( + len(done) < n + and steps < max_steps + and (now < horizon or (session_mode and (running or waiting or fe_heap))) + ): steps += 1 # 1) Ingest arrivals due by ``now`` into the FCFS waiting queue. while next_arrival < n and pending[next_arrival].arrival_ms <= now + 1e-9: waiting.append(pending[next_arrival]) next_arrival += 1 + while arrivals_heap and arrivals_heap[0][0] <= now + 1e-9: + q_ = heapq.heappop(arrivals_heap)[2] + if fe_us > 0: + fe_free = max(fe_free, q_.arrival_ms) + fe_us * 1e-3 * q_.prompt_len + heapq.heappush(fe_heap, (fe_free, q_.idx, q_)) + else: + waiting.append(q_) + while fe_heap and fe_heap[0][0] <= now + 1e-9: + waiting.append(heapq.heappop(fe_heap)[2]) # 2) Nothing resident and nothing waiting → jump to the next arrival. if not running and not waiting: + if arrivals_heap or fe_heap: + now = min(h[0][0] for h in (arrivals_heap, fe_heap) if h) + continue if next_arrival < n and math.isfinite(pending[next_arrival].arrival_ms): now = pending[next_arrival].arrival_ms continue @@ -599,6 +950,9 @@ def _schedule_tokens(r: _Req, budget: float) -> int: if not scheduled: # Budget/KV starved this step with nothing runnable; advance to the # next arrival if possible, else we are stuck (bound will trip). + if arrivals_heap or fe_heap: + now = min(h[0][0] for h in (arrivals_heap, fe_heap) if h) + continue if next_arrival < n: now = pending[next_arrival].arrival_ms continue @@ -612,7 +966,7 @@ def _schedule_tokens(r: _Req, budget: float) -> int: if prefill_q > 0: prefill_kv = int(sum(kv + q for _, q, kv in pref) / len(pref)) decode_ctx = int(sum(kv for _, _, kv in dec) / len(dec)) if dec else input_len - step_dt = kernel.mixed_step_ms(num_decode, prefill_q, decode_ctx, prefill_kv) + step_dt = kernel.mixed_step_ms(num_decode, prefill_q, decode_ctx, prefill_kv, len(pref)) else: decode_ctx = int(sum(kv for _, _, kv in dec) / len(dec)) if dec else input_len step_dt = kernel.decode_step_ms(num_decode, decode_ctx) @@ -629,7 +983,7 @@ def _schedule_tokens(r: _Req, budget: float) -> int: r.prefill_done = True r.first_token_ms = now r.generated = 1 # last prefill chunk emits token 1 - r.itls.append(step_dt) + r.last_emit_ms = now if r.generated >= r.output_len: r.status = "FINISHED" r.finish_ms = now @@ -639,7 +993,9 @@ def _schedule_tokens(r: _Req, budget: float) -> int: continue acc = _sample_accepted(rng, spec_k, accept, cap) r.generated += acc - per_tok = step_dt / acc + gap = (now - r.last_emit_ms) if r.last_emit_ms >= 0 else step_dt + r.last_emit_ms = now + per_tok = gap / acc r.itls.extend([per_tok] * acc) if r.generated >= r.output_len: r.status = "FINISHED" @@ -653,12 +1009,30 @@ def _schedule_tokens(r: _Req, budget: float) -> int: if r.status == "FINISHED": kv_used -= r.reserved_kv done.append(r) - if closed_loop and next_unissued < n: - pending[next_unissued].arrival_ms = now + if session_mode: + for w in waiters.get(r.idx, ()): + need[w.idx] -= 1 + ready_at[w.idx] = max(ready_at[w.idx], now) + if need[w.idx] == 0: + _release(w, ready_at[w.idx]) + sess_left[r.sess] -= 1 + if sess_left[r.sess] == 0: + _start_session(now) + elif closed_loop and next_unissued < n: + nxt = pending[next_unissued] + nxt.arrival_ms = now + (think_ms if nxt.think_ms is None else nxt.think_ms) + heapq.heappush(arrivals_heap, (nxt.arrival_ms, next_unissued, nxt)) next_unissued += 1 else: still.append(r) running = still + if idle_cap_ms > 0 and not running and not waiting and not fe_heap and arrivals_heap: + lead = arrivals_heap[0][0] - (now + idle_cap_ms) + if lead > 0: + # A uniform shift keeps the heap ordered. + arrivals_heap = [(a - lead, i, q) for a, i, q in arrivals_heap] + for _a, _i, q in arrivals_heap: + q.arrival_ms -= lead # 8) Bookkeeping: waiting depth + packing stats (+ optional records). backlog.append((now, len(waiting))) @@ -671,6 +1045,7 @@ def _schedule_tokens(r: _Req, budget: float) -> int: pk_qtokens += prefill_q + num_decode * q_len pk_prefill_steps += 1 if pref else 0 pk_kv_peak = max(pk_kv_peak, kv_used) + pk_kv_sum += kv_used if record_steps: step_records.append( { @@ -694,10 +1069,15 @@ def _schedule_tokens(r: _Req, budget: float) -> int: } ) - # ---- aggregate latency metrics (drop warmup by completion order) ---- + # ---- aggregate latency metrics (drop the opening transient) ---- done.sort(key=lambda r: r.finish_ms) - drop = int(len(done) * max(0.0, min(0.9, warmup_frac))) - sample = done[drop:] if len(done) - drop >= 8 else done + if session_mode: + # The session trace starts at the profiling boundary (the harness's + # warmup is already applied to it), so every request it completed is + # scored. + sample = [r for r in done if r.arrival_ms < horizon] + else: + sample = _scored_sample(done, warmup_frac, warmup_requests) # TTFT is measured from *admission* (server start), not arrival, so the # client-side wait for a concurrency slot is excluded -- matching the vLLM / @@ -705,23 +1085,13 @@ def _schedule_tokens(r: _Req, budget: float) -> int: # request acquires its semaphore. The queue wait and the arrival-relative # TTFT are reported separately (queue_wait, ttft_arrival) so the deployment # view is not lost. - # Host prompt-tokenization cost (per prompt token). The prompt is sent as - # text and tokenized server-side after the TTFT clock starts, so it lands in - # TTFT. Latency-only -- the server scheduler/makespan above are unchanged. - tok_ms_pt = max(0.0, req.tokenize_overhead_us) / 1000.0 def _admit(r): return r.admit_ms if r.admit_ms >= 0 else r.arrival_ms - ttft = [ - (r.first_token_ms - _admit(r)) + tok_ms_pt * r.prompt_len - for r in sample - if r.first_token_ms >= 0 - ] + ttft = [(r.first_token_ms - _admit(r)) + r.host_ms for r in sample if r.first_token_ms >= 0] ttft_arrival = [ - (r.first_token_ms - r.arrival_ms) + tok_ms_pt * r.prompt_len - for r in sample - if r.first_token_ms >= 0 + (r.first_token_ms - r.arrival_ms) + r.host_ms for r in sample if r.first_token_ms >= 0 ] queue_wait = [_admit(r) - r.arrival_ms for r in sample if r.first_token_ms >= 0] # Per-output-token detokenization + streaming (client-side host cost). @@ -730,7 +1100,7 @@ def _admit(r): # serving harness, which measures ITL client-side. detok_ms = max(0.0, req.detokenize_overhead_us) / 1000.0 e2e = [ - (r.finish_ms - r.arrival_ms) + tok_ms_pt * r.prompt_len + detok_ms * r.generated + (r.finish_ms - r.arrival_ms) + r.host_ms + detok_ms * r.generated for r in sample if r.finish_ms >= 0 ] @@ -745,7 +1115,16 @@ def _admit(r): makespan = max((r.finish_ms for r in done), default=0.0) total_out = sum(r.generated for r in done) - if closed_loop and len(sample) >= 2: + if session_mode and sample: + # Rated over the profiling window plus the drain, the way the harness + # divides by its measured duration. + span_ms = max(horizon if math.isfinite(horizon) else 0.0, max(r.finish_ms for r in sample)) + out_sample = sum(r.generated for r in sample) + achieved_rate = (len(sample) * 1000.0 / span_ms) if span_ms > 0 else 0.0 + sys_tps = (out_sample * 1000.0 / span_ms) if span_ms > 0 else 0.0 + tot_tokens = out_sample + sum(r.prompt_len for r in sample) + tot_tps = (tot_tokens * 1000.0 / span_ms) if span_ms > 0 else 0.0 + elif closed_loop and len(sample) >= 2: # The tail of a closed run drains: once the last request is issued the # population falls below ``clients`` and throughput with it. Rating the # whole run would charge that drain against the engine, so the rate is @@ -754,9 +1133,13 @@ def _admit(r): out_sample = sum(r.generated for r in sample) achieved_rate = (len(sample) * 1000.0 / span_ms) if span_ms > 0 else 0.0 sys_tps = (out_sample * 1000.0 / span_ms) if span_ms > 0 else 0.0 + tot_tokens = out_sample + sum(r.prompt_len for r in sample) + tot_tps = (tot_tokens * 1000.0 / span_ms) if span_ms > 0 else 0.0 else: achieved_rate = (len(done) * 1000.0 / makespan) if makespan > 0 else 0.0 sys_tps = (total_out * 1000.0 / makespan) if makespan > 0 else 0.0 + tot_tokens = total_out + sum(r.prompt_len for r in done) + tot_tps = (tot_tokens * 1000.0 / makespan) if makespan > 0 else 0.0 utilization = (busy_ms / makespan) if makespan > 0 else 0.0 # Saturation: the waiting queue diverges (grows ~linearly) while arrivals are @@ -799,6 +1182,13 @@ def dist(xs: list[float]) -> dict[str, float]: "prefill_step_fraction": (pk_prefill_steps / pk_steps) if pk_steps else 0.0, "kv_peak_tokens": float(pk_kv_peak), "kv_utilization": (pk_kv_peak / kv_pool) if kv_pool > 0 else 0.0, + # Occupancy averaged over steps, not the high-water mark. The peak is + # set by the opening transient: the cache is cold, nothing is a hit, + # and the first batch reserves whole prompts, so the peak pins to the + # pool on any run whose cold working set exceeds it and says nothing + # about the run that follows. What the cache has to live alongside is + # the steady occupancy. + "kv_mean_tokens": (pk_kv_sum / pk_steps) if pk_steps else 0.0, "closed_loop_clients": float(clients), } @@ -821,6 +1211,7 @@ def dist(xs: list[float]) -> dict[str, float]: num_requests=len(done), makespan_ms=makespan, system_throughput_tps=sys_tps, + system_total_throughput_tps=tot_tps, saturated=saturated, ttft=dist(ttft), ttft_arrival=dist(ttft_arrival), @@ -834,6 +1225,10 @@ def dist(xs: list[float]) -> dict[str, float]: ) +# Engine-side cap on concurrently running sequences (vLLM's default +# ``max_num_seqs``), for workloads whose client concurrency does not bound it. +_SERVER_MAX_RUNNING = 256 + _ROUTING_POLICIES = ("round_robin", "random", "prefix_aware", "kv") # Default KV block size (tokens per paged block). 512 matches the Mooncake @@ -842,45 +1237,109 @@ def dist(xs: list[float]) -> dict[str, float]: class _BlockCache: - """Per-instance paged-KV block store (content-addressed, LRU-evicted). + """Per-instance paged-KV block store (content-addressed, leaf-first eviction). Models engine-style automatic prefix caching: a prompt is an ordered sequence of block-hash ids; a **hit** is the longest *contiguous prefix* of that sequence already resident (prefix caching only reuses a leading run of - matching blocks). Capacity is finite in blocks; least-recently-used blocks - are evicted under pressure. This is the reuse store the KV-aware router scores - routes against, and what a single engine hits across a sequential stream. + matching blocks). This is the reuse store the KV-aware router scores routes + against, and what a single engine hits across a sequential stream. + + Eviction is leaf-first, least-recently-used among the leaves, which is what + a radix prefix cache does -- a block cannot be dropped while a resident + block continues from it, so the shared head of a corpus survives pressure + and only the divergent tails are reclaimed. + + Modelling it as a flat LRU instead was not a smaller approximation of that, + it was the LRU pathology. Blocks are touched in prefix order, so a corpus + whose working set exceeds the pool is a cyclic sweep, and a cyclic sweep + evicts every block exactly before it is reused. Replaying the AgentX c256 + trace against DeepSeek-V4's real 231,336-block pool, flat LRU returned a + 0.0% hit rate over 10.5M evictions -- not a degraded rate, every single + access -- where the same trace and pool under leaf-first eviction return + 90.6%, against the corpus's own published reuse of ~92%. Downstream that + was throughput at 0.08x of measured on MI355X/SGLang at C=256, and a + throughput rank correlation of 0.11 on a curve that otherwise orders at + 0.93. + + Blocks are content-addressed over their prefix, so a block id implies its + predecessor and the parent recorded on first insert is stable. """ def __init__(self, capacity_blocks: int = 0) -> None: self.capacity = int(capacity_blocks or 0) # 0 = unbounded - self._lru: OrderedDict[int, None] = OrderedDict() + self._parent: dict[int, int | None] = {} + self._children: dict[int, int] = {} # resident children; 0 => a leaf + self._access: dict[int, int] = {} + # Candidate leaves by access time. Entries go stale when a block is + # touched again or stops being a leaf, so they are filtered on pop + # rather than removed in place. + self._leaves: list[tuple[int, int]] = [] + self._tick = 0 self.evictions = 0 def prefix_match(self, blocks: list[int]) -> int: """Number of leading blocks already resident (contiguous from the head).""" m = 0 for b in blocks: - if b in self._lru: + if b in self._children: m += 1 else: break return m def insert(self, blocks: list[int]) -> None: - """Warm a request's blocks (mark MRU); evict LRU beyond capacity.""" + """Warm a request's blocks, then reclaim leaves beyond capacity.""" + prev: int | None = None for b in blocks: - if b in self._lru: - self._lru.move_to_end(b) + self._tick += 1 + if b in self._children: + self._access[b] = self._tick + # Only leaves are eviction candidates, so only a leaf's touch + # needs re-filing; an interior block is not reclaimable anyway. + if self._children[b] == 0: + heapq.heappush(self._leaves, (self._tick, b)) else: - self._lru[b] = None - if self.capacity > 0: - while len(self._lru) > self.capacity: - self._lru.popitem(last=False) - self.evictions += 1 + self._parent[b] = prev + self._children[b] = 0 + self._access[b] = self._tick + if prev is not None and prev in self._children: + self._children[prev] += 1 + heapq.heappush(self._leaves, (self._tick, b)) + prev = b + self._evict() + + def _evict(self) -> None: + if self.capacity <= 0: + return + while len(self._children) > self.capacity: + victim = None + while self._leaves: + t, b = heapq.heappop(self._leaves) + if b not in self._children: + continue # already reclaimed + if self._access[b] != t: + continue # touched since; a newer entry stands + if self._children[b]: + continue # no longer a leaf + victim = b + break + if victim is None: + # Every resident block is interior. Nothing is reclaimable + # without orphaning a continuation, which is the one thing a + # radix cache will not do. + return + p = self._parent.pop(victim) + del self._children[victim] + del self._access[victim] + self.evictions += 1 + if p is not None and p in self._children: + self._children[p] -= 1 + if self._children[p] == 0: + heapq.heappush(self._leaves, (self._access[p], p)) def __len__(self) -> int: # pragma: no cover - trivial - return len(self._lru) + return len(self._children) class _BlockHasher: @@ -949,16 +1408,60 @@ def _g(r: dict, *names: str, default=0): return low[nm] return default - rows: list[tuple[float, int, int, list[int]]] = [] + rows: list[tuple] = [] for r in records: ts = float(_g(r, "timestamp", "arrival", "arrival_ms", "time", default=0.0)) isl = int(_g(r, "input_length", "isl", "input_len", "prompt_len", default=1)) osl = int(_g(r, "output_length", "osl", "output_len", default=1)) hids = _g(r, "hash_ids", "block_hashes", "blocks", default=[]) or [] - rows.append((ts, max(1, isl), max(1, osl), [int(x) for x in hids])) + if "hp" in r: + # ``hp = [k, n]``: the first ``n`` blocks are those of file row ``k``. + k, m = r["hp"] + hids = list(rows[int(k)][3][: int(m)]) + [int(x) for x in hids] + think = _g(r, "think_ms", "delay_ms", default=None) + row = ( + ts, + max(1, isl), + max(1, osl), + [int(x) for x in hids], + None if think is None else max(0.0, float(think)), + ) + if "sess" in r: + row = row + ( + int(r["sess"]), + int(r.get("dep", -1)), + tuple(int(j) for j in (r.get("joins") or ())), + ) + rows.append(row) return rows +def _reqs_from_rows(rows) -> list[_Req]: + """``_Req`` list from Mooncake rows, in timestamp order. + + Rows carry an optional fifth field, the per-request closed-loop think time. + """ + out: list[_Req] = [] + for i, row in enumerate(sorted(rows or [], key=lambda x: x[0])): + a, isl, osl, hids = row[:4] + think = row[4] if len(row) > 4 else None + sess, dep, joins = (row[5], row[6], row[7]) if len(row) > 7 else (-1, -1, ()) + out.append( + _Req( + idx=i, + arrival_ms=a, + prompt_len=isl, + output_len=osl, + blocks=list(hids), + think_ms=think, + sess=sess, + dep=dep, + joins=joins, + ) + ) + return out + + def _draw_prefix_ids(n: int, num_prefixes: int, zipf: float, rng: random.Random) -> list[int]: """Assign each of ``n`` requests a shared-prefix id in ``[0, num_prefixes)``. @@ -1005,22 +1508,50 @@ def _route_and_warm( cache_blocks: int, rng: random.Random, overlap_weight: float = 1.0, + waiting_depth: int = 0, + resident_cap: int = 0, ) -> tuple[list[list[_Req]], dict[str, float]]: """Route requests across instances and derive per-request prefix-cache hits from a content-addressed block cache (as real serving engines do). - Each instance owns a :class:`_BlockCache`. Requests are processed in arrival - order; for each one the router picks a target, the number of resident leading - blocks (longest contiguous prefix match) becomes the cache hit -- its - ``cached_prefix`` tokens are seeded into ``num_computed`` so the scheduler - only prefills the uncached suffix -- and the request's blocks are then warmed - into that instance (LRU-evicted under ``cache_blocks`` capacity). + Each instance owns a :class:`_BlockCache`. For each request the router picks + a target, the number of resident leading blocks (longest contiguous prefix + match) becomes the cache hit -- its ``cached_prefix`` tokens are seeded into + ``num_computed`` so the scheduler only prefills the uncached suffix -- and + the request's blocks are then warmed into that instance (evicted under + ``cache_blocks`` capacity). Routing policies: ``kv`` trades cache overlap against load by the serving router's own cost function, with ``overlap_weight`` as the dial between them (``0`` routes purely by load); ``prefix_aware`` consistently hashes the leading block so same-prefix requests co-locate; ``round_robin``/``random`` ignore locality. + + ``waiting_depth`` is how many requests the engine has outstanding at once, + and it sets the order they are resolved in. With it unset they are resolved + oldest-first, which is only the order an engine admits in while everything + offered fits; once the waiting queue is deep a scheduler chooses from it by + longest resident prefix rather than by age (SGLang's default policy, and + what vLLM's prefix-caching scheduler approximates). The distinction does + not matter below the pressure point -- with room for every outstanding + context both orders touch the same blocks -- but above it the two diverge + completely: oldest-first walks the offered lanes round-robin, so a lane's + context is always evicted before its next turn and every request pays a + full reprefill, while prefix-first keeps working on what is already + resident and reuse falls off gradually instead of to zero. + + ``resident_cap`` is how many of those requests fit in the KV pool at once, + and it is what the reordering window has to be measured against: a + scheduler can only reorder requests that are actually queued. Below the + cap nothing waits, so there is no choice to make and admission is just the + order the lanes arrived in -- which is the regime where reuse degrades, + because C lanes all resident at once evict each other. Handing the policy + the full client count instead let it reorder an empty queue and kept reuse + high everywhere: on kimik3 at C=14 the hardware thrashed to 0.67 and the + replay reported 0.93, and the band where that happens is exactly where C + meets the cap. Above the cap the queue is genuinely deep, prefix-first is + the real policy, and measured reuse recovers -- 0.92 by C=44 -- which the + replay only gets right if the window grows with the queue and not with C. """ per_inst: list[list[_Req]] = [[] for _ in range(num_instances)] caches = [_BlockCache(cache_blocks) for _ in range(num_instances)] @@ -1036,7 +1567,52 @@ def _route_and_warm( inst_hits = [0] * num_instances inst_reqs = [0] * num_instances - for r in sorted(reqs, key=lambda x: (x.arrival_ms, x.idx)): + ordered = sorted(reqs, key=lambda x: (x.arrival_ms, x.idx)) + + def _admission_order(): + """The order the engine takes requests off its waiting queue. + + A closed loop holds one request per client, so the stream is the + clients interleaved: client ``i`` owns arrival positions ``i``, + ``i + depth``, ``i + 2 * depth``. Serving a request frees its own + client, and the turn that client issues next is the one continuing the + context just served -- the whole reason its prefix is worth keeping. + Refilling the slot from a global pointer instead hands it to some other + client and quietly turns the loop into a sliding window over arrival + order, which throws that reuse away. + """ + # Only the backlog is reorderable: what fits in the pool is already + # running and was admitted in the order it arrived. + depth = int(waiting_depth or 0) - max(0, int(resident_cap or 0)) + if depth <= 1: + yield from ordered + return + window = list(range(min(depth, len(ordered)))) + # A session turn's prompt extends the context of the request it waits + # on; resolving it first would hand that request a full hit and this one + # both prefills. Such a turn is not a candidate until its wait resolves. + known = {r.idx for r in ordered} + seen: set[int] = set() + + def _ready(i: int) -> bool: + r = ordered[i] + return all(d in seen or d not in known for d in (r.dep, *r.joins) if d >= 0) + + while window: + # Longest resident prefix anywhere in the fleet: the scheduler is + # choosing what to run next, not where to run it, so the routing + # decision below is still the router's to make. + pick = max( + [i for i in window if _ready(i)] or [min(window)], + key=lambda i: max(c.prefix_match(ordered[i].blocks or []) for c in caches), + ) + seen.add(ordered[pick].idx) + window.remove(pick) + if pick + depth < len(ordered): + window.append(pick + depth) + yield ordered[pick] + + for r in _admission_order(): blocks = r.blocks or [] if num_instances <= 1: inst = 0 @@ -1081,6 +1657,12 @@ def _route_and_warm( caches[inst].insert(blocks) per_inst[inst].append(r) + # Resolving prefixes in schedule order says how much each request has to + # prefill; it does not say when each one arrived. The per-instance loops + # below are arrival-driven, so hand them back the stream they expect. + for sub in per_inst: + sub.sort(key=lambda x: (x.arrival_ms, x.idx)) + n = len(reqs) summary = { "num_instances": float(num_instances), @@ -1148,6 +1730,7 @@ def dist(xs: list[float]) -> dict[str, float]: num_requests=sum(r.num_requests for r in results), makespan_ms=makespan, system_throughput_tps=sum(r.system_throughput_tps for r in results), + system_total_throughput_tps=sum(r.system_total_throughput_tps for r in results), saturated=any(r.saturated for r in results), ttft=dist(_pool("ttft")), ttft_arrival=dist(_pool("ttft_arrival")), @@ -1168,13 +1751,17 @@ def simulate_disaggregated( arrival_model: str = "poisson", num_requests: int = 400, warmup_frac: float = 0.1, + warmup_requests: int = 0, seed: int = 0, burstiness: float = 1.0, range_ratio: float = 1.0, kv_cache_tokens: int = 0, prebuilt: list[_Req] | None = None, closed_loop_clients: int = 0, + closed_loop_think_ms: float = 0.0, return_samples: bool = False, + duration_ms: float = 0.0, + closed_loop_idle_cap_ms: float = 0.0, ) -> DESResult: """Run a prefill pool and a decode pool as two stations on one clock. @@ -1217,6 +1804,9 @@ def simulate_disaggregated( input_len = max(1, req.input_seq_len) output_len = max(1, req.output_seq_len) max_running = max(1, req.resolved_max_concurrency()) + cap_seqs = int(getattr(req, "max_num_seqs", 0) or 0) + if cap_seqs > 0: + max_running = min(max_running, cap_seqs) token_budget = int(req.max_num_batched_tokens or 0) # 0 = unlimited long_prefill = int(req.chunked_prefill_size or 0) max_model_len = max(2, int(req.resolved_max_context_len())) @@ -1236,6 +1826,7 @@ def simulate_disaggregated( rng = random.Random(seed) clients = max(0, int(closed_loop_clients)) closed_loop = clients > 0 + think_ms = max(0.0, float(closed_loop_think_ms)) # ---- workload ---- if prebuilt is not None: @@ -1271,6 +1862,8 @@ def simulate_disaggregated( pending.sort(key=lambda r: r.arrival_ms) next_unissued = n + _set_host_ms(pending, req) + # The KV pool is per *pool*, not per fleet. A prefill pool holds a prompt # only until the handoff, while a decode replica holds the whole sequence # for the generation's duration, so splitting the budget evenly would @@ -1293,6 +1886,14 @@ def simulate_disaggregated( done: list[_Req] = [] handoff: dict[int, float] = {} next_arrival = 0 + # Closed loop: issued requests keyed on arrival (see ``simulate_once``). + idle_cap_ms = max(0.0, float(closed_loop_idle_cap_ms or 0.0)) + arrivals_heap: list[tuple[float, int, _Req]] = [] + if closed_loop: + arrivals_heap = [(pending[i].arrival_ms, i, pending[i]) for i in range(clients)] + heapq.heapify(arrivals_heap) + next_arrival = n + last_retire_ms = 0.0 p_steps = 0 d_steps = 0 pk_p_batch = 0 @@ -1319,6 +1920,9 @@ def _release_arrivals(t: float) -> int: p_waiting.append(pending[next_arrival]) next_arrival += 1 got += 1 + while arrivals_heap and arrivals_heap[0][0] <= t: + p_waiting.append(heapq.heappop(arrivals_heap)[2]) + got += 1 return got def _deliver_handoffs(t: float) -> int: @@ -1337,6 +1941,19 @@ def _step_prefill() -> bool: """One prefill-only batch on the prefill pool.""" nonlocal now_p, p_kv, p_busy, p_steps, pk_p_batch, pk_maxbatch, pk_kv_peak nonlocal pk_q_tokens + # An idle prefill pool advances to meet the next request instead of + # holding its clock where its last batch left it. The two stations keep + # separate clocks and in a split it is decode that leads: a closed-loop + # client retires on the decode clock and submits its replacement stamped + # with that time, so a prefill pool that has drained its queue is behind + # by however long it sat idle. The decode station already has this rule + # -- ``_deliver_handoffs`` moves ``now_d`` up to the handoff it is given + # -- and without the match here every request reissued while prefill was + # idle stays invisible to it until its clock crawls forward a step at a + # time, which is what collapsed a 32-client closed loop to one request + # in flight and pinned the decode batch at 1. + if not p_running and p_waiting: + now_p = max(now_p, p_waiting[0].arrival_ms) budget: float = token_budget if token_budget > 0 else math.inf scheduled: list[tuple[_Req, int, int]] = [] for r in p_running: @@ -1348,6 +1965,11 @@ def _step_prefill() -> bool: budget -= q while p_waiting and budget >= 1 and len(p_running) < max_running: head = p_waiting[0] + if head.arrival_ms > now_p: + # Queued against the simulation clock but not yet submitted on + # this station's, so it waits rather than being prefilled before + # the client asked for it. + break # A prefill pool reserves the prompt it is about to compute; it hands # the KV off and frees it, so it is not charged for the generation. need_kv = head.prompt_len @@ -1450,9 +2072,12 @@ def _reissue(t: float) -> None: would let a client hold two requests at once and report a concurrency the harness never ran. """ - nonlocal next_unissued + nonlocal next_unissued, last_retire_ms + last_retire_ms = max(last_retire_ms, t) if next_unissued < n: - pending[next_unissued].arrival_ms = t + nxt = pending[next_unissued] + nxt.arrival_ms = t + (think_ms if nxt.think_ms is None else nxt.think_ms) + heapq.heappush(arrivals_heap, (nxt.arrival_ms, next_unissued, nxt)) next_unissued += 1 # ---- event loop ---- @@ -1474,9 +2099,17 @@ def _unstall() -> None: guard = 0 bound = 4000 * max(1, n) + horizon = duration_ms if duration_ms and duration_ms > 0 else math.inf while guard < bound: guard += 1 - moved = _release_arrivals(now_p) + _deliver_handoffs(max([now_p, *now_d])) + # Both stations are fed against the simulation clock rather than the + # prefill station's own. Gating arrivals on ``now_p`` hid every request + # a client reissued at a decode-clock time from the pool that has to + # prefill it. + now_any = max([now_p, *now_d]) + if now_any >= horizon: + break + moved = _release_arrivals(now_any) + _deliver_handoffs(now_any) if moved: _unstall() @@ -1487,11 +2120,26 @@ def _unstall() -> None: if (d_running[k] or d_waiting[k]) and not stalled_d[k]: runnable.append((now_d[k], k)) if not runnable: + if ( + idle_cap_ms > 0 + and arrivals_heap + and not xfer + and not p_running + and not p_waiting + and not any(d_running) + and not any(d_waiting) + ): + lead = arrivals_heap[0][0] - (last_retire_ms + idle_cap_ms) + if lead > 0: + arrivals_heap = [(a - lead, i, q) for a, i, q in arrivals_heap] + for _a, _i, q in arrivals_heap: + q.arrival_ms -= lead # Nothing resident anywhere: jump to the next thing that can happen. nxt = [ t for t in ( pending[next_arrival].arrival_ms if next_arrival < n else None, + arrivals_heap[0][0] if arrivals_heap else None, min((ts for ts, _ in xfer), default=None), ) if t is not None and t < math.inf @@ -1517,27 +2165,19 @@ def _unstall() -> None: # ---- aggregate ---- done.sort(key=lambda r: r.finish_ms) - drop = int(len(done) * max(0.0, min(0.9, warmup_frac))) - sample = done[drop:] if len(done) - drop >= 8 else done - tok_ms_pt = max(0.0, req.tokenize_overhead_us) / 1000.0 + sample = _scored_sample(done, warmup_frac, warmup_requests) detok_ms = max(0.0, req.detokenize_overhead_us) / 1000.0 def _admit(r): return r.admit_ms if r.admit_ms >= 0 else r.arrival_ms - ttft = [ - (r.first_token_ms - _admit(r)) + tok_ms_pt * r.prompt_len - for r in sample - if r.first_token_ms >= 0 - ] + ttft = [(r.first_token_ms - _admit(r)) + r.host_ms for r in sample if r.first_token_ms >= 0] ttft_arrival = [ - (r.first_token_ms - r.arrival_ms) + tok_ms_pt * r.prompt_len - for r in sample - if r.first_token_ms >= 0 + (r.first_token_ms - r.arrival_ms) + r.host_ms for r in sample if r.first_token_ms >= 0 ] queue_wait = [_admit(r) - r.arrival_ms for r in sample if r.first_token_ms >= 0] e2e = [ - (r.finish_ms - r.arrival_ms) + tok_ms_pt * r.prompt_len + detok_ms * r.generated + (r.finish_ms - r.arrival_ms) + r.host_ms + detok_ms * r.generated for r in sample if r.finish_ms >= 0 ] @@ -1560,9 +2200,13 @@ def _admit(r): out_sample = sum(r.generated for r in sample) achieved_rate = (len(sample) * 1000.0 / span_ms) if span_ms > 0 else 0.0 sys_tps = (out_sample * 1000.0 / span_ms) if span_ms > 0 else 0.0 + tot_tokens = out_sample + sum(r.prompt_len for r in sample) + tot_tps = (tot_tokens * 1000.0 / span_ms) if span_ms > 0 else 0.0 else: achieved_rate = (len(done) * 1000.0 / makespan) if makespan > 0 else 0.0 sys_tps = (total_out * 1000.0 / makespan) if makespan > 0 else 0.0 + tot_tokens = total_out + sum(r.prompt_len for r in done) + tot_tps = (tot_tokens * 1000.0 / makespan) if makespan > 0 else 0.0 # Fleet utilisation, GPU-weighted. An unweighted mean of the two stations # would rate a 4-GPU prefill pool and a 12-GPU decode pool as equals. @@ -1627,6 +2271,7 @@ def dist(xs: list[float]) -> dict[str, float]: num_requests=len(done), makespan_ms=makespan, system_throughput_tps=sys_tps, + system_total_throughput_tps=tot_tps, saturated=saturated, ttft=dist(ttft), ttft_arrival=dist(ttft_arrival), @@ -1648,6 +2293,7 @@ def simulate_multi_instance( num_requests: int, seed: int, warmup_frac: float, + warmup_requests: int, burstiness: float, range_ratio: float, kv_cache_tokens: int, @@ -1659,8 +2305,15 @@ def simulate_multi_instance( prefix_zipf: float, block_size: int, cache_blocks: int, + admit_backlog_only: bool = False, mooncake_rows: list[tuple[float, int, int, list[int]]] | None = None, closed_loop_clients: int = 0, + duration_ms: float = 0.0, + prefill_exclusive: bool = False, + cache_shares_pool: bool = False, + closed_loop_think_ms: float = 0.0, + whole_context_residency: bool = False, + closed_loop_idle_cap_ms: float = 0.0, ) -> DESResult: """Route one arrival stream across ``num_instances`` replicas and pool. @@ -1679,13 +2332,12 @@ def simulate_multi_instance( bs = int(block_size) if block_size and block_size > 0 else _DEFAULT_BLOCK_SIZE hasher = _BlockHasher() + def _trace_requests() -> list[_Req]: + return _reqs_from_rows(mooncake_rows) + if mooncake_rows is not None: # Trace-driven: arrivals, lengths and block hashes all come from the file. - rows = sorted(mooncake_rows, key=lambda x: x[0]) - reqs = [ - _Req(idx=i, arrival_ms=a, prompt_len=isl, output_len=osl, blocks=list(hids)) - for i, (a, isl, osl, hids) in enumerate(rows) - ] + reqs = _trace_requests() else: arrivals = _generate_arrivals(num_requests, rate_per_s, arrival_model, rng, burstiness) reqs = _build_workload(len(arrivals), arrivals, input_len, output_len, range_ratio, rng) @@ -1697,43 +2349,75 @@ def simulate_multi_instance( r.prefix_id = pid r.blocks = _blocks_from_prefix(r.idx, r.prompt_len, pid, eff_prefix_len, bs, hasher) - per_inst, prefix_summary = _route_and_warm( - reqs, - policy=routing, - num_instances=num_instances, - block_size=bs, - cache_blocks=cache_blocks, - rng=rng, - overlap_weight=overlap_weight, - ) - prefix_summary["routing"] = ( - float(_ROUTING_POLICIES.index(routing)) if routing in _ROUTING_POLICIES else -1.0 - ) - prefix_summary["trace_driven"] = 1.0 if mooncake_rows is not None else 0.0 + def _warm_and_run(pool_reqs: list[_Req], cap_blocks: int) -> DESResult: + per_inst, prefix_summary = _route_and_warm( + pool_reqs, + policy=routing, + num_instances=num_instances, + block_size=bs, + cache_blocks=cap_blocks, + rng=random.Random(seed), + overlap_weight=overlap_weight, + waiting_depth=closed_loop_clients, + resident_cap=_resident_cap( + pool_reqs, + kv_cache_tokens, + inference_config.request_config.resolved_max_concurrency(), + admit_backlog_only, + ), + ) + prefix_summary["routing"] = ( + float(_ROUTING_POLICIES.index(routing)) if routing in _ROUTING_POLICIES else -1.0 + ) + prefix_summary["trace_driven"] = 1.0 if mooncake_rows is not None else 0.0 - results: list[DESResult] = [] - for i, sub in enumerate(per_inst): - if not sub: - continue - # Per-instance offered rate ≈ its share of the global stream. - inst_rate = rate_per_s * (len(sub) / len(reqs)) if reqs else rate_per_s - results.append( - simulate_once( - inference_config, - projector, - rate_per_s=inst_rate, - arrival_model=arrival_model, - seed=seed + i, - warmup_frac=warmup_frac, - kv_cache_tokens=kv_cache_tokens, - prebuilt=sub, - return_samples=True, - # Concurrency is per engine, so every replica runs the full - # client count rather than a share of it. - closed_loop_clients=closed_loop_clients, + results: list[DESResult] = [] + for i, sub in enumerate(per_inst): + if not sub: + continue + # Per-instance offered rate ≈ its share of the global stream. + inst_rate = rate_per_s * (len(sub) / len(pool_reqs)) if pool_reqs else rate_per_s + results.append( + simulate_once( + inference_config, + projector, + rate_per_s=inst_rate, + arrival_model=arrival_model, + seed=seed + i, + warmup_frac=warmup_frac, + warmup_requests=warmup_requests, + kv_cache_tokens=kv_cache_tokens, + prebuilt=sub, + return_samples=True, + # Concurrency is per engine, so every replica runs the full + # client count rather than a share of it. + closed_loop_clients=closed_loop_clients, + closed_loop_think_ms=closed_loop_think_ms, + closed_loop_idle_cap_ms=closed_loop_idle_cap_ms, + whole_context_residency=whole_context_residency, + # Every replica stops on the same clock: the window is the + # harness's, so it is not divided across instances the way + # the request stream is. + duration_ms=duration_ms, + # Whether prefill excludes decode is a property of the engine, + # so every replica schedules the same way. + prefill_exclusive=prefill_exclusive, + ) ) - ) - agg = _aggregate_instances(results, prefix_summary) + return _aggregate_instances(results, prefix_summary) + + agg = _warm_and_run(reqs, cache_blocks) + # The cache and the running set are the same memory, and the warm pass runs + # before there is a running set to measure. So measure one, then re-warm + # against what it left free. Residency is independent of reuse, so the + # second pass is the answer and not a step towards it. + if cache_shares_pool and mooncake_rows is not None and kv_cache_tokens > 0: + live = float((agg.packing or {}).get("kv_mean_tokens") or 0.0) + free_blocks = _free_cache_blocks(kv_cache_tokens, live, bs, cache_blocks) + if free_blocks != cache_blocks: + agg = _warm_and_run(_trace_requests(), free_blocks) + agg.prefix["cache_shares_pool"] = 1.0 + agg.prefix["cache_free_tokens"] = float(max(0.0, kv_cache_tokens - live)) return agg @@ -1746,6 +2430,7 @@ def run_des( num_requests: int = 400, seed: int = 0, warmup_frac: float = 0.1, + warmup_requests: int = 0, sweep: bool = False, burstiness: float = 1.0, range_ratio: float = 1.0, @@ -1762,9 +2447,15 @@ def run_des( block_size: int = 0, cache_blocks: int = 0, mooncake_trace: str | None = None, + duration_ms: float = 0.0, + admit_backlog_only: bool = False, prefill_exclusive: bool = False, new_seqs_per_step: int = 0, closed_loop: bool = False, + closed_loop_think_ms: float = 0.0, + cache_shares_pool: bool = False, + whole_context_residency: bool = False, + closed_loop_idle_cap_ms: float = 0.0, ) -> dict[str, object]: """Run the DES at the configured load and (optionally) a load sweep. @@ -1790,17 +2481,14 @@ def run_des( # engine has no way to say it is doing. if getattr(getattr(inference_config, "disaggregation_config", None), "enabled", False): reqs = None + prefix_summary: dict[str, float] | None = None if mooncake_rows is not None: hasher = _BlockHasher() bs = int(block_size) if block_size and block_size > 0 else _DEFAULT_BLOCK_SIZE - rows = sorted(mooncake_rows, key=lambda x: x[0]) - reqs = [ - _Req(idx=i, arrival_ms=a, prompt_len=isl, output_len=osl, blocks=list(hids)) - for i, (a, isl, osl, hids) in enumerate(rows) - ] + reqs = _reqs_from_rows(mooncake_rows) # The prefill pool is what owns a prefix cache here, so the hits are # warmed against one station rather than routed across a fleet. - _route_and_warm( + _, prefix_summary = _route_and_warm( reqs, policy="kv", num_instances=1, @@ -1808,7 +2496,18 @@ def run_des( cache_blocks=eff_cache_blocks, rng=random.Random(seed), overlap_weight=overlap_weight, + waiting_depth=( + inference_config.request_config.resolved_max_concurrency() if closed_loop else 0 + ), + resident_cap=_resident_cap( + reqs, + kv_cache_tokens, + inference_config.request_config.resolved_max_concurrency(), + admit_backlog_only, + ), ) + prefix_summary["routing"] = float(_ROUTING_POLICIES.index("kv")) + prefix_summary["trace_driven"] = 1.0 del hasher eff_rate = rate_per_s if mooncake_rows and not closed_loop: @@ -1824,6 +2523,7 @@ def run_des( num_requests=num_requests, seed=seed, warmup_frac=warmup_frac, + warmup_requests=warmup_requests, burstiness=burstiness, range_ratio=range_ratio, kv_cache_tokens=kv_cache_tokens, @@ -1831,7 +2531,16 @@ def run_des( closed_loop_clients=( inference_config.request_config.resolved_max_concurrency() if closed_loop else 0 ), + closed_loop_think_ms=closed_loop_think_ms, + closed_loop_idle_cap_ms=closed_loop_idle_cap_ms, + duration_ms=duration_ms, ) + # The split warms a prefix cache exactly as the colocated path does, + # but discarded the summary, so every disaggregated row reported no + # hit rate at all -- indistinguishable, to anyone reading the output, + # from a split that never reuses anything. + if prefix_summary is not None: + out["point"].prefix = prefix_summary return out if closed_loop: clients = inference_config.request_config.resolved_max_concurrency() @@ -1846,6 +2555,7 @@ def run_des( num_requests=num_requests, seed=seed, warmup_frac=warmup_frac, + warmup_requests=warmup_requests, burstiness=burstiness, range_ratio=range_ratio, kv_cache_tokens=kv_cache_tokens, @@ -1857,8 +2567,15 @@ def run_des( prefix_zipf=max(0.0, prefix_zipf or 0.0), block_size=max(0, block_size or 0), cache_blocks=eff_cache_blocks, + admit_backlog_only=admit_backlog_only, mooncake_rows=mooncake_rows, closed_loop_clients=clients, + closed_loop_think_ms=closed_loop_think_ms, + closed_loop_idle_cap_ms=closed_loop_idle_cap_ms, + duration_ms=duration_ms, + prefill_exclusive=prefill_exclusive, + cache_shares_pool=cache_shares_pool, + whole_context_residency=whole_context_residency, ) return out out["point"] = simulate_once( @@ -1869,12 +2586,17 @@ def run_des( num_requests=num_requests, seed=seed, warmup_frac=warmup_frac, + warmup_requests=warmup_requests, range_ratio=range_ratio, kv_cache_tokens=kv_cache_tokens, record_steps=record_steps, closed_loop_clients=clients, + closed_loop_think_ms=closed_loop_think_ms, + closed_loop_idle_cap_ms=closed_loop_idle_cap_ms, + whole_context_residency=whole_context_residency, prefill_exclusive=prefill_exclusive, new_seqs_per_step=new_seqs_per_step, + duration_ms=duration_ms, ) return out # A block cache is modelled whenever there is a fleet, a synthetic prefix @@ -1895,6 +2617,7 @@ def run_des( num_requests=num_requests, seed=seed, warmup_frac=warmup_frac, + warmup_requests=warmup_requests, burstiness=burstiness, range_ratio=range_ratio, kv_cache_tokens=kv_cache_tokens, @@ -1906,7 +2629,12 @@ def run_des( prefix_zipf=max(0.0, prefix_zipf or 0.0), block_size=max(0, block_size or 0), cache_blocks=eff_cache_blocks, + admit_backlog_only=admit_backlog_only, mooncake_rows=mooncake_rows, + duration_ms=duration_ms, + prefill_exclusive=prefill_exclusive, + cache_shares_pool=cache_shares_pool, + whole_context_residency=whole_context_residency, ) return out out["point"] = simulate_once( @@ -1917,11 +2645,13 @@ def run_des( num_requests=num_requests, seed=seed, warmup_frac=warmup_frac, + warmup_requests=warmup_requests, burstiness=burstiness, range_ratio=range_ratio, kv_cache_tokens=kv_cache_tokens, workload_file=workload_file, record_steps=record_steps, + duration_ms=duration_ms, ) if sweep and not workload_file: @@ -1945,6 +2675,7 @@ def run_des( num_requests=sweep_n, seed=seed, warmup_frac=warmup_frac, + warmup_requests=warmup_requests, burstiness=burstiness, range_ratio=range_ratio, kv_cache_tokens=kv_cache_tokens, diff --git a/infera/projection/core/projection/inference_projection/kv_cache.py b/infera/projection/core/projection/inference_projection/kv_cache.py index 4763d2a72..8ceb43795 100644 --- a/infera/projection/core/projection/inference_projection/kv_cache.py +++ b/infera/projection/core/projection/inference_projection/kv_cache.py @@ -56,6 +56,28 @@ def attention_dp_size(inference_config: InferenceConfig) -> int: return max(1, int(getattr(mp, "attention_data_parallel_size", 1) or 1)) +def decode_context_parallel_size(inference_config: InferenceConfig) -> int: + """Ranks a single sequence's context is split across during decode. + + Decode context parallelism (vLLM ``--decode-context-parallel-size``, SGLang + ``--dcp-size``) is the one axis that shrinks an MLA cache. Tensor + parallelism cannot: the latent is shared across heads, so every rank keeps + a whole copy. DCP instead gives each rank a slice of the *tokens*, so a + rank stores ``context / dcp`` per sequence and the replica holds ``dcp`` + times as many tokens. + + This is not a small effect on the models that use it. The published Kimi-K3 + rows move from a 2.9M-token pool at ``dcp_size: 1`` to 30.3M at + ``dcp_size: 8`` on the same eight GPUs. + """ + mp = inference_config.model_parallel_config + for name in ("decode_context_parallel_size", "context_parallel_size"): + v = getattr(mp, name, None) + if v: + return max(1, int(v)) + return 1 + + def kv_bytes_per_token_per_layer(inference_config: InferenceConfig) -> float: """Bytes of KV cache stored per token, per transformer layer, *per rank*.""" mc = inference_config.model_config @@ -92,6 +114,26 @@ def kv_bytes_per_token_per_layer(inference_config: InferenceConfig) -> float: return 2.0 * kv_heads_per_rank * head_dim * kv_bytes +def index_bytes_per_token_per_layer(inference_config: InferenceConfig) -> float: + """Bytes the sparse-attention indexer stores per token, per indexed layer, per rank. + + Native sparse attention has to decide what to attend to before it can + attend, and it decides by scoring the query against one small index key per + past token. Those keys are a per-token cache in their own right, sitting + beside the KV cache. + + One key per token per indexed layer, not one per head: the key is + head-shared the way an MLA latent is, and for the same reason it is not + divided by tensor parallelism either -- every rank scoring its own query + heads needs the whole key. The dtype is the stack's, not + ``kv_cache_dtype``, which is what makes the term separable in a memory log. + """ + mc = inference_config.model_config + if not mc.sparse_index_layer_count(): + return 0.0 + return int(mc.sparse_index_head_dim) * dtype_num_bytes(mc.sparse_index_dtype or "bf16") + + def linear_state_bytes_per_layer(inference_config: InferenceConfig) -> float: """Bytes of recurrent state one linear-attention layer stores, per rank, per sequence. @@ -177,7 +219,19 @@ def _alloc(ctx: int) -> int: # adding that state would let search fit batches the HBM cannot hold. full_frac = mc.full_attention_layer_fraction() per_token = per_token_per_layer * max(1, layers_on_rank) * full_frac - per_sequence = per_token * effective_context + # Native sparse attention's index keys are a second per-token cache, on + # however many layers carry an indexer. They are not scaled by + # ``full_frac``: a model is either hybrid-linear or sparse, and where a + # layer has an indexer it also has the KV cache the indexer selects from. + n_index = mc.sparse_index_layer_count() + if n_index: + index_frac = n_index / max(1, int(mc.num_layers or 0) or n_index) + per_token += ( + index_bytes_per_token_per_layer(inference_config) * max(1, layers_on_rank) * index_frac + ) + # Decode context parallelism splits the tokens of one sequence across + # ranks, so a rank holds its slice of the context rather than all of it. + per_sequence = per_token * effective_context / decode_context_parallel_size(inference_config) lin_frac = 1.0 - full_frac if lin_frac > 0.0: per_sequence += ( diff --git a/infera/projection/core/projection/inference_projection/launcher.py b/infera/projection/core/projection/inference_projection/launcher.py index c64a0aa1e..1daa92f60 100644 --- a/infera/projection/core/projection/inference_projection/launcher.py +++ b/infera/projection/core/projection/inference_projection/launcher.py @@ -36,6 +36,8 @@ "detokenize_overhead_us": "detokenize_overhead_us", "tokenize_overhead_us": "tokenize_overhead_us", "request_overhead_ms": "request_overhead_ms", + "uncached_prompt_latency_us": "uncached_prompt_latency_us", + "uncached_prompt_latency_max_tokens": "uncached_prompt_latency_max_tokens", "prefill_rate_us_per_token": "prefill_rate_us_per_token", "prefill_rate_lo_tokens": "prefill_rate_lo_tokens", "prefill_rate_hi_tokens": "prefill_rate_hi_tokens", @@ -47,13 +49,18 @@ "kv_block_size": "kv_block_size", "kv_offload_gb_per_gpu": "kv_offload_gb_per_gpu", "kv_offload_bw_gbps": "kv_offload_bw_gbps", + "workload_resident_tokens": "workload_resident_tokens", + "kv_pool_tokens": "kv_pool_tokens", "max_num_batched_tokens": "max_num_batched_tokens", + "max_num_seqs": "max_num_seqs", "ep_load_balance": "ep_load_balance", "redundant_experts": "redundant_experts", "request_rate": "request_rate", "arrival_model": "arrival_model", "attention_backend": "attention_backend", + "serving_engine": "serving_engine", "sparse_attention_topk": "sparse_attention_topk", + "sparse_indexer_cost_scale": "sparse_indexer_cost_scale", "sliding_window": "sliding_window", "sliding_window_layer_fraction": "sliding_window_layer_fraction", "moe_expert_dtype": "moe_expert_dtype", @@ -112,6 +119,10 @@ def _collect_inference_overrides(args) -> dict[str, object]: attn_dp = getattr(args, "attention_dp_size", None) if attn_dp is not None: overrides["attention_data_parallel_size"] = int(attn_dp) + # Decode context parallelism is the same kind of serving-only axis. + dcp = getattr(args, "decode_context_parallel_size", None) + if dcp is not None: + overrides["decode_context_parallel_size"] = int(dcp) # --comm-model {explicit,builtin} → collective_enabled. comm_model = getattr(args, "comm_model", None) if comm_model is not None: @@ -195,8 +206,18 @@ def _print_performance(inference_config, perf, gpu_cost_per_hour=None) -> None: feats.append(f"redundant_experts={req.redundant_experts}") if getattr(req, "attention_backend", None): feats.append(f"attn_backend={req.attention_backend}") + if getattr(req, "serving_engine", None): + feats.append(f"engine={req.serving_engine}") if getattr(req, "sparse_attention_topk", 0): feats.append(f"sparse_attn_topk={req.sparse_attention_topk}") + idx_heads = int(getattr(mc, "sparse_index_n_heads", 0) or 0) + if idx_heads: + cost = float(getattr(req, "sparse_indexer_cost_scale", 1.0) or 1.0) + feats.append( + f"indexer={idx_heads}x{mc.sparse_index_head_dim}" + f"@{mc.sparse_index_layers or mc.num_layers}L" + + (f"(cost={cost:g}x)" if cost != 1.0 else "") + ) n_lin = mc.linear_attention_layer_count() if n_lin: feats.append( @@ -655,7 +676,7 @@ def _anchor_from_store(args, inference_config): from .search.regime import recipe_from_inference_config store = AnchorStore(root) - recipe = recipe_from_inference_config(inference_config) + recipe = recipe_from_inference_config(inference_config, getattr(args, "gpu_arch", None)) model = _anchor_model_filter(store, args) entry, distance = store.nearest(recipe, model=model) except Exception as exc: # noqa: BLE001 - a broken store must not fail a projection @@ -945,6 +966,8 @@ def _btp(blob): num_requests=int(getattr(args, "des_num_requests", 400) or 400), seed=int(getattr(args, "des_seed", 0) or 0), warmup_frac=float(getattr(args, "des_warmup_frac", 0.1) or 0.0), + warmup_requests=int(getattr(args, "des_warmup_requests", 0) or 0), + admit_backlog_only=bool(getattr(args, "des_admit_backlog_only", False)), sweep=bool(getattr(args, "des_sweep", False)), burstiness=float(getattr(args, "des_burstiness", 1.0) or 1.0), range_ratio=float(getattr(args, "des_range_ratio", 1.0) or 1.0), @@ -961,7 +984,12 @@ def _btp(blob): block_size=int(getattr(args, "des_block_size", 0) or 0), cache_blocks=int(getattr(args, "des_kv_blocks", 0) or 0), mooncake_trace=mooncake_trace, + duration_ms=float(getattr(args, "des_duration_s", 0.0) or 0.0) * 1000.0, closed_loop=closed_loop, + closed_loop_think_ms=float(getattr(args, "des_client_think_ms", 0.0) or 0.0), + closed_loop_idle_cap_ms=float(getattr(args, "des_client_idle_cap_ms", 0.0) or 0.0), + cache_shares_pool=bool(getattr(args, "des_cache_shares_pool", False)), + whole_context_residency=bool(getattr(args, "des_whole_context_residency", False)), prefill_exclusive=bool(getattr(args, "des_exclusive_prefill", False)), new_seqs_per_step=int(getattr(args, "des_new_seqs_per_step", 0) or 0), ) diff --git a/infera/projection/core/projection/inference_projection/performance.py b/infera/projection/core/projection/inference_projection/performance.py index 112a16d0e..6c90b87a6 100644 --- a/infera/projection/core/projection/inference_projection/performance.py +++ b/infera/projection/core/projection/inference_projection/performance.py @@ -231,6 +231,18 @@ def _usable_packed_probe(packed, seq_rate_ms_per_tok=0.0, seq_cost_at=None): # warns and asks for a re-harvest rather than quietly standing in for one. +def _kv_bytes(dtype) -> float | None: + """Bytes per KV element for a cache dtype name; None when unstated.""" + d = str(dtype or "").lower() + if not d: + return None + if "fp4" in d or "int4" in d: + return 0.5 + if "fp8" in d or "int8" in d: + return 1.0 + return 2.0 + + def _safe_forward(profiler, batch: int, seq_len: int) -> float: """Forward time of a sub-profiler, or 0 if it does not implement timing. @@ -616,6 +628,7 @@ def __init__( # 0 => flat (no grid), preserving prior behaviour. self._decode_kv_slope_ms: float = 0.0 # ms per KV token (batch-independent) self._decode_ctx_ref: float = 0.0 # context the batch curve was measured at + self._decode_ctx_mid: float = 0.0 self._decode_ctx_max: float = 0.0 # largest measured context (guard) self._decode_kv_slope_by_batch: list = [] # (batch, ms per KV token) self._meas_prefill_rate_ms_per_tok: float = 0.0 # for sub-prompt prefill pieces @@ -652,6 +665,16 @@ def __init__( # per-output-token time with acceptance already folded in, harvested # without it is a single-token step. 0 = no speculation on the anchor. self._bench_spec_k = 0 + # (batch, context, width) -> verify/single step ratio, or None where the + # analytical path could not be evaluated. Memoised because the replay + # prices a step per simulated event and the ratio moves only with the + # step shape. + self._verify_ratio_cache: dict = {} + # (batch, ref_batch, context, phase) -> modelled step ratio, for giving + # a single-point anchor a slope. Same memoisation reason. + self._batch_shape_cache: dict = {} + # Set only while a single-point anchor's batch shape is being modelled. + self._shape_decode_schedule = False # phase -> batch -> tp -> (ms, ep, pp), and the split fitted from it. self._bench_scaling_raw: dict = {} self._bench_scaling_fit: dict = {} @@ -789,17 +812,25 @@ def _loglog_transport(batch: int, pts: list) -> float: slope = (y1 - y0) / (x1 - x0) if x1 != x0 else 0.0 return math.exp(y0 + slope * (lb - x0)) - def _transport_batch(self, batch: int, pts: list) -> float: + def _transport_batch(self, batch: int, pts: list, phase: str = "decode") -> float: """Transport a measured ``(batch -> ms)`` curve to an arbitrary ``batch`` — MEASUREMENT-ONLY. The benchmark protocol always sweeps batch within a single run, so - ``pts`` carries >= 2 measured points and we interpolate/extrapolate the - real curve in log-log space (:meth:`_loglog_transport`) — the analytical - (origami) simulator is never consulted for the batch shape. A lone - anchor (a degenerate, non-swept artifact) holds its measured value - rather than falling back to the simulator, so a benchmark-calibrated - projection stays free of simulator bias by construction. + ``pts`` normally carries >= 2 measured points and we interpolate/ + extrapolate the real curve in log-log space + (:meth:`_loglog_transport`) — the analytical (origami) simulator is + never consulted for the batch shape. + + A lone point is a degenerate, non-swept artifact, and holding it flat + is not the neutral choice it looks like. Step latency grows with batch + on every measured curve here, so a flat hold asserts a zero slope -- + the one shape the hardware never has -- and it asserts it over the + whole ladder a sweep is ranked on, which collapses the ordering the + projection exists to produce. The measured point still sets the level; + the shape around it comes from the model, which is the same split + :meth:`_restore_whole` uses to carry an anchor across widths. Falls + back to the flat hold only where that ratio cannot be evaluated. Returns the exact measured value when ``batch`` is itself measured.""" if not pts: @@ -810,7 +841,56 @@ def _transport_batch(self, batch: int, pts: list) -> float: return v0 if len(P) >= 2: return self._loglog_transport(batch, P) - return P[0][1] + b0, v0 = P[0] + shape = self._batch_shape_ratio(batch, b0, phase) + return v0 if shape is None else v0 * shape + + def _batch_shape_ratio(self, batch: int, ref_batch: int, phase: str) -> float | None: + """How a step at ``batch`` compares with one at ``ref_batch``, modelled. + + Used only to give a single-point anchor a slope. ``None`` when the + analytical path cannot be evaluated, which leaves the caller holding + the measured value flat. + """ + b, b0 = max(1, int(batch)), max(1, int(ref_batch)) + if b == b0: + return 1.0 + kv = max(1, int(self._meas_ref_input or self.cfg.request_config.input_seq_len or 1024)) + key = (b, b0, kv, phase) + if key in self._batch_shape_cache: + return self._batch_shape_cache[key] + try: + if phase == "prefill": + # Prefill work is per token, so the comparable pair is the same + # per-sequence prompt carried at each batch. + here = self._forward_times(b, kv, "prefill", kv).total_ms + there = self._forward_times(b0, kv, "prefill", kv).total_ms + else: + # The whole step, at the width the target verifies at, and not + # the bare forward pass. A measured step carries its fixed + # costs -- dispatch, occupancy, the launch floor -- and at the + # small batches anchors come from those are most of it, so + # scaling the measurement by a compute-only ratio grows the + # fixed part with batch too. Carried up from a C=1 anchor that + # read DeepSeek-V4-Pro at 25.7 ms/token at C=16 against 8.1 + # measured, while the model's own step grows only 2.1x. + spec_k = int(self.cfg.request_config.speculative_num_tokens or 0) + q = spec_k + 1 if spec_k > 0 else 1 + self._shape_decode_schedule = True + try: + here = self._analytic_decode_step_ms(b, kv, q) + there = self._analytic_decode_step_ms(b0, kv, q) + finally: + self._shape_decode_schedule = False + except Exception: # noqa: BLE001 - simulator availability is arch-dependent + self._batch_shape_cache[key] = None + return None + if here <= 0.0 or there <= 0.0: + self._batch_shape_cache[key] = None + return None + ratio = here / there + self._batch_shape_cache[key] = ratio + return ratio # -- measured-time accessors (benchmark-based projection) ------------------ @@ -826,6 +906,50 @@ def _decode_kv_slope_at(self, batch: int) -> float: return self._decode_kv_slope_ms return self._loglog_transport(batch, pts) if len(pts) >= 2 else pts[0][1] + def _kv_dtype_delta_ms(self, batch: int, context: float) -> float: + """Decode-step cost of the target's resident KV over the anchor's -- + another dtype, another context, or both -- priced by the analytic step + at the target view. + + For anchors whose sweep is at one context, so no KV slope was fitted + and the measured step carries the anchor's KV, at the anchor's width + and length, with nothing to scale it by. The anchor's decode ran from + its prompt length to prompt plus output, so that midpoint is what it + read.""" + key = (int(batch), int(context) // 256) + memo = self._kv_delta_memo + if key in memo: + return memo[key] + rc = self.cfg.request_config + tgt = rc.kv_cache_dtype + ref = int(self._decode_ctx_mid or context) + delta = 0.0 + try: + t = self._forward_times(max(1, batch), 1, "decode", max(1, int(context))).total_ms + rc.kv_cache_dtype = self._bench_kv_dtype + a = self._forward_times(max(1, batch), 1, "decode", max(1, ref)).total_ms + delta = t - a + except Exception: + delta = 0.0 + finally: + rc.kv_cache_dtype = tgt + memo[key] = delta + return delta + + def _sparse_decode_reads(self) -> bool: + """Whether decode attention selects or compresses what it reads. + + Such a step reads a bounded set rather than the resident context, and + the analytic step does not know how bounded: moving an ISL-8192 anchor + to ISL 1024 by it took 3-24% off SGLang's DeepSeek-V4-Flash and + GLM-5.2 steps that measure within 7% of the anchor.""" + mc = self.cfg.model_config + return bool( + int(getattr(mc, "sparse_index_n_heads", 0) or 0) + or getattr(mc, "compress_ratios", None) + or int(getattr(self.cfg.request_config, "sparse_attention_topk", 0) or 0) + ) + def _measured_decode_step_ms(self, batch: int, context: float | None = None) -> float: """Measured whole-model / composed decode *step* latency at ``batch``. @@ -837,7 +961,7 @@ def _measured_decode_step_ms(self, batch: int, context: float | None = None) -> if self._decode_pad_to_capture: base = self._bucket_up(batch, pts) else: - base = self._transport_batch(batch, pts) + base = self._transport_batch(batch, pts, "decode") if ( context is not None and self._decode_kv_slope_ms > 0.0 @@ -857,8 +981,21 @@ def _measured_decode_step_ms(self, batch: int, context: float | None = None) -> # reference it eventually crosses the context-free cost of the # step, which is weights and compute and does not go away. slope = self._decode_kv_slope_at(batch) + r = getattr(self, "_kv_bytes_ratio", 1.0) + base = max( + base + slope * (r * float(context) - self._decode_ctx_ref), + self._decode_floor_ms(batch), + ) + elif context is not None and ( + getattr(self, "_kv_bytes_ratio", 1.0) != 1.0 + or ( + self._decode_ctx_mid > 0.0 + and not self._sparse_decode_reads() + and abs(float(context) - self._decode_ctx_mid) > 0.1 * self._decode_ctx_mid + ) + ): base = max( - base + slope * (float(context) - self._decode_ctx_ref), + base + self._kv_dtype_delta_ms(batch, float(context)), self._decode_floor_ms(batch), ) return base @@ -873,7 +1010,7 @@ def _measured_decode_step_ms(self, batch: int, context: float | None = None) -> def _measured_full_prefill_ms(self, batch: int) -> float: """Measured whole-model / composed prefill latency for the full prompt.""" if self._meas_whole.get("prefill"): - return self._transport_batch(batch, self._meas_whole["prefill"]) + return self._transport_batch(batch, self._meas_whole["prefill"], "prefill") tok = self._meas_ref_input or 1 d = self._restore_per_layer( "dense", self._meas_layer.get(("prefill", "dense"), 0.0), batch, tok @@ -883,7 +1020,7 @@ def _measured_full_prefill_ms(self, batch: int) -> float: ) return self._n_dense * d + self._n_moe * m + self._restore_pp_ms(batch, tok) - def _measured_prefill_tokens_ms(self, total_tokens: int) -> float: + def _measured_prefill_tokens_ms(self, total_tokens: int, kv_len: int = 0) -> float: """Measured prefill time for an arbitrary token count (chunk pieces). Prefill is compute-bound and ~linear in total processed tokens, so we @@ -893,9 +1030,11 @@ def _measured_prefill_tokens_ms(self, total_tokens: int) -> float: if rate <= 0: # Decode-only artifact (or an untrusted prefill, see # set_benchmark_calibration): simulate the chunk rather than bill it - # as free. + # as free. The chunk attends over everything already resident, so + # it is simulated at that context: at the chunk's own length an 8k + # suffix on a 200k agentic prompt cost what an 8k prompt does. return self._forward_times( - 1, max(1, total_tokens), "prefill", max(1, total_tokens) + 1, max(1, total_tokens), "prefill", max(1, kv_len or total_tokens) ).total_ms # Charged once per step, not per token and not per request: a step # packing eight 1024-token prompts pays it once, which is what makes @@ -951,9 +1090,24 @@ def set_benchmark_calibration(self, benchmark_layer_times: dict) -> None: self._bench_ep = int(meta.get("benchmark_ep") or meta.get("ep") or 1) self._bench_pp = int(meta.get("benchmark_pp") or meta.get("pp") or 1) _attn_dp = meta.get("attention_data_parallel_size") - self._bench_attn_dp = int(_attn_dp) if _attn_dp else None + self._bench_attn_dp = ( + int(_attn_dp) + if _attn_dp + else _attention_dp_from_server_args(meta.get("server_args"), tp=self._bench_tp) + ) self._bench_spec_k = int(meta.get("speculative_num_tokens") or 0) self._decode_pad_to_capture = bool(meta.get("decode_pad_to_capture")) + # The fitted KV slope is a stream out of the anchor's cache, so a target + # that stores KV at another width reads proportionally more or less. + # Without it GLM-5.2-MXFP4 on SGLang at ISL 8192 projected fp8 and bf16 + # KV identically where they measured 1771 and 1381 tok/s at C128. + _ab = _kv_bytes(meta.get("kv_cache_dtype")) + _tb = _kv_bytes(getattr(self.cfg.request_config, "kv_cache_dtype", None)) + self._kv_bytes_ratio = (_tb / _ab) if (_ab and _tb) else 1.0 + _akv = str(meta.get("kv_cache_dtype") or "") + self._bench_kv_dtype = "bf16" if _akv.lower() in ("", "auto") else _akv + self._kv_delta_memo: dict = {} + self._decode_ctx_mid = float(ref_input) + 0.5 * float(meta.get("output_len") or 0) self._meas_ref_input = ref_input @@ -1039,6 +1193,54 @@ def set_benchmark_calibration(self, benchmark_layer_times: dict) -> None: "hit fraction to 1.0 for repeated-prefix traffic." ) pre_pts = [] + # A prefill curve also has to be transportable to the prompt length + # being asked about, and only the length probe makes it so. Four or + # more probed lengths fit a quadratic, which is read at the target + # (held at the last probed length past the ladder). Two lengths fit + # only a chord, and what survives is the sweep's flat per-token rate + # at the anchor's own ``input_len`` -- a number that is then sharded + # across TP width and billed unchanged at any prompt length. + # + # For a 16x reach that is not a small error, because prefill + # per-token cost is not context-free: attention grows with context. + # DeepSeek-V4-Pro's TP8 MI355X anchor probes 4096 and 8192 only, so + # at ISL 130000 it charges the 8192-token rate, which prices the + # prompt as very nearly free -- TTFT 531 ms against a measured + # 15570, and 108005 tok/s/gpu against a best measured 49370, a 2.2x + # over-read on the headline metric. Simulated prefill lands within + # 0.3% of the same measurement, so the fallback is strictly better + # than the flat rate here. + # + # Refused rather than scaled: transporting it would mean inventing + # the context term from the analytical model and calling the result + # measured, which is the fallback's job and is reported as such. + # Decode is unaffected and stays calibrated. + _prefill_probed = sorted( + int(p.get("input_len") or 0) + for p in ((meta.get("prefill_anchor") or {}).get("points") or []) + ) + _has_curve = bool( + ((meta.get("prefill_anchor") or {}).get("curve_fit") or {}).get("ms_per_token") + ) + _target_n = int(self.cfg.request_config.input_seq_len or ref_input or 0) + if ( + pre_pts + and not _has_curve + and _prefill_probed + and _target_n > 1.5 * max(_prefill_probed) + ): + print( + f"[inferasim:Inference] WARNING: PREFILL IS NOT CALIBRATED. This " + f"anchor probed prefill at {min(_prefill_probed)}..{max(_prefill_probed)} " + f"tokens, too few lengths to fit a context curve, and the target " + f"runs {_target_n}. Its flat per-token rate carries no context " + f"growth, so billing it here would price a " + f"{_target_n / max(_prefill_probed):.0f}x longer prompt at the same " + f"cost per token. Decode is calibrated as usual; prefill and TTFT " + f"are simulated. Re-harvest with --prefill-anchor-points 4 or more " + f"over a ladder that brackets this prompt length." + ) + pre_pts = [] pre_pts_bench = list(pre_pts) dec_pts_bench = list(dec_pts) if self._restore: @@ -1163,6 +1365,37 @@ def set_benchmark_calibration(self, benchmark_layer_times: dict) -> None: f"batch point — batch transport will hold it flat. Re-run the " f"benchmark with a batch sweep for an accurate {_ph} batch curve." ) + # And having a sweep is not the same as having one that reaches the + # batch being asked about. DeepSeek-V4-Pro's MI355X anchor sweeps + # batch 1..64 and the agentic recipe runs concurrency 256, so the + # curve is read 4x past its widest measured rung. + # + # The sign of that error is not stated here because it is not a + # property of the anchor alone: against the same measured runs, the + # closed-form path comes out 2.57x high at concurrency 256 while + # trace replay of the same candidate comes out 0.74x, and both use + # this curve. What is common to them is that the step cost past + # batch 64 is modelled rather than measured. Which is the thing to + # say, and the reason the fix is a harvest rather than a factor. + _tgt_b = max( + int(self.cfg.request_config.batch_size or 0), + int(getattr(self.cfg.request_config, "max_concurrency", 0) or 0), + ) + for _ph in ("prefill", "decode"): + _pts = self._meas_whole.get(_ph, []) + if len(_pts) < 2 or _tgt_b <= 0: + continue + _widest = max(int(b) for b, _ in _pts) + if _tgt_b > _widest: + print( + f"[inferasim:Inference] WARNING: {_ph} was swept to batch " + f"{_widest} and this config runs {_tgt_b}. The curve is being " + f"read {_tgt_b / _widest:.1f}x past its widest measured rung, " + f"so the step cost out here is modelled rather than measured. " + f"Harvest with --inference-batch-size {_tgt_b} (or " + f"--benchmark-batches up to {_tgt_b}) before trusting " + f"throughput or TPOT from this run." + ) # Per-token prefill rate (for sub-prompt chunk pieces): full-prompt # prefill of ``b`` seqs processes ``b * ref_input`` tokens. if pre_pts and ref_input > 0: @@ -2032,6 +2265,11 @@ def _anchor_floor_from(blob: dict) -> float: measurement is not a floor, it is just that measurement, and treating it as one would pin the whole curve to it. """ + floor = float(((blob.get("meta") or {}).get("decode_floor_ms")) or 0.0) + if floor > 0.0: + # Measured elsewhere: a single-point anchor can carry the fixed part + # of its step as solved from the same stack at two parallelisms. + return floor pts = [float(e["decode_ms"]) for e in (blob.get("sweep") or []) if e.get("decode_ms")] return min(pts) if len(pts) >= 2 else 0.0 @@ -2159,7 +2397,9 @@ def _restore_pp_ms(self, batch: int, tokens: int) -> float: return 0.0 return self._comm_tgt.pp_p2p_ms(batch, tokens) - self._comm_bench.pp_p2p_ms(batch, tokens) - def _origami_steps(self, batch: int, tokens: int, phase: str) -> tuple | None: + def _origami_steps( + self, batch: int, tokens: int, phase: str, bench_batch: int | None = None + ) -> tuple | None: """Simulated whole-step time at the target and bench views, ``(tgt, bench)``. Reuses the analytical ``_forward_times`` at the bench and target views by @@ -2191,7 +2431,7 @@ def _origami_steps(self, batch: int, tokens: int, phase: str) -> tuple | None: # message and genuinely exposed, so it stays in the ratio. comm_free = phase == "decode" - def _step(lm, comm, view, imb) -> float: + def _step(lm, comm, view, imb, batch) -> float: saved = (self._lm, self._comm, self._view, self._gemm, self._sdpa, self._moe_imbalance) self._lm, self._comm, self._view = lm, comm, view self._gemm, self._sdpa = self._gemm_sim, self._sdpa_sim @@ -2214,8 +2454,14 @@ def _step(lm, comm, view, imb) -> float: ) try: - s_tgt = _step(self._lm_ratio_tgt, comm_tgt, self._view_tgt, self._imb_tgt) - s_bench = _step(self._lm_ratio_bench, comm_bench, self._view_bench, self._imb_bench) + s_tgt = _step(self._lm_ratio_tgt, comm_tgt, self._view_tgt, self._imb_tgt, batch) + s_bench = _step( + self._lm_ratio_bench, + comm_bench, + self._view_bench, + self._imb_bench, + batch if bench_batch is None else bench_batch, + ) except Exception: return None if s_bench <= 0.0 or s_tgt <= 0.0: @@ -2461,11 +2707,95 @@ def _forward_times(self, batch: int, q_len: int, phase: str, kv_len: int) -> Pha # from 0.53x of measured TPOT to 0.83x on MI355X and 0.94x on GB300, # while the fixed-shape 8k sweep, where the floor barely binds, does not # move at all. - sparse_scale = ( - 1.0 - if phase == "decode" - else self.cfg.request_config.resolved_sparse_attention_scale(kv_len) - ) + # A model that records a per-layer compression schedule is costed from + # it, weighting each branch by how many layers actually run it, rather + # than from one top-k scale standing for the whole stack. Prefill only, + # for the reason just given: decode's measured step does not shrink with + # the selection, so the schedule is not what sets its cost either. + # A model that selects the same way on every layer is costed from its + # indexer geometry instead, for the same reason: one number cannot + # hold a selection term that falls away with context and an indexing + # term that does not. The floored top-k below stays as the fallback + # for a model that declares a window but no indexer to choose with. + sparse_scale = 1.0 + if phase != "decode" and os.getenv("INFERASIM_PREFILL_DENSE") == "1": + pass + elif phase != "decode": + from infera.projection.core.projection.training_config import ( + hybrid_attention_scale, + uniform_sparse_attention_scale, + ) + + sparse_scale = hybrid_attention_scale(self.cfg.model_config, kv_len) + if sparse_scale is None: + sparse_scale = uniform_sparse_attention_scale( + self.cfg.model_config, + kv_len, + self.cfg.request_config.sparse_attention_topk, + self.cfg.request_config.sparse_indexer_cost_scale, + ) + if sparse_scale is None: + sparse_scale = self.cfg.request_config.resolved_sparse_attention_scale(kv_len) + else: + # Decode keeps its dense charge -- the step does not shrink with + # the selection, which is what the measured TPOT says -- but a + # stack without a fused selection kernel pays for choosing on top + # of it, and that is not in the dense charge. Only the excess + # over a fused kernel is added, so a stack that has one is priced + # exactly as the calibration that set the dense charge left it. + from infera.projection.core.projection.training_config import ( + sparse_indexer_decode_overhead, + ) + + indexer_over = sparse_indexer_decode_overhead( + self.cfg.model_config, + self.cfg.request_config.sparse_attention_topk, + self.cfg.request_config.sparse_indexer_cost_scale, + ) + sparse_scale = 1.0 + indexer_over + if self._shape_decode_schedule: + # The dense charge sets a level, and in benchmark mode the + # anchor sets that instead; what is left for the model is how + # the step grows per resident sequence. Under a compression + # schedule a sequence adds its compressed pools and window, + # not a dense read of its context: charged dense, one extra + # DeepSeek-V4-Pro sequence at 205k context costs 2.7 ms of + # step against the ~0.5 ms the measured TPOT grows by. + from infera.projection.core.projection.training_config import ( + hybrid_attention_scale, + ) + + h = hybrid_attention_scale(self.cfg.model_config, kv_len) + if h is not None: + sparse_scale += h - 1.0 + from .kv_cache import attention_dp_size + + if os.getenv("INFERASIM_DPA_DECODE_DENSE") != "1" and attention_dp_size(self.cfg) > 1: + # The dense charge above was calibrated where a rank holds + # heads/tp and the step is a cache read. Under attention-DP a + # rank runs every head of its own sequences, so dense attention + # over the full context turns compute-bound -- and that compute + # is exactly what selection or compression removes. Charged + # dense, DeepSeek-V4-Pro atom TP8/DPA8 at C=48 projects a 108 ms + # ITL against ~21 ms measured, while the same load without DPA + # projects within 10%. So DPA decode is costed like prefill. + from infera.projection.core.projection.training_config import ( + hybrid_attention_scale, + uniform_sparse_attention_scale, + ) + + s = hybrid_attention_scale(self.cfg.model_config, kv_len) + if s is None: + s = uniform_sparse_attention_scale( + self.cfg.model_config, + kv_len, + self.cfg.request_config.sparse_attention_topk, + self.cfg.request_config.sparse_indexer_cost_scale, + ) + if s is None: + s = self.cfg.request_config.resolved_sparse_attention_scale(kv_len) + if s is not None and s < 1.0: + sparse_scale = s + indexer_over # Attention-DP: the memory model has always known that a rank under DP # attention owns a subset of the *requests* rather than a slice of every # request's heads, but the time model did not, and charged every rank @@ -2483,7 +2813,42 @@ def _forward_times(self, batch: int, q_len: int, phase: str, kv_len: int) -> Pha dp = attention_dp_size(self.cfg) attn_batch = max(1, math.ceil(batch / dp)) if dp > 1 else batch factor = self._attn_backend_mult * sparse_scale - if factor != 1.0 or attn_batch != batch: + if phase == "decode" and os.getenv("INFERASIM_DCP_DECODE") == "1": + # Decode context parallelism splits each sequence's cache across + # the DCP ranks, so a rank's decode attention reads its slice. + dcp = int( + getattr(self.cfg.model_parallel_config, "decode_context_parallel_size", 1) or 1 + ) + if dcp > 1: + factor /= dcp + # Opt-in, for engines whose decode kernels read the compressed / top-k + # cache: the dense charge sets the one-sequence level, and each further + # resident sequence adds its pools, window and indexer pass rather than + # a dense read of its context. Charged dense, the step's growth with + # concurrency is overstated several-fold at long agentic contexts. + growth = None + if ( + phase == "decode" + and dp <= 1 + and not self._shape_decode_schedule + and os.getenv("INFERASIM_DECODE_SPARSE_GROWTH") == "1" + ): + from infera.projection.core.projection.training_config import ( + hybrid_attention_scale, + uniform_sparse_attention_scale, + ) + + g = hybrid_attention_scale(self.cfg.model_config, kv_len) + if g is None: + g = uniform_sparse_attention_scale( + self.cfg.model_config, + kv_len, + self.cfg.request_config.sparse_attention_topk, + self.cfg.request_config.sparse_indexer_cost_scale, + ) + if g is not None and g < 1.0: + growth = g + if factor != 1.0 or attn_batch != batch or growth is not None: for has, prof, is_moe in ((has_dense, dense_p, False), (has_moe, moe_p, True)): if not has or not hasattr(prof, "get_sub_profiler"): continue @@ -2493,7 +2858,11 @@ def _forward_times(self, batch: int, q_len: int, phase: str, kv_len: int) -> Pha charged = sub.measured_forward_time(batch, q_len) actual = ( sub.measured_forward_time(attn_batch, q_len) if attn_batch != batch else charged - ) * factor + ) + if growth is not None and attn_batch > 1: + one = sub.measured_forward_time(1, q_len) + actual = one + max(0.0, actual - one) * growth + actual *= factor if is_moe: moe_compute = max(0.0, moe_compute + actual - charged) else: @@ -2840,7 +3209,9 @@ def _draft_overhead_ms(self, per_token_step_ms: float) -> float: return dcf * spec_k * max(0.0, per_token_step_ms) return 0.0 - def _measured_verify_step_scale(self) -> float: + def _measured_verify_step_scale( + self, batch: int | None = None, context: float | None = None + ) -> float: """Factor turning one measured decode number into one verify-step time. What the anchor measured decides this. Harvested *with* speculation, the @@ -2856,31 +3227,92 @@ def _measured_verify_step_scale(self) -> float: perfect acceptance -- and understated throughput by the same ratio, so speculation read worse the better it was accepted. - Harvested *without* speculation there is nothing folded in, and a verify - pass is not derivable from a single-token step: that is why speculation - is regime-defining. The anchor store refuses that pairing, but - ``--load-benchmark`` does not, so the old width scaling stays for it - rather than inventing a number that would look calibrated. + Harvested *without* speculation there is nothing folded in, and the + anchor is a single-token step. Charging the verify width ``k + 1`` for + it treats decode as compute-bound, which at these contexts it is not: + the pass reads the same KV once however many positions ride along, so + ``_verify_width_ratio`` prices the widening from the model and leaves + the level to the anchor. The blind width remains the fallback for hosts + where that ratio cannot be evaluated. + + Speculation stays regime-defining and the anchor store still refuses + the pairing; this is what ``--load-benchmark`` and a relaxed axis get + instead of a factor that scales with the draft depth. """ spec_k = int(self.cfg.request_config.speculative_num_tokens or 0) if spec_k <= 0: return 1.0 if self._bench_spec_k > 0: return max(1e-6, self._spec_tokens_per_step()) - return float(spec_k + 1) + ratio = self._verify_width_ratio(spec_k + 1, batch, context) + return float(spec_k + 1) if ratio is None else ratio + + def _verify_width_ratio( + self, width: int, batch: int | None, context: float | None + ) -> float | None: + """What a verify pass costs relative to the single-token step measured. + + The verify pass puts ``width`` query positions through the stack, but it + reads the KV cache once for the same sequences. At the contexts these + anchors are read at the read dominates, so the pass costs a little more + than one step and nowhere near ``width`` of them; charging the width is + what made a speculating target read slower the deeper its draft. + + Only the ratio between two step shapes comes from the analytical model + -- the anchor still sets the level. That is the division of labour + ``_restore_whole`` already uses to carry an anchor across tensor-parallel + widths, applied to query width instead. + + ``None`` when the analytical path cannot be evaluated on this host, + which leaves the caller on the blind width rather than on a ratio + nothing computed. + """ + if width <= 1: + return 1.0 + b = max(1, int(batch or 1)) + kv = max( + 1, + int(context or self._meas_ref_input or self.cfg.request_config.input_seq_len or 1024), + ) + key = (b, kv, int(width)) + if key in self._verify_ratio_cache: + return self._verify_ratio_cache[key] + try: + one = self._forward_times(b, 1, "decode", kv).total_ms + wide = self._forward_times(b, int(width), "decode", kv).total_ms + except Exception: # noqa: BLE001 - simulator availability is arch-dependent + self._verify_ratio_cache[key] = None + return None + if one <= 0.0 or wide <= 0.0: + self._verify_ratio_cache[key] = None + return None + # Bracketed by the two things it sits between: a pass carrying the extra + # positions cannot beat the single-token step inside it, and cannot cost + # more than running that step once per position. + ratio = min(float(width), max(1.0, wide / one)) + self._verify_ratio_cache[key] = ratio + return ratio def _decode_step_latency_ms(self, batch: int, kv_len: int, q_len: int = 1) -> float: # Benchmark-based: use the measured decode step directly (memory-bound, # ~flat in context over a generation, so no simulator context-scaling). if self._measured_mode: per_token = self._measured_decode_step_ms(batch, kv_len) - step = per_token * self._measured_verify_step_scale() + step = per_token * self._measured_verify_step_scale(batch, kv_len) # A speculation-harvested anchor already paid for the draft pass, so # adding the modelled overhead would bill it a second time. if self._bench_spec_k <= 0: step += self._draft_overhead_ms(per_token) step += self._decode_step_overhead_ms() + if self._meas_whole.get("decode"): + # A whole-step sweep is transported by width only; the stage + # boundaries a deeper target PP adds are charged here, once. + step += self._restore_pp_ms(batch, q_len) return max(step, self._decode_floor_ms(batch)) + return self._analytic_decode_step_ms(batch, kv_len, q_len) + + def _analytic_decode_step_ms(self, batch: int, kv_len: int, q_len: int = 1) -> float: + """The pure-simulate decode step, whatever mode the projector is in.""" ft = self._forward_times(batch, q_len, "decode", kv_len) per_token = ft.total_ms / max(1, q_len) step = ft.total_ms + self._draft_overhead_ms(per_token) + self._decode_step_overhead_ms() @@ -2910,27 +3342,52 @@ def mixed_step_latency_ms( decode_ctx: int, prefill_kv_len: int, q_len: int = 1, + prefill_batch: int = 1, ) -> float: """One scheduler step carrying a prefill chunk plus ``num_decode`` - concurrent decodes (``num_decode == 0`` → a pure prefill-chunk step).""" + concurrent decodes (``num_decode == 0`` → a pure prefill-chunk step). + + ``prefill_batch`` chunks of ``chunk_tokens`` ride the step together, one + per attention-DP rank.""" penalty = max(0.0, self.cfg.request_config.resolved_mixed_batch_penalty()) ov = self._decode_step_overhead_ms() chunk_tokens = max(1, int(chunk_tokens)) num_decode = max(0, int(num_decode)) + pb = max(1, int(prefill_batch)) if self._measured_mode: # See ``_measured_verify_step_scale``: a speculation-harvested decode # number is per output token, so the step scales by the tokens it # emits rather than by the verify width ``q_len``. - spec = self._measured_verify_step_scale() - prefill_piece = self._measured_prefill_tokens_ms(chunk_tokens) + spec = self._measured_verify_step_scale(max(1, num_decode), decode_ctx) + prefill_piece = self._measured_prefill_tokens_ms(chunk_tokens * pb, prefill_kv_len) + # The prefill rate is sharded by TP width alone, so an attention + # layout change is moved by the simulator's difference: the anchor's + # step for the one chunk a rank attends over, plus what the target + # layout adds on top of it (the other ranks' chunks through the + # MoE, every head of its own chunk). Not the ratio, and not the + # step for all ``pb`` chunks: on DeepSeek-V4-Flash under SGLang the + # measured TP4 prefill runs 10x the simulated one, almost all of it + # work each rank repeats per token it holds, and the DP4 step over + # four ranks' chunks measures 1.07x the TP4 step over one. Not when + # TP moved as well: the rate already carries that part of it. + if getattr(self, "_restore_layout_moved", False) and self._tgt_tp == self._bench_tp: + steps = self._origami_steps(pb, chunk_tokens, "prefill", bench_batch=1) + if steps is not None: + own = self._measured_prefill_tokens_ms(chunk_tokens, prefill_kv_len) + prefill_piece = max(own, own + steps[0] - steps[1]) dec_piece = ( self._measured_decode_step_ms(num_decode, decode_ctx) * spec if num_decode > 0 else 0.0 ) - return (prefill_piece + dec_piece) * (1.0 + penalty) + ov + pp = ( + self._restore_pp_ms(max(1, num_decode), chunk_tokens) + if self._meas_whole.get("decode") or self._meas_whole.get("prefill") + else 0.0 + ) + return (prefill_piece + dec_piece) * (1.0 + penalty) + ov + pp prefill_piece = ( - self._forward_times(1, chunk_tokens, "prefill", max(1, prefill_kv_len)).total_ms + self._forward_times(pb, chunk_tokens, "prefill", max(1, prefill_kv_len)).total_ms * self._prefill_rate_scale() ) dec_piece = ( @@ -3055,7 +3512,6 @@ def _continuous_decode_metrics( # Same anchor semantics as ``_decode_step_latency_ms``: a # speculation-harvested decode number is per output token, so the # step scales by the tokens a step emits, not by the verify width. - spec = self._measured_verify_step_scale() draft_billed_by_anchor = self._bench_spec_k > 0 n_samples = min(8, max(2, int(OSL))) # Benchmark mode prices prefill from measured kernels, so the anchor @@ -3066,9 +3522,11 @@ def _continuous_decode_metrics( ctx = int(ISL + frac * OSL) d_pure = self._measured_decode_step_ms(C, ctx) draft = 0.0 if draft_billed_by_anchor else self._draft_overhead_ms(d_pure) - pure.append(d_pure * spec + draft + ov) + pure.append(d_pure * self._measured_verify_step_scale(C, ctx) + draft + ov) prefill_piece = self._measured_prefill_tokens_ms(chunk_tokens) - dec_piece = self._measured_decode_step_ms(max(1, C - 1), ctx) * spec + dec_piece = self._measured_decode_step_ms( + max(1, C - 1), ctx + ) * self._measured_verify_step_scale(max(1, C - 1), ctx) mixed.append((prefill_piece + dec_piece) * (1.0 + penalty) + ov) pf_only.append(prefill_piece * (1.0 + penalty)) t_pure = sum(pure) / len(pure) @@ -3445,6 +3903,96 @@ def _sustainable_concurrency(self) -> tuple[int | None, float, str]: except Exception: return None, hbm_gb, source + def _admission_ceiling(self) -> int: + """How many requests of the target workload the KV pool can admit. + + Differs from ``_sustainable_concurrency`` in the context it sizes a + resident request at. That one asks "how many sequences fit *at the + context ceiling*", which is the right question for whether a config is + legal and the wrong one for how many requests are resident under load: + on an agentic replay whose prompts run from 400 tokens to 950k, sizing + residency at the longest possible request puts the bound around 3-7 + where the hardware admits tens. This sizes it at what a request + actually holds -- see ``resolved_resident_tokens``. + + The pool itself is taken from ``kv_pool_tokens`` when the caller knows + what the engine allocated, and otherwise from the memory model. + """ + req = self.cfg.request_config + # If the engine told us how big its pool is, believe it rather than + # reconstructing it. The bound is a ratio of pool to residency, so this + # is the numerator measured instead of predicted. + # + # A KV offload tier deliberately does *not* raise this. The connectors + # these runs use (vLLM's SimpleCPUOffloadConnector, LMCache) put evicted + # blocks in host DRAM so a later request can re-read them; a request + # being decoded still needs its blocks in HBM. Offload buys cache + # capacity, not admission capacity. + if req.kv_pool_tokens: + resident = max(1, req.resolved_resident_tokens()) + return max(1, int(req.kv_pool_tokens) // resident) + + hbm_gb, _ = self._resolve_hbm_gb() + try: + from .kv_cache import max_concurrent_sequences + from .memory import project_inference_memory + + mem = project_inference_memory(self.cfg, hbm_capacity_gb=hbm_gb, verbose=False) + fraction = req.kv_cache_memory_fraction + usable = mem.hbm_capacity_bytes or 0 + if fraction: + usable = int(usable * float(fraction)) + free_for_kv = usable - mem.weight_bytes - mem.activation_bytes + free_for_kv += req.kv_offload_gb_per_gpu * (1024.0**3) + return max_concurrent_sequences( + self.cfg, + mem.layers_on_rank, + free_for_kv, + context_len=req.resolved_resident_tokens(), + ) + except Exception: + return 0 + + def _admission_wait_ms( + self, request_latency_ms: float, per_replica: int = 1 + ) -> tuple[float, int]: + """How long a request waits to be admitted, and the bound it waits on. + + Below the bound this is zero: the pool holds every concurrent request, + the closed-loop prefill queue already prices the contention, and TTFT is + a service time. Above it the engine cannot hold all ``C`` clients, so + only ``C_crit`` are resident and the rest wait outside for a resident + request to *retire* -- not for a prefill slot to open. What they wait on + is therefore a whole request latency, prefill plus the entire decode + span, which is why measured TTFT past the knee grows faster than + linearly in concurrency: across the published AgentX rows the median + local slope ``d(log TTFT)/d(log C)`` climbs 0.18, 0.57, 1.35, 2.33 + through KV utilization bands of 0.60, 0.90, 0.99 and above. + + Throughput past the bound is pinned at ``C_crit / latency``, so by + Little's law a request's sojourn is ``C * latency / C_crit`` and the + admission wait is the excess over the latency it would have seen + inside. Nothing here is fitted. + + Scored against 176 published AgentX measurements -- 6 models on 16 + model/stack curves -- this takes median per-curve TTFT rank correlation + from 0.44 to 0.76 and the regret of picking on it from 3% to 0%, and it + beats the discrete-event replay at 0.62. It does not cost accuracy in + either direction: median absolute error is unchanged at ~0.8, because + what the term fixes is the shape of TTFT in concurrency rather than its + level, and the term is identically zero wherever the pool does not + bind, which includes every fixed-sequence workpoint we run. + """ + ceiling = self._admission_ceiling() + requested = self.cfg.request_config.resolved_max_concurrency() + # ``per_replica`` spreads the clients over a pool of replicas, so what + # competes for one replica's pool is its share of them. + if per_replica > 1: + requested = max(1, int(math.ceil(requested / per_replica))) + if ceiling <= 0 or requested <= ceiling: + return 0.0, ceiling + return max(0.0, request_latency_ms) * (requested / ceiling - 1.0), ceiling + def _effective_concurrency(self) -> dict: """Concurrency that drives throughput, reconciled against the KV-feasible ceiling (cap + report): @@ -3592,6 +4140,22 @@ def _project_colocated(self) -> InferencePerfResult: per_req_decode_tps = (1000.0 / itl) if itl > 0 else 0.0 request_latency = ttft + decode_total + + # Admission. Everything above prices the engine's own queue, which + # assumes every concurrent request is resident; the KV pool decides + # whether it can be. Where the pool binds, the excess clients wait + # outside for a resident request to retire, and that wait is the + # dominant term in TTFT -- it is not a correction to it. + admission_ms, admission_ceiling = self._admission_wait_ms(request_latency) + if admission_ms > 0.0: + ttft += admission_ms + request_latency += admission_ms + extras["admission_ceiling"] = float(admission_ceiling) + extras["admission_wait_ms"] = float(admission_ms) + extras["resident_tokens_per_request"] = float( + self.cfg.request_config.resolved_resident_tokens() + ) + decode_tps_per_gpu = decode_tps / replica_gpus if replica_gpus else 0.0 prefill_tps = (batch * input_len * 1000.0 / prefill_full_ms) if prefill_full_ms > 0 else 0.0 @@ -3804,6 +4368,32 @@ def _project_disaggregated(self) -> InferencePerfResult: if output_len > 1: itl = decode_total / (output_len - 1) request_latency = ttft + decode_total + + # Admission, which this path did not price at all. The KV pool binds on + # the decode side -- that is where a sequence's blocks live for the + # whole generation, and the memory pass already re-binds to that pool -- + # so a split is no less capable of over-subscribing it than a colocated + # engine is. Charged per decode replica, because the clients are spread + # across the pool and it is one replica's residency that binds. + # + # On the published disaggregated DeepSeek rows this term is identically + # zero, and that is the finding rather than a failure of it: moving + # prefill off the decode GPUs leaves a pool that admits ~770 sequences + # against the 256 ever offered, so the split never reaches its knee. + # Their TTFT still grows 9x over that range, which means for a split + # the binding queue is the *prefill* pool's and not the KV pool's -- + # the closed-form knee ranks those seven rows at rho 0.286 where the + # replay, which simulates that queue, gets 0.71-0.96. Price a split's + # TTFT from the replay; this term is here for the case where the decode + # pool is genuinely over-subscribed, which one replica with a small + # pool still can be. + admission_ms, admission_ceiling = self._admission_wait_ms( + request_latency, per_replica=max(1, int(disagg.decode_replicas or 1)) + ) + if admission_ms > 0.0: + ttft += admission_ms + request_latency += admission_ms + per_req_decode_tps = (1000.0 / itl) if itl > 0 else 0.0 # Per-replica decode throughput, scaled by the decode-pool replica count. @@ -3854,6 +4444,9 @@ def _project_disaggregated(self) -> InferencePerfResult: if self.is_benchmark_calibrated: extras["benchmark_calibrated"] = 1.0 extras["prefill_compute_ttft_ms"] = ttft_compute + extras["admission_ceiling"] = float(admission_ceiling) + extras["admission_wait_ms"] = float(admission_ms) + extras["resident_tokens_per_request"] = float(req.resolved_resident_tokens()) extras["prefill_replicas"] = float(disagg.prefill_replicas) extras["decode_replicas"] = float(disagg.decode_replicas) extras["prefill_occupancy"] = float(n_prefill) diff --git a/infera/projection/core/projection/inference_projection/search/anchor_store.py b/infera/projection/core/projection/inference_projection/search/anchor_store.py index da5d495d2..a3c62fd7b 100644 --- a/infera/projection/core/projection/inference_projection/search/anchor_store.py +++ b/infera/projection/core/projection/inference_projection/search/anchor_store.py @@ -184,8 +184,16 @@ def add_result(self, result: dict[str, Any], artifact_path: str) -> dict[str, An def _make_entry(self, path: str, art: dict[str, Any]) -> dict[str, Any]: meta = art.get("meta", {}) - recipe = regime.recipe_from_meta(meta) - sig = meta.get("regime_signature") or regime.regime_signature(recipe) + recipe = regime.recipe_from_meta(meta, engine=art.get("backend")) + # Recomputed, not read back from ``meta``. A recorded signature was + # hashed by whichever build harvested the artifact, over whichever axes + # that build had, so trusting it pins the store to an older definition + # of "same regime" -- and the axis set does grow. Adding ``engine`` + # found seven DeepSeek-V4-Flash anchors sharing one recorded signature + # across SGLang and vLLM: re-indexing could not separate them while the + # stale hash was preferred, which is the merge the axis exists to stop. + # The recorded value stays in ``meta`` as provenance. + sig = regime.regime_signature(recipe) sweep = art.get("sweep") or [] batches = sorted({int(e["batch"]) for e in sweep if e.get("batch") is not None}) return { diff --git a/infera/projection/core/projection/inference_projection/search/regime.py b/infera/projection/core/projection/inference_projection/search/regime.py index dfbccbcde..2366d5304 100644 --- a/infera/projection/core/projection/inference_projection/search/regime.py +++ b/infera/projection/core/projection/inference_projection/search/regime.py @@ -47,6 +47,16 @@ # target that speculates needs its own measurement. REGIME_AXES = ( "model", + # The serving engine. Named as regime-defining in the split above and in + # the design note, but absent from this tuple until now, so vLLM, SGLang + # and Atom anchors for one checkpoint all hashed to one signature: the + # bench cache could return a vLLM measurement for an SGLang run, and + # ``regime_distance`` called the two a perfect match. They are not one + # regime. The engines differ in scheduler, paging and kernel selection -- + # on this store the same DeepSeek-V4 recipe under `mori-sglang` and + # `sglang` reports KV pools 16x apart, because one shards the MLA latent + # across ranks and the other replicates it. + "engine", "weight_dtype", "kv_cache_dtype", "moe_expert_dtype", @@ -55,6 +65,25 @@ "aiter", "aiter_ops", "speculative", + # The accelerator the anchor was measured on. Absent from this tuple until + # now, and absent from the anchors themselves, which is the more serious + # half: nothing in the store records which GPU produced a timing, so the + # restore cannot form the sim(target)/sim(bench) ratio that would carry + # the change, and does not try. A measurement is simply reused. On this + # store every anchor is gfx950 -- the GLM-5.2-MXFP4 Quark build only + # loads there -- and one target, GLM-5.2 on MI325X, is gfx942, so it was + # priced on kernels from a part with roughly twice the compute and a + # third more bandwidth. It read 4.3x the measured throughput and a + # twentieth of the measured TTFT, the worst curve in the matrix, while + # reporting itself calibrated. + # + # Regime-defining rather than transportable, which is the conservative + # of the two readings. Hardware is transportable in principle -- the + # analytical model knows both parts -- but only once the anchor says + # which part it ran on, and until a harvest records that, refusing is + # the honest answer. ``IX_RELAX_AXES`` can still relax it deliberately, + # which is the difference between a known approximation and a silent one. + "gpu_arch", ) # Canonical value of the ``speculative`` axis when speculation is off. Distinct @@ -93,6 +122,27 @@ ) +# Engine spellings of an attention backend, folded onto the kernel family the +# axis is actually about. AITER ships its MLA, unified-attention and +# DeepSeek-V4 backends under separate flag values, and an engine records +# whichever one its command line named -- so an anchor harvested with +# ``--attention-backend dsv4`` and a target the projector describes as +# ``aiter`` are the same kernels under two names. Compared raw, that pair is a +# regime mismatch that no harvest can clear: the anchor is refused, the +# projection silently falls back to analytical, and the axis rejects the very +# measurement it asked for. The projector has always modelled these as one +# family (ix_recipe.ATTENTION_BACKEND maps all three onto "aiter"); this is +# where the store learns the same thing. +_ATTENTION_FAMILY = { + "dsv4": "aiter", + "aiter": "aiter", + "rocm_aiter_mla": "aiter", + "rocm_aiter_unified_attn": "aiter", + "triton_attn": "triton", + "triton": "triton", +} + + def _canon(v: Any) -> str: """Canonical, comparison-stable string for a single axis value.""" if v is None: @@ -109,6 +159,19 @@ def _canon(v: Any) -> str: return str(v).strip().lower() +def _canon_axis(k: str, v: Any) -> str: + """Canonical value for axis ``k``, folding engine spellings onto families. + + Kept separate from ``_canon`` because the folding is per-axis: only the + attention backend has several names for one kernel family, and applying a + rename table to every axis would let unrelated values collide. + """ + s = _canon(v) + if k == "attention_backend" and s: + return _ATTENTION_FAMILY.get(s, s) + return s + + def normalise_model_id(name: Any) -> str: """Reduce a model name to a form comparable across its spellings. @@ -144,7 +207,7 @@ def leaf(v: Any) -> str: def _sig(recipe: dict[str, Any], axes: Iterable[str]) -> str: - payload = {k: _canon(recipe.get(k)) for k in axes} + payload = {k: _canon_axis(k, recipe.get(k)) for k in axes} blob = json.dumps(payload, sort_keys=True) return hashlib.sha256(blob.encode()).hexdigest()[:16] @@ -159,7 +222,7 @@ def config_key(recipe: dict[str, Any], extra: dict[str, Any] | None = None) -> s """Exact-run identity: hash over regime + transport axes, plus any ``extra`` (measurement knobs that change the number but not the regime, e.g. decode-steps). Used by the benchmark result cache.""" - payload = {k: _canon(recipe.get(k)) for k in (*REGIME_AXES, *TRANSPORT_AXES)} + payload = {k: _canon_axis(k, recipe.get(k)) for k in (*REGIME_AXES, *TRANSPORT_AXES)} if extra: for k, v in extra.items(): payload[f"x_{k}"] = _canon(v) @@ -202,7 +265,7 @@ def regime_distance(a: dict[str, Any], b: dict[str, Any], *, ignore_missing: boo continue if ignore_missing and missing: continue - if _canon(av) != _canon(bv): + if _canon_axis(k, av) != _canon_axis(k, bv): d += 1 return d @@ -258,7 +321,34 @@ def aiter_ops_axis(env: dict[str, str] | None) -> str | None: # -------------------------------------------------------------------------- -def recipe_from_meta(meta: dict[str, Any], *, model: str | None = None) -> dict[str, Any]: +#: Names a *quantizer toolchain* rather than a numeric format. Both vLLM and +#: SGLang resolve ``quantization_config.quant_method`` to one of these and the +#: benchmark records it verbatim, so an anchor measured on an MXFP4 checkpoint +#: comes back labelled ``quark``. Quark emits fp8 and fp4 alike, so the label +#: does not determine the dtype and cannot be rewritten to one here. +_QUANTIZER_NAMES = frozenset( + {"quark", "compressed-tensors", "compressed_tensors", "modelopt", "torchao"} +) + + +def _weight_dtype_or_unknown(quant: Any) -> Any: + """Drop quantizer-toolchain names, which are not weight dtypes. + + Returning None follows the policy the caller already documents: an unknown + axis is skipped by ``regime_distance`` while a wrong one forces a mismatch. + ``quark`` is a wrong one. Every GLM-5.2-MXFP4 anchor in the store was + measured on MI355X against the same MXFP4 checkpoint the matrix projects, + and all of them were being refused on ``weight_dtype('mxfp4'->'quark')`` -- + a naming artifact, not a difference in what ran. + """ + if quant is None: + return None + return None if str(quant).strip().lower() in _QUANTIZER_NAMES else quant + + +def recipe_from_meta( + meta: dict[str, Any], *, model: str | None = None, engine: str | None = None +) -> dict[str, Any]: """Canonical recipe from a benchmark artifact's ``meta`` block. The *benchmark* parallelism (what it actually ran at) is recorded on the transport axes so the anchor's coverage is described in benchmark space; @@ -275,9 +365,13 @@ def recipe_from_meta(meta: dict[str, Any], *, model: str | None = None) -> dict[ # requested quantization; otherwise leave it unknown rather than assert a # value, since an unknown axis is skipped by ``regime_distance`` while a # wrong one forces a mismatch. - quant = meta.get("weight_dtype") or meta.get("quantization") + quant = _weight_dtype_or_unknown(meta.get("weight_dtype") or meta.get("quantization")) return { "model": model or meta.get("model"), + # The engine lives on the artifact rather than in ``meta``, so the + # caller that opened the artifact passes it; the ``meta`` keys are read + # as a fallback for harnesses that record it inline. + "engine": engine or meta.get("engine") or meta.get("backend"), "weight_dtype": quant if quant else None, "kv_cache_dtype": meta.get("kv_cache_dtype") or "bf16", "moe_expert_dtype": meta.get("moe_expert_dtype"), @@ -295,6 +389,13 @@ def recipe_from_meta(meta: dict[str, Any], *, model: str | None = None) -> dict[ if "speculative_method" in meta else None ), + # Absent key => unknown (pre-tracking artifact), which + # ``regime_distance`` skips. That is the permissive reading and it is + # deliberate: the alternative refuses every anchor harvested before + # the axis existed. What makes it safe is that a store can be told + # which part it ran on after the fact -- see ``ix_anchor_store`` -- + # so an unknown here means "nobody has said yet", not "any part". + "gpu_arch": (str(meta.get("gpu_arch") or meta.get("gpu") or "").lower().strip() or None), # transport (benchmark space) "tp": meta.get("benchmark_tp") or meta.get("tp"), "pp": meta.get("benchmark_pp") or meta.get("pp"), @@ -322,6 +423,9 @@ def recipe_from_bench_args(args: Any, env: dict[str, str] | None = None) -> dict ep = int(getattr(args, "tp", 1) or 1) if getattr(args, "enable_expert_parallel", False) else 1 return { "model": getattr(args, "model", None), + # ``benchmark_serving`` names the engine it drove; the in-container + # vLLM harness has no such flag because it is only ever vLLM. + "engine": getattr(args, "serving_backend", None) or "vllm", "weight_dtype": getattr(args, "quantization", None) or "bf16", "kv_cache_dtype": getattr(args, "kv_cache_dtype", None) or "bf16", "moe_expert_dtype": None, @@ -337,6 +441,7 @@ def recipe_from_bench_args(args: Any, env: dict[str, str] | None = None) -> dict getattr(args, "speculative_method", None) or "", getattr(args, "speculative_num_tokens", None), ), + "gpu_arch": (str(getattr(args, "gpu_arch", "") or "").lower().strip() or None), "tp": getattr(args, "tp", 1), "pp": getattr(args, "pp", 1), "ep": ep, @@ -348,7 +453,7 @@ def recipe_from_bench_args(args: Any, env: dict[str, str] | None = None) -> dict } -def recipe_from_inference_config(cfg: Any) -> dict[str, Any]: +def recipe_from_inference_config(cfg: Any, gpu_arch: Any = None) -> dict[str, Any]: """Canonical recipe from an ``InferenceConfig`` (the reconstruction target). Structural configs carry no HF model *name*, so ``model`` is left ``None`` @@ -366,6 +471,12 @@ def g(o, name, default=None): ep = int(g(mp, "expert_model_parallel_size", 1) or 1) return { "model": None, + # Unknown unless the target says so. A projection has no engine of its + # own -- it is asking what some engine would do -- so the recipe has to + # state which one, and ``--serving-engine`` is how. Left unset the axis + # is skipped, which keeps every existing caller matching the anchors it + # matched before. + "engine": g(req, "serving_engine"), "weight_dtype": g(req, "weight_dtype", "bf16"), "kv_cache_dtype": g(req, "kv_cache_dtype", "bf16"), "moe_expert_dtype": g(req, "moe_expert_dtype"), @@ -378,6 +489,11 @@ def g(o, name, default=None): "spec" if g(req, "speculative_num_tokens") else "", g(req, "speculative_num_tokens"), ), + # The target always knows its part; the caller has to pass it, + # because it arrives on the command line rather than on the config. + # Left None the axis is skipped, which keeps a caller that does not + # know its hardware matching the anchors it matched before. + "gpu_arch": str(gpu_arch).lower().strip() if gpu_arch else None, "tp": tp, "pp": int(g(mp, "pipeline_model_parallel_size", 1) or 1), "ep": ep, diff --git a/infera/projection/core/projection/training_config.py b/infera/projection/core/projection/training_config.py index 0a53aa7bc..e0a6d4d1d 100644 --- a/infera/projection/core/projection/training_config.py +++ b/infera/projection/core/projection/training_config.py @@ -28,6 +28,11 @@ class ModelParallelConfig: # tensor/expert parallel. Standard for MLA models, whose latent KV is # replicated by tensor parallelism rather than sharded by it. attention_data_parallel_size: int = 1 + # Serving-only axis: one sequence's *context* is split across this many + # ranks during decode, so a rank stores a slice of each request's KV + # instead of all of it. The only axis that shrinks an MLA cache, since + # tensor parallelism replicates the latent rather than sharding it. + decode_context_parallel_size: int = 1 use_torch_fsdp2: bool = False use_distributed_optimizer: bool = False overlap_grad_reduce: bool = True @@ -79,6 +84,21 @@ class ModelConfig: # full-attention), i.e. half the layers. sink_sliding_window: int = 0 sink_window_even_layers_only: bool = False + # Hybrid compressed attention (DeepSeek-V4). ``compress_ratios`` is a + # per-layer schedule naming which attention branch each layer runs, using + # the same 0 / 4 / 128 encoding the reference implementation accepts (see + # ``_COMPRESS_RATIO_LAYER_TYPES``). It is a list or the ``"[...]"`` string + # form the checkpoint ships. The remaining fields size the branches: a + # compressed layer attends over its local window plus a pool of one entry + # per ``compress_rate`` source tokens, and a CSA layer additionally runs a + # lightning indexer (``index_n_heads`` x ``index_head_dim``) over that whole + # pool to pick ``index_topk`` of it. + hybrid_attention_enabled: bool = False + compress_ratios: object = None + index_topk: int = 0 + index_n_heads: int = 0 + index_head_dim: int = 0 + attn_sliding_window: int = 0 # Hybrid linear attention (Kimi Delta Attention / gated delta net): a # count of layers that keep a fixed-size recurrent state instead of a # per-token KV cache. ``linear_attention_layers`` is the explicit count @@ -91,6 +111,37 @@ class ModelConfig: linear_attention_head_dim: int = 0 linear_attention_conv_kernel: int = 0 linear_attention_freq: int = 0 + # Native sparse attention's indexer keeps a *second* per-token cache. To + # choose which blocks to attend to it scores the query against one small + # index key per past token, and that key has to be stored, so a token's + # real footprint is the KV cache plus an index cache. + # + # The key is head-shared, like an MLA latent: ``index_n_heads`` in the + # published configs counts the *query* heads that score against it, not + # copies of it, so it does not multiply the storage and is not recorded + # here. What sets the size is the key width (``index_head_dim``, 128 on + # every model that has one), how many layers carry an indexer, and the + # dtype the serving stack keeps it in. ``sparse_index_layers`` is 0 for + # "all layers"; models that leave the first few layers dense, or that + # share one layer's index keys with the layers after it, pay for fewer. + # + # Leaving this out is not a rounding error. On MiniMax-M3 it is a third of + # the footprint: GQA alone predicts 15,360 bytes per token at TP8 with an + # fp8 cache where vLLM allocates 30,023. Because the index cache does not + # follow ``kv_cache_dtype``, the two terms can be read apart from two runs + # of the same model -- with a bf16 cache the same engine allocates 45,398, + # and both are matched to under 0.5% once the indexer is charged. + sparse_index_head_dim: int = 0 + sparse_index_layers: int = 0 + sparse_index_dtype: str = "bf16" + # The query-head count above, recorded for *time* rather than memory. It + # does not size the cache -- one key per token serves every head -- but it + # does set the indexer's arithmetic, because each head scores the whole + # pool. Published values differ by a factor of sixteen (64 on + # DeepSeek-V4-Pro, 32 on GLM-5.2, 4 on MiniMax-M3), so a model priced + # without it is not merely missing a constant, it is missing the term that + # separates these three. + sparse_index_n_heads: int = 0 # FFN & MoE swiglu: bool = False num_experts: int = 0 @@ -178,6 +229,16 @@ def linear_attention_state_len(self) -> int: d = int(self.linear_attention_head_dim or 0) or int(self.kv_channels or 0) return max(1, d) + def sparse_index_layer_count(self) -> int: + """How many layers store a sparse-attention index cache.""" + if self.sparse_index_head_dim <= 0: + return 0 + n = int(self.num_layers or 0) + declared = int(self.sparse_index_layers or 0) + if declared > 0: + return min(declared, n) if n else declared + return n + def full_attention_layer_fraction(self) -> float: """Fraction of layers that store and read a per-token KV cache.""" n_lin = self.linear_attention_layer_count() @@ -382,6 +443,46 @@ class InferenceRequestConfig: # the offload is cheap enough to leave on. kv_offload_bw_gbps: float = 64.0 + # Mean KV tokens one request of the target workload occupies while it is + # resident, which is what decides how many requests the pool can admit at + # once. ``None`` = use the configured context, which is exact for a + # fixed-length workpoint and is the only thing a single-point projection + # can know. + # + # For a workload with a spread of lengths this is E[L^2]/E[L] and *not* the + # arithmetic mean, because a request holds the pool for a time proportional + # to its length: the set resident at a random instant is sampled + # proportional to length, so a time-average occupancy is length-biased. + # Against the published AgentX rows -- 572 of them, 7 models, 9 hardware + # platforms -- predicting occupancy as ``concurrency * this / pool`` matches + # the reported ``gpu_kv_cache_usage_pct`` with a slope of 0.91 pooled and a + # median of 0.90 across 48 individual curves, against 1.14 and 1.19 for the + # arithmetic mean. A statistic that needs no multiplier has slope 1, so this + # is the one the model uses and it carries no fitted constant. + # + # It must NOT be discounted by the prefix-cache hit rate. A 95% block-reuse + # rate means 95% of the prefill FLOPs are skipped; it does not mean 95% of + # the blocks are free. Cached blocks live in the same pool as active ones + # and are held by reference from every request reading them -- used blocks + # are total minus free, and a shared prefix is counted once but counted. + # Discounting residency that way needs a 6.1x multiplier to match measured + # occupancy (interquartile range 2.0 to 13.4) and predicts it worse than + # assuming the mean. + workload_resident_tokens: int | None = None + + # The KV pool the engine actually allocated, in tokens, when that is known. + # Every engine prints it at startup, so anyone validating a projection + # against a running deployment -- or against a published benchmark that + # reports it -- has the real number and should not have to get it back out + # of an HBM figure, a memory fraction and a weight estimate. + # + # When it is set the admission bound is evaluated against it directly. That + # separates the two things that can make the bound wrong, which otherwise + # arrive together: the admission law itself, and the memory model's answer + # to how many tokens fit. It is left unset for planning, where the whole + # point is that no engine has been started yet. + kv_pool_tokens: int | None = None + # ---- Precision ---- weight_dtype: str = "bf16" # weights kept resident (bf16 | fp8 | ...) kv_cache_dtype: str = "bf16" # KV cache precision (bf16 | fp8 | int8 | ...) @@ -519,6 +620,17 @@ class InferenceRequestConfig: # invent one. Latency-only, like the two terms above -- it is host work that # overlaps the GPU, so no throughput moves. request_overhead_ms: float = 0.0 + # Per-token TTFT latency for the part of a prompt the prefix cache did not + # serve (microseconds/token). On the AgentX MI355X rows TTFT climbs + # ~40-130 us for every token a request has to prefill, while ITL on the + # same stacks at C>=4 shows the GPU is not held for anything like that + # long, so it is not charged as GPU time. Added to TTFT and end-to-end + # latency only (not throughput), like the terms above. Default 0. + uncached_prompt_latency_us: float = 0.0 + # Tokens past which the term above stops growing (0 = no cap). The same rows + # put a ceiling on it: at C=4 a run whose trace holds several 300-460k-token + # cold prompts still has a TTFT spread of ~0.6 s. + uncached_prompt_latency_max_tokens: int = 0 # The *other* half of that same fit: the per-token slope. Differencing TTFT # across two prompt lengths at concurrency 1 yields ``TTFT = C + r*L``; ``C`` # is the constant above and ``r`` is this -- a measured prefill rate, in @@ -583,6 +695,11 @@ class InferenceRequestConfig: # the number of prefill steps (and the mixed-step fraction) → higher TPOT / # lower throughput. ``0`` = unlimited (legacy behaviour). max_num_batched_tokens: int = 0 + # Engine cap on concurrently running sequences (vLLM ``--max-num-seqs``, + # SGLang ``--max-running-requests``). Clients beyond it wait in the + # server queue, so it bounds the decode batch and adds queue time to TTFT. + # ``0`` = no cap beyond the client concurrency. + max_num_seqs: int = 0 # MoE expert routing imbalance. On an EP-sharded model the MoE step time is # set by the BUSIEST rank, not the average: ``ep_load_balance`` is the ratio # of the hottest rank's token load to the average (1.0 = perfectly @@ -614,6 +731,20 @@ class InferenceRequestConfig: # baseline, 1.0). Values: aiter | triton | ck | hip. attention_backend: str | None = None + # ---- Serving engine ---- + # Which engine this projection is asking about: vllm | sglang | atom, or a + # build of one (`mori-sglang`). Distinct from ``attention_backend``, which + # names a kernel library inside an engine. + # + # Nothing in the analytical model reads this -- simulate mode prices an + # architecture on a device and returns the same number whichever engine is + # named. It is a regime axis, so what it decides is which measured anchor + # benchmark mode is allowed to calibrate from, and there the engine is the + # whole question: an SGLang anchor does not describe vLLM's scheduler or + # its paging. Left unset the axis is unknown and matching behaves as it did + # before, which is why it is optional rather than required. + serving_engine: str | None = None + # ---- Native sparse attention (DeepSeek V3.2 / V4 NSA) ---- # Number of KV tokens each query attends to under native sparse attention # (NSA + indexer top-k selection). ``0`` = dense attention (legacy). When @@ -622,6 +753,17 @@ class InferenceRequestConfig: # attention roughly constant in context length. sparse_attention_topk: int = 0 + # What the indexer's selection costs on this stack relative to a fused + # kernel. The arithmetic is fixed by the model, but the kernel is not: + # the same GLM-5.2 top-k runs fused through aiter on gfx950 and unfused + # through Torch on gfx942, and the selection is parallelism-starved at + # the batch sizes a latency-band run uses, so the fallback costs a + # multiple rather than a margin. 1.0 prices the fused path; raise it to + # charge a stack serving without one. Only reaches models priced by + # ``uniform_sparse_attention_scale`` -- a schedule-priced model carries + # its indexer per layer and is not governed by one number. + sparse_indexer_cost_scale: float = 1.0 + # ---- Sliding-window / local attention ---- # Override the model's sliding-window size (KV tokens each query attends to; # 0 = full attention). ``None`` follows the model config's @@ -693,6 +835,18 @@ def resolved_max_concurrency(self) -> int: return int(self.max_concurrency) return int(self.batch_size) + def resolved_resident_tokens(self) -> int: + """KV tokens one request holds while resident. + + Falls back to the configured context (prompt plus half the generation, + the mean over a request's decode span), which is what a fixed-length + workpoint occupies. A trace-driven planner passes the workload's + ``E[L^2]/E[L]`` instead; see ``workload_resident_tokens``. + """ + if self.workload_resident_tokens: + return max(1, int(self.workload_resident_tokens)) + return max(1, int(self.input_seq_len) + int(self.output_seq_len) // 2) + def resolved_prefix_cache_hit_rate(self) -> float: """Prefix-cache hit rate clamped to ``[0, 0.999]``. @@ -906,8 +1060,265 @@ def resolved_act_quant_dtype(self, model_fp8=None) -> str | None: } # Floor on the sparse-attention scale (projections + indexer don't shrink). +# Only reached by models that declare a top-k without a per-layer schedule; +# ``hybrid_attention_scale`` supersedes it where a schedule exists, and states +# in one place why a floor is the wrong shape for the work it stands in for. _SPARSE_ATTENTION_FLOOR = 0.15 +# Per-layer attention branch, keyed by the compression ratio the checkpoint +# records for that layer. Mirrors the reference implementation's own mapping so +# a schedule read here means what it means to the model; an unrecognised ratio +# is a schedule this cost model has not been taught, and is refused rather than +# guessed at (see ``hybrid_attention_scale``). +_COMPRESS_RATIO_LAYER_TYPES = { + 0: "sliding_attention", + 4: "compressed_sparse_attention", + 128: "heavily_compressed_attention", +} + +# Schedules whose length has already been reported as unexplained by the layer +# count, so the warning is said once per shape rather than once per forward. +_SCHEDULE_LENGTH_WARNED: set = set() + + +def parse_compress_ratios(model_config) -> list | None: + """Per-layer compression schedule as a list of ints, or ``None``. + + Accepts the list form and the ``"[128, 128, 4, ...]"`` string form that the + checkpoint ships. Returns ``None`` when the model declares no schedule. + """ + raw = getattr(model_config, "compress_ratios", None) + if raw is None: + return None + if isinstance(raw, str): + raw = raw.strip().strip("[]") + if not raw: + return None + try: + return [int(float(x)) for x in raw.split(",") if x.strip()] + except ValueError: + return None + try: + return [int(x) for x in raw] + except (TypeError, ValueError): + return None + + +def hybrid_attention_scale(model_config, context_len: int) -> float | None: + """Fraction of dense attention cost under a per-layer compression schedule. + + ``None`` when the model declares no schedule, or declares one this mapping + does not recognise, so the caller keeps its previous behaviour. + + A compressed layer does not attend over its context. It attends over a local + window plus a pool holding one entry per ``compress_rate`` source tokens, so + what it reads is set by the schedule and the window, not by the prompt. The + two branches differ in how they use that pool: HCA has full visibility over + it, while CSA runs a lightning indexer across the whole pool to select + ``index_topk`` entries and attends to those. + + That selection is why a single top-k scale reads this model wrong at long + context. CSA's *attention* stops growing once the pool exceeds the top-k + (above 4k context for DeepSeek-V4-Pro, which selects 1024 of ``context/4``), + but its *indexer* still has to score every entry in the pool, so the term + that keeps growing with context is the indexer rather than the attention. + At 8k context the indexer is 9% of a CSA layer and at 130k it is 62%. + + A multiplicative floor cannot express that. The work it stands in for is + real -- the indexer genuinely does not shrink with the window -- but it + grows like ``context / compress_rate`` scaled by the indexer's much smaller + head geometry, whereas flooring the scale makes it grow like dense attention + over the full context. On DeepSeek-V4-Pro at 130k the schedule costs 0.0157 + of dense against the 0.15 floor, a 9.6x over-read, while at the 8k the floor + was set against it is only 1.7x and does not show. + """ + sched = parse_compress_ratios(model_config) + if not sched: + return None + if any(cr not in _COMPRESS_RATIO_LAYER_TYPES for cr in sched): + return None + ctx = int(context_len or 0) + if ctx <= 0: + return None + + # The schedule covers the MTP head as well as the stack: DeepSeek-V4-Pro + # declares 61 layers and schedules 62, Flash declares 43 and schedules 44, + # and both carry ``mtp_num_layers: 1``. In both the surplus entry is the + # trailing 0, a dense windowed layer, which is what an MTP head is. Only the + # stack runs in the forward being costed here, so the average is taken over + # that and the speculative head is left to whoever models speculation. + declared = int(getattr(model_config, "num_layers", 0) or 0) + if declared and len(sched) > declared: + mtp = int(getattr(model_config, "mtp_num_layers", 0) or 0) + if len(sched) != declared + mtp: + key = (declared, len(sched)) + if key not in _SCHEDULE_LENGTH_WARNED: + _SCHEDULE_LENGTH_WARNED.add(key) + print( + f"[inferasim:Inference] WARNING: this model's compression " + f"schedule names {len(sched)} layers against {declared} in the " + f"stack and {mtp} predicting ahead of it. The first {declared} " + f"entries are used; the rest are unexplained. Re-transcribe " + f"compress_ratios from the checkpoint's config.json." + ) + sched = sched[:declared] + elif declared and len(sched) < declared: + key = (declared, len(sched)) + if key not in _SCHEDULE_LENGTH_WARNED: + _SCHEDULE_LENGTH_WARNED.add(key) + print( + f"[inferasim:Inference] WARNING: this model's compression schedule " + f"names only {len(sched)} of {declared} layers. The branch mix is " + f"weighted over the entries present, which assumes the rest look " + f"like them. Re-transcribe compress_ratios from the checkpoint." + ) + + heads = int(getattr(model_config, "num_attention_heads", 0) or 0) + latent = int(getattr(model_config, "kv_channels", 0) or 0) + rope = int(getattr(model_config, "qk_pos_emb_head_dim", 0) or 0) + # Absorbed MLA reads the shared latent plus the positional part for the + # scores and the latent again for the values, per KV entry. + main_per_entry = heads * ((latent + rope) + latent) + if main_per_entry <= 0: + return None + idx_per_entry = int(getattr(model_config, "index_n_heads", 0) or 0) * int( + getattr(model_config, "index_head_dim", 0) or 0 + ) + window = int(getattr(model_config, "attn_sliding_window", 0) or 0) + topk = int(getattr(model_config, "index_topk", 0) or 0) + + local = min(ctx, window) if window > 0 else ctx + total = 0.0 + for cr in sched: + kind = _COMPRESS_RATIO_LAYER_TYPES[cr] + if kind == "sliding_attention": + total += local * main_per_entry + continue + pool = ctx / cr + if kind == "heavily_compressed_attention": + # Full visibility over the pool, and no indexer to select with. + total += (local + pool) * main_per_entry + else: + selected = min(topk, pool) if topk > 0 else pool + total += (local + selected) * main_per_entry + pool * idx_per_entry + return total / (len(sched) * ctx * main_per_entry) + + +def uniform_sparse_attention_scale( + model_config, + context_len: int, + topk: int, + indexer_cost_scale: float = 1.0, +) -> float | None: + """Fraction of dense prefill attention for a uniform top-k sparse model. + + The counterpart to ``hybrid_attention_scale`` for the models that select + the same way on every layer instead of running a per-layer compression + schedule: GLM-5.2 and MiniMax-M3 both keep one window, ``index_topk``, + and differ only in how many layers compute an index and how wide the + indexer is. + + Two terms, and the reason to spend a function on this rather than a scale + is that they move in opposite directions. *Selection* is what attention + reads once the indexer has chosen -- ``topk`` entries regardless of how + long the prompt is -- so as a fraction of dense it falls away like + ``topk / context``. *Indexing* is what choosing costs: every head scores + every token in the pool, on the layers that carry an indexer, so it grows + with context exactly as dense attention does and its fraction is flat. + Past a few times the window the second term is the whole cost, and it is + the one a top-k scale cannot express. + + That is not a small correction on this workload. GLM-5.2 served the + agentic corpus between 136k and 365k tokens, where the floored top-k + charges a constant 0.15 of dense: nine to twenty-seven times the + selection it stands for, while the indexer it is standing in for is + charged nothing. Pricing the terms apart puts GLM at 0.032 of dense at + 131k falling to 0.021 at 365k -- and, unlike the floor, it moves with the + prompt, which is what a ladder whose context swings by 2.7x needs. + + ``indexer_cost_scale`` prices the kernel rather than the arithmetic. The + indexer's top-k is fused on some stacks and not others -- SGLang runs it + through aiter's one-block kernel on gfx950 but falls back to an unfused + Torch path on gfx942 -- and the selection is parallelism-starved at low + concurrency (batch 4 with 4 draft tokens is 16 workgroups on a 256-CU + MI355X), so the same arithmetic costs several times more on the fallback. + 1.0 is the fused path; raise it to charge a stack that does not have one. + + Returns ``None`` when the model is not a uniform top-k sparse model, so + the caller can fall through to whatever it did before. + """ + ctx = int(context_len or 0) + topk = int(topk or 0) + if ctx <= 0 or topk <= 0: + return None + layers = int(getattr(model_config, "num_layers", 0) or 0) + idx_width = int(getattr(model_config, "sparse_index_head_dim", 0) or 0) + idx_heads = int(getattr(model_config, "sparse_index_n_heads", 0) or 0) + if layers <= 0 or idx_width <= 0 or idx_heads <= 0: + return None + + heads = int(getattr(model_config, "num_attention_heads", 0) or 0) + if heads <= 0: + return None + if getattr(model_config, "multi_latent_attention", False): + # Absorbed MLA reads the latent plus the positional part for the + # scores and the latent again for the values, per KV entry. + latent = int(getattr(model_config, "kv_lora_rank", 0) or 0) or int( + getattr(model_config, "kv_channels", 0) or 0 + ) + rope = int(getattr(model_config, "qk_pos_emb_head_dim", 0) or 0) + main_per_token = heads * ((latent + rope) + latent) + else: + # GQA reads a head-width key and a head-width value per query head. + width = int(getattr(model_config, "kv_channels", 0) or 0) + main_per_token = heads * 2 * width + if main_per_token <= 0: + return None + + # ``sparse_index_layers`` is 0 for "every layer"; models that leave the + # first few dense, or that share one layer's index across the layers + # after it, run the indexer on fewer than they have. + indexed = int(getattr(model_config, "sparse_index_layers", 0) or 0) or layers + indexed = min(indexed, layers) + + selection = min(topk, ctx) / ctx + indexing = (indexed / layers) * (idx_heads * idx_width) / main_per_token + return selection + indexing * max(0.0, float(indexer_cost_scale)) + + +def sparse_indexer_decode_overhead( + model_config, topk: int, indexer_cost_scale: float = 1.0 +) -> float: + """Extra decode attention a stack pays for an unfused indexer selection. + + Decode is charged dense, and that charge is calibrated: withdrawing the + top-k discount is what moved DeepSeek-V4 from 0.53x of measured TPOT to + 0.83-0.94x. So the fused path must stay exactly where that calibration + left it, and only the excess over it is added here -- 0.0 at a cost scale + of 1.0, by construction. + + What the excess is proportional to is the indexer's share of the layer, + not the whole step, because the selection is the only part a different + top-k kernel changes. Bounded by construction too: GLM-5.2's indexer is + 1.6% of its attention arithmetic, so even an eight-fold slower kernel is + a ninth of the step rather than a multiple of it. That agrees with the + only end-to-end number available -- routing the ROCm decode top-k through + aiter's one-block kernel is 3.3x on the kernel and 5.6% on the trace -- + and it is worth stating plainly that a term this size cannot be what + separates a 4x throughput miss from a correct one. + """ + over = max(0.0, float(indexer_cost_scale) - 1.0) + if over <= 0.0: + return 0.0 + # Reuse the geometry by asking for the scale at a context long enough + # that selection has fallen away and only the indexing term is left. + ctx = max(1, int(topk or 0)) * 4096 + fused = uniform_sparse_attention_scale(model_config, ctx, topk, 1.0) + if fused is None: + return 0.0 + return max(0.0, fused - min(int(topk or 0), ctx) / ctx) * over + + # Representative expert grouped-GEMM compute multipliers by expert dtype. _MOE_EXPERT_DTYPE_SPEEDUP = { "bf16": 1.0, @@ -968,6 +1379,22 @@ class InferenceCollectiveConfig: # Whether to charge pipeline-stage P2P (send/recv) latency. Only nonzero # when pipeline_model_parallel_size > 1. include_pp_p2p: bool = True + # Engine time per pipeline hop per forward, beyond the transfer: the next + # stage's scheduling and synchronisation on the host. A PP-n forward makes + # n hops, the n-1 activation boundaries and the sampled tokens' return to + # the first stage. SGLang on MI355X at one client measured 920-1130 us per + # hop over PP 2/4/8 on Llama-3.1-8B, Qwen3-14B-FP8 and Qwen3.6-35B-A3B + # (TP1), about 55 us of it the send/recv. + pp_stage_overhead_us: float = 970.0 + # Scheduler host time per iteration that a stage pays in series with its + # forward once the pipeline is full: with micro-batches queued upstream a + # stage never waits on its neighbour, so the batch preparation it hides in + # that wait at one client lands on every stage. SGLang forces the overlap + # schedule off under PP; its TP1 overlap-off minus overlap-on step on + # MI355X read 329 us + 6.3 us per running request, within 0.08 ms across + # Llama-3.1-8B, Qwen3-14B-FP8 and Qwen3.6-35B-A3B at 1-64 requests. + pp_host_us: float = 329.0 + pp_host_per_req_us: float = 6.3 # Optional hardware overrides forwarded to ``get_default_args`` (node_bw, # pod_bw, bw_eff, latencies, ...). ``None`` uses the model defaults. hardware_config: dict | None = None @@ -1378,6 +1805,12 @@ def convert_config_to_inference_config( ) if attn_dp is not None: mp.attention_data_parallel_size = int(attn_dp) + # Decode context parallelism arrives the same way, and for the same reason. + dcp = overrides.get( + "decode_context_parallel_size", yaml_inf.get("decode_context_parallel_size") + ) + if dcp is not None: + mp.decode_context_parallel_size = int(dcp) return InferenceConfig( model_config=training_config.model_config, diff --git a/tests/unit/projection/conftest.py b/tests/unit/projection/conftest.py index 18021cd68..82b2d268d 100644 --- a/tests/unit/projection/conftest.py +++ b/tests/unit/projection/conftest.py @@ -137,6 +137,7 @@ def project_spec(**overrides): ("--attention-dp-size", "attn_dp"), ("--sparse-attention-topk", "sparse_attention_topk"), ("--prefix-cache-hit-rate", "prefix_cache_hit_rate"), + ("--kv-pool-tokens", "kv_pool_tokens"), ): if spec.get(key) is not None: argv += [flag, str(spec[key])] @@ -214,4 +215,7 @@ def project_spec(**overrides): # What the run printed, for the parts of the report that are computed # nowhere else. "report": report.getvalue(), + # The whole extras bag, for the projected quantities that are reported + # through it and not through the printed report -- admission being one. + "extras": extras, } diff --git a/tests/unit/projection/test_anchor_reuse.py b/tests/unit/projection/test_anchor_reuse.py index 588c91a86..f3590746f 100644 --- a/tests/unit/projection/test_anchor_reuse.py +++ b/tests/unit/projection/test_anchor_reuse.py @@ -19,8 +19,10 @@ from infera.projection.core.projection.inference_projection.search.regime import ( aiter_ops_axis, + recipe_from_inference_config, recipe_from_meta, regime_distance, + regime_signature, ) @@ -113,6 +115,130 @@ def test_an_unrecorded_op_set_matches_defaults_but_not_a_swap(): assert regime_distance(old, swapped) >= 1 +def test_two_engines_measuring_one_checkpoint_are_two_regimes(): + """An SGLang anchor does not describe vLLM, and the axes have to say so. + + The engine was named as regime-defining in the design split from the start + but was missing from ``REGIME_AXES``, so every engine's anchor for a given + checkpoint hashed to one signature: the bench cache would hand a vLLM + measurement to an SGLang run, and ``regime_distance`` called the pair a + perfect match. + """ + vllm = recipe_from_meta({"model": "m"}, engine="vllm") + sglang = recipe_from_meta({"model": "m"}, engine="sglang") + + assert regime_distance(vllm, sglang) >= 1 + assert regime_signature(vllm) != regime_signature(sglang) + + +def test_a_rebuilt_engine_does_not_inherit_the_engine_it_forked_from(): + """``mori-sglang`` runs SGLang's launch script and is not SGLang. + + It is the case that motivates the axis: the two share a recipe script, and + their reported KV pools differ 16x because one shards the MLA latent across + ranks where the other replicates it. Keyed on anything coarser than the + engine string, the fork silently inherits the parent's anchor. + """ + parent = recipe_from_meta({"model": "m"}, engine="sglang") + fork = recipe_from_meta({"model": "m"}, engine="mori-sglang") + + assert regime_distance(parent, fork) >= 1 + + +def test_the_engine_is_read_off_the_artifact_not_only_the_meta_block(): + """``benchmark_serving`` records the backend beside ``meta``, not inside it. + + So the store passes it explicitly. An artifact that does record it inline + is still honoured, which is what lets a future harness stop threading it. + """ + assert recipe_from_meta({"model": "m"}, engine="atom")["engine"] == "atom" + assert recipe_from_meta({"model": "m", "backend": "atom"})["engine"] == "atom" + assert recipe_from_meta({"model": "m", "engine": "atom"})["engine"] == "atom" + + +def test_an_anchor_indexed_before_the_axis_existed_still_matches(): + """Unknown stays skippable, or adding the axis would orphan the store. + + Every artifact already in the store predates the axis. Counting unknown as + a difference would refuse all of them at once, which is a worse failure + than the one being fixed -- and the engine filter in the harness's own + ``resolve`` is what actually gates reuse today. + """ + old = recipe_from_meta({"model": "m"}) + assert old["engine"] is None + assert regime_distance(old, recipe_from_meta({"model": "m"}, engine="vllm")) == 0 + + +def test_a_projection_names_the_engine_it_is_asking_about(): + """The target side is unknown unless ``--serving-engine`` says otherwise. + + A projection has no engine of its own, so without this the axis could never + bite on the path that actually consumes anchors. + """ + from infera.projection.core.projection.training_config import InferenceRequestConfig + + cfg = SimpleNamespace( + request_config=InferenceRequestConfig(serving_engine="sglang"), + model_config=None, + model_parallel_config=None, + ) + assert recipe_from_inference_config(cfg)["engine"] == "sglang" + assert InferenceRequestConfig().serving_engine is None + + +def test_naming_an_engine_does_not_change_what_simulate_projects(): + """It is a matching key, not a cost input. + + Simulate mode prices an architecture on a device analytically; if naming an + engine moved the number, the axis would be quietly inventing engine + performance it has no measurement for. + """ + from infera.projection.core.projection.training_config import InferenceRequestConfig + + base = InferenceRequestConfig() + named = InferenceRequestConfig(serving_engine="sglang") + assert ( + base.resolved_attention_backend_multiplier() + == named.resolved_attention_backend_multiplier() + ) + + +def test_the_store_rehashes_rather_than_trusting_a_recorded_signature(tmp_path): + """A signature in ``meta`` was hashed over whichever axes that build had. + + Preferring it pins the store to an older definition of "same regime", so + adding an axis cannot separate anchors that the older one merged -- which + is what happened to seven DeepSeek-V4-Flash anchors sharing one recorded + signature across SGLang and vLLM. + """ + import json + + from infera.projection.core.projection.inference_projection.search.anchor_store import ( + AnchorStore, + ) + + paths = [] + for engine in ("sglang", "vllm"): + p = tmp_path / f"{engine}.json" + p.write_text( + json.dumps( + { + "backend": engine, + # The same stale hash on both, as the real artifacts carry. + "meta": {"model": "m", "regime_signature": "staleaaaaaaaaaaa"}, + "sweep": [{"batch": 1, "decode_ms": 1.0}], + } + ) + ) + paths.append(str(p)) + + store = AnchorStore(str(tmp_path / "store")) + sigs = [store.add_artifact(p)["regime_signature"] for p in paths] + + assert "staleaaaaaaaaaaa" not in sigs, "a recorded signature is provenance, not the key" + assert sigs[0] != sigs[1], "two engines must not land in one regime" + + class _FakeStore: def __init__(self, models): self._models = models @@ -318,3 +444,34 @@ def test_discovery_stops_before_walking_an_entire_filesystem(tmp_path): deep = "a/b/c/d/e" _artifact(tmp_path, "far.json", sub=deep) assert AnchorStore(str(tmp_path)).entries() == [] + + +def test_an_anchor_from_another_part_is_not_the_same_regime(): + """gfx950 timings do not describe a gfx942 deployment. + + Nothing in the store recorded which accelerator produced a timing, so the + axis could not be compared and every anchor matched every target. The one + place that bit in this matrix is GLM-5.2: its only anchors are the + MXFP4 Quark build, which loads on gfx950 alone, and they were being used + to price an MI325X deployment -- a part with roughly half the compute. + It read 4.3x the measured throughput while reporting itself calibrated. + """ + anchor = recipe_from_meta( + {"model": "/models/GLM-5.2-MXFP4", "gpu_arch": "mi355x"}, engine="sglang" + ) + assert anchor["gpu_arch"] == "mi355x" + assert regime_distance(anchor, dict(anchor, gpu_arch="mi325x")) == 1 + assert regime_distance(anchor, dict(anchor, gpu_arch="mi355x")) == 0 + + +def test_an_anchor_that_never_recorded_its_part_is_still_usable(): + """Adding the axis must not retire every anchor harvested before it. + + An unrecorded part is unknown rather than wrong, and ``regime_distance`` + skips an axis neither side pins. The store is what closes the gap, by + stamping the part at build time for the pairs whose part can be argued + for -- so silence here means "nobody has said yet". + """ + legacy = recipe_from_meta({"model": "/models/DeepSeek-V4-Pro"}, engine="vllm") + assert legacy["gpu_arch"] is None + assert regime_distance(legacy, dict(legacy, gpu_arch="mi300x")) == 0 diff --git a/tests/unit/projection/test_des_cache_pool_sharing.py b/tests/unit/projection/test_des_cache_pool_sharing.py new file mode 100644 index 000000000..52ad345af --- /dev/null +++ b/tests/unit/projection/test_des_cache_pool_sharing.py @@ -0,0 +1,75 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The prefix cache and the running requests are the same pool. + +``_Req.reserved_kv`` does not charge a request for the leading blocks it hit +on, because an earlier request is holding them. That is the right account, and +it only balances if whatever is holding them is charged instead. The holder is +the block cache, and it was being handed the whole pool at the same time as the +running set, so the two structures booked the same memory and neither ever saw +it run out. + +Where the pool is roomy against the working set the correction is nothing, and +that is most of the corpus. Where it is not, it is the entire behaviour: on the +AgentX ladders GLM-5.2 on MI325X has a pool that holds 5.5 of these +conversations, and the hardware's reuse falls 0.79 -> 0.30 -> 0.12 across +C=4,5,6 while TTFT goes from 3s to 198s. Double-booked, the replay reported +reuse of 0.96 and a TTFT flat within a factor of two across that whole +collapse, so the ladder came out ranked by nothing in particular. +""" + +from __future__ import annotations + +from infera.projection.core.projection.inference_projection import des as des_mod + +_BLOCK = 64 + + +def test_a_roomy_pool_leaves_the_cache_what_it_was_given(): + """The correction has to be silent where it does not apply.""" + cap = 1000 + got = des_mod._free_cache_blocks( + kv_cache_tokens=10_000_000, live_tokens=100_000, block_size=_BLOCK, cap_blocks=cap + ) + assert got == cap + + +def test_a_full_pool_leaves_the_cache_nothing(): + """No room is no room, however the request was charged for it.""" + got = des_mod._free_cache_blocks( + kv_cache_tokens=1_048_576, live_tokens=1_500_000, block_size=_BLOCK, cap_blocks=16_384 + ) + assert got == 1 + + +def test_no_room_is_never_read_as_unlimited_room(): + """``_BlockStore`` takes a zero capacity as unbounded. + + Returning the honest zero would hand a pool with nothing left in it the + one cache that never evicts, which is the failure this is here to fix + rather than a smaller version of it. + """ + for live in (1_048_576, 2_000_000, 10**9): + got = des_mod._free_cache_blocks( + kv_cache_tokens=1_048_576, live_tokens=live, block_size=_BLOCK, cap_blocks=16_384 + ) + assert got >= 1 + + +def test_the_cache_shrinks_as_the_running_set_grows(): + """Reuse degrades with pressure rather than falling off a single step.""" + pool = 1_048_576 + caps = [ + des_mod._free_cache_blocks(pool, c * 190_000, _BLOCK, 16_384) for c in (1, 2, 3, 4, 5, 6) + ] + assert caps == sorted(caps, reverse=True), caps + assert caps[0] > caps[3] > caps[-1] + + +def test_an_unmeasured_pool_is_left_alone(): + """Without a pool figure there is nothing to divide, so nothing is claimed.""" + assert des_mod._free_cache_blocks(0, 500_000, _BLOCK, 4096) == 4096 + assert des_mod._free_cache_blocks(1_048_576, 500_000, 0, 4096) == 4096 diff --git a/tests/unit/projection/test_des_closed_loop.py b/tests/unit/projection/test_des_closed_loop.py index 0291f5d54..2c79da286 100644 --- a/tests/unit/projection/test_des_closed_loop.py +++ b/tests/unit/projection/test_des_closed_loop.py @@ -74,7 +74,13 @@ def mixed_step_latency_ms(self, num_decode, prefill_tokens, ctx, prefill_kv, q_l return _BASE_MS + _PER_SEQ_MS * num_decode + _PER_PREFILL_TOKEN_MS * prefill_tokens -def _run(concurrency: int, chunk: int = 256, requests_per_client: int = 6, **kw): +def _run( + concurrency: int, + chunk: int = 256, + requests_per_client: int = 6, + think_ms: float = 0.0, + **kw, +): cfg = _Cfg(_Req(max_concurrency=concurrency, chunked_prefill_size=chunk, **kw)) return des_mod.simulate_once( cfg, @@ -87,6 +93,7 @@ def _run(concurrency: int, chunk: int = 256, requests_per_client: int = 6, **kw) # which is a separate finding; here it is dropped. warmup_frac=0.5, closed_loop_clients=concurrency, + closed_loop_think_ms=think_ms, ) @@ -321,3 +328,98 @@ def test_run_des_takes_the_client_count_from_the_configured_concurrency(): assert point.arrival_model == "closed" assert point.packing["closed_loop_clients"] == 32.0 assert point.ttft["mean"] > 0.0 + + +def test_exclusive_prefill_reaches_the_path_a_block_cache_dispatches_to(): + """Whether prefill excludes decode has to survive the trip through ``run_des``. + + The mechanism is implemented in ``simulate_once``, but a closed-loop run + with a block cache to model -- a prefix pool or a trace to replay, which is + every agentic replay -- is dispatched through ``simulate_multi_instance``, + and that driver did not take the flag at all. A parameter dropped one frame + up is indistinguishable from one that was never passed: SGLang and Atom + replays scheduled a unified batch while the flag that was meant to + serialise them sat unused, which dissolves the herd a closed-loop + population forms and reads TTFT early. Assert the effect at the entry point + the harness calls, not at the loop that implements it. + """ + clients = 8 + kw = dict( + arrival_model="closed", + rate_per_s=0.0, + num_requests=clients * 8, + closed_loop=True, + warmup_frac=0.0, + # Enough of a prefix pool to put the run on the multi-instance path. + num_prefixes=4, + prefix_len=256, + block_size=16, + cache_blocks=1 << 14, + ) + cfg = _Cfg(_Req(max_concurrency=clients, chunked_prefill_size=256)) + unified = des_mod.run_des(cfg, _Kernel(), **kw)["point"] + exclusive = des_mod.run_des(cfg, _Kernel(), prefill_exclusive=True, **kw)["point"] + + # Serialising prefill against the resident decodes can only make a request + # wait longer for its first token. + assert exclusive.ttft["mean"] > unified.ttft["mean"], ( + exclusive.ttft["mean"], + unified.ttft["mean"], + ) + + +# -- client think time -------------------------------------------------------- +# +# A closed loop that resubmits the instant its last token lands is a load +# generator, not a user. An agentic harness ends a turn, runs the tool the turn +# asked for, and sends the next turn when that returns; for the run the replay +# is scored against, reconstructing each client's cycle from the run's own +# completed-request count and wall clock leaves half to three quarters of it +# unserved at the rungs below saturation. +# +# What makes it worth modelling rather than absorbing into a level correction +# is that it does not shift the curve, it bends it: the same delay is most of +# the cycle when the server is idle and none of it once the server is the +# bottleneck and the next turn queues anyway. + + +def test_a_thinking_client_leaves_the_server_idle(): + """Without a think time the server is busy by construction. + + Every client resubmits into the instant it was served in, so there is no + interval in which the engine has nothing to run, and the replay reports a + saturated server for a load that never saturated one. + """ + c = 16 + eager = _run(c) + # Comparable to the stub's own service time, which is the regime the + # AgentX ladders sit in: the delay is the same order as the turn. + thinking = _run(c, think_ms=500.0) + assert eager.utilization > 0.99, eager.utilization + assert thinking.utilization < 0.8 * eager.utilization, ( + eager.utilization, + thinking.utilization, + ) + + +def test_thinking_costs_throughput_only_while_the_server_has_slack(): + """The bend, which is the reason this is a workload input and not a factor. + + At low concurrency the delay is most of a client's cycle and throughput + falls with it. At high concurrency the client's next turn waits on the + server regardless, so the same delay costs progressively less -- a + correction applied as a constant could not do both ends at once. + """ + think = 500.0 + lo_loss = _run(4, think_ms=think).achieved_rate / _run(4).achieved_rate + hi_loss = _run(256, think_ms=think).achieved_rate / _run(256).achieved_rate + assert lo_loss < 1.0, lo_loss + assert hi_loss > lo_loss, (lo_loss, hi_loss) + + +def test_an_unset_think_time_is_the_loop_as_it_was(): + """The default has to leave every existing closed-loop answer untouched.""" + a = _run(16) + b = _run(16, think_ms=0.0) + assert a.achieved_rate == b.achieved_rate + assert a.ttft["mean"] == b.ttft["mean"] diff --git a/tests/unit/projection/test_des_disaggregated.py b/tests/unit/projection/test_des_disaggregated.py index b55ebbfba..d94d8657d 100644 --- a/tests/unit/projection/test_des_disaggregated.py +++ b/tests/unit/projection/test_des_disaggregated.py @@ -327,3 +327,94 @@ def test_the_loop_stops_instead_of_spinning_when_nothing_can_progress(): ) assert res.num_requests == 0 assert res.makespan_ms == 0.0 + + +def _run_split( + concurrency: int, + *, + requests_per_client: int = 4, + prefill_scale: float = 0.1, + chunk: int = 256, + output_len: int = 512, + range_ratio: float = 0.0, +): + """A split whose prefill pool finishes well ahead of its decode pool. + + That ordering is the one that matters below: prefill drains its queue and + goes idle while decode is still working, which is exactly when the two + stations' clocks come apart. + """ + cfg = _Cfg( + _Req(max_concurrency=concurrency, chunked_prefill_size=chunk, output_seq_len=output_len), + _Disagg(decode_replicas=1), + ) + proj = _Projector(_Pool(prefill_scale), _Pool(1.0), handoff_ms=1.0) + return des_mod.simulate_disaggregated( + cfg, + proj, + rate_per_s=0.0, + num_requests=concurrency * requests_per_client, + warmup_frac=0.0, + range_ratio=range_ratio, + closed_loop_clients=concurrency, + seed=7, + ) + + +def test_a_closed_loop_keeps_its_clients_in_flight(): + """The two stations keep separate clocks, and only one was feeding itself. + + A closed-loop client retires on the *decode* clock and submits its + replacement stamped with that time, but arrivals were released against the + *prefill* clock. Whenever prefill drained its queue its clock stopped while + decode's kept moving, so every request reissued in that window was invisible + to the pool that has to prefill it. The population in flight collapsed from + C to about one, and with nothing to batch with the decode pool ran at batch + 1 however many clients were offered -- which reads as a split that cannot + use concurrency rather than as a bug in the loop. + """ + res = _run_split(32) + # Nearly all of a request's life is decode here, so most of the 32 clients + # should be resident there at any moment. This was 1.29. + assert res.packing["avg_decode_reqs"] > 16.0 + assert res.packing["closed_loop_clients"] == 32.0 + + +def test_the_decode_batch_grows_with_the_client_count(): + """Offering more clients has to put more of them in the batch. + + The failure this guards is specifically a *flat* response: the collapsed + loop returned about the same batch, and so about the same throughput, at + every concurrency, which is indistinguishable from a deployment that is + genuinely saturated. Before the fix a fourfold rise in clients moved the + batch from 1.08 to 1.29. + """ + few = _run_split(8) + many = _run_split(32) + assert many.packing["avg_decode_reqs"] > 2.5 * few.packing["avg_decode_reqs"] + assert many.system_throughput_tps > 2.0 * few.system_throughput_tps + + +def test_the_stall_does_not_depend_on_the_lengths_being_uniform(): + """Staggered lengths keep decode occupied, which is the harder case. + + With every request the same size the pools empty together, and an empty + system takes the loop's "nothing can run" branch, which jumps both clocks + forward and hides the drift. A spread of lengths keeps decode busy across + the gap, so the drift never gets collected. + """ + res = _run_split(32, range_ratio=0.8) + assert res.packing["avg_decode_reqs"] > 16.0 + + +def test_no_request_is_prefilled_before_its_client_submitted_it(): + """Feeding prefill from the simulation clock must not outrun arrival. + + Releasing against the global clock is what fixes the stall, and the risk it + introduces is the mirror image: serving a request at a time the prefill pool + has not reached yet, which would surface as a negative queue wait and an + understated TTFT. + """ + res = _run_split(8) + assert res.queue_wait["p50"] >= 0.0 + assert res.queue_wait["p99"] >= 0.0 diff --git a/tests/unit/projection/test_des_measurement_window.py b/tests/unit/projection/test_des_measurement_window.py new file mode 100644 index 000000000..7d89f82ab --- /dev/null +++ b/tests/unit/projection/test_des_measurement_window.py @@ -0,0 +1,298 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""What the replay reports over, and what it is allowed to reorder. + +Three decisions about the *measurement window* rather than about physics, each +of which moved a validated result further than any cost-model change did. + +1. A fixed-concurrency run is bounded by a clock, not by a request count. The + distinction bites under a closed loop because a request budget divided + among C lanes fixes how far each lane walks into its conversation -- deep + at low concurrency, shallow at high -- and an agentic turn's prompt grows + with its position, so a budget-bounded replay offers a prompt-length trend + with concurrency that the hardware never had. + +2. If the opening transient is dropped at all, it has to be dropped by + *issue* order. Every client fires at once, so the first C requests queue + against each other and carry the run's longest waits; because they waited, + they are also the last to retire, so dropping a fraction of the earliest + *completions* keeps every one of them. (Whether to drop it is a separate + question, and for this harness the answer is no -- its lane-advance + requests are one-token requests, so the profiled window still opens with + every lane firing together and the burst is in the reported average.) + +3. A scheduler can only reorder requests that are queued. Admitting by + longest resident prefix is the right policy for a deep queue and the wrong + one for an empty one. + +The cost kernel is a stub so that every number below follows from it. +""" + +from __future__ import annotations + +import random +from dataclasses import dataclass + +from infera.projection.core.projection.inference_projection import des as des_mod + +_BASE_MS = 2.0 +_PER_SEQ_MS = 0.05 +_PER_PREFILL_TOKEN_MS = 0.002 + + +@dataclass +class _Req: + input_seq_len: int = 1024 + output_seq_len: int = 64 + max_concurrency: int = 8 + max_num_batched_tokens: int = 8192 + chunked_prefill_size: int = 256 + speculative_num_tokens: int = 0 + speculative_acceptance_rate: float = 0.0 + tokenize_overhead_us: float = 0.0 + detokenize_overhead_us: float = 0.0 + + def resolved_max_concurrency(self) -> int: + return self.max_concurrency + + def resolved_max_context_len(self) -> int: + return self.input_seq_len + self.output_seq_len + + def resolved_prefix_cache_hit_rate(self) -> float: + return 0.0 + + +@dataclass +class _Cfg: + request_config: _Req + + +class _Kernel: + def decode_step_latency_ms(self, batch, ctx, q_len): + return _BASE_MS + _PER_SEQ_MS * batch + + def mixed_step_latency_ms(self, num_decode, prefill_tokens, ctx, prefill_kv, q_len): + return _BASE_MS + _PER_SEQ_MS * num_decode + _PER_PREFILL_TOKEN_MS * prefill_tokens + + +def _run( + concurrency: int = 8, + requests_per_client: int = 20, + warmup_frac: float = 0.0, + pool_tokens: int = 0, + **kw, +): + cfg = _Cfg(_Req(max_concurrency=concurrency)) + return des_mod.simulate_once( + cfg, + _Kernel(), + rate_per_s=0.0, + arrival_model="closed", + num_requests=concurrency * requests_per_client, + warmup_frac=warmup_frac, + kv_cache_tokens=pool_tokens, + closed_loop_clients=concurrency, + **kw, + ) + + +# --- 1. the clock bounds the run ------------------------------------------ + + +def test_a_duration_stops_the_run_before_the_request_budget_does(): + """With a horizon short of the budget, the budget stops mattering.""" + full = _run() + horizon = full.makespan_ms / 4.0 + cut = _run(duration_ms=horizon) + assert cut.num_requests < full.num_requests, (cut.num_requests, full.num_requests) + assert cut.makespan_ms <= horizon + 1e-6, (cut.makespan_ms, horizon) + + +def test_the_horizon_reaches_the_path_the_harness_runs_on(): + """The clock has to survive the trip through ``run_des``. + + The horizon is read in ``simulate_once``, but a closed-loop run with a + block cache to model is dispatched through ``simulate_multi_instance``, + and a parameter that stops at the driver is indistinguishable from one + that was never passed: every replay reported over its whole trace while + the flag that was supposed to bound it sat unused two frames up. Assert + the effect at the entry point the harness calls, not at the loop that + implements it. + """ + clients = 4 + kw = dict( + arrival_model="closed", + rate_per_s=0.0, + num_requests=clients * 40, + closed_loop=True, + num_prefixes=4, + prefix_len=256, + block_size=16, + cache_blocks=1 << 14, + warmup_frac=0.0, + ) + cfg = _Cfg(_Req(max_concurrency=clients)) + full = des_mod.run_des(cfg, _Kernel(), **kw)["point"] + cut = des_mod.run_des(cfg, _Kernel(), duration_ms=full.makespan_ms / 4.0, **kw)["point"] + assert cut.num_requests < full.num_requests, (cut.num_requests, full.num_requests) + + +def test_a_horizon_longer_than_the_run_changes_nothing(): + """The clock is a bound, not a target: a slack horizon is inert.""" + full = _run() + slack = _run(duration_ms=full.makespan_ms * 10.0) + assert slack.num_requests == full.num_requests + assert slack.ttft["mean"] == full.ttft["mean"] + + +def test_the_lane_walk_is_what_a_request_budget_fixes(): + """Why the clock matters: a budget sets per-lane depth from the lane count. + + This is the artifact the horizon exists to remove. Splitting one budget + across C lanes means each lane advances ``budget / C`` turns, so the depth + a lane reaches -- and with it the prompt it offers, for a workload whose + turns grow -- is decided by the concurrency rather than by the workload. + """ + budget = 400 + depths = {c: budget // c for c in (1, 8, 64)} + assert depths[1] == 400 and depths[8] == 50 and depths[64] == 6 + # A 66x spread in how far a lane walks, for one unchanged workload. + assert depths[1] / depths[64] > 60 + + +# --- 2. if the opening transient is dropped, drop it by issue order -------- + + +def _done(n: int, opening: int, opening_wait: float = 500.0): + """A completed run whose opening requests waited long and finished last. + + This is the shape a closed loop actually produces. The first ``opening`` + requests are issued together at t=0 and queue against each other, so they + carry the run's longest waits -- and because they waited, they are also + the last to retire. Everything issued afterwards arrives as a slot frees + and waits for nothing. + """ + out = [] + for i in range(n): + r = des_mod._Req(idx=i, arrival_ms=0.0, prompt_len=1024, output_len=8) + burst = i < opening + r.admit_ms = opening_wait if burst else 0.0 + r.first_token_ms = r.admit_ms + 10.0 + # Finish order is the point: the burst retires at the end of the run. + r.finish_ms = (1e6 + i) if burst else float(i) + r.generated = r.output_len + out.append(r) + out.sort(key=lambda x: x.finish_ms) + return out + + +def test_a_completion_ordered_drop_retains_the_opening_burst(): + """Why ``warmup_frac`` cannot remove the transient it was meant to. + + It drops the earliest *completions*, and a request that waited a long + time for a slot is among the last to finish. Dropping a quarter of the + run by completion keeps all eight of the burst. + """ + done = _done(n=80, opening=8) + kept = des_mod._scored_sample(done, warmup_frac=0.25, warmup_requests=0) + assert sum(1 for r in kept if r.idx < 8) == 8 + + +def test_an_issue_ordered_drop_removes_it(): + """``warmup_requests`` drops by the order the clients issued, which works.""" + done = _done(n=80, opening=8) + kept = des_mod._scored_sample(done, warmup_frac=0.0, warmup_requests=8) + assert sum(1 for r in kept if r.idx < 8) == 0 + assert len(kept) == 72 + + +def test_the_two_drops_report_very_different_queue_waits(): + """The size of the correction, on the shape the harness actually produces.""" + done = _done(n=80, opening=8, opening_wait=500.0) + + def mean_wait(sample): + return sum(r.admit_ms - r.arrival_ms for r in sample) / len(sample) + + by_completion = mean_wait(des_mod._scored_sample(done, warmup_frac=0.25, warmup_requests=0)) + by_issue = mean_wait(des_mod._scored_sample(done, warmup_frac=0.0, warmup_requests=8)) + assert by_issue == 0.0 + assert by_completion > 60.0, by_completion + + +def test_the_issue_order_drop_keeps_at_least_half_the_run(): + """A trace shorter than the requested warmup still reports over something.""" + done = _done(n=20, opening=4) + kept = des_mod._scored_sample(done, warmup_frac=0.0, warmup_requests=10_000) + assert len(kept) >= 10 + + +# --- 3. only a real queue is reorderable ---------------------------------- + + +def _reuse(waiting_depth: int, resident_cap: int, blocks_per_req: int = 4): + """Reuse ``_route_and_warm`` derives for a synthetic interleaved stream. + + The resolved order is not observable from the outside -- the function + deliberately hands the stream back in arrival order, because the loops + that consume it are arrival-driven -- so the ordering shows up in the one + place it has an effect: how much of each prompt the cache covered. + """ + convs = 4 + reqs = [] + for i in range(convs * 5): + r = des_mod._Req(idx=i, arrival_ms=float(i), prompt_len=1024, output_len=16) + # More live conversations than the cache holds, round-robin the way C + # lanes arrive. Consecutive turns of one conversation share blocks, so + # whether that reuse survives depends entirely on how many other + # conversations get admitted in between. + conv = i % convs + r.blocks = [conv * 1000 + b for b in range(blocks_per_req)] + reqs.append(r) + _, summary = des_mod._route_and_warm( + reqs, + policy="kv", + num_instances=1, + block_size=64, + cache_blocks=blocks_per_req + 1, + rng=random.Random(0), + overlap_weight=1.0, + waiting_depth=waiting_depth, + resident_cap=resident_cap, + ) + return summary["hit_rate"] + + +def test_reuse_is_lower_when_nothing_is_queued_to_reorder(): + """The finding: an empty queue leaves reuse to the arrival interleaving. + + Below the cap every lane is resident at once and they evict each other, + which is the regime where measured reuse collapses. Handing the policy + the whole client count let it reorder that empty queue and held modelled + reuse high through exactly the band where the hardware thrashed. + """ + unqueued = _reuse(waiting_depth=20, resident_cap=20) + queued = _reuse(waiting_depth=20, resident_cap=1) + assert queued > unqueued, (queued, unqueued) + + +def test_the_cap_counts_what_the_pool_holds(): + """``_resident_cap`` divides the pool by a length-biased context. + + Length-biased rather than the plain mean: a request occupies the pool for + a time proportional to its own length, so what is resident at any instant + is sampled in proportion to length, and the arithmetic mean overstates how + many fit. + """ + reqs = [ + des_mod._Req(idx=i, arrival_ms=0.0, prompt_len=n, output_len=0) + for i, n in enumerate([100] * 9 + [900]) + ] + # Plain mean is 180, so 1800 tokens would look like room for 10. + assert des_mod._resident_cap(reqs, 1800, 0, enabled=True) < 10 + # No pool to divide leaves the policy unconstrained. + assert des_mod._resident_cap(reqs, 0, 0, enabled=True) == 0 + # And it is opt-in: by default the window is not constrained at all. + assert des_mod._resident_cap(reqs, 1800, 0) == 0 diff --git a/tests/unit/projection/test_des_prefix_cache.py b/tests/unit/projection/test_des_prefix_cache.py index 902091816..a36bf645d 100644 --- a/tests/unit/projection/test_des_prefix_cache.py +++ b/tests/unit/projection/test_des_prefix_cache.py @@ -81,3 +81,254 @@ def test_the_scheduler_only_charges_for_the_uncached_suffix(): assert need == 10400 # And the cache blocks count as resident context immediately. assert r.kv_len == 119600 + + +def test_a_shared_prefix_is_charged_to_the_pool_once_and_not_per_sharer(): + """Admission has to read the hit too, not just prefill. + + The pool gate charged every request its whole context, so a corpus whose + requests all continue from the same conversation paid for the one resident + copy of that prefix once per sharer. At the 96-98% reuse an agentic replay + carries that inflates a request's footprint 25-50x: a measured 7.67M-token + pool looked full at ~19 concurrent and the replay queued from there, while + the MI355X ladder reports its pool at 0.1-0.8% utilization at every + concurrency and never binds. The uncached suffix plus what the request + generates is what it actually adds. + """ + r = _seed(0.92)[0] + assert r.reserved_kv == (r.prompt_len - r.cached_prefix) + r.output_len == 11300 + # A cold request is untouched, which is every fixed-sequence workpoint. + cold = _seed(0.0)[0] + assert cold.reserved_kv == cold.prompt_len + cold.output_len + + +def _sweep(cache: des_mod._BlockCache, n_convs: int, depth: int, rounds: int): + """Interleave ``n_convs`` conversations that share a head and diverge. + + The shape every long-context corpus has: a common prefix, then per-thread + continuation. Each round revisits every thread, so a thread's own blocks + are reused a full sweep after they were written. + """ + head = [0, 1, 2] + hits = 0 + for _ in range(rounds): + for c in range(n_convs): + blocks = head + [1000 + c * depth + d for d in range(depth)] + hits += cache.prefix_match(blocks) + cache.insert(blocks) + return hits + + +def test_the_shared_head_of_a_corpus_survives_a_working_set_it_cannot_hold(): + """Leaf-first eviction is the whole difference between 0% reuse and most of it. + + Blocks are touched in prefix order, so a corpus larger than the pool is a + cyclic sweep -- the access pattern that makes a flat LRU evict every block + exactly before it is reused. A radix cache cannot do that: it may only + reclaim a block no resident block continues from, so the shared head stays + and the divergent tails are what go. + """ + # Room for the head and roughly half the threads, so the sweep cannot fit. + cache = des_mod._BlockCache(capacity_blocks=3 + 10 * 4) + hits = _sweep(cache, n_convs=20, depth=4, rounds=4) + + assert hits > 0, "a flat LRU scores exactly zero here; that was the bug" + assert cache.prefix_match([0, 1, 2]) == 3, "the shared head must still be resident" + + +def test_a_block_is_never_dropped_while_something_continues_from_it(): + """The invariant that makes it a radix cache rather than a priority order.""" + cache = des_mod._BlockCache(capacity_blocks=4) + cache.insert([0, 1, 2, 3]) + cache.insert([0, 1, 2, 4]) # diverges at the last block, forcing eviction + + assert cache.prefix_match([0, 1, 2]) == 3, "an interior block was reclaimed" + assert len(cache) <= 4 + + +def test_the_cache_still_honours_its_capacity(): + """Protecting interior blocks must not become "never evict".""" + cache = des_mod._BlockCache(capacity_blocks=8) + _sweep(cache, n_convs=30, depth=4, rounds=2) + assert len(cache) <= 8 + + +def test_an_unbounded_cache_evicts_nothing(): + cache = des_mod._BlockCache(capacity_blocks=0) + _sweep(cache, n_convs=10, depth=4, rounds=2) + assert cache.evictions == 0 + assert cache.prefix_match([0, 1, 2]) == 3 + + +def _lanes(n_lanes: int, turns: int, blocks_per_turn: int = 6): + """A closed-loop corpus: ``n_lanes`` conversations taking turns round-robin. + + Each lane extends its own context every turn, so turn ``t`` of a lane can + reuse everything that lane wrote at ``t - 1`` -- if it is still resident a + full sweep of the other lanes later. This is the shape of every agentic + trace, and the distance between a lane's turns is the offered concurrency. + """ + reqs = [] + for t in range(turns): + for lane in range(n_lanes): + blocks = [lane * 10_000 + b for b in range((t + 1) * blocks_per_turn)] + reqs.append( + des_mod._Req( + idx=len(reqs), + arrival_ms=float(len(reqs)), + prompt_len=len(blocks) * 16, + output_len=16, + blocks=blocks, + ) + ) + return reqs + + +def _warm(reqs, capacity: int, depth: int) -> float: + import random + + _, summary = des_mod._route_and_warm( + [ + des_mod._Req( + idx=r.idx, + arrival_ms=r.arrival_ms, + prompt_len=r.prompt_len, + output_len=r.output_len, + blocks=list(r.blocks), + ) + for r in reqs + ], + policy="kv", + num_instances=1, + block_size=16, + cache_blocks=capacity, + rng=random.Random(0), + waiting_depth=depth, + ) + return summary["block_hit_rate"] + + +def test_reuse_does_not_fall_off_a_cliff_once_the_offered_load_stops_fitting(): + """Oldest-first admission turns "does not fit" into "reuses nothing". + + Past the point where the pool cannot hold every outstanding context, taking + the oldest waiting request walks the lanes round-robin, which guarantees a + lane's context is evicted before its next turn: every request then pays a + full reprefill and reuse is not merely degraded, it is exactly zero. Real + hardware degrades gradually instead, because a scheduler picks off its + waiting queue by longest resident prefix and keeps working on whatever is + still resident. The measured ladders show exactly that -- reuse holding + near 0.93 up to the admission ceiling and never collapsing past it. + """ + reqs = _lanes(n_lanes=24, turns=6) + # Room for a few lanes' contexts, nowhere near all 24. + capacity = 6 * 21 + + oldest_first = _warm(reqs, capacity, depth=0) + prefix_first = _warm(reqs, capacity, depth=24) + + assert oldest_first == 0.0, "the cliff being fixed; if this moves, re-derive it" + assert prefix_first > 0.6, prefix_first + + +def test_admission_order_changes_nothing_while_the_load_still_fits(): + """The correction has to be invisible below the pressure point. + + With room for every outstanding context nothing is evicted between a lane's + turns, so both orders touch the same blocks and must score identically. + Were that not so this would be a tuning knob rather than a missing + behaviour, and it would be moving results in regimes it has no business in. + """ + reqs = _lanes(n_lanes=8, turns=6) + capacity = 0 # unbounded: everything offered fits + + assert _warm(reqs, capacity, depth=0) == _warm(reqs, capacity, depth=8) + + +def test_the_slot_a_served_request_frees_goes_back_to_its_own_client(): + """Which client gets the freed slot is most of the correction. + + Serving a request frees the client that issued it, and that client's next + turn continues the context just served. Hand the slot to whichever request + is next in arrival order instead and the closed loop degrades into a + sliding window: the reuse that was about to be collected is dropped, and + on the dsv4 ladder that is the difference between reproducing the measured + 0.93 at C=128 and reporting 0.83. + """ + reqs = _lanes(n_lanes=24, turns=6) + capacity = 6 * 21 + + # Same policy, same capacity; the only difference is that a depth which + # does not match the interleave misattributes the freed slot. + tied = _warm(reqs, capacity, depth=24) + untied = _warm(reqs, capacity, depth=23) + + assert tied > untied, (tied, untied) + + +def test_the_arrival_stream_each_instance_replays_is_still_in_arrival_order(): + """Schedule order says how much to prefill, not when a request showed up. + + The per-instance loops downstream are arrival-driven, so resolving prefixes + in a different order must not leak into the stream they replay. + """ + import random + + reqs = _lanes(n_lanes=8, turns=4) + per_inst, _ = des_mod._route_and_warm( + reqs, + policy="kv", + num_instances=1, + block_size=16, + cache_blocks=32, + rng=random.Random(0), + waiting_depth=8, + ) + arrivals = [r.arrival_ms for r in per_inst[0]] + assert arrivals == sorted(arrivals) + + +def test_a_split_reports_the_prefix_cache_it_actually_warmed(tmp_path): + """The disaggregated branch warmed a cache and threw the summary away. + + Every split row then reported no hit rate, which reads as a split that + reuses nothing rather than one whose reuse was never recorded -- and it + was the disaggregated rows that the reuse question was being asked about. + """ + import json + + from .test_des_disaggregated import _Cfg, _Pool, _Projector, _Req + + # Two passes over the same prompts, so there is reuse to find. + trace = tmp_path / "trace.jsonl" + with trace.open("w") as f: + for turn in range(2): + for lane in range(4): + f.write( + json.dumps( + { + "timestamp": turn * 1000 + lane, + "input_length": 2048, + "output_length": 64, + "hash_ids": [0, 1, 2, 100 + lane, 200 + lane], + } + ) + + "\n" + ) + + point = des_mod.run_des( + _Cfg(_Req(max_concurrency=4, chunked_prefill_size=256, output_seq_len=64)), + _Projector(_Pool(), _Pool()), + arrival_model="closed", + rate_per_s=0.0, + num_requests=8, + warmup_frac=0.0, + closed_loop=True, + mooncake_trace=str(trace), + block_size=512, + cache_blocks=64, + )["point"] + + assert point.packing["disaggregated"] == 1.0, "this must be the split path" + assert point.prefix, "a split reported no prefix-cache summary at all" + assert point.prefix["hit_rate"] > 0.0 diff --git a/tests/unit/projection/test_disaggregated_report.py b/tests/unit/projection/test_disaggregated_report.py index 856412270..f07acfbcd 100644 --- a/tests/unit/projection/test_disaggregated_report.py +++ b/tests/unit/projection/test_disaggregated_report.py @@ -133,3 +133,55 @@ def test_a_colocated_run_still_reports_its_own_replica_gpus(): """The regression guard: the split's line must not displace this one.""" m = parse_inference_metrics(_colocated(tp=8)["report"]) assert m.get("replica_gpus") == 8 + + +def test_a_split_prices_admission_against_its_decode_pool(): + """A split could over-subscribe KV and reported no admission wait at all. + + The pool binds on the decode side, so the split is as capable of + over-subscribing it as a colocated engine. The term was only ever applied + on the colocated path, which left a split's TTFT as pure service time -- + and silently, because an absent extra reads the same as a zero one. + """ + e = _split(kv_pool_tokens=4 * 2048, input_len=2048, concurrency=32)["extras"] + assert e.get("admission_ceiling") == pytest.approx(4, abs=1), ( + f"a pool holding four 2048-token residents admits 4; got {e.get('admission_ceiling')!r}" + ) + assert (e.get("admission_wait_ms") or 0.0) > 0.0, ( + "32 clients against a pool that admits 4 must wait to be admitted" + ) + + +def test_a_roomy_decode_pool_leaves_a_split_with_no_admission_wait(): + """Zero is the answer for the published splits, not a missing term. + + Moving prefill off the decode GPUs is what buys the room: the measured + DeepSeek splits have a decode pool admitting ~770 sequences against the + 256 ever offered, so they never reach a knee. Their TTFT still grows 9x + across that range, which is the prefill pool's queue and not this term. + """ + e = _split(kv_pool_tokens=4096 * 2048, input_len=2048, concurrency=32)["extras"] + assert e.get("admission_ceiling", 0) > 32 + assert (e.get("admission_wait_ms") or 0.0) == pytest.approx(0.0) + + +def test_admission_is_charged_per_decode_replica(): + """Clients spread over the pool, so one replica faces its share of them.""" + pool = 8 * 2048 # admits seven of these residents: prompt plus generation + one = _split(kv_pool_tokens=pool, decode_replicas=1, concurrency=32)["extras"] + four = _split(kv_pool_tokens=pool, decode_replicas=4, concurrency=32)["extras"] + eight = _split(kv_pool_tokens=pool, decode_replicas=8, concurrency=32)["extras"] + assert (one.get("admission_wait_ms") or 0.0) > 0.0 + assert four["admission_wait_ms"] < one["admission_wait_ms"] / 3, ( + "32 clients over four replicas is 8 each and not 32, so one replica's " + "over-subscription falls roughly four-fold" + ) + assert (eight.get("admission_wait_ms") or 0.0) == pytest.approx(0.0), ( + "four clients each is inside a pool that admits seven" + ) + + +def test_a_colocated_run_still_prices_its_own_admission(): + """The regression guard: the new per-pool argument defaults to one pool.""" + e = _colocated(kv_pool_tokens=4 * 2048, input_len=2048, concurrency=32)["extras"] + assert (e.get("admission_wait_ms") or 0.0) > 0.0 diff --git a/tests/unit/projection/test_hybrid_compressed_attention.py b/tests/unit/projection/test_hybrid_compressed_attention.py new file mode 100644 index 000000000..b2980e90a --- /dev/null +++ b/tests/unit/projection/test_hybrid_compressed_attention.py @@ -0,0 +1,212 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""A layer that compresses its KV should not be billed for context it never reads. + +DeepSeek-V4 records a per-layer schedule saying which attention branch each +layer runs: a local window only, a heavily-compressed pool with full visibility +over it, or a compressed pool that an indexer selects a fixed top-k from. The +checkpoint ships it, and ``deepseek_v4_base.yaml`` ships the geometry that sizes +it, but nothing read either -- the whole 61-layer stack was priced from one +top-k scale floored at 0.15. + +A floor cannot say what that schedule says. Attention under selection stops +growing once the pool exceeds the top-k, so what keeps growing with context is +the indexer scoring the pool, which is 9% of a CSA layer at 8k and 62% at 130k. +Flooring the scale instead makes the non-shrinking part grow like dense +attention over the whole prompt: at 130k the schedule costs 0.0157 of dense +against the floor's 0.15, a 9.6x over-read, and the resulting prefill chunk was +1,905 ms where the model's own calibrated anchor says 311 ms. + +These assert the shape of the schedule rather than exact milliseconds: that the +cost falls as context grows, that the indexer is the term which does not, and +that a schedule this cost model has not been taught is refused instead of +guessed at. +""" + +from __future__ import annotations + +import pytest + +from infera.projection.core.projection.training_config import ( + _SPARSE_ATTENTION_FLOOR, + hybrid_attention_scale, + parse_compress_ratios, +) + +from .conftest import project_spec + +# DeepSeek-V4-Pro, verbatim from its yaml: 2 HCA to bootstrap, then CSA/HCA +# interleaved, then a trailing CSA and one dense windowed layer. That is 62 +# entries against 61 declared layers, because the last one is the MTP head. +PRO_SCHEDULE = "[128, 128" + ", 4, 128" * 29 + ", 4, 0]" + + +class _Model: + """Only the fields the scale reads, at DeepSeek-V4-Pro's values.""" + + def __init__(self, **over): + self.compress_ratios = PRO_SCHEDULE + self.num_layers = 61 + self.mtp_num_layers = 1 + self.num_attention_heads = 128 + self.kv_channels = 512 + self.qk_pos_emb_head_dim = 64 + self.index_n_heads = 64 + self.index_head_dim = 128 + self.index_topk = 1024 + self.attn_sliding_window = 128 + for k, v in over.items(): + setattr(self, k, v) + + +def _only(schedule: list, **over) -> _Model: + """A model whose whole stack is ``schedule``, for isolating one branch.""" + return _Model(compress_ratios=schedule, num_layers=len(schedule), mtp_num_layers=0, **over) + + +def test_the_schedule_the_checkpoint_ships_is_read_as_written(): + sched = parse_compress_ratios(_Model()) + assert len(sched) == 62, "61 layers plus the MTP head the schedule also names" + assert sched.count(128) == 31 + assert sched.count(4) == 30 + assert sched.count(0) == 1 + + +def test_the_list_form_and_the_string_form_agree(): + as_list = parse_compress_ratios(_Model(compress_ratios=[128, 4, 0])) + as_text = parse_compress_ratios(_Model(compress_ratios="[128, 4, 0]")) + assert as_list == as_text == [128, 4, 0] + + +def test_the_speculative_head_is_not_averaged_into_the_stack(): + """The schedule's last entry is the MTP head, which this forward doesn't run. + + It is also the one dense windowed layer in V4-Pro, and the cheapest at long + context, so averaging it in would quietly make the stack look cheaper than + the 61 layers that actually run. + """ + with_head = hybrid_attention_scale(_Model(num_layers=62, mtp_num_layers=0), 130_000) + stack_only = hybrid_attention_scale(_Model(), 130_000) + assert stack_only > with_head, "dropping a near-free layer must raise the mean" + # 61 of 62 layers, so the two are close; the point is which one is used. + assert stack_only == pytest.approx(with_head, rel=0.05) + + +@pytest.mark.parametrize( + "context,ceiling", + [(8192, 0.10), (32768, 0.04), (130000, 0.02), (262144, 0.02)], +) +def test_a_compressed_stack_costs_less_of_dense_the_longer_the_prompt(context, ceiling): + """The pools grow far slower than the prompt, so the fraction has to fall.""" + assert hybrid_attention_scale(_Model(), context) < ceiling + + +def test_the_long_context_cost_is_nothing_like_the_floor_it_replaces(): + """At 130k the floor charges 9.6x what the schedule does.""" + scale = hybrid_attention_scale(_Model(), 130_000) + assert scale < _SPARSE_ATTENTION_FLOOR / 5, ( + f"schedule says {scale:.4f}; the floor charged {_SPARSE_ATTENTION_FLOOR}" + ) + + +def test_the_floor_was_set_where_it_barely_showed(): + """At 8k the two are close, which is why an 8k sweep never caught this.""" + assert hybrid_attention_scale(_Model(), 8192) > _SPARSE_ATTENTION_FLOOR / 3 + + +def _indexer_share(ctx: int, schedule=None) -> float: + build = _Model if schedule is None else (lambda **kw: _only(schedule, **kw)) + with_idx = hybrid_attention_scale(build(), ctx) + without = hybrid_attention_scale(build(index_n_heads=0), ctx) + return (with_idx - without) / with_idx + + +def test_the_indexer_is_the_term_that_keeps_growing_with_context(): + """Selected attention saturates; scoring the pool to select it does not. + + Measured on a CSA layer alone, because only 29 of the 61 layers carry an + indexer at all and averaging over the HCA layers that don't would dilute + the very thing being asserted. + """ + short, long = _indexer_share(8192, [4]), _indexer_share(130_000, [4]) + assert short < 0.15, f"indexer is a small part of an 8k step, got {short:.2f}" + assert long > 0.5, f"indexer dominates a 130k step, got {long:.2f}" + assert long > 5 * short + + +def test_over_the_whole_stack_the_indexer_grows_from_a_rounding_error_to_half(): + """Diluted by the 31 HCA layers that have no indexer, but the trend holds.""" + assert _indexer_share(8192) < 0.12 + assert _indexer_share(262_144) > 0.45 + + +def test_selected_attention_stops_growing_once_the_pool_exceeds_the_top_k(): + """Past the crossover a CSA layer reads ``top_k`` entries at any context. + + With the indexer removed the only context term left is the pool it would + have scored, so the cost per layer goes flat and the fraction falls as + exactly 1/context. + """ + csa_only = _only([4] * 8, index_n_heads=0) + # Crossover is at index_topk * compress_rate = 4096 tokens. + a = hybrid_attention_scale(csa_only, 32768) * 32768 + b = hybrid_attention_scale(csa_only, 262144) * 262144 + assert a == pytest.approx(b, rel=0.01) + + +def test_a_schedule_this_cost_model_has_not_been_taught_is_refused(): + """An unknown branch means the caller keeps its old behaviour, not a guess.""" + assert hybrid_attention_scale(_only([128, 7, 0]), 130_000) is None + + +def test_a_model_without_a_schedule_is_left_alone(): + assert hybrid_attention_scale(_Model(compress_ratios=None), 130_000) is None + assert hybrid_attention_scale(_Model(compress_ratios="[]"), 130_000) is None + + +def test_a_dense_model_is_never_costed_from_a_schedule(): + """gpt-oss declares no schedule, so nothing here can change its attention.""" + assert parse_compress_ratios(_Model(compress_ratios=None)) is None + + +def _ttft(input_len: int) -> float: + return project_spec( + model="deepseek_v4_pro", + tp=8, + ep=1, + pp=1, + attn_dp=8, + input_len=input_len, + output_len=900, + concurrency=1, + weight_dtype="fp4", + kv_cache_dtype="fp8", + prefix_cache_hit_rate=0.0, + )["ttft_ms"] + + +def test_prefill_grows_with_the_prompt_and_not_with_the_square_of_it(): + """End to end, through the real yaml, on the parallelism the fleet serves. + + TTFT is *supposed* to grow with the prompt -- the projections and the MLP + are linear in tokens and dominate here -- so the thing to check is not that + it stays flat but that nothing grows faster than the prompt does. Attention + is the only term that could, and under the schedule its pools grow 128x and + 4x slower than the prompt, so the excess over linear should be slight. + + One client, so TTFT is the prefill alone. Under load it also holds the wait + behind other clients' prefills, and that wait grows faster than the prompt + for queueing reasons, not attention ones: at 64 clients it reads 1.31x + superlinear while the prefill underneath reads 0.92x. + """ + length_ratio = 130_000 / 8192 + short, long = _ttft(8192), _ttft(130_000) + assert long > short, "a longer prompt still has to cost more" + superlinear = (long / short) / length_ratio + assert superlinear < 1.3, ( + f"TTFT grew {superlinear:.1f}x faster than the prompt did, which is the " + f"dense-attention shape the schedule is supposed to remove" + ) diff --git a/tests/unit/projection/test_prefill_curve_extrapolation.py b/tests/unit/projection/test_prefill_curve_extrapolation.py index dbc195733..245652b43 100644 --- a/tests/unit/projection/test_prefill_curve_extrapolation.py +++ b/tests/unit/projection/test_prefill_curve_extrapolation.py @@ -138,8 +138,7 @@ def test_the_held_rate_is_the_one_the_anchor_measured_at_its_longest_probe(tmp_p want_tok_s = 1000.0 / (a + b * max(PROBED)) got = _project(tmp_path, input_len=130000)["prefill_throughput_tps"] assert got == pytest.approx(want_tok_s, rel=0.05), ( - f"expected the fit's rate at {max(PROBED)} tokens (~{want_tok_s:.0f} " - f"tok/s), got {got:.0f}" + f"expected the fit's rate at {max(PROBED)} tokens (~{want_tok_s:.0f} tok/s), got {got:.0f}" ) diff --git a/tests/unit/projection/test_serving_anchor.py b/tests/unit/projection/test_serving_anchor.py index c1653a5d1..9feb70e09 100644 --- a/tests/unit/projection/test_serving_anchor.py +++ b/tests/unit/projection/test_serving_anchor.py @@ -343,6 +343,8 @@ def test_the_client_is_not_a_second_reason_a_model_cannot_be_measured(monkeypatc monkeypatch.setattr(benchmark_serving.shutil, "which", lambda _: None) assert benchmark_serving.client_kind(spec(serving_backend="sglang")) == "sglang" + # The ATOM image has no SGLang either; it vendors vLLM's legacy client. + assert benchmark_serving.client_kind(spec(serving_backend="atom")) == "atom" # An engine with no client of its own says so, rather than failing later # inside a subprocess that was never going to exist. with pytest.raises(RuntimeError, match="no load generator"): diff --git a/tests/unit/projection/test_speculative_anchor.py b/tests/unit/projection/test_speculative_anchor.py index 4454c4d93..7f0d6903e 100644 --- a/tests/unit/projection/test_speculative_anchor.py +++ b/tests/unit/projection/test_speculative_anchor.py @@ -32,23 +32,45 @@ class _Anchored: - """The two methods under test, over just the state they read. + """The methods under test, over just the state they read. Constructing a real projector needs a model preset, a parallel layout and a GEMM backend, none of which this arithmetic touches. + + ``step_ms`` stands in for the analytical model the width ratio is taken + from: a mapping from query width to what a step of that width costs. Left + unset, the model is unreachable -- which is the case on a host without the + simulator, and the case the blind fallback exists for. """ _measured_verify_step_scale = InferencePerformanceProjector._measured_verify_step_scale _spec_tokens_per_step = InferencePerformanceProjector._spec_tokens_per_step - - def __init__(self, *, spec_k: int, accept: float, anchor_spec_k: int): + _verify_width_ratio = InferencePerformanceProjector._verify_width_ratio + + def __init__( + self, + *, + spec_k: int, + accept: float, + anchor_spec_k: int, + step_ms: dict[int, float] | None = None, + ): self.cfg = SimpleNamespace( request_config=SimpleNamespace( speculative_num_tokens=spec_k, speculative_acceptance_rate=accept, + input_seq_len=8192, ) ) self._bench_spec_k = anchor_spec_k + self._meas_ref_input = 8192 + self._verify_ratio_cache: dict = {} + self._step_ms = step_ms + + def _forward_times(self, batch, q_len, phase, kv_len): + if self._step_ms is None: + raise RuntimeError("analytical backend unavailable on this host") + return SimpleNamespace(total_ms=self._step_ms[q_len]) # (k, acceptance) pairs spanning the useful range. k=1/0.84 is GLM-5.2-MXFP4's @@ -115,19 +137,51 @@ def test_perfect_acceptance_is_where_the_two_scales_agree(): @pytest.mark.parametrize("spec_k,accept", _SPEC_CASES) -def test_an_anchor_without_speculation_keeps_the_width_scaling(spec_k, accept): +def test_a_non_speculative_anchor_falls_back_to_width_with_no_model(spec_k, accept): """Nothing was folded into a non-speculative measurement, so there is - nothing to undo. + nothing to undo, and without a model to price the widening there is + nothing better than the width to charge. A single-token step cannot be turned into a verify pass by arithmetic -- which is why speculation is regime-defining and the anchor store refuses - this pairing. ``--load-benchmark`` does not refuse it, so the old behaviour - stays rather than returning a number that would look calibrated. + this pairing. ``--load-benchmark`` does not refuse it, so the blind width + stands rather than returning a number that would look calibrated. """ p = _Anchored(spec_k=spec_k, accept=accept, anchor_spec_k=0) assert p._measured_verify_step_scale() == pytest.approx(float(spec_k + 1)) +@pytest.mark.parametrize("spec_k,accept", _SPEC_CASES) +def test_a_non_speculative_anchor_prices_the_widening_from_the_model(spec_k, accept): + """With a model reachable, the verify pass costs what it costs. + + The width is an upper bound, not an estimate: a verify pass puts ``k + 1`` + query positions through the stack but reads the KV cache once for the same + sequences, and at the contexts these anchors are read at that read is most + of the step. Charging the width made a target read slower the deeper its + draft, which is backwards. + """ + width = spec_k + 1 + # A step that widens by a tenth per extra position: memory-bound decode, + # where the query dimension is nearly free. + steps = {1: 10.0, width: 10.0 * (1.0 + 0.1 * spec_k)} + p = _Anchored(spec_k=spec_k, accept=accept, anchor_spec_k=0, step_ms=steps) + scale = p._measured_verify_step_scale() + + assert scale == pytest.approx(1.0 + 0.1 * spec_k) + assert scale < width, "the modelled widening must undercut the blind width" + + +@pytest.mark.parametrize("wide_ms,expected", [(5.0, 1.0), (1000.0, 4.0)]) +def test_the_modelled_widening_stays_between_one_step_and_one_per_position(wide_ms, expected): + """A verify pass cannot beat the single-token step inside it, and cannot + cost more than running that step once per position. A model that says + otherwise is answering about a different shape than the one being priced. + """ + p = _Anchored(spec_k=3, accept=0.8, anchor_spec_k=0, step_ms={1: 10.0, 4: wide_ms}) + assert p._measured_verify_step_scale() == pytest.approx(expected) + + def test_no_speculation_leaves_the_measured_step_alone(): p = _Anchored(spec_k=0, accept=0.0, anchor_spec_k=0) assert p._measured_verify_step_scale() == pytest.approx(1.0) diff --git a/tests/unit/projection/test_tuning_disaggregated_pools.py b/tests/unit/projection/test_tuning_disaggregated_pools.py index 616659f9a..883be68d7 100644 --- a/tests/unit/projection/test_tuning_disaggregated_pools.py +++ b/tests/unit/projection/test_tuning_disaggregated_pools.py @@ -186,9 +186,7 @@ def test_every_proposed_split_fits_and_spends_the_whole_cluster(world): splits = disagg_splits(world, [1, 2, 4, 8]) assert splits, f"no legal split offered for a {world}-GPU cluster" for p_tp, d_tp, reps in splits: - assert p_tp + d_tp * reps == world, ( - f"prefill {p_tp} + decode {d_tp}x{reps} != {world} GPUs" - ) + assert p_tp + d_tp * reps == world, f"prefill {p_tp} + decode {d_tp}x{reps} != {world} GPUs" def test_a_split_is_never_proposed_that_leaves_decode_homeless(): @@ -255,15 +253,11 @@ def test_the_topology_is_scored_at_the_default_seed_budget(): """ from infera.projection.agents.tuning_agent.cli import _parse_args - default_budget = _parse_args( - ["--workload", "w.yaml", "--target-cluster", "c.yaml"] - ).seed_budget + default_budget = _parse_args(["--workload", "w.yaml", "--target-cluster", "c.yaml"]).seed_budget plan = _plan(budget=default_budget) disagg = [c for c in plan.candidates if c.disaggregate] - assert disagg, ( - f"no disaggregated candidate within the default budget of {default_budget}" - ) + assert disagg, f"no disaggregated candidate within the default budget of {default_budget}" def test_the_reserved_slots_do_not_displace_the_head_of_the_plan(): @@ -312,9 +306,7 @@ def test_the_budget_is_spent_on_different_splits_rather_than_one(): """ plan = _plan(_Cluster(2, 8), budget=12) shapes = { - (c.prefill_tp, c.decode_tp, c.decode_replicas) - for c in plan.candidates - if c.disaggregate + (c.prefill_tp, c.decode_tp, c.decode_replicas) for c in plan.candidates if c.disaggregate } assert len(shapes) > 1, f"only one pool shape priced: {shapes}" @@ -367,7 +359,8 @@ class _AgentCfg: assert flag in cmd, f"{flag} never reaches the projector" known = { - s for a in build_parser()._subparsers._group_actions[0].choices["inference"]._actions + s + for a in build_parser()._subparsers._group_actions[0].choices["inference"]._actions for s in a.option_strings } for tok in cmd: diff --git a/tests/unit/projection/test_uniform_sparse_attention.py b/tests/unit/projection/test_uniform_sparse_attention.py new file mode 100644 index 000000000..26f01c663 --- /dev/null +++ b/tests/unit/projection/test_uniform_sparse_attention.py @@ -0,0 +1,131 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""A model that selects on every layer still has to pay for choosing. + +DeepSeek-V4 records a per-layer compression schedule and is costed from it. +GLM-5.2 and MiniMax-M3 do not: they keep one window on every layer and differ +only in how wide the indexer is and how many layers compute one. Having no +schedule, both fell through to a top-k floored at 0.15, which on the agentic +corpus -- served between 136k and 365k tokens -- is a constant: nine to +twenty-seven times the selection it names, and nothing at all for the +indexing it stands in for. + +The two terms are the point. Selection reads ``topk`` entries however long the +prompt is, so its share falls away like ``topk / context``. Indexing scores +every token in the pool on the layers that carry an indexer, so it grows with +context exactly as dense attention does and its share is flat. A single +multiplicative floor can be either of those and is neither. + +These assert the shape -- that the total falls with context, that what it +falls towards is the indexing term rather than zero, and that the two models +are separated by geometry the floor could not see -- rather than exact +milliseconds. +""" + +from __future__ import annotations + +import pytest + +from infera.projection.core.projection.training_config import ( + _SPARSE_ATTENTION_FLOOR, + sparse_indexer_decode_overhead, + uniform_sparse_attention_scale, +) + + +class _GLM: + """Only the fields the scale reads, at GLM-5.2's published values.""" + + def __init__(self, **over): + self.num_layers = 78 + self.num_attention_heads = 64 + self.multi_latent_attention = True + self.kv_lora_rank = 512 + self.qk_pos_emb_head_dim = 64 + self.sparse_index_head_dim = 128 + self.sparse_index_n_heads = 32 + # IndexShare: 21 of the 78 layers compute an index, the rest reuse it. + self.sparse_index_layers = 21 + self.__dict__.update(over) + + +class _M3: + """MiniMax-M3: GQA rather than MLA, and a four-head indexer.""" + + def __init__(self, **over): + self.num_layers = 60 + self.num_attention_heads = 64 + self.multi_latent_attention = False + self.kv_channels = 128 + self.sparse_index_head_dim = 128 + self.sparse_index_n_heads = 4 + self.sparse_index_layers = 57 + self.__dict__.update(over) + + +def test_selection_falls_away_as_the_prompt_grows(): + """Reading 2048 of 8k is a quarter of the work; 2048 of 365k is not.""" + near = uniform_sparse_attention_scale(_GLM(), 8192, 2048) + far = uniform_sparse_attention_scale(_GLM(), 364672, 2048) + assert far < near + # The floor it replaces cannot do this: past 14k it is the same at both. + assert max(_SPARSE_ATTENTION_FLOOR, 2048 / 32768) == _SPARSE_ATTENTION_FLOOR + assert max(_SPARSE_ATTENTION_FLOOR, 2048 / 364672) == _SPARSE_ATTENTION_FLOOR + + +def test_what_it_falls_towards_is_the_indexer_not_zero(): + """Selection vanishes at long context; the cost that remains is indexing.""" + scale = uniform_sparse_attention_scale(_GLM(), 4_000_000, 2048) + assert scale > 0.0 + # Far past the window, selection is negligible and what is left is the + # indexer's own share of a layer -- flat, because it grows with context + # exactly as the dense attention it is measured against does. + wider = uniform_sparse_attention_scale(_GLM(), 16_000_000, 2048) + assert wider == pytest.approx(scale, rel=0.05) + + +def test_the_floor_over_reads_the_corpus_this_model_serves(): + """At the contexts GLM-5.2 was served, the floor is the wrong constant.""" + for ctx in (131072, 262144, 364672): + priced = uniform_sparse_attention_scale(_GLM(), ctx, 2048) + assert priced < _SPARSE_ATTENTION_FLOOR + # And the gap widens with context, which is what a constant cannot track. + near = _SPARSE_ATTENTION_FLOOR / uniform_sparse_attention_scale(_GLM(), 131072, 2048) + far = _SPARSE_ATTENTION_FLOOR / uniform_sparse_attention_scale(_GLM(), 364672, 2048) + assert far > near + + +def test_indexer_width_separates_two_models_the_floor_could_not(): + """Thirty-two heads against four is a difference; 0.15 against 0.15 is not.""" + glm = uniform_sparse_attention_scale(_GLM(), 364672, 2048) + m3 = uniform_sparse_attention_scale(_M3(), 364672, 2048) + assert glm != pytest.approx(m3, rel=0.05) + + +def test_index_share_is_cheaper_than_indexing_every_layer(): + """Reusing one layer's index across the next three is 21 layers, not 78.""" + shared = uniform_sparse_attention_scale(_GLM(), 131072, 2048) + every = uniform_sparse_attention_scale(_GLM(sparse_index_layers=78), 131072, 2048) + assert shared < every + + +def test_a_model_with_no_indexer_geometry_is_left_alone(): + """Declaring a window but no indexer is not enough to be priced here.""" + assert uniform_sparse_attention_scale(_GLM(sparse_index_n_heads=0), 131072, 2048) is None + assert uniform_sparse_attention_scale(_GLM(sparse_index_head_dim=0), 131072, 2048) is None + assert uniform_sparse_attention_scale(_GLM(), 131072, 0) is None + + +def test_a_fused_selection_leaves_the_calibrated_decode_step_alone(): + """Decode's dense charge is measured; the fused path must not disturb it.""" + assert sparse_indexer_decode_overhead(_GLM(), 2048, 1.0) == 0.0 + + +def test_an_unfused_selection_costs_decode_a_bounded_amount(): + """A slower top-k kernel is a percentage of the step, not a multiple.""" + over = sparse_indexer_decode_overhead(_GLM(), 2048, 3.3) + assert 0.0 < over < 0.10 + assert sparse_indexer_decode_overhead(_GLM(), 2048, 8.0) > over