feat(kda): export up to four retained prefill checkpoints on GB10

b12x/sequence/kda_prefill/__init__.py

@@ -11,9 +11,18 @@
 ``[slot, head, value_dim, key_dim]`` in fp32, so a prefill and a decode of the
 same request share one pool without conversion. State slots are addressed by
 index rather than gathered: a request names its initial slot, its final slot,
-and optionally one checkpoint slot with a chunk-aligned token offset, and the
-op reads and writes those slots directly. ``Caps.null_state_index`` may reserve
-one index meaning "zero initial state" and "do not write".
+and optionally checkpoint slots with chunk-aligned token offsets. The op reads
+and writes those slots directly. ``Caps.max_checkpoints`` defaults to one with
+one-checkpoint vector metadata. Explicit ``max_checkpoints=2`` or ``4`` uses contiguous
+``[sequence_capacity, max_checkpoints]`` checkpoint indices/offsets, requires
+checkpoint export and transactional validation, and exports the enabled states
+during the same recurrence. Multi-checkpoint export is research-only: policy
+validation accepts NVIDIA GB10 (SM121, 48 SMs). Twelve four-checkpoint GPU
+tests passed on GB10, covering checkpoint values, graph replay, invalid metadata,
+and pool addressing. This coverage is not a measured component performance
+profile; the embedded registry has no measured KDA-prefill profile.
+``Caps.null_state_index``
+may reserve one index meaning "zero initial state" and "do not write".
 
 Requests are packed. Request ``r`` covers tokens
 ``cu_seqlens[r]:cu_seqlens[r + 1]``; ``num_seqs`` and ``num_tokens`` are device
@@ -25,8 +34,8 @@
 use caller-owned scratch, allocate no tensor storage, and are capture safe.
 Device-side validation is transactional: bit 0 reports a duplicate or
 conflicting write slot, bit 1 malformed packed metadata, bit 2 an invalid state
-slot, and bit 3 an unusable checkpoint offset. Any error poisons the live
-output rows without mutating recurrent state.
+slot, and bit 3 an unusable checkpoint offset. Any error poisons the full bound
+output capacity without mutating recurrent state.
 """
 
 from __future__ import annotations

b12x/sequence/kda_prefill/_cute_kernels.py

@@ -294,6 +294,7 @@ def __init__(
         flag_count: int,
         max_state_slots: int,
         validate: bool,
+        max_checkpoints: int,
         null_state_index: int | None,
         index_type: type[cutlass.Numeric],
     ) -> None:
@@ -305,6 +306,7 @@ def __init__(
         self.flag_count = int(flag_count)
         self.max_state_slots = int(max_state_slots)
         self.validate = bool(validate)
+        self.max_checkpoints = int(max_checkpoints)
         self.has_null = null_state_index is not None
         self.null_state_index = 0 if null_state_index is None else int(null_state_index)
         self.index_type = index_type
@@ -542,8 +544,6 @@ def kernel(
                         flags[1] = Int32(1)
                     initial = Int64(initial_indices[seq])
                     final = Int64(final_indices[seq.to(Int64) * final_stride])
-                    checkpoint = Int64(checkpoint_indices[seq])
-                    offset = checkpoint_offsets[seq].to(Int32)
                     slot_limit = Int64(self.max_state_slots)
                     if not self._is_null(initial):
                         if (initial < Int64(0)) | (initial >= slot_limit):
@@ -553,16 +553,28 @@ def kernel(
                             flags[2] = Int32(1)
                         elif self._insert(table, final) != Int32(0):
                             flags[0] = Int32(1)
-                    if offset > length:
-                        flags[3] = Int32(1)
-                    if (offset > Int32(0)) & ((offset % Int32(_CHUNK)) != Int32(0)):
-                        flags[3] = Int32(1)
-                    if offset > Int32(0):
-                        if not self._is_null(checkpoint):
-                            if (checkpoint < Int64(0)) | (checkpoint >= slot_limit):
-                                flags[2] = Int32(1)
-                            elif self._insert(table, checkpoint) != Int32(0):
-                                flags[0] = Int32(1)
+                    for cp in cutlass.range_constexpr(self.max_checkpoints):
+                        cp_index = seq * Int32(self.max_checkpoints) + Int32(cp)
+                        checkpoint = Int64(checkpoint_indices[cp_index])
+                        offset = checkpoint_offsets[cp_index].to(Int32)
+                        if offset > length:
+                            flags[3] = Int32(1)
+                        if (offset > Int32(0)) & ((offset % Int32(_CHUNK)) != Int32(0)):
+                            flags[3] = Int32(1)
+                        if offset > Int32(0):
+                            if not self._is_null(checkpoint):
+                                if (checkpoint < Int64(0)) | (checkpoint >= slot_limit):
+                                    flags[2] = Int32(1)
+                                elif self._insert(table, checkpoint) != Int32(0):
+                                    flags[0] = Int32(1)
+                                if cutlass.const_expr(cp > 0):
+                                    for previous in cutlass.range_constexpr(cp):
+                                        previous_index = seq * Int32(self.max_checkpoints) + Int32(previous)
+                                        previous_slot = Int64(checkpoint_indices[previous_index])
+                                        previous_offset = checkpoint_offsets[previous_index].to(Int32)
+                                        if not self._is_null(previous_slot):
+                                            if previous_offset == offset:
+                                                flags[3] = Int32(1)
                 cute.arch.atomic_add(hist.iterator + count, Int32(1))
             counts[seq] = count
             seq += Int32(_PROLOGUE_THREADS)
@@ -1053,6 +1065,7 @@ def __init__(
         k_split: int,
         stages: int,
         checkpoint_export: bool,
+        max_checkpoints: int,
         null_state_index: int | None,
         index_type: type[cutlass.Numeric],
     ) -> None:
@@ -1076,6 +1089,7 @@ def __init__(
         self.kb_steps = self.cols // 16
         self.nb_blocks = self.cols // 8
         self.checkpoint_export = bool(checkpoint_export)
+        self.max_checkpoints = int(max_checkpoints)
         self.has_null = null_state_index is not None
         self.null_state_index = 0 if null_state_index is None else int(null_state_index)
         self.index_type = index_type
@@ -1505,8 +1519,6 @@ def kernel(
                 if has_tiles:
                     initial = Int64(initial_indices[seq])
                     final = Int64(final_indices[seq.to(Int64) * final_stride])
-                    checkpoint = Int64(checkpoint_indices[seq])
-                    offset = checkpoint_offsets[seq].to(Int32)
                     for nb in cutlass.range_constexpr(self.nb_blocks):
                         acc[nb, 0] = Float32(0.0)
                         acc[nb, 1] = Float32(0.0)
@@ -1740,12 +1752,16 @@ def kernel(
                             count += Int32(1)
 
                             if cutlass.const_expr(self.checkpoint_export):
-                                if (offset > Int32(0)) & ((local + Int32(1)) * Int32(_CHUNK) == offset):
-                                    if not self._is_null(checkpoint):
-                                        self._store_state(
-                                            recurrent_state, acc, checkpoint * slot_stride + head_base,
-                                            row0, row1, col_base, tid,
-                                        )
+                                for cp in cutlass.range_constexpr(self.max_checkpoints):
+                                    cp_index = seq * Int32(self.max_checkpoints) + Int32(cp)
+                                    checkpoint = Int64(checkpoint_indices[cp_index])
+                                    offset = checkpoint_offsets[cp_index].to(Int32)
+                                    if (offset > Int32(0)) & ((local + Int32(1)) * Int32(_CHUNK) == offset):
+                                        if not self._is_null(checkpoint):
+                                            self._store_state(
+                                                recurrent_state, acc, checkpoint * slot_stride + head_base,
+                                                row0, row1, col_base, tid,
+                                            )
                         # Final state, or the running state for the next window.
                         if not self._is_null(final):
                             self._store_state(
@@ -1795,6 +1811,7 @@ def _recurrence_key(binding: Binding) -> tuple[object, ...]:
         plan.k_split,
         plan.stages,
         caps.checkpoint_export,
+        caps.max_checkpoints,
         caps.null_state_index,
         binding.initial_state_indices.dtype,
     )
@@ -1816,6 +1833,7 @@ def _compile_recurrence(binding: Binding) -> tuple[tuple[object, ...], Callable[
         k_split=binding.plan.k_split,
         stages=binding.plan.stages,
         checkpoint_export=caps.checkpoint_export,
+        max_checkpoints=caps.max_checkpoints,
         null_state_index=caps.null_state_index,
         index_type=index_type,
     )
@@ -1844,7 +1862,7 @@ def _compile_recurrence(binding: Binding) -> tuple[tuple[object, ...], Callable[
         Int32(1),
         Int32(0),
         current_cuda_stream(),
-        compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.recurrence", 8, key),
+        compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.recurrence", 9, key),
     )
 
     def launch(active: Binding, window: int) -> None:
@@ -1907,6 +1925,7 @@ def _prologue_key(binding: Binding) -> tuple[object, ...]:
         2 * binding.plan.window_tiles * caps.heads,
         caps.max_state_slots,
         caps.metadata_validation,
+        caps.max_checkpoints,
         caps.null_state_index,
         binding.initial_state_indices.dtype,
     )
@@ -1928,6 +1947,7 @@ def _compile_prologue(binding: Binding) -> tuple[tuple[object, ...], Callable[..
         flag_count=2 * binding.plan.window_tiles * caps.heads,
         max_state_slots=caps.max_state_slots,
         validate=caps.metadata_validation == "transactional",
+        max_checkpoints=caps.max_checkpoints,
         null_state_index=caps.null_state_index,
         index_type=index_type,
     )
@@ -1955,7 +1975,7 @@ def _compile_prologue(binding: Binding) -> tuple[tuple[object, ...], Callable[..
         Int32(1),
         Int32(1),
         current_cuda_stream(),
-        compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.prologue", 5, key),
+        compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.prologue", 6, key),
     )
 
     def launch(active: Binding, launched_tiles: int) -> None:

b12x/sequence/kda_prefill/_impl.py

@@ -53,6 +53,7 @@ class Caps:
     state_dtype: torch.dtype = torch.float32
     qk_l2norm: bool = True
     checkpoint_export: bool = False
+    max_checkpoints: int = 1
     null_state_index: int | None = None
     metadata_validation: MetadataValidation = "transactional"
     chunk_tokens: int = 16
@@ -79,6 +80,10 @@ def __post_init__(self) -> None:
             raise ValueError("metadata_validation must be 'transactional' or 'trusted'")
         object.__setattr__(self, "qk_l2norm", bool(self.qk_l2norm))
         object.__setattr__(self, "checkpoint_export", bool(self.checkpoint_export))
+        if type(self.max_checkpoints) is not int or self.max_checkpoints not in (1, 2, 4):
+            raise ValueError("max_checkpoints must be 1, 2 or 4")
+        if self.max_checkpoints > 1 and (not self.checkpoint_export or self.metadata_validation != "transactional"):
+            raise ValueError("multiple checkpoints require checkpoint_export and transactional validation")
         if self.null_state_index is not None:
             null = int(self.null_state_index)
             if null < 0 or null >= self.max_state_slots:
@@ -199,6 +204,7 @@ def _query(caps: Caps) -> KdaPrefillQuery:
         state_dtype=str(caps.state_dtype).removeprefix("torch."),
         qk_l2norm=caps.qk_l2norm,
         checkpoint_export=caps.checkpoint_export,
+        max_checkpoints=caps.max_checkpoints,
         max_tokens=caps.max_tokens,
         max_seqs=caps.max_seqs,
     )
@@ -220,7 +226,7 @@ def _materialize_plan(
     window_tiles = max(1, min(int(window_tiles), tiles))
     max_windows = -(-tiles // window_tiles)
     ring_records = 2 * window_tiles * heads
-    duplicate_table_size = _next_power_of_two(4 * caps.max_seqs)
+    duplicate_table_size = _next_power_of_two(2 * (1 + caps.max_checkpoints) * caps.max_seqs)
     regions = (
         ("error_code", 1, torch.int32),
         ("duplicate_slots", duplicate_table_size, torch.int32),
@@ -342,11 +348,10 @@ def bind(
     )
     require_tensor("cu_seqlens", cu_seqlens, shape=(seq_capacity + 1,), device=device, dtypes=(torch.int32,))
     index_dtypes = (torch.int32, torch.int64)
-    for name, tensor in (
-        ("initial_state_indices", initial_state_indices),
-        ("checkpoint_state_indices", checkpoint_state_indices),
-    ):
-        require_tensor(name, tensor, shape=(seq_capacity,), device=device, dtypes=index_dtypes)
+    require_tensor(
+        "initial_state_indices", initial_state_indices,
+        shape=(seq_capacity,), device=device, dtypes=index_dtypes,
+    )
     require_tensor(
         "final_state_indices",
         final_state_indices,
@@ -357,10 +362,13 @@ def bind(
     )
     if final_state_indices.stride(0) <= 0:
         raise ValueError("final_state_indices must have a positive stride")
+    checkpoint_shape = (seq_capacity,) if caps.max_checkpoints == 1 else (seq_capacity, caps.max_checkpoints)
+    require_tensor("checkpoint_state_indices", checkpoint_state_indices, shape=checkpoint_shape,
+                   device=device, dtypes=index_dtypes)
     if not (initial_state_indices.dtype == final_state_indices.dtype == checkpoint_state_indices.dtype):
         raise TypeError("state index tensors must share one dtype")
     require_tensor(
-        "checkpoint_offsets", checkpoint_offsets, shape=(seq_capacity,), device=device, dtypes=(torch.int32,)
+        "checkpoint_offsets", checkpoint_offsets, shape=checkpoint_shape, device=device, dtypes=(torch.int32,)
     )
     for name, tensor in (("num_seqs", num_seqs), ("num_tokens", num_tokens)):
         require_tensor(name, tensor, shape=(1,), device=device, dtypes=(torch.int32,))

b12x/sequence/kda_prefill/_policy.py

@@ -6,14 +6,21 @@
 
 from b12x.policy import ComponentPolicy
 from b12x.policy.components import KDA_PREFILL
-from b12x.policy.types import FrozenMapping
+from b12x.policy.types import DeviceIdentity, FrozenMapping
 
 BACKEND = "cutedsl"
 V_SPLIT_CHOICES = (16, 32, 64, 128)
 K_SPLIT_CHOICES = (1, 2, 4)
 STAGE_CHOICES = (2, 3, 4)
 CHUNK_TOKENS = 16
 
+# Multi-checkpoint device eligibility is checked once by the component policy.
+# This identity restricts planning; it is not a measured GPU profile.
+_MULTI_CHECKPOINT_TARGET = DeviceIdentity(
+    vendor="nvidia", product_name="NVIDIA GB10",
+    compute_capability=(12, 1), sm_count=48,
+)
+
 
 class WorkspaceRecord:
     """Byte layout of one prepared (tile, head) record in the workspace ring.
@@ -65,6 +72,7 @@ class KdaPrefillQuery:
     checkpoint_export: bool
     max_tokens: int
     max_seqs: int
+    max_checkpoints: int = 1
 
     def profile_fields(self) -> dict[str, object]:
         return {
@@ -76,6 +84,7 @@ def profile_fields(self) -> dict[str, object]:
             "checkpoint_export": bool(self.checkpoint_export),
             "max_tokens": int(self.max_tokens),
             "max_seqs": int(self.max_seqs),
+            "max_checkpoints": int(self.max_checkpoints),
         }
 
 
@@ -144,7 +153,16 @@ def _heuristic(query: KdaPrefillQuery, device) -> KdaPrefillConfig:
 
 
 def _validate(query: KdaPrefillQuery, config: KdaPrefillConfig, device) -> None:
-    del device
+    if type(query.max_checkpoints) is not int or query.max_checkpoints not in (1, 2, 4):
+        raise ValueError("max_checkpoints must be 1, 2 or 4")
+    if query.max_checkpoints > 1:
+        if not query.checkpoint_export:
+            raise ValueError("multiple checkpoints require checkpoint_export")
+        if device != _MULTI_CHECKPOINT_TARGET:
+            raise ValueError(
+                "multi-checkpoint KDA prefill supports only NVIDIA GB10 "
+                "(SM121, 48 SMs); use max_checkpoints=1 on other devices"
+            )
     if config.backend != BACKEND:
         raise ValueError(f"unsupported {KDA_PREFILL} backend {config.backend!r}")
     if config.v_split not in V_SPLIT_CHOICES:
@@ -177,7 +195,7 @@ def _validate(query: KdaPrefillQuery, config: KdaPrefillConfig, device) -> None:
 
 KDA_PREFILL_POLICY = ComponentPolicy(
     component_id=KDA_PREFILL,
-    query_schema_version=1,
+    query_schema_version=3,
     config_schema_version=1,
     query_fields=frozenset(
         {
@@ -189,6 +207,7 @@ def _validate(query: KdaPrefillQuery, config: KdaPrefillConfig, device) -> None:
             "checkpoint_export",
             "max_tokens",
             "max_seqs",
+            "max_checkpoints",
         }
     ),
     config_fields=frozenset({"backend", "v_split", "k_split", "stages", "window_tiles"}),