Repository navigation
[DSV4][ROCm] Reuse native operators for Indexer state and MXFP4 cache (task 22) - #626
inaniloquentee wants to merge 6 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 9 issue(s) in this PR.
📄
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| dummy = torch.zeros((1, 128), dtype=torch.bfloat16, device=self.device) | ||
| self._rotate(dummy) |
There was a problem hiding this comment.
Constructing the backend launches a native GPU kernel (self._rotate(dummy)) during __init__. C4Compressor.__init__ may be called from inside a CUDA graph capture region or before the runtime installed the intended device as current; the parent eager() graph-capture guard is only applied on bind/__call__, not here. A construction inside capture would pollute the captured graph with a JIT-bound kernel launch, and construction on a different device index than torch.cuda.current_device() at import time would warn/behave inconsistently. Consider guarding construction the same way eager() does, or deferring the priming launch to the first forward.
| unrotated_grad = self._rotate(grad.to(torch.bfloat16)) | ||
| return super().vjp(bound, ape, norm, cos, sin, unrotated_grad, contract) |
There was a problem hiding this comment.
The incoming gradient is unconditionally truncated to BF16 before the rotation and before the FP32 mLite VJP. The parent AiterC4Backend.vjp fuses grad.float() into the FP32 reference arithmetic (grad[index:index+1].float() on the padded grad_out), so this cast happens earlier than in the main-C4 path and silently discards FP32 precision for any caller passing an FP32 dy. No validation exists for the FP32 case (ROCm tests only pass .bfloat16() grads), so numerical regressions in the FP32 path would pass unnoticed. Either enforce BF16 as a documented precondition of vjp or keep the accumulated gradient dtype and cast only where the native rotation requires it.
| self.implementation_fingerprint = hashlib.sha256( | ||
| json.dumps( | ||
| { | ||
| "shared": self.implementation_fingerprint, | ||
| "indexer_model": source_hash(__file__), | ||
| }, | ||
| sort_keys=True, | ||
| ).encode() | ||
| ).hexdigest() |
There was a problem hiding this comment.
The Indexer fingerprint is re-derived from source_hash(__file__) of this bridge module, but the shared parent fingerprint already mixes source_hash(kv_compressor_c4.py) and the backend provenance. The Indexer backend's own provenance (indexer_sources, indexer_bridge, rotation) is stored on self.backend.provenance but is not included here beyond whatever the parent AiterC4Backend.__init__ provenance captured. Since AiterIndexerC4Backend.__init__ mutates self.provenance after super().__init__() (which builds the parent identity dict), any consumer that pins expected_implementation will not detect a change to the indexer provider sources (e.g. a different aiter.ops.quant revision) without also re-deriving via the backend provenance. Verify the parent C4Compressor.__init__ reads backend.provenance after the indexer subclass finishes mutating it; if it does not, the fingerprint can miss provider changes.
| def __call__(self, bound, ape, norm, cos, sin, *, version, parameter_id=""): | ||
| out = super().__call__( |
There was a problem hiding this comment.
The quantized columns are produced outside the recorded autograd path and outside the canonical backward budget: backend.quantize is invoked with torch.no_grad() on out.values.detach() after super().__call__ returns. If the provider returns a per-column-major or per-group scale layout (the FP8 path passes column_major_scales=False but nothing here verifies what the installed vLLM actually returns), the require(...) below raises C4Error only after the parent already performed a full forward and registered canonical backward slots, so the failure surfaces as a raising exception mid-session with dangling registered rows. More importantly, the empty-input branch of quantize returns 2-D (0, 1)/(0, 4) scales while the non-empty FP8 branch returns whatever per_token_group_quant_fp8 yields; if that is 1-D (rows,), the scales.shape == (rows, 1) check fails only for non-empty rows. The empty path is the only one exercised by the CPU tests, so this inconsistency is not covered. Normalize the returned scale to (rows, 1)/(rows, 4) explicitly at the backend boundary.
| assert train is rollout | ||
| op = operator() | ||
| assert train( | ||
| contract=op.contract, | ||
| backend=ReferenceIndexerC4Backend(), | ||
| expected_implementation=op.implementation_fingerprint, | ||
| ) |
There was a problem hiding this comment.
expected_implementation is validated against a freshly constructed ReferenceIndexerC4Backend(), which has no provenance attribute (the parent falls back to {"test_provider": type(backend).__name__}). This makes the fingerprint comparison trivially reproducible on CPU but does not exercise the case where the expected value is pinned from a real GPU backend artifact; a backend-provenance regression on the ROCm path would not be caught by this test. The negative case below (expected_implementation="old") only verifies the mismatch message. If the contract intends expected_implementation to pin a real deployment fingerprint, add a GPU-side check (or document the CPU-only scope).
|
@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
Roadmap task 22 ??independent Indexer state and MXFP4 cache
Dependencies: task 20 / #624 and task 21 / #625 (and their #618 prerequisite). Target remains
dsv4-rocm. Task-22-only commit:f011f5206255125e2d8009c73e599fe11d8d8728; the base diff includes prerequisites until they merge. Refs #512.Reuse
Cache contract
Unshuffled row-major uint8 E2M1
[pages,64,64], low nibble first, and uint8 E8M0[pages,64,4]group-32 scales. Inactive/released payloads are zero and scales are 127; validity comes from the committed cursor, not the scale byte. Padded/partial rows never write persistent cache. A failure between payload/scale writes poisons the arena without publishing a cursor.Real MI300X acceptance
Docs:
docs/models/deepseek_v4/indexer-state.md.Report:
reports/indexer-state-mi300x-20261029.md.This is an explicit unshuffled cache boundary, not an implicit replacement for swizzled engine layouts. No cache-write autograd, full quantized VJP, actual distributed CP/RCCL, multi-node, complete framework registration or full-model acceptance is claimed.