feat(ple): export prepared internal prefill checkpoints

b12x/sequence/ple/STATE.md

@@ -35,6 +35,23 @@ persists the newest `state_length` normalized inputs, and clears the speculative
 tail. `state_is_fresh[r]` makes an existing physical slot read as zero without
 requiring the integration to clear recycled storage first.
 
+`export_checkpoint` copies an internal prefill window from the immediately
+preceding `run_mixed` invocation. Its int32 offsets and int64 destination slots
+have the plan's request capacity. An offset is relative to the request's query
+start and must lie strictly inside that query. Zero offsets, negative slots,
+decode rows and empty rows leave state unchanged. Each enabled destination must
+be a valid, distinct pool slot separate from every live input/output slot.
+The export kernel checks these destination constraints on the device before
+copying. An out-of-range destination or one matching any live state-slot ID
+is skipped. If enabled exports share a destination, all of those exports are
+skipped; independent valid exports still run. Disabled and inactive exports
+do not reserve destinations. These checks read runtime metadata on every
+CUDA graph replay and require no additional storage or host readback.
+Offsets shorter than the convolution window include the saved input history;
+the speculative tail is cleared. Export must finish before the mixed binding's
+scratch is reused. Kernel preparation includes the export path, and runtime
+offsets and slot IDs may change during CUDA graph replay.
+
 A mixed plan binds a fixed-capacity device boolean `request_is_prefill` with one
 entry per request row. `run_mixed` applies prefill semantics to true live rows
 and decode semantics to false live rows without partitioning or reordering the
@@ -50,7 +67,7 @@ A `state_slot_ids[r]` value of `-1` is a dummy sink for CUDA-graph padding. Its
 tokens produce zero output and no state mutation. Other negative values and
 values at or above `max_state_slots` are invalid.
 
-The kernels read `num_seqs`, `num_tokens`, `query_start_loc`,
+The mixed, prefill and decode kernels read `num_seqs`, `num_tokens`, `query_start_loc`,
 `state_slot_ids`, and `num_accepted_tokens` from the device without checking
 them against the planned capacities or against each other; the only runtime
 masks are the live token count, the live request count, and the `-1` slot

b12x/sequence/ple/__init__.py

@@ -34,6 +34,7 @@
         "run_decode",
         "run_mixed",
         "run_prefill",
+        "export_checkpoint",
         "is_supported",
     ),
     dtypes=("bf16", "int64"),
@@ -71,6 +72,7 @@
         run_decode,
         run_mixed,
         run_prefill,
+        export_checkpoint,
     )
 
 install_lazy_api(globals(), META)

b12x/sequence/ple/_kernels.py

@@ -410,6 +410,64 @@ def _update_state_kernel(
 
 
 
+@triton.jit(do_not_specialize=["max_state_slots"])
+def _export_checkpoint_kernel(
+    normalized_u_ptr, gathered_state_ptr, query_start_loc_ptr,
+    checkpoint_offsets_ptr, checkpoint_slots_ptr, request_is_prefill_ptr,
+    num_seqs_ptr, state_slot_ids_ptr, conv_state_ptr, max_state_slots,
+    CHANNELS: tl.constexpr, STATE_LENGTH: tl.constexpr,
+    STATE_CAPACITY: tl.constexpr, STATE_STRIDE: tl.constexpr,
+    MAX_SEQS: tl.constexpr, BLOCK: tl.constexpr,
+):
+    """Copy an internal prefill window into an independent state slot."""
+    request = tl.program_id(0)
+    element = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    num_seqs = tl.load(num_seqs_ptr)
+    live = request < num_seqs
+    offset = tl.load(checkpoint_offsets_ptr + request, live, other=0)
+    slot = tl.load(checkpoint_slots_ptr + request, live, other=-1).to(tl.int64)
+    prefill = tl.load(request_is_prefill_ptr + request, live, other=False)
+    start = tl.load(query_start_loc_ptr + request, live, other=0)
+    end = tl.load(query_start_loc_ptr + request + 1, live, other=0)
+    valid = live & prefill & (slot >= 0) & (slot < max_state_slots) & (offset > 0) & (offset < end - start)
+    # Check runtime ownership in every copying program so conflicting rows
+    # cannot race, including when metadata changes during graph replay.
+    peers = tl.arange(0, triton.next_power_of_2(MAX_SEQS))
+    peer_live = (peers < MAX_SEQS) & (peers < num_seqs)
+    live_slots = tl.load(state_slot_ids_ptr + peers, peer_live, other=-1)
+    peer_slots = tl.load(checkpoint_slots_ptr + peers, peer_live, other=-1)
+    peer_offsets = tl.load(checkpoint_offsets_ptr + peers, peer_live, other=0)
+    peer_prefill = tl.load(request_is_prefill_ptr + peers, peer_live, other=False)
+    peer_start = tl.load(query_start_loc_ptr + peers, peer_live, other=0)
+    peer_end = tl.load(query_start_loc_ptr + peers + 1, peer_live, other=0)
+    peer_enabled = peer_prefill & (peer_offsets > 0) & (peer_offsets < peer_end - peer_start)
+    conflicts = peer_live & (
+        (slot == live_slots)
+        | (peer_enabled & (peers != request) & (slot == peer_slots))
+    )
+    valid = valid & (tl.sum(conflicts.to(tl.int32), 0) == 0)
+    channel = element // STATE_CAPACITY
+    position = element % STATE_CAPACITY
+    relative = offset - STATE_LENGTH + position
+    payload = (channel < CHANNELS) & (position < STATE_LENGTH)
+    token = start.to(tl.int64) + relative.to(tl.int64)
+    query = tl.load(
+        normalized_u_ptr + token * CHANNELS + channel.to(tl.int64),
+        valid & payload & (relative >= 0), other=0,
+    )
+    history = tl.load(
+        gathered_state_ptr
+        + (request.to(tl.int64) * CHANNELS + channel.to(tl.int64)) * STATE_LENGTH
+        + (offset + position).to(tl.int64),
+        valid & payload & (relative < 0), other=0,
+    )
+    value = tl.where(relative >= 0, query, history)
+    tl.store(
+        conv_state_ptr + slot * STATE_STRIDE + element.to(tl.int64),
+        value, valid & (element < CHANNELS * STATE_CAPACITY),
+    )
+
+
 @torch.library.custom_op(
     "b12x::ple_layer_pipeline",
     mutates_args=(

b12x/sequence/ple/_preparation.py

@@ -108,6 +108,15 @@ def compile_layer(query_payload, config_payload, ordinal):
                 p["normalized_u"], p["gathered_state"], p["query_start_loc"], p["state_slot_ids"],
                 p["request_is_prefill"], p["num_seqs"], p["conv_state"], **state_options,
             ),
+            kernels._export_checkpoint_kernel.warmup(
+                p["normalized_u"], p["gathered_state"], p["query_start_loc"],
+                _CompilePointer(torch.int32, 4), _CompilePointer(torch.int64, 8),
+                p["request_is_prefill"], p["num_seqs"], p["state_slot_ids"],
+                p["conv_state"], query.max_state_slots,
+                CHANNELS=channels, STATE_LENGTH=length, STATE_CAPACITY=capacity,
+                STATE_STRIDE=query.state_strides[0], MAX_SEQS=n, BLOCK=256,
+                num_warps=4, grid=(n, triton.cdiv(channels * capacity, 256)),
+            ),
         )
 
 
@@ -176,6 +185,25 @@ def run(self, binding, *, eps, token_count=None):
         self.run_tensors(*_binding_tensors(binding, token_count), eps=eps)
         return binding.out[:token_count]
 
+    def export_checkpoint(self, binding, offsets, slots):
+        if binding._state is not self.layout or not self.mixed:
+            raise ValueError("checkpoint export requires a binding from this mixed PLE plan")
+        for name, tensor, dtype in (
+            ("checkpoint offsets", offsets, torch.int32),
+            ("checkpoint slots", slots, torch.int64),
+        ):
+            if (tensor.shape != (self.query.max_seqs,) or tensor.dtype != dtype
+                    or tensor.device != self.layout.caps.device or not tensor.is_contiguous()):
+                raise ValueError(f"PLE {name} must match the planned request capacity, dtype and device")
+        with torch.cuda.device(self.layout.caps.device):
+            self.programs[5][(self.query.max_seqs, triton.cdiv(self.channels * self.state_capacity, 256), 1)](
+                binding.normalized_u, binding.gathered_state, binding.query_start_loc,
+                offsets, slots, binding.request_is_prefill, binding.num_seqs,
+                binding.state_slot_ids, binding.conv_state, self.query.max_state_slots,
+                self.channels, self.state_length, self.state_capacity,
+                self.query.state_strides[0], self.query.max_seqs, 256,
+            )
+
     def run_tensors(
         self, residual, key, value, k_norm_weight, q_norm_weight, u_norm_weight,
         conv_weight, query_start_loc, state_slot_ids, state_is_fresh, num_accepted_tokens,

b12x/sequence/ple/api.py

@@ -17,6 +17,21 @@
 from ._tuning import PleConfig, PleQuery
 
 
+def export_checkpoint(binding: Binding, *, offsets, slots) -> None:
+    """Save internal prefill windows after the matching mixed PLE invocation.
+
+    Offsets are relative to each request's query start. Negative slots and
+    zero offsets disable export. Call before reusing the binding's scratch.
+    Out-of-range destinations, destinations shared by enabled exports, and
+    destinations overlapping live input/output state slots are skipped on
+    the device, including during CUDA graph replay.
+    """
+    from b12x.preparation.types import require_prepared
+
+    state = require_prepared(binding.plan, "sequence.ple")
+    state.export_checkpoint(binding, offsets, slots)
+
+
 def is_supported(device=None) -> bool:
     """True on supported b12x devices with Triton available."""
     return default_is_supported(device, requires=("triton",))
@@ -34,5 +49,6 @@ def is_supported(device=None) -> bool:
     "run_decode",
     "run_mixed",
     "run_prefill",
+    "export_checkpoint",
     "is_supported",
 ]

tests/sequence/test_ple.py

@@ -1832,3 +1832,179 @@ def test_ple_state_slot_past_int32_element_offset_matches_oracle() -> None:
         del binding
         del conv_state
         torch.cuda.empty_cache()
+
+
+@pytest.mark.parametrize("high_slot", [False, True])
+@pytest.mark.parametrize("other_device", [False, True])
+@torch.inference_mode()
+def test_ple_internal_checkpoint_replays_offsets_and_preserves_other_slots(high_slot, other_device):
+    device = require_b12x()
+    original_device = torch.cuda.current_device()
+    if other_device:
+        if torch.cuda.device_count() < 2:
+            pytest.skip("Checkpoint export with another active device requires two CUDA GPUs")
+        device = torch.device("cuda", (original_device + 1) % torch.cuda.device_count())
+        with torch.cuda.device(device):
+            require_b12x()
+    tokens, streams, hidden = 17, 4, 2560
+    length, speculative, kernel_size, dilation = 9, 3, 4, 3
+    channels = streams * hidden
+    stride = channels * (length + speculative)
+    destination = (2**31 // stride + 2) if high_slot else 4
+    pool = torch.empty(destination + 3, channels, length + speculative,
+                       dtype=torch.bfloat16, device=device)
+    pool[:3].normal_()
+    pool[destination:].fill_(91)
+    untouched = pool[0].clone()
+    prior = pool[1].clone()
+    residual, key, value, weights, generator = _cuda_projected_inputs(
+        tokens, streams, hidden, device=device, seed=1626,
+    )
+    conv_weight = torch.randn(channels, kernel_size, dtype=torch.bfloat16,
+                             device=device, generator=generator)
+    _, binding = _bind_cuda_layer(
+        mode="mixed", residual=residual, key=key, value=value, weights=weights,
+        conv_weight=conv_weight,
+        query_start_loc=torch.tensor([0, 16, 17, 17], dtype=torch.int32, device=device),
+        state_slot_ids=torch.tensor([1, 2, -1], dtype=torch.int64, device=device),
+        state_is_fresh=torch.tensor([False, False, True], device=device),
+        num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=device),
+        num_seqs=3, num_tokens=tokens, conv_state=pool,
+        max_speculative_tokens=speculative, dilation=dilation,
+        request_is_prefill=torch.tensor([True, False, True], device=device),
+    )
+    ple.run_mixed(binding, eps=1e-6)
+    _, normalized = ple_projected_u_reference(
+        residual, key, value, k_norm_weight=weights[0], q_norm_weight=weights[1],
+        u_norm_weight=weights[2], eps=1e-6,
+    )
+    history = torch.cat((prior[:, :length], normalized[:16].reshape(16, channels).T), dim=1)
+    offsets = torch.tensor([2, 1, 1], dtype=torch.int32, device=device)
+    slots = torch.tensor([destination, destination + 1, destination + 2], dtype=torch.int64, device=device)
+    assert torch.cuda.current_device() == original_device
+    ple.export_checkpoint(binding, offsets=offsets, slots=slots)
+    assert torch.cuda.current_device() == original_device
+    graph = torch.cuda.CUDAGraph()
+    capture_stream = torch.cuda.Stream(device=device)
+    with torch.cuda.device(device), torch.cuda.graph(graph, stream=capture_stream):
+        ple.export_checkpoint(binding, offsets=offsets, slots=slots)
+    for offset in (2, 12, 0, 16):
+        offsets[0] = offset
+        pool[destination:].fill_(91)
+        before = torch.cuda.memory_allocated(device)
+        graph.replay()
+        torch.cuda.synchronize(device)
+        assert torch.cuda.memory_allocated(device) == before
+        if 0 < offset < 16:
+            torch.testing.assert_close(pool[destination, :, :length],
+                                       history[:, offset:offset + length], rtol=0, atol=0)
+            assert bool((pool[destination, :, length:] == 0).all())
+        else:
+            assert bool((pool[destination] == 91).all())
+        assert bool((pool[destination + 1:] == 91).all())
+        torch.testing.assert_close(pool[0], untouched, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("slot_padding", [0, 11])
+@torch.inference_mode()
+def test_ple_internal_checkpoint_rejects_unsafe_slots_on_replay(slot_padding, monkeypatch):
+    from triton.runtime.jit import JITFunction
+    from b12x.sequence.ple import _kernels
+
+    device = require_b12x()
+    tokens, streams, hidden = 15, 2, 32
+    length, speculative, kernel_size, dilation = 9, 3, 4, 3
+    channels, capacity, pool_slots = streams * hidden, length + speculative, 10
+    stride = channels * capacity + slot_padding
+    storage = torch.randn(pool_slots * stride, dtype=torch.bfloat16, device=device)
+    pool = storage.as_strided((pool_slots, channels, capacity), (stride, capacity, 1))
+    initial_storage = storage.clone()
+    initial_pool = pool.clone()
+    residual, key, value, weights, generator = _cuda_projected_inputs(
+        tokens, streams, hidden, device=device, seed=1627,
+    )
+    conv_weight = torch.randn(channels, kernel_size, dtype=torch.bfloat16,
+                             device=device, generator=generator)
+    starts = [0, 4, 8, 11, 15, 15]
+    live_slots = [0, 1, 2, 3, 4]
+    _, binding = _bind_cuda_layer(
+        mode="mixed", residual=residual, key=key, value=value, weights=weights,
+        conv_weight=conv_weight,
+        query_start_loc=torch.tensor(starts, dtype=torch.int32, device=device),
+        state_slot_ids=torch.tensor(live_slots, dtype=torch.int64, device=device),
+        state_is_fresh=torch.zeros(5, dtype=torch.bool, device=device),
+        num_accepted_tokens=torch.ones(5, dtype=torch.int32, device=device),
+        num_seqs=4, num_tokens=tokens, conv_state=pool,
+        max_speculative_tokens=speculative, dilation=dilation,
+        request_is_prefill=torch.tensor([True, True, False, True, True], device=device),
+    )
+    _, normalized = ple_projected_u_reference(
+        residual, key, value, k_norm_weight=weights[0], q_norm_weight=weights[1],
+        u_norm_weight=weights[2], eps=1e-6,
+    )
+    offsets = torch.tensor([1, 2, 1, 1, 1], dtype=torch.int32, device=device)
+    slots = torch.tensor([5, 6, 5, -1, 5], dtype=torch.int64, device=device)
+    ple.run_mixed(binding, eps=1e-6)
+
+    def reject_resolution(*args, **kwargs):
+        pytest.fail("Checkpoint export must reuse the prepared kernels")
+
+    for kernel in vars(_kernels).values():
+        if isinstance(kernel, JITFunction):
+            monkeypatch.setattr(kernel, "_do_compile", reject_resolution)
+
+    ple.export_checkpoint(binding, offsets=offsets, slots=slots)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        ple.export_checkpoint(binding, offsets=offsets, slots=slots)
+
+    # Each row specifies destinations, offsets, live count, live state slots,
+    # and the request rows whose exports must succeed.
+    cases = [
+        ([5, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (0, 1)),
+        ([10, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (1,)),
+        ([2**40 + 5, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (1,)),
+        ([-1, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (1,)),
+        ([0, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (1,)),
+        ([2, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (1,)),
+        ([3, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (1,)),
+        ([5, 5, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, ()),
+        ([5, 5, 5, 5, 5], [1, 2, 1, 1, 1], 4, live_slots, ()),
+        ([5, 5, 5, -1, 5], [1, 0, 1, 1, 1], 4, live_slots, (0,)),
+        ([5, 5, 5, -1, 5], [1, -1, 1, 1, 1], 4, live_slots, (0,)),
+        ([5, 5, 5, -1, 5], [1, 4, 1, 1, 1], 4, live_slots, (0,)),
+        ([5, 5, 5, -1, 5], [1, 5, 1, 1, 1], 4, live_slots, (0,)),
+        ([5, 5, 5, -1, 5], [1, 2, 1, 1, 1], 1, live_slots, (0,)),
+        ([5, 6, 5, -1, 5], [1, 2, 1, 1, 1], 0, live_slots, ()),
+        ([4, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, live_slots, (0, 1)),
+        ([4, 6, 5, -1, 5], [1, 2, 1, 1, 1], 5, live_slots, (1,)),
+        ([5, 6, 5, -1, 5], [1, 2, 1, 1, 1], 4, [0, 1, 2, 5, 4], (1,)),
+        ([5, 6, 5, 9, 5], [1, 2, 1, 3, 1], 5, live_slots, (0, 1, 3)),
+    ]
+    for destinations, boundaries, num_seqs, state_slots, exported in cases:
+        slots.copy_(torch.tensor(destinations, dtype=slots.dtype, device=device))
+        offsets.copy_(torch.tensor(boundaries, dtype=offsets.dtype, device=device))
+        binding.state_slot_ids.copy_(torch.tensor(state_slots, dtype=torch.int64, device=device))
+        binding.num_seqs.fill_(num_seqs)
+        binding.num_tokens.fill_(starts[num_seqs])
+        storage.copy_(initial_storage)
+        ple.run_mixed(binding, eps=1e-6)
+        before_export = storage.clone()
+        expected_storage = storage.clone()
+        expected = expected_storage.as_strided(pool.shape, pool.stride())
+        for request in exported:
+            history = torch.cat((
+                initial_pool[state_slots[request], :, :length],
+                normalized[starts[request]:starts[request + 1]].reshape(-1, channels).T,
+            ), dim=1)
+            destination, offset = destinations[request], boundaries[request]
+            expected[destination, :, :length] = history[:, offset:offset + length]
+            expected[destination, :, length:] = 0
+        ple.export_checkpoint(binding, offsets=offsets, slots=slots)
+        torch.testing.assert_close(storage, expected_storage, rtol=0, atol=0)
+        storage.copy_(before_export)
+        allocated = torch.cuda.memory_allocated(device)
+        graph.replay()
+        torch.cuda.synchronize(device)
+        assert torch.cuda.memory_allocated(device) == allocated
+        torch.testing.assert_close(storage, expected_storage, rtol=0, atol=0)