Repository navigation
[DSV4][ROCm] Add contract-bound C4 KV compressor (task 19) - #618
inaniloquentee wants to merge 2 commits into
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the
 |
|
🔍 OpenCodeReview found 11 issue(s) in this PR.
📄
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| def timed(fn, repeats): | ||
| for _ in range(5): | ||
| fn() | ||
| torch.cuda.synchronize() | ||
| start = time.perf_counter() | ||
| for _ in range(repeats): | ||
| fn() | ||
| torch.cuda.synchronize() | ||
| return (time.perf_counter() - start) * 1e6 / repeats |
There was a problem hiding this comment.
This times host wall clock around an async device loop only. Peers in this repo report GPU-event and wall time together (see benchmarks/backends/rocm/benchmark_rocm_singleton_lse.py::_elapsed_ms). Without the event pair, forward_us / provider_forward_us are dominated by Python-side work in the timed region: torch.empty + two zeros_like in AiterC4Backend.forward, the (blocks, 64, 512) BF16 cache allocation, the clone() of context_values, the Python loop over logical_keys, and (for bind_us) a blocking token_ids.cpu().tolist(). Recording kernel time alongside wall time makes it explicit which part is host/launch overhead, so the JSON is not read as kernel latency.
Also, this torch.cuda.synchronize() and the device="cuda" fixtures in main() are unconditional, so the script cannot run (or clearly skip) on a non-CUDA/ROCm host; the repo's guarded pattern is if device.type == "cuda" (benchmark_rocm_logp.py::_device_synchronize).
Suggestion:
| def timed(fn, repeats): | |
| for _ in range(5): | |
| fn() | |
| torch.cuda.synchronize() | |
| start = time.perf_counter() | |
| for _ in range(repeats): | |
| fn() | |
| torch.cuda.synchronize() | |
| return (time.perf_counter() - start) * 1e6 / repeats | |
| def timed(fn, repeats): | |
| for _ in range(5): | |
| fn() | |
| torch.cuda.synchronize() | |
| start = torch.cuda.Event(enable_timing=True) | |
| end = torch.cuda.Event(enable_timing=True) | |
| wall = time.perf_counter() | |
| start.record() | |
| for _ in range(repeats): | |
| fn() | |
| end.record() | |
| end.synchronize() | |
| return ( | |
| start.elapsed_time(end) * 1e3 / repeats, | |
| (time.perf_counter() - wall) * 1e6 / repeats, | |
| ) |
| if group > 0: | ||
| grad_out = torch.cat((torch.zeros_like(y[:1]), grad[index : index + 1].float())) | ||
| dx, da, dn = torch.autograd.grad(y, (context, a, n), grad_outputs=grad_out) | ||
| else: | ||
| dx, da, dn = torch.autograd.grad( | ||
| y, (source, a, n), grad_outputs=grad[index : index + 1].float() | ||
| ) |
There was a problem hiding this comment.
The group > 0 branch builds grad_out with a leading zero row plus the incoming grad, i.e. it assumes compressed_compact_graph returns two output rows per group (a first row to be masked like the previous-window slot). The group == 0 branch passes grad[index:index+1] (one row) for the truncated source. In the CPU reference (c4_helpers.context_reference) the functional returns a single pooled row and the whole y receives the incoming grad, with the previous-window masking handled by the -inf gate rather than by zeroing a gradient row. If the pinned mLite functional also emits one row per group, grad_out would be length 2 against a length-1 y and torch.autograd.grad raises, or (if it emits 2 rows) the first branch silently drops half of the group output's gradient. The two branches are mutually inconsistent about y's output arity — please verify against the pinned backward.py and make both branches use the same output-row convention.
|
@inaniloquentee The automated OCR review has completed. Please review and address all automated review bot findings, then push your updates. Maintainers will begin manual review after the bot findings have been addressed. |
Flink-ddd
added
platform: rocm
| for index, (_request, group) in enumerate(bound.logical_keys): | ||
| context = bound.context_values[index].float().detach().requires_grad_() |
There was a problem hiding this comment.
For group == 0, source = context[4:] is a view of the leaf tensor context, and torch.autograd.grad(y, (source, a, n), ...) returns the gradient with respect to that view's shape (4 rows). That is then padded with torch.zeros_like(context[:4]) to 8 rows. This relies on the predecessor halves being exactly zeroed for group 0 (context[:, :4] is zeroed at bind time). If the predecessor rows were ever non-zero for a group-0 binding, the returned dx would silently drop the predecessor contribution to ape/norm (the predecessor rows' effect on the softmax window) because only source is differentiated there. Cross-checking against c4_helpers.context_reference, group 0 uses torch.zeros_like(previous) for content and -inf gate — i.e. the predecessor rows are constant, so they contribute zero gradient through context but still contribute to da/dn in the reference implementation (which differentiates the full 8-row context). Here the VJP differentiates only 4 rows for group 0, so da/dn for the group-0 window differ from the reference: the reference's ape/norm gradient includes the (zero-content, -inf-gated) predecessor rows' weight dependence, whereas source=context[4:] excludes them by construction. This produces silently different APE/RMS gradients for the first group of every request; it is only masked in tests because group 0 has zeroed content and the accumulated test tolerances are loose (rtol=4e-3). Recommend differentiating the full context for group 0 (masking the same way forward does) and taking dx[4:], matching the reference semantics.
|
Pushed Correctness fixes: use aiter's native negative-position padding mask instead of the old finite gate floor; validate predecessor context even when not emitted; preserve real FP64 in the independent oracle (the old helper downcast internally, so the previous all-FP64 description was inaccurate). New environment: Torch 2.11.0+gitd0c8b1f, HIP 7.2.53211, Triton 3.6.0, aiter 0.1.16.post3, FlyDSL 0.2.0. Exact provider source hashes distinguish the admitted ABI; no silent fallback. Validation: 85 C4 tests passed and 227 full targeted tests passed, zero skips, one warning. Lint/format passed; benchmark rerun. Details and limitations: |
Summary
Implements roadmap item 19, the contract-bound DSV4 main-KV C4 compressor on the
dsv4-rocmintegration branch.emit_group_ids.k_per_block=64).flydsl_fused_compress_attnwithoverlap=True, ratio 4, BF16 output and fixed one-wave arithmetic; no new numerical GPU kernel was added.compressed_compact_graphfor VJP and RL KernelCanonicalBackwardSessionfor fixed-order parameter-row reduction.Validation on MI300X
Environment: AMD Instinct MI300X VF
gfx942, ROCm/HIP 7.14.60850, PyTorch2.12.0+rocm7.14.0a20260608, Triton 3.7.0, FlyDSL 0.3.4.1, aiter commit326cffa38c9eda16bef19375ba55334e4d3acd59, mLite commit327a238239448de6771beea185dc6902a8de3eaa.Detailed evidence and benchmark:
reports/kv-compressor-c4-mi300x-20261028.md.Explicit non-claims
This PR does not claim RCCL/multi-GPU CP acceptance, Graph replay, live cache commit, prefill/decode framework integration or full-model acceptance. Those remain separate roadmap work.