diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..99cc475
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,10 @@
+__pycache__/
+*.pyc
+.venv/
+/data/
+cascadia_*/
+checkpoint_*/
+logs/
+*.egg-info/
+build/
+dist/
diff --git a/cascadia/__pycache__/__init__.cpython-310.pyc b/cascadia/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 32201d0..0000000
Binary files a/cascadia/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/__pycache__/augment.cpython-310.pyc b/cascadia/__pycache__/augment.cpython-310.pyc
deleted file mode 100644
index fc024ed..0000000
Binary files a/cascadia/__pycache__/augment.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/__pycache__/cascadia.cpython-310.pyc b/cascadia/__pycache__/cascadia.cpython-310.pyc
deleted file mode 100644
index 18e82f9..0000000
Binary files a/cascadia/__pycache__/cascadia.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/__pycache__/model.cpython-310.pyc b/cascadia/__pycache__/model.cpython-310.pyc
deleted file mode 100644
index 9dbaac4..0000000
Binary files a/cascadia/__pycache__/model.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/__pycache__/version.cpython-310.pyc b/cascadia/__pycache__/version.cpython-310.pyc
deleted file mode 100644
index e0d2466..0000000
Binary files a/cascadia/__pycache__/version.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/augment.py b/cascadia/augment.py
index 34046f7..598a1c3 100644
--- a/cascadia/augment.py
+++ b/cascadia/augment.py
@@ -133,27 +133,32 @@ def write_asf(outfile, prec_to_spec, scan_width, max_pep_length, max_charge):
         ms1_rts = ms1_rts[sorted_ms1_rt_idxs]
         ms1_scans = ms1_scans[sorted_ms1_rt_idxs]
         
-        for charge in range(1,max_charge+1):
-            count += 1
+        # The peak list is identical for every candidate charge state, so
+        # format it once and reuse it.
+        peak_lines = []
+        for scan, cur_rt in zip(scans,rts):
+            for mz, intensity in scan:
+                peak_lines.append(f"{mz}\t{intensity}\t{cur_rt}\t2\n")
 
-            out.write("BEGIN IONS\n")
-            out.write(f"TITLE={str(count)}\n")
-            out.write(f"PEPMASS={prec}\n")
-            out.write(f"CHARGE={charge}\n")
-            out.write(f"SCAN={count}\n")
-            out.write(f"RT={rt}\n")
-            out.write(f"SEQ={'K'*max_pep_length}\n")
-            
-            for scan, cur_rt in zip(scans,rts):
-                for mz, intensity in scan:
-                    out.write(f"{mz}\t{intensity}\t{cur_rt}\t2\n")
-            
-            for scan, cur_rt in zip(ms1_scans, ms1_rts):
-                for mz, intensity in scan:
-                    if np.abs(mz - prec) < window_width + 1:
-                        out.write(f"{mz}\t{intensity}\t{cur_rt}\t1\n")
+        for scan, cur_rt in zip(ms1_scans, ms1_rts):
+            for mz, intensity in scan:
+                if np.abs(mz - prec) < window_width + 1:
+                    peak_lines.append(f"{mz}\t{intensity}\t{cur_rt}\t1\n")
+        peak_block = ''.join(peak_lines)
 
-            out.write("END IONS\n")   
+        for charge in range(1,max_charge+1):
+            count += 1
+            out.write(
+                "BEGIN IONS\n"
+                f"TITLE={str(count)}\n"
+                f"PEPMASS={prec}\n"
+                f"CHARGE={charge}\n"
+                f"SCAN={count}\n"
+                f"RT={rt}\n"
+                f"SEQ={'K'*max_pep_length}\n"
+                + peak_block +
+                "END IONS\n"
+            )
     out.close()
 
 def augment_spectra(mzml_file, temp_path, top_n = 150, scan_width = 1, max_pep_length = 30, max_charge = 3):
diff --git a/cascadia/cascadia.py b/cascadia/cascadia.py
index aae6554..6814a79 100755
--- a/cascadia/cascadia.py
+++ b/cascadia/cascadia.py
@@ -3,7 +3,7 @@
 from .depthcharge.tokenizers import PeptideTokenizer
 import torch
 import numpy as np
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import os
 import sys
 import argparse
@@ -11,11 +11,78 @@
 from .depthcharge.utils import *
 from .model import AugmentedSpec2Pep
 from .augment import *
+from .utils import write_results
 from datetime import datetime
+import time
 import warnings
 import json
 warnings.filterwarnings('ignore')
 
+MODEL_HPARAMS = dict(
+    d_model=512,
+    n_layers=9,
+    n_head=8,
+    dim_feedforward=1024,
+    dropout=0,
+)
+
+
+def select_accelerator(requested='auto'):
+  """Pick the Lightning accelerator: CUDA, then Apple MPS, then CPU.
+
+  ``requested`` may be 'auto', 'gpu' (CUDA), 'mps', or 'cpu'.
+  """
+  if requested != 'auto':
+    print(f'Using requested device: {requested}')
+    return requested
+  if torch.cuda.is_available():
+    print('GPU found')
+    return 'gpu'
+  if torch.backends.mps.is_available():
+    print('Apple MPS GPU found')
+    return 'mps'
+  print('No GPU found - running on cpu')
+  return 'cpu'
+
+
+def dataloader_workers(n_workers):
+  """Number of DataLoader worker processes to use.
+
+  macOS and Windows start worker processes with ``spawn``, which requires
+  pickling the dataset; the tokenizer's numba typed dictionaries are not
+  picklable, so spectra are loaded in the main process there. Linux forks
+  the workers instead. Reading batches from the HDF5 index is cheap compared
+  to inference, so this costs little.
+  """
+  return n_workers if sys.platform.startswith('linux') else 0
+
+
+def load_tokenizer(mods):
+  if mods == 'mskb':
+    return PeptideTokenizer.from_massivekb(reverse=False, replace_isoleucine_with_leucine=True)
+  with open(mods, 'r') as f:
+    proforma = json.load(f)
+  return PeptideTokenizer.from_proforma(proforma, reverse=False, replace_isoleucine_with_leucine=True)
+
+
+def load_model(model_ckpt_path, tokenizer, rt_width, max_charge, lr=1e-4):
+  """Build an AugmentedSpec2Pep and load weights from a Lightning checkpoint.
+
+  Cascadia checkpoints contain numpy scalars in their optimizer/scheduler
+  state, which the ``weights_only=True`` default of ``torch.load`` (PyTorch
+  >= 2.6) refuses to unpickle, so the weights are loaded explicitly here.
+  """
+  checkpoint = torch.load(model_ckpt_path, map_location='cpu', weights_only=False)
+  model = AugmentedSpec2Pep(
+      rt_width=rt_width,
+      tokenizer=tokenizer,
+      max_charge=max_charge,
+      lr=lr,
+      **MODEL_HPARAMS,
+  )
+  model.load_state_dict(checkpoint['state_dict'])
+  return model
+
 
 def sequence():
   parser=argparse.ArgumentParser()
@@ -23,10 +90,11 @@ def sequence():
   parser.add_argument("model", type= str, help="A path to a trained Cascadia model checkpoint.")
   parser.add_argument("-o", "--outfile", type= str, default= 'cascadia_results', help="Output file for inference.")
   parser.add_argument("-t", "--score_threshold", type= float, default= 0.8, help="Score threshold for Cascadia predictions.")
-  parser.add_argument("-b", "--batch_size", type= int, default= 32, help="Number of spectra to include in a batch.")
+  parser.add_argument("-b", "--batch_size", type= int, default= 128, help="Number of spectra to include in a batch. Larger batches are faster on a GPU as long as they fit in memory.")
   parser.add_argument("-w", "--width", type= int, default= 2, help="Number of adjacent scans to use when constructing each augmented spectrum.")
   parser.add_argument("-c", "--max_charge", type= int, default= 4, help="Maximum precursor charge state to consider")
   parser.add_argument("-p", "--modifications", type= str, default= 'mskb', help="A path to a json file containing a list of the PTMs to consider.")
+  parser.add_argument("-d", "--device", type= str, default= 'auto', choices=['auto', 'gpu', 'mps', 'cpu'], help="Device to run inference on. 'auto' picks CUDA, then Apple MPS, then CPU.")
   
   args = parser.parse_args(args=sys.argv[2:])
 
@@ -43,44 +111,29 @@ def sequence():
   os.mkdir(temp_path)
   train_index_filename = temp_path + "/index.hdf5"
 
+  t_start = time.perf_counter()
   print("Augmenting spectra from:", spectrum_file)
   asf_file, isolation_window_size, cycle_time = augment_spectra(spectrum_file, temp_path, max_charge=max_charge)
-  
-  if mods == 'mskb':
-    tokenizer = PeptideTokenizer.from_massivekb(reverse=False, replace_isoleucine_with_leucine=True)
-  else:
-    with open('ptms.json', 'r') as f:
-      proforma = json.load(f)
-    tokenizer = PeptideTokenizer.from_proforma(proforma, reverse=False, replace_isoleucine_with_leucine=True)
-  
+  t_augment = time.perf_counter()
+
+  tokenizer = load_tokenizer(mods)
+
+  device = select_accelerator(args.device)
+
   train_dataset = AnnotatedSpectrumDataset(tokenizer, asf_file, index_path=train_index_filename, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
-  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=4, pin_memory=True)
+  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=dataloader_workers(4), pin_memory=(device == 'gpu'))
 
   if os.path.exists(asf_file):
       os.remove(asf_file)
+  t_index = time.perf_counter()
 
-  model = AugmentedSpec2Pep.load_from_checkpoint(
-      model_ckpt_path,
-      d_model = 512,
-      n_layers = 9,
-      n_head = 8,
-      dim_feedforward = 1024,
-      dropout = 0,
-      rt_width = 2,
-      tokenizer=tokenizer,
-      max_charge=10,
-  )
+  model = load_model(model_ckpt_path, tokenizer, rt_width=2, max_charge=10)
 
   print("Running inference on augmented spectra from:", spectrum_file)
-  if torch.cuda.is_available():
-    device = 'gpu'
-    print('GPU found')
-  else:
-    device = 'cpu'
-    print(f'No GPU found - running inference on cpu')
 
-  trainer = pl.Trainer(max_epochs=50, log_every_n_steps=1, accelerator=device, devices = 1)
+  trainer = pl.Trainer(max_epochs=50, log_every_n_steps=1, accelerator=device, devices = 1, logger=False)
   preds = trainer.predict(model, dataloaders=train_loader)
+  t_predict = time.perf_counter()
 
   print("Writing results to:", results_file + '.ssl')
 
@@ -88,6 +141,12 @@ def sequence():
   
   os.remove(train_index_filename)
   os.rmdir(temp_path)
+  t_end = time.perf_counter()
+  print(
+      f"Done in {t_end - t_start:.1f}s "
+      f"(augmentation {t_augment - t_start:.1f}s, indexing {t_index - t_augment:.1f}s, "
+      f"model loading + inference {t_predict - t_index:.1f}s, writing results {t_end - t_predict:.1f}s)"
+  )
   
   return parser
 
@@ -116,11 +175,8 @@ def train():
   lr = args.learning_rate
   mods = args.modifications
   
-  if torch.cuda.is_available():
-    device = 'gpu'
-    print('GPU found')
-  else:
-    device = 'cpu'
+  device = select_accelerator()
+  if device == 'cpu':
     print("No GPU Available - training on CPU will be extremely slow!")
 
   print("Training on spectra from:", train_spectrum_file)
@@ -136,13 +192,8 @@ def train():
   if os.path.exists(val_index_filename):
     os.remove(val_index_filename)
 
-  if mods == 'mskb':
-    tokenizer = PeptideTokenizer.from_massivekb(reverse=False, replace_isoleucine_with_leucine=True)
-  else:
-    with open(mods, 'r') as f:
-      proforma = json.load(f)
-    tokenizer = PeptideTokenizer.from_proforma(proforma, reverse=False, replace_isoleucine_with_leucine=True)
-  
+  tokenizer = load_tokenizer(mods)
+
   if '.hdf5' in train_spectrum_file:
     train_dataset = AnnotatedSpectrumDataset(tokenizer, index_path=train_spectrum_file, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
     val_dataset = AnnotatedSpectrumDataset(tokenizer, index_path=val_spectrum_file, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
@@ -150,37 +201,22 @@ def train():
     train_dataset = AnnotatedSpectrumDataset(tokenizer, train_spectrum_file, index_path=train_index_filename, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm]) 
     val_dataset = AnnotatedSpectrumDataset(tokenizer, val_spectrum_file, index_path=val_index_filename, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
   
-  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=10, pin_memory=True, shuffle=True)
-  val_loader = val_dataset.loader(batch_size=batch_size, num_workers=10, pin_memory=True)
+  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=dataloader_workers(10), pin_memory=(device == 'gpu'), shuffle=True)
+  val_loader = val_dataset.loader(batch_size=batch_size, num_workers=dataloader_workers(10), pin_memory=(device == 'gpu'))
 
   if model_ckpt_path is None:
     print("Training model from scratch") 
     model = AugmentedSpec2Pep(
-        d_model = 512,
-        n_layers = 9,
-        n_head = 8,
-        dim_feedforward = 1024,
-        dropout = 0,
         rt_width = augmentation_width,
         tokenizer=tokenizer,
         max_charge=max_charge,
-        lr=lr
+        lr=lr,
+        **MODEL_HPARAMS,
     )
     
   else:
     print("Loading model from checkpoint:", model_ckpt_path) 
-    model = AugmentedSpec2Pep.load_from_checkpoint(
-      model_ckpt_path, 
-      d_model = 512,
-      n_layers = 9,
-      n_head = 8,
-      dim_feedforward = 1024,
-      dropout = 0,
-      rt_width = 2,
-      tokenizer=tokenizer,
-      max_charge=max_charge,
-      lr=lr
-    )
+    model = load_model(model_ckpt_path, tokenizer, rt_width=2, max_charge=max_charge, lr=lr)
 
   tb_logger = pl_loggers.TensorBoardLogger(save_dir="logs/")
   ckpt_callback = pl.callbacks.ModelCheckpoint(
diff --git a/cascadia/depthcharge/__pycache__/__init__.cpython-310.pyc b/cascadia/depthcharge/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 3bd40ad..0000000
Binary files a/cascadia/depthcharge/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/__pycache__/constants.cpython-310.pyc b/cascadia/depthcharge/__pycache__/constants.cpython-310.pyc
deleted file mode 100644
index 6b2e7ff..0000000
Binary files a/cascadia/depthcharge/__pycache__/constants.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/__pycache__/feedforward.cpython-310.pyc b/cascadia/depthcharge/__pycache__/feedforward.cpython-310.pyc
deleted file mode 100644
index 5ba382e..0000000
Binary files a/cascadia/depthcharge/__pycache__/feedforward.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/__pycache__/primitives.cpython-310.pyc b/cascadia/depthcharge/__pycache__/primitives.cpython-310.pyc
deleted file mode 100644
index 33cc6a4..0000000
Binary files a/cascadia/depthcharge/__pycache__/primitives.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/__pycache__/utils.cpython-310.pyc b/cascadia/depthcharge/__pycache__/utils.cpython-310.pyc
deleted file mode 100644
index 2a11328..0000000
Binary files a/cascadia/depthcharge/__pycache__/utils.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/__pycache__/version.cpython-310.pyc b/cascadia/depthcharge/__pycache__/version.cpython-310.pyc
deleted file mode 100644
index 734f4af..0000000
Binary files a/cascadia/depthcharge/__pycache__/version.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/data/__pycache__/__init__.cpython-310.pyc b/cascadia/depthcharge/data/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index da309c8..0000000
Binary files a/cascadia/depthcharge/data/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/data/__pycache__/parsers.cpython-310.pyc b/cascadia/depthcharge/data/__pycache__/parsers.cpython-310.pyc
deleted file mode 100644
index 5bed688..0000000
Binary files a/cascadia/depthcharge/data/__pycache__/parsers.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/data/__pycache__/peptide_datasets.cpython-310.pyc b/cascadia/depthcharge/data/__pycache__/peptide_datasets.cpython-310.pyc
deleted file mode 100644
index deb22b9..0000000
Binary files a/cascadia/depthcharge/data/__pycache__/peptide_datasets.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/data/__pycache__/preprocessing.cpython-310.pyc b/cascadia/depthcharge/data/__pycache__/preprocessing.cpython-310.pyc
deleted file mode 100644
index f395048..0000000
Binary files a/cascadia/depthcharge/data/__pycache__/preprocessing.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/data/__pycache__/spectrum_datasets.cpython-310.pyc b/cascadia/depthcharge/data/__pycache__/spectrum_datasets.cpython-310.pyc
deleted file mode 100644
index 01e8fd2..0000000
Binary files a/cascadia/depthcharge/data/__pycache__/spectrum_datasets.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/data/parsers.py b/cascadia/depthcharge/data/parsers.py
index 9137abe..026327f 100644
--- a/cascadia/depthcharge/data/parsers.py
+++ b/cascadia/depthcharge/data/parsers.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import logging
+import re
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable
 from os import PathLike
@@ -390,6 +391,87 @@ def parse_spectrum(self, spectrum: dict) -> MassSpectrum:
     def open(self) -> Iterable[dict]:
         pass
     
+    @staticmethod
+    def _is_peak_line(fields: list[str]) -> bool:
+        """Whether a whitespace-split line is a 4-column peak row."""
+        return len(fields) == 4 and "." in fields[0] and "." in fields[1]
+
+    def _parse_block(self, lines: list[str]) -> None:
+        """Parse the lines between BEGIN IONS and END IONS."""
+        precursor_mz = None
+        precursor_charge = 0
+        scan_id = 0
+        label = None
+        rt = None
+
+        # Header lines come first; everything after the first peak row is
+        # parsed as a block of numbers in one go.
+        first_peak = len(lines)
+        for idx, line in enumerate(lines):
+            fields = line.split()
+            if not fields:
+                continue
+
+            if self._is_peak_line(fields):
+                first_peak = idx
+                break
+
+            key = fields[0]
+            if key.startswith("PEPMASS"):
+                precursor_mz = float(key.split("=")[1])
+            elif key.startswith("CHARGE"):
+                precursor_charge = float(
+                    re.match(r"\d+", key.split("=")[1]).group()
+                )
+            elif key.startswith("SCAN"):
+                scan_id = float(key.split("=")[1])
+            elif key.startswith("SEQ"):
+                label = key.split("=")[1]
+            elif key.startswith("RT="):
+                rt = float(key.split("=")[1])
+
+        peaks = np.array("".join(lines[first_peak:]).split(), dtype=np.float64)
+        if peaks.size % 4:
+            # Something other than 4-column peak rows; fall back to the
+            # conservative row-by-row filter.
+            rows = [
+                fields
+                for fields in (line.split() for line in lines[first_peak:])
+                if self._is_peak_line(fields)
+            ]
+            peaks = np.array(rows, dtype=np.float64).reshape(-1, 4)
+        else:
+            peaks = peaks.reshape(-1, 4)
+
+        spectrum = MassSpectrum(
+            filename=str(self.path),
+            scan_id=scan_id,
+            mz=np.ascontiguousarray(peaks[:, 0]),
+            intensity=np.ascontiguousarray(peaks[:, 1]),
+            rt=np.ascontiguousarray(peaks[:, 2]),
+            level=np.ascontiguousarray(peaks[:, 3]),
+            precursor_mz=precursor_mz,
+            precursor_charge=precursor_charge,
+            label=label,
+            fragment_label=np.zeros(peaks.shape[0], dtype=np.int64),
+        )
+
+        if self.preprocessing_fn is not None:
+            for processor in self.preprocessing_fn:
+                spectrum = processor(spectrum)
+
+        self.global_rts.append(rt)
+        self.mz_arrays.append(spectrum.mz)
+        self.intensity_arrays.append(spectrum.intensity)
+        self.rt_arrays.append(spectrum.rt)
+        self.level_arrays.append(spectrum.level)
+        self.fragment_arrays.append(spectrum.fragment_label)
+        self.precursor_mz.append(spectrum.precursor_mz)
+        self.precursor_charge.append(spectrum.precursor_charge)
+        self.scan_id.append(_parse_scan_id(spectrum.scan_id))
+        if self.annotations is not None:
+            self.annotations.append(spectrum.label)
+
     def read(self) -> BaseParser:
         """Read the ms data file.
 
@@ -397,87 +479,17 @@ def read(self) -> BaseParser:
         -------
         Self
         """
-        n_skipped = 0
+        block = None
         with open(self.path) as asf:
             for line in tqdm(asf, desc=str(self.path), unit="lines"):
-                line = line.split()
-                if line[0] == 'BEGIN':
-                    mzs = []
-                    intensities = []
-                    rts = []
-                    levels = []
-                    precursor_mz = None
-                    precursor_charge = 0
-                    scan_id = 0
-                    label = None
-                    rt = None
-                elif line[0][:7] == 'PEPMASS':
-                    precursor_mz = float(line[0].split('=')[1])
-                elif line[0][:6] == 'CHARGE':
-                    precursor_charge = float(line[0][7])
-                elif line[0][:4] == 'SCAN':
-                    scan_id = float(line[0].split('=')[1])
-                elif line[0][:3] == 'SEQ':
-                    label = line[0].split('=')[1]
-                elif line[0][:3] == 'RT=':
-                    rt = float(line[0].split('=')[1])
-                elif len(line) == 4 and '.' in line[0] and '.' in line[1]:
-                    mzs.append(float(line[0]))
-                    intensities.append(float(line[1]))
-                    rts.append(float(line[2]))
-                    levels.append(float(line[3]))
-                elif line[0] == 'END':
-
-                    seq = ''.join(filter(str.isalpha, label))
-                    fragments = []
-                    fragment_labels = []
-                    #for i in range(1, len(seq)):
-                        # fragments.append(mass.fast_mass(seq[:i], ion_type='b', charge=1) + 57.021 * seq.count('C'))
-                        # fragments.append(mass.fast_mass(seq[i:], ion_type='y', charge=1) + 57.021 * seq.count('C'))
-                        # fragments.append(mass.fast_mass(seq[:i], ion_type='b', charge=2) + 57.021/2 * seq.count('C'))
-                        # fragments.append(mass.fast_mass(seq[i:], ion_type='y', charge=2) + 57.021/2 * seq.count('C'))
-
-                    for peak in mzs:
-                        peak_label = 0
-                    #     for frag in fragments:
-                    #         if np.abs(peak - frag) < (20 * frag / 1e6):
-                    #             peak_label = 1
-                    #             break
-                        fragment_labels.append(peak_label)
-
-                    spectrum = MassSpectrum(
-                        filename=str(self.path),
-                        scan_id=scan_id,
-                        mz=mzs,
-                        intensity=intensities,
-                        rt=np.array(rts),
-                        level=np.array(levels),
-                        precursor_mz=precursor_mz,
-                        precursor_charge=precursor_charge,
-                        label=label,
-                        fragment_label=np.array(fragment_labels)
-                    )
-                    
-                    if self.preprocessing_fn is not None:
-                        for processor in self.preprocessing_fn:
-                            spectrum = processor(spectrum)
-
-                    self.global_rts.append(rt)
-                    self.mz_arrays.append(spectrum.mz)
-                    self.intensity_arrays.append(spectrum.intensity)
-                    self.rt_arrays.append(spectrum.rt)
-                    self.level_arrays.append(spectrum.level)
-                    self.fragment_arrays.append(spectrum.fragment_label)
-                    self.precursor_mz.append(spectrum.precursor_mz)
-                    self.precursor_charge.append(spectrum.precursor_charge)
-                    self.scan_id.append(_parse_scan_id(spectrum.scan_id))
-                    if self.annotations is not None:
-                        self.annotations.append(spectrum.label)
-
-        if n_skipped:
-            LOGGER.warning(
-                "Skipped %d spectra with invalid precursor info", n_skipped
-            )
+                if line.startswith("BEGIN IONS"):
+                    block = []
+                elif line.startswith("END IONS"):
+                    if block is not None:
+                        self._parse_block(block)
+                    block = None
+                elif block is not None:
+                    block.append(line)
 
         self.precursor_mz = np.array(self.precursor_mz, dtype=np.float64)
         self.precursor_charge = np.array(
diff --git a/cascadia/depthcharge/encoders/__pycache__/__init__.cpython-310.pyc b/cascadia/depthcharge/encoders/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index ff5c07f..0000000
Binary files a/cascadia/depthcharge/encoders/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/encoders/__pycache__/sinusoidal.cpython-310.pyc b/cascadia/depthcharge/encoders/__pycache__/sinusoidal.cpython-310.pyc
deleted file mode 100644
index 00f7850..0000000
Binary files a/cascadia/depthcharge/encoders/__pycache__/sinusoidal.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/primitives.py b/cascadia/depthcharge/primitives.py
index 3ef7f83..926a00c 100644
--- a/cascadia/depthcharge/primitives.py
+++ b/cascadia/depthcharge/primitives.py
@@ -154,17 +154,20 @@ def from_proforma(
 
         seq = [None] * len(pep)
         mods = [None] * (len(pep) + 2)
-        mods[0] = meta["n_term"]
-        mods[-1] = meta["c_term"]
+        # pyteomics >= 5 returns an empty list (rather than None) for absent
+        # terminal modifications; normalize so they are not tokenized as a
+        # spurious "[+0.000000]-" modification.
+        mods[0] = meta["n_term"] or None
+        mods[-1] = meta["c_term"] or None
         for idx, (res, var_mods) in enumerate(pep):
             seq[idx] = res
             if res in static_mods:
-                if var_mods is None:
+                if not var_mods:
                     var_mods = []
 
                 var_mods.insert(0, static_mods[res])
 
-            mods[idx + 1] = var_mods
+            mods[idx + 1] = var_mods or None
 
         return cls("".join(seq), mods, charge)
 
diff --git a/cascadia/depthcharge/tokenizers/__pycache__/__init__.cpython-310.pyc b/cascadia/depthcharge/tokenizers/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index a1e54cf..0000000
Binary files a/cascadia/depthcharge/tokenizers/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/tokenizers/__pycache__/peptides.cpython-310.pyc b/cascadia/depthcharge/tokenizers/__pycache__/peptides.cpython-310.pyc
deleted file mode 100644
index f1991fe..0000000
Binary files a/cascadia/depthcharge/tokenizers/__pycache__/peptides.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/tokenizers/__pycache__/tokenizer.cpython-310.pyc b/cascadia/depthcharge/tokenizers/__pycache__/tokenizer.cpython-310.pyc
deleted file mode 100644
index f4fc871..0000000
Binary files a/cascadia/depthcharge/tokenizers/__pycache__/tokenizer.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/transformers/__pycache__/__init__.cpython-310.pyc b/cascadia/depthcharge/transformers/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 9de82fc..0000000
Binary files a/cascadia/depthcharge/transformers/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/transformers/__pycache__/peptides.cpython-310.pyc b/cascadia/depthcharge/transformers/__pycache__/peptides.cpython-310.pyc
deleted file mode 100644
index b7947e5..0000000
Binary files a/cascadia/depthcharge/transformers/__pycache__/peptides.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/transformers/__pycache__/spectra.cpython-310.pyc b/cascadia/depthcharge/transformers/__pycache__/spectra.cpython-310.pyc
deleted file mode 100644
index db4231e..0000000
Binary files a/cascadia/depthcharge/transformers/__pycache__/spectra.cpython-310.pyc and /dev/null differ
diff --git a/cascadia/depthcharge/transformers/peptides.py b/cascadia/depthcharge/transformers/peptides.py
index a11456a..8cc359c 100644
--- a/cascadia/depthcharge/transformers/peptides.py
+++ b/cascadia/depthcharge/transformers/peptides.py
@@ -1,5 +1,6 @@
 """Transformer models for peptides."""
 import torch
+import torch.nn.functional as F
 
 from ..encoders import FloatEncoder, PositionalEncoder
 from ..tokenizers import PeptideTokenizer
@@ -278,6 +279,200 @@ def forward(
         return self.final(preds)
 
 
+    @torch.no_grad()
+    def greedy_decode(
+        self,
+        precursors: torch.Tensor,
+        memory: torch.Tensor,
+        memory_key_padding_mask: torch.Tensor,
+        max_length: int,
+        stop_int: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Greedy autoregressive decoding with key/value caching.
+
+        This computes exactly the same sequence of next-token distributions
+        as calling ``forward`` repeatedly with a growing prefix, but each
+        decoding step only processes the newest position: the self-attention
+        keys/values of earlier positions and the cross-attention keys/values
+        of the spectrum memory are computed once and cached.
+
+        Parameters
+        ----------
+        precursors : torch.Tensor of size (batch_size, 2)
+            The precursor mass (axis 0) and charge (axis 1).
+        memory : torch.Tensor of shape (batch_size, n_peaks, d_model)
+            The representations from the spectrum encoder.
+        memory_key_padding_mask : torch.Tensor of shape (batch_size, n_peaks)
+            The mask that indicates which elements of ``memory`` are padding.
+        max_length : int
+            The maximum number of tokens to decode.
+        stop_int : int
+            The integer index of the stop token. Decoding halts early once
+            every sequence in the batch has emitted it.
+
+        Returns
+        -------
+        tokens : torch.Tensor of shape (batch_size, n_steps)
+            The greedily decoded token indices.
+        confidences : torch.Tensor of shape (batch_size, n_steps)
+            The softmax probability of each decoded token.
+        """
+        device = memory.device
+        batch_size = memory.shape[0]
+        layers = self.transformer_decoder.layers
+        n_head = layers[0].self_attn.num_heads
+        memory_key_padding_mask = memory_key_padding_mask.to(device)
+        mem_attn_mask = _padding_to_attn_mask(memory_key_padding_mask)
+
+        # Cross-attention keys/values of the memory never change: cache them.
+        mem_kv = [
+            (
+                _heads(_in_proj(layer.multihead_attn, memory, 1), n_head),
+                _heads(_in_proj(layer.multihead_attn, memory, 2), n_head),
+            )
+            for layer in layers
+        ]
+        self_kv = [None] * len(layers)
+
+        # Position 0 is the precursor "token".
+        masses = self.mass_encoder(precursors[:, None, 0])
+        charges = self.charge_encoder(precursors[:, 1].int() - 1)
+        x = masses + charges[:, None, :]
+        key_padding = torch.zeros(
+            (batch_size, 0), dtype=torch.bool, device=device
+        )
+
+        tokens = []
+        confidences = []
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        for pos in range(max_length):
+            # Padding mask matching forward(): all-zero embeddings are padding.
+            key_padding = torch.cat(
+                [key_padding, (x.sum(dim=2) == 0)], dim=1
+            )
+            self_attn_mask = _padding_to_attn_mask(key_padding)
+            x = _positional_encoding(self.positional_encoder, pos, x)
+
+            for idx, layer in enumerate(layers):
+                x = _decoder_layer_step(
+                    layer,
+                    x,
+                    self_kv,
+                    idx,
+                    self_attn_mask,
+                    mem_kv[idx],
+                    mem_attn_mask,
+                    n_head,
+                )
+
+            scores = torch.softmax(self.final(x[:, -1, :]), dim=1)
+            conf, next_token = scores.max(dim=1)
+            tokens.append(next_token)
+            confidences.append(conf)
+            finished |= next_token == stop_int
+            if bool(finished.all()):
+                break
+
+            x = self.aa_encoder(next_token)[:, None, :]
+
+        return torch.stack(tokens, dim=1), torch.stack(confidences, dim=1)
+
+
+def _in_proj(
+    attn: torch.nn.MultiheadAttention, x: torch.Tensor, which: int
+) -> torch.Tensor:
+    """Apply the query (0), key (1), or value (2) input projection."""
+    d = attn.embed_dim
+    weight = attn.in_proj_weight[which * d : (which + 1) * d]
+    bias = attn.in_proj_bias
+    if bias is not None:
+        bias = bias[which * d : (which + 1) * d]
+
+    return F.linear(x, weight, bias)
+
+
+def _heads(x: torch.Tensor, n_head: int) -> torch.Tensor:
+    """Reshape (batch, length, d_model) to (batch, n_head, length, d_head)."""
+    batch_size, length, d_model = x.shape
+    return x.view(batch_size, length, n_head, d_model // n_head).transpose(
+        1, 2
+    )
+
+
+def _padding_to_attn_mask(key_padding_mask: torch.Tensor) -> torch.Tensor:
+    """Convert a boolean key padding mask (True = pad) to an SDPA mask."""
+    if not bool(key_padding_mask.any()):
+        return None
+
+    return ~key_padding_mask[:, None, None, :]
+
+
+def _attention(
+    attn: torch.nn.MultiheadAttention,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None,
+    n_head: int,
+) -> torch.Tensor:
+    """Multi-head attention with pre-projected keys and values."""
+    batch_size, length, d_model = query.shape
+    q = _heads(_in_proj(attn, query, 0), n_head)
+    out = F.scaled_dot_product_attention(q, key, value, attn_mask=attn_mask)
+    out = out.transpose(1, 2).reshape(batch_size, length, d_model)
+    return attn.out_proj(out)
+
+
+def _decoder_layer_step(
+    layer: torch.nn.TransformerDecoderLayer,
+    x: torch.Tensor,
+    self_kv: list,
+    idx: int,
+    self_attn_mask: torch.Tensor | None,
+    mem_kv: tuple[torch.Tensor, torch.Tensor],
+    mem_attn_mask: torch.Tensor | None,
+    n_head: int,
+) -> torch.Tensor:
+    """One post-norm TransformerDecoderLayer applied to the newest position."""
+    if layer.norm_first:
+        raise NotImplementedError("Only post-norm decoder layers are supported.")
+
+    k = _heads(_in_proj(layer.self_attn, x, 1), n_head)
+    v = _heads(_in_proj(layer.self_attn, x, 2), n_head)
+    if self_kv[idx] is not None:
+        k = torch.cat([self_kv[idx][0], k], dim=2)
+        v = torch.cat([self_kv[idx][1], v], dim=2)
+
+    self_kv[idx] = (k, v)
+
+    x = layer.norm1(
+        x + _attention(layer.self_attn, x, k, v, self_attn_mask, n_head)
+    )
+    x = layer.norm2(
+        x
+        + _attention(
+            layer.multihead_attn, x, mem_kv[0], mem_kv[1], mem_attn_mask, n_head
+        )
+    )
+    x = layer.norm3(
+        x + layer.linear2(layer.activation(layer.linear1(x)))
+    )
+    return x
+
+
+def _positional_encoding(
+    encoder: torch.nn.Module, position: int, x: torch.Tensor
+) -> torch.Tensor:
+    """Add the positional encoding for a single sequence position."""
+    if isinstance(encoder, torch.nn.Identity):
+        return x
+
+    pos = torch.tensor(float(position)).type_as(encoder.sin_term)
+    sin_pos = torch.sin(pos / encoder.sin_term)
+    cos_pos = torch.cos(pos / encoder.cos_term)
+    return x + torch.cat([sin_pos, cos_pos])[None, None, :]
+
+
 def generate_tgt_mask(sz: int) -> torch.Tensor:
     """Generate a square mask for the sequence.
 
diff --git a/cascadia/depthcharge/transformers/spectra.py b/cascadia/depthcharge/transformers/spectra.py
index a896eab..dc23a12 100644
--- a/cascadia/depthcharge/transformers/spectra.py
+++ b/cascadia/depthcharge/transformers/spectra.py
@@ -57,9 +57,13 @@ def __init__(
             dropout=dropout,
         )
 
+        # The nested-tensor "sparsity fast path" relies on an operator that
+        # is not implemented for Apple MPS, and spectra in a batch have
+        # similar lengths anyway, so disable it.
         self.transformer_encoder = torch.nn.TransformerEncoder(
             layer,
             num_layers=n_layers,
+            enable_nested_tensor=False,
         )
 
     def forward(
diff --git a/cascadia/model.py b/cascadia/model.py
index 3c32f41..a837589 100644
--- a/cascadia/model.py
+++ b/cascadia/model.py
@@ -3,7 +3,7 @@
 from typing import Any, Dict, Iterable, List, Tuple
 import torch
 import numpy as np
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 
 class AugmentedPeakEncoder(torch.nn.Module):
     """Encode an augmented m/z, intensity,
@@ -256,36 +256,66 @@ def validation_step(self, batch, _):
         return loss + frag_loss, pep_acc
 
     def predict_step(self, batch, *args):
-        torch.set_grad_enabled(True)
+        """Greedy autoregressive decoding of a batch of augmented spectra.
 
+        The spectrum encoder is run exactly once per unique spectrum in the
+        batch (augmented spectra are emitted once per candidate charge state
+        with identical peaks, and the charge only enters the decoder). The
+        decoder is then applied step by step, stopping as soon as every
+        sequence in the batch has produced a stop token.
+        """
         spectra, precursors, sequences, frag_labels, rts = batch
-        device = spectra.get_device()
-        if device < 0:
-            device = 'cpu'
-
-        cur_sequences = torch.empty((len(sequences),0), dtype=torch.int32, device=device)
-        aa_conf = torch.empty((len(sequences),0), dtype=torch.int32, device=device)
-        for i in range(len(sequences[0]) + 1):
-            preds, _, _ = self._forward_step(spectra, precursors, cur_sequences)
-            next_aa_scores = torch.softmax(preds[:,-1,:], 1)
-            next_aas = torch.argmax(next_aa_scores, 1)
-            aa_conf = torch.cat([aa_conf, torch.reshape(next_aa_scores[range(len(next_aas)),next_aas], (-1, 1))], dim=1)
-            cur_sequences = torch.cat([cur_sequences, torch.reshape(next_aas, (-1, 1))], dim=1)
-
-        preds_seqs = cur_sequences
-        true_seqs = batch[2]
-
-        pred_seqs = self.tokenizer.detokenize(preds_seqs)
-        true_seqs = self.tokenizer.detokenize(true_seqs)
-
-        pep_conf = []
-        trimmed_pred_seqs = []
-        for pred_pep, pep_aa_conf in zip(pred_seqs, aa_conf):
-            trimmed_pep = pred_pep.split('$')[0]
-            trimmed_pred_seqs.append(trimmed_pep)
-            pep_conf.append(torch.sum(torch.log(pep_aa_conf[:len(trimmed_pep) + 1])))
-
-        return trimmed_pred_seqs, true_seqs, pep_conf, aa_conf, rts, precursors
+        device = spectra.device
+        n_spectra = spectra.shape[0]
+        max_length = sequences.shape[1] + 1
+
+        with torch.no_grad():
+            # Encode each unique spectrum only once.
+            flat = spectra.reshape(n_spectra, -1)
+            unique_spectra, inverse = torch.unique(
+                flat.cpu(), dim=0, return_inverse=True
+            )
+            if unique_spectra.shape[0] < n_spectra:
+                unique_spectra = unique_spectra.reshape(-1, *spectra.shape[1:])
+                memory, mem_mask = self.spectrum_encoder(
+                    unique_spectra.to(device)
+                )
+                inverse = inverse.to(device)
+                memory, mem_mask = memory[inverse], mem_mask[inverse]
+            else:
+                memory, mem_mask = self.spectrum_encoder(spectra)
+
+            cur_sequences, aa_conf = self.decoder.greedy_decode(
+                precursors,
+                memory,
+                mem_mask,
+                max_length=max_length,
+                stop_int=self.tokenizer.stop_int,
+            )
+
+        cur_sequences = cur_sequences.cpu()
+        aa_conf = aa_conf.cpu()
+        pred_seqs = self.tokenizer.detokenize(cur_sequences)
+        # Move to the CPU first: element-wise iteration over a GPU tensor
+        # forces one device synchronization per token.
+        true_seqs = self.tokenizer.detokenize(batch[2].cpu())
+
+        # Number of decoded tokens up to and including the stop token.
+        is_stop = cur_sequences == self.tokenizer.stop_int
+        has_stop = is_stop.any(dim=1)
+        n_tokens = torch.where(
+            has_stop,
+            is_stop.float().argmax(dim=1) + 1,
+            torch.full_like(has_stop, cur_sequences.shape[1], dtype=torch.long),
+        )
+        log_conf = torch.log(aa_conf)
+        token_mask = (
+            torch.arange(cur_sequences.shape[1])[None, :] < n_tokens[:, None]
+        )
+        pep_conf = (log_conf * token_mask).sum(dim=1)
+
+        trimmed_pred_seqs = [pep.split("$")[0] for pep in pred_seqs]
+        return trimmed_pred_seqs, true_seqs, list(pep_conf), aa_conf, rts, precursors
 
     def configure_optimizers(self):
         """
diff --git a/docs/source/getting_started.md b/docs/source/getting_started.md
index b0eae40..6588d9a 100755
--- a/docs/source/getting_started.md
+++ b/docs/source/getting_started.md
@@ -1,6 +1,6 @@
 # Getting Started
 
-Cascadia is currently only compatible with Linux, but we are actively working on support for Windows and macOS. 
+Cascadia runs on Linux and macOS. On Linux, an NVIDIA GPU is used automatically when available; on Apple silicon Macs, the built-in GPU is used through PyTorch's MPS backend.
 
 ### Installation 
 
@@ -21,15 +21,25 @@ Finally, you can install Cascadia and all of its dependencies with:
 ```sh
 pip install cascadia
 ```
-The installation should take no longer than a couple of minutes.
+The installation should take no longer than a couple of minutes. Cascadia supports Python 3.10 and newer, PyTorch 2.2 and newer, and Lightning 2.x.
+
+### Running on Google Colab
+
+Google Colab provides a free GPU, which is a convenient way to try Cascadia. In a notebook with a GPU runtime, install Cascadia and download the demo data and model checkpoints, then run the demo:
+
+```sh
+pip install cascadia gdown
+gdown --folder https://drive.google.com/drive/folders/1UTrZIrCdUqYqscbqga_KdX8kc8ZjMMfr
+cascadia sequence Cascadia/demo.mzML Cascadia/cascadia.ckpt --out demo_results
+```
 
 ### Run de novo sequencing on new data with a trained model 
 
 ```{note}
-We recommend using linux and a dedicated GPU to achieve optimal runtime performance.
+We recommend a dedicated GPU (NVIDIA on Linux, or an Apple silicon Mac) to achieve optimal runtime performance. Inference on the CPU works but is considerably slower.
 ```
 
-Most users will want to use a pretrained Cascadia model to perform de novo sequencing on a new dataset. Cascadia takes input MS data in the [mzML](file_formats.md) format. A small demo dataset, along with the pretrained model checkpoints from the paper, are available [here](https://drive.google.com/drive/folders/1UTrZIrCdUqYqscbqga_KdX8kc8ZjMMfr?usp=sharing). The following example on the provided demo dataset should take approximately 1 minute to run on a GPU:
+Most users will want to use a pretrained Cascadia model to perform de novo sequencing on a new dataset. Cascadia takes input MS data in the [mzML](file_formats.md) format. A small demo dataset, along with the pretrained model checkpoints from the paper, are available [here](https://drive.google.com/drive/folders/1UTrZIrCdUqYqscbqga_KdX8kc8ZjMMfr?usp=sharing). The following example on the provided demo dataset takes well under a minute on a GPU (including on an Apple silicon Mac) and a few minutes on the CPU:
 
 ```sh
     cascadia sequence \
diff --git a/docs/source/usage.md b/docs/source/usage.md
index 522ecf7..37f13e5 100755
--- a/docs/source/usage.md
+++ b/docs/source/usage.md
@@ -10,11 +10,12 @@ Argument | Description
 |spectrum_file |  __(required)__ The mzML file to perform de novo sequencing on, or a list of multiple mzML files |
 |-o, --outfile | The output file to save de novo sequencing results to. (default: cascadia_results.ssl)|
 |-t, --score_threshold] | The score threshold applied to predictions. (default: 0.8)|
-|-b, --batch_size | The batch size for inference. For the fastest inference the largest batch size that fits in GPU memory is recommended.  (default: 32)|
+|-b, --batch_size | The batch size for inference. For the fastest inference the largest batch size that fits in GPU memory is recommended.  (default: 128)|
 |-w, --width | The number of adjacent scans to use when construcing augmented spectra. (default: 2)|
 |-c, --max_charge | The maximum precursor charge to consider when making predictions. (default 4) |
 |-f, --format | One of 'ssl' or 'csv'. The file format to write results to. Select 'ssl' to generate output which can be loaded into Skyline. (default csv) |
 |-p, --modifications | A path to the json file containing a list of PTMs in [Proforma](https://github.com/HUPO-PSI/ProForma) format. If not provided, the PTMs present in the [Massive-KB](https://massive.ucsd.edu/ProteoSAFe/static/massive-kb-libraries.jsp) dataset will be used by default. The list of PTMs needs to match those used to train the model. (default mskb) |
+|-d, --device | One of 'auto', 'gpu', 'mps', or 'cpu'. The device to run inference on; 'auto' uses an NVIDIA GPU if available, then an Apple silicon GPU (MPS), then the CPU. (default auto) |
 
 ## Train
 ```
diff --git a/pyproject.toml b/pyproject.toml
index a19f8e1..c2dda8e 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,14 +11,13 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Operating System :: OS Independent",
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 dependencies = [
-    "setuptools<70.0.0",
-    "lightning>=2.0,<2.1",
-    "pytorch-lightning>=1.9,<2.0",
+    "lightning>=2.0,<3.0",
     "pyteomics>=4.6",
+    "psims",
     "torch>=2.2.0",
-    "numpy<2.0",
+    "numpy>=1.24",
     "numba>=0.48.0",
     "lxml>=4.9.1",
     "h5py>=3.7.0",