[ATOM-vLLM] Add support for GLM-5.3-Flash with LMCache offload

atom/model_ops/glm5_next/indexer.py

@@ -75,10 +75,14 @@ def _kpool_write_completed_pools(
         )
         from_tail = abs_slot < chunk_start[req_idx][:, None]
         read_slots = state_slot_idx if state_slot_idx_in is None else state_slot_idx_in
-        safe_slots = read_slots[req_idx].clamp_min(0)
-        stash = tail_cache[safe_slots]
-        pool_k = torch.where(from_tail[..., None], stash[:, 0], pool_k)
-        pool_gate = torch.where(from_tail[..., None], stash[:, 1], pool_gate)
+        safe_slots = read_slots[req_idx].clamp_min(0)[:, None]
+        rows = abs_slot.clamp_min(0) % tail_cache.shape[2]
+        pool_k = torch.where(
+            from_tail[..., None], tail_cache[safe_slots, 0, rows], pool_k
+        )
+        pool_gate = torch.where(
+            from_tail[..., None], tail_cache[safe_slots, 1, rows], pool_gate
+        )
 
     pooled = kpool.pool_and_rotate(pool_k, pool_gate, compress_ape)
     abs_pos = positions.to(torch.int64)

atom/model_ops/glm5_next/kpool.py

@@ -587,12 +587,18 @@ def expand_pools_and_append_tail(
 # The tail buffer: the in-progress pool's raw K and gate, per request
 # --------------------------------------------------------------------------
 #
-# Layout, per indexer layer: [num_slots, 2, POOL, HEAD_DIM] bf16, where index 0
+# Layout, per indexer layer: [num_slots, 2, RING, HEAD_DIM] bf16, where index 0
 # on the second axis is K and 1 is the gate score. A token at absolute position
-# `p` occupies row `p % POOL`. It rides the same per-request state slots KDA's
+# `p` occupies row `p % RING`. It rides the same per-request state slots KDA's
 # recurrent state uses, so it inherits their lifetime, forking and relocation.
 
 
+def _tail_ring_size(tail: torch.Tensor, pool_size: int) -> int:
+    ring = tail.shape[2]
+    assert tail.is_contiguous() and ring >= pool_size, (tail.shape, pool_size)
+    return ring
+
+
 @triton.jit
 def _kpool_seed_tail_kernel(
     k_ptr,
@@ -604,6 +610,7 @@ def _kpool_seed_tail_kernel(
     tail_ptr,
     HEAD_DIM: tl.constexpr,
     POOL: tl.constexpr,
+    RING: tl.constexpr,
 ):
     """Persist each prefill request's trailing incomplete pool.
 
@@ -629,20 +636,21 @@ def _kpool_seed_tail_kernel(
     if src_slot < 0 or dst_slot < 0:
         return
     offs = tl.arange(0, HEAD_DIM)
-    src_base = src_slot * (2 * POOL * HEAD_DIM) + j * HEAD_DIM
-    dst_base = dst_slot * (2 * POOL * HEAD_DIM) + j * HEAD_DIM
+    row = (seq_len - tail_count + j) % RING
+    src_base = src_slot * (2 * RING * HEAD_DIM) + row * HEAD_DIM
+    dst_base = dst_slot * (2 * RING * HEAD_DIM) + row * HEAD_DIM
     from_chunk = i >= q_start
     kval = tl.load(k_ptr + i * HEAD_DIM + offs, mask=from_chunk, other=0.0)
     gval = tl.load(gate_ptr + i * HEAD_DIM + offs, mask=from_chunk, other=0.0)
     kval = tl.where(from_chunk, kval, tl.load(tail_ptr + src_base + offs))
     gval = tl.where(
         from_chunk,
         gval,
-        tl.load(tail_ptr + src_base + POOL * HEAD_DIM + offs),
+        tl.load(tail_ptr + src_base + RING * HEAD_DIM + offs),
     )
     tl.store(tail_ptr + dst_base + offs, kval)
     tl.store(
-        tail_ptr + dst_base + POOL * HEAD_DIM + offs,
+        tail_ptr + dst_base + RING * HEAD_DIM + offs,
         gval,
     )
 
@@ -672,6 +680,7 @@ def kpool_seed_tail(
         tail,
         HEAD_DIM=k.shape[-1],
         POOL=pool_size,
+        RING=_tail_ring_size(tail, pool_size),
     )
 
 
@@ -687,6 +696,7 @@ def _kpool_decode_stash_and_pool_kernel(
     out_ptr,
     HEAD_DIM: tl.constexpr,
     POOL: tl.constexpr,
+    RING: tl.constexpr,
 ):
     """Stash one decode token into the tail, then pool the (possibly complete)
     pool.
@@ -706,8 +716,9 @@ def _kpool_decode_stash_and_pool_kernel(
     pos = tl.load(positions_ptr + r).to(tl.int32)
     phase = pos % POOL
     offs = tl.arange(0, HEAD_DIM)
-    src_base = src_slot * (2 * POOL * HEAD_DIM)
-    dst_base = dst_slot * (2 * POOL * HEAD_DIM)
+    src_base = src_slot * (2 * RING * HEAD_DIM)
+    dst_base = dst_slot * (2 * RING * HEAD_DIM)
+    pool_start = pos - phase
     cur_k = tl.load(k_ptr + r * HEAD_DIM + offs)
     cur_g = tl.load(gate_ptr + r * HEAD_DIM + offs)
 
@@ -716,29 +727,31 @@ def _kpool_decode_stash_and_pool_kernel(
     # into a newly allocated slot.
     m = tl.full((HEAD_DIM,), float("-inf"), tl.float32)
     for s in tl.static_range(POOL):
+        row = ((pool_start + s) % RING) * HEAD_DIM
         g = tl.where(
             s == phase,
             cur_g,
-            tl.load(tail_ptr + src_base + POOL * HEAD_DIM + s * HEAD_DIM + offs),
+            tl.load(tail_ptr + src_base + RING * HEAD_DIM + row + offs),
         )
-        tl.store(tail_ptr + dst_base + POOL * HEAD_DIM + s * HEAD_DIM + offs, g)
+        tl.store(tail_ptr + dst_base + RING * HEAD_DIM + row + offs, g)
         g = g.to(tl.float32)
         a = tl.load(ape_ptr + s * HEAD_DIM + offs).to(tl.float32)
         m = tl.maximum(m, g + a)
     acc = tl.zeros((HEAD_DIM,), tl.float32)
     den = tl.zeros((HEAD_DIM,), tl.float32)
     for s in tl.static_range(POOL):
+        row = ((pool_start + s) % RING) * HEAD_DIM
         g = tl.where(
             s == phase,
             cur_g,
-            tl.load(tail_ptr + src_base + POOL * HEAD_DIM + s * HEAD_DIM + offs),
+            tl.load(tail_ptr + src_base + RING * HEAD_DIM + row + offs),
         ).to(tl.float32)
         kv = tl.where(
             s == phase,
             cur_k,
-            tl.load(tail_ptr + src_base + s * HEAD_DIM + offs),
+            tl.load(tail_ptr + src_base + row + offs),
         )
-        tl.store(tail_ptr + dst_base + s * HEAD_DIM + offs, kv)
+        tl.store(tail_ptr + dst_base + row + offs, kv)
         a = tl.load(ape_ptr + s * HEAD_DIM + offs).to(tl.float32)
         w = tl.exp(g + a - m)
         acc += w * kv.to(tl.float32)
@@ -786,5 +799,6 @@ def kpool_decode_stash_and_pool(
         out,
         HEAD_DIM=head_dim,
         POOL=pool_size,
+        RING=_tail_ring_size(tail, pool_size),
     )
     return out

atom/plugin/vllm/attention/layer_sparse_mla.py

@@ -137,6 +137,7 @@ def generate_sparse_seqlen_kernel(
     cu_query_lens_ptr,
     out_ptr,
     topk_token: tl.constexpr,
+    index_kpool: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     seq_id = tl.program_id(0)
@@ -151,10 +152,12 @@ def generate_sparse_seqlen_kernel(
     if seq_len == 0:
         return
     context_start_point = seq_len - query_len
-    sparse_seqlen = context_start_point + query_offset
-    sparse_seqlen_masked = tl.where(
-        sparse_seqlen + 1 < topk_token, sparse_seqlen + 1, topk_token
-    )
+    causal_len = context_start_point + query_offset + 1
+    if index_kpool > 1:
+        pooled = tl.minimum(causal_len // index_kpool, topk_token // index_kpool)
+        sparse_seqlen_masked = pooled * index_kpool + causal_len % index_kpool
+    else:
+        sparse_seqlen_masked = tl.minimum(causal_len, topk_token)
     tl.store(
         out_ptr + query_start + query_offset,
         sparse_seqlen_masked,
@@ -169,6 +172,7 @@ def generate_sparse_seqlen_triton(
     topk_token: int,
     num_tokens: int,
     max_query_len: int,
+    index_kpool: int = 1,
 ):
     num_seqs = query_lens.size(0)
     out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device)
@@ -180,6 +184,7 @@ def generate_sparse_seqlen_triton(
         cu_query_lens,
         out,
         topk_token,
+        index_kpool,
         block_size,
     )
     return out

atom/plugin/vllm/attention/metadata.py

@@ -18,6 +18,7 @@
 from atom.config import get_current_atom_config
 from atom.distributed.dcp_utils import dcp_persistent_supported
 from atom.model_ops.attention_mla import _MLA_MIN_HEADS, mla_dcp_kernel_num_heads
+from atom.model_ops.glm5_next import geometry as kpool_geometry
 from atom.plugin.vllm.attention.layer_mla import (
     disabled_mla_persistent_metadata,
     mla_fold_kv_metadata_triton,
@@ -413,6 +414,7 @@ class AiterMlaSparseMetadataForVllm:
     reduce_indptr: torch.Tensor | None = None
     reduce_final_map: torch.Tensor | None = None
     reduce_partial_map: torch.Tensor | None = None
+    kpool: object | None = None
 
 
 @dataclass
@@ -2034,6 +2036,15 @@ def build(
         return attn_metadata
 
 
+def sparse_selection_counts(
+    causal_lens: torch.Tensor, index_topk: int, index_kpool: int
+) -> torch.Tensor:
+    if index_kpool <= 1:
+        return torch.clamp(causal_lens, max=index_topk)
+    pooled = torch.clamp(causal_lens // index_kpool, max=index_topk // index_kpool)
+    return pooled * index_kpool + causal_lens % index_kpool
+
+
 class AiterMlaSparseMetadataBuilder(AttentionMetadataBuilder):
     """vLLM-only metadata builder for sparse MLA main attention."""
 
@@ -2070,7 +2081,15 @@ def __init__(
         self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
         self.padded_num_heads = max(self.num_heads, _MLA_MIN_HEADS)
         self.mla_dims = get_mla_dims(self.model_config)
-        self.topk_tokens = config.model_config.hf_config.index_topk
+        text_config = self.model_config.hf_text_config
+        self.index_topk = int(text_config.index_topk)
+        self.index_kpool = kpool_geometry.effective_kpool_size(
+            int(getattr(text_config, "index_kpool", 1) or 1)
+        )
+        self.topk_tokens = kpool_geometry.topk_output_width(
+            self.index_topk, self.index_kpool
+        )
+        self._kpool_tail_store = None
         self.max_model_len_tensor = torch.tensor(
             [self.model_config.max_model_len], device=device, dtype=torch.int32
         )
@@ -2274,9 +2293,10 @@ def build(
             query_lens,
             seq_lens,
             common_attn_metadata.query_start_loc,
-            self.topk_tokens,
+            self.index_topk,
             num_tokens,
             common_attn_metadata.max_query_len,
+            index_kpool=self.index_kpool,
         )
         torch.cumsum(
             sparse_seqlen,
@@ -2327,13 +2347,13 @@ def build(
                 seq_lens_cpu = getattr(common_attn_metadata, "_seq_lens_cpu", None)
             if seq_lens_cpu is None:
                 seq_lens_cpu = seq_lens.cpu()
-            clamped_seq_lens = torch.clamp(
-                seq_lens_cpu[:num_reqs], max=self.topk_tokens
+            selected_counts = sparse_selection_counts(
+                seq_lens_cpu[:num_reqs], self.index_topk, self.index_kpool
             )
             metadata_key = (
                 num_tokens,
                 self.padded_num_heads,
-                clamped_seq_lens.to(torch.int32).numpy().tobytes(),
+                selected_counts.to(torch.int32).numpy().tobytes(),
             )
         if metadata_key is None or metadata_key != self._prev_metadata_key:
             get_mla_metadata_v1(
@@ -2386,9 +2406,25 @@ def build(
             reduce_final_map=self._mla_reduce_final_map,
             reduce_partial_map=self._mla_reduce_partial_map,
         )
+        if self.index_kpool > 1:
+            attn_metadata.kpool = self._build_kpool_metadata(common_attn_metadata)
 
         return attn_metadata
 
+    def _build_kpool_metadata(self, common_attn_metadata):
+        from atom.plugin.vllm import glm5_kpool
+
+        if self._kpool_tail_store is None:
+            self._kpool_tail_store = glm5_kpool.find_tail_store(
+                self.vllm_config, self.model_config.hf_text_config
+            )
+        return glm5_kpool.build_kpool_metadata(
+            common_attn_metadata,
+            decode_threshold=self.reorder_batch_threshold,
+            index_kpool=self.index_kpool,
+            tail_store=self._kpool_tail_store,
+        )
+
     def build_for_drafting(self, common_attn_metadata, draft_index):
         """Sync-free per-draft-step build for sparse MLA main attention.