glm53-kindling: rebase the 23 drifted experimental/ overlays onto the v0.31.0 base (10 files carry real drift; 21 auto-merged, 2 hand-merged: rejection_sampler accept_drafts after the gather_draft_sampled() refactor, sm90 _TRITON early-return after the 3-tuple _kv_lens_host). Unrebased overlays reintroduce ddd6fbca-era import blocks on the v0.31.0 tree -> ImportError at boot

docker-images/glm53-kindling/experimental/draft-trunc/grouped_topk_router.py

@@ -285,7 +285,6 @@ class GroupedTopKRouter(BaseRouter):
             renormalize=self.renormalize,
             num_expert_group=self.num_expert_group,
             has_e_score_bias=self.e_score_correction_bias is not None,
-            routed_scaling_factor=self.routed_scaling_factor,
         )
 
     def _compute_routing(

docker-images/glm53-kindling/experimental/draft-trunc/model_runner.py

@@ -21,6 +21,7 @@ import gc
 import time
 from contextlib import AbstractContextManager
 from copy import deepcopy
+from dataclasses import replace
 from typing import Any, NamedTuple
 
 import numpy as np
@@ -495,7 +496,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     ),
                 )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
-                self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
+                self.max_num_reqs,
+                self.device,
+                logprobs_mode=self.model_config.logprobs_mode,
             )
             self.structured_outputs_worker = StructuredOutputsWorker(
                 max_num_logits=self.max_num_reqs * self.decode_query_len,
@@ -563,7 +566,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         return encoder_runner.get_encoder_timing_stats()
 
     def get_kv_cache_spec(self):
-        return get_kv_cache_spec(self.vllm_config)
+        kv_cache_spec = get_kv_cache_spec(self.vllm_config)
+        # A draft running at DCP=1 under a DCP target keeps a replicated cache.
+        if (
+            isinstance(self.speculator, DraftModelSpeculator)
+            and self.speculator.dcp_size == 1
+            and self.dcp_size > 1
+        ):
+            for name in self.speculator.draft_attn_layer_names & kv_cache_spec.keys():
+                kv_cache_spec[name] = replace(kv_cache_spec[name], dcp_sharded=False)
+        return kv_cache_spec
 
     def initialize_kv_cache(
         self,
@@ -589,6 +601,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         block_sizes = []
         max_num_blocks_per_group = []
         slot_mapping_enabled = []
+        dcp_sharded = []
         for kv_cache_group in kv_cache_config.kv_cache_groups:
             spec = kv_cache_group.kv_cache_spec
             block_sizes.append(spec.block_size)
@@ -596,6 +609,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
             )
             slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
+            dcp_sharded.append(spec.dcp_sharded)
             # Let each cache type account for CP. Attention KV is DCP-sharded,
             # while Mamba/GDN recurrent state is replicated across DCP ranks.
             max_num_blocks = spec.max_num_blocks_per_req(
@@ -642,7 +656,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 self.speculator, "enable_adaptive_verification", False
             ),
             attn_groups=self.attn_groups,
-            attn_cg_support=attn_cg_support,
             req_states=self.req_states,
             query_start_loc=self.input_buffers.query_start_loc,
             num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
@@ -660,6 +673,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             device=self.device,
             kernel_block_sizes=self.kernel_block_sizes,
             slot_mapping_enabled=slot_mapping_enabled,
+            dcp_sharded=dcp_sharded,
             cp_size=self.dcp_size,
             cp_rank=self.dcp_rank,
             cp_interleave=self.cp_interleave,
@@ -780,6 +794,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         skip_eplb: bool = False,
         is_profile: bool = False,
         valid_dummy_state_slots: bool = False,
+        randomize_inputs: bool = False,
         **kwargs,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         if skip_attn and not is_profile:
@@ -787,7 +802,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 "skip_attn must only be True for initial memory profiling."
             )
 
-        # Create a dummy scheduler output.
+        # Create a dummy scheduler output. Plain draft-model speculation adds
+        # one correction slot per request during prefill. The scheduler
+        # accounts for these slots, while dummy runs bypass the scheduler.
+        # Adjust the dummy token count only when the expanded draft batch
+        # would exceed the token budget.
         num_reqs = min(num_tokens, self.max_num_reqs)
         if uniform_decode:
             # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
@@ -801,8 +820,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             num_reqs = min(
                 num_reqs, self.max_num_tokens // self.speculator.num_query_per_req
             )
+            num_tokens = min(
+                num_tokens,
+                self.max_num_tokens
+                - num_reqs * self.speculator.num_extra_query_per_req,
+            )
+
         # Distribute the remainder evenly so no dummy request exceeds
         # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+
         num_tokens_per_request = [
             num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
             for i in range(num_reqs)
@@ -842,6 +868,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 is_profile=is_profile,
                 context_len=context_len,
                 valid_dummy_state_slots=valid_dummy_state_slots,
+                randomize_inputs=randomize_inputs,
             )
         self.kv_connector.set_disabled(False)
 
@@ -933,7 +960,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.pooling_runner.dummy_pooler_run(hidden_states)
 
     @torch.inference_mode()
-    def profile_run(self) -> None:
+    def profile_run(self, randomize_inputs: bool = False) -> None:
         if self.supports_mm_inputs and self.is_first_pp_rank:
             mm_config = self.model_config.multimodal_config
             if mm_config is not None and not mm_config.skip_mm_profiling:
@@ -946,7 +973,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 )
 
         hidden_states, sample_hidden_states = self._dummy_run(
-            self.max_num_tokens, skip_attn=True, is_profile=True
+            self.max_num_tokens,
+            skip_attn=True,
+            is_profile=True,
+            randomize_inputs=randomize_inputs,
         )
 
         # Only run sampler/pooler on last PP rank (non-last ranks return None).
@@ -1219,7 +1249,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         if dummy_run:
             # Dummy batches are uniform by construction.
             return None, get_uniform_decode_token_count(
-                num_reqs, num_toks, max_query_len, has_prefill=False
+                num_reqs, num_toks, max_query_len, decode_graph_eligible=True
             )
 
         draft_tokens = scheduler_output.scheduled_spec_decode_tokens
@@ -1238,24 +1268,53 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             idx_mapping_np
         ]
         is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
+        has_prefill = bool(is_prefilling_np.any())
+        max_seq_len_np = None
+        if self.use_pp and has_prefill:
+            # consumed by the PP `compute_need_sampled_mask` for batches with prefill.
+            max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
+
+        num_draft_tokens_np = None
+        if draft_tokens:
+            num_drafts_iter = (len(draft_tokens.get(req_id, ())) for req_id in req_ids)
+            num_draft_tokens_np = np.fromiter(
+                num_drafts_iter, dtype=np.int32, count=num_reqs
+            )
+            if self.adaptive_verification is not None:
+                num_toks = self.adaptive_verification.get_num_tokens(
+                    num_tokens_per_req, draft_tokens
+                )
 
-        if self.adaptive_verification is not None and draft_tokens:
-            num_toks = self.adaptive_verification.get_num_tokens(
-                num_tokens_per_req, draft_tokens
+        prefill_runs_as_decode_np = None
+        decode_graph_eligible = not has_prefill
+        if has_prefill and self.pcp_manager is None:
+            # One new prompt token over existing context (possibly padded with
+            # placeholder drafts) runs as a decode. PCP partitions prefill rows.
+            num_new_tokens = num_scheduled_tokens
+            if num_draft_tokens_np is not None:
+                num_new_tokens = num_new_tokens - num_draft_tokens_np
+            one_token_ext = (num_new_tokens == 1) & (num_computed_prefill_tokens_np > 0)
+            prefill_runs_as_decode_np = is_prefilling_np & one_token_ext
+            decode_graph_eligible = bool(
+                (prefill_runs_as_decode_np | ~is_prefilling_np).all()
             )
 
         batch_state = BatchReqState(
             req_ids=req_ids,
             num_scheduled_tokens=num_scheduled_tokens,
             num_tokens=num_toks,
+            num_draft_tokens_np=num_draft_tokens_np,
             idx_mapping_np=idx_mapping_np,
             prefill_len_np=prefill_len_np,
             num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
             is_prefilling_np=is_prefilling_np,
-            has_prefill=bool(is_prefilling_np.any()),
+            has_prefill=has_prefill,
+            prefill_runs_as_decode_np=prefill_runs_as_decode_np,
+            decode_graph_eligible=decode_graph_eligible,
+            max_seq_len_np=max_seq_len_np,
         )
         return batch_state, get_uniform_decode_token_count(
-            num_reqs, num_toks, max_query_len, batch_state.has_prefill
+            num_reqs, num_toks, max_query_len, decode_graph_eligible
         )
 
     def prepare_inputs(
@@ -1297,11 +1356,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 num_reqs, dtype=torch.int32, device=self.device
             )
         else:
-            num_draft_tokens_per_req = np.fromiter(
-                (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
-                dtype=np.int32,
-                count=num_reqs,
-            )
+            num_draft_tokens_per_req = batch_req_state.num_draft_tokens_np
+            if num_draft_tokens_per_req is None:
+                num_draft_tokens_per_req = np.zeros(num_reqs, dtype=np.int32)
             num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
             total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
             total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
@@ -1445,6 +1502,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
             is_prefilling_np=batch_req_state.is_prefilling_np,
             has_prefill=batch_req_state.has_prefill,
+            prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np,
+            decode_graph_eligible=batch_req_state.decode_graph_eligible,
+            max_seq_len_np=batch_req_state.max_seq_len_np,
             input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
             positions=self.input_buffers.positions[:num_tokens_after_padding],
             is_padding=is_padding,
@@ -1629,6 +1689,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         is_profile: bool = False,
         context_len: int = 0,
         valid_dummy_state_slots: bool = False,
+        randomize_inputs: bool = False,
     ) -> ModelRunnerOutput | IntermediateTensors | None:
         if not dummy_run:
             # Update the request states.
@@ -1653,12 +1714,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # Get batch descriptor and sync across DP ranks.
         num_reqs = len(scheduler_output.num_scheduled_tokens)
         num_toks = scheduler_output.total_num_scheduled_tokens
-        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
+        max_query_len: int | None = max(scheduler_output.num_scheduled_tokens.values())
         batch_req_state, uniform_tok_count = self.gather_batch_req_state(
             scheduler_output, dummy_run
         )
         if batch_req_state is not None:
             num_toks = batch_req_state.num_tokens
+            if batch_req_state.has_prefill:
+                # Varlen decode graphs replay decodes only, and their bound
+                # alone would admit a short prefill.
+                max_query_len = None
             if self.pcp_manager is not None:
                 num_toks = self.pcp_manager.get_num_tokens_for_dispatch(
                     batch_req_state.num_scheduled_tokens,
@@ -1743,6 +1808,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 # so MoE memory is measured and MoE kernels are exercised.
                 is_padding=not is_profile,
             )
+            if randomize_inputs:
+                # All-zero input_ids route every token to the same experts.
+                input_batch.input_ids.random_(0, self.vocab_size)
             if self.pcp_manager is not None:
                 input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
             if skip_attn_for_dummy_run:
@@ -2054,6 +2122,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.req_states.num_computed_tokens.gpu,
             self.req_states.prompt_len.np,
         )
+        prompt_token_id_logprobs_dict = (
+            self.prompt_logprobs_worker.compute_prompt_token_id_logprobs(
+                self.model.compute_logits,
+                hidden_states,
+                input_batch,
+                self.req_states.prompt_len.np,
+            )
+        )
 
         # Prepare the model runner output.
         model_runner_output = ModelRunnerOutput(
@@ -2063,6 +2139,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
             sampled_token_ids=None,  # type: ignore
             prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
+            prompt_token_id_logprobs_dict=prompt_token_id_logprobs_dict,
             cudagraph_stats=cudagraph_stats,
         )
         pending_aux_output = None
@@ -2320,11 +2397,15 @@ class BatchReqState(NamedTuple):
     # May be less than scheduler_output.total_num_scheduled_tokens:
     # adaptive verification trims the draft budget before running.
     num_tokens: int
+    num_draft_tokens_np: np.ndarray | None  # [num_reqs], None if no drafts
     idx_mapping_np: np.ndarray  # [num_reqs]
     prefill_len_np: np.ndarray  # [num_reqs]
     num_computed_prefill_tokens_np: np.ndarray  # [num_reqs]
     is_prefilling_np: np.ndarray  # [num_reqs]
+    max_seq_len_np: np.ndarray | None  # [num_reqs], None in non-pp case
     has_prefill: bool
+    prefill_runs_as_decode_np: np.ndarray | None  # [num_reqs]
+    decode_graph_eligible: bool
 
 
 def sort_batch_req_ids(

docker-images/glm53-kindling/experimental/draft-trunc/rejection_sampler.py

@@ -76,6 +76,76 @@ def _flatten_sampled_kernel(
         tl.store(flat_sampled_ptr + start_idx + i, token_id)
 
 
+@triton.jit
+def _gather_draft_sampled_kernel(
+    # [num_logits]
+    draft_sampled_ptr,
+    # [num_logits]
+    pos_ptr,
+    # [num_tokens]
+    input_ids_ptr,
+    # [num_tokens]
+    positions_ptr,
+    # [num_logits]
+    logits_indices_ptr,
+    # [num_logits]
+    expanded_idx_mapping_ptr,
+    # [num_logits]
+    expanded_local_pos_ptr,
+    # [max_num_reqs]
+    prefill_len_ptr,
+    num_logits,
+    BLOCK_SIZE: tl.constexpr,
+):
+    block = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = block < num_logits
+    token_idx = tl.load(logits_indices_ptr + block, mask=mask)
+    token_id = tl.load(input_ids_ptr + token_idx, mask=mask)
+    pos = tl.load(positions_ptr + token_idx, mask=mask)
+    req_state_idx = tl.load(expanded_idx_mapping_ptr + block, mask=mask)
+    local_pos = tl.load(expanded_local_pos_ptr + block, mask=mask)
+    prefill_len = tl.load(prefill_len_ptr + req_state_idx, mask=mask)
+    # Drafts are only proposed after sampling past the prefill, so any draft
+    # slots in a step starting within the prefill are placeholders (e.g. the
+    # padded first step after a P/D remote KV load).
+    is_placeholder = (local_pos > 0) & (pos - local_pos < prefill_len)
+    token_id = tl.where(is_placeholder, -1, token_id)
+    tl.store(draft_sampled_ptr + block, token_id, mask=mask)
+    tl.store(pos_ptr + block, pos, mask=mask)
+
+
+def gather_draft_sampled(
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    logits_indices: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    expanded_local_pos: torch.Tensor,
+    prefill_len: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather the input token and position of each logits row.
+
+    Draft rows of requests that have not yet sampled past their prefill are
+    set to -1 so that the rejection kernels reject them.
+    """
+    num_logits = logits_indices.shape[0]
+    draft_sampled = input_ids.new_empty(num_logits)
+    pos = positions.new_empty(num_logits)
+    BLOCK_SIZE = 1024
+    _gather_draft_sampled_kernel[(triton.cdiv(num_logits, BLOCK_SIZE),)](
+        draft_sampled,
+        pos,
+        input_ids,
+        positions,
+        logits_indices,
+        expanded_idx_mapping,
+        expanded_local_pos,
+        prefill_len,
+        num_logits,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return draft_sampled, pos
+
+
 class RejectionSampler:
     def __init__(
         self,
@@ -301,12 +371,18 @@ class RejectionSampler:
         # that num_nans is computed before applying penalties and temperature.
         num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
 
-        draft_sampled = input_batch.input_ids[input_batch.logits_indices]
+        draft_sampled, pos = gather_draft_sampled(
+            input_batch.input_ids,
+            input_batch.positions,
+            input_batch.logits_indices,
+            input_batch.expanded_idx_mapping,
+            input_batch.expanded_local_pos,
+            self.sampler.req_states.prefill_len.gpu,
+        )
         # A dead draft (draft_trunc) is -1 for acceptance only, so it never matches.
         accept_drafts = draft_sampled
         if draft_trunc.ENABLED:
             accept_drafts = draft_sampled.masked_fill(input_batch.is_padding[input_batch.logits_indices], -1)
-        pos = input_batch.positions[input_batch.logits_indices]
 
         max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
             input_batch.idx_mapping_np

docker-images/glm53-kindling/experimental/draft-trunc/speculator.py

@@ -60,6 +60,8 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
 
 class BaseSpeculator(ABC):
     num_query_per_req: int = 1
+    # Extra query slots reserved per request outside the regular queries.
+    num_extra_query_per_req: int = 0
 
     @abstractmethod
     def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
@@ -140,6 +142,7 @@ class DraftModelSpeculator(BaseSpeculator):
         # DP configuration
         self.dp_size = vllm_config.parallel_config.data_parallel_size
         self.dp_rank = vllm_config.parallel_config.data_parallel_rank
+        self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
 
         self.eplb_state: EplbState | None = None
 
@@ -223,13 +226,13 @@ class DraftModelSpeculator(BaseSpeculator):
         self.model = self.load_draft_model(target_model, target_attn_layer_names)
         self._validate_local_argmax_reduction()
 
-        all_attn_layers = set[str](
-            get_layers_from_vllm_config(
-                self.vllm_config,
-                AttentionLayerBase,  # type: ignore[type-abstract]
-            ).keys()
+        all_attn_layers = get_layers_from_vllm_config(
+            self.vllm_config,
+            AttentionLayerBase,  # type: ignore[type-abstract]
         )
-        self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names
+        self.draft_attn_layer_names = set(all_attn_layers) - target_attn_layer_names
+        for layer_name in self.draft_attn_layer_names:
+            all_attn_layers[layer_name].is_draft_layer = True
 
         target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs
         draft_supports_mm = supports_multimodal_embeddings(self.model)
@@ -307,6 +310,7 @@ class DraftModelSpeculator(BaseSpeculator):
         step: int,
         causal: bool | Mapping[int, bool] = True,
         dcp_local_seq_lens: torch.Tensor | None = None,
+        slot_mappings: torch.Tensor | None = None,
     ) -> dict[str, Any] | None:
         num_reqs_padded = batch_desc.num_reqs or num_reqs
         # A FULL graph replays a captured shape whose padded requests each hold
@@ -328,7 +332,8 @@ class DraftModelSpeculator(BaseSpeculator):
         block_tables = [
             x[:num_reqs_padded] for x in self.block_tables.input_block_tables
         ]
-        slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
+        if slot_mappings is None:
+            slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
         draft_seq_lens_cpu_upper_bound = torch.zeros(
             num_reqs_padded, dtype=torch.int32, device="cpu"
         )
@@ -338,14 +343,14 @@ class DraftModelSpeculator(BaseSpeculator):
             out=draft_seq_lens_cpu_upper_bound[:num_reqs],
         )
         draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
-        if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
+        if dcp_local_seq_lens is None and self.dcp_size > 1:
             # Draft steps advance and rewind their own global sequence lengths,
             # so the target model's DCP-local lengths may already be stale.
             dcp_local_seq_lens = prepare_dcp_local_seq_lens(
                 self.input_buffers.dcp_local_seq_lens,
                 self.input_buffers.seq_lens,
                 num_reqs,
-                self.block_tables.cp_size,
+                self.dcp_size,
                 self.block_tables.cp_rank,
                 self.block_tables.cp_interleave,
             )

docker-images/glm53-kindling/experimental/fixes/flashinfer_mla_sparse_sm90.py

@@ -17,14 +17,14 @@ encoded by the indexer's selection, so ``causal=False``.
 CUDA-graph handling: ``plan()`` copies its inputs to host unconditionally,
 so it must stay outside graph capture. Each metadata builder owns a wrapper,
 reserved capture-stable device buffers, and the plan parameters. The wrapper
-bakes the per-row ``kv_len`` into its int schedule
-at plan() time — ``run()`` never reads the device-side buffer — so the
-metadata builder replans every step (outside capture) with exact host-side
-lengths derived from the batch's sequence lengths; a full-width schedule
-would send the kernel past each row's valid count into the -1 tail of the
-converted index buffer (illegal address). Per-step content (top-k slots)
-is written into the reserved buffers by kernels inside the captured
-forward, and captured runs read the refreshed plan buffers on replay.
+bakes the per-row ``kv_len`` into its int schedule at plan() time — ``run()``
+never reads the device-side buffer. The builder plans (outside capture) from
+sync-free host upper bounds and clamps each scheduled work item's ``kv_end``
+on device to the exact valid count, so the kernel never reads past a row's
+valid prefix (the -1 tail of the converted index buffer) and async scheduling
+needs no D2H sync. Per-step content (top-k slots) is written into the
+reserved buffers by kernels inside the captured forward, and captured runs
+read the refreshed plan buffers on replay.
 
 KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]``
 (uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The
@@ -64,6 +64,26 @@ from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
 
 _FP8_KV_DTYPES = ("fp8", "fp8_e4m3")
 _WORKSPACE_BYTES = 128 * 1024 * 1024
+# TODO: FlashInfer bakes per-row kv_len into the plan on the host, which
+# forces a D2H sync under async scheduling. The plan-info layout constants
+# below and the int-workspace clamp in _SM90State work around this by
+# planning from sync-free upper bounds; delete them once FlashInfer accepts
+# an upper-bound plan() that clamps kv_end internally, or exposes a
+# GPU-side plan API (upstream ask:
+# https://github.com/flashinfer-ai/flashinfer/issues/5559, see
+# vllm-project/vllm#58684).
+# FlashInfer MLAPlanInfo layout (scheduler.cuh) field indices.
+_PLAN_INFO_LEN = 18
+_PI_NUM_BLKS_Y = 1
+_PI_Q_INDPTR = 2
+_PI_KV_INDPTR = 3
+_PI_KV_START = 13
+_PI_KV_END = 14
+_PI_WORK_INDPTR = 15
+# Planned lengths exceed the host upper bound by up to this many tokens so a
+# plan stays reusable across draft steps and consecutive decode steps; the
+# device-side clamp restores the exact lengths.
+_PLAN_SLACK = 32
 
 # GB10: a Triton sparse MLA kernel in place of FlashInfer's FA2 path, which
 # runs at ~10 TFLOPS on SM121 (gb10_sparse_mla.py; about 3.3x faster).
@@ -87,6 +107,40 @@ def _pack_topk_indices(
     tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)
 
 
+@triton.jit
+def _clamp_work_kv_end(
+    int_ws,
+    saved_kv_end,
+    seq_lens,
+    query_start_loc,
+    req_id_per_token,
+    num_works_idx,
+    q_indptr_off,
+    kv_start_off,
+    kv_end_off,
+    TOPK: tl.constexpr,
+    KPOOL: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    # Clamp each planned work item to its row's exact valid count, derived
+    # from the device batch layout exactly as the convert kernel counts it.
+    works = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = works < tl.load(int_ws + num_works_idx)
+    row = tl.load(int_ws + q_indptr_off + works, mask=mask, other=0)
+    req = tl.load(req_id_per_token + row, mask=mask, other=0)
+    seq_len = tl.load(seq_lens + req, mask=mask, other=0)
+    query_end = tl.load(query_start_loc + req + 1, mask=mask, other=0)
+    ctx = seq_len - query_end + row + 1
+    valid = tl.where(ctx <= TOPK, ctx, TOPK + ctx % KPOOL)
+    start = tl.load(int_ws + kv_start_off + works, mask=mask, other=0)
+    end = tl.load(saved_kv_end + works, mask=mask, other=0)
+    tl.store(
+        int_ws + kv_end_off + works,
+        tl.maximum(tl.minimum(end, valid), start),
+        mask=mask,
+    )
+
+
 class FlashInferMLASparseSM90Backend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -195,6 +249,8 @@ class _SM90State:
         kv_lora_rank: int,
         qk_rope_head_dim: int,
         sm_scale: float,
+        index_topk: int,
+        index_kpool: int,
     ) -> None:
         from flashinfer.mla import BatchMLAPagedAttentionWrapper
 
@@ -207,6 +263,10 @@ class _SM90State:
         self.kv_lora_rank = kv_lora_rank
         self.qk_rope_head_dim = qk_rope_head_dim
         self.sm_scale = sm_scale
+        self.index_topk = index_topk
+        self.index_kpool = max(index_kpool, 1)
+        # Largest valid count the convert kernel can produce for any context.
+        self.max_valid = min(self.index_topk + self.index_kpool - 1, topk_width)
         # User-reserved buffers: with use_cuda_graph=True plan() refreshes
         # these in place, so run()'s captured kernels always read them.
         self.kv_indices = torch.zeros(
@@ -231,24 +291,54 @@ class _SM90State:
         self._lens_cpu = torch.full(
             (self.max_tokens,), self.topk_width, dtype=torch.int32
         )
-
-    def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
-        """Replan with exact per-row KV lengths (CPU int32, ``[num_tokens]``).
-
-        The wrapper bakes kv_len into its int schedule from host values;
-        run() never reads the device kv_len_arr buffer. Scheduling with the
-        full buffer width while kv_indices rows carry a -1 tail past each
-        row's valid count makes the kernel compute ``-1 * ckv_stride_page``
-        (illegal address), so the lengths must be exact at plan time and
-        replanned every step as contexts grow. Must run outside CUDA graph
-        capture: the in-place refreshed plan_info/indptr buffers are what
-        captured runs read.
+        # Planned per-work kv_end of the current plan, restored by each clamp;
+        # sized on first plan to the planner's work capacity.
+        self._saved_kv_end = torch.empty(0, dtype=torch.int32, device=device)
+        self._planned_num_tokens = -1
+        self._int_ws: torch.Tensor | None = None
+        self._plan_offsets: tuple[int, int, int, int] | None = None
+
+    def plan(
+        self,
+        num_tokens: int,
+        kv_lens: torch.Tensor,
+        cam: CommonAttentionMetadata | None,
+        req_id_per_token: torch.Tensor | None,
+    ) -> None:
+        """Plan per-row KV lengths (CPU int32, ``[num_tokens]``).
+
+        The wrapper bakes kv_len into its int schedule from host values, so
+        rows past their valid count would send the kernel into the -1 tail of
+        the converted index buffer. With ``cam`` and ``req_id_per_token``
+        (device batch layout), ``kv_lens`` only needs to upper-bound the
+        valid counts: the plan is padded by ``_PLAN_SLACK`` and reused while
+        the bound still fits, and every call clamps each work item's kv_end
+        on device to the exact count, so no D2H sync is needed. Without them
+        the lengths must be exact. Must run outside CUDA graph capture: the
+        in-place refreshed plan_info/indptr buffers are what captured runs
+        read.
         """
         if torch.cuda.is_current_stream_capturing():
             raise RuntimeError(
                 "FlashInferMLASparseSM90 plan() called inside CUDA graph "
                 "capture; lengths must be planned host-side before capture."
             )
+        kv_lens = kv_lens.to(torch.int32)
+        if cam is None or req_id_per_token is None:
+            self._plan(num_tokens, kv_lens)
+            self._planned_num_tokens = -1
+            return
+        planned = self._lens_cpu[:num_tokens]
+        if num_tokens != self._planned_num_tokens or not bool(
+            ((planned >= kv_lens) & (planned <= kv_lens + _PLAN_SLACK)).all()
+        ):
+            self._plan(
+                num_tokens, torch.clamp(kv_lens + _PLAN_SLACK, max=self.max_valid)
+            )
+            self._planned_num_tokens = num_tokens
+        self._clamp_kv_end(cam, req_id_per_token)
+
+    def _plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
         # CPU staging buffers are filled in place: plan() runs per step
         # (once per draft/verify metadata build), so per-call allocations and
         # device round trips are on the hot path. Passing CPU tensors lets
@@ -262,13 +352,16 @@ class _SM90State:
         # matching the exact lengths, including zero-length padded rows.
         torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
         self._lens_cpu.zero_()
-        self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
+        self._lens_cpu[:num_tokens] = kv_lens
         self._kv_cpu[0] = 0
         torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])
+        # Only the live prefix of kv_indices: the wrapper snapshots and
+        # copies whatever is passed.
+        num_indices = max(int(self._kv_cpu[-1]), 1)
         self.wrapper.plan(
             self._qo_cpu,
             self._kv_cpu,
-            self.kv_indices,
+            self.kv_indices[:num_indices],
             self._lens_cpu,
             self.num_heads,
             self.kv_lora_rank,  # head_dim_ckv
@@ -279,6 +372,42 @@ class _SM90State:
             q_data_type=torch.bfloat16,
             kv_data_type=self.kv_dtype,
         )
+        plan_info = [int(x) for x in self.wrapper._plan_info]
+        assert len(plan_info) == _PLAN_INFO_LEN, plan_info
+        # Per-work arrays are laid out back to back at the planner's capacity.
+        max_works = (plan_info[_PI_KV_INDPTR] - plan_info[_PI_Q_INDPTR]) // 4
+        if self._saved_kv_end.numel() != max_works:
+            self._saved_kv_end = torch.empty(
+                max_works, dtype=torch.int32, device=self.device
+            )
+        self._int_ws = self.wrapper._planned_backend._int_workspace_buffer.view(
+            torch.int32
+        )
+        self._plan_offsets = (
+            plan_info[_PI_WORK_INDPTR] // 4 + plan_info[_PI_NUM_BLKS_Y],
+            plan_info[_PI_Q_INDPTR] // 4,
+            plan_info[_PI_KV_START] // 4,
+            plan_info[_PI_KV_END] // 4,
+        )
+        kv_end_off = self._plan_offsets[3]
+        self._saved_kv_end.copy_(self._int_ws[kv_end_off : kv_end_off + max_works])
+
+    def _clamp_kv_end(
+        self, cam: CommonAttentionMetadata, req_id_per_token: torch.Tensor
+    ) -> None:
+        assert self._int_ws is not None and self._plan_offsets is not None
+        block = 1024
+        _clamp_work_kv_end[(triton.cdiv(self._saved_kv_end.numel(), block),)](
+            self._int_ws,
+            self._saved_kv_end,
+            cam.seq_lens,
+            cam.query_start_loc,
+            req_id_per_token,
+            *self._plan_offsets,
+            TOPK=self.index_topk,
+            KPOOL=self.index_kpool,
+            BLOCK_SIZE=block,
+        )
 
     def pack_indices(self, topk_slots: torch.Tensor) -> None:
         num_tokens, width = topk_slots.shape
@@ -317,6 +446,8 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
             )
         topk_indices_buffer = impl.topk_indices_buffer
         assert topk_indices_buffer is not None
+        hf_config = vllm_config.model_config.hf_text_config
+        assert hf_config.index_topk is not None
         self.state = _SM90State(
             device,
             impl.num_heads,
@@ -327,20 +458,18 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
             kv_lora_rank=impl.kv_lora_rank,
             qk_rope_head_dim=impl.qk_rope_head_dim,
             sm_scale=impl.scale,
+            index_topk=int(hf_config.index_topk),
+            index_kpool=int(getattr(hf_config, "index_kpool", 1) or 1),
+        )
+        spec_config = vllm_config.speculative_config
+        self._adaptive_verification = bool(
+            spec_config is not None and spec_config.enable_adaptive_verification
         )
-        # seq_lens_cpu_upper_bound is optimistic on decode rows under async
-        # spec decode, so the fast sync-free path is only safe without it;
-        # under async scheduling the exact (device) positions are used at
-        # the cost of one D2H sync per metadata build.
-        self._async_scheduling = bool(vllm_config.scheduler_config.async_scheduling)
-        hf_config = vllm_config.model_config.hf_text_config
-        assert hf_config.index_topk is not None
-        self._index_topk = int(hf_config.index_topk)
-        self._index_kpool = int(getattr(hf_config, "index_kpool", 1) or 1)
 
-    def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]:
-        """Exact per-row KV lengths, host-side (the flashinfer wrapper bakes
-        them into its schedule at plan time; there is no device-side path).
+    def _kv_lens_host(
+        self, cam: CommonAttentionMetadata
+    ) -> tuple[int, torch.Tensor, bool]:
+        """Host per-row KV lengths for the plan and whether they are exact.
 
         A row for the j-th query token of request i attends
         ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection
@@ -349,53 +478,41 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
         ``index_topk`` pool-expanded tokens plus the trailing incomplete
         pool (valid == ``index_topk + context % index_kpool``). Both match
         the count of non -1 entries the convert kernel produces.
+
+        With ``seq_lens_cpu_upper_bound`` (optimistic under async spec
+        decode) returns sync-free upper bounds, which the plan clamps on
+        device. Otherwise returns exact lengths via a D2H copy.
         """
         num_reqs = cam.num_reqs
         qsl = cam.query_start_loc_cpu[: num_reqs + 1]
         num_rows = int(qsl[-1])
         if num_rows == 0:
-            return 0, torch.zeros(0, dtype=torch.int32)
-        # Row context == position + 1. Without async scheduling the
-        # maintained host upper bound is exact, giving a sync-free path;
-        # under async scheduling it is optimistic on decode rows, so fall
-        # back to the exact device positions (one D2H sync per build). The
-        # seq_lens derivation equals position + 1 because positions are
-        # contiguous per request.
+            return 0, torch.zeros(0, dtype=torch.int32), True
+        state = self.state
         sl_host = cam.seq_lens_cpu_upper_bound
-        positions = cam.positions
-        if not self._async_scheduling and sl_host is not None:
-            seq_lens = sl_host[:num_reqs].to(torch.int32)
-            q_lens = qsl[1:] - qsl[:-1]
-            first_pos = seq_lens - q_lens
-            req_of_row = torch.repeat_interleave(
-                torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
-            )
-            rows = torch.arange(num_rows, dtype=torch.int32)
-            ctx = (
-                first_pos.to(torch.int64)[req_of_row]
-                + rows.to(torch.int64)
-                - qsl.to(torch.int64)[req_of_row]
-                + 1
-            )
-        elif positions is not None and num_rows <= positions.shape[0]:
-            ctx = positions[:num_rows].cpu().to(torch.int64) + 1
+        if sl_host is None:
+            ctx = self._row_contexts(qsl, cam.seq_lens[:num_reqs].cpu())
+            topk, kpool = state.index_topk, state.index_kpool
+            lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)
+            return num_rows, lens.to(torch.int32), True
+        if self._adaptive_verification:
+            # The host query split may not match the device one, so bound
+            # every row by the longest request.
+            ctx = torch.full((num_rows,), int(sl_host[:num_reqs].max()))
         else:
-            seq_lens = cam.seq_lens[:num_reqs].cpu().to(torch.int32)
-            q_lens = qsl[1:] - qsl[:-1]
-            first_pos = seq_lens - q_lens
-            req_of_row = torch.repeat_interleave(
-                torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
-            )
-            rows = torch.arange(num_rows, dtype=torch.int32)
-            ctx = (
-                first_pos.to(torch.int64)[req_of_row]
-                + rows.to(torch.int64)
-                - qsl.to(torch.int64)[req_of_row]
-                + 1
-            )
-        topk = self._index_topk
-        lens = torch.where(ctx <= topk, ctx, topk + ctx % self._index_kpool)
-        return num_rows, lens.to(torch.int32)
+            ctx = self._row_contexts(qsl, sl_host[:num_reqs])
+        return num_rows, torch.clamp(ctx, max=state.max_valid), False
+
+    @staticmethod
+    def _row_contexts(qsl: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor:
+        num_reqs = qsl.shape[0] - 1
+        q_lens = (qsl[1:] - qsl[:-1]).to(torch.int64)
+        first_ctx = seq_lens.to(torch.int64) - q_lens + 1
+        req_of_row = torch.repeat_interleave(
+            torch.arange(num_reqs, dtype=torch.int64), q_lens
+        )
+        rows = torch.arange(int(qsl[-1]), dtype=torch.int64)
+        return first_ctx[req_of_row] + rows - qsl.to(torch.int64)[req_of_row]
 
     def build(
         self,
@@ -411,9 +528,9 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
             # scheduling (it stalls the CPU until the previous step ends).
             metadata.state = self.state
             return metadata
-        # Replan every step outside any CUDA graph capture with this step's
-        # exact per-row lengths; captured runs read the refreshed buffers.
-        num_rows, kv_lens = self._kv_lens_host(common_attn_metadata)
+        # Plan outside any CUDA graph capture; captured runs read the
+        # refreshed buffers.
+        num_rows, kv_lens, exact = self._kv_lens_host(common_attn_metadata)
         # MHA prefills are removed from q before forward_mqa. Plan only the
         # decode prefix in that case, using the same routing decision as MLA.
         if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha(
@@ -421,7 +538,12 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
         ):
             num_rows = metadata.num_decode_tokens
             kv_lens = kv_lens[:num_rows]
-        self.state.plan(num_rows, kv_lens)
+        if exact:
+            self.state.plan(num_rows, kv_lens, None, None)
+        else:
+            self.state.plan(
+                num_rows, kv_lens, common_attn_metadata, metadata.req_id_per_token
+            )
         metadata.state = self.state
         return metadata
 

docker-images/glm53-kindling/experimental/fixes/gdn_attn.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Backend for GatedDeltaNet attention."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Literal
 
 import os
@@ -62,6 +62,7 @@ class GDNAttentionMetadata:
         None  # shape: [batch - num_spec_decodes,]
     )
     spec_sequence_masks: torch.Tensor | None = None  # shape: [batch,]
+    spec_sequence_masks_cpu: torch.Tensor | None = None  # shape: [batch,]
     spec_token_indx: torch.Tensor | None = None
     non_spec_token_indx: torch.Tensor | None = None
 
@@ -125,6 +126,18 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         self.use_full_cuda_graph: bool = (
             self.compilation_config.cudagraph_mode.has_full_cudagraphs()
         )
+        # update_block_table() keeps the source group's batch-level FULL graph
+        # buffers, so only MRV2, which also reuses metadata at capture, may use
+        # it.
+        self.supports_update_block_table = (
+            vllm_config.use_v2_model_runner
+            and device.type == "cuda"
+            # Not isinstance: KDA's RecoverSSM/checkpoint metadata is per group.
+            and type(self) is GDNAttentionMetadataBuilder
+        )
+        if self.supports_update_block_table:
+            # Opts into MRV2's CUDA-only aligned-index precompute.
+            self.mamba_aligned_state_indices: torch.Tensor | None = None
 
         self.decode_cudagraph_max_bs: int = (
             self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1)
@@ -248,12 +261,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         query_start_loc = m.query_start_loc
         query_start_loc_cpu = m.query_start_loc_cpu
         nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
-        block_table_tensor = mamba_get_block_table_tensor(
-            m.block_table_tensor,
-            m.seq_lens,
-            self.kv_cache_spec,
-            self.vllm_config.cache_config.mamba_cache_mode,
-        )
+        aligned_state_indices = getattr(self, "mamba_aligned_state_indices", None)
+        if aligned_state_indices is not None:
+            block_table_tensor = aligned_state_indices[: m.num_reqs]
+        else:
+            block_table_tensor = mamba_get_block_table_tensor(
+                m.block_table_tensor,
+                m.seq_lens,
+                self.kv_cache_spec,
+                self.vllm_config.cache_config.mamba_cache_mode,
+            )
 
         uniform_spec_sequence_length = None
         spec_sequence_masks_cpu: torch.Tensor | None = None
@@ -488,12 +505,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         # metadata below is indexed by request.
         batch_size = m.num_reqs
 
-        if (
-            self.use_full_cuda_graph
-            and num_prefills == 0
-            and num_decodes == 0
-            and num_spec_decodes <= self.decode_cudagraph_max_bs
-            and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
+        if self._stage_spec_decode(
+            num_prefills, num_decodes, num_spec_decodes, num_spec_decode_tokens
         ):
             assert spec_sequence_masks is not None
             _spec_step_cache.clear()
@@ -545,12 +558,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
             num_accepted_tokens = self.num_accepted_tokens[:batch_size]
             num_accepted_tokens[num_spec_decodes:].fill_(1)
 
-        if (
-            self.use_full_cuda_graph
-            and num_prefills == 0
-            and num_spec_decodes == 0
-            and num_decodes <= self.decode_cudagraph_max_bs
-        ):
+        if self._stage_decode(num_prefills, num_decodes, num_spec_decodes):
             self.non_spec_state_indices_tensor[:num_decodes].copy_(
                 non_spec_state_indices_tensor, non_blocking=True
             )
@@ -585,6 +593,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
             spec_state_indices_tensor=spec_state_indices_tensor,
             non_spec_state_indices_tensor=non_spec_state_indices_tensor,
             spec_sequence_masks=spec_sequence_masks,
+            spec_sequence_masks_cpu=spec_sequence_masks_cpu,
             spec_token_indx=spec_token_indx,
             non_spec_token_indx=non_spec_token_indx,
             num_accepted_tokens=num_accepted_tokens,
@@ -595,6 +604,83 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         )
         return attn_metadata
 
+    def _stage_spec_decode(
+        self,
+        num_prefills: int,
+        num_decodes: int,
+        num_spec_decodes: int,
+        num_spec_decode_tokens: int,
+    ) -> bool:
+        """Whether spec-decode metadata goes into the FULL cudagraph buffers."""
+        return (
+            self.use_full_cuda_graph
+            and num_prefills == 0
+            and num_decodes == 0
+            and num_spec_decodes <= self.decode_cudagraph_max_bs
+            and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
+        )
+
+    def _stage_decode(
+        self, num_prefills: int, num_decodes: int, num_spec_decodes: int
+    ) -> bool:
+        """Whether decode metadata goes into the FULL cudagraph buffers."""
+        return (
+            self.use_full_cuda_graph
+            and num_prefills == 0
+            and num_spec_decodes == 0
+            and num_decodes <= self.decode_cudagraph_max_bs
+        )
+
+    def update_block_table(
+        self,
+        metadata: GDNAttentionMetadata,
+        blk_table: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> GDNAttentionMetadata:
+        """Re-gather this group's state indices. The other fields are
+        batch-level and stay shared with ``metadata``."""
+        m = metadata
+        if self.vllm_config.cache_config.mamba_cache_mode == "align":
+            assert self.mamba_aligned_state_indices is not None
+            blk_table = self.mamba_aligned_state_indices
+        masks = m.spec_sequence_masks_cpu
+        spec_indices = non_spec_indices = prefill_indices = None
+        if masks is None:
+            non_spec_indices = blk_table[:, 0]
+            if m.num_prefills > 0:
+                prefill_indices = non_spec_indices[m.num_decodes :]
+        else:
+            spec_indices = blk_table[masks, : self.num_spec + 1]
+            if m.num_prefills > 0:
+                non_spec_indices = prefill_indices = blk_table[~masks, 0]
+
+        if self._stage_spec_decode(
+            m.num_prefills, m.num_decodes, m.num_spec_decodes, m.num_spec_decode_tokens
+        ):
+            assert m.spec_state_indices_tensor is not None
+            batch_size = m.spec_state_indices_tensor.shape[0]
+            self.spec_state_indices_tensor[: m.num_spec_decodes].copy_(
+                spec_indices, non_blocking=True
+            )
+            spec_indices = self.spec_state_indices_tensor[:batch_size]
+            spec_indices[m.num_spec_decodes :].fill_(NULL_BLOCK_ID)
+
+        if self._stage_decode(m.num_prefills, m.num_decodes, m.num_spec_decodes):
+            assert m.non_spec_state_indices_tensor is not None
+            batch_size = m.non_spec_state_indices_tensor.shape[0]
+            self.non_spec_state_indices_tensor[: m.num_decodes].copy_(
+                non_spec_indices, non_blocking=True
+            )
+            non_spec_indices = self.non_spec_state_indices_tensor[:batch_size]
+            non_spec_indices[m.num_decodes :].fill_(NULL_BLOCK_ID)
+
+        return replace(
+            m,
+            spec_state_indices_tensor=spec_indices,
+            non_spec_state_indices_tensor=non_spec_indices,
+            prefill_state_indices=prefill_indices,
+        )
+
     def _build_spec_from_cache(self, m: CommonAttentionMetadata, c: dict) -> GDNAttentionMetadata:
         """The pure spec-decode full-graph case of build(), from this step's cache."""
         block_table_tensor = mamba_get_block_table_tensor(

docker-images/glm53-kindling/experimental/fixes/kv_cache_coordinator.py

@@ -11,7 +11,6 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     KVCacheBlock,
-    dcp_world_size_for_kv_cache_spec,
 )
 from vllm.v1.core.single_type_kv_cache_manager import (
     CrossAttentionManager,
@@ -20,6 +19,7 @@ from vllm.v1.core.single_type_kv_cache_manager import (
     get_manager_for_kv_cache_spec,
 )
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheSpec,
@@ -175,9 +175,7 @@ class KVCacheCoordinator(ABC):
                 role=kv_cache_group.role,
                 enable_caching=enable_caching,
                 kv_cache_group_id=i,
-                dcp_world_size=dcp_world_size_for_kv_cache_spec(
-                    kv_cache_group.kv_cache_spec, dcp_world_size
-                ),
+                dcp_world_size=dcp_world_size,
                 pcp_world_size=pcp_world_size,
                 scheduler_block_size=self.scheduler_block_size,
                 needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing,
@@ -212,6 +210,7 @@ class KVCacheCoordinator(ABC):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         """Get the number of device blocks needed to be allocated for the request.
 
@@ -233,12 +232,19 @@ class KVCacheCoordinator(ABC):
                 per-request admission cap (SWA / chunked-local). Set only by
                 the full-sequence admission gate; per-step allocation must
                 leave it False so the predictor matches `allocate_new_blocks`.
+            prefill_end: The token index the request's prefill ends at, the
+                same value the scheduler splits chunks against. Under sparse
+                retention Mamba reserves a prefill checkpoint only on the chunk
+                reaching it; under dense retention every chunk publishes a
+                state, so it is ignored (0).
 
         Returns:
             The number of blocks to allocate.
 
         """
         num_blocks_to_allocate = 0
+        if self.retention_interval != 0:
+            prefill_end = 0
         for i, manager in enumerate(self.single_type_managers):
             if isinstance(manager, CrossAttentionManager):
                 # For cross-attention, we issue a single static allocation
@@ -261,6 +267,7 @@ class KVCacheCoordinator(ABC):
                     num_local_computed_tokens,
                     num_tokens_main_model,
                     apply_admission_cap=apply_admission_cap,
+                    prefill_end=prefill_end,
                 )
         return num_blocks_to_allocate
 
@@ -700,16 +707,18 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
             f"hash_block_size. block_sizes={cacheable_block_sizes}, "
             f"hash_block_size={hash_block_size}"
         )
-        assert pcp_world_size == 1, "PCP not support hybrid attn now."
-        if dcp_world_size > 1:
-            # DCP shards full-attention KV across ranks and replicates Mamba
-            # state; other spec types (e.g. sliding window) have no DCP-aware
-            # handling yet, so reject them explicitly.
-            for g in kv_cache_config.kv_cache_groups:
-                assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), (
+        for g in kv_cache_config.kv_cache_groups:
+            spec = g.kv_cache_spec
+            replicated = not spec.dcp_sharded
+            if pcp_world_size > 1:
+                assert isinstance(spec, FullAttentionSpec) or (
+                    isinstance(spec, AttentionSpec) and replicated
+                ), "PCP only supports full attention and replicated draft groups."
+            if dcp_world_size > 1:
+                assert isinstance(spec, FullAttentionSpec) or replicated, (
                     "DCP with hybrid KV cache layouts only supports "
-                    "full-attention and Mamba groups, got: "
-                    f"{type(g.kv_cache_spec).__name__}."
+                    "full-attention, Mamba, and replicated draft groups, got: "
+                    f"{type(spec).__name__}."
                 )
         # Fine-grained hash hits require Mamba "align" and compatible cache
         # managers in every group. TP needs hashing finer than the Mamba block;

docker-images/glm53-kindling/experimental/fixes/kv_cache_interface.py

@@ -158,6 +158,8 @@ class KVCacheSpec:
 
     # number of tokens in a block
     block_size: int
+    dcp_sharded: bool = field(default=False, kw_only=True)
+    """Whether DCP shards this cache's token positions across ranks."""
 
     block_stride_alignment: int | None = field(default=None, kw_only=True)
     """Required byte alignment between physical blocks, including packed layers."""
@@ -470,9 +472,11 @@ class HiSparseResidentSpec(KVCacheSpec):
         return cdiv(num_tokens, self.block_size)
 
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
-        return cdiv(vllm_config.model_config.max_model_len, self.block_size) * (
-            self.page_size
+        max_blocks = self.max_admission_blocks_per_request(
+            max_in_flight_tokens=vllm_config.max_in_flight_tokens,
+            max_model_len=vllm_config.model_config.max_model_len,
         )
+        return max_blocks * self.page_size
 
     @property
     def has_layer_views(self) -> bool:
@@ -481,6 +485,10 @@ class HiSparseResidentSpec(KVCacheSpec):
 
 @dataclass(frozen=True, kw_only=True)
 class AttentionSpec(KVCacheSpec):
+    dcp_sharded: bool = True
+    max_tp_shards: int | None = None
+    """Distinct shards this cache splits into across TP; TP ranks beyond this
+    hold replicas. None: unknown, treated as not replicated."""
     num_kv_heads: int
     head_size: int
     dtype: torch.dtype
@@ -536,7 +544,9 @@ class AttentionSpec(KVCacheSpec):
 
     def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
         parallel_config = vllm_config.parallel_config
-        kv_shard_count = parallel_config.decode_context_parallel_size
+        kv_shard_count = (
+            parallel_config.decode_context_parallel_size if self.dcp_sharded else 1
+        )
         return cdiv(max_len, self.block_size * kv_shard_count)
 
 
@@ -568,7 +578,7 @@ class FullAttentionSpec(AttentionSpec):
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         max_model_len = vllm_config.model_config.max_model_len
         dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
-        if dcp_world_size > 1:
+        if self.dcp_sharded and dcp_world_size > 1:
             max_model_len = cdiv(max_model_len, dcp_world_size)
         return cdiv(max_model_len, self.block_size) * self.page_size_bytes
 
@@ -612,6 +622,8 @@ class FullAttentionSpec(AttentionSpec):
             head_size_v=specs[0].head_size_v,
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
+            dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -698,6 +710,8 @@ class MLAAttentionSpec(FullAttentionSpec):
             head_size=specs[0].head_size,
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
+            dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -760,6 +774,8 @@ class RSWASpec(FullAttentionSpec):
             head_size_v=base.head_size_v,
             dtype=base.dtype,
             kv_quant_mode=base.kv_quant_mode,
+            dcp_sharded=base.dcp_sharded,
+            max_tp_shards=base.max_tp_shards,
             page_size_padded=base.page_size_padded,
             num_head_slots=base.num_head_slots,
             state_content_bytes=base.state_content_bytes,
@@ -850,9 +866,10 @@ class SlidingWindowSpec(AttentionSpec):
         return cdiv(num_tokens, self.block_size) + 1
 
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
-        assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
-            "DCP not support sliding window."
-        )
+        assert (
+            not self.dcp_sharded
+            or vllm_config.parallel_config.decode_context_parallel_size == 1
+        ), "DCP not support sliding window."
         max_blocks = self.max_admission_blocks_per_request(
             max_in_flight_tokens=vllm_config.max_in_flight_tokens,
             max_model_len=vllm_config.model_config.max_model_len,
@@ -950,6 +967,8 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
         extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
         bounded_replay_set = set(spec.bounded_replay for spec in specs)
         block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
+        assert len({spec.dcp_sharded for spec in specs}) == 1
+        assert len({spec.max_tp_shards for spec in specs}) == 1
         assert (
             len(cache_dtype_str_set) == 1
             and len(tokens_per_state_set) == 1
@@ -969,6 +988,8 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
             num_kv_heads=specs[0].num_kv_heads,
             head_size=specs[0].head_size,
             dtype=specs[0].dtype,
+            dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -1185,6 +1206,8 @@ class SinkFullAttentionSpec(FullAttentionSpec):
             block_stride_alignment=specs[0].block_stride_alignment,
             dtype=specs[0].dtype,
             kv_quant_mode=specs[0].kv_quant_mode,
+            dcp_sharded=specs[0].dcp_sharded,
+            max_tp_shards=specs[0].max_tp_shards,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
@@ -1217,6 +1240,11 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
 
     kv_cache_specs: dict[str, KVCacheSpec]
 
+    def __post_init__(self):
+        super().__post_init__()
+        if self.kv_cache_specs:
+            object.__setattr__(self, "dcp_sharded", self.first_spec.dcp_sharded)
+
     @property
     def prefix_cacheable(self) -> bool:
         return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
@@ -1270,6 +1298,8 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
         if len(block_sizes) > 1:
             # Different block sizes, not uniform.
             return False
+        if len({spec.dcp_sharded for spec in kv_cache_specs.values()}) > 1:
+            return False
         first_spec = next(iter(kv_cache_specs.values()))
         return first_spec.is_uniform_with_collection(kv_cache_specs)
 
@@ -1464,6 +1494,9 @@ class KVCacheConfig:
     hisparse_shared_host_pool: bool = False
     """Whether local TP ranks share one physical HiSparse host pool."""
 
+    kv_tp_replicas: int = 1
+    """Consecutive TP ranks holding identical KV for every layer (1: none)."""
+
     draft_num_blocks: int | None = None
     """GLM-5-Next with a drafter: blocks in the drafter groups' own pool, sized
     for their sliding windows instead of one draft page per target block."""

docker-images/glm53-kindling/experimental/fixes/kv_cache_utils.py

@@ -9,7 +9,7 @@ import os
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from dataclasses import dataclass, field, replace
-from functools import partial
+from functools import partial, reduce
 from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
 
 from vllm import envs
@@ -21,7 +21,6 @@ from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import format_gib
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.hisparse.layout import (
-    get_hisparse_gpu_memory_usage,
     get_hisparse_host_pool_bytes,
     get_hisparse_kv_cache_config,
     get_hisparse_kv_cache_groups,
@@ -502,8 +501,8 @@ def _gen_mm_extra_hash_keys(
 ) -> tuple[list[Any], int]:
     """Generate extra keys related to MultiModal request for block hash
     computation. For multi-modal inputs, the extra keys are
-    (mm_hash, start_offset) that indicate a mm input contained in the
-    block and its starting offset in the block tokens.
+    ("mm", mm_hash, start_offset) tuples that indicate a mm input contained
+    in the block and its starting offset in the block tokens.
 
     Args:
         request: The request object.
@@ -549,7 +548,7 @@ def _gen_mm_extra_hash_keys(
             # relative to the start of the block so prefix-cache keys stay
             # distinct when the same MM item appears at different positions
             # within otherwise-identical placeholder blocks.
-            extra_keys.append((mm_feature.identifier, offset - start_token_idx))
+            extra_keys.append(("mm", mm_feature.identifier, offset - start_token_idx))
 
             if end_token_idx >= offset + length:
                 # If this block contains the end of the current mm input,
@@ -565,25 +564,29 @@ def _gen_mm_extra_hash_keys(
     return extra_keys, curr_mm_idx
 
 
-def _gen_lora_extra_hash_keys(request: Request) -> list[str]:
+def _gen_lora_extra_hash_keys(request: Request) -> list[tuple[str, str, str]]:
     """Generate extra keys related to LoRA for block hash computation.
 
+    The adapter path is included so that re-pointing a LoRA name at a different
+    adapter does not reuse KV computed with the previous one.
+
     Args:
         request: The request object.
 
     Returns:
-        Return LoRA name of the request if it is a LoRA request. Return empty
-        list otherwise.
+        Return the LoRA name and path of the request if it is a LoRA request.
+        Return empty list otherwise.
 
     """
-    if not request.lora_request:
+    lora_request = request.lora_request
+    if not lora_request:
         return []
-    return [request.lora_request.lora_name]
+    return [("lora", lora_request.lora_name, lora_request.lora_path)]
 
 
 def _gen_prompt_embeds_extra_hash_keys(
     request: Request, start_token_idx: int, end_token_idx: int
-) -> list[bytes]:
+) -> list[tuple[str, bytes]]:
     """Generate extra keys related to prompt embeds for block hash computation.
 
     Args:
@@ -605,7 +608,7 @@ def _gen_prompt_embeds_extra_hash_keys(
         # Hash prompt embeds once per block and cache on request
         embeds_hash = hashlib.sha256(tensor_data(block_prompt_embeds)).digest()
         request._prompt_embeds_per_block_hashes[block_range] = embeds_hash
-    return [embeds_hash]
+    return [("prompt_embeds", embeds_hash)]
 
 
 def generate_block_hash_extra_keys(
@@ -629,9 +632,11 @@ def generate_block_hash_extra_keys(
     mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(
         request, start_token_idx, end_token_idx, start_mm_idx
     )
-    lora_extra_keys: list[str] = _gen_lora_extra_hash_keys(request)
-    cache_salt_keys: list[str] = (
-        [request.cache_salt] if (start_token_idx == 0 and request.cache_salt) else []
+    lora_extra_keys = _gen_lora_extra_hash_keys(request)
+    cache_salt_keys: list[tuple[str, str]] = (
+        [("cache_salt", request.cache_salt)]
+        if (start_token_idx == 0 and request.cache_salt)
+        else []
     )
     prompt_embeds_keys = _gen_prompt_embeds_extra_hash_keys(
         request, start_token_idx, end_token_idx
@@ -647,6 +652,35 @@ def generate_block_hash_extra_keys(
     return tuple(extra_keys), new_start_mm_idx
 
 
+def to_event_extra_keys(
+    extra_keys: Iterable[tuple[Any, ...] | None] | None,
+) -> list[tuple[Any, ...] | None] | None:
+    """Convert block-hash extra keys to the untagged per-block list published
+    in KV events.
+
+    External KV event consumers parse the pre-tagging shapes: bare LoRA names
+    and cache salts, ``(mm_identifier, offset)`` pairs and bare prompt-embeds
+    digests. Events keep that format until consumers handle the tagged keys.
+
+    Args:
+        extra_keys: One entry per block, each as returned by
+            `generate_block_hash_extra_keys`, or None.
+
+    Returns:
+        One untagged entry per block, or None if there are no entries.
+
+    """
+    if not extra_keys:
+        return None
+    event_keys = [
+        None
+        if keys is None
+        else tuple(key[1:] if key[0] == "mm" else key[1] for key in keys)
+        for keys in extra_keys
+    ]
+    return event_keys or None
+
+
 def hash_block_tokens(
     hash_function: Callable[[Any], bytes],
     parent_block_hash: BlockHash | None,
@@ -682,12 +716,7 @@ def hash_block_tokens(
 
 def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
     """Return the token span of a cache block under DCP."""
-    layer_specs = iter_layer_specs(spec)
-    if len(layer_specs) > 0 and all(
-        isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs
-    ):
-        return spec.block_size * dcp_world_size
-    return spec.block_size
+    return spec.block_size * (dcp_world_size if spec.dcp_sharded else 1)
 
 
 def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:
@@ -707,28 +736,6 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache
     return replace(spec, block_size=block_size)
 
 
-def dcp_world_size_for_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> int:
-    """Return the DCP size that owns this group's block geometry.
-
-    Full-attention KV (including MLA) is sharded across DCP ranks, so prefix
-    hashing and manager ``block_size`` use the process DCP size. Other specs
-    keep replicated per-rank state (Mamba, sliding window, chunked-local) and
-    must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
-
-    Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
-    ``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated
-    draft group would need a different spec, not this helper.
-    """
-    if dcp_world_size <= 1:
-        return 1
-    inner = spec
-    if isinstance(spec, UniformTypeKVCacheSpecs):
-        inner = next(iter(spec.kv_cache_specs.values()))
-    if isinstance(inner, FullAttentionSpec):
-        return dcp_world_size
-    return 1
-
-
 def resolve_kv_cache_block_sizes(
     kv_cache_config: KVCacheConfig,
     vllm_config: VllmConfig,
@@ -752,6 +759,8 @@ def resolve_kv_cache_block_sizes(
     groups = kv_cache_config.kv_cache_groups
 
     if len(groups) <= 1:
+        if groups and not groups[0].kv_cache_spec.dcp_sharded:
+            dcp = 1
         bs = cache_config.block_size * dcp
         return bs, bs
 
@@ -1133,7 +1142,9 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
     `available_memory` into `num_blocks`. Used to compute the effective KV cache
     capacity once `num_gpu_blocks_override` is applied.
     """
-    return _get_kv_cache_bytes_per_block(kv_cache_groups)
+    return _get_kv_cache_bytes_per_block(
+        [group for group in kv_cache_groups if not group.host_resident]
+    )
 
 
 def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
@@ -2132,18 +2143,21 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
         kv_cache_spec: The kv cache spec of each attention layer in the model
 
     """
-    if is_kv_cache_spec_uniform(
-        kv_cache_spec
-    ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):
-        return
-
-    logger.warning(
-        "Hybrid KV cache manager is disabled for this hybrid model, "
-        "This means we do not enable any optimizations for saving KV cache "
-        "memory (e.g., dropping the KV cache outside the sliding window). "
-        "The compute of layers like sliding window is still saved."
-    )
-    kv_cache_spec.update(_promote_local_kv_cache_specs(kv_cache_spec))
+    groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict)
+    for name, spec in kv_cache_spec.items():
+        replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded
+        groups[replicated][name] = spec
+    for specs in groups.values():
+        promoted_specs = _promote_local_kv_cache_specs(specs)
+        if promoted_specs == specs:
+            continue
+        logger.warning(
+            "Hybrid KV cache manager is disabled for this hybrid model, "
+            "This means we do not enable any optimizations for saving KV cache "
+            "memory (e.g., dropping the KV cache outside the sliding window). "
+            "The compute of layers like sliding window is still saved."
+        )
+        kv_cache_spec.update(promoted_specs)
 
 
 def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
@@ -2473,7 +2487,7 @@ def _ensure_min_page_size(
     scaled: list[KVCacheGroupSpec] = []
     for g in groups:
         s = g.kv_cache_spec
-        kw: dict[str, int] = {"block_size": s.block_size * scale}
+        kw: dict[str, Any] = {"block_size": s.block_size * scale}
         if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None:
             kw["page_size_padded"] = s.page_size_padded * scale
         scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw)))
@@ -2580,6 +2594,26 @@ def get_kv_cache_groups(
     return groups
 
 
+def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int:
+    if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None:
+        return 1
+    if spec.dcp_sharded and dcp_size > 1:
+        return 1
+    return max(1, tp_size // spec.max_tp_shards)
+
+
+def kv_cache_groups_tp_replicas(
+    groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1
+) -> int:
+    """Consecutive TP ranks holding identical KV for every layer."""
+    specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)]
+    if not specs:
+        return 1
+    return reduce(
+        math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size
+    )
+
+
 def generate_scheduler_kv_cache_config(
     kv_cache_configs: list[KVCacheConfig],
 ) -> KVCacheConfig:
@@ -2647,12 +2681,10 @@ def _max_memory_usage_bytes_from_groups(
     Each group independently claims blocks from the shared pool, so a request consumes
     the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``.
     """
+    kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident]
     if not kv_cache_groups:
         return 0
 
-    if vllm_config.attention_config.hisparse_config is not None:
-        return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups)
-
     if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
         (
             attn_group,
@@ -2727,12 +2759,11 @@ def _estimate_max_model_len_from_groups(
         vllm_config.model_config.max_model_len = model_len
         if hisparse_enabled:
             try:
-                config = get_kv_cache_config_from_groups(
+                get_kv_cache_config_from_groups(
                     vllm_config, kv_cache_groups, available_memory
                 )
             except ValueError:
                 return False
-            return get_max_concurrency_for_kv_cache_config(vllm_config, config) >= 1
         return (
             _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
             <= available_memory
@@ -2842,6 +2873,9 @@ def _project_kv_cache_groups_to_worker(
         worker_layer_names = [
             layer_name for layer_name in group.layer_names if layer_name in worker_spec
         ]
+        if len(worker_layer_names) == len(group.layer_names):
+            projected_groups.append(group)
+            continue
         group_spec = group.kv_cache_spec
         if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs):
             group_spec = UniformTypeKVCacheSpecs(
@@ -2957,9 +2991,6 @@ def get_kv_cache_configs(
             adjusted_memory.append(override * bytes_per_block)
         available_memory = adjusted_memory
 
-    if vllm_config.attention_config.hisparse_config is not None:
-        available_memory = [min(available_memory)] * len(available_memory)
-
     # Reserve the null block BlockPool permanently holds back, so auto-fit and
     # the capacity check both plan against usable blocks. Allocation below
     # still uses the full memory.
@@ -3012,6 +3043,13 @@ def get_kv_cache_configs(
             min_num_blocks * _pool_bytes_per_block(groups) + _glm5_next_draft_pool_bytes(vllm_config, groups),
         )
 
+    for kv_cache_config in kv_cache_configs:
+        kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
+            kv_cache_config.kv_cache_groups,
+            vllm_config.parallel_config.tensor_parallel_size,
+            vllm_config.parallel_config.decode_context_parallel_size,
+        )
+
     return kv_cache_configs
 
 

docker-images/glm53-kindling/experimental/fixes/worker_utils.py

@@ -394,7 +394,13 @@ def select_common_block_size(
     for size in sorted(candidates, reverse=True):
         if block_size_is_supported(backends, size):
             return size
-    raise ValueError(f"No common block size for {kv_manager_block_size}.")
+    raise ValueError(
+        f"No common block size for {kv_manager_block_size} ("
+        + "; ".join(
+            f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" for b in backends
+        )
+        + ")."
+    )
 
 
 def allocate_kv_cache(
@@ -752,10 +758,12 @@ def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) ->
 
 
 def get_uniform_decode_token_count(
-    num_reqs: int, num_tokens: int, max_query_len: int, has_prefill: bool
+    num_reqs: int, num_tokens: int, max_query_len: int, decode_graph_eligible: bool
 ) -> int | None:
     """Per-request token count of a uniform decode batch, or None."""
-    if not has_prefill and is_uniform_query_len(num_reqs, num_tokens, max_query_len):
+    if decode_graph_eligible and is_uniform_query_len(
+        num_reqs, num_tokens, max_query_len
+    ):
         return max_query_len
     return None
 

docker-images/glm53-kindling/image/Dockerfile

@@ -33,6 +33,15 @@
 # unchanged, except routed_experts.py where vllm#58720 landed verbatim upstream
 # (the patch is dropped below), and flashinfer 0.7.0 -> 0.7.0.post1.
 # torch/NCCL/cutlass-dsl pins unchanged.
+#
+# The experimental/ overlays (baked by build.sh) were written against the
+# ddd6fbca tree. Rebased 2026-10-15 onto v0.31.0: for every overlay whose
+# base file drifted, new = v0.31.0 file + (overlay - ddd6fbca) delta; 21
+# applied cleanly, 2 needed a hand merge (rejection_sampler: the
+# gather_draft_sampled() refactor; flashinfer_mla_sparse_sm90: the
+# 3-tuple _kv_lens_host + exact flag). Unrebaked overlays reintroduce
+# ddd6fbca-era import blocks on top of the v0.31.0 tree and die at import
+# time (e.g. hisparse get_hisparse_gpu_memory_usage).
 #   kept:
 #     glm53-flash_SM121.py (MiaAI)  5 of 7 edits apply unchanged; see below
 #     gb10_plugin_backend.py        MiaAI vs LibertAI sparse-MLA switch