diff --git a/reasyn/models/decoder.py b/reasyn/models/decoder.py index e551ef8..c133e84 100644 --- a/reasyn/models/decoder.py +++ b/reasyn/models/decoder.py @@ -99,5 +99,8 @@ def forward( tgt_mask=causal_mask, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=code_padding_mask, + # Pass is_causal explicitly so nn.TransformerDecoder skips _detect_is_causal_mask, + # whose GPU-to-host scalar check adds latency and prevents CUDA-graph capture. + tgt_is_causal=bool(self.use_causal_mask), ) # (bsz, seq_len, d_model) return y diff --git a/reasyn/sampler/parallel.py b/reasyn/sampler/parallel.py index 77ac24b..df9cf9f 100644 --- a/reasyn/sampler/parallel.py +++ b/reasyn/sampler/parallel.py @@ -87,14 +87,16 @@ def __init__( def run(self) -> None: os.sched_setaffinity(0, range(os.cpu_count() or 1)) - + device = torch.device(f"cuda:{self._gpu_id}") + torch.cuda.set_device(device) + assert isinstance(self._model_path, list) and len(self._model_path) == 2 self._model = [] for _model_path in self._model_path: ckpt = torch.load(_model_path, map_location="cpu") config = OmegaConf.create(ckpt["hyper_parameters"]["config"]) - _model = ReaSyn(config.model).to(f"cuda:{self._gpu_id}") + _model = ReaSyn(config.model).to(device) _model.load_state_dict({k[6:]: v for k, v in ckpt["state_dict"].items()}) _model.eval() self._model.append(_model) @@ -117,11 +119,16 @@ def run(self) -> None: break try: result_df = self.process(next_task) - self._task_queue.task_done() - self._result_queue.put((next_task, result_df)) except KeyboardInterrupt: print(f"{self.name}: Exiting due to KeyboardInterrupt") return + except Exception as e: + print(f'{self.name}: Failed {next_task.csmiles}') + print(e) + result_df = None + finally: + self._task_queue.task_done() + self._result_queue.put((next_task, result_df)) def process(self, mol: Molecule): sampler = Sampler( @@ -151,6 +158,8 @@ def process(self, mol: Molecule): except Exception as e: print(f'{mol.csmiles}') print(e) + finally: + sampler._clear_graph_cache() class WorkerPool: @@ -163,6 +172,12 @@ def __init__( **worker_opt, ) -> None: super().__init__() + graph_enabled = (os.environ.get('REASYN_CUDAGRAPH', '0') == '1' + and os.environ.get('REASYN_BATCHED_AR', '1') != '0' + and os.environ.get('REASYN_BUCKET', '1') != '0' + and os.environ.get('REASYN_BF16', '1') != '0') + if graph_enabled and num_workers_per_gpu != 1: + raise ValueError('REASYN_CUDAGRAPH=1 requires num_workers_per_gpu=1') self._task_queue: TaskQueueType = mp.JoinableQueue(task_qsize) self._result_queue: ResultQueueType = mp.Queue(result_qsize) self._gpu_ids = [str(d) for d in gpu_ids] @@ -404,12 +419,15 @@ def run_sampling_one( ) tl = TimeLimit(time_limit) t_start = time.time() - sampler.evolve(gpu_lock=None, time_limit=tl, - num_cycles=num_cycles, - max_evolve_steps=max_evolve_steps, - num_editflow_samples=num_editflow_samples, - num_editflow_steps=num_editflow_steps) - t_elapsed = time.time() - t_start - df = sampler.get_dataframe()[: max_results] - df['time'] = t_elapsed - return df + try: + sampler.evolve(gpu_lock=None, time_limit=tl, + num_cycles=num_cycles, + max_evolve_steps=max_evolve_steps, + num_editflow_samples=num_editflow_samples, + num_editflow_steps=num_editflow_steps) + t_elapsed = time.time() - t_start + df = sampler.get_dataframe()[: max_results] + df['time'] = t_elapsed + return df + finally: + sampler._clear_graph_cache() diff --git a/reasyn/sampler/sampler.py b/reasyn/sampler/sampler.py index b2e613b..ac739fc 100644 --- a/reasyn/sampler/sampler.py +++ b/reasyn/sampler/sampler.py @@ -37,6 +37,7 @@ RXN_PATTERN = re.compile('R\d+') +MAX_CUDAGRAPH_CACHE_SIZE = 3 class Sampler: @@ -77,6 +78,13 @@ def __init__( # Pad AR forwards to power-of-2 batch / 64-multiple seq so cuBLAS/cuDNN reuse # kernels across the search's constantly-varying shapes (~7x fewer re-selections). self._bucket_ar = os.environ.get('REASYN_BUCKET', '1') != '0' + # Capture a CUDA graph per bucket shape (encoder `code` is constant within a + # Sampler, so it is baked in). Requires bucketing + bf16. Opt-in. + self._cudagraph = (os.environ.get('REASYN_CUDAGRAPH', '0') == '1' + and self._batched_ar and self._bucket_ar and self._use_bf16 + and self.device.type == 'cuda') + self._graph_cache: dict = {} + self._graph_stream = None # input filtering if self._mols_to_filter is not None and \ @@ -250,10 +258,15 @@ def _forward_ar_batched(self, code, code_padding_mask, seq_lists: list[list[int] gather = torch.tensor([l - 1 for l in lens] + [0] * (Bp - B), dtype=torch.long, device=device) + if getattr(self, '_cudagraph', False): + logits = self._run_ar_graph(code, code_padding_mask, batch, gather, Bp, Lp) + with torch.cuda.device(self.device): + return logits[:B].to(torch.float32, copy=True) + codeB = code.expand(Bp, *code.shape[1:]) maskB = code_padding_mask.expand(Bp, *code_padding_mask.shape[1:]) if self._use_bf16: - with torch.autocast('cuda', dtype=torch.bfloat16): + with torch.autocast('cuda', dtype=torch.bfloat16, cache_enabled=False): logits = self.model.sample(code=codeB, code_padding_mask=maskB, tokens=batch, token_padding_mask=None, gather_idx=gather) @@ -261,7 +274,92 @@ def _forward_ar_batched(self, code, code_padding_mask, seq_lists: list[list[int] logits = self.model.sample(code=codeB, code_padding_mask=maskB, tokens=batch, token_padding_mask=None, gather_idx=gather) - return logits[:B].float() + return logits[:B].to(torch.float32, copy=True) + + def _run_ar_graph(self, code, code_padding_mask, batch, gather, Bp, Lp) -> torch.Tensor: + """Replay a captured CUDA graph for this (Bp, Lp) shape. The encoder memory + (code/mask) is constant within a Sampler so it is baked into the graph; only + the token / gather-index static buffers are updated per call. Hot cached shapes + amortize capture over the hundreds of forwards the search issues at that shape.""" + key = (Bp, Lp) + e = self._graph_cache.pop(key, None) + if e is not None: + # Plain dicts preserve insertion order: reinsert hits for LRU eviction. + self._graph_cache[key] = e + if e is None: + with torch.cuda.device(self.device): + if len(self._graph_cache) >= MAX_CUDAGRAPH_CACHE_SIZE: + torch.cuda.synchronize(self.device) + oldest_key = next(iter(self._graph_cache)) + oldest = self._graph_cache.pop(oldest_key) + self._reset_graph_entry(oldest) + torch.cuda.empty_cache() + + s_code = code.expand(Bp, *code.shape[1:]) + s_mask = code_padding_mask.expand(Bp, *code_padding_mask.shape[1:]) + s_tok = torch.zeros((Bp, Lp), dtype=torch.long, device=self.device) + s_gi = torch.zeros((Bp,), dtype=torch.long, device=self.device) + s_tok.copy_(batch); s_gi.copy_(gather) + + def _run(): + return self.model.sample(code=s_code, code_padding_mask=s_mask, + tokens=s_tok, token_padding_mask=None, gather_idx=s_gi) + + if self._graph_stream is None: + self._graph_stream = torch.cuda.Stream(device=self.device) + + # Torch 2.7 only documents shared-pool safety when replay order matches + # capture order. Search order is arbitrary, so use independent pools and + # cap them with the LRU above. Captures still share one device-bound stream. + current = torch.cuda.current_stream(self.device) + self._graph_stream.wait_stream(current) + with torch.cuda.stream(self._graph_stream): + with torch.no_grad(), torch.autocast( + 'cuda', dtype=torch.bfloat16, cache_enabled=False): + for _ in range(3): + _run() + current.wait_stream(self._graph_stream) + g = torch.cuda.CUDAGraph() + with torch.no_grad(), torch.autocast( + 'cuda', dtype=torch.bfloat16, cache_enabled=False), torch.cuda.graph( + g, stream=self._graph_stream): + s_out = _run() + e = dict(g=g, s_tok=s_tok, s_gi=s_gi, s_out=s_out, + s_code=s_code, s_mask=s_mask) + self._graph_cache[key] = e + + with torch.cuda.device(self.device): + e['s_tok'].copy_(batch) + e['s_gi'].copy_(gather) + e['g'].replay() + return e['s_out'] + + @staticmethod + def _reset_graph_entry(entry: dict) -> None: + graph = entry.pop('g', None) + if graph is not None: + graph.reset() + entry.clear() + + def _clear_graph_cache(self) -> None: + """Release per-molecule graphs and return their pools to the CUDA allocator.""" + has_graph_state = bool(self._graph_cache) or self._graph_stream is not None + if not has_graph_state: + return + if self.device.type != 'cuda': + self._graph_cache.clear() + self._graph_stream = None + return + with torch.cuda.device(self.device): + torch.cuda.synchronize(self.device) + while self._graph_cache: + _, entry = self._graph_cache.popitem() + self._reset_graph_entry(entry) + del entry + self._graph_stream = None + # One graph-enabled worker owns each GPU, so returning unused cached + # blocks here cannot disrupt another worker on the same device. + torch.cuda.empty_cache() @torch.no_grad() def _predict_ar_batched( @@ -367,7 +465,15 @@ def _evolve_ar_singlestep( # The legacy path retains its original per-state deadline checks below. if self._batched_ar and time_limit is not None and time_limit.exceeded(): return - + + if gpu_lock is not None: + gpu_lock.acquire() + try: + return self._evolve_ar_singlestep( + gpu_lock=None, time_limit=time_limit, sampling_direction=sampling_direction) + finally: + gpu_lock.release() + feat_list = [ featurize_stack( state.stack, @@ -376,15 +482,6 @@ def _evolve_ar_singlestep( ) for state in self._active ] - - if gpu_lock is not None: - gpu_lock.acquire() - - # Acquiring a shared GPU lock may itself cross the deadline. - if self._batched_ar and time_limit is not None and time_limit.exceeded(): - if gpu_lock is not None: - gpu_lock.release() - return code, code_padding_mask = self.code @@ -426,17 +523,17 @@ def _evolve_ar_singlestep( if sampling_direction == 'td': if sampled_type == 'END': finished.append(base_state) - + elif sampled_type == 'BB' or sampled_type == 'RXN': mol_or_rxn, idx, score = sampled_item[i] new_state = copy.deepcopy(base_state) new_state.stack.push_topdown(mol_or_rxn, idx) new_state.scores.append(score) next.append(new_state) - + else: self._aborted.append(base_state) - + else: if sampled_type == 'END': finished.append(base_state) @@ -481,14 +578,11 @@ def _evolve_ar_singlestep( del self._active self._active = next self._sort_states() - + if sampling_direction == 'td': [state.stack.final_seq_topdown() for state in finished] self._add_finished_states(finished) - if gpu_lock is not None: - gpu_lock.release() - @torch.inference_mode() def _predict_editflow( self, diff --git a/tests/cuda_graph_testlib.py b/tests/cuda_graph_testlib.py new file mode 100644 index 0000000..c6b69a0 --- /dev/null +++ b/tests/cuda_graph_testlib.py @@ -0,0 +1,596 @@ +"""Adversarial checkpoint-backed regression harness for CUDA-graph decoding. + +Run on an A100 with the AR and Edit-Bridge checkpoint environment variables. +The pytest suite imports the focused checks from this file; the CLI also retains +the longer diagnostics used during review. +""" + +import argparse +import gc +import multiprocessing as mp +import os +import pathlib +import queue +import sys +import time + +import torch +from omegaconf import OmegaConf + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) + +from reasyn.chem.featurize import TokenType +from reasyn.chem.mol import Molecule +from reasyn.models.reasyn import ReaSyn +from reasyn.sampler.sampler import MAX_CUDAGRAPH_CACHE_SIZE, Sampler +from reasyn.utils.sample_utils import State + + +AR_CKPT = os.environ.get( + "REASYN_AR_CKPT", "/workspace/shared/ckpt/nv-reasyn-ar-166m-v2.ckpt" +) +EB_CKPT = os.environ.get( + "REASYN_EB_CKPT", "/workspace/shared/ckpt/nv-reasyn-eb-174m-v2.ckpt" +) +DEV = "cuda" + + +def load_model(path: str, device: str = DEV) -> ReaSyn: + ckpt = torch.load(path, map_location="cpu", weights_only=False) + cfg = OmegaConf.create(ckpt["hyper_parameters"]["config"]) + model = ReaSyn(cfg.model).to(device) + model.load_state_dict({k[6:]: v for k, v in ckpt["state_dict"].items()}) + model.eval() + return model + + +def decoder_without_hints( + decoder, code, code_padding_mask, tokens, token_padding_mask=None, t=None +): + """The parent-branch Decoder.forward, kept literal for A/B comparison.""" + bsz, seqlen = tokens.size() + x = decoder.embed(tokens) + if t is not None and decoder.use_time_embed: + x = x + decoder.te_dec(t) + causal_mask = ( + torch.nn.Transformer.generate_square_subsequent_mask( + x.size(1), dtype=x.dtype, device=x.device + ) + if decoder.use_causal_mask + else None + ) + tgt_key_padding_mask = ( + torch.zeros((bsz, seqlen), dtype=x.dtype, device=x.device).masked_fill_( + token_padding_mask, -torch.finfo(x.dtype).max + ) + if token_padding_mask is not None + else None + ) + return decoder.dec( + tgt=x, + memory=code, + tgt_mask=causal_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=code_padding_mask, + ) + + +def padded_inputs(model: ReaSyn, batch: int = 3, src_len: int = 40, tgt_len: int = 24): + ntt = model.encoder.smiles_emb.num_embeddings + smiles = torch.randint(1, ntt, (batch, src_len), device=DEV) + for row, length in enumerate((src_len, src_len - 7, src_len - 13)[:batch]): + smiles[row, length:] = 0 + code, code_mask = model.encoder(smiles) + tokens = torch.randint(1, min(model.vocab_size, 100), (batch, tgt_len), device=DEV) + lengths = torch.tensor((tgt_len, tgt_len - 5, tgt_len - 11)[:batch], device=DEV) + token_mask = torch.arange(tgt_len, device=DEV)[None] >= lengths[:, None] + tokens[token_mask] = 0 + return code, code_mask, tokens, token_mask + + +@torch.inference_mode() +def check_decoder(ar: ReaSyn, eb: ReaSyn): + assert ar.model_type == "autoregressive" and ar.decoder.use_causal_mask + assert eb.model_type == "editflow" and not eb.decoder.use_causal_mask + + for name, model in (("AR", ar), ("EB", eb)): + code, code_mask, tokens, token_mask = padded_inputs(model) + t = torch.full((tokens.size(0), 1), 0.37, device=DEV) if name == "EB" else None + old = decoder_without_hints(model.decoder, code, code_mask, tokens, token_mask, t) + new = model.decoder(code, code_mask, tokens, token_mask, t) + diff = (old - new).abs() + print( + f"DECODER_{name} equal={torch.equal(old, new)} " + f"max_abs={diff.max().item():.9g} padded=True" + ) + assert torch.equal(old, new) + + if name == "EB": + ut, ins, sub = model.sample(code, code_mask, tokens, token_mask, t) + assert ut.shape[:2] == tokens.shape and ins.shape[:2] == tokens.shape + assert sub.shape == ins.shape + assert torch.isfinite(ut).all() and torch.isfinite(ins).all() and torch.isfinite(sub).all() + print(f"EDIT_BRIDGE_SAMPLE_OK shapes={tuple(ut.shape)},{tuple(ins.shape)}") + + # A future-token perturbation must not affect an AR prefix, but must affect EB. + for name, model in (("AR", ar), ("EB", eb)): + code, code_mask, tokens, _ = padded_inputs(model) + changed = tokens.clone() + changed[:, 8:] = (changed[:, 8:] + 17) % min(model.vocab_size, 100) + t = torch.full((tokens.size(0), 1), 0.37, device=DEV) if name == "EB" else None + # Edit-Bridge's custom MHA requires the slow path used in production, where + # _predict_editflow always supplies a token padding mask. + semantic_mask = torch.zeros_like(tokens, dtype=torch.bool) if name == "EB" else None + y1 = model.decoder(code, code_mask, tokens, semantic_mask, t) + y2 = model.decoder(code, code_mask, changed, semantic_mask, t) + prefix_delta = (y1[:, :8] - y2[:, :8]).abs().max().item() + print(f"CAUSAL_SEMANTICS_{name} future_to_prefix_max_abs={prefix_delta:.9g}") + if name == "AR": + assert prefix_delta == 0.0 + else: + assert prefix_delta > 0.0 + + # On the tested runtime the implicit detector performs a scalar extraction. + code, code_mask, tokens, token_mask = padded_inputs(ar) + try: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CPU]) as prof: + decoder_without_hints(ar.decoder, code, code_mask, tokens, token_mask) + torch.cuda.synchronize() + sync_ops = { + event.key: event.count + for event in prof.key_averages() + if "local_scalar" in event.key or "equal" in event.key + } + print(f"IMPLICIT_CAUSAL_DETECT_SYNC_OPS {sync_ops}") + except Exception as exc: + print(f"IMPLICIT_CAUSAL_DETECT_PROFILER_SKIPPED {type(exc).__name__}: {exc}") + + +def bare_sampler(model: ReaSyn, use_graph: bool) -> Sampler: + sampler = Sampler.__new__(Sampler) + sampler.model = model + sampler.device = next(model.parameters()).device + sampler._bucket_ar = True + sampler._use_bf16 = True + sampler._cudagraph = use_graph + sampler._graph_cache = {} + sampler._graph_stream = None + return sampler + + +def production_sampler(model: ReaSyn, use_graph: bool) -> Sampler: + """Construct through Sampler.__init__ with the real environment toggle.""" + class EditflowStub: + model_type = "editflow" + + keys = ("REASYN_CUDAGRAPH", "REASYN_BATCHED_AR", "REASYN_BUCKET", "REASYN_BF16") + old = {key: os.environ.get(key) for key in keys} + try: + os.environ.update( + REASYN_CUDAGRAPH="1" if use_graph else "0", + REASYN_BATCHED_AR="1", + REASYN_BUCKET="1", + REASYN_BF16="1", + ) + return Sampler(None, None, Molecule("CCO"), [model, EditflowStub()]) + finally: + for key, value in old.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def make_sequences(seed: int, count: int, max_len: int = 60) -> list[list[int]]: + generator = torch.Generator().manual_seed(seed) + lengths = torch.randint(max(1, max_len - 20), max_len + 1, (count,), generator=generator) + return [ + torch.randint(1, 100, (int(length),), generator=generator).tolist() + for length in lengths + ] + + +def static_batch(seqs: list[list[int]], bp: int, lp: int, device): + padded = [seq + [0] * (lp - len(seq)) for seq in seqs] + padded += [[0] * lp for _ in range(bp - len(seqs))] + batch = torch.tensor(padded, dtype=torch.long, device=device) + gather = torch.tensor( + [len(seq) - 1 for seq in seqs] + [0] * (bp - len(seqs)), + dtype=torch.long, + device=device, + ) + return batch, gather + + +@torch.inference_mode() +def check_graph(ar: ReaSyn): + torch.manual_seed(123) + smiles1 = torch.randint(1, ar.encoder.smiles_emb.num_embeddings, (1, 48), device=DEV) + smiles2 = torch.randint(1, ar.encoder.smiles_emb.num_embeddings, (1, 48), device=DEV) + code1, mask1 = ar.encoder(smiles1) + code2, mask2 = ar.encoder(smiles2) + eager = production_sampler(ar, False) + graph = production_sampler(ar, True) + assert not eager._cudagraph and graph._cudagraph + + first_capture_s = None + equal_cases = 0 + max_graph_eager = 0.0 + max_prob_delta = 0.0 + argmax_matches = 0 + sample_matches = 0 + repeated_draw_matches = 0 + repeated_draws = 0 + rows = 0 + for seed in range(20): + # Counts 5..8 and lengths 40..60 all hit exactly cache key (8, 64). + seqs = make_sequences(seed, 5 + seed % 4) + ref = eager._forward_ar_batched(code1, mask1, seqs) + start = time.perf_counter() + got = graph._forward_ar_batched(code1, mask1, seqs) + torch.cuda.synchronize() + if first_capture_s is None: + first_capture_s = time.perf_counter() - start + max_abs = (ref - got).abs().max().item() + equal_cases += int(torch.equal(ref, got)) + assert torch.equal(ref, got), f"graph/eager mismatch for same-shape seed {seed}" + max_graph_eager = max(max_graph_eager, max_abs) + ref_prob = torch.softmax(ref / 0.1, dim=-1) + got_prob = torch.softmax(got / 0.1, dim=-1) + max_prob_delta = max(max_prob_delta, (ref_prob - got_prob).abs().max().item()) + argmax_matches += int((ref.argmax(-1) == got.argmax(-1)).sum()) + state = torch.cuda.get_rng_state(graph.device) + ref_sample = torch.multinomial(ref_prob, 1) + torch.cuda.set_rng_state(state, graph.device) + got_sample = torch.multinomial(got_prob, 1) + assert torch.equal(ref_sample, got_sample) + sample_matches += int((ref_sample == got_sample).sum()) + state = torch.cuda.get_rng_state(graph.device) + ref_draws = torch.multinomial(ref_prob, 128, replacement=True) + torch.cuda.set_rng_state(state, graph.device) + got_draws = torch.multinomial(got_prob, 128, replacement=True) + assert torch.equal(ref_draws, got_draws) + repeated_draw_matches += int((ref_draws == got_draws).sum()) + repeated_draws += ref_draws.numel() + rows += len(seqs) + entry = graph._graph_cache[(8, 64)] + assert entry["s_code"].stride(0) == 0 + assert entry["s_mask"].stride(0) == 0 + assert graph._graph_stream.device == graph.device + print( + f"GRAPH_REPLAY_EQ cases=20 bit_equal_cases={equal_cases}/20 " + f"max_abs={max_graph_eager:.9g} max_prob_at_T0.1={max_prob_delta:.9g} " + f"argmax_match={argmax_matches}/{rows} same_rng_sample_match={sample_matches}/{rows} " + f"same_rng_repeated_draw_match={repeated_draw_matches}/{repeated_draws} " + f"first_capture_s={first_capture_s:.6f} cache_keys={list(graph._graph_cache)}" + ) + + # Raw graph output aliases and is overwritten; public output must not. + seqs1, seqs2 = make_sequences(77, 7), make_sequences(88, 7) + batch1, gather1 = static_batch(seqs1, 8, 64, graph.device) + batch2, gather2 = static_batch(seqs2, 8, 64, graph.device) + raw1 = graph._run_ar_graph(code1, mask1, batch1, gather1, 8, 64) + torch.cuda.synchronize() + raw_snapshot = raw1.clone() + raw2 = graph._run_ar_graph(code1, mask1, batch2, gather2, 8, 64) + torch.cuda.synchronize() + raw_overwritten = not torch.equal(raw1, raw_snapshot) + assert raw1.data_ptr() == raw2.data_ptr() == entry["s_out"].data_ptr() + assert raw_overwritten + + public1 = graph._forward_ar_batched(code1, mask1, seqs1) + torch.cuda.synchronize() + public_snapshot = public1.clone() + public2 = graph._forward_ar_batched(code1, mask1, seqs2) + torch.cuda.synchronize() + assert torch.equal(public1, public_snapshot) + assert public1.data_ptr() != entry["s_out"].data_ptr() + print( + f"ALIAS_TRAP raw_dtype={entry['s_out'].dtype} raw_reused=True " + f"raw_overwritten={raw_overwritten} public_dtype={public1.dtype} " + f"public_distinct_storage=True public_survived_replay=True " + f"new_output_changed={not torch.equal(public1, public2)}" + ) + + # Deliberately violate the baked-code contract to expose the failure mode. + seqs = make_sequences(99, 6) + graph1 = graph._forward_ar_batched(code1, mask1, seqs) + eager2 = eager._forward_ar_batched(code2, mask2, seqs) + stale = graph._forward_ar_batched(code2, mask2, seqs) + fresh_graph = production_sampler(ar, True) + fresh = fresh_graph._forward_ar_batched(code2, mask2, seqs) + torch.cuda.synchronize() + assert torch.equal(stale, graph1) + assert not torch.equal(stale, eager2) + assert torch.equal(fresh, eager2) + print( + f"BAKED_CODE same_cache_ignores_new_code=True " + f"stale_vs_new_max_abs={(stale-eager2).abs().max().item():.9g} " + f"fresh_graph_vs_eager_max_abs={(fresh-eager2).abs().max().item():.9g} " + f"fresh_graph_vs_eager_bit_equal=True cache_objects_distinct=" + f"{graph._graph_cache is not fresh_graph._graph_cache}" + ) + + # Independent graph pools must be safe under arbitrary serial replay order. + # Keep a copied result alive while replaying the other shape to catch aliasing. + seqs_64 = make_sequences(501, 7, 60) + seqs_128 = make_sequences(502, 3, 100) + for index in range(12): + first_seqs, second_seqs = ( + (seqs_64, seqs_128) if index % 2 == 0 else (seqs_128, seqs_64) + ) + expected = eager._forward_ar_batched(code1, mask1, first_seqs) + saved = graph._forward_ar_batched(code1, mask1, first_seqs) + saved_snapshot = saved.clone() + graph._forward_ar_batched(code1, mask1, second_seqs) + torch.cuda.synchronize() + assert torch.equal(saved, expected) + assert torch.equal(saved, saved_snapshot) + pools = {entry["g"].pool() for entry in graph._graph_cache.values()} + assert len(pools) == len(graph._graph_cache) + print( + f"PRIVATE_POOL_ALTERNATING pass=True keys={sorted(graph._graph_cache)} " + f"pool_count={len(pools)} stream_device={graph._graph_stream.device}" + ) + + # Steady-state timing at the key used above, including Python/tensor copies. + cases = [make_sequences(1000 + i, 8) for i in range(30)] + for seqs in cases[:3]: + eager._forward_ar_batched(code1, mask1, seqs) + graph._forward_ar_batched(code1, mask1, seqs) + torch.cuda.synchronize() + start = time.perf_counter() + for seqs in cases: + eager._forward_ar_batched(code1, mask1, seqs) + torch.cuda.synchronize() + eager_s = time.perf_counter() - start + start = time.perf_counter() + for seqs in cases: + graph._forward_ar_batched(code1, mask1, seqs) + torch.cuda.synchronize() + graph_s = time.perf_counter() - start + speedup = eager_s / graph_s + breakeven = first_capture_s / max(eager_s / len(cases) - graph_s / len(cases), 1e-12) + print( + f"SPEED key=(8,64) n=30 eager_s={eager_s:.6f} graph_s={graph_s:.6f} " + f"speedup={speedup:.3f}x capture_s={first_capture_s:.6f} " + f"estimated_replays_to_amortize={breakeven:.1f}" + ) + + # Exceed the cache limit, verify LRU eviction, then revisit an evicted shape. + lru_cases = [ + make_sequences(601, 2, 180), + make_sequences(602, 1, 240), + make_sequences(603, 16, 60), + make_sequences(604, 4, 100), + ] + for seqs in lru_cases: + expected = eager._forward_ar_batched(code1, mask1, seqs) + got = graph._forward_ar_batched(code1, mask1, seqs) + assert torch.equal(got, expected) + assert len(graph._graph_cache) <= MAX_CUDAGRAPH_CACHE_SIZE + pools = {entry["g"].pool() for entry in graph._graph_cache.values()} + assert len(graph._graph_cache) == MAX_CUDAGRAPH_CACHE_SIZE + assert len(pools) == MAX_CUDAGRAPH_CACHE_SIZE + assert list(graph._graph_cache) == [(1, 256), (16, 64), (4, 128)] + print( + f"GRAPH_LRU pass=True resident={len(graph._graph_cache)} " + f"limit={MAX_CUDAGRAPH_CACHE_SIZE} keys={list(graph._graph_cache)}" + ) + passed = equal_cases == 20 + graph._clear_graph_cache() + fresh_graph._clear_graph_cache() + return passed + + +def check_sampler_contract_and_guards(ar: ReaSyn, eb: ReaSyn): + keys = ("REASYN_CUDAGRAPH", "REASYN_BATCHED_AR", "REASYN_BUCKET", "REASYN_BF16") + old = {key: os.environ.get(key) for key in keys} + try: + os.environ.update( + REASYN_CUDAGRAPH="1", REASYN_BATCHED_AR="1", + REASYN_BUCKET="1", REASYN_BF16="1", + ) + first = Sampler(None, None, Molecule("CCO"), [ar, eb]) + second = Sampler(None, None, Molecule("CCN"), [ar, eb]) + code_a = first.code + assert code_a[0].data_ptr() == first.code[0].data_ptr() + assert first._graph_cache is not second._graph_cache + print( + "SAMPLER_CONTRACT code_cached=True graph_cache_per_sampler=True " + f"graph_enabled={first._cudagraph}" + ) + + os.environ["REASYN_BUCKET"] = "0" + no_bucket = Sampler(None, None, Molecule("CCO"), [ar, eb]) + os.environ.update(REASYN_BUCKET="1", REASYN_BF16="0") + no_bf16 = Sampler(None, None, Molecule("CCO"), [ar, eb]) + os.environ.update(REASYN_BF16="1", REASYN_BATCHED_AR="0") + no_batched = Sampler(None, None, Molecule("CCO"), [ar, eb]) + assert not no_bucket._cudagraph and not no_bf16._cudagraph + assert not no_batched._cudagraph + print("GRAPH_GUARDS bucket0=False bf16_0=False batched0=False") + + class Dummy(torch.nn.Module): + def __init__(self, model_type): + super().__init__() + self.model_type = model_type + self.p = torch.nn.Parameter(torch.zeros(())) + + os.environ.update(REASYN_BUCKET="1", REASYN_BF16="1", REASYN_BATCHED_AR="1") + cpu_sampler = Sampler( + None, + None, + Molecule("CCO"), + [Dummy("autoregressive"), Dummy("editflow")], + ) + print( + f"CPU_GUARD_ON_GPU_HOST graph_enabled={cpu_sampler._cudagraph} " + f"model_device={cpu_sampler.device} cuda_available={torch.cuda.is_available()}" + ) + assert not cpu_sampler._cudagraph + finally: + for key, value in old.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +@torch.inference_mode() +def run_core(): + print( + f"RUNTIME torch={torch.__version__} cuda={torch.version.cuda} " + f"device={torch.cuda.get_device_name()} current_device={torch.cuda.current_device()}" + ) + torch.manual_seed(7) + ar = load_model(AR_CKPT) + eb = load_model(EB_CKPT) + check_decoder(ar, eb) + check_sampler_contract_and_guards(ar, eb) + graph_equal = check_graph(ar) + print(f"CORE_GRAPH_EAGER_EQUIVALENCE pass={graph_equal}") + assert graph_equal, "CUDA-graph output is not bit-identical to eager output" + + +@torch.inference_mode() +def run_growth(full: bool = False): + ar = load_model(AR_CKPT) + src_len = 256 if full else 64 + smiles = torch.randint(1, ar.encoder.smiles_emb.num_embeddings, (1, src_len), device=DEV) + code, mask = ar.encoder(smiles) + sampler = bare_sampler(ar, True) + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + baseline_alloc = torch.cuda.memory_allocated() + baseline_res = torch.cuda.memory_reserved() + shapes = ( + sorted( + ((bp, lp) for bp in (1, 2, 4, 8, 16, 32, 64, 128, 256) + for lp in (64, 128, 192, 256, 320, 384, 448, 512)), + key=lambda shape: shape[0] * shape[1] ** 2, + ) + if full + else [ + (1, 64), (2, 64), (4, 64), (8, 64), (16, 64), (32, 64), + (1, 128), (2, 128), (4, 128), (8, 128), (1, 256), (2, 256), + ] + ) + print( + f"GROWTH_BASE allocated_mib={baseline_alloc/2**20:.1f} " + f"reserved_mib={baseline_res/2**20:.1f}" + ) + captured = 0 + for bp, lp in shapes: + seqs = make_sequences(bp * 1000 + lp, bp, lp) + start = time.perf_counter() + try: + sampler._forward_ar_batched(code, mask, seqs) + except torch.OutOfMemoryError as exc: + print(f"GROWTH_OOM key=({bp},{lp}) cache={len(sampler._graph_cache)} error={exc}") + break + captured += 1 + torch.cuda.synchronize() + elapsed = time.perf_counter() - start + alloc = torch.cuda.memory_allocated() + reserved = torch.cuda.memory_reserved() + print( + f"GROWTH key=({bp},{lp}) capture_s={elapsed:.4f} cache={len(sampler._graph_cache)} " + f"allocated_mib={alloc/2**20:.1f} delta_alloc_mib={(alloc-baseline_alloc)/2**20:.1f} " + f"reserved_mib={reserved/2**20:.1f} delta_reserved_mib={(reserved-baseline_res)/2**20:.1f}" + ) + completed = captured == len(shapes) + reserved_growth = torch.cuda.memory_reserved() - baseline_res + print( + f"GROWTH_COMPLETE pass={completed} captured={captured}/{len(shapes)} " + f"resident={len(sampler._graph_cache)}/{MAX_CUDAGRAPH_CACHE_SIZE} " + f"peak_allocated_mib={torch.cuda.max_memory_allocated()/2**20:.1f} " + f"delta_reserved_mib={reserved_growth/2**20:.1f}" + ) + assert completed, "shape-cache growth run did not capture every requested key" + assert len(sampler._graph_cache) <= MAX_CUDAGRAPH_CACHE_SIZE + assert reserved_growth < 12 * 2**30, "bounded graph cache exceeded 12 GiB" + sampler._clear_graph_cache() + gc.collect() + torch.cuda.synchronize() + print( + f"GROWTH_AFTER_CLEAR allocated_mib={torch.cuda.memory_allocated()/2**20:.1f} " + f"reserved_mib={torch.cuda.memory_reserved()/2**20:.1f}" + ) + assert torch.cuda.memory_reserved() <= baseline_res + 512 * 2**20 + + +def worker_capture(rank: int, lock, result_queue): + try: + torch.cuda.set_device(0) + ar = load_model(AR_CKPT) + smiles = torch.randint(1, ar.encoder.smiles_emb.num_embeddings, (1, 32), device=DEV) + with torch.inference_mode(): + code, mask = ar.encoder(smiles) + sampler = bare_sampler(ar, True) + with lock: + out = sampler._forward_ar_batched(code, mask, make_sequences(rank + 300, 4)) + torch.cuda.synchronize() + result_queue.put((rank, "ok", float(out.sum()), len(sampler._graph_cache))) + except Exception as exc: + result_queue.put((rank, "error", f"{type(exc).__name__}: {exc}", 0)) + + +def run_multiprocess(): + ctx = mp.get_context("spawn") + lock = ctx.Lock() + result_queue = ctx.Queue() + workers = [ctx.Process(target=worker_capture, args=(rank, lock, result_queue)) for rank in range(2)] + for worker in workers: + worker.start() + results = [] + for _ in workers: + try: + results.append(result_queue.get(timeout=240)) + except queue.Empty: + results.append((-1, "timeout", "no result", 0)) + for worker in workers: + worker.join(timeout=30) + print(f"MULTIPROCESS_RESULTS {sorted(results)} exitcodes={[w.exitcode for w in workers]}") + assert all(item[1] == "ok" for item in results) + assert all(worker.exitcode == 0 for worker in workers) + print("MULTIPROCESS_PASS workers=2 captures=2 serialized_by_gpu_lock=True") + + +def run_lock_failure(): + class FailingModel: + @staticmethod + def encoder(_): + raise RuntimeError("synthetic post-acquire failure") + + sampler = Sampler.__new__(Sampler) + sampler._active = [State()] + sampler._batched_ar = True + sampler.model = FailingModel() + sampler._smiles = torch.zeros((1, 1), dtype=torch.long) + lock = mp.Lock() + try: + sampler._evolve_ar_singlestep(gpu_lock=lock) + except RuntimeError as exc: + print(f"LOCK_INJECTED_ERROR {exc}") + reacquired = lock.acquire(block=False) + print(f"LOCK_AFTER_EXCEPTION released={reacquired}") + assert reacquired + lock.release() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "mode", choices=("core", "growth", "fullgrowth", "multiprocess", "lock") + ) + args = parser.parse_args() + { + "core": run_core, + "growth": run_growth, + "fullgrowth": lambda: run_growth(full=True), + "multiprocess": run_multiprocess, + "lock": run_lock_failure, + }[args.mode]() diff --git a/tests/test_cuda_graphs.py b/tests/test_cuda_graphs.py new file mode 100644 index 0000000..c762fc8 --- /dev/null +++ b/tests/test_cuda_graphs.py @@ -0,0 +1,179 @@ +"""Regression tests for CUDA-graph autoregressive decoding. + +These tests require the real AR checkpoint and a CUDA device. They are skipped +cleanly on ordinary CPU-only developer machines. +""" + +import gc +import os +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) + +from cuda_graph_testlib import ( # noqa: E402 + AR_CKPT, + bare_sampler, + check_graph, + check_sampler_contract_and_guards, + load_model, + make_sequences, + production_sampler, + run_lock_failure, +) +from reasyn.chem.mol import Molecule # noqa: E402 +from reasyn.sampler.parallel import Worker, WorkerPool # noqa: E402 +from reasyn.sampler.sampler import MAX_CUDAGRAPH_CACHE_SIZE, Sampler # noqa: E402 + + +@pytest.fixture(scope="module") +def ar_model(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if not hasattr(torch.cuda, "CUDAGraph") or not torch.cuda.is_bf16_supported(): + pytest.skip("CUDA graphs and BF16 are required") + if not pathlib.Path(AR_CKPT).is_file(): + pytest.skip(f"AR checkpoint not found: {AR_CKPT}") + torch.cuda.set_device(0) + model = load_model(AR_CKPT) + yield model + del model + gc.collect() + torch.cuda.empty_cache() + + +def test_actual_env_graph_eager_parity_many_same_shape_cases(ar_model): + """Exercise real Sampler init/env paths, replay, RNG, and pool ordering.""" + assert check_graph(ar_model) + + +def test_graph_guards_and_fp32_output_ownership(ar_model): + class EditflowStub: + model_type = "editflow" + + check_sampler_contract_and_guards(ar_model, EditflowStub()) + + sampler = production_sampler(ar_model, True) + static = torch.randn( + (8, ar_model.vocab_size), dtype=torch.float32, device=sampler.device + ) + + def fake_graph(*_args, **_kwargs): + return static + + sampler._run_ar_graph = fake_graph + output = sampler._forward_ar_batched( + torch.empty(0, device=sampler.device), + torch.empty(0, device=sampler.device), + make_sequences(917, 7), + ) + snapshot = output.clone() + static.add_(1) + assert output.dtype == torch.float32 + assert output.data_ptr() != static.data_ptr() + assert torch.equal(output, snapshot) + + +def test_cpu_model_never_enables_graphs_on_a_gpu_host(monkeypatch): + class Dummy(torch.nn.Module): + def __init__(self, model_type): + super().__init__() + self.model_type = model_type + self.p = torch.nn.Parameter(torch.zeros(())) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setenv("REASYN_CUDAGRAPH", "1") + monkeypatch.setenv("REASYN_BATCHED_AR", "1") + monkeypatch.setenv("REASYN_BUCKET", "1") + monkeypatch.setenv("REASYN_BF16", "1") + sampler = Sampler( + None, + None, + Molecule("CCO"), + [Dummy("autoregressive"), Dummy("editflow")], + ) + assert sampler.device.type == "cpu" + assert not sampler._cudagraph + + +def test_bounded_graph_cache_has_bounded_growth_and_releases(ar_model): + if torch.cuda.get_device_properties(0).total_memory < 16 * 2**30: + pytest.skip("large-shape graph memory regression requires at least 16 GiB") + torch.manual_seed(2026) + smiles = torch.randint( + 1, ar_model.encoder.smiles_emb.num_embeddings, (1, 128), device="cuda" + ) + with torch.inference_mode(): + code, mask = ar_model.encoder(smiles) + + sampler = bare_sampler(ar_model, True) + shapes = [ + (1, 64), (8, 64), (32, 64), (128, 64), (256, 64), + (1, 256), (8, 256), (32, 256), (128, 256), (256, 256), + (1, 512), (8, 512), (32, 512), (128, 512), (256, 512), + ] + torch.cuda.empty_cache() + baseline = torch.cuda.memory_reserved() + try: + with torch.inference_mode(): + for batch_size, length in shapes: + sampler._forward_ar_batched( + code, + mask, + make_sequences(batch_size * 1000 + length, batch_size, length), + ) + torch.cuda.synchronize() + growth = torch.cuda.memory_reserved() - baseline + print( + f"MEMORY_GROWTH keys={len(shapes)} reserved_delta_mib={growth / 2**20:.1f}" + ) + assert len(sampler._graph_cache) == MAX_CUDAGRAPH_CACHE_SIZE + assert len({entry["g"].pool() for entry in sampler._graph_cache.values()}) == ( + MAX_CUDAGRAPH_CACHE_SIZE + ) + # This selected grid includes the largest legal key. Pool identity catches + # accidental sharing; the LRU and bound catch runaway aggregate reservation. + assert growth < 12 * 2**30, f"graph cache reserved {growth / 2**30:.2f} GiB" + finally: + sampler._clear_graph_cache() + torch.cuda.synchronize() + released = torch.cuda.memory_reserved() + print( + f"MEMORY_AFTER_CLEAR reserved_mib={released / 2**20:.1f} " + f"baseline_mib={baseline / 2**20:.1f}" + ) + assert released <= baseline + 512 * 2**20 + + +def test_worker_device_binding_and_single_worker_guard(monkeypatch): + class StopBeforeCheckpointLoad(Exception): + pass + + selected = [] + + def stop_after_device(device): + selected.append(device) + raise StopBeforeCheckpointLoad + + monkeypatch.setattr(os, "sched_setaffinity", lambda *_args: None, raising=False) + monkeypatch.setattr(torch.cuda, "set_device", stop_after_device) + worker = Worker( + [pathlib.Path("ar.ckpt"), pathlib.Path("eb.ckpt")], None, None, "3", None + ) + with pytest.raises(StopBeforeCheckpointLoad): + worker.run() + assert selected == [torch.device("cuda:3")] + + monkeypatch.setenv("REASYN_CUDAGRAPH", "1") + monkeypatch.setenv("REASYN_BATCHED_AR", "1") + monkeypatch.setenv("REASYN_BUCKET", "1") + monkeypatch.setenv("REASYN_BF16", "1") + with pytest.raises(ValueError, match="num_workers_per_gpu=1"): + WorkerPool([0], 2, 1, 1) + + +def test_gpu_lock_is_released_after_failure(): + run_lock_failure()