docs/fixed_budget_decode_audit.md

Fixed-budget decode audit — 2026-09-11

Which implementation was tested

The existing optimized KVPath-shared is DecodeSession(..., backend="triton_grouped"): projected 128-dimensional base plus 32-dimensional per-layer delta, a fused split-score kernel, static append cache and growing CUDA graphs. It does not perform matrix absorption. inference/decode.py explicitly rejects kvpath_absorb_base=True, and supports only gqa and projected-base kvpath-shared.

The separate kvpath_absorb model uses eager PyTorch absorption and caches anchor hidden states: 4608 values/token, not the fixed 6144 budget. mla also uses eager absorption, with 448-dimensional latent + 64-dimensional RoPE, and has 6144 values/token. There is no optimized MLA session or absorbed-KVPath-shared Triton kernel in the current project-local sources.

The quoted prior-agent handover describes GQA versus MLA. Its warning about unequal backend optimization is correct; it does not identify another absorbed KVPath-shared kernel.

Setup and interpretation

• NVIDIA A100-PG509-200 40 GiB, BF16, context 2048, batches 1/8/32/64/128.
• GPUs 0–3 were reported idle before the tests. Full-model measurements used GPU 0; group attention prototypes used GPU 1; diagnostic variants used GPU 2; the hardware-counter attempt used GPU 3. Other users' processes were untouched.
• GQA and KVPath-shared: exactly 98,710,016 parameters, 6144 cache values/token. MLA: 98,703,104 parameters, 6144 values/token.
• Random initialized weights, reproducible prompt/decode IDs (seed 123). These are implementation tests, not quality tests or trained-checkpoint tests.
• Full-model runs perform actual chunked prefill (chunk 32), then 8 growing decode steps, with 8 warmup steps. Static capacity is 2056 tokens. CUDA-event loop times exclude prefill/JIT/capture/sampling, but can include CPU submission gaps. Eager has no graphs; optimized runs use graphs.
• Five timed rounds for the initial static/optimized sweep, seven for the final eager sweep. All rounds are preserved; variants are sequential in seeded shuffled order. This is not a randomized interleaved full-model trial, and small differences should not be overinterpreted.
• Absolute PyTorch allocated peaks include weights and cache (graph pools for graphed variants). Reserved memory is also saved. These are not NVML totals.

Two new standalone drivers leave the original batch<=8/shared-GPU guards and production model/kernel files unchanged. The user explicitly requested batch128 on idle devices. The full-model driver caps PyTorch allocation at 16 GiB; the group prototype driver caps it at 8 GiB and checks free memory first.

Optimized full-model results

All entries below are the existing grouped Triton implementations with graphs, split size 128, at identical parameter and cache budgets. Units: ms/decode step.

Batch	GQA	KVPath-shared	KVPath-shared/GQA
1	0.694	0.738	1.06
8	1.089	1.308	1.20
32	1.298	2.206	1.70
64	1.998	3.899	1.95
128	3.289	6.916	2.10

At batch128 GQA FlashAttention measured 3.724 ms, so the faster measured GQA control was grouped Triton. Absolute decode allocated peaks were 3305.1 MiB for GQA grouped and 3314.2 MiB for KVPath-shared grouped. Their actual static-cache allocations are equal, including unused capacity.

Sources: inference/results/audit_batch_sweep_ctx2048.json and inference/results/audit_b128_ctx2048.json. The dynamic-cache peak entries in these initial files are superseded by the corrected eager sweep below.

Like-for-like eager absorption comparison

Final corrected batch128 results, all eager + dynamic cache on GPU 0:

Model	Cache values/token	ms/step	Absolute decode peak, MiB
GQA	6144	17.406	3413.9
KVPath-shared projected	6144	25.625	3551.8
KVPath-shared absorbed	4608	23.265	2775.3
MLA absorbed	6144	22.492	3510.7

Absorption lowers KVPath-shared cache capacity and measured peak in this implementation; it did not beat GQA in this eager run. This does not establish how equally optimized absorbed implementations would perform. Do not compare these eager times against optimized GQA as if only architecture differed.

The corrected sweep covers all five batch sizes: inference/results/audit_eager_corrected_ctx2048.json (measurement_version=2). Eager times do not grow proportionally with batch; small/medium-batch timings are visibly affected by submission overhead and run-to-run system conditions. Batch size alone does not prove that attention HBM bandwidth dominates.

Measurement correction: the initial dynamic-cache driver reset peak stats before the entire rounds loop. Its harness-only rewind can clone several cache tensors while retaining the old dictionary, inflating the recorded peak. Version 2 resets stats after each rewind, before each timed decode sequence, then takes the maximum actual decode peak. Rewind was always outside timing. Use only version 2 for dynamic-cache peak claims. The old files remain available as development records rather than silently deleting inconvenient results.

What theory can and cannot establish

For a four-layer group, per historical token, counting the required input elements with once-per-KV-head reuse within each layer:

• GQA: 4 * 2 * 4 * 64 = 2048.
• Current KVPath-shared: 4 * 2 * 4 * (128 + 32) = 5120.
• Joint group scan: 2 * 4 * 128 + 4 * 2 * 4 * 32 = 2048.

These are logical input volumes, not measured DRAM bytes. L2 reuse, transactions, scratch, cache append, GEMM dispatch and kernel launch costs change the observed behavior. Likewise useful mathematical FLOPs are not the same as issued/padded tensor-core work. The current grouped kernel has two useful query rows per KV head but pads its row tile to at least 16. Grouping four layers exposes eight useful query rows for the shared base, improving that utilization in addition to shared input loads.

Nsight Compute was attempted using actual DRAM/L2 byte and throughput metrics. It failed with ERR_NVGPUCTRPERM: the current user cannot access performance counters. No GPU permissions/settings were changed to bypass this restriction. See inference/results/audit_ncu_gqa_b128.csv. No hardware bandwidth number, L2 hit rate, or definitive memory-bound classification is claimed in this audit.

The original benchmarking/bench.py computes achieved_bw from assumed bytes divided by time. Its automatic "kernel-launch bound" message is therefore not a measured bottleneck classification. It also reports incremental peaks.

A/B feasibility: implemented group-attention prototypes

Source: benchmarking/attention_audit.py. Production architecture is unchanged. The benchmark allocates identical physical KV storage for GQA and KVPath-shared: 2048 BF16 elements/token/group (1 GiB at batch128/context2048). A/B use the same base128 + delta32 tensors as current KVPath-shared, with no latent compression, unused padding cache, token omission or reduced precision.

A: group-entry queries and one joint history scan

The prototype kernel loads a base K/V tile once for four layers' available queries, adds each layer's own delta logits before its own softmax, and computes four separate outputs. It was compared against an independent FP32 explicitly concatenated attention oracle; BF16 max absolute errors in the tested batches were below 0.000465.

The proposed architecture changes follower query inputs to the group-entry hidden state. Historical delta already exists; current-token delta can still come from each sequential layer. Its current-token contribution is merged using alpha = sigmoid(current_score - history_logsumexp) and output = (1-alpha)*history_output + alpha*current_value. An independent CPU FP64 test verifies this merge identity to 1e-12 tolerance.

The kernel benchmark excludes query projection, current-token append/merge, output projection, residuals and MLP. It does not implement or time a complete modified A model. It retains the current mathematical QK/PV operation count for the supplied queries; it changes data reuse and tensor-core tile utilization.

B: one shared base retrieval plus four delta attentions

The prototype runs one base128 attention and four delta32 attentions, each with its own softmax, using existing grouped kernels. Separate-softmax outputs are checked against independent FP32 references. All cache entries remain used.

B changes model semantics: delta no longer reweights the shared base history inside the same softmax. Output fusion/gating, projections and MLP are excluded from timing; no trained model quality or complete B decoder has been tested.

Measured group-attention times

Nine randomly interleaved rounds, 20 graph replays/round, fixed history length 2048, split128. These are four-layer history-attention-only milliseconds, not full-model growing decode times.

Batch	GQA four layers	Current KVPath-shared four layers	A joint history	B base once + deltas
8	0.0817	0.1876	0.0810	0.0961
32	0.2387	0.5437	0.2333	0.2457
128	0.8559	2.0831	0.9148	0.8586

At batch128, A is ~2.28x faster than the current KVPath-shared history-attention prototype, B ~2.43x, but A is ~6.9% slower than the GQA kernel control and B approximately ties it. These establish implementation feasibility and a route to reducing KVPath-shared's current gap, not an end-to-end win over GQA. Extra per-token summaries and split reduction scratch must be included before making any A/B peak-memory claim.

Raw samples: inference/results/audit_ab_attention_ctx2048.json.

Reproduction

Run from the repository root on an explicitly selected idle GPU:

bash benchmarking/run_fixed_budget_audit.sh 0 \
  --batch 1,8,32,64,128 \
  --variants gqa:flash,gqa:triton_grouped,kvpath-shared:triton_grouped \
  --output inference/results/audit_optimized_reproduction.json

bash benchmarking/run_fixed_budget_audit.sh 0 --rounds 7 \
  --variants gqa:dynamic_eager,kvpath-shared:dynamic_eager,kvpath_absorb:dynamic_eager,mla:dynamic_eager \
  --output inference/results/audit_eager_reproduction.json

bash benchmarking/run_fixed_budget_audit.sh 1 attention \
  --output inference/results/audit_ab_reproduction.json

The user's remembered batch128 win has not been reproduced with the current sources and tested paths. Its original command/output or exact code revision is still needed to identify the discrepancy; no claim is made that the earlier observation was impossible.