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