[wangye805]

What

Phase A1 of the projected-performance estimator (the throughput/memory layer on top of Wen's parameter oracle). Replaces scaleout.py's flat NON_PARAM_HEADROOM_GIB = 82.0 with a component-wise activation model as f(seq, mbs, CP, recompute, layers_on_stage, mla_on_stage, K).

Two files, both under kimi_k3/tools/:

• activation_mem_model.py -- pure-Python model, no torch. Each activation term is one function with its formula + scaling + a MEASURED/ESTIMATED tag.
• activation_mem_report.py -- driver importing the repo's own mem_budget.params_per_gpu + scaleout.aligned_layout (the reconciliation point). Run: python -m kimi_k3.tools.activation_mem_report.

Why the flat 82 GiB is wrong

It is a single G28 point: 4 L on one node, no recompute, K=1. It is reused flat for every 93 L candidate, blind to seq, mbs, CP, recompute, and layers-per-stage.

Validation (why the numbers are trustworthy)

• Reproduces G46 to the decimal: params/gpu = 19.25 B, state = 141.1 GiB (dist_muon@dp8).
• Reproduces G6 AttnRes payload exactly: first boundary (1+1) x 8192 x 7168 x 2 = 224.0 MiB.

Findings

1. Under production recompute=full, the 93 L per-stage peak is ~59-65 GiB, below the flat 82 -- so the flat constant is conservative for the recompute regime and the 28-node floor (G46) is not threatened by activation. The binding stage is the last one (LM-head logits), not the mid-pipeline payload stage.
2. But mbs=2 @ seq8192 or seq=16384 pushes a single stage over 82 GiB -- there the flat constant instead under-provisions.
3. The back-solved 37.5 GiB residual (frag + NCCL) is now the load-bearing unknown. One two-seq trainer mem report splits constant from seq-scaling.
4. Corrected an earlier claim: MLA [B,nh,S,S] score is query-chunked, so it is linear in S in memory (quadratic stays in time only). MoE dispatch is the larger transient at production width.

Not in scope / next

• Not yet folded into mem_budget.py/scaleout.py -- kept standalone for review first.
• Residual calibration needs the mi355 node (a two-seq mem report).
• Next: A2 kernel-shape enumerator.

Base branch dev/yewang12/kimi-k3-proj_perf_estimator is Wen's latest a2c2ad1 with no extra commits, so this diff is exactly the two new files.

Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com

[wangye805]

[Agent reply] Pushed v2 (b077e32) implementing all 7 review threads -- the model is now fully bottom-up, no measured value is a model INPUT. Measurements moved to crosscheck_* helpers (validation only), and the gaps are reported as findings rather than absorbed:

• state bytes/param: analytic 2 + 4 + 8/dp (mem_budget.OPTIMIZER_BYTES_PER_PARAM, = 7.00 @ dp8) vs G5 measured 7.87 -> gap +0.87
• AttnRes mixer: derived from the (K+1)-slot fp32 upcast shape (2.58 GB fwd) vs G6 7.1 GB -> gap x2.75, flagged as uncounted live fp32 temporaries in the CK mixer (open item: source read, NOT tuned away)
• 4 L headroom: analytic 38.8 GiB vs G28 82.0 GiB -> gap +43.2 GiB not-yet-itemized (candidates: aiter/hipBLASLt workspaces, deep_ep buffers, CUDA context)
• MLA now flash-style by default (no [B,nh,S,S] score, linear in S); QK-clip rescore is a default-off optional term
• calibrate_residual removed; fragmentation is an explicit 7% of the itemized set
• StageSpec -> PipelineStageInputs

The +43 GiB 4 L gap is the honest headline: the bottom-up itemization only explains ~39 of the measured 82 GiB, so there is real memory we have not modeled yet. That is now visible instead of hidden inside a back-solved residual. For the StageSpec rename I went with PipelineStageInputs; say the word if you prefer StageGeometry.

[wangye805]

[Agent reply]

Empirical 4L memory attribution -- the +43 GiB "unexplained headroom" is fully accounted for (nothing unmodelled)

Method (Ye's): capture the real PyTorch allocator trace of the 4L proxy, sweep (seq, mbs, recompute), and cross-reference every bucket against the bottom-up model. Two new tools on this branch: mem_probe.py (history capture around one warm step, no profiler active) and mem_snapshot_attrib.py (buckets segments by python alloc frames, regresses vs S*B). Commit c0dc404.

The +43 GiB was two confounds on a single-point anchor, not unmodelled memory

(1) 44.79 GiB of LAZY Muon state. The G28 "82 GiB headroom" was defined as peak - after_optimizer. But after_optimizer (135.93 GiB) is read before the first step(), and Muon allocates its master/momentum lazily on that first step -- so it undercounts the true resident (181.18 GiB) by exactly 44.79 GiB. That delta was silently folded into what the model was asked to explain as activation. The resident set reconciles exactly:

param_grad_buf 60.75 + opt_state 44.80 + muon_state 44.79 (lazy) + bf16_cast 30.38 = 181.18 GiB

(2) An eager-KDA backend artifact. The proxy ran --kda-backend eager, whose chunked linear-attention recurrence materializes large per-chunk state. Peak-decode shows the transient is KDA-dominated and linear in seq (25 -> 50 GiB as 2048 -> 4096). The analytic model targets the fused fla production path, which does not materialize those chunks.

Direct fla measurement closes the gap

Re-ran the identical sweep with --kda-backend fla (fused production path). The transient collapses and stops scaling with seq:

transient = peak - resident   (rank0, rc=full, mbs1)
seq         1024     2048     4096
eager      17.97    31.34    58.07     slope 0.0137 GiB/tok (KDA chunk-state artifact)
fla         9.71     9.70     9.70     FLAT -- matches analytic 9.1 GiB to 0.6 GiB

The +43 GiB residual closes to ~0. The analytic activation model was correct for the production path; the eager proxy was the confound.

Root cause of the ~9.7 GiB fla transient: it is NOT activation

Replaying device_traces to the global peak-live point names it exactly:

seq4096 peak-live 192.60 GiB, transient blocks:
   7.98 GiB   layer_wise_optimizer.py:223  _allgather_helper   (per-rank recv buffers)
   1.73 GiB   layer_wise_optimizer.py:207  _allgather_helper   (local shard flatten)
   ------
   9.71 GiB   = the entire fla transient

At the peak essentially no fwd/bwd activation is live (kda.py:160 forward 0.07 GiB). Source: layer_wise_optimizer.allgather_params._allgather_helper -- after step() each rank flattens its own param shard and allocates a receive buffer for every other rank's shard, all-gathering the full un-sharded bf16 params, then frees. It is parameter-sized and token-independent (hence flat vs seq). Once fla makes activation cheap, the per-step memory peak is optimizer-step-bound, not activation-bound.

Consequences for the model

• crosscheck_headroom now compares analytic activation against the measured transient (peak - resident) in the correct rc=full regime, anchored on the fla production transient (9.1 vs 9.7, gap +0.6); the eager number is available only to expose the artifact.
• The dist-optimizer all-gather buffer must be modelled as a separate step-time transient (size ~ the largest single all-gather = the bf16 dp_cp param group per rank), not folded into the activation slope. In 93L it scales with the sharded param count per rank (dp_cp / expert-dp), so it belongs in the state/optimizer oracle and must be tracked as dp/ep change.
• No calibrate_residual / back-solved catch-all was added; every term is either derived from shapes or a named measured physical quantity.

[wangye805]

What is your payload vs mixer? Can you change to specific input/output or internal state naming

[wangye805]

[Agent reply] Done in 673f2ca. Renamed to input/output vs internal-state: attn_res_payload_bytes -> attn_res_io_bytes (the block's bf16 INPUT/OUTPUT, the (1+K) packed residual-stream slots that cross a boundary) and attn_res_mixer_bytes -> attn_res_internal_bytes (the live fp32 temporaries INTERNAL to the mix op). Constant MIXER_ACCUM_FP32_TEMPS -> ATTN_RES_INTERNAL_FP32_TEMPS, crosscheck_mixer -> crosscheck_attn_res_internal, and the part keys are now attn_res_io / attn_res_internal.

[wangye805]

Is this for q, k, v tensor in fwd? If so, it should be activation. But you name it workspace? What workspace is it for?

[wangye805]

[Agent reply] Correct -- it is the per-head Q/K/V/O attention activation the flash/CK/AITER kernel reads and writes, so it IS activation, not a workspace. That was a bad name. Renamed flash_ws -> qkvo_activation in 673f2ca and fixed the docstring to say 'This is activation, not a scratch workspace.' The full [B,nh,S,S] score is never materialized (flash-style), so it is linear in S.

[wangye805]

Let's get rid of this estimated thing. For all things we cannot explain, we should report a gap for future investigation, instead of inventing some factors to make us feel like we understand but we actually not

[wangye805]

[Agent reply] Agreed. Removed the FRAG_FRACTION=0.07 factor entirely in 673f2ca -- stage_activation_gib no longer takes frag_fraction and TOTAL is now the plain sum of the itemized components. Anything we cannot derive from shapes is surfaced as an explicit GAP by the crosscheck helpers, never absorbed by an invented factor. (See the follow-up note below on the known unmodeled MLA buffers -- softmax LSE, dq_acc -- which we deliberately leave as a measured gap rather than guess.)