refactor(kv-cache): one declaration per fact about the KV pools

.claude/hooks/auto_lint.py

@@ -1,5 +1,18 @@
 #!/usr/bin/env python3
-"""PostToolUse hook: auto-format Python files with black + ruff after Edit/Write."""
+"""PostToolUse hook: black + ruff on the file just edited, and say what is left.
+
+Fixing silently is only half of it. `ruff --fix` cannot fix every rule -- RUF012
+is one it only advises on -- and a violation it leaves behind used to vanish
+here, because this hook captured ruff's output and never read it. CI does not
+vanish: reviewdog filters to the pull request's diff context, so any finding on
+a line this edit moved becomes a red check, whatever the repository's total was
+before.
+
+So whatever survives `--fix` goes back to Claude on stderr with exit 2, which
+for a PostToolUse hook means "show this, the tool already ran". Pre-existing
+findings in the file are reported too, on purpose: they are exactly the ones CI
+picks up once an edit lands beside them.
+"""
 
 import json
 import shutil
@@ -10,15 +23,38 @@
 info = json.loads(data)
 file_path = info.get("tool_input", {}).get("file_path", "")
 
+remaining = ""
 if file_path.endswith(".py"):
     black = shutil.which("black")
     ruff = shutil.which("ruff")
+    # `check=False` throughout: a formatter that finds nothing to do and a
+    # linter that finds something both exit non-zero, and neither is this
+    # hook's failure -- the second one is its whole output.
     if black:
-        subprocess.run([black, "--quiet", file_path], capture_output=True)
+        subprocess.run([black, "--quiet", file_path], capture_output=True, check=False)
     if ruff:
         subprocess.run(
-            [ruff, "check", "--fix", "--quiet", file_path], capture_output=True
+            [ruff, "check", "--fix", "--quiet", file_path],
+            capture_output=True,
+            check=False,
         )
+        left = subprocess.run(
+            [ruff, "check", "--output-format=concise", file_path],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if left.returncode:
+            remaining = left.stdout.strip()
 
 # Output original data to stdout per hook protocol
 print(data)
+
+if remaining:
+    print(
+        f"ruff still reports these in {file_path} after --fix. CI reviewdog "
+        "flags any of them that land in the diff context, so fix them now "
+        "rather than after the push:\n" + remaining,
+        file=sys.stderr,
+    )
+    sys.exit(2)

atom/kv_transfer/disaggregation/types.py

@@ -157,10 +157,9 @@ class KVTransferTensors:
     compressor-state PD staging pool and are invalid as sidecar SLOT sources.
     """
 
-    # Block-indexed PAGE regions, indexed forward by physical block id.
+    # Block-indexed PAGE regions, indexed forward by block id.
     block_regions: list[KVTransferRegion]
     slot_regions: list[KVTransferRegion]
-    num_blocks: int
     num_slots: int = 0
     # Optional producer-local -> consumer-global mapping for non-uniform block
     # region layouts. Uniform per-layer groups leave this unset and use the
@@ -185,6 +184,38 @@ class KVTransferTensors:
     # loose runtime attribute so the field the connector reads is part of the
     # contract, not an undocumented assignment two layers away.
     state_backend: object | None = None
+    # Scheduler blocks the PAGE regions are addressed in. `init=False` because
+    # a backend cannot answer it: `req.block_ids` is the scheduler's id space,
+    # and a backend counts in its own page -- a different unit even where it is
+    # the same number. Set through `set_block_count`.
+    num_blocks: int = field(init=False, default=0)
+
+    def set_block_count(self, num_blocks: int) -> None:
+        """Fix the block id space, and check every region is in it.
+
+        Called once the last contributor's regions are in the list; a draft
+        appends its own after construction, so this cannot be a `__post_init__`.
+
+        A region that does not divide into exactly `num_blocks` units was
+        registered in some other unit. Both ends would still agree on
+        `base + id * unit_bytes` and disagree on the stride, so the wrong bytes
+        move and nothing reports it.
+        """
+        for i, region in enumerate(self.block_regions):
+            name = region.semantic_role or i
+            if not region.unit_bytes or region.total_bytes % region.unit_bytes:
+                raise ValueError(
+                    f"PAGE region {name} does not divide into whole blocks: "
+                    f"{region.total_bytes} B in units of {region.unit_bytes} B"
+                )
+            held = region.total_bytes // region.unit_bytes
+            if held != num_blocks:
+                raise ValueError(
+                    f"PAGE region {name} holds {held} blocks but the scheduler "
+                    f"addresses {num_blocks}; it is registered in some unit "
+                    "other than the scheduler's block"
+                )
+        self.num_blocks = num_blocks
 
 
 @dataclass

atom/kv_transfer/offload/dense/connector.py

@@ -98,9 +98,8 @@ def register_kv_caches(
         rank, world = pp_aware_rank_and_world(self._config, tp)
         self._rank = rank
 
-        # num_blocks is the physical block count (num_physical_kvcache_blocks),
-        # threaded from the model runner. MLA stores its KV token-major, so the
-        # codec can't infer the block count from tensor.shape[0]; pass it.
+        # Scheduler blocks, threaded from the model runner. MLA stores its KV
+        # token-major, so the codec cannot infer the count from shape[0].
         self._codec = DenseKVByteCodec(
             kv_caches,
             num_blocks=num_blocks,

atom/model_engine/ipc_utils.py

@@ -20,7 +20,7 @@
 import logging
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 logger = logging.getLogger("atom")
 
@@ -58,30 +58,27 @@ def _import_tensor(meta: dict) -> torch.Tensor:
 # ---------------------------------------------------------------------------
 
 
-def export_kv_cache_handle(
-    kv_cache: torch.Tensor, kv_scale: torch.Tensor | None = None
-) -> dict:
-    """Export kv_cache (and optionally kv_scale for fp8) as CUDA IPC handles.
+def export_kv_cache_handle(kv_cache: torch.Tensor) -> dict:
+    """Export the paged pool as a CUDA IPC handle.
+
+    One handle because there is one allocation: the dequantization scales, any
+    indexer key cache and a draft's sibling pool are regions of it, and the
+    decode side finds them by carving with the same declarations rather than by
+    being sent a tensor per region.
 
     Must be called from the process that allocated the tensor (prefill).
     Returns a dict that can be pickled and sent over ZMQ to the decode process.
     """
-    result = {"kv_cache": _export_tensor(kv_cache)}
-    if kv_scale is not None:
-        result["kv_scale"] = _export_tensor(kv_scale)
-    return result
+    return {"kv_cache": _export_tensor(kv_cache)}
 
 
-def import_kv_cache(meta: dict) -> tuple[torch.Tensor, torch.Tensor | None]:
-    """Reconstruct kv_cache (and kv_scale if present) from CUDA IPC handles.
+def import_kv_cache(meta: dict) -> torch.Tensor:
+    """Reconstruct the paged pool from its CUDA IPC handle.
 
-    Must be called from the consumer process (decode).
-    Returns (kv_cache, kv_scale) — kv_scale is None when not fp8.
-    The returned tensors share GPU memory with prefill's allocation — no copy.
+    Must be called from the consumer process (decode). The returned tensor
+    shares GPU memory with prefill's allocation — no copy.
     """
-    kv_cache = _import_tensor(meta["kv_cache"])
-    kv_scale = _import_tensor(meta["kv_scale"]) if "kv_scale" in meta else None
-    return kv_cache, kv_scale
+    return _import_tensor(meta["kv_cache"])
 
 
 # ---------------------------------------------------------------------------