[Qwen 3.8] LMCache offload for QSA + GDN + PLE state

atom/config.py

@@ -66,6 +66,9 @@ class KVCacheTensor:
     # without their exponents, and every index it computes stays in bounds while
     # doing it. Whatever moves one must move the other.
     index_scale: torch.Tensor | None = None
+    # Per-block planes a layer's pages carry beyond the fields above, in stored
+    # byte order (Qwen3.8's QSA side caches). Empty on most models.
+    extra_planes: tuple[torch.Tensor, ...] = ()
     # ReplaySSM record buffers for linear-attention layers: this layer's slice
     # of the (k, u, g) pools.  None for every other attention type.  Carried
     # here because the layer-id -> linear-attn-index mapping already lives in

atom/kv_transfer/disaggregation/factory.py

@@ -93,6 +93,10 @@ class KVConnectorFactory:
     # Whether a backend is a KV offload tier, whose scheduler half rewrites the
     # completion sets it is handed (`process_completions`). Absent means no.
     _offload: ClassVar[dict[str, bool]] = {}
+    # Whether a backend stores and restores a PAGE checkpoint image, the
+    # per-request state behind a hybrid's pages (Qwen3.8's GDN and PLE).
+    # Absent means no: a model whose state rides that image refuses it.
+    _reads_page_images: ClassVar[dict[str, bool]] = {}
 
     @classmethod
     def register(
@@ -108,6 +112,7 @@ def register(
         copies_whole_block_regions: bool = False,
         requires_pd_staging: bool = True,
         offload: bool = False,
+        reads_page_images: bool = False,
     ) -> None:
         """Register a KV connector backend.
 
@@ -119,6 +124,8 @@ def register(
             scheduler_class: Class name within *scheduler_module*.
             offload: The backend is a KV offload tier; a ``multi`` runs at most
                 one (`is_offload_backend`).
+            reads_page_images: The backend carries PAGE checkpoint images
+                (`reads_page_images`).
         """
         cls._registry[name] = {
             "worker_module": worker_module,
@@ -143,6 +150,7 @@ def register(
         cls._reads_block_regions[name] = bool(reads_block_regions)
         cls._copies_whole_block_regions[name] = bool(copies_whole_block_regions)
         cls._offload[name] = bool(offload)
+        cls._reads_page_images[name] = bool(reads_page_images)
 
     @classmethod
     def canonical_name(cls, value: object, *, path: str = "kv_transfer_config") -> str:
@@ -170,6 +178,19 @@ def is_offload_backend(cls, name: str) -> bool:
 
         return cls._offload.get(name, False)
 
+    @classmethod
+    def reads_page_images(cls, name: str) -> bool:
+        """Whether the registered backend ``name`` (canonical) carries PAGE
+        checkpoint images."""
+
+        return cls._reads_page_images.get(name, False)
+
+    @classmethod
+    def page_image_readers(cls) -> list[str]:
+        """Every registered backend that carries PAGE checkpoint images."""
+
+        return sorted(name for name, reads in cls._reads_page_images.items() if reads)
+
     @classmethod
     def connector_name(cls, kv_transfer_config: dict[str, Any] | None) -> str | None:
         """Canonical name of the configured connector; None without a transfer config."""

atom/kv_transfer/offload/README.md

@@ -25,8 +25,10 @@ For the in-process path, the public configuration remains
 layouts: `m3` for MiniMax-M3 PAGE regions
 (including its NSA index cache), `kimi_k3` when the text config has
 `model_type == "kimi_linear"` (dense paged MLA KV plus a KDA per-request state
-tier), `qwen` for GDN hybrids (`qwen3_next`, `qwen3_5_*`; a subclass of the
-K3 tier over an MHA page image), `hybrid` when `hf_config.compress_ratios` is
+tier), `qwen` for GDN hybrids (`qwen3_next`, `qwen3_5_*`, `qwen4_exp`; a
+subclass of the K3 tier over a PAGE image in the MHA pool, or in the QSA arena
+for Qwen3.8, whose image also carries the PLE windows), `hybrid` when
+`hf_config.compress_ratios` is
 present (DSV4 PAGE+SLOT), and `dense` otherwise.
 `kv_transfer_config.offload_layout` can override compatible choices without
 giving scheduler and worker different connector names. MiniMax-M3 cannot be
@@ -97,7 +99,7 @@ Four rules carry the module:
 | `hybrid/dsv4/codec.py` | `DSV4PageSlotCodec`, `DSV4CheckpointCodec`, and `DSV4CheckpointStore`: unified GPU layout plans plus AOS1 framing/storage. |
 | `hybrid/dsv4/triton_page_slot.py` | Raw-`uint8` PAGE/SLOT gather/scatter kernels; PAGE is forward-indexed and SLOT is reverse-indexed. |
 | `hybrid/kimi_k3/connector.py` | Kimi-K3 worker and scheduler: the dense paged-KV path plus one extra leg for the KDA per-request state tier. |
-| `hybrid/qwen/connector.py` | Qwen GDN worker and scheduler. Subclasses the K3 tier; the pages come from the MHA pool. |
+| `hybrid/qwen/connector.py` | Qwen GDN worker and scheduler. Subclasses the K3 tier; the pages come from the MHA pool or Qwen3.8's QSA arena. |
 | `hybrid/kimi_k3/staging.py` | Single-entry bounded GPU staging buffer, D2H/H2D copy stream, and producer event for one flat state entry per transfer. |
 | `hybrid/kimi_k3/state_object.py` | One state checkpoint as a single opaque object keyed by ATOM's own hash, bypassing LMCache's `ChunkedTokenDatabase` (state bytes are not token-sliceable). |
 | `hybrid/kimi_k3/state_tier.py` | Worker-side store/load driver for the state tier on its own executor; reports store/finished/failed hash sets for the engine-side `StateOffloadIndex` to apply. |
@@ -1178,6 +1180,16 @@ Under `--kv_cache_dtype fp8`, each KV block carries its own `k_scale` / `v_scale
 (`dense/kv_byte_codec.py`). The scales travel with the quantized data, so a
 reloaded fp8 block dequantizes identically; no scale is recomputed or dropped.
 
+### Multimodal prompts: seeded keys
+
+Image placeholder ids are the same for every picture, so the lookup, save and
+load requests `ChunkedOffloadSchedulerBase` builds (in-process
+`lmcache_offload` and `lmcache_mp`, including the native-state lookup) key
+their ids through `tier_token_ids(seq, end)` (`_offload_common.py`): the first
+4 ids of a prompt with a `cache_seed` become digests of
+`(cache_seed, position, id)` with bit 31 set, off every vocabulary id. Text
+prompts key unchanged.
+
 ### Invariants enforced in code
 
 | Invariant | Where | Why |
@@ -1409,7 +1421,8 @@ env var. This per-thread isolation is exactly why load and save never contend.
   `lookup_server_worker_ids=[0]` so rank 0 is the authoritative lookup answerer at
   TP>1. `build_lmcache_metadata()` fills `kv_shape` from `hf_config` and pins a
   shared `engine_id` so the scheduler's lookup client and the workers' lookup
-  servers derive the **same** ZMQ socket path.
+  servers derive the **same** ZMQ socket path, one per server
+  (`data_parallel_base_port`) and DP rank.
 - **`metadata.py`** — `ATOMRawBytesLMCacheMetadata` overrides LMCache's allocation
   to hand out opaque `uint8` MemoryObjs (`get_shapes` returns
   `nblocks × bytes_per_block`) and asserts

atom/kv_transfer/offload/__init__.py

@@ -32,6 +32,7 @@
     aliases=("LMCacheOffloadConnector", "LMCacheConnectorV1"),
     requires_pd_staging=False,
     offload=True,
+    reads_page_images=True,
     # The dense codec builds from `KVCacheTensor`s and ignores the region map;
     # the hybrid, m3 and kimi_k3 layouts do read it, and which one applies comes
     # from the model, so `topology_reads_block_regions` refines this by layout.
@@ -47,6 +48,7 @@
     aliases=("LMCacheMPConnector",),
     requires_pd_staging=False,
     offload=True,
+    reads_page_images=True,
     # Groups regions by per-block shape and copies whole blocks, so a layer may
     # publish extra planes (the FP4 indexer's e8m0 scale) as more regions.
     copies_whole_block_regions=True,

atom/kv_transfer/offload/_offload_common.py

@@ -12,9 +12,11 @@
 from __future__ import annotations
 
 import collections
+import hashlib
 import json
 import logging
 import os
+import struct
 import threading
 import time
 import weakref
@@ -63,6 +65,38 @@ def tokens_to_tensor(tokens: list[int]) -> torch.Tensor:
     return torch.from_numpy(np.asarray(tokens, dtype=np.int64))
 
 
+# 4 x 31 bits covers the 64-bit `cache_seed`. The high bit keeps a keyed id off
+# every vocabulary id; the LMCache MP server packs an id into 32 bits.
+_MEDIA_KEY_TOKENS = 4
+_MEDIA_KEY_BIT = 1 << 31
+_MEDIA_KEY_PERSON = b"atom-lmc-media"
+
+
+def tier_token_ids(seq, end: int) -> list[int]:
+    """``seq``'s first ``end`` token ids, as LMCache keys them.
+
+    Image placeholder ids are the same for every image, so the media identity
+    (``Sequence.cache_seed``, which also seeds the state images' hashes) is
+    folded into the first ids. A text prompt keys exactly as before, and a
+    prefix's keys are those of the whole prompt.
+    """
+    tokens = list(seq.token_ids[:end])
+    seed = getattr(seq, "cache_seed", -1)
+    if seed is None or int(seed) == -1:
+        return tokens
+    salt = int(seed).to_bytes(16, "little", signed=True)
+    for i in range(min(_MEDIA_KEY_TOKENS, len(tokens))):
+        digest = hashlib.blake2b(
+            salt + struct.pack("<Iq", i, int(tokens[i])),
+            digest_size=4,
+            person=_MEDIA_KEY_PERSON,
+        ).digest()
+        tokens[i] = _MEDIA_KEY_BIT | (
+            int.from_bytes(digest, "little") & (_MEDIA_KEY_BIT - 1)
+        )
+    return tokens
+
+
 def validated_kv_role(kvc: dict) -> str:
     role = kvc.get("kv_role", "offload")
     if role not in _VALID_KV_ROLES:
@@ -679,13 +713,21 @@ def _chunk_floor(self, tokens: int) -> int:
         chunk = int(self.chunk_size or 256)
         return (max(0, int(tokens)) // chunk) * chunk
 
-    def _loadable_hit(self, hit: int, num_prompt: int) -> int:
+    def _withholds_last_prompt_token(self, seq) -> bool:
+        """Whether a full-prompt hit gives back its last token (`_loadable_hit`):
+        a load that ends at the prompt's end leaves nothing to compute, unless
+        the layout's claim stops short of the load on its own."""
+        del seq
+        return True
+
+    def _loadable_hit(self, hit: int, num_prompt: int, seq=None) -> int:
         """Turn a lookup hit into a length the external tier can actually serve.
 
         Two steps, in this order:
 
         * A hit covering the whole prompt leaves nothing to compute, so step
-          back one token.
+          back one token -- unless `seq` is given and
+          `_withholds_last_prompt_token(seq)` says its claim stops short.
         * Floor to a chunk. `retrieve` resolves the tier at chunk granularity,
           so a `hit` that is not a chunk multiple names tokens the tier does
           not hold and the load's `ret_mask[hbm:lmc].all()` check can never
@@ -698,7 +740,9 @@ def _loadable_hit(self, hit: int, num_prompt: int) -> int:
         """
 
         hit = int(hit)
-        if hit == int(num_prompt):
+        if hit == int(num_prompt) and (
+            seq is None or self._withholds_last_prompt_token(seq)
+        ):
             hit -= 1
         return self._chunk_floor(hit)
 

atom/kv_transfer/offload/chunked_scheduler.py

@@ -19,6 +19,7 @@
 from atom.kv_transfer.offload import config as offcfg
 from atom.kv_transfer.offload._offload_common import (
     OffloadSchedulerMixin,
+    tier_token_ids,
     validated_kv_role,
 )
 from atom.kv_transfer.offload.metadata import (
@@ -172,7 +173,7 @@ def _begin_load_lifecycle(self, seq) -> None:
 
     def _lookup_token_ids(self, seq) -> list[int]:
         """The prompt extent this layout can safely ask the tier to resume."""
-        return list(seq.token_ids[: seq.num_prompt_tokens])
+        return tier_token_ids(seq, seq.num_prompt_tokens)
 
     def _lookup_len(self, seq) -> int:
         """`len(self._lookup_token_ids(seq))` without copying the prompt."""
@@ -511,7 +512,7 @@ def _answer_from_tier_hit(self, seq, sid: str, hit: int) -> tuple[int, bool]:
             return 0, False
         num_prompt = seq.num_prompt_tokens
         frontier = int(seq.num_cached_tokens)
-        hit = self._loadable_hit(hit, num_prompt)
+        hit = self._loadable_hit(hit, num_prompt, seq)
         if self._hit_cap_hook is not None:
             # After `_loadable_hit`, so the hook caps the hit that will
             # actually be requested; before the save floor is recorded, so a
@@ -656,7 +657,7 @@ def _build_save_request(
         """Admit any layout-specific sources before advancing the watermark."""
         return LMCacheReqMeta(
             req_id=seq.id,
-            token_ids=list(seq.token_ids[:aligned]),
+            token_ids=tier_token_ids(seq, aligned),
             block_ids=block_ids,
             save_spec=SaveSpec(skip_leading_tokens=saved, can_save=True),
             is_last_prefill=is_last_prefill,
@@ -677,7 +678,7 @@ def _build_load_request(
         """
         return LMCacheReqMeta(
             req_id=seq.id,
-            token_ids=list(seq.token_ids[:transfer_end]),
+            token_ids=tier_token_ids(seq, transfer_end),
             block_ids=list(seq.block_table),
             load_spec=load_spec,
             load_operation=load_operation,

atom/kv_transfer/offload/config.py

@@ -61,7 +61,7 @@
 }
 # GDN/linear model types carrying a per-request recurrent state but NOT in the
 # `kimi_linear` family `kimi_k3` owns. Same set as `attn_family` resolves to
-# `Family.GDN`, but note the attribute skew: the guard below reads
+# `Family.GDN` or `Family.QSA_GDN`, but note the attribute skew: the guard below reads
 # `hf_config.model_type` while `attn_family` reads the text config. On the
 # native path they agree only because `get_hf_config` flattens the text type
 # up; `is_vllm()` drops `qwen3_5`/`qwen3_5_moe` from that flatten, so if
@@ -75,7 +75,13 @@
 # delivered only through transfer_tensors.block_regions, so the dense codec --
 # which builds from per-layer KVCacheTensor K/V -- would silently drop it.
 _GDN_LINEAR_MODEL_TYPES = frozenset(
-    {"qwen3_next", "qwen3_next_mtp", "qwen3_5_text", "qwen3_5_moe_text"}
+    {
+        "qwen3_next",
+        "qwen3_next_mtp",
+        "qwen3_5_text",
+        "qwen3_5_moe_text",
+        "qwen4_exp_text",
+    }
 )
 
 # Layouts that own a tier for a model's per-request *recurrent* state.
@@ -379,6 +385,13 @@ def build_page_namespace(
             getattr(config, "speculative_config", None)
         ),
     }
+    from atom.utils import get_hf_text_config
+    from atom.utils.selector import Family, attn_family
+
+    if attn_family(get_hf_text_config(hf)) is Family.QSA_GDN:
+        from atom.model_ops.attentions.qwen4_exp_attn import qsa_page_segments
+
+        document.update(qsa_page_segments(config))
     canonical = json.dumps(
         document,
         sort_keys=True,
@@ -448,7 +461,7 @@ def build_lmcache_config(
 
 
 def lmcache_engine_id(config) -> str:
-    """Return the LMCache engine id for this process, namespaced by DP rank.
+    """Return the LMCache engine id for this process: one per server and DP rank.
 
     A scheduler's ``LookupClient`` and its workers' ``LookupServer`` derive the
     same ZMQ ipc socket path from this id, so it must be identical inside one
@@ -460,6 +473,13 @@ def lmcache_engine_id(config) -> str:
     Under ``--enable-dp-attention`` this is the normal case, not a corner case:
     ``CoreManager`` folds TP into DP, so each GPU becomes its own replica with
     ``tensor_parallel_size == 1``.
+
+    The socket lives under a fixed host directory, so two servers on one host
+    would both bind ``dp0`` and a scheduler would follow the other server's
+    hits. ``data_parallel_base_port`` (the runners' TCPStore port, bound by
+    ``CoreManager`` for the server's lifetime) tells servers apart; unresolved
+    (0) or not a port, the DP rank alone is used. The id is not part of a
+    chunk's key.
     """
     dp_rank = int(
         getattr(
@@ -469,6 +489,12 @@ def lmcache_engine_id(config) -> str:
         )
         or 0
     )
+    port = getattr(
+        getattr(config, "parallel_config", None), "data_parallel_base_port", None
+    )
+    # `type(...) is int`: a bool is an int, and True would name port 1.
+    if type(port) is int and port > 0:
+        return f"atom-offload-r{port}-dp{dp_rank}"
     return f"atom-offload-dp{dp_rank}"
 
 

atom/kv_transfer/offload/dense/kv_byte_codec.py

@@ -138,6 +138,8 @@ def __init__(
                 # `build_page_namespace` has no segment order in it to notice --
                 # so a reorder has to bump `PAGE_LAYOUT_VERSION`.
                 getattr(kvt, "index_scale", None),
+                # Same rule: new tails, never a reorder.
+                *getattr(kvt, "extra_planes", ()),
             ):
                 if t is not None and isinstance(t, torch.Tensor) and t.numel() > 0:
                     self._segments.append(t)

atom/kv_transfer/offload/hybrid/qwen/connector.py

@@ -5,7 +5,7 @@
 
 Qwen has no second codec. ``KimiK3OffloadConnector`` already packs
 ``page_unit_views`` and unpacks ``state_entry_views``, and it does not assume
-KDA shapes. Qwen publishes those views from its MHA pool
+KDA shapes. Qwen publishes those views from its MHA pool, Qwen3.8 from its QSA arena
 (``pool_layout/mha_page_unit_geometry``) under the ``gdn-paged-state`` layout id
 spelled in ``GDNAttentionMetadataBuilder.state_transfer``. These
 subclasses exist so the selected layout is ``qwen`` rather than ``kimi_k3``.
@@ -20,7 +20,7 @@
 
 
 class QwenOffloadConnector(KimiK3OffloadConnector):
-    """Worker side: dense MHA KV plus the GDN page-image state tier."""
+    """Worker side: the paged KV (MHA or QSA) plus the GDN page-image state tier."""
 
 
 class QwenOffloadScheduler(KimiK3OffloadScheduler):
@@ -29,3 +29,19 @@ class QwenOffloadScheduler(KimiK3OffloadScheduler):
     Inherits ``StateOffloadFace`` through K3, so ``isinstance`` routing still
     sends state stores and loads here.
     """
+
+    def _withholds_last_prompt_token(self, seq) -> bool:
+        """Not when the joint boundary lies in the withheld chunk.
+
+        A hybrid claims the joint boundary, which already leaves a token to
+        compute. Withholding would floor a chunk-aligned prompt a whole chunk
+        below the rungs it kept near its end, so no boundary forms. A hit no
+        boundary backs is still withheld: `_decide_load_after_alloc` refuses it.
+        """
+        if not getattr(seq, "has_per_req_cache", False):
+            return super()._withholds_last_prompt_token(seq)
+        num_prompt = int(seq.num_prompt_tokens)
+        withheld = self._chunk_floor(num_prompt - 1)
+        if self._block_manager is None or withheld == self._chunk_floor(num_prompt):
+            return True
+        return self._block_manager.joint_boundary_tokens(seq, num_prompt) <= withheld