From ce4729f5602341363f12aebe5b0ede2734eaa52f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 26 Oct 2026 04:42:36 +0000 Subject: [PATCH] Fuse DSpark draft scoring and sparse support restoration Read strided BF16/FP32 draft-head rows and optional bias directly with FP32 addition, then combine policy scoring and local candidate generation. Emit canonical support bytes during selection and restore only locally written support after each broadcast. Preserve score inspection and selector-dependent materialization, proposal probabilities, RNG reservations and rejection inputs. Implemented on the public dependency and stable-toolchain prerequisites. Markov projections remain normally prepared; target arithmetic, graph families and native metadata are unchanged. Re-export packages to capture the fusion. Qualified here with 66 CPU contracts. Historical TP2 native proposal-pipeline checks retain byte-exact dense q and sampled drafts with stable allocations; their timings exclude model forwards, Markov GEMM and target verification. GPU/native qualification of this separated branch remains pending; no isolated HTTP throughput gain is claimed. Design-Signed-Off-By: Codex (GPT-6); design-review=babb --- ignition/components/draft_sampling.py | 176 +++++++++++++--- ignition/components/parallel_draft.py | 26 ++- .../test_bound_draft_scores_cuda.py | 196 ++++++++++++++++++ tests/components/test_draft_support_cuda.py | 117 +++++++++++ .../test_sampling_broadcast_cuda.py | 28 ++- 5 files changed, 510 insertions(+), 33 deletions(-) create mode 100644 tests/components/test_bound_draft_scores_cuda.py create mode 100644 tests/components/test_draft_support_cuda.py diff --git a/ignition/components/draft_sampling.py b/ignition/components/draft_sampling.py index 93912c8b..036d588b 100644 --- a/ignition/components/draft_sampling.py +++ b/ignition/components/draft_sampling.py @@ -81,7 +81,7 @@ def _block_candidates(Scores, Probabilities, TopK, Candidates, V: tl.constexpr, @triton.jit def _select_and_sample(Candidates, Selected, Probabilities, Valid, TopK, TopP, Params, Slots, Seed, Offset, Sampled, V: tl.constexpr, BLOCKS: tl.constexpr, POOL: tl.constexpr, K: tl.constexpr, - SLOTS: tl.constexpr): + SLOTS: tl.constexpr, SupportIds, SupportProbabilities, WRITE_SUPPORT: tl.constexpr): row = tl.program_id(0) index = tl.arange(0, POOL) pool = tl.load(Candidates + row * BLOCKS * K + index, index < BLOCKS * K, other=-9223372036854775808) @@ -138,20 +138,52 @@ def _select_and_sample(Candidates, Selected, Probabilities, Valid, TopK, TopP, P tl.store(Probabilities + row * V + sampled, 1., ~valid) tl.store(Valid + row, valid) tl.store(Sampled + row, sampled.to(tl.int64)) + if WRITE_SUPPORT: + # Keep the original Selected order and exact dense-q bits, including padding at token zero. + tl.debug_barrier() + local_keys = tl.load(Selected + row * K + rank) + support_ids = tl.where(local_keys != -9223372036854775808, + 2147483647 - (local_keys & 2147483647).to(tl.int32), 0) + support_probabilities = tl.load(Probabilities + row * V + support_ids) + tl.store(SupportIds + row * K + rank, tl.where(valid, support_ids, 0)) + tl.store(SupportProbabilities + row * K + rank, tl.where(valid, support_probabilities, 1.)) @triton.jit -def _draft_scores(Logits, Histogram, Candidates, Slots, Params, GrammarMask, GrammarEnabled, Scores, BlockMax, - hist_stride, mask_stride, candidate_stride, CAPACITY: tl.constexpr, V: tl.constexpr, - PREFIX: tl.constexpr, BLOCK: tl.constexpr, BLOCKS: tl.constexpr): +def _restore_draft_support(Selected, SupportIds, SupportProbabilities, Probabilities, + V: tl.constexpr, K: tl.constexpr): + row = tl.program_id(0) + rank = tl.arange(0, K) + keys = tl.load(Selected + row * K + rank) + local_ids = 2147483647 - (keys & 2147483647).to(tl.int32) + # Candidate generation zeroed the dense row; selection only wrote its local support or token zero. + tl.store(Probabilities + row * V + local_ids, 0., + (keys != -9223372036854775808) & (local_ids > 0) & (local_ids < V)) + tl.store(Probabilities + row * V, 0.) + tl.debug_barrier() + ids = tl.load(SupportIds + row * K + rank) + values = tl.load(SupportProbabilities + row * K + rank) + # Padding repeats zero with identical bits. Write it once without floating-point arithmetic. + zero_bits = tl.max(tl.where(ids == 0, values.to(tl.uint32, bitcast=True), 0), 0) + tl.store(Probabilities + row * V, zero_bits.to(tl.float32, bitcast=True)) + tl.store(Probabilities + row * V + ids, values, ids != 0) + + +@triton.jit +def _draft_score_values(Logits, Bias, Histogram, Candidates, Slots, Params, GrammarMask, GrammarEnabled, + hist_stride, mask_stride, candidate_stride, logit_stride, bias_stride, + ADD_BIAS: tl.constexpr, CAPACITY: tl.constexpr, V: tl.constexpr, + PREFIX: tl.constexpr, BLOCK: tl.constexpr): # One vocabulary block of one row: the committed histogram plus the row's earlier drafts give the penalty # counts; repetition, presence and frequency penalties, temperature and the grammar mask follow the - # unfused order of operations. Each block records its maximum. + # unfused order of operations. The caller reduces each block's maximum. row, block = tl.program_id(0), tl.program_id(1) safe = tl.minimum(tl.maximum(tl.load(Slots + row).to(tl.int64), 0), CAPACITY - 1) token = block * BLOCK + tl.arange(0, BLOCK) inside = token < V - score = tl.load(Logits + row * V + token, mask=inside, other=0.).to(tl.float32) + score = tl.load(Logits + row * logit_stride + token, mask=inside, other=0.).to(tl.float32) + if ADD_BIAS: + score = score + tl.load(Bias + row * bias_stride + token, mask=inside, other=0.).to(tl.float32) temperature = tl.load(Params + safe * 6) repetition = tl.load(Params + safe * 6 + 3) presence = tl.load(Params + safe * 6 + 4) @@ -173,12 +205,48 @@ def _draft_scores(Logits, Histogram, Candidates, Slots, Params, GrammarMask, Gra if tl.load(GrammarEnabled + safe) != 0: word = tl.load(GrammarMask + safe * mask_stride + token // 32, mask=inside, other=0) score = tl.where(((word >> (token % 32)) & 1) != 0, score, -float("inf")) + return token, inside, score + + +@triton.jit +def _draft_scores(Logits, Histogram, Candidates, Slots, Params, GrammarMask, GrammarEnabled, Scores, BlockMax, + hist_stride, mask_stride, candidate_stride, CAPACITY: tl.constexpr, V: tl.constexpr, + PREFIX: tl.constexpr, BLOCK: tl.constexpr, BLOCKS: tl.constexpr): + token, inside, score = _draft_score_values(Logits, Logits, Histogram, Candidates, Slots, Params, + GrammarMask, GrammarEnabled, hist_stride, mask_stride, candidate_stride, V, V, + False, CAPACITY, V, PREFIX, BLOCK) + row, block = tl.program_id(0), tl.program_id(1) tl.store(Scores + row * V + token, score, mask=inside) score = tl.where(inside, score, -float("inf")) maximum = tl.max(score, 0) tl.store(BlockMax + row * BLOCKS + block, maximum) +@triton.jit +def _draft_score_candidates(Logits, Bias, Histogram, PrefixIds, Slots, Params, GrammarMask, GrammarEnabled, + TopK, Probabilities, LocalCandidates, BlockMax, hist_stride, mask_stride, + candidate_stride, logit_stride, bias_stride, ADD_BIAS: tl.constexpr, + CAPACITY: tl.constexpr, V: tl.constexpr, PREFIX: tl.constexpr, + BLOCK: tl.constexpr, BLOCKS: tl.constexpr, K: tl.constexpr): + tokens, inside, scores = _draft_score_values(Logits, Bias, Histogram, PrefixIds, Slots, Params, + GrammarMask, GrammarEnabled, hist_stride, mask_stride, candidate_stride, logit_stride, + bias_stride, ADD_BIAS, CAPACITY, V, PREFIX, BLOCK) + row, block = tl.program_id(0), tl.program_id(1) + # Validity depends on every score, including nonfinite entries outside the selected support. + tl.store(BlockMax + row * BLOCKS + block, tl.max(tl.where(inside, scores, -float("inf")), 0)) + tl.store(Probabilities + row * V + tokens, tl.zeros((BLOCK,), tl.float32), inside) + safe = tl.minimum(tl.maximum(tl.load(Slots + row).to(tl.int64), 0), CAPACITY - 1) + top_k = tl.load(TopK + safe) + k = tl.minimum(tl.maximum(tl.where(top_k > 0, top_k, V), 1), K) + ordered = tl.where(inside, _orderable(scores), -2147483648) + position, take, count = _candidate_positions(ordered, _kth_largest(ordered, k), inside, K) + keys = (ordered.to(tl.int64) << 32) | (2147483647 - tokens).to(tl.int64) + base = LocalCandidates + (row * BLOCKS + block) * K + slots = tl.arange(0, K) + tl.store(base + slots, tl.full((K,), -9223372036854775808, tl.int64), slots >= count) + tl.store(base + position, keys, take) + + @triton.jit def _selected_row(Slots, Generations, Generation, live_rows, row, CAPACITY: tl.constexpr, BATCH: tl.constexpr): slot = tl.load(Slots + row).to(tl.int64) @@ -217,7 +285,7 @@ def _draft_scores_finish(BlockMax, Candidates, Slots, Generations, LiveRows, Gen tl.store(PackedTopP + row, tl.load(Params + safe * 6 + 1)) -_SCORE_BLOCK = 4096 +_SCORE_BLOCK = BLOCK class PreparedDraftProposal: @@ -243,14 +311,32 @@ def __init__(self, scores, probabilities, valid, top_k, top_p, params, slots, se seed, offset, sampled, vocab, self.blocks, pool, DRAFT_CANDIDATES, capacity) self._grids = ((rows, self.blocks, 1), (rows, 1, 1)) self._kernels = None + self._support = None + + def bind_support(self, ids, probabilities): + """Bind canonical broadcast payloads before graph capture, preserving candidate order.""" + shape = (self.selected.shape[0], DRAFT_CANDIDATES) + for tensor, dtype in ((ids, torch.int64), (probabilities, torch.float32)): + if (tensor.shape != shape or tensor.dtype != dtype or tensor.device != self.selected.device + or not tensor.is_contiguous()): + raise ValueError("draft support requires contiguous Int64 IDs and FP32 probabilities") + self._support = ids, probabilities + if self._kernels is not None: + self._kernels = self._kernels[:-1] + (self._prepare_select(),) + + def _prepare_select(self): + from ignition.preparation.triton import prepare_triton - def prepare(self): + support = (self.selected, self._block_args[1]) if self._support is None else self._support + return prepare_triton(_select_and_sample, self._select_args + support, self._grids[1], + WRITE_SUPPORT=self._support is not None, num_warps=4) + + def prepare(self, *, candidates_prepared=False): from ignition.preparation.triton import prepare_triton - self._kernels = ( + self._kernels = (() if candidates_prepared else ( prepare_triton(_block_candidates, self._block_args, self._grids[0], num_warps=8), - prepare_triton(_select_and_sample, self._select_args, self._grids[1], num_warps=4), - ) + )) + (self._prepare_select(),) self() def __call__(self): @@ -292,6 +378,23 @@ def __init__(self, logits, requests, slots, generations, live_rows, candidate_pr candidates=self.proposal_candidates, selected=self.proposal_selected) self.sampled = sampled self.broadcast = None + self.source_logits, self.source_bias = None, None + self._candidate_kernels = None + self._restore_support = None + + def bind_logits(self, logits, bias=None): + """Read strided BF16/FP32 head rows and bias with FP32 addition, without staging logits.""" + if self._prepared: + raise RuntimeError("draft logits must be bound before prepare()") + shape = (self.logits.shape[0], self.logits.shape[-1]) + if (self.logits.dtype != torch.float32 or logits.shape != shape + or logits.device != self.logits.device or logits.dtype not in (torch.bfloat16, torch.float32) + or logits.stride(-1) != 1 or logits.stride(0) < shape[1]): + raise ValueError("bound draft logits require FP32 staging and BF16/FP32 contiguous vocabulary rows") + if bias is not None and (bias.shape != shape or bias.dtype not in (torch.bfloat16, torch.float32) + or bias.device != logits.device or bias.stride(-1) != 1 or bias.stride(0) < shape[1]): + raise ValueError("draft bias requires BF16/FP32 contiguous vocabulary rows on the logit device") + self.source_logits, self.source_bias = logits, bias def prepare_broadcast(self, *, runtime, session, coordinator, name, buffers=None): if runtime is None: @@ -305,9 +408,19 @@ def allocate(shape, dtype): layout = self.broadcast_layout(self.logits.shape[0]) self.support_ids = allocate(*layout["support_ids"]) self.support_probabilities = allocate(*layout["support_probabilities"]) + self.proposal.bind_support(self.support_ids, self.support_probabilities) self.broadcast = PreparedBroadcast( (self.support_ids, self.support_probabilities, self.sampled, self.valid), runtime=runtime, session=session, coordinator=coordinator, name=name, buffers=buffers) + if self._prepared: + self._prepare_support_restore() + + def _prepare_support_restore(self): + from ignition.preparation.triton import prepare_triton + + self._restore_support = prepare_triton(_restore_draft_support, + (self.proposal_selected, self.support_ids, self.support_probabilities, self.probabilities), + (self.logits.shape[0], 1, 1), V=self.logits.shape[-1], K=DRAFT_CANDIDATES, num_warps=4) def prepare(self): if self.logits.device.type != "cuda": @@ -330,31 +443,46 @@ def prepare(self): self.packed_top_p, self.candidates.stride(0)), (batch, 1, 1), CAPACITY=state.capacity, BATCH=batch, V=vocab, PREFIX=self.prefix_offset, BLOCKS=blocks, BLOCK_COUNT=triton.next_power_of_2(blocks), RUNNING=RUNNING, num_warps=1)) + if self.source_logits is not None: + bias = self.source_logits if self.source_bias is None else self.source_bias + self._candidate_kernels = ( + prepare_triton(_draft_score_candidates, (self.source_logits, bias, state.histogram, + self.candidates, self.slots, state.parameters, state.grammar_mask, state.grammar_enabled, + state.top_k, self.probabilities, self.proposal_candidates, self.block_max, + state.histogram.stride(0), state.grammar_mask.stride(0), self.candidates.stride(0), + self.source_logits.stride(0), bias.stride(0)), (batch, blocks, 1), + ADD_BIAS=self.source_bias is not None, CAPACITY=state.capacity, V=vocab, + PREFIX=self.prefix_offset, BLOCK=BLOCK, BLOCKS=blocks, K=DRAFT_CANDIDATES, + num_warps=8, enable_fp_fusion=False), self._score_kernels[1]) self._prepared = True self.prepare_scores() - self.proposal.prepare() + if self._candidate_kernels is not None: + for kernel in self._candidate_kernels: + kernel() + self.proposal.prepare(candidates_prepared=self._candidate_kernels is not None) + if self.broadcast is not None: + self._prepare_support_restore() def prepare_scores(self): """One-position scores in two kernels, without materialized penalty counts or a separate maximum.""" + if self.source_logits is not None: + # Explicit score inspection retains the ordinary materialized-logit contract. + self.logits[:, 0].copy_(self.source_logits.float()) + if self.source_bias is not None: + self.logits[:, 0].add_(self.source_bias) for kernel in self._score_kernels: kernel() def __call__(self): if not self._prepared: raise RuntimeError("draft proposal requires prepare() before capture") - self.prepare_scores() + if self._candidate_kernels is None: + self.prepare_scores() + else: + for kernel in self._candidate_kernels: + kernel() self.proposal() if self.broadcast is not None: - keys = self.proposal_selected - present = keys != -9223372036854775808 - ids = 2147483647 - (keys & 2147483647) - self.support_ids.copy_(torch.where(present, ids, 0)) - probabilities = self.probabilities[:, 0].gather(1, self.support_ids) - # Padding repeats token zero with its exact probability, so duplicate stores agree. - # Invalid proposals use the sampler's point mass at token zero. - self.support_probabilities.copy_(torch.where(self.valid, probabilities, 1.)) - self.support_ids.masked_fill_(~self.valid, 0) self.broadcast() - self.probabilities.zero_() - self.probabilities[:, 0].scatter_(1, self.support_ids, self.support_probabilities) + self._restore_support() return self.probabilities diff --git a/ignition/components/parallel_draft.py b/ignition/components/parallel_draft.py index 68396f62..ef2a1e58 100644 --- a/ignition/components/parallel_draft.py +++ b/ignition/components/parallel_draft.py @@ -218,7 +218,13 @@ def prepare(self): ) if target is not None: target.prepare() - for proposal in proposals: + for index, proposal in enumerate(proposals): + bias = None if self.markov is None else getattr(self.markov, "logits", None) + if (self.draft is not None and self.selector is None + and (self.markov is None or bias is not None)): + block = self.draft_inputs.shape[1] + proposal.bind_logits(self.draft.scratch.logits.view(b, block, -1)[:, index + o.logit_offset], + bias) proposal.prepare() if self.proposal_only: prepare_sampling_broadcasts(proposals, None, collectives=o.collectives, session=o.session, @@ -305,13 +311,17 @@ def propose(self, advance_rng=True, observer=None, allowed=None): self.probabilities.zero_() self.probabilities.scatter_(2, self.drafts[:, :, None], 1.0) for i, proposal in enumerate(proposals): - self.logits[:, 0].copy_( - self.draft.scratch.logits.view(b, block, -1)[:, i + o.logit_offset].float() - ) - if self.markov is not None: - self.logits[:, 0].add_(self.markov()) - if self.selector is not None: - self.selector(self.logits[:, 0], self.markov_token, i + o.logit_offset) + if proposal.source_logits is not None: + if self.markov is not None: + self.markov() + else: + self.logits[:, 0].copy_( + self.draft.scratch.logits.view(b, block, -1)[:, i + o.logit_offset].float() + ) + if self.markov is not None: + self.logits[:, 0].add_(self.markov()) + if self.selector is not None: + self.selector(self.logits[:, 0], self.markov_token, i + o.logit_offset) if observer is not None: observer.record_embedding(i, self.markov.embedding) self.offset.copy_(self.base_offset + 4 * i) diff --git a/tests/components/test_bound_draft_scores_cuda.py b/tests/components/test_bound_draft_scores_cuda.py new file mode 100644 index 00000000..b0a10e84 --- /dev/null +++ b/tests/components/test_bound_draft_scores_cuda.py @@ -0,0 +1,196 @@ +"""Direct draft-head/bias inputs preserve the canonical proposal and RNG outputs.""" +import os + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not (os.environ.get('FLARE_RUN_CUDA_TESTS') and torch.cuda.is_available()), + reason='set FLARE_RUN_CUDA_TESTS=1 on a CUDA host') + + +@pytest.mark.parametrize('batch,vocab,dtype,compiled,prefix', [ + (1, 17, torch.float32, False, 0), + (4, 10000, torch.bfloat16, True, 3), + (32, 129280, torch.bfloat16, True, 7), + (64, 257, torch.float32, False, 7), +]) +@pytest.mark.parametrize('bias_dtype', [None, torch.bfloat16, torch.float32]) +@torch.no_grad() +def test_bound_inputs_match_canonical_native_replay(batch, vocab, dtype, compiled, prefix, bias_dtype): + from ignition.components.draft_sampling import PreparedDraftDistribution + from ignition.components.sampling import DeviceRequestState, RUNNING + from ignition.preparation.compilation import compile_forward + from ignition.preparation.graphs import PreparedGraphBundle + + generator = torch.Generator().manual_seed(881 + prefix) + capacity = max(4, batch) + state = DeviceRequestState.allocate('cuda:0', vocab_size=vocab, capacity=capacity) + state.generation.fill_(7) + state.status.fill_(RUNNING) + state.histogram.copy_(torch.randint(0, 4, state.histogram.shape, generator=generator, dtype=torch.int32)) + state.parameters.copy_(torch.tensor([1., 1., 0., 1., 0., 0.], device='cuda')) + state.top_k.zero_() + slots = torch.arange(batch, dtype=torch.int32, device='cuda') + epochs = torch.full((batch,), 7, dtype=torch.int64, device='cuda') + live = torch.tensor([batch], dtype=torch.int32, device='cuda') + candidates = torch.tensor([[7, 7, 0, 3, 4, 5, 6]] * batch, device='cuda') + head = torch.randn((batch, 8, vocab), generator=generator).to(device='cuda', dtype=dtype) + source = head[:, 2] + bias = (torch.randn((batch, vocab), generator=generator).to(device='cuda', dtype=bias_dtype) + if bias_dtype is not None else None) + head_input = head.clone() + bias_input = None if bias is None else bias.clone() + source_input = head_input[:, 2] + saved_head, saved_bias = head_input.clone(), None if bias_input is None else bias_input.clone() + seed = torch.arange(batch, dtype=torch.int64, device='cuda') + 931 + offset = torch.zeros_like(seed) + operations, calls, graphs = [], [], {} + for bound in (False, True): + logits = torch.zeros((batch, 1, vocab), device='cuda') + op = PreparedDraftDistribution(logits, state, slots, epochs, live, candidates, + seed=seed, offset=offset, sampled=torch.zeros_like(seed), prefix_offset=prefix) + if bound: + op.bind_logits(source, bias) + op.prepare() + def run(op=op, bound=bound): + # The producer writes the full tensor; the fused reader holds its pre-created strided view. + head.copy_(head_input) + if bias is not None: + bias.copy_(bias_input) + if not bound: + op.logits[:, 0].copy_(source.float()) + if bias is not None: + op.logits[:, 0].add_(bias) + op() + call = compile_forward(run) if compiled else run + call() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + operations.append(op) + calls.append(call) + graphs[str(bound)] = graph + bundle = PreparedGraphBundle(device=0, graphs=graphs, buffers={}, resources=(operations, calls, state, head, bias, head_input, bias_input)) + worker = bundle.start_worker() + outputs = lambda op: (op.probabilities, op.sampled, op.valid, op.maximum, + op.packed_top_k, op.packed_top_p, op.proposal_candidates, op.proposal_selected) + expected = [torch.empty_like(tensor) for tensor in outputs(operations[0])] + modes = ('neutral', 'penalties', 'grammar', 'ties', 'signed_zero', 'tiny_temperature', + 'nonfinite', 'invalid_prefix', 'inactive', 'duplicate_or_stale', 'rng_wrap', 'neutral') + try: + for phase, mode in enumerate(modes): + head_input.copy_(saved_head) + if bias is not None: + bias_input.copy_(saved_bias) + state.status.fill_(RUNNING) + state.parameters.copy_(torch.tensor([1., 1., 0., 1., 0., 0.], device='cuda')) + state.top_k.copy_(torch.tensor([(-1, 0, 1, min(20, vocab), min(64, vocab))[i % 5] + for i in range(capacity)], dtype=torch.int32, device='cuda')) + state.grammar_enabled.zero_() + slots.copy_(torch.arange(batch, dtype=torch.int32, device='cuda')) + epochs.fill_(7) + live.fill_(batch) + candidates.copy_(torch.tensor([[7, 7, 0, 3, 4, 5, 6]] * batch, device='cuda')) + expected_offset = phase * 4 + (2**32 if mode == 'rng_wrap' else 0) + offset.fill_(expected_offset) + if mode == 'penalties': + state.parameters.copy_(torch.tensor([.7, .8, .1, 1.3, -.2, -.1], device='cuda')) + if mode == 'grammar': + state.grammar_enabled.fill_(1) + state.grammar_mask.copy_(torch.randint(-2**31, 2**31-1, state.grammar_mask.shape, + generator=generator, dtype=torch.int64).to(torch.int32)) + if mode in ('ties', 'signed_zero'): + source_input.zero_() + source_input[:, 1::2] = -0. if mode == 'signed_zero' else 1. + if bias is not None: + bias_input.zero_() + if mode == 'signed_zero': + state.parameters[:, 4] = -0. + state.parameters[:, 5] = -0. + if mode == 'tiny_temperature': + state.parameters[:, 0] = torch.finfo(torch.float32).tiny + if mode == 'nonfinite': + source_input[:, :3] = torch.tensor([float('nan'), float('inf'), -float('inf')], device='cuda') + if mode == 'invalid_prefix': + candidates[:, 0] = -1 + if mode == 'inactive': + live.fill_(batch // 2) + if mode == 'duplicate_or_stale': + epochs[-1] = 99 + if batch > 1: + slots[-1] = slots[0] + for index, op in enumerate(operations): + op.probabilities.fill_(float('nan')) + op.block_max.fill_(float('nan')) + op.proposal_candidates.fill_(71) + op.proposal_selected.fill_(73) + torch.cuda.synchronize() + allocations = torch.cuda.memory_stats()['allocation.all.allocated'] + ticket = worker.submit(index) + completed = worker.poll(timeout_ms=30000) + assert completed is not None and completed.ticket == ticket and completed.error is None + assert torch.cuda.memory_stats()['allocation.all.allocated'] == allocations + if index == 0: + for destination, tensor in zip(expected, outputs(op), strict=True): + destination.copy_(tensor) + else: + for actual, reference in zip(outputs(op), expected, strict=True): + if actual is op.maximum: + torch.testing.assert_close(actual, reference, rtol=0, atol=0, equal_nan=True) + else: + assert torch.equal(actual, reference), (mode, actual.shape) + # The explicit score-inspection API remains materialized and bit exact. + for op in operations: + op.prepare_scores() + valid = operations[0].valid[:, 0] + assert torch.equal(operations[0].scores[valid].view(torch.int32), + operations[1].scores[valid].view(torch.int32)) + assert torch.equal(offset, torch.full_like(offset, expected_offset)) + finally: + worker.close() + for graph in graphs.values(): + graph.reset() + + +@pytest.mark.parametrize('bias_dtype', [None, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize('selector', [False, True]) +@torch.no_grad() +def test_parallel_draft_preparation_preserves_bias_and_selector_contract(bias_dtype, selector): + from types import SimpleNamespace + from ignition.components.parallel_draft import PreparedParallelDraftRound + from ignition.components.sampling import DeviceRequestState, RUNNING + + batch, width, vocab, block = 2, 3, 257, 4 + state = DeviceRequestState.allocate('cuda:0', vocab_size=vocab, capacity=4) + state.generation.fill_(7) + state.status.fill_(RUNNING) + state.parameters.copy_(torch.tensor([1., 1., 0., 1., 0., 0.], device='cuda')) + inputs = SimpleNamespace(row_slot_ids=torch.arange(batch, dtype=torch.int32, device='cuda'), + row_generations=torch.full((batch,), 7, dtype=torch.int64, device='cuda'), + live_rows=torch.tensor([batch], dtype=torch.int32, device='cuda')) + operation = PreparedParallelDraftRound.__new__(PreparedParallelDraftRound) + operation.owner = SimpleNamespace(requests=state, inputs=inputs, logit_offset=1, session=None, + collectives=dict(vocabulary_runtime=None, coordinator=None)) + operation.batch, operation.width, operation.name = batch, width, 'bias-preparation' + operation.proposal_only, operation.generate_proposals, operation.argmax = True, True, None + operation.operations, operation.buffers, operation.sampling_scratch = None, None, {} + operation.draft_inputs = torch.zeros((batch, block), dtype=torch.int64, device='cuda') + operation.draft = SimpleNamespace(scratch=SimpleNamespace( + logits=torch.zeros((batch * block, vocab), dtype=torch.bfloat16, device='cuda'))) + operation.markov = None if bias_dtype is None else SimpleNamespace( + logits=torch.zeros((batch, vocab), dtype=bias_dtype, device='cuda')) + operation.selector = object() if selector else None + operation.logits = torch.zeros((batch, 1, vocab), device='cuda') + operation.drafts = torch.zeros((batch, width - 1), dtype=torch.int64, device='cuda') + operation.seed = torch.zeros(batch, dtype=torch.int64, device='cuda') + operation.offset, operation.sampled = torch.zeros_like(operation.seed), torch.zeros_like(operation.seed) + operation.prepare() + for index, proposal in enumerate(operation.operations[1]): + if selector: + assert proposal.source_logits is None + else: + assert proposal.source_logits.data_ptr() == operation.draft.scratch.logits.view(batch, block, vocab)[:, index + 1].data_ptr() + assert proposal.source_logits.stride(0) == block * vocab + assert proposal.source_bias is (None if operation.markov is None else operation.markov.logits) + proposal() + assert torch.isfinite(proposal.probabilities).all() diff --git a/tests/components/test_draft_support_cuda.py b/tests/components/test_draft_support_cuda.py new file mode 100644 index 00000000..4db4a581 --- /dev/null +++ b/tests/components/test_draft_support_cuda.py @@ -0,0 +1,117 @@ +"""Sparse broadcast support preserves dense proposal bytes and clears divergent local support.""" +import os + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not (os.environ.get('FLARE_RUN_CUDA_TESTS') and torch.cuda.is_available()), + reason='set FLARE_RUN_CUDA_TESTS=1 on a CUDA host') + + +@pytest.mark.parametrize('batch,vocab', [(1, 257), (4, 129280)]) +@torch.no_grad() +def test_support_emission_and_sparse_restore_native_replay(batch, vocab): + from ignition.components.draft_sampling import PreparedDraftDistribution, DRAFT_CANDIDATES + from ignition.components.sampling import DeviceRequestState, RUNNING + from ignition.preparation.graphs import PreparedGraphBundle + + state = DeviceRequestState.allocate('cuda:0', vocab_size=vocab, capacity=4) + state.generation.fill_(7) + state.status.fill_(RUNNING) + state.parameters.copy_(torch.tensor([1., 1., 0., 1., 0., 0.], device='cuda')) + slots = torch.arange(batch, dtype=torch.int32, device='cuda') + epochs = torch.full((batch,), 7, dtype=torch.int64, device='cuda') + live = torch.tensor([batch], dtype=torch.int32, device='cuda') + candidates = torch.zeros((batch, 7), dtype=torch.int64, device='cuda') + logits = torch.randn((batch, 1, vocab), device='cuda') + seed = torch.arange(batch, dtype=torch.int64, device='cuda') + offset = torch.zeros(batch, dtype=torch.int64, device='cuda') + op = PreparedDraftDistribution(logits, state, slots, epochs, live, candidates, + seed=seed, offset=offset, + sampled=torch.zeros(batch, dtype=torch.int64, device='cuda')) + op.prepare() + reference = PreparedDraftDistribution(logits, state, slots, epochs, live, candidates, + seed=seed, offset=offset, sampled=torch.zeros_like(seed)) + reference.prepare() + op.support_ids = torch.zeros((batch, DRAFT_CANDIDATES), dtype=torch.int64, device='cuda') + op.support_probabilities = torch.zeros((batch, DRAFT_CANDIDATES), device='cuda') + op.proposal.bind_support(op.support_ids, op.support_probabilities) + op._prepare_support_restore() + for mode in ('ordinary', 'ties', 'all_invalid', 'filtered', 'token_zero_padding'): + logits.normal_() + state.top_k.zero_() + state.parameters[:, 1] = 1. + if mode == 'ties': + logits.zero_() + if mode == 'all_invalid': + logits.fill_(float('nan')) + if mode == 'filtered': + state.top_k.fill_(3) + state.parameters[:, 1] = .7 + if mode == 'token_zero_padding': + state.top_k.fill_(3) + logits.fill_(-10.) + logits[:, :, 0] = 2. + logits[:, :, 4] = 1. + logits[:, :, 7] = .5 + reference() + op() + assert torch.equal(op.probabilities.view(torch.int32), reference.probabilities.view(torch.int32)) + assert torch.equal(op.sampled, reference.sampled) and torch.equal(op.valid, reference.valid) + keys = op.proposal_selected + expected_ids = torch.where(keys != -9223372036854775808, 2147483647-(keys & 2147483647), 0) + expected_probs = op.probabilities[:, 0].gather(1, expected_ids) + expected_probs = torch.where(op.valid, expected_probs, 1.) + expected_ids = expected_ids.masked_fill(~op.valid, 0) + assert torch.equal(op.support_ids, expected_ids) + assert torch.equal(op.support_probabilities.view(torch.int32), expected_probs.view(torch.int32)) + + root_ids = torch.zeros_like(op.support_ids) + root_probabilities = torch.zeros_like(op.support_probabilities) + expected = torch.zeros_like(op.probabilities) + def receive_and_restore(): + op() + # A delayed, divergent canonical payload arrives after local sampling has written q. + torch.cuda._sleep(10000) + op.support_ids.copy_(root_ids) + op.support_probabilities.copy_(root_probabilities) + op._restore_support() + receive_and_restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + receive_and_restore() + bundle = PreparedGraphBundle(device=0, graphs={'support': graph}, buffers={}, resources=(op, state, root_ids, root_probabilities)) + worker = bundle.start_worker() + try: + for phase, mode in enumerate(('ordinary', 'local_invalid', 'root_invalid', 'signed_zero', 'subnormal', 'ordinary')): + logits.normal_() + state.top_k.zero_() + state.parameters[:, 1] = 1. + if mode == 'local_invalid': + logits.fill_(float('nan')) + root_ids.zero_() + root_ids[:, :4] = torch.tensor([0, 4, 7, 11], device='cuda') + root_probabilities.zero_() + root_probabilities[:, :4] = torch.tensor([.1, .2, .3, .4], device='cuda') + if mode == 'root_invalid': + root_ids.zero_() + root_probabilities.fill_(1.) + else: + if mode == 'signed_zero': + root_probabilities[:, 0] = -0. + if mode == 'subnormal': + root_probabilities[:, 0].view(torch.int32).fill_(1) + root_probabilities[:, 4:] = root_probabilities[:, :1] + expected.zero_() + expected[:, 0].scatter_(1, root_ids, root_probabilities) + op.probabilities.fill_(float('nan')) + torch.cuda.synchronize() + allocations = torch.cuda.memory_stats()['allocation.all.allocated'] + ticket = worker.submit(0, repetitions=3) + completion = worker.poll(timeout_ms=30000) + assert completion is not None and completion.ticket == ticket and completion.error is None + assert torch.cuda.memory_stats()['allocation.all.allocated'] == allocations + assert torch.equal(op.probabilities.view(torch.int32), expected.view(torch.int32)), mode + finally: + worker.close() + graph.reset() diff --git a/tests/components/test_sampling_broadcast_cuda.py b/tests/components/test_sampling_broadcast_cuda.py index 3178536a..4cdee4f3 100644 --- a/tests/components/test_sampling_broadcast_cuda.py +++ b/tests/components/test_sampling_broadcast_cuda.py @@ -36,7 +36,7 @@ def results(self): self.verify.state.seed.view(torch.int64), self.verify.state.rng_counter.view(torch.int64)) -def _rank(rank, world, rendezvous, result_channel=False): +def _rank(rank, world, rendezvous, result_channel=False, divergent_invalid=False): from ignition.components.draft_sampling import PreparedDraftDistribution from ignition.components.parallel_draft import prepare_sampling_broadcasts, prepare_sampling_scratch from ignition.components.sampling import DeviceRequestState, RequestAdmission @@ -95,6 +95,23 @@ def _rank(rank, world, rendezvous, result_channel=False): proposals = tuple(PreparedDraftDistribution(logits, state, slots[:batch], epochs[:batch], live, ids, prefix_offset=index, seed=seed, offset=offset, sampled=sampled, workspace=workspace["proposal"]) for index in range(width - 1)) + if divergent_invalid: + if rank == 0: + logits[0].fill_(float("nan")) + if batch > 1: + logits[1].zero_() + if batch > 2: + logits[2].fill_(-10.) + logits[2, :, 0] = 2. + state.top_k[2] = 3 + else: + if batch > 1: + logits[1].fill_(float("nan")) + if batch > 2: + logits[2].fill_(-10.) + logits[2, :, 11] = 2. + for proposal in proposals: + proposal.bind_logits(logits[:, 0]) for proposal in proposals: proposal.prepare() target = torch.softmax(torch.randn((batch, width, 257), @@ -185,6 +202,15 @@ def test_native_sampling_result_channel_catalog(tmp_path, world): mp.spawn(_rank, args=(world, (tmp_path / "rendezvous").as_uri(), True), nprocs=world, join=True) + +@pytest.mark.parametrize("result_channel", [False, True]) +def test_native_sampling_divergent_support_tp2(tmp_path, result_channel): + assert os.environ.get("CUDA_VISIBLE_DEVICES") + assert torch.cuda.device_count() >= 2 + mp.spawn(_rank, args=(2, (tmp_path / "rendezvous").as_uri(), result_channel, True), + nprocs=2, join=True) + + def _byte_rank(rank, world, rendezvous): from ignition.components.collectives import PreparedBroadcast from ignition.preparation.collectives import CollectivePreparation, create_result_broadcast_runtime