diff --git a/reasyn/models/reasyn.py b/reasyn/models/reasyn.py index 64f5dfd..a175bca 100644 --- a/reasyn/models/reasyn.py +++ b/reasyn/models/reasyn.py @@ -87,6 +87,7 @@ def sample( tokens: torch.Tensor, token_padding_mask: torch.Tensor | None, t: torch.Tensor | None = None, # for EditFlow + gather_idx: torch.Tensor | None = None, # AR: per-row index of the last real token ) -> torch.Tensor: h = self.decoder( code=code, @@ -101,6 +102,11 @@ def sample( sub_logits = self.sub_out(h) return ut, ins_logits, sub_logits if self.model_type == 'autoregressive': - h = h[:, -1] # (1, h_dim) + if gather_idx is not None: + # Batched decoding: right-padded rows have their last real token at + # position (length-1), not at index -1. Gather per row. + h = h[torch.arange(h.size(0), device=h.device), gather_idx] # (B, h_dim) + else: + h = h[:, -1] # (1, h_dim) — unbatched / full-length prefix logits = self.token_head(h) return logits diff --git a/reasyn/sampler/sampler.py b/reasyn/sampler/sampler.py index 6617c9e..b2e613b 100644 --- a/reasyn/sampler/sampler.py +++ b/reasyn/sampler/sampler.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import re import copy from collections.abc import Iterable @@ -69,7 +70,14 @@ def __init__( self.exact_break = exact_break self._mols_to_filter = mols_to_filter self._filter_sim = filter_sim - + # Batched AR decoding (env-toggleable for A/B against the original per-state loop). + # REASYN_BATCHED_AR=0 -> original one-state-at-a-time path; REASYN_BF16=1 -> bf16 forward. + self._batched_ar = os.environ.get('REASYN_BATCHED_AR', '1') != '0' + self._use_bf16 = os.environ.get('REASYN_BF16', '0') != '0' + # 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' + # input filtering if self._mols_to_filter is not None and \ max([self._mol.sim(f) for f in self._mols_to_filter] or [-1]) > self._filter_sim: @@ -216,6 +224,136 @@ def _sample_token(tokens, enforce_td=False): return PredictResult(sampled_type, sampled_item) + def _forward_ar_batched(self, code, code_padding_mask, seq_lists: list[list[int]]) -> torch.Tensor: + """One batched decoder forward over CPU int-list prefixes. Right-pads on CPU, + moves to device in a single H2D copy (no per-lane GPU<->CPU syncs), and gathers + per-row last-real-token logits. Returns (B, vocab) in fp32. + + Shape bucketing (REASYN_BUCKET): the search feeds a new (batch, seq) shape on + almost every call, so cuBLAS/cuDNN re-selects algorithms per shape (~7x overhead + measured). Padding batch to a power-of-2 and seq to a 64-multiple collapses the + distinct-shape count so kernels are reused. Extra rows/cols are ignored via gather + + causal masking, so results are identical (the pad mask is redundant: the gathered + last-real-token position is causally isolated from the right-padded future).""" + device = self.device + lens = [len(s) for s in seq_lists] + B = len(seq_lists) + Lmax = max(lens) + if self._bucket_ar: + Bp = 1 << max(0, (B - 1)).bit_length() # next power of two + Lp = min(((Lmax + 63) // 64) * 64, self.model.max_len) # next multiple of 64 + else: + Bp, Lp = B, Lmax + padded = [s + [0] * (Lp - len(s)) for s in seq_lists] + padded += [[0] * Lp for _ in range(Bp - B)] # dummy rows -> ignored + batch = torch.tensor(padded, dtype=torch.long, device=device) # one H2D + gather = torch.tensor([l - 1 for l in lens] + [0] * (Bp - B), + dtype=torch.long, device=device) + + 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): + logits = self.model.sample(code=codeB, code_padding_mask=maskB, + tokens=batch, token_padding_mask=None, + gather_idx=gather) + else: + logits = self.model.sample(code=codeB, code_padding_mask=maskB, + tokens=batch, token_padding_mask=None, + gather_idx=gather) + return logits[:B].float() + + @torch.no_grad() + def _predict_ar_batched( + self, + code: torch.Tensor | None, + code_padding_mask: torch.Tensor | None, + tokens_list: list[list[int]], + temperature_token: float = 0.1, + use_edit_distance: bool = True, + sampling_direction: str = 'bu', + ) -> list[PredictResult]: + """Batched equivalent of _predict_ar over all active states at once. + + Batches the GPU forwards (outer draw + inner BB-SMILES loop) while keeping the + chemistry callbacks (get_reactants/get_reactions) per-state. Token bookkeeping + stays on CPU as int-lists so each character-step incurs exactly one H2D (batch + build) and one D2H (sample readout) sync regardless of batch size. Numerically + equivalent to _predict_ar up to RNG-draw-order and bf16 (validated distributionally).""" + MOL_START, MOL_END = int(TokenType.MOL_START), int(TokenType.MOL_END) + RXN_MIN, RXN_MAX = int(TokenType.RXN_MIN), int(TokenType.RXN_MAX) + max_len = self.model.max_len + N = len(tokens_list) + results: list[PredictResult | None] = [None] * N + + # CPU int-lists; immediate aborts (prefix over budget) excluded from the batch. + act: list[int] = [] + seqs: list[list[int]] = [] + for k in range(N): + L = len(tokens_list[k]) + if L > max_len: + results[k] = PredictResult('ABORTED', None) + else: + act.append(k) + seqs.append(tokens_list[k]) + if not act: + return results + + def _sample(logits, td_first_idx=None) -> list[int]: + if td_first_idx: + logits = logits.clone() + for i in td_first_idx: + logits[i, :RXN_MIN] = -torch.inf # TD first token must be a reaction + probs = torch.nn.functional.softmax(logits / temperature_token, dim=-1) + return torch.multinomial(probs, num_samples=1).squeeze(-1).tolist() # one D2H + + # ---- Outer draw (one batched forward over all active states) ---- + logits0 = self._forward_ar_batched(code, code_padding_mask, seqs) # (A, vocab) + td_first_idx = [i for i, s in enumerate(seqs) + if sampling_direction == 'td' and len(s) == 1] + tok0 = _sample(logits0, td_first_idx or None) # list[int] + + # ---- Classify each active row (pure CPU) ---- + is_bb = [tok0[i] == MOL_START or seqs[i][-1] == MOL_START for i in range(len(seqs))] + is_rxn = [(not is_bb[i]) and tok0[i] >= RXN_MIN for i in range(len(seqs))] + + # ---- Batched inner BB-SMILES loop (lane-masked, CPU token bookkeeping) ---- + bb_rows = [i for i in range(len(seqs)) if is_bb[i]] + bb_tokens: dict[int, list[int]] = {i: [tok0[i]] for i in bb_rows} + ctx: dict[int, list[int]] = {i: list(seqs[i]) for i in bb_rows} + prev: dict[int, int] = {i: tok0[i] for i in bb_rows} + + def _cont(i): + return prev[i] != MOL_END and len(ctx[i]) < max_len - 2 + + active_bb = [i for i in bb_rows if _cont(i)] + while active_bb: + for i in active_bb: + ctx[i].append(prev[i]) # append previous token to context + logits = self._forward_ar_batched(code, code_padding_mask, [ctx[i] for i in active_bb]) + samp = _sample(logits) # list[int] + for j, i in enumerate(active_bb): + prev[i] = samp[j] + bb_tokens[i].append(samp[j]) + active_bb = [i for i in bb_rows if _cont(i)] + + # ---- Per-state chemistry (CPU): assemble PredictResult per active row ---- + for i in range(len(seqs)): + k = act[i] + if is_bb[i]: + smiles = decode_smiles(bb_tokens[i]) + item = get_reactants(smiles, fpindex=self._fpindex, topk=100, + use_edit_distance=use_edit_distance) + results[k] = PredictResult('BB' if item is not None else 'ABORTED', item) + elif is_rxn[i]: + reaction_logits = logits0[i, RXN_MIN:RXN_MAX + 1] # (115,) + item = get_reactions(reaction_logits, rxn_matrix=self._rxn_matrix) + results[k] = PredictResult('RXN', item) + else: + results[k] = PredictResult('END', None) + + return results + def _evolve_ar_singlestep( self, gpu_lock: Lock | None = None, @@ -224,6 +362,11 @@ def _evolve_ar_singlestep( ) -> None: if len(self._active) == 0: return + + # Batched prediction is atomic: do not start it after the deadline. + # 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 feat_list = [ featurize_stack( @@ -237,26 +380,48 @@ def _evolve_ar_singlestep( 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 + # Build per-state token prefixes (with the BU MOL_START seed) once. + tokens_list: list[list[int]] = [] + for feat in feat_list: + tokens = feat['tokens'].tolist() + if sampling_direction == 'bu' and len(tokens) == 1: # for BU + tokens.append(int(TokenType.MOL_START)) + tokens_list.append(tokens) + + # Batched path: one set of batched forwards for all active states. + results_list = ( + self._predict_ar_batched( + code=code, code_padding_mask=code_padding_mask, + tokens_list=tokens_list, sampling_direction=sampling_direction) + if self._batched_ar else None + ) + finished: list[State] = [] next: list[State] = [] - for feat, base_state in zip(feat_list, self._active): - if time_limit is not None and time_limit.exceeded(): + for k, base_state in enumerate(self._active): + if results_list is None and time_limit is not None and time_limit.exceeded(): break - - tokens = feat['tokens'].to(self.device) - if sampling_direction == 'bu' and len(tokens) == 1: # for BU - tokens = torch.hstack([tokens, torch.tensor([TokenType.MOL_START]).to(self.device)]) - - result = self._predict_ar( - code=code, - code_padding_mask=code_padding_mask, - tokens=tokens, - sampling_direction=sampling_direction - ) + + if results_list is not None: + result = results_list[k] + else: + tokens = torch.tensor(tokens_list[k], dtype=torch.long, device=self.device) + result = self._predict_ar( + code=code, + code_padding_mask=code_padding_mask, + tokens=tokens, + sampling_direction=sampling_direction, + ) sampled_type, sampled_item = result.sampled_type, result.sampled_item - + for i in range(self._factor): if sampling_direction == 'td': if sampled_type == 'END': diff --git a/tests/test_batched_ar.py b/tests/test_batched_ar.py new file mode 100644 index 0000000..79b6821 --- /dev/null +++ b/tests/test_batched_ar.py @@ -0,0 +1,345 @@ +"""Collected regression tests for batched autoregressive decoding.""" + +import os +import sys +import types +from contextlib import contextmanager + +import pytest +import torch +from omegaconf import OmegaConf +from torch.nn.attention import SDPBackend, sdpa_kernel + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import reasyn.sampler.sampler as sampler_module +from reasyn.chem.featurize import TokenType +from reasyn.chem.mol import Molecule +from reasyn.models.reasyn import ReaSyn +from reasyn.sampler.sampler import Sampler +from reasyn.utils.sample_utils import PredictResult, State + + +CKPT = os.environ.get( + "REASYN_AR_CKPT", + "/workspace/shared/ckpt/nv-reasyn-ar-166m-v2.ckpt", +) +TARGET = "C[NH+](C=C1SC(=O)NC1=S)C1(CO)CCC(OCC2CCCCC2)C1" +TEMPERATURE = 0.1 +TOPK = 16 + + +@contextmanager +def controlled_math(): + """Make batch-shape comparisons use the same arithmetic path.""" + old_matmul = torch.backends.cuda.matmul.allow_tf32 + old_cudnn = torch.backends.cudnn.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + with sdpa_kernel(SDPBackend.MATH): + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_matmul + torch.backends.cudnn.allow_tf32 = old_cudnn + + +@pytest.fixture(scope="module") +def ar_context(): + if not torch.cuda.is_available(): + pytest.skip("AR checkpoint test requires CUDA") + if not os.path.exists(CKPT): + pytest.skip(f"AR checkpoint not found: {CKPT}") + + checkpoint = torch.load(CKPT, map_location="cpu") + cfg = OmegaConf.create(checkpoint["hyper_parameters"]["config"]) + model = ReaSyn(cfg.model).cuda() + model.load_state_dict({k[6:]: v for k, v in checkpoint["state_dict"].items()}) + model.eval() + assert model.max_len >= 512 + + smiles = Molecule(TARGET).tokenize_csmiles()[None].cuda() + with torch.inference_mode(): + code, code_padding_mask = model.encoder(smiles) + + sampler = Sampler.__new__(Sampler) + sampler.model = model + sampler.device = torch.device("cuda") + sampler._bucket_ar = True + sampler._use_bf16 = False + sampler._fpindex = object() + sampler._rxn_matrix = object() + return sampler, code, code_padding_mask + + +def make_prefix(length: int, seed: int) -> list[int]: + generator = torch.Generator().manual_seed(seed) + prefix = torch.randint(4, int(TokenType.RXN_MIN), (length,), generator=generator).tolist() + prefix[0] = int(TokenType.START) + return prefix + + +def reaction_order(logits: torch.Tensor) -> torch.Tensor: + return logits[:, int(TokenType.RXN_MIN):int(TokenType.RXN_MAX) + 1].topk(TOPK).indices + + +@pytest.mark.parametrize("lmax", [1, 63, 64, 65, 511, 512]) +@torch.inference_mode() +def test_production_forward_matches_per_row(ar_context, monkeypatch, lmax): + sampler, code, code_padding_mask = ar_context + mid = max(1, lmax // 2) + lengths = sorted({1, mid, lmax}) + prefixes = {length: make_prefix(length, 1000 + lmax + length) for length in lengths} + + with controlled_math(): + reference = { + length: sampler.model.sample( + code=code, + code_padding_mask=code_padding_mask, + tokens=torch.tensor(prefix, device=sampler.device)[None], + token_padding_mask=None, + )[0].float() + for length, prefix in prefixes.items() + } + + original_sample = sampler.model.sample + production_calls = [] + + def record_sample(**kwargs): + production_calls.append(kwargs) + return original_sample(**kwargs) + + monkeypatch.setattr(sampler.model, "sample", record_sample) + + for batch_size in (1, 3, 9): + if batch_size == 1: + row_lengths = [lmax] + else: + row_lengths = ([1, mid, lmax] * 3)[:batch_size] + seqs = [prefixes[length] for length in row_lengths] + ref = torch.stack([reference[length] for length in row_lengths]) + batched = sampler._forward_ar_batched(code, code_padding_mask, seqs) + + call = production_calls[-1] + expected_batch = 1 << max(0, batch_size - 1).bit_length() + expected_length = min(((lmax + 63) // 64) * 64, sampler.model.max_len) + assert call["token_padding_mask"] is None + assert tuple(call["tokens"].shape) == (expected_batch, expected_length) + + ref_probs = torch.softmax(ref / TEMPERATURE, dim=-1) + batched_probs = torch.softmax(batched / TEMPERATURE, dim=-1) + tv = 0.5 * (ref_probs - batched_probs).abs().sum(dim=-1) + assert tv.max().item() < 1e-4 + + ref_top = reaction_order(ref) + batched_top = reaction_order(batched) + assert torch.equal(batched_top, ref_top) + + +@torch.inference_mode() +def test_predict_ar_batched_uses_production_reaction_order(ar_context, monkeypatch): + sampler, code, code_padding_mask = ar_context + seqs = [make_prefix(length, 2000 + length) for length in (1, 63, 65)] + + def force_reaction(probs, num_samples): + return torch.full( + (probs.shape[0], num_samples), + int(TokenType.RXN_MIN), + dtype=torch.long, + device=probs.device, + ) + + monkeypatch.setattr(torch, "multinomial", force_reaction) + monkeypatch.setattr( + sampler_module, + "get_reactions", + lambda logits, rxn_matrix: logits.topk(TOPK).indices.tolist(), + ) + + with controlled_math(): + expected = reaction_order(sampler._forward_ar_batched(code, code_padding_mask, seqs)) + results = sampler._predict_ar_batched( + code, + code_padding_mask, + seqs, + sampling_direction="td", + ) + + assert [result.sampled_type for result in results] == ["RXN"] * len(seqs) + assert [result.sampled_item for result in results] == expected.tolist() + + +def test_predict_ar_batched_lane_control_flow(monkeypatch): + sampler = Sampler.__new__(Sampler) + sampler.model = types.SimpleNamespace(max_len=8) + sampler._fpindex = object() + sampler._rxn_matrix = object() + vocab_size = int(max(TokenType)) + + base_lengths = {10: 2, 11: 3, 12: 2, 13: 6} + + def forward(self, code, code_padding_mask, seqs): + logits = torch.full((len(seqs), vocab_size), -1000.0) + for row, seq in enumerate(seqs): + marker, length = seq[0], len(seq) + if marker == 10: + token = 4 if length == base_lengths[marker] else int(TokenType.MOL_END) + elif marker == 11: + token = 4 if length == 3 else 5 if length == 4 else int(TokenType.MOL_END) + elif marker == 12: + token = int(TokenType.MOL_END) + else: + token = 6 + logits[row, token] = 1000.0 + return logits + + sampler._forward_ar_batched = types.MethodType(forward, sampler) + decoded = [] + monkeypatch.setattr( + sampler_module, + "decode_smiles", + lambda tokens: decoded.append([int(token) for token in tokens]) or "ok", + ) + monkeypatch.setattr(sampler_module, "get_reactants", lambda *args, **kwargs: ["reactant"]) + + results = sampler._predict_ar_batched( + None, + None, + [ + [10, int(TokenType.MOL_START)], + [11, 7, int(TokenType.MOL_START)], + [12, int(TokenType.MOL_START)], + [13, 4, 4, 4, 4, int(TokenType.MOL_START)], + ], + ) + + assert [result.sampled_type for result in results] == ["BB"] * 4 + assert decoded == [[4, 3], [4, 5, 3], [3], [6]] + + +class RecordingLock: + def __init__(self): + self.acquire_calls = 0 + self.release_calls = 0 + self.locked = False + + def acquire(self): + self.acquire_calls += 1 + self.locked = True + + def release(self): + self.release_calls += 1 + self.locked = False + + +@pytest.mark.parametrize( + ("deadline_answers", "expected_acquires", "expected_releases"), + [([True], 0, 0), ([False, True], 1, 1)], +) +def test_expired_deadline_skips_batched_prediction( + deadline_answers, + expected_acquires, + expected_releases, +): + sampler = Sampler.__new__(Sampler) + state = State() + sampler._active = [state] + sampler._batched_ar = True + lock = RecordingLock() + answers = iter(deadline_answers) + time_limit = types.SimpleNamespace(exceeded=lambda: next(answers)) + + sampler._predict_ar_batched = lambda *args, **kwargs: pytest.fail( + "batched prediction ran after the deadline" + ) + sampler._evolve_ar_singlestep( + gpu_lock=lock, + time_limit=time_limit, + sampling_direction="td", + ) + + assert lock.acquire_calls == expected_acquires + assert lock.release_calls == expected_releases + assert not lock.locked + assert sampler._active == [state] + + +def test_legacy_deadline_check_remains_per_state(): + sampler = Sampler.__new__(Sampler) + sampler._active = [State(), State()] + sampler._finished = [] + sampler._aborted = [] + sampler._batched_ar = False + sampler._factor = 1 + sampler._max_active_states = 2 + sampler.device = torch.device("cpu") + sampler.__dict__["code"] = (None, None) + lock = RecordingLock() + answers = iter([False, True]) + time_limit = types.SimpleNamespace(exceeded=lambda: next(answers)) + calls = [] + + def predict(self, **kwargs): + calls.append(kwargs["tokens"].tolist()) + return PredictResult("ABORTED", None) + + sampler._predict_ar = types.MethodType(predict, sampler) + sampler._evolve_ar_singlestep( + gpu_lock=lock, + time_limit=time_limit, + sampling_direction="td", + ) + + assert calls == [[int(TokenType.START)]] + assert lock.acquire_calls == lock.release_calls == 1 + assert not lock.locked + + +def test_batched_evolve_keeps_prefixes_on_cpu_lists(): + sampler = Sampler.__new__(Sampler) + sampler._active = [State()] + sampler._finished = [] + sampler._aborted = [] + sampler._batched_ar = True + sampler._factor = 1 + sampler._max_active_states = 1 + sampler.device = torch.device("cpu") + sampler.__dict__["code"] = (None, None) + captured = [] + + def predict(self, code, code_padding_mask, tokens_list, sampling_direction): + captured.extend(tokens_list) + return [PredictResult("ABORTED", None)] + + sampler._predict_ar_batched = types.MethodType(predict, sampler) + sampler._evolve_ar_singlestep(sampling_direction="bu") + + assert captured == [[int(TokenType.START), int(TokenType.MOL_START)]] + assert all(isinstance(token, int) for token in captured[0]) + + +def test_bf16_is_opt_in(monkeypatch): + class FakeModel: + def __init__(self, model_type): + self.model_type = model_type + + def parameters(self): + yield torch.empty(0) + + molecule = types.SimpleNamespace( + tokenize_csmiles=lambda: torch.tensor([4]), + csmiles="C", + ) + models = [FakeModel("autoregressive"), FakeModel("editflow")] + + monkeypatch.delenv("REASYN_BF16", raising=False) + sampler = Sampler(object(), object(), molecule, models) + assert not sampler._use_bf16 + + monkeypatch.setenv("REASYN_BF16", "1") + sampler = Sampler(object(), object(), molecule, models) + assert sampler._use_bf16 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"]))