| # 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 | |
| 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. | |
Xet Storage Details
- Size:
- 10.7 kB
- Xet hash:
- d9f384ca450d63ab805d3db763074ab4a2437329ed2e9267d79ba5cc5df9c754
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.