diff --git a/b12x/sequence/kda_prefill/__init__.py b/b12x/sequence/kda_prefill/__init__.py index ba7859b38..e102cf312 100644 --- a/b12x/sequence/kda_prefill/__init__.py +++ b/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 diff --git a/b12x/sequence/kda_prefill/_cute_kernels.py b/b12x/sequence/kda_prefill/_cute_kernels.py index 17be79c69..b830c47b0 100644 --- a/b12x/sequence/kda_prefill/_cute_kernels.py +++ b/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: diff --git a/b12x/sequence/kda_prefill/_impl.py b/b12x/sequence/kda_prefill/_impl.py index 4c2731bca..6f42b9960 100644 --- a/b12x/sequence/kda_prefill/_impl.py +++ b/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,)) diff --git a/b12x/sequence/kda_prefill/_policy.py b/b12x/sequence/kda_prefill/_policy.py index 96695cdba..551ef9775 100644 --- a/b12x/sequence/kda_prefill/_policy.py +++ b/b12x/sequence/kda_prefill/_policy.py @@ -6,7 +6,7 @@ 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) @@ -14,6 +14,13 @@ 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"}), diff --git a/b12x/sequence/kda_prefill/metadata.py b/b12x/sequence/kda_prefill/metadata.py new file mode 100644 index 000000000..7d73ed827 --- /dev/null +++ b/b12x/sequence/kda_prefill/metadata.py @@ -0,0 +1,87 @@ +"""GPU-free checkpoint metadata oracle for transactional KDA prefill.""" + +from __future__ import annotations + + +def validate_metadata( + *, + cu_seqlens, + initial_state_indices, + final_state_indices, + checkpoint_state_indices, + checkpoint_offsets, + num_seqs, + num_tokens, + token_capacity, + seq_capacity, + state_slots, + chunk=16, + null_state_index=None, + max_checkpoints=1, +): + """Validate packed ownership and return spans without reading state tensors. + + Nonpositive offsets disable an entry. A null checkpoint never owns storage; + its positive offset still must be aligned and within the request. Enabled, + non-null exports need distinct offsets and globally unique destinations. + Initial may alias only own final. + """ + if type(max_checkpoints) is not int or max_checkpoints not in (1, 2, 4): + raise ValueError("max_checkpoints must be 1, 2 or 4") + if not 0 <= num_seqs <= seq_capacity or not 0 <= num_tokens <= token_capacity: + raise ValueError("live counts exceed capacities") + if int(cu_seqlens[0]) != 0 or int(cu_seqlens[num_seqs]) != num_tokens: + raise ValueError("packed boundaries do not match live tokens") + spans = [] + writes = set() + + def null(slot): + return null_state_index is not None and slot == null_state_index + + def check_slot(slot): + if not null(slot) and not 0 <= slot < state_slots: + raise IndexError("state slot out of range") + + def write(slot): + if null(slot): + return + check_slot(slot) + if slot in writes: + raise ValueError("duplicate checkpoint/final write slot") + writes.add(slot) + + for seq in range(num_seqs): + start, end = int(cu_seqlens[seq]), int(cu_seqlens[seq + 1]) + if not 0 <= start <= end <= num_tokens: + raise ValueError("invalid packed sequence interval") + spans.append((start, end)) + check_slot(int(initial_state_indices[seq])) + write(int(final_state_indices[seq])) + seen_offsets = set() + for cp in range(max_checkpoints): + slot = int( + checkpoint_state_indices[seq] + if max_checkpoints == 1 + else checkpoint_state_indices[seq][cp] + ) + offset = int( + checkpoint_offsets[seq] + if max_checkpoints == 1 + else checkpoint_offsets[seq][cp] + ) + if offset > end - start or (offset > 0 and offset % chunk): + raise ValueError("checkpoint offset out of bounds or unaligned") + if offset > 0 and not null(slot): + if offset in seen_offsets: + raise ValueError("active checkpoints have the same offset") + seen_offsets.add(offset) + write(slot) + for seq in range(num_seqs): + initial, final = int(initial_state_indices[seq]), int(final_state_indices[seq]) + if not null(initial) and initial in writes - ( + {final} if not null(final) else set() + ): + raise ValueError( + "initial slot conflicts with a checkpoint or another request write" + ) + return spans diff --git a/b12x/sequence/kda_prefill/reference.py b/b12x/sequence/kda_prefill/reference.py index cb7b86f24..c7e17d5a1 100644 --- a/b12x/sequence/kda_prefill/reference.py +++ b/b12x/sequence/kda_prefill/reference.py @@ -51,17 +51,21 @@ def recurrent_kda( lower_bound: float, initial_state: torch.Tensor, checkpoint_offset: int = -1, + checkpoint_offsets: tuple[int, ...] | None = None, scale: float | None = None, eps: float = 1e-6, qk_l2norm: bool = True, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | dict[int, torch.Tensor] | None]: """Run the fp32 token recurrence for one sequence. ``q, k, v, raw_g`` are ``[T, heads, 128]``, ``raw_beta`` is ``[T, heads]``, ``initial_state`` is ``[heads, 128, 128]`` in ``[value_dim, key_dim]`` order. Returns the bf16 output ``[T, heads, 128]``, the fp32 final state in the same orientation, and the state after ``checkpoint_offset`` tokens - (``None`` unless ``0 <= checkpoint_offset <= T``). + (``None`` unless ``0 <= checkpoint_offset <= T``). With explicit plural + ``checkpoint_offsets``, the third result is an offset-to-state mapping + captured by that same recurrence. Without plural offsets, the third result + is a single checkpoint tensor or None. """ tokens = int(q.shape[0]) heads = int(q.shape[1]) @@ -76,18 +80,25 @@ def recurrent_kda( output = torch.empty( (tokens, heads, KDA_HEAD_DIM), dtype=torch.bfloat16, device=q.device ) + checkpoints = {} + requested = (checkpoint_offset,) if checkpoint_offsets is None else checkpoint_offsets checkpoint = None if checkpoint_offset == 0: checkpoint = state.transpose(-1, -2).contiguous() + if checkpoint_offsets is not None and 0 in requested: + checkpoints[0] = state.transpose(-1, -2).contiguous().clone() for t in range(tokens): state = state * torch.exp(log_decay[t])[:, :, None] k_t = kf[t] delta = vf[t] - torch.einsum("hk,hkv->hv", k_t, state) state = state + (beta[t][:, None] * k_t)[:, :, None] * delta[:, None, :] output[t] = torch.einsum("hk,hkv->hv", qf[t], state).to(torch.bfloat16) - if t + 1 == checkpoint_offset: - checkpoint = state.transpose(-1, -2).contiguous() - return output, state.transpose(-1, -2).contiguous(), checkpoint + if t + 1 in requested: + saved = state.transpose(-1, -2).contiguous().clone() + checkpoints[t + 1] = saved + if t + 1 == checkpoint_offset: + checkpoint = saved + return output, state.transpose(-1, -2).contiguous(), (checkpoint if checkpoint_offsets is None else checkpoints) def _validate_packed( @@ -104,8 +115,16 @@ def _validate_packed( state_slots: int, chunk: int, null_state_index: int | None, + max_checkpoints: int = 1, ) -> list[tuple[int, int]]: """Raise on every condition the device validator flags; return spans.""" + if max_checkpoints != 1: + from .metadata import validate_metadata + return validate_metadata(cu_seqlens=cu_seqlens, initial_state_indices=initial_state_indices, + final_state_indices=final_state_indices, checkpoint_state_indices=checkpoint_state_indices, + checkpoint_offsets=checkpoint_offsets, num_seqs=num_seqs, num_tokens=num_tokens, + token_capacity=token_capacity, seq_capacity=seq_capacity, state_slots=state_slots, + chunk=chunk, null_state_index=null_state_index, max_checkpoints=max_checkpoints) if num_seqs < 0 or num_seqs > seq_capacity: raise ValueError(f"num_seqs={num_seqs} exceeds capacity {seq_capacity}") if num_tokens < 0 or num_tokens > token_capacity: @@ -183,6 +202,7 @@ def prefill_kda( qk_l2norm: bool = True, null_state_index: int | None = None, chunk: int = 16, + max_checkpoints: int = 1, output: torch.Tensor | None = None, ) -> torch.Tensor: """Run the fp32 recurrence for a packed batch over a state pool. @@ -215,6 +235,7 @@ def prefill_kda( state_slots=int(recurrent_state.shape[0]), chunk=chunk, null_state_index=null_state_index, + max_checkpoints=max_checkpoints, ) if output is None: output = torch.zeros( @@ -227,8 +248,8 @@ def is_null(slot: int) -> bool: for request, (start, end) in enumerate(spans): initial = int(initial_state_indices[request]) final = int(final_state_indices[request]) - checkpoint_slot = int(checkpoint_state_indices[request]) - offset = int(checkpoint_offsets[request]) + slots = [int(checkpoint_state_indices[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_state_indices[request]] + offsets = [int(checkpoint_offsets[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_offsets[request]] if is_null(initial): state = torch.zeros( (heads, KDA_HEAD_DIM, KDA_HEAD_DIM), dtype=torch.float32, device=q.device @@ -245,14 +266,15 @@ def is_null(slot: int) -> bool: dt_bias, lower_bound=lower_bound_value, initial_state=state, - checkpoint_offset=offset if offset > 0 else -1, + checkpoint_offsets=tuple(offset for offset in offsets if offset > 0), scale=scale, eps=eps, qk_l2norm=qk_l2norm, ) output[start:end] = out - if checkpoint is not None and not is_null(checkpoint_slot): - recurrent_state[checkpoint_slot].copy_(checkpoint.to(recurrent_state.dtype)) + for checkpoint_slot, offset in zip(slots, offsets, strict=True): + if offset in checkpoint and not is_null(checkpoint_slot): + recurrent_state[checkpoint_slot].copy_(checkpoint[offset].to(recurrent_state.dtype)) if not is_null(final): recurrent_state[final].copy_(final_state.to(recurrent_state.dtype)) return output @@ -277,7 +299,7 @@ class MirrorTrace: k1: dict[tuple[int, int], dict[str, torch.Tensor]] = field(default_factory=dict) k2: dict[tuple[int, int], dict[str, torch.Tensor]] = field(default_factory=dict) - checkpoints: dict[int, torch.Tensor] = field(default_factory=dict) + checkpoints: dict[int | tuple[int, int], torch.Tensor] = field(default_factory=dict) def _neumann_inverse(lower: torch.Tensor, chunk: int) -> torch.Tensor: @@ -432,6 +454,7 @@ def prefill_kda_chunk_mirror( qk_l2norm: bool = True, null_state_index: int | None = None, chunk: int = 16, + max_checkpoints: int = 1, policy: MirrorPolicy | None = None, trace: bool = False, output: torch.Tensor | None = None, @@ -466,6 +489,7 @@ def prefill_kda_chunk_mirror( state_slots=int(recurrent_state.shape[0]), chunk=chunk, null_state_index=null_state_index, + max_checkpoints=max_checkpoints, ) if output is None: output = torch.zeros( @@ -487,8 +511,8 @@ def padded(x: torch.Tensor, start: int, rows: int) -> torch.Tensor: for request, (start, end) in enumerate(spans): initial = int(initial_state_indices[request]) final = int(final_state_indices[request]) - checkpoint_slot = int(checkpoint_state_indices[request]) - offset = int(checkpoint_offsets[request]) + slots = [int(checkpoint_state_indices[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_state_indices[request]] + offsets = [int(checkpoint_offsets[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_offsets[request]] if is_null(initial): state = torch.zeros( (heads, KDA_HEAD_DIM, KDA_HEAD_DIM), @@ -528,11 +552,13 @@ def padded(x: torch.Tensor, start: int, rows: int) -> torch.Tensor: if trace: record.k1[(request, local)] = prep record.k2[(request, local)] = step - if offset > 0 and (local + 1) * chunk == offset: - if trace: - record.checkpoints[request] = state.clone() - if not is_null(checkpoint_slot): - recurrent_state[checkpoint_slot].copy_(state.to(recurrent_state.dtype)) + for checkpoint_index, (checkpoint_slot, offset) in enumerate(zip(slots, offsets, strict=True)): + if offset > 0 and (local + 1) * chunk == offset: + if trace: + key = request if max_checkpoints == 1 else (request, checkpoint_index) + record.checkpoints[key] = state.clone() + if not is_null(checkpoint_slot): + recurrent_state[checkpoint_slot].copy_(state.to(recurrent_state.dtype)) if not is_null(final): recurrent_state[final].copy_(state.to(recurrent_state.dtype)) finally: diff --git a/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/README.md b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/README.md new file mode 100644 index 000000000..15ef5cfcb --- /dev/null +++ b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/README.md @@ -0,0 +1,15 @@ +# GB10 four-checkpoint KDA evidence + +[The serving report](report.md) and [sanitized measurements](evidence.json) contain the four-configuration GLM-5.3-Flash experiment: 36 cold prefill samples, 16 decode cells and 20 exact-answer/cache smoke checks. They compare neither feature, continuation coalescing, token-sharded mHC and both features. The report discloses sequential arm order, both recovery reboots, diagnostic overhead and the limits of these observations. + +The B12X component exports up to four recurrent states during one traversal. The serving integration owns checkpoint planning, physical-page retention, convolution history and pass coalescing. Serving throughput changes therefore do not establish a standalone B12X kernel speedup. mHC is a separate optimization. + +The tested component revision is `70fe41974ef4b18f61caaa2579c81cdc05d1265f`, based on `06b4de7c723e6f166d65abf5909c5b7d0f8acc68`. The tested serving composition is vLLM `abb715f132bdccb592a34b2596a3d3a8d757ffbc`. Runtime source IDs and the immutable image digest are included in the JSON. Documentation packaging does not change the tested kernel source. + +[Component-check evidence](component-checks.json) records **12 GPU tests passed** and **41 CPU tests passed**, their source-log digests and coverage. The committed [GPU suite](../../../../tests/sequence/test_kda_prefill_two_checkpoints_gpu.py) checks independent FP32 oracles, exact H16 8K checkpoint positions, graph replay, invalid metadata and high pool addresses. The [CPU suite](../../../../tests/sequence/test_kda_prefill_two_checkpoints_cpu.py) covers capacity, policy, metadata and reference contracts. The GPU selection enables `B12X_RUN_LARGE_POOL_TESTS=1`; no selected cases were skipped. + +Eligibility remains NVIDIA GB10. These execution/correctness checks are not a measured component-policy profile. KDA catalog/offline-provider integration and embedded measured profiles remain unresolved. The serving smoke tests do not establish full numerical or model-quality equivalence. + +Both native vLLM split-page settings were held at 512 in every serving arm: `VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE` and `VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE`. They preserve physical/lookup/scheduler alignment `(512, 512, 2048)` and do not require SparkCache. Enabled feature flags were accompanied by request-associated checkpoint/mHC dispatch evidence and completed API requests. + +The companion [serving reproduction instructions](https://github.com/FujitsuPolycom/vllm/blob/feat/gb10-continuation-prefill/docs/benchmarking/glm-kda-checkpoints-20260907/reproduction.md) identify the model snapshot, client inputs and public runtime composition. diff --git a/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/component-checks.json b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/component-checks.json new file mode 100644 index 000000000..646082523 --- /dev/null +++ b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/component-checks.json @@ -0,0 +1,31 @@ +{ + "schema": "b12x.kda-prefill.bounded-correctness.v1", + "base_revision": "06b4de7c723e6f166d65abf5909c5b7d0f8acc68", + "component_revision": "70fe41974ef4b18f61caaa2579c81cdc05d1265f", + "serving_composition_revision": "abb715f132bdccb592a34b2596a3d3a8d757ffbc", + "image_digest": "sha256:52b207e716a285c16e5e1b14ec2a41f6208b9450c617e7d1cde5a507ea879d7f", + "gpu": { + "passed": 12, + "failed": 0, + "skipped_selected": 0, + "selection": "tests/sequence/test_kda_prefill_two_checkpoints_gpu.py -k four_checkpoint", + "large_pool_enabled": true, + "coverage": [ + "H1/H16 independent FP32 recurrence oracle", + "separate and in-place final state", + "H16 8192 tokens with states at 4096,6144,7168,7680", + "frozen graph replay with changing live counts and offsets", + "invalid fourth-slot aliases, nonadjacent offsets and bounds", + "H16 high pool addresses beyond signed 32-bit element offsets" + ], + "source_log_sha256": "766e4f526f87e2b83e05e45b300870c045591c800d5a34bc1a2bcfba3338fbed" + }, + "cpu": { + "passed": 41, + "failed": 0, + "selection": "tests/sequence/test_kda_prefill_two_checkpoints_cpu.py", + "source_log_sha256": "1aca8fea89d6d21712311751d0f8ebfa846493328f7bb392791b2377347cf008" + }, + "scope": "bounded component correctness; no standalone component throughput or measured policy-profile claim", + "policy_dependency": "KDA catalog/offline-provider and measured profile integration remains unresolved; no embedded KDA profile is claimed." +} diff --git a/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/evidence.json b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/evidence.json new file mode 100644 index 000000000..aa390048c --- /dev/null +++ b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/evidence.json @@ -0,0 +1,1534 @@ +{ + "schema": "glm-prefill-four-arm-public-evidence/v1", + "exporter_sha256": "290fede850bc4f296bd6b689e4dfcc3ef71a641bee23184a1b64ef61b16bf285", + "status": "complete", + "qualification": "bounded semantic/cache smoke checks and performance observations; not full numerical/model-quality qualification", + "hardware": "four NVIDIA GB10 Sparks", + "model_family": "GLM-5.3-Flash-NVFP4-Spark", + "runtime_sources": { + "vllm": "abb715f132bdccb592a34b2596a3d3a8d757ffbc", + "b12x": "70fe41974ef4b18f61caaa2579c81cdc05d1265f" + }, + "image_digest": "sha256:52b207e716a285c16e5e1b14ec2a41f6208b9450c617e7d1cde5a507ea879d7f", + "settings": { + "tp": 4, + "dcp": 4, + "pp": 1, + "max_batched_tokens": 8192, + "max_model_len": 1048576, + "block_size": 512, + "mamba_block_size": 512, + "mtp_tokens": 3, + "kv_bytes_per_rank": 25769803776, + "prefix_cache_retention_interval": 0, + "recurrent_checkpoint_policy": "aligned", + "speculative_method": "mtp" + }, + "native_split_settings": { + "VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE": "512", + "VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE": "512" + }, + "diagnostics_enabled": true, + "operational_comparison": { + "recovery_events": [ + { + "event": "one rank rebooted between experiment arms", + "after_arm": "neither", + "before_arm": "continuation", + "reason": "recover memory allocation capacity after a failed startup preflight", + "recorded_settings_matched_before_after": true, + "recorded_settings": { + "driver": "580.173.02", + "maximum_sm_clock_mhz": 3003, + "application_graphics_clock_mhz": 2418, + "persistence": "enabled", + "cpu0_governor": "performance" + }, + "receipt_sha256": "64254bb4ee8a54428b5695ef4f4096754d8921a685dc9bddf3db5294e9a2fc45" + }, + { + "event": "one rank rebooted between experiment arms", + "after_arm": "continuation", + "before_arm": "mhc", + "reason": "recover memory allocation capacity after a failed startup preflight", + "recorded_settings_matched_before_after": true, + "recorded_settings": { + "driver": "580.173.02", + "maximum_sm_clock_mhz": 3003, + "application_graphics_clock_mhz": 2418, + "persistence": "enabled", + "cpu0_governor": "performance" + }, + "receipt_sha256": "093455e00436e5b0f87b6de3d1972dffc4e34e18d9631b416a28d48bca1a620e" + } + ], + "continuous_operating_condition_equivalence_verified": false + }, + "measurement_order": [ + "both", + "neither", + "continuation", + "mhc" + ], + "arms": [ + { + "arm": "neither", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "0", + "VLLM_GLM53_MHC_PREFILL_SHARD": "0" + }, + "started_at_utc": "2026-09-07T13:01:15.745976+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 4.196101099980297, + "tokens_per_second": 1952.2885185103064 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 4.203437799995299, + "tokens_per_second": 1948.8809849902289 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 4.201881200016942, + "tokens_per_second": 1949.6029540213965 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 6.943599899997935, + "tokens_per_second": 2359.582959266543 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 6.9457052999932785, + "tokens_per_second": 2358.8677164313112 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 6.966370399983134, + "tokens_per_second": 2351.8703513151795 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 12.541768900002353, + "tokens_per_second": 2612.709599520196 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 12.532975899986923, + "tokens_per_second": 2614.5426482495823 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 12.522141299996292, + "tokens_per_second": 2616.8048431149473 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 4.201881200016942, + "min_ttft_seconds": 4.196101099980297, + "max_ttft_seconds": 4.203437799995299, + "tokens_per_second": 1949.6029540213965, + "throughput_change_vs_neither_pct": 0.0, + "ttft_reduction_vs_neither_pct": 0.0 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 6.9457052999932785, + "min_ttft_seconds": 6.943599899997935, + "max_ttft_seconds": 6.966370399983134, + "tokens_per_second": 2358.8677164313112, + "throughput_change_vs_neither_pct": 0.0, + "ttft_reduction_vs_neither_pct": 0.0 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 12.532975899986923, + "min_ttft_seconds": 12.522141299996292, + "max_ttft_seconds": 12.541768900002353, + "tokens_per_second": 2614.5426482495823, + "throughput_change_vs_neither_pct": 0.0, + "ttft_reduction_vs_neither_pct": 0.0 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": false, + "checkpoint_scheduled": false, + "checkpoint_dispatch_shape": null, + "checkpoint_capacity": 1, + "configured_checkpoint_grid": null, + "mhc_enabled": false, + "mhc_collectives_per_dispatch": null, + "mhc_request_dispatches_by_rank": [ + 0, + 0, + 0, + 0 + ], + "mhc_common_dispatch_count": 0, + "mhc_rows": null, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 49.16553901666799, + "mtp_accept_length": 2.752808988764045, + "mtp_normalized_steps_per_second": 17.86013458156511, + "reported_server_steps_per_second": 17.892046308818017, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.953, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 981, + "server_output_tokens": 981, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 20.578999999997905, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.3900000000139698, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 151, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 122.18713977842667, + "mtp_accept_length": 2.7958715596330275, + "mtp_normalized_steps_per_second": 43.702701348149326, + "reported_server_steps_per_second": 43.702701348149326, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.953, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 2438, + "server_output_tokens": 2438, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3589999999967404, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 711, + "completed": false + }, + { + "ttft": 3.2030000000086147, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 707, + "completed": false + }, + { + "ttft": 3.2030000000086147, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 725, + "completed": false + }, + { + "ttft": 3.2030000000086147, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 758, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 46.18463847889143, + "mtp_accept_length": 2.591036414565826, + "mtp_normalized_steps_per_second": 17.82477398590729, + "reported_server_steps_per_second": 17.86339754817361, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.985, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 923, + "server_output_tokens": 923, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 21.34299999999348, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.4070000000065193, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 163, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 123.43883232178055, + "mtp_accept_length": 2.7839366515837103, + "mtp_normalized_steps_per_second": 44.339669960363274, + "reported_server_steps_per_second": 44.339669960363274, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.937, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2461, + "server_output_tokens": 2461, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4530000000086147, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 764, + "completed": false + }, + { + "ttft": 3.312000000005355, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 786, + "completed": false + }, + { + "ttft": 3.312000000005355, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 709, + "completed": false + }, + { + "ttft": 3.312000000005355, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 734, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + } + ], + "artifact_sha256": { + "conditions": "106930faaea7c14960599399e280bd5a2930124ebb9ed81a181ece2373bf5746", + "prefill": "c01c9662b422ccd7926589c69136b41a26871ea4f8401a32cb7da23ad6f8d65d", + "activation": "af7926a719fcc0b52535bc850d639daf681bf2ad843eb9ccff582377e7100372", + "complete": "7fe568efe23bc6c96dec974c0f01678b5d6f11daa2e4b8b29b5621e87ee5eb3d", + "decode": "1ff4deea8f1d2d9d4bbf18fd1b28a40a48f340b6f1fcf3ccf55d4fe3933872ea", + "decode_receipt": "bd636073241c3003bd1fcc955c7bba8b79b758c4b131686c3b4edf1c67e3f0b7" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + }, + { + "arm": "continuation", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "1", + "VLLM_GLM53_MHC_PREFILL_SHARD": "0" + }, + "started_at_utc": "2026-09-07T13:24:22.409490+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 2.7674208999960683, + "tokens_per_second": 2960.1568738646292 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 2.7794476999843027, + "tokens_per_second": 2947.3481368425337 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 2.7874553999863565, + "tokens_per_second": 2938.8811028295186 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 5.595004900009371, + "tokens_per_second": 2928.3263004778705 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 5.563162100006593, + "tokens_per_second": 2945.087650776989 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 5.603458099998534, + "tokens_per_second": 2923.908719867877 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 11.21369129998493, + "tokens_per_second": 2922.1421495742475 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 11.208513199992012, + "tokens_per_second": 2923.4921184750315 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 11.190499399992405, + "tokens_per_second": 2928.1981821135028 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 2.7794476999843027, + "min_ttft_seconds": 2.7674208999960683, + "max_ttft_seconds": 2.7874553999863565, + "tokens_per_second": 2947.3481368425337, + "throughput_change_vs_neither_pct": 51.176839918256874, + "ttft_reduction_vs_neither_pct": 33.85230167923129 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 5.595004900009371, + "min_ttft_seconds": 5.563162100006593, + "max_ttft_seconds": 5.603458099998534, + "tokens_per_second": 2928.3263004778705, + "throughput_change_vs_neither_pct": 24.141183504265463, + "ttft_reduction_vs_neither_pct": 19.446554981035757 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 11.208513199992012, + "min_ttft_seconds": 11.190499399992405, + "max_ttft_seconds": 11.21369129998493, + "tokens_per_second": 2923.4921184750315, + "throughput_change_vs_neither_pct": 11.81657795608302, + "ttft_reduction_vs_neither_pct": 10.567822922218273 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": true, + "checkpoint_scheduled": true, + "checkpoint_dispatch_shape": { + "tokens": 8192, + "interior_checkpoints": 4 + }, + "checkpoint_capacity": 4, + "configured_checkpoint_grid": [ + 512, + 512, + 2048 + ], + "mhc_enabled": false, + "mhc_collectives_per_dispatch": null, + "mhc_request_dispatches_by_rank": [ + 0, + 0, + 0, + 0 + ], + "mhc_common_dispatch_count": 0, + "mhc_rows": null, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 47.3734288147036, + "mtp_accept_length": 2.6452513966480447, + "mtp_normalized_steps_per_second": 17.908856933119207, + "reported_server_steps_per_second": 17.92778807152631, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.969, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 946, + "server_output_tokens": 946, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3440000000118744, + "latency": 20.875, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.360000000015134, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 144, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -3.6450535025291497, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.27279946481695294, + "reported_server_steps_per_second_change_vs_neither_pct": 0.19976341493526117 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 120.17856247175385, + "mtp_accept_length": 2.698198198198198, + "mtp_normalized_steps_per_second": 44.540301951134154, + "reported_server_steps_per_second": 44.54030195113415, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.937, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 2396, + "server_output_tokens": 2396, + "num_errors": 0, + "request_samples": [ + { + "ttft": 3.172000000020489, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 747, + "completed": false + }, + { + "ttft": 1.3590000000258442, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 732, + "completed": false + }, + { + "ttft": 3.172000000020489, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 685, + "completed": false + }, + { + "ttft": 3.172000000020489, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 729, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -1.6438532813806472, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 1.9165877100187467, + "reported_server_steps_per_second_change_vs_neither_pct": 1.9165877100187245 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 45.92117782560242, + "mtp_accept_length": 2.5903954802259888, + "mtp_normalized_steps_per_second": 17.72747759025437, + "reported_server_steps_per_second": 17.72747759025437, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.969, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 917, + "server_output_tokens": 917, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 22.312999999994645, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.4219999999913853, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 107, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -0.570450829466651, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -0.5458492530106884, + "reported_server_steps_per_second_change_vs_neither_pct": -0.7608852546257983 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 125.2001601281229, + "mtp_accept_length": 2.7434210526315788, + "mtp_normalized_steps_per_second": 45.63650920737334, + "reported_server_steps_per_second": 45.63650920737334, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.984, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 2502, + "server_output_tokens": 2502, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 773, + "completed": false + }, + { + "ttft": 3.25, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 762, + "completed": false + }, + { + "ttft": 3.25, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 781, + "completed": false + }, + { + "ttft": 3.25, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 728, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 1.4268830749718298, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 2.924783265570885, + "reported_server_steps_per_second_change_vs_neither_pct": 2.924783265570885 + } + ], + "artifact_sha256": { + "conditions": "f412918e008077386ab15b64593acb2ece6d1e236388b1eeac1a8ad7f9f64031", + "prefill": "802eae6d9bca97cc529933d186a50c3ae5d2d82e711888f8889986a2789d86ce", + "activation": "311fd5b30308386959bd719f5684b1e667af8ab454dafa9db02562dded25a310", + "complete": "0a01ca5a333609d6c3aa99a124b695ba203bba47c42d813e1a2cf5e2e14a1ac0", + "decode": "fccb47fc2141ec0cf0e1e1ba412706b2db96719d48301b38a9fb0c3a716a12bf", + "decode_receipt": "218a37c92ed7f70a40acf571a0fd39688602f50a6a5702c0201cce89fdd9b872" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + }, + { + "arm": "mhc", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "0", + "VLLM_GLM53_MHC_PREFILL_SHARD": "1" + }, + "started_at_utc": "2026-09-07T13:51:05.058507+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 4.192637299973285, + "tokens_per_second": 1953.9014262102278 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 4.173298399982741, + "tokens_per_second": 1962.9557282637347 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 4.182778299989877, + "tokens_per_second": 1958.5068613413782 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 6.833160300011514, + "tokens_per_second": 2397.7192515112506 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 6.852885000000242, + "tokens_per_second": 2390.8178818117362 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 6.870854299981147, + "tokens_per_second": 2384.5651915577596 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 12.249805500003276, + "tokens_per_second": 2674.981247660727 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 12.213526299980003, + "tokens_per_second": 2682.9270429502126 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 12.218640999984927, + "tokens_per_second": 2681.8039747661314 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 4.182778299989877, + "min_ttft_seconds": 4.173298399982741, + "max_ttft_seconds": 4.192637299973285, + "tokens_per_second": 1958.5068613413782, + "throughput_change_vs_neither_pct": 0.4567036227358878, + "ttft_reduction_vs_neither_pct": 0.4546273232805209 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 6.852885000000242, + "min_ttft_seconds": 6.833160300011514, + "max_ttft_seconds": 6.870854299981147, + "tokens_per_second": 2390.8178818117362, + "throughput_change_vs_neither_pct": 1.3544704163725552, + "ttft_reduction_vs_neither_pct": 1.3363696843447403 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 12.218640999984927, + "min_ttft_seconds": 12.213526299980003, + "max_ttft_seconds": 12.249805500003276, + "tokens_per_second": 2681.8039747661314, + "throughput_change_vs_neither_pct": 2.5725847907503097, + "ttft_reduction_vs_neither_pct": 2.5080627499038166 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": false, + "checkpoint_scheduled": false, + "checkpoint_dispatch_shape": null, + "checkpoint_capacity": 1, + "configured_checkpoint_grid": null, + "mhc_enabled": true, + "mhc_collectives_per_dispatch": { + "reduce_scatter": 90, + "all_gather": 90, + "auxiliary_gathers": 0 + }, + "mhc_request_dispatches_by_rank": [ + 20, + 20, + 20, + 20 + ], + "mhc_common_dispatch_count": 20, + "mhc_rows": { + "first_pre": { + "8192": 1 + }, + "attention_post_pre": { + "2048": 44 + }, + "ffn_post_pre": { + "2048": 45 + }, + "final_post": { + "2048": 1 + } + }, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 48.463652260841044, + "mtp_accept_length": 2.7142857142857144, + "mtp_normalized_steps_per_second": 17.85502978030986, + "reported_server_steps_per_second": 17.886585061217623, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 20.015, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 970, + "server_output_tokens": 970, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 20.60999999998603, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.3589999999967404, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 161, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -1.4275990253844184, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -0.028582098482721197, + "reported_server_steps_per_second_change_vs_neither_pct": -0.030523325874143303 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 121.17257303488826, + "mtp_accept_length": 2.6933333333333334, + "mtp_normalized_steps_per_second": 44.98981672087435, + "reported_server_steps_per_second": 44.98981672087435, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.922, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2414, + "server_output_tokens": 2424, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 777, + "completed": false + }, + { + "ttft": 3.187999999994645, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 765, + "completed": false + }, + { + "ttft": 3.187999999994645, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 728, + "completed": false + }, + { + "ttft": 3.187999999994645, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 730, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -0.83033840171578, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 2.9451620449533866, + "reported_server_steps_per_second_change_vs_neither_pct": 2.9451620449533866 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 47.31118127597817, + "mtp_accept_length": 2.684659090909091, + "mtp_normalized_steps_per_second": 17.622789216025733, + "reported_server_steps_per_second": 17.641457424941013, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.953, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 944, + "server_output_tokens": 944, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 21.280999999988126, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.422000000020489, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 135, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 2.439215362921221, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -1.1331687573780669, + "reported_server_steps_per_second_change_vs_neither_pct": -1.2424295133894425 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 123.89911929545656, + "mtp_accept_length": 2.7030567685589517, + "mtp_normalized_steps_per_second": 45.836669335475854, + "reported_server_steps_per_second": 45.83666933547585, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.984, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2476, + "server_output_tokens": 2476, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 761, + "completed": false + }, + { + "ttft": 3.2339999999967404, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 734, + "completed": false + }, + { + "ttft": 3.2339999999967404, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 763, + "completed": false + }, + { + "ttft": 3.2339999999967404, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 775, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.37288668810162573, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 3.3762077535777646, + "reported_server_steps_per_second_change_vs_neither_pct": 3.3762077535777646 + } + ], + "artifact_sha256": { + "conditions": "c16558f03cd14b7978ae4a1ef3c30a743777a268c7957c4e31f0b19728db98d8", + "prefill": "dece9ad530835a6cb49a128995379156b688864bc343d3c012c71933506fb92e", + "activation": "a7776dc5d9ea45771af66b50d6bd7dba330f2b5633ad868fdca75c1eee2b1236", + "complete": "097d63814b64631fe9e7ddb226b85671a7b5cef4503613e4dde1251931b1b895", + "decode": "804885007b0ccef0ef113543b02800459f3354516be6b19d57dbf507b3c9148f", + "decode_receipt": "23938dd51cc74b63ffa659c7deb0ebc968b856810e068f3b6f9f4e212db50cf9" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + }, + { + "arm": "both", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "1", + "VLLM_GLM53_MHC_PREFILL_SHARD": "1" + }, + "started_at_utc": "2026-09-07T12:41:23.139314+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 2.6545916999748442, + "tokens_per_second": 3085.973635824157 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 2.6740572999988217, + "tokens_per_second": 3063.509521655953 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 2.6768195999902673, + "tokens_per_second": 3060.348183355272 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 5.342545599996811, + "tokens_per_second": 3066.7028841101105 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 5.352795799990417, + "tokens_per_second": 3060.83037952416 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 5.355028399993898, + "tokens_per_second": 3059.554268660586 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 10.720633199991425, + "tokens_per_second": 3056.5358770064263 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 10.74581809999654, + "tokens_per_second": 3049.3722948847003 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 10.745053999999072, + "tokens_per_second": 3049.5891411995535 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 2.6740572999988217, + "min_ttft_seconds": 2.6545916999748442, + "max_ttft_seconds": 2.6768195999902673, + "tokens_per_second": 3063.509521655953, + "throughput_change_vs_neither_pct": 57.13504718162894, + "ttft_reduction_vs_neither_pct": 36.360473494870824 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 5.352795799990417, + "min_ttft_seconds": 5.342545599996811, + "max_ttft_seconds": 5.355028399993898, + "tokens_per_second": 3060.83037952416, + "throughput_change_vs_neither_pct": 29.75845818750851, + "ttft_reduction_vs_neither_pct": 22.933732878134094 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 10.745053999999072, + "min_ttft_seconds": 10.720633199991425, + "max_ttft_seconds": 10.74581809999654, + "tokens_per_second": 3049.5891411995535, + "throughput_change_vs_neither_pct": 16.639487339831007, + "ttft_reduction_vs_neither_pct": 14.265741147636902 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": true, + "checkpoint_scheduled": true, + "checkpoint_dispatch_shape": { + "tokens": 8192, + "interior_checkpoints": 4 + }, + "checkpoint_capacity": 4, + "configured_checkpoint_grid": [ + 512, + 512, + 2048 + ], + "mhc_enabled": true, + "mhc_collectives_per_dispatch": { + "reduce_scatter": 90, + "all_gather": 90, + "auxiliary_gathers": 0 + }, + "mhc_request_dispatches_by_rank": [ + 35, + 35, + 35, + 35 + ], + "mhc_common_dispatch_count": 35, + "mhc_rows": { + "first_pre": { + "8192": 1 + }, + "attention_post_pre": { + "2048": 44 + }, + "ffn_post_pre": { + "2048": 45 + }, + "final_post": { + "2048": 1 + } + }, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 46.28702962370651, + "mtp_accept_length": 2.5865921787709496, + "mtp_normalized_steps_per_second": 17.894985534867097, + "reported_server_steps_per_second": 17.894985534867097, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.984, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 925, + "server_output_tokens": 926, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3590000000258442, + "latency": 21.90600000001723, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.4059999999881256, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 101, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -5.854729655227853, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.19513264663726382, + "reported_server_steps_per_second_change_vs_neither_pct": 0.01642755668271878 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 119.36552554952085, + "mtp_accept_length": 2.7767441860465114, + "mtp_normalized_steps_per_second": 42.98758457813565, + "reported_server_steps_per_second": 42.98758457813565, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.922, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 2378, + "server_output_tokens": 2388, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 722, + "completed": false + }, + { + "ttft": 3.1719999999913853, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 751, + "completed": false + }, + { + "ttft": 3.1719999999913853, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 646, + "completed": false + }, + { + "ttft": 3.1719999999913853, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 697, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -2.309256304732654, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -1.636321664230389, + "reported_server_steps_per_second_change_vs_neither_pct": -1.636321664230389 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 47.10834920317223, + "mtp_accept_length": 2.6818181818181817, + "mtp_normalized_steps_per_second": 17.565825126606597, + "reported_server_steps_per_second": 17.565825126606597, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.954, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 940, + "server_output_tokens": 944, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 21.655999999988126, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.422000000020489, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 127, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 2.0000388759197074, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -1.4527469436943385, + "reported_server_steps_per_second_change_vs_neither_pct": -1.665822085437696 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 124.25, + "mtp_accept_length": 2.761111111111111, + "mtp_normalized_steps_per_second": 45.0, + "reported_server_steps_per_second": 45.0, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 20.0, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 2485, + "server_output_tokens": 2485, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4219999999913853, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 709, + "completed": false + }, + { + "ttft": 3.265999999974156, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 732, + "completed": false + }, + { + "ttft": 3.265999999974156, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 733, + "completed": false + }, + { + "ttft": 3.265999999974156, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 761, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.6571414059595826, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 1.4892533936924135, + "reported_server_steps_per_second_change_vs_neither_pct": 1.4892533936924135 + } + ], + "artifact_sha256": { + "conditions": "984fe189c370b3e4d49904d8dac1fffbee22694c627bc5e6e79b8112bae58c37", + "prefill": "4a79c47b00dd8d20768a728b75b2f4214275829be171d75ab996a8e388313b3e", + "activation": "f65297db47d7da94ff2ce93bcc0667dc71b317f55535f12dc9e2ba2a85644f67", + "complete": "ce055a3f82345da35462862d5157fa42d39d01cd019a6445f8f25048d4ff972c", + "decode": "62ccbd2cfbd7e3b8f15115fbff8c1a38f508d8045edcb474ff53d25d606127d5", + "decode_receipt": "31c66c9129073ba86c4f16769834bb34db0415626f353f55ad01c6b9392fa4b4" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + } + ], + "limitations": [ + "Three cold samples per prefill size; one 20-second decode observation per cell.", + "Sequential arm order and stochastic decode acceptance limit small-difference conclusions.", + "Two recovery reboots occurred between arms: one between Neither and Coalescing, and another between Coalescing and Sharded mHC. Each receipt records matching driver, GPU clock settings, persistence and CPU0 governor before and after its reboot. Reboots can change allocator/cache state and thermal conditions; this was not an interleaved A/B experiment.", + "Prefill throughput is prompt tokens divided by first-token latency, including request overhead.", + "MTP-normalized steps/s is aggregate tok/s divided by measured acceptance length; server steps/s is retained separately.", + "Dispatch logs plus completed requests corroborate asynchronous execution without per-kernel GPU completion events.", + "Exact-answer and cache-reuse tests are smoke coverage, not full model-quality or numerical-equivalence evaluation.", + "Diagnostic logging remains enabled in all arms and adds host work.", + "Deployment readiness/source checks and exclusive access are recorded attestations; public output excludes private deployment identifiers." + ] +} diff --git a/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/manifest.json b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/manifest.json new file mode 100644 index 000000000..c2ac99e2c --- /dev/null +++ b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/manifest.json @@ -0,0 +1,11 @@ +{ + "schema": "b12x.kda-prefill.public-evidence-package.v1", + "component_revision": "70fe41974ef4b18f61caaa2579c81cdc05d1265f", + "serving_revision": "abb715f132bdccb592a34b2596a3d3a8d757ffbc", + "files": { + "component-checks.json": "b632ab944de952219f4bc56b01d4940e27e209f25bfd2dff1aa7ffbe4a43fb91", + "evidence.json": "20f4bab99a7c8bcd1ac564d8f62aca3c4628da014458be48829ba7ea65d7091b", + "README.md": "115e3d3a2edbec3b2ab197b823b743771476f36cf23c747d204b518314534678", + "report.md": "063a667f62b9b5da1de115ba2f55eeb1ac4f92ea08f52c3e2a8f09ba636ec3e8" + } +} diff --git a/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/report.md b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/report.md new file mode 100644 index 000000000..da4dcc3af --- /dev/null +++ b/docs/evidence/kda-prefill/gb10-four-checkpoint-20260907/report.md @@ -0,0 +1,51 @@ +# GLM prefill: four configurations + +bounded semantic/cache smoke checks and performance observations; not full numerical/model-quality qualification. + +Four GB10 Sparks, TP4/DCP4, MTP3, fixed 8192-token budget. Runtime vLLM `abb715f132bdccb592a34b2596a3d3a8d757ffbc`, B12X `70fe41974ef4b18f61caaa2579c81cdc05d1265f`. + +Both native split-page settings are 512 in every arm. Diagnostics are enabled throughout. + +## Cold prefill + +Three samples per size; prompt tokens / median first-token latency. Positive percentages indicate higher throughput than Neither. + +| Prompt | Neither tok/s | Coalescing tok/s | Sharded mHC tok/s | Both tok/s | +| --- | ---: | ---: | ---: | ---: | +| 8192 | 1949.6 (+0.0%) | 2947.3 (+51.2%) | 1958.5 (+0.5%) | 3063.5 (+57.1%) | +| 16384 | 2358.9 (+0.0%) | 2928.3 (+24.1%) | 2390.8 (+1.4%) | 3060.8 (+29.8%) | +| 32768 | 2614.5 (+0.0%) | 2923.5 (+11.8%) | 2681.8 (+2.6%) | 3049.6 (+16.6%) | + +## Decode + +One 20-second observation per cell. Each entry is aggregate tok/s / acceptance-normalized steps/s. + +| Context / concurrency | Neither | Coalescing | Sharded mHC | Both | +| --- | ---: | ---: | ---: | ---: | +| 8192 / C1 | 49.17 / 17.86 | 47.37 / 17.91 | 48.46 / 17.86 | 46.29 / 17.89 | +| 8192 / C4 | 122.19 / 43.70 | 120.18 / 44.54 | 121.17 / 44.99 | 119.37 / 42.99 | +| 32768 / C1 | 46.18 / 17.82 | 45.92 / 17.73 | 47.31 / 17.62 | 47.11 / 17.57 | +| 32768 / C4 | 123.44 / 44.34 | 125.20 / 45.64 | 123.90 / 45.84 | 124.25 / 45.00 | + +## Execution and correctness evidence + +| Configuration | Exact-answer/cache checks | Four-checkpoint dispatch | mHC common dispatches across four ranks | +| --- | ---: | --- | ---: | +| Neither | 5/5 passed | Disabled | 0 | +| Coalescing | 5/5 passed | Observed | 0 | +| Sharded mHC | 5/5 passed | Disabled | 20 | +| Both | 5/5 passed | Observed | 35 | + +The JSON contains every prefill sample, sanitized decode request measurements, source/artifact hashes and deltas. No prompts, answers, private addresses or container identifiers are included. + +## Limits + +- Three cold samples per prefill size; one 20-second decode observation per cell. +- Sequential arm order and stochastic decode acceptance limit small-difference conclusions. +- Two recovery reboots occurred between arms: one between Neither and Coalescing, and another between Coalescing and Sharded mHC. Each receipt records matching driver, GPU clock settings, persistence and CPU0 governor before and after its reboot. Reboots can change allocator/cache state and thermal conditions; this was not an interleaved A/B experiment. +- Prefill throughput is prompt tokens divided by first-token latency, including request overhead. +- MTP-normalized steps/s is aggregate tok/s divided by measured acceptance length; server steps/s is retained separately. +- Dispatch logs plus completed requests corroborate asynchronous execution without per-kernel GPU completion events. +- Exact-answer and cache-reuse tests are smoke coverage, not full model-quality or numerical-equivalence evaluation. +- Diagnostic logging remains enabled in all arms and adds host work. +- Deployment readiness/source checks and exclusive access are recorded attestations; public output excludes private deployment identifiers. diff --git a/docs/kda-checkpoint-export.md b/docs/kda-checkpoint-export.md new file mode 100644 index 000000000..5a4cedd32 --- /dev/null +++ b/docs/kda-checkpoint-export.md @@ -0,0 +1,153 @@ +# Bounded recurrent checkpoints in one KDA prefill + +The KDA prefill operation, `b12x.sequence.kda_prefill`, saves recurrent state at +up to four token offsets during one sequence traversal. A serving scheduler can use +those saved states to satisfy checkpoint boundaries without splitting the +model's forward pass at each boundary. + +Status: research-only. Forty-one CPU contract tests and twelve selected +four-checkpoint GPU tests passed on NVIDIA GB10, SM121, 48 SMs. The GPU checks +cover one and sixteen local KDA heads, including checkpoint values, graph +replay, invalid metadata, and high pool addresses. These are bounded correctness +results; standalone checkpoint-export throughput and an embedded measured +KDA-prefill policy profile remain unqualified. Exact selections and source +identities are in the [component evidence](evidence/kda-prefill/gb10-four-checkpoint-20260907/component-checks.json). + +## Checkpoint contract + +A checkpoint is a saved recurrent state after a specified number of request +tokens. The planned checkpoint capacity, `Caps.max_checkpoints`, defaults to +`1`, with vector metadata `[sequence_capacity]`. A capacity of `2` or `4` uses +contiguous matrix metadata `[sequence_capacity, max_checkpoints]` and requires +`checkpoint_export=True` with transactional metadata validation. + +State indices use the same `int32` or `int64` dtype as initial/final state +indices. Offsets are `int32`, relative to each request's packed token start. +An active offset must be positive, no larger than the live sequence length, +and divisible by the recurrence tile size of 16 tokens. The tile size is +independent of a serving scheduler's token budget. + +Nonpositive offsets disable an entry. A null destination does not own storage; +its positive offset must still be in bounds and aligned. Active non-null +exports require distinct offsets and globally unique destinations. Offsets +may be unordered. An export at the final token boundary is legal if its +destination differs from the final-state destination. + +Initial state may alias its own final slot. Checkpoints may not overwrite an +initial read, another checkpoint or a final write. Invalid device metadata +poisons the full bound output capacity, sets the error bits documented by the operation, and +leaves the recurrent pool unchanged. A request spanning pipeline windows +requires a non-null final slot for its running recurrence. + +Activations are BF16. Recurrent state is FP32 with physical layout +`[slot, head, value_dim, key_dim]` and head dimension 128. The recurrence +exports its FP32 accumulator without recomputing a prefix. Checkpoint stores +do not alter prepare arithmetic, recurrence arithmetic or final-state layout. + +Each saved state writes `heads * 128 * 128 * 4` bytes: 1 MiB at 16 heads and +2 MiB at 32 heads. The kernel does not select scheduler chunk sizes, allocate +cache pages, export convolution state or publish cache entries. Whole-model +pass removal requires a serving integration that owns those operations. + +H16 denotes sixteen local KDA heads processed by one GPU; the 64-head +GLM-5.3-Flash model has sixteen heads per rank at TP4. Four-slot capacity reserves +two more potential state destinations than two-slot capacity. At H16, two +additional enabled exports write 2 MiB per sequence. +Scratch duplicate-detection capacity grows with the planned count; request +counts and offsets remain runtime metadata. A vLLM DCP4 integration can retain +both 512-token and 2048-token replay boundaries without reducing prefix reuse. + +## Planning and device eligibility + +Construct `kda_prefill.Caps` with `max_checkpoints=2` or `4` and pass it to +`kda_prefill.plan`. The typed component policy validates the normalized device +identity `nvidia / nvidia gb10 / (12, 1) / 48 SMs` once during planning. Other +devices receive a planning error for multiple checkpoints. One-checkpoint planning +retains its device eligibility. + +The device restriction defines the validation target. There is no TP-count, +ring, HCA or network condition. Enabling another GPU requires a reviewed device +eligibility change and hardware evidence for that device. Live request values +do not enter policy queries or kernel compile/cache keys: checkpoint capacity +is planned, while indices, offsets and live token/request counts are runtime +data. Bind and replay perform no device-eligibility lookup. + +Query schema 3 permits `max_checkpoints` values 1, 2 and 4 for vector and matrix metadata. +Config schema 1 describes the launch-config fields. A schema-1 or schema-2 KDA profile is +incompatible with the schema-3 query contract and is rejected. AUTO and +HEURISTIC_ONLY resolve through the component heuristic when the registry has +no matching KDA profile. PREPLANNED_ONLY rejects that missing qualification. +Explicit configuration overrides pass the same target validator. + +The base source revision, `06b4de7c723e6f166d65abf5909c5b7d0f8acc68`, lacks a +KDA-prefill registration and generator in the planned-op catalog. The catalog +completeness test fails on that source and on this implementation. The +component therefore requires catalog/provider integration, schema-3 coverage, +and measured device profiles, or an explicitly reviewed policy for +unqualified components. This inherited policy-inventory dependency remains +separate from the bounded GPU correctness results below. + +## Verification commands and coverage + +CPU ownership, reference and policy contracts: + +```sh +.venv/bin/python -m pytest tests/sequence/test_kda_prefill_two_checkpoints_cpu.py -q +``` + +The GPU suite exercises public plan/bind/run operations. The twelve selected +four-slot cases passed with no selected skips. They cover 8192-token exports +at 4096, 6144, 7168 and 7680, nonadjacent duplicate offsets, frozen replay with +changing live metadata, and high pool indices. This command reproduces that +selection, including its large-pool case: + +```sh +B12X_RUN_LARGE_POOL_TESTS=1 .venv/bin/python -m pytest \ + tests/sequence/test_kda_prefill_two_checkpoints_gpu.py -k four_checkpoint -q -rs +``` + +The two-slot cases compare outputs +and saved states with an independent FP32 oracle; check in-place initial/final +storage; compare 8192-token exports at 6144 and 7168 with one-checkpoint +full/prefix runs; replay CUDA graphs with changing counts, destinations and +offsets under frozen kernel resolution; check allocator counters and stable +addresses; and verify transactional rejection. Those fifteen two-slot cases +were deselected in the recorded four-slot run. The full-suite command below +expands coverage beyond the recorded twelve cases: + +```sh +B12X_RUN_LARGE_POOL_TESTS=1 .venv/bin/python -m pytest \ + tests/sequence/test_kda_prefill_two_checkpoints_gpu.py -q -rs +``` + +The high-offset pool case reserves over 8 GiB, leaves unrelated pages +uninitialized, and places every live slot beyond the signed 32-bit +element-offset boundary. Run it on an explicitly selected idle GB10 with +sufficient free memory. Record a skipped case as missing coverage. + +| Geometry | Purpose | GPU evidence | +|---|---|---| +| H1, D128 | One local head: FP32 oracle, invalid metadata and frozen replay | Four-slot cases passed | +| H16, D128 | Sixteen local heads: FP32 oracle, exact 8K prefix states and high pool offsets | Four-slot cases passed | +| H32, D128 | Thirty-two local heads: additional component geometry | Unrun; no four-slot qualification | + +The serving validation target is GLM-5.3-Flash on four Sparks at TP4. Its KDA +configuration has 64 heads with dimension 128; the model adapter partitions +those heads over the TP ranks. The H32 test geometry does not establish TP2 +serving support. The [serving report](evidence/kda-prefill/gb10-four-checkpoint-20260907/report.md) +records the separate four-Spark TP4/DCP4 comparison. + +## Switch-connected Spark qualification + +A four-Spark deployment connected through a switch can run the single-GPU +component checks without SparkRing. Whole-model pass removal also requires a +vLLM integration for scheduler checkpoint boundaries, convolution checkpoint +metadata, recurrent-state ownership and cache publication. Keep its collective +backend, switch topology and serving configuration fixed across measurements. + +A performance record must identify the source and toolchain, physical device, +correctness gates, warmup and graph state, allocation behavior, raw timings and +ratio direction. Measure checkpoint-export overhead separately from a paired +whole-model coalescing comparison. The component checks establish the listed +correctness coverage. The serving measurements belong to the recorded +B12X/vLLM composition and do not measure standalone checkpoint-export latency. diff --git a/tests/sequence/test_kda_prefill.py b/tests/sequence/test_kda_prefill.py index 7593b7043..0c1838f8b 100644 --- a/tests/sequence/test_kda_prefill.py +++ b/tests/sequence/test_kda_prefill.py @@ -351,7 +351,7 @@ def pad_rows(t: torch.Tensor) -> torch.Tensor: return out def pad_seqs(t: torch.Tensor, extra: int = 0) -> torch.Tensor: - out = torch.zeros(max_seqs + extra, dtype=t.dtype, device=device) + out = torch.zeros((max_seqs + extra, *t.shape[1:]), dtype=t.dtype, device=device) out[: t.shape[0]] = t return out diff --git a/tests/sequence/test_kda_prefill_two_checkpoints_cpu.py b/tests/sequence/test_kda_prefill_two_checkpoints_cpu.py new file mode 100644 index 000000000..f22511d5d --- /dev/null +++ b/tests/sequence/test_kda_prefill_two_checkpoints_cpu.py @@ -0,0 +1,470 @@ +"""CPU contracts for bounded checkpoint export; no GPU qualification claims.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from b12x.policy import ( + ComponentProfile, + DeviceIdentity, + FrozenMapping, + GpuProfile, + InvalidPreplannedPolicyError, + PolicyContext, + PolicyMode, + PolicySource, + PreplannedPolicyNotFoundError, + ProfileRegistry, + ProfileRule, +) +from b12x.sequence.kda_prefill import _impl as impl +from b12x.sequence.kda_prefill._policy import ( + KDA_PREFILL_POLICY, + KdaPrefillConfig, + KdaPrefillQuery, +) +from b12x.sequence.kda_prefill.metadata import validate_metadata +from b12x.sequence.kda_prefill.reference import prefill_kda_chunk_mirror, recurrent_kda + +from .test_kda_prefill import PURE_FP32, make_inputs, run_oracle + + +GB10 = DeviceIdentity( + vendor="nvidia", + product_name="NVIDIA GB10", + compute_capability=(12, 1), + sm_count=48, +) + + +def query(checkpoints: int = 2) -> KdaPrefillQuery: + return KdaPrefillQuery( + heads=16, + head_dim=128, + model_dtype="bfloat16", + state_dtype="float32", + qk_l2norm=True, + checkpoint_export=True, + max_tokens=8192, + max_seqs=4, + max_checkpoints=checkpoints, + ) + + +def metadata() -> dict: + return dict( + cu_seqlens=[0, 64, 160], + initial_state_indices=[1, 5], + final_state_indices=[2, 6], + checkpoint_state_indices=[[3, 4], [7, 8]], + checkpoint_offsets=[[16, 48], [32, 80]], + num_seqs=2, + num_tokens=160, + token_capacity=192, + seq_capacity=2, + state_slots=10, + null_state_index=0, + max_checkpoints=2, + ) + + +@pytest.mark.parametrize( + "device", + [ + None, + replace(GB10, product_name="Unknown SM121"), + replace(GB10, compute_capability=(12, 0)), + replace(GB10, sm_count=47), + ], +) +def test_two_checkpoint_policy_rejects_unsupported_targets_even_with_override(device): + context = PolicyContext.for_identity(device, mode=PolicyMode.HEURISTIC_ONLY) + for override in (False, True): + kwargs = {"override": KdaPrefillConfig()} if override else {} + with pytest.raises(ValueError, match="supports only NVIDIA GB10"): + context.resolve(KDA_PREFILL_POLICY, query(), **kwargs) + # One-checkpoint planning accepts these device identities. + assert ( + context.resolve(KDA_PREFILL_POLICY, query(1)).source is PolicySource.HEURISTIC + ) + + +def test_gb10_two_checkpoint_activation_is_explicit_and_not_qualification(): + context = PolicyContext.for_identity(GB10, registry=ProfileRegistry()) + one = context.resolve(KDA_PREFILL_POLICY, query(1)) + two = context.resolve(KDA_PREFILL_POLICY, query(2)) + assert one.config == two.config + assert two.source is PolicySource.HEURISTIC + assert two.profile_id is None and two.evidence is None + assert context.resolve(KDA_PREFILL_POLICY, query(2)) is two + strict = PolicyContext.for_identity( + GB10, + mode=PolicyMode.PREPLANNED_ONLY, + registry=ProfileRegistry(), + ) + with pytest.raises(PreplannedPolicyNotFoundError): + strict.resolve(KDA_PREFILL_POLICY, query()) + + +def test_query_schema_includes_planned_checkpoint_capacity_only(): + assert KDA_PREFILL_POLICY.query_schema_version == 3 + assert KDA_PREFILL_POLICY.config_schema_version == 1 + assert set(query().profile_fields()) == KDA_PREFILL_POLICY.query_fields + assert query(1).profile_fields()["max_checkpoints"] == 1 + assert query(2).profile_fields()["max_checkpoints"] == 2 + assert { + "num_tokens", + "num_seqs", + "checkpoint_offsets", + "checkpoint_state_indices", + }.isdisjoint(KDA_PREFILL_POLICY.query_fields) + + +@pytest.mark.parametrize("schema", [1, 2]) +def test_schema1_and_schema2_profiles_rejected_by_schema3_query_contract(schema): + registry = ProfileRegistry() + registry.register( + GpuProfile( + profile_id="test.kda.schema1", + targets=(GB10,), + metadata=FrozenMapping(), + components=( + ComponentProfile( + component_id="sequence.kda_prefill", + query_schema_version=schema, + config_schema_version=1, + rules=( + ProfileRule.create( + name="synthetic-only", + exact={"heads": 16}, + ranges={}, + config=KdaPrefillConfig().to_dict(), + evidence="unit-test-not-gpu-evidence", + ), + ), + ), + ), + ) + ) + context = PolicyContext.for_identity(GB10, registry=registry) + with pytest.raises(InvalidPreplannedPolicyError, match="query schema mismatch"): + context.resolve(KDA_PREFILL_POLICY, query()) + + +@pytest.mark.parametrize("count", [0, 3, True, 2.0]) +def test_caps_reject_invalid_checkpoint_capacity(count): + with pytest.raises(ValueError, match="max_checkpoints"): + impl.Caps( + device="cuda:0", + max_tokens=32, + max_seqs=1, + max_state_slots=8, + heads=1, + checkpoint_export=True, + max_checkpoints=count, + ) + + +@pytest.mark.parametrize( + "extra", [{"checkpoint_export": False}, {"metadata_validation": "trusted"}] +) +def test_two_checkpoint_caps_require_transactional_export(extra): + kwargs = dict(checkpoint_export=True, metadata_validation="transactional") + kwargs.update(extra) + with pytest.raises(ValueError, match="multiple checkpoints require"): + impl.Caps( + device="cuda:0", + max_tokens=32, + max_seqs=1, + max_state_slots=8, + heads=1, + max_checkpoints=2, + **kwargs, + ) + + +@pytest.mark.parametrize("slot", [1, 2, 3, 5, 6, 7, 8]) +def test_checkpoint_destinations_cannot_alias_other_owners(slot): + args = metadata() + args["checkpoint_state_indices"][0][1] = slot + with pytest.raises(ValueError): + validate_metadata(**args) + + +@pytest.mark.parametrize("offset", [17, 80, 16]) +def test_second_checkpoint_rejects_unaligned_past_end_or_duplicate_offset(offset): + args = metadata() + args["checkpoint_offsets"][0][1] = offset + with pytest.raises(ValueError): + validate_metadata(**args) + + +@pytest.mark.parametrize("slot", [-1, 10]) +def test_second_checkpoint_rejects_invalid_active_slot(slot): + args = metadata() + args["checkpoint_state_indices"][0][1] = slot + with pytest.raises(IndexError): + validate_metadata(**args) + + +def test_null_disabled_reversed_and_final_boundary_exports(): + args = metadata() + args["final_state_indices"][0] = 1 # Own initial/final alias is legal. + args["checkpoint_offsets"] = [[64, 16], [96, 32]] + assert validate_metadata(**args) == [(0, 64), (64, 160)] + args["checkpoint_state_indices"][0][1] = 0 + args["checkpoint_offsets"][0][1] = 64 # A null writer does not own an offset. + validate_metadata(**args) + args["checkpoint_offsets"][0][1] = 17 + with pytest.raises(ValueError, match="unaligned"): + validate_metadata(**args) + args["checkpoint_state_indices"][0][1] = 1 + for offset in (0, -1): + args["checkpoint_offsets"][0][1] = offset + validate_metadata(**args) + + +def test_inactive_metadata_is_ignored_but_live_counts_are_bounded(): + args = metadata() + args.update(num_seqs=1, num_tokens=64) + args["checkpoint_state_indices"][1] = [-999, -999] + args["checkpoint_offsets"][1] = [17, 999] + assert validate_metadata(**args) == [(0, 64)] + args["num_seqs"] = 3 + with pytest.raises(ValueError, match="capacities"): + validate_metadata(**args) + + +@pytest.mark.parametrize("inplace", [False, True]) +def test_two_exports_equal_independent_prefix_recurrences_and_final_state(inplace): + inputs = make_inputs(lengths=[64], heads=1, seed=830, state_slots=6) + if inplace: + inputs["final"][0] = inputs["initial"][0] + one_checkpoint_output, one_checkpoint_pool = run_oracle(inputs) + inputs["checkpoint_slots"] = torch.tensor([[3, 4]], dtype=torch.int32) + inputs["checkpoint_offsets"] = torch.tensor([[48, 16]], dtype=torch.int32) + output, pool = run_oracle(inputs, max_checkpoints=2) + torch.testing.assert_close(output, one_checkpoint_output, rtol=0, atol=0) + torch.testing.assert_close( + pool[int(inputs["final"][0])], + one_checkpoint_pool[int(inputs["final"][0])], + rtol=0, + atol=0, + ) + for slot, offset in ((3, 48), (4, 16)): + _, expected, _ = recurrent_kda( + *(inputs[name][:offset] for name in ("q", "k", "v", "raw_g", "raw_beta")), + inputs["A_log"], + inputs["dt_bias"], + lower_bound=-5.0, + initial_state=inputs["pool"][int(inputs["initial"][0])], + ) + torch.testing.assert_close(pool[slot], expected, rtol=0, atol=0) + _, mirror_pool = run_oracle( + inputs, + fn=prefill_kda_chunk_mirror, + max_checkpoints=2, + policy=PURE_FP32, + ) + for slot in (int(inputs["final"][0]), 3, 4): + torch.testing.assert_close(mirror_pool[slot], pool[slot], rtol=2e-4, atol=2e-5) + + +def test_reference_rejection_is_transactional_for_pool_and_output(): + inputs = make_inputs(lengths=[64], heads=1, seed=831, state_slots=6) + inputs["checkpoint_slots"] = torch.tensor([[3, 1]], dtype=torch.int32) + inputs["checkpoint_offsets"] = torch.tensor([[16, 48]], dtype=torch.int32) + before = inputs["pool"].clone() + output = torch.full_like(inputs["q"], 7) + from b12x.sequence.kda_prefill.reference import prefill_kda + + with pytest.raises(ValueError, match="duplicate"): + prefill_kda( + *(inputs[name] for name in ("q", "k", "v", "raw_g", "raw_beta")), + inputs["A_log"], + inputs["dt_bias"], + inputs["pool"], + inputs["cu_seqlens"], + inputs["initial"], + inputs["final"], + inputs["checkpoint_slots"], + inputs["checkpoint_offsets"], + 1, + 64, + max_checkpoints=2, + output=output, + ) + assert torch.equal(inputs["pool"], before) + assert torch.all(output == 7) + + +def test_plural_reference_zero_offset_returns_an_independent_initial_snapshot(): + inputs = make_inputs(lengths=[32], heads=1, seed=832) + initial = inputs["pool"][0] + _, _, snapshots = recurrent_kda( + *(inputs[name] for name in ("q", "k", "v", "raw_g", "raw_beta")), + inputs["A_log"], + inputs["dt_bias"], + lower_bound=-5.0, + initial_state=initial, + checkpoint_offsets=(0, 16), + ) + assert isinstance(snapshots, dict) and set(snapshots) == {0, 16} + assert torch.equal(snapshots[0], initial) + assert snapshots[0].data_ptr() != initial.data_ptr() + + +@pytest.mark.parametrize("fault", ["shape", "dtype"]) +def test_bind_reports_checkpoint_tensor_contract_before_cross_index_dtype(fault): + from torch._subclasses.fake_tensor import FakeTensorMode + + caps = impl.Caps( + device="cuda:0", + max_tokens=32, + max_seqs=1, + max_state_slots=4, + heads=1, + checkpoint_export=True, + max_checkpoints=2, + ) + # Materialize only the static layout; fake CUDA tensors allocate no storage. + plan = impl._materialize_plan( + caps, + v_split=64, + k_split=1, + stages=3, + window_tiles=4, + policy_resolution=None, + ) + with FakeTensorMode(): + + def tensor(shape, dtype=torch.bfloat16): + return torch.empty(shape, device="cuda:0", dtype=dtype) + + inputs = {name: tensor((32, 1, 128)) for name in ("q", "k", "v", "raw_g")} + checkpoint_indices = tensor( + (2,) if fault == "shape" else (1, 2), + torch.int16, + ) + with pytest.raises( + (ValueError, TypeError), match="checkpoint_state_indices must have" + ): + impl.bind( + plan, + scratch=tensor(plan.scratch_specs()[0].shape, torch.uint8), + **inputs, + raw_beta=tensor((32, 1)), + A_log=tensor((1,), torch.float32), + dt_bias=tensor((1, 128), torch.float32), + recurrent_state=tensor((4, 1, 128, 128), torch.float32), + cu_seqlens=tensor((2,), torch.int32), + initial_state_indices=tensor((1,), torch.int64), + final_state_indices=tensor((1,), torch.int64), + checkpoint_state_indices=checkpoint_indices, + checkpoint_offsets=tensor((1, 2), torch.int32), + num_seqs=tensor((1,), torch.int32), + num_tokens=tensor((1,), torch.int32), + output=tensor((32, 1, 128)), + ) + + +@pytest.mark.parametrize("capacity", [2, 4]) +def test_multi_checkpoint_caps_validate_matrix_width_before_live_scalar_dtype(capacity): + from torch._subclasses.fake_tensor import FakeTensorMode + + caps = impl.Caps( + device="cuda:0", + max_tokens=64, + max_seqs=2, + max_state_slots=16, + heads=1, + checkpoint_export=True, + max_checkpoints=capacity, + ) + plan = impl._materialize_plan( + caps, v_split=64, k_split=1, stages=3, window_tiles=4, policy_resolution=None + ) + with FakeTensorMode(): + + def tensor(shape, dtype=torch.bfloat16): + return torch.empty(shape, device="cuda:0", dtype=dtype) + + # Fake tensors have no distinct device addresses. Stop before alias checks. + with pytest.raises((ValueError, TypeError), match="num_tokens must have"): + impl.bind( + plan, + scratch=tensor(plan.scratch_specs()[0].shape, torch.uint8), + **{name: tensor((64, 1, 128)) for name in ("q", "k", "v", "raw_g")}, + raw_beta=tensor((64, 1)), + A_log=tensor((1,), torch.float32), + dt_bias=tensor((1, 128), torch.float32), + recurrent_state=tensor((16, 1, 128, 128), torch.float32), + cu_seqlens=tensor((3,), torch.int32), + initial_state_indices=tensor((2,), torch.int32), + final_state_indices=tensor((2,), torch.int32), + checkpoint_state_indices=tensor((2, capacity), torch.int32), + checkpoint_offsets=tensor((2, capacity), torch.int32), + num_seqs=tensor((1,), torch.int32), + num_tokens=tensor((1,), torch.int16), + output=tensor((64, 1, 128)), + ) + resolved = PolicyContext.for_identity(GB10, registry=ProfileRegistry()).resolve( + KDA_PREFILL_POLICY, query(capacity) + ) + assert resolved.source is PolicySource.HEURISTIC + with pytest.raises(ValueError, match="supports only NVIDIA GB10"): + PolicyContext.for_identity(None).resolve(KDA_PREFILL_POLICY, query(capacity)) + + +@pytest.mark.parametrize("inplace", [False, True]) +def test_four_exports_equal_independent_prefix_recurrences(inplace): + inputs = make_inputs(lengths=[80], heads=1, seed=842, state_slots=8) + if inplace: + inputs["final"][0] = inputs["initial"][0] + baseline_output, baseline_pool = run_oracle(inputs) + inputs["checkpoint_slots"] = torch.tensor([[2, 3, 4, 5]], dtype=torch.int32) + inputs["checkpoint_offsets"] = torch.tensor([[64, 16, 48, 32]], dtype=torch.int32) + output, pool = run_oracle(inputs, max_checkpoints=4) + torch.testing.assert_close(output, baseline_output, rtol=0, atol=0) + torch.testing.assert_close( + pool[int(inputs["final"][0])], + baseline_pool[int(inputs["final"][0])], + rtol=0, + atol=0, + ) + for slot, offset in zip((2, 3, 4, 5), (64, 16, 48, 32), strict=True): + _, expected, _ = recurrent_kda( + *(inputs[name][:offset] for name in ("q", "k", "v", "raw_g", "raw_beta")), + inputs["A_log"], + inputs["dt_bias"], + lower_bound=-5.0, + initial_state=inputs["pool"][int(inputs["initial"][0])], + ) + torch.testing.assert_close(pool[slot], expected, rtol=1e-6, atol=1e-9) + + +@pytest.mark.parametrize("fault", ["offset", "slot", "initial"]) +def test_four_checkpoint_nonadjacent_aliases_are_rejected(fault): + args = metadata() + args.update( + cu_seqlens=[0, 80], + initial_state_indices=[0], + final_state_indices=[1], + checkpoint_state_indices=[[2, 3, 4, 5]], + checkpoint_offsets=[[16, 32, 48, 64]], + num_seqs=1, + num_tokens=80, + null_state_index=None, + max_checkpoints=4, + ) + if fault == "offset": + args["checkpoint_offsets"][0][3] = 16 + elif fault == "slot": + args["checkpoint_state_indices"][0][3] = 2 + else: + args["checkpoint_state_indices"][0][3] = 0 + with pytest.raises(ValueError): + validate_metadata(**args) diff --git a/tests/sequence/test_kda_prefill_two_checkpoints_gpu.py b/tests/sequence/test_kda_prefill_two_checkpoints_gpu.py new file mode 100644 index 000000000..4343c3b4b --- /dev/null +++ b/tests/sequence/test_kda_prefill_two_checkpoints_gpu.py @@ -0,0 +1,552 @@ +"""GB10 GPU qualification tests for two recurrent checkpoint exports. + +H16 and H32 correspond to GLM-5.3-Flash's 64 KDA heads at TP4 and TP2. +These are per-rank component tests, not multi-GPU serving qualification. +""" + +from __future__ import annotations + +import os + +import pytest +import torch + +from .test_kda_prefill import ( + HEAD_DIM, + _run as run_op, + assert_kda_close, + make_binding, + make_inputs, + run_oracle, +) + + +def require_gb10(): + from ..conftest import require_b12x + from b12x.policy import PolicyContext + + device = require_b12x() + identity = PolicyContext.for_device(device).device + if ( + identity is None + or identity.vendor != "nvidia" + or identity.product_name != "nvidia gb10" + or identity.compute_capability != (12, 1) + or identity.sm_count != 48 + ): + pytest.skip( + "two-checkpoint GPU qualification requires NVIDIA GB10 / SM121 / 48 SMs" + ) + return device + + +def checkpoint_inputs( + lengths, *, heads=1, offsets=None, seed=900, state_slots=16, capacity=2, device +): + count = len(lengths) + inputs = make_inputs( + lengths=lengths, + heads=heads, + seed=seed, + state_slots=state_slots, + device=device, + ) + inputs["checkpoint_slots"] = torch.arange( + 2 * count, + (2 + capacity) * count, + dtype=torch.int32, + device=device, + ).reshape(count, capacity) + inputs["checkpoint_offsets"] = torch.tensor( + offsets or [[16, length // 16 * 16] for length in lengths], + dtype=torch.int32, + device=device, + ) + return inputs + + +def assert_checkpoint_oracle(binding, tensors, inputs): + capacity = inputs["checkpoint_slots"].shape[1] + expected_out, expected_pool = run_oracle(inputs, max_checkpoints=capacity) + tokens = inputs["num_tokens"] + assert binding.error_code.item() == 0 + assert_kda_close( + "output", expected_out[:tokens], binding.output[:tokens], ratio=1e-2 + ) + writes = set(inputs["final"].tolist()) + for slots, offsets in zip( + inputs["checkpoint_slots"].tolist(), + inputs["checkpoint_offsets"].tolist(), + strict=True, + ): + writes.update( + slot for slot, offset in zip(slots, offsets, strict=True) if offset > 0 + ) + for slot in writes: + assert_kda_close( + f"state[{slot}]", + expected_pool[slot], + tensors["recurrent_state"][slot], + ratio=5e-3, + ) + untouched = sorted(set(range(inputs["pool"].shape[0])) - writes) + torch.testing.assert_close( + tensors["recurrent_state"][untouched], inputs["pool"][untouched], rtol=0, atol=0 + ) + + +@pytest.mark.parametrize( + "heads", [1, 16, 32], ids=["minimal", "glm-tp4-heads", "glm-tp2-heads"] +) +@pytest.mark.parametrize("inplace", [False, True]) +def test_two_checkpoint_gpu_independent_fp32_oracle(heads, inplace): + device = require_gb10() + inputs = checkpoint_inputs( + [80, 64], offsets=[[64, 16], [16, 64]], heads=heads, device=device + ) + if inplace: + inputs["final"][0] = inputs["initial"][0] + binding, tensors = make_binding( + inputs, + max_tokens=256, + max_seqs=4, + final_stride=3, + checkpoint_export=True, + max_checkpoints=2, + ) + + run_op(binding, inputs) + torch.cuda.synchronize(device) + assert_checkpoint_oracle(binding, tensors, inputs) + + +@pytest.mark.parametrize("heads", [16, 32], ids=["glm-tp4-heads", "glm-tp2-heads"]) +def test_two_checkpoint_gpu_8k_matches_one_checkpoint_prefixes(heads): + device = require_gb10() + + inputs = checkpoint_inputs( + [8192], heads=heads, offsets=[[6144, 7168]], device=device + ) + binding, tensors = make_binding( + inputs, + max_tokens=8192, + max_seqs=1, + checkpoint_export=True, + max_checkpoints=2, + ) + run_op(binding, inputs) + torch.cuda.synchronize(device) + assert binding.error_code.item() == 0 + for length, slot in ((8192, 1), (6144, 2), (7168, 3)): + prefix = dict(inputs) + for name in ("q", "k", "v", "raw_g", "raw_beta"): + prefix[name] = inputs[name][:length] + prefix.update( + num_tokens=length, + cu_seqlens=torch.tensor([0, length], dtype=torch.int32, device=device), + checkpoint_slots=torch.tensor([0], dtype=torch.int32, device=device), + checkpoint_offsets=torch.tensor([0], dtype=torch.int32, device=device), + ) + one_checkpoint, one_checkpoint_tensors = make_binding( + prefix, max_tokens=8192, max_seqs=1 + ) + run_op(one_checkpoint, prefix) + torch.cuda.synchronize(device) + assert one_checkpoint.error_code.item() == 0 + torch.testing.assert_close( + tensors["recurrent_state"][slot], + one_checkpoint_tensors["recurrent_state"][1], + rtol=0, + atol=0, + ) + if length == 8192: + torch.testing.assert_close( + binding.output, one_checkpoint.output, rtol=0, atol=0 + ) + + +def copy_live(tensors, live): + for name in ("q", "k", "v", "raw_g", "raw_beta"): + tensors[name].zero_() + tensors[name][: live["num_tokens"]].copy_(live[name]) + for destination, source in ( + ("cu_seqlens", "cu_seqlens"), + ("initial_state_indices", "initial"), + ("final_state_indices", "final"), + ("checkpoint_state_indices", "checkpoint_slots"), + ("checkpoint_offsets", "checkpoint_offsets"), + ): + tensors[destination].zero_() + tensors[destination][: live[source].shape[0]].copy_(live[source]) + for name in ("A_log", "dt_bias"): + tensors[name].copy_(live[name]) + tensors["recurrent_state"].copy_(live["pool"]) + tensors["num_tokens"].fill_(live["num_tokens"]) + tensors["num_seqs"].fill_(live["num_seqs"]) + + +def test_two_checkpoint_gpu_frozen_replay_changes_live_counts_and_offsets(): + device = require_gb10() + from b12x._lib.runtime_control import ( + freeze_kernel_resolution, + unfreeze_kernel_resolution, + ) + from b12x.sequence.kda_prefill import _cute_kernels as kernels + + first = checkpoint_inputs([96], device=device) + binding, tensors = make_binding( + first, max_tokens=256, max_seqs=4, checkpoint_export=True, max_checkpoints=2 + ) + run_op(binding, first) + torch.cuda.synchronize(device) + launchers = ( + kernels._PROLOGUE_CACHE[kernels._prologue_key(binding)], + kernels._PREPARE_CACHE[kernels._prepare_key(binding)], + kernels._RECURRENCE_CACHE[kernels._recurrence_key(binding)], + ) + addresses = tuple(t.data_ptr() for t in (*tensors.values(), binding.scratch)) + freeze_kernel_resolution("two-checkpoint GB10 qualification") + try: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_op(binding, first) + for lengths, offsets in ( + ([128], [[112, 32]]), + ([64, 80, 96], [[16, 48], [64, 32], [32, 80]]), + ([48], [[32, 0]]), + ): + live = checkpoint_inputs( + lengths, offsets=offsets, seed=sum(lengths), device=device + ) + copy_live(tensors, live) + binding.output.fill_(float("nan")) + binding.scratch.fill_(0xFF) + torch.cuda.synchronize(device) + before = torch.cuda.memory_stats(device)["allocation.all.allocated"] + graph.replay() + torch.cuda.synchronize(device) + assert torch.cuda.memory_stats(device)["allocation.all.allocated"] == before + assert addresses == tuple( + t.data_ptr() for t in (*tensors.values(), binding.scratch) + ) + assert_checkpoint_oracle(binding, tensors, live) + assert torch.isnan(binding.output[live["num_tokens"] :].float()).all() + finally: + unfreeze_kernel_resolution() + assert launchers == ( + kernels._PROLOGUE_CACHE[kernels._prologue_key(binding)], + kernels._PREPARE_CACHE[kernels._prepare_key(binding)], + kernels._RECURRENCE_CACHE[kernels._recurrence_key(binding)], + ) + + +@pytest.mark.parametrize( + "fault", + [ + "duplicate-slot", + "initial-alias", + "duplicate-offset", + "unaligned", + "out-of-range", + ], +) +def test_two_checkpoint_gpu_invalid_metadata_preserves_state(fault): + device = require_gb10() + + inputs = checkpoint_inputs([64], device=device) + binding, tensors = make_binding( + inputs, max_tokens=96, max_seqs=2, checkpoint_export=True, max_checkpoints=2 + ) + if fault == "duplicate-slot": + tensors["checkpoint_state_indices"][0, 1] = tensors["checkpoint_state_indices"][ + 0, 0 + ] + elif fault == "initial-alias": + tensors["checkpoint_state_indices"][0, 1] = 0 + elif fault == "duplicate-offset": + tensors["checkpoint_offsets"][0, 1] = 16 + elif fault == "unaligned": + tensors["checkpoint_offsets"][0, 1] = 17 + else: + tensors["checkpoint_state_indices"][0, 1] = 16 + before = tensors["recurrent_state"].clone() + binding.output.fill_(7) + run_op(binding, inputs) + torch.cuda.synchronize(device) + assert binding.error_code.item() != 0 + assert torch.equal(tensors["recurrent_state"], before) + # Transactional failure poisons the full bound output capacity. + assert torch.isnan(binding.output.float()).all() + + +def test_two_checkpoint_gpu_high_pool_offsets(): + if os.environ.get("B12X_RUN_LARGE_POOL_TESTS") != "1": + pytest.skip( + "set B12X_RUN_LARGE_POOL_TESTS=1 on an idle GB10; requires over 8 GiB" + ) + device = require_gb10() + + slot_stride = HEAD_DIM * HEAD_DIM + 2048 + high = (1 << 31) // slot_stride + 8 + storage_elements = (high + 4) * slot_stride + free, _ = torch.cuda.mem_get_info(device) + if free < storage_elements * 4 + (2 << 30): + pytest.skip( + "insufficient free memory for an 8 GiB high-offset pool plus scratch" + ) + inputs = checkpoint_inputs([64], offsets=[[16, 48]], state_slots=4, device=device) + compact, compact_tensors = make_binding( + inputs, max_tokens=64, max_seqs=1, checkpoint_export=True, max_checkpoints=2 + ) + run_op(compact, inputs) + torch.cuda.synchronize(device) + assert compact.error_code.item() == 0 + storage = torch.empty(storage_elements, dtype=torch.float32, device=device) + pool = torch.as_strided( + storage, + (high + 4, 1, HEAD_DIM, HEAD_DIM), + (slot_stride, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), + ) + for slot in range(4): + pool[high + slot].copy_(inputs["pool"][slot]) + large_inputs = dict(inputs) + large_inputs["initial"] = torch.tensor([high], dtype=torch.int64, device=device) + large_inputs["final"] = torch.tensor([high + 1], dtype=torch.int64, device=device) + large_inputs["checkpoint_slots"] = torch.tensor( + [[high + 2, high + 3]], dtype=torch.int64, device=device + ) + large, _ = make_binding( + large_inputs, + max_tokens=64, + max_seqs=1, + recurrent_state=pool, + checkpoint_export=True, + max_checkpoints=2, + ) + run_op(large, large_inputs) + torch.cuda.synchronize(device) + assert large.error_code.item() == 0 + torch.testing.assert_close(large.output, compact.output, rtol=0, atol=0) + for slot in range(4): + torch.testing.assert_close( + pool[high + slot], compact_tensors["recurrent_state"][slot], rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("heads", [1, 16], ids=["minimal", "glm-tp4-heads"]) +@pytest.mark.parametrize("inplace", [False, True]) +def test_four_checkpoint_gpu_independent_fp32_oracle(heads, inplace): + device = require_gb10() + inputs = checkpoint_inputs( + [80, 96], + heads=heads, + capacity=4, + offsets=[[64, 16, 48, 32], [16, 80, 48, 96]], + device=device, + ) + if inplace: + inputs["final"][0] = inputs["initial"][0] + binding, tensors = make_binding( + inputs, max_tokens=256, max_seqs=4, checkpoint_export=True, max_checkpoints=4 + ) + run_op(binding, inputs) + torch.cuda.synchronize(device) + assert_checkpoint_oracle(binding, tensors, inputs) + + +def test_four_checkpoint_gpu_8k_matches_fine_and_coarse_prefixes(): + device = require_gb10() + positions = (4096, 6144, 7168, 7680) + inputs = checkpoint_inputs( + [8192], heads=16, capacity=4, offsets=[list(positions)], device=device + ) + binding, tensors = make_binding( + inputs, max_tokens=8192, max_seqs=1, checkpoint_export=True, max_checkpoints=4 + ) + run_op(binding, inputs) + torch.cuda.synchronize(device) + assert binding.error_code.item() == 0 + for length, slot in ((8192, 1), *zip(positions, (2, 3, 4, 5), strict=True)): + prefix = dict(inputs) + for name in ("q", "k", "v", "raw_g", "raw_beta"): + prefix[name] = inputs[name][:length] + prefix.update( + num_tokens=length, + cu_seqlens=torch.tensor([0, length], dtype=torch.int32, device=device), + checkpoint_slots=torch.tensor([0], dtype=torch.int32, device=device), + checkpoint_offsets=torch.tensor([0], dtype=torch.int32, device=device), + ) + single, single_tensors = make_binding(prefix, max_tokens=8192, max_seqs=1) + run_op(single, prefix) + torch.cuda.synchronize(device) + assert single.error_code.item() == 0 + torch.testing.assert_close( + tensors["recurrent_state"][slot], + single_tensors["recurrent_state"][1], + rtol=0, + atol=0, + ) + if length == 8192: + torch.testing.assert_close(binding.output, single.output, rtol=0, atol=0) + + +def test_four_checkpoint_gpu_frozen_replay_changes_live_counts_and_offsets(): + device = require_gb10() + from b12x._lib.runtime_control import ( + freeze_kernel_resolution, + unfreeze_kernel_resolution, + ) + from b12x.sequence.kda_prefill import _cute_kernels as kernels + + first = checkpoint_inputs( + [96], capacity=4, state_slots=32, offsets=[[16, 32, 64, 80]], device=device + ) + binding, tensors = make_binding( + first, max_tokens=256, max_seqs=4, checkpoint_export=True, max_checkpoints=4 + ) + run_op(binding, first) + torch.cuda.synchronize(device) + launchers = ( + kernels._PROLOGUE_CACHE[kernels._prologue_key(binding)], + kernels._PREPARE_CACHE[kernels._prepare_key(binding)], + kernels._RECURRENCE_CACHE[kernels._recurrence_key(binding)], + ) + addresses = tuple(t.data_ptr() for t in (*tensors.values(), binding.scratch)) + freeze_kernel_resolution("four-checkpoint GB10 qualification") + try: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_op(binding, first) + for lengths, offsets in ( + ([128], [[112, 32, 64, 16]]), + ([64, 80, 96], [[16, 32, 48, 64], [64, 32, 16, 80], [32, 80, 16, 64]]), + ([48], [[32, 0, 16, 0]]), + ): + live = checkpoint_inputs( + lengths, + capacity=4, + offsets=offsets, + state_slots=32, + seed=sum(lengths), + device=device, + ) + copy_live(tensors, live) + binding.output.fill_(float("nan")) + binding.scratch.fill_(0xFF) + torch.cuda.synchronize(device) + before = torch.cuda.memory_stats(device)["allocation.all.allocated"] + graph.replay() + torch.cuda.synchronize(device) + assert torch.cuda.memory_stats(device)["allocation.all.allocated"] == before + assert addresses == tuple( + t.data_ptr() for t in (*tensors.values(), binding.scratch) + ) + assert_checkpoint_oracle(binding, tensors, live) + assert torch.isnan(binding.output[live["num_tokens"] :].float()).all() + finally: + unfreeze_kernel_resolution() + assert launchers == ( + kernels._PROLOGUE_CACHE[kernels._prologue_key(binding)], + kernels._PREPARE_CACHE[kernels._prepare_key(binding)], + kernels._RECURRENCE_CACHE[kernels._recurrence_key(binding)], + ) + + +@pytest.mark.parametrize( + "fault", + [ + "duplicate-slot", + "initial-alias", + "duplicate-offset", + "unaligned", + "out-of-range", + ], +) +def test_four_checkpoint_gpu_invalid_fourth_metadata_preserves_state(fault): + device = require_gb10() + inputs = checkpoint_inputs( + [80], capacity=4, offsets=[[16, 32, 48, 64]], device=device + ) + binding, tensors = make_binding( + inputs, max_tokens=96, max_seqs=2, checkpoint_export=True, max_checkpoints=4 + ) + if fault == "duplicate-slot": + tensors["checkpoint_state_indices"][0, 3] = tensors["checkpoint_state_indices"][ + 0, 0 + ] + elif fault == "initial-alias": + tensors["checkpoint_state_indices"][0, 3] = 0 + elif fault == "duplicate-offset": + tensors["checkpoint_offsets"][0, 3] = 16 + elif fault == "unaligned": + tensors["checkpoint_offsets"][0, 3] = 17 + else: + tensors["checkpoint_state_indices"][0, 3] = 16 + before = tensors["recurrent_state"].clone() + binding.output.fill_(7) + run_op(binding, inputs) + torch.cuda.synchronize(device) + assert binding.error_code.item() != 0 + assert torch.equal(tensors["recurrent_state"], before) + assert torch.isnan(binding.output.float()).all() + + +def test_four_checkpoint_gpu_high_pool_offsets(): + if os.environ.get("B12X_RUN_LARGE_POOL_TESTS") != "1": + pytest.skip( + "set B12X_RUN_LARGE_POOL_TESTS=1 on an idle GB10; requires over 8 GiB" + ) + device = require_gb10() + heads = 16 + slot_stride = heads * HEAD_DIM * HEAD_DIM + 2048 + high = (1 << 31) // slot_stride + 8 + storage_elements = (high + 6) * slot_stride + free, _ = torch.cuda.mem_get_info(device) + if free < storage_elements * 4 + (2 << 30): + pytest.skip( + "insufficient free memory for an 8 GiB high-offset pool plus scratch" + ) + inputs = checkpoint_inputs( + [80], + heads=heads, + capacity=4, + state_slots=6, + offsets=[[16, 32, 48, 64]], + device=device, + ) + compact, compact_tensors = make_binding( + inputs, max_tokens=80, max_seqs=1, checkpoint_export=True, max_checkpoints=4 + ) + run_op(compact, inputs) + torch.cuda.synchronize(device) + assert compact.error_code.item() == 0 + storage = torch.empty(storage_elements, dtype=torch.float32, device=device) + pool = torch.as_strided( + storage, + (high + 6, heads, HEAD_DIM, HEAD_DIM), + (slot_stride, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), + ) + for slot in range(6): + pool[high + slot].copy_(inputs["pool"][slot]) + large_inputs = dict(inputs) + large_inputs["initial"] = torch.tensor([high], dtype=torch.int64, device=device) + large_inputs["final"] = torch.tensor([high + 1], dtype=torch.int64, device=device) + large_inputs["checkpoint_slots"] = torch.tensor( + [[high + i for i in (2, 3, 4, 5)]], dtype=torch.int64, device=device + ) + large, _ = make_binding( + large_inputs, + max_tokens=80, + max_seqs=1, + recurrent_state=pool, + checkpoint_export=True, + max_checkpoints=4, + ) + run_op(large, large_inputs) + torch.cuda.synchronize(device) + assert large.error_code.item() == 0 + torch.testing.assert_close(large.output, compact.output, rtol=0, atol=0) + for slot in range(6): + torch.testing.assert_close( + pool[high + slot], compact_tensors["recurrent_state"][slot], rtol=0, atol=0 + )