Autotune complete sparse-candidate exchange and selection

docs/model-preparation.md

@@ -1004,6 +1004,19 @@ of two; key scales use `max(absmax, 1e-4) / 448` in FP32.
 MLA query/KV normalization and indexer LayerNorm use the checkpoint's
 `rms_norm_eps` value of `1e-5`.
 
+With compiled forwards, speculation enabled, `--collective-tuning static`
+(the default), and `--autotune-effort default` or `audit`, preparation measures
+two complete sparse-candidate exchanges for single-request decode and
+verification. One gathers every rank's candidates on every rank; the other
+sends each query's candidates to one owner, selects there, and broadcasts only
+the selected IDs. Both preserve FP32 score words, cutoff ties and ID order.
+The measurement includes packing, transport, selection and replication, and
+binds one implementation per row capacity before model capture. The choice is
+recorded under `candidate_compositions` in preparation measurements. Prefill,
+multiple-request graphs, runtime collective tuning and disabled autotuning
+retain the replicated exchange. No serving flag or runtime graph change is
+required.
+
 With the default `--prefill-chunk-size 4096`, the 4096-token prefill graph
 stages owner-encoded latent and index history locally, distributes index-query
 rows across context ranks, and gathers final selected IDs once per layer.

ignition/components/candidate_selection.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exact rank-local candidate merging with prepared communication alternatives."""
+
+import torch
+
+from .candidate_rows import PreparedCandidateRows
+from .collectives import PreparedAllGather
+from .ordered_row_topk import PreparedOrderedRowTopK
+from .owner_topk import PreparedOwnerTopK
+
+
+class PreparedCandidateSelection:
+    """Normalize local IDs and select global rows without changing score bits.
+
+    The replicated composition gathers every candidate on every rank. The
+    owner composition selects each query on one rank and broadcasts final IDs.
+    Both use the same candidate preprocessing, cutoff ties and emission order.
+    Preparation binds one composition before model compilation; serving has
+    no runtime choice. Serial calls may share storage and transport channels.
+    """
+
+    def __init__(self, indices, scores, positions, rank, output, *, world, gather_runtime,
+                 owner_runtime, reduction_runtime, session, coordinator, name, buffers,
+                 dependencies=(), alternatives=("replicated", "owner")):
+        rows, topk = indices.shape
+        if (indices.dtype != torch.int32 or scores.shape != indices.shape or scores.dtype != torch.float32
+                or positions.shape != (rows,) or positions.dtype != torch.int64
+                or output.shape != (rows, topk) or output.dtype != torch.int32
+                or (rank is None and world != 1)
+                or (rank is not None and (rank.shape != (1,) or rank.dtype != torch.int64))
+                or any(not value.is_contiguous() or value.device != indices.device
+                       for value in (indices, scores, positions, rank, output) if value is not None)):
+            raise ValueError("candidate merging requires matching contiguous local rows and global output IDs")
+        if not alternatives or not set(alternatives) <= {"replicated", "owner"}:
+            raise ValueError("candidate merging requires supported prepared alternatives")
+        self.indices, self.scores, self.positions, self.output = indices, scores, positions, output
+        self.rank = 0 if rank is None else rank
+        self.rows, self.world, self.topk = rows, world, topk
+        self.dependencies = list(dependencies)
+        self.alternatives = tuple(alternatives)
+        self.implementation = None
+        self.active = buffers.zeros((rows,), torch.bool)
+        self.gather = self.owner = None
+        if "replicated" in alternatives:
+            self.packed = buffers.zeros((rows, topk, 2), torch.float32)
+            self.gathered = buffers.zeros((rows, topk*world, 2), torch.float32)
+            self.values = buffers.zeros((rows, topk*world), torch.float32)
+            self.global_ids = buffers.zeros((rows, topk*world), torch.int32)
+            self.lengths = buffers.zeros((rows,), torch.int32)
+            self.selected_scores = buffers.zeros((rows, topk), torch.float32)
+            self.gather = PreparedAllGather(self.packed.view(torch.bfloat16).view(rows, -1),
+                self.gathered.view(torch.bfloat16).view(rows, -1), runtime=gather_runtime,
+                session=session, coordinator=coordinator, name=name+".replicated.transport")
+            self.gather.dependencies = self.dependencies
+            self.candidates = PreparedCandidateRows(self.gathered, self.active, self.values,
+                self.global_ids, self.lengths, topk=topk, buffers=buffers)
+            self.fold = PreparedOrderedRowTopK(self.values, self.lengths, self.global_ids,
+                self.selected_scores, output, session=session, name=name+".replicated.fold", buffers=buffers)
+        if "owner" in alternatives:
+            if owner_runtime is None or reduction_runtime is None or owner_runtime.world_size != world:
+                raise ValueError("owner selection requires matching prepared transport owners")
+            self.planar = buffers.zeros((2*rows, topk), torch.float32)
+            self.owner_scores, self.owner_ids = self.planar[:rows], self.planar[rows:].view(torch.int32)
+            self.owner = PreparedOwnerTopK(self.owner_ids, self.owner_scores, self.active, rank, output,
+                owner_runtime=owner_runtime, reduction_runtime=reduction_runtime, session=session,
+                coordinator=coordinator, name=name+".owner", buffers=buffers)
+            self.owner.transport.dependencies = self.dependencies
+
+    def bind(self, implementation):
+        if implementation not in self.alternatives:
+            raise ValueError("candidate selection was not prepared for this implementation")
+        self.implementation = implementation
+
+    def plans(self, implementation):
+        if implementation == "replicated" and self.gather is not None:
+            return (("transport", self.gather.plan), ("selection", self.fold.plan))
+        if implementation == "owner" and self.owner is not None:
+            return (("transport", self.owner.transport.plan), ("selection", self.owner.fold.plan),
+                    ("replication", self.owner.broadcast.reduction.plan))
+        raise ValueError("candidate selection implementation has no prepared plans")
+
+    def pack(self, implementation=None):
+        implementation = self.implementation if implementation is None else implementation
+        live = self.indices >= 0
+        if implementation == "owner":
+            self.owner_ids.copy_(torch.where(live, self.indices*self.world+self.rank, -1).to(torch.int32))
+            self.owner_scores.copy_(torch.where(live, self.scores, -torch.inf))
+            return self.planar
+        if implementation == "replicated":
+            self.packed[:, :, 0].copy_(torch.where(live, self.scores, -torch.inf))
+            self.packed.view(torch.int32)[:, :, 1].copy_(
+                torch.where(live, self.indices*self.world+self.rank, -1).to(torch.int32))
+            return self.packed
+        raise RuntimeError("candidate selection must bind a composition before execution")
+
+    def exchange(self, implementation=None):
+        implementation = self.implementation if implementation is None else implementation
+        self.active.copy_(self.positions >= 0)
+        if implementation == "owner":
+            return self.owner()
+        if implementation == "replicated":
+            self.gather(self.packed.view(torch.bfloat16).view(self.rows, -1))
+            self.candidates()
+            self.fold()
+            return self.output
+        raise RuntimeError("candidate selection must bind a composition before execution")
+
+
+class CandidateSelectionForward:
+    """One immutable composition used by preparation-time compiled trials."""
+
+    def __init__(self, selection, implementation):
+        self.selection, self.implementation = selection, implementation
+
+    def __call__(self):
+        self.selection.pack(self.implementation)
+        return self.selection.exchange(self.implementation)

ignition/components/owner_topk.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Distribute exact candidate selection by query row and replicate final IDs."""
+
+import torch
+import triton
+import triton.language as tl
+
+from ignition.components.candidate_rows import PreparedCandidateRows
+from ignition.components.ordered_row_topk import PreparedOrderedRowTopK
+from ignition.components.row_ownership import PreparedCandidateOwnerStage, PreparedOwnedRowBroadcast
+
+
+@triton.jit
+def _owner_candidates(Indices, ScoreBits, Active, Rank, Packed, OwnerActive,
+                      ROWS: tl.constexpr, OWNER_ROWS: tl.constexpr,
+                      WIDTH: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    column = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    global_row = tl.load(Rank) * OWNER_ROWS + row
+    live = tl.load(Active + global_row, global_row < ROWS, other=0) != 0
+    ids = tl.load(Indices + row * WIDTH + column, live & (column < WIDTH), other=-1)
+    bits = tl.load(ScoreBits + row * WIDTH + column, live & (column < WIDTH), other=-8388608)
+    tl.store(Packed + (row * WIDTH + column) * 2, bits, column < WIDTH)
+    tl.store(Packed + (row * WIDTH + column) * 2 + 1, ids, column < WIDTH)
+    if tl.program_id(1) == 0:
+        tl.store(OwnerActive + row, live)
+
+
+class PreparedOwnerTopK:
+    """Select the same rank-major candidates on one owner per query.
+
+    Input labels are already global int32 IDs, and score words are FP32.
+    Candidate preprocessing and ordered cutoff selection match the replicated
+    composition, including underfull rows, exceptional values and cutoff ties.
+    Only selected IDs are broadcast; no scores are reduced or requantized.
+    """
+
+    def __init__(self, local_indices, local_scores, active, rank, output, *,
+                 owner_runtime, reduction_runtime, session, coordinator, name, buffers, dependencies=()):
+        rows, local_topk = local_indices.shape
+        world = owner_runtime.world_size
+        if (output.ndim != 2 or output.shape[0] != rows or output.dtype != torch.int32
+                or active.shape != (rows,) or active.dtype != torch.bool
+                or rank.shape != (1,) or rank.dtype != torch.int64
+                or local_scores.shape != local_indices.shape
+                or local_indices.dtype != torch.int32 or local_scores.dtype != torch.float32
+                or reduction_runtime.world_size != world
+                or reduction_runtime.rank != owner_runtime.rank
+                or any(t.device != local_indices.device or not t.is_contiguous()
+                       for t in (local_indices, local_scores, active, rank, output))):
+            raise ValueError("owner selection requires matching candidate rows and communication ranks")
+        owner_rows, width, topk = -(-rows // world), local_topk * world, output.shape[1]
+        self.rows, self.owner_rows, self.width = rows, owner_rows, width
+        self.rank, self.active, self.output = rank, active, output
+        self.transport = PreparedCandidateOwnerStage(local_indices, local_scores,
+            runtime=owner_runtime, session=session, coordinator=coordinator, name=name + ".stage",
+            dependencies=dependencies)
+        self.packed = buffers.zeros((owner_rows, width, 2), torch.float32)
+        self.owner_active = buffers.zeros((owner_rows,), torch.bool)
+        self.global_scores = buffers.zeros((owner_rows, width), torch.float32)
+        self.global_ids = buffers.zeros((owner_rows, width), torch.int32)
+        self.lengths = buffers.zeros((owner_rows,), torch.int32)
+        self.selected_scores = buffers.zeros((owner_rows, topk), torch.float32)
+        self.selected_ids = buffers.zeros((owner_rows, topk), torch.int32)
+        self.candidates = PreparedCandidateRows(self.packed, self.owner_active,
+            self.global_scores, self.global_ids, self.lengths, topk=topk, buffers=buffers)
+        self.fold = PreparedOrderedRowTopK(self.global_scores, self.lengths, self.global_ids,
+            self.selected_scores, self.selected_ids, session=session, name=name + ".fold", buffers=buffers)
+        self.broadcast = PreparedOwnedRowBroadcast(self.selected_ids, output, rank,
+            runtime=reduction_runtime, session=session, coordinator=coordinator,
+            name=name + ".selected", buffers=buffers)
+
+    def __call__(self):
+        indices, scores = self.transport()
+        torch.library.wrap_triton(_owner_candidates)[(self.owner_rows, triton.cdiv(self.width, 1024))](
+            indices, scores.view(torch.int32), self.active, self.rank, self.packed.view(torch.int32), self.owner_active,
+            ROWS=self.rows, OWNER_ROWS=self.owner_rows, WIDTH=self.width, BLOCK=1024, num_warps=4)
+        self.candidates()
+        self.fold()
+        return self.broadcast()

ignition/components/row_ownership.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Prepare exact candidate transport and byte-preserving row-owner broadcasts."""
+
+import itertools
+import weakref
+
+import torch
+import triton
+import triton.language as tl
+from b12x.comm.pcie import _owner_preparation
+
+from ignition.components.collectives import PreparedAllReduce
+from ignition.preparation.tuning import declare
+
+_OWNERS = weakref.WeakValueDictionary()
+_HANDLES = itertools.count()
+
+
+@torch.library.custom_op("flare::candidate_owner_stage", mutates_args=("indices", "scores"))
+def _stage(local_indices: torch.Tensor, local_scores: torch.Tensor,
+           indices: torch.Tensor, scores: torch.Tensor, dependencies: list[torch.Tensor], handle: int) -> None:
+    operation = _OWNERS[handle]
+    if (indices.data_ptr() != operation.indices.data_ptr()
+            or scores.data_ptr() != operation.scores.data_ptr()):
+        raise ValueError("candidate-stage outputs must alias the reserved channel views")
+    with operation.runtime.capture(plan=operation.plan):
+        operation.runtime.stage_candidates(local_indices, local_scores, plan=operation.plan)
+
+
+@_stage.register_fake
+def _fake_stage(local_indices, local_scores, indices, scores, dependencies, handle):
+    pass
+
+
+class PreparedCandidateOwnerStage:
+    """Deliver rank-major score/ID rows to their contiguous row partitions.
+
+    Every rank declares and executes the same call. Outputs alias fixed channel
+    storage; their consumers must complete before another call reuses it.
+    """
+
+    def __init__(self, local_indices, local_scores, *, runtime, session, coordinator, name, dependencies=()):
+        self.runtime = runtime
+        self.local_indices, self.local_scores = local_indices, local_scores
+        self.dependencies = list(dependencies)
+        self.indices, self.scores = runtime.graph_candidate_views(local_indices.shape[0])
+        call = dict(local_indices=local_indices, local_scores=local_scores)
+        self.plan = _owner_preparation.plan(
+            _owner_preparation.query_from_runtime(runtime, call=call), runtime=runtime)
+        declare(session, (self.plan.request(name=name,
+            prepare_call=lambda state: _owner_preparation.prepared_call(state, **call),
+            collective=_owner_preparation.collective(runtime, key=name)),), coordinator=coordinator)
+        self.handle = next(_HANDLES)
+        _OWNERS[self.handle] = self
+
+    def __call__(self):
+        _stage(self.local_indices, self.local_scores, self.indices, self.scores, self.dependencies, self.handle)
+        return self.indices, self.scores
+
+
+@triton.jit
+def _encode_owned_rows(Source, Rank, Send, ROWS: tl.constexpr, OWNER_ROWS: tl.constexpr,
+                       ROW_BYTES: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    column = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    local_row = row - tl.load(Rank) * OWNER_ROWS
+    value = tl.full((BLOCK,), 0, tl.uint8)
+    if (local_row >= 0) & (local_row < OWNER_ROWS):
+        value = tl.load(Source + local_row * ROW_BYTES + column, column < ROW_BYTES, other=0)
+    tl.store(Send + row * ROW_BYTES + column, value.to(tl.float32), column < ROW_BYTES)
+
+
+@triton.jit
+def _decode_owned_rows(Received, Output, ROW_BYTES: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    column = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    value = tl.load(Received + row * ROW_BYTES + column, column < ROW_BYTES, other=0)
+    tl.store(Output + row * ROW_BYTES + column, value.to(tl.uint8), column < ROW_BYTES)
+
+
+class PreparedOwnedRowBroadcast:
+    """Copy contiguous row partitions to every rank without changing payload bits.
+
+    Source rows belong to the channel rank: rank r supplies rows starting at
+    r*ceil(total_rows/world). The supplied int64 rank tensor must identify that
+    rank on the device. Each byte is represented as a BF16 integer in [0,255];
+    all other ranks contribute zero. Reduction is therefore exact, including
+    payloads that represent negative integers, signed zero or NaNs.
+    """
+
+    def __init__(self, source, output, rank, *, runtime, session, coordinator, name, buffers=None):
+        world = 1 if runtime is None else runtime.world_size
+        if (source.ndim != 2 or output.ndim != 2 or min(output.shape) <= 0
+                or source.shape != (-(-output.shape[0] // world), output.shape[1])
+                or source.dtype != output.dtype or source.device != output.device
+                or not source.is_contiguous() or not output.is_contiguous()
+                or rank.shape != (1,) or rank.dtype != torch.int64 or rank.device != source.device):
+            raise ValueError("owned-row broadcast requires contiguous matching row partitions and an int64 group rank")
+        self.source, self.output, self.rank = source, output, rank
+        self.rows, self.owner_rows = output.shape[0], source.shape[0]
+        self.row_bytes = output.shape[1] * output.element_size()
+        if self.row_bytes % 16:
+            raise ValueError("owned-row payload bytes must be a multiple of 16")
+        def allocate():
+            shape = (self.rows, self.row_bytes)
+            return (torch.empty(shape, dtype=torch.bfloat16, device=source.device) if buffers is None
+                    else buffers.zeros(shape, torch.bfloat16))
+        self.send, self.received = allocate(), allocate()
+        self.reduction = PreparedAllReduce(self.send, self.received, runtime=runtime,
+            session=session, coordinator=coordinator, name=name)
+
+    def __call__(self):
+        grid = (self.rows, triton.cdiv(self.row_bytes, 1024))
+        torch.library.wrap_triton(_encode_owned_rows)[grid](
+            self.source.view(torch.uint8), self.rank, self.send,
+            ROWS=self.rows, OWNER_ROWS=self.owner_rows, ROW_BYTES=self.row_bytes, BLOCK=1024, num_warps=4)
+        self.reduction(self.send)
+        torch.library.wrap_triton(_decode_owned_rows)[grid](
+            self.received, self.output.view(torch.uint8), ROW_BYTES=self.row_bytes, BLOCK=1024, num_warps=4)
+        return self.output