diff --git a/bergson/config/config_io.py b/bergson/config/config_io.py
index da50616d..105accc7 100644
--- a/bergson/config/config_io.py
+++ b/bergson/config/config_io.py
@@ -120,6 +120,31 @@ def read_config(path: str | Path) -> dict[str, Any]:
     return config
 
 
+def read_first_step_config(path: str | Path) -> dict[str, Any] | None:
+    """The first step's configuration from a run's ``config.yaml``.
+
+    ``path`` may be the run directory or an artifact inside it (``scores.pt``),
+    so callers holding an output path need not resolve it. Returns ``None``
+    when there is no readable config, which is the signal that a caller must
+    fall back to inspecting the artifact itself.
+    """
+    p = Path(path)
+    cfg_path = (p if p.is_dir() else p.parent) / CONFIG_FILENAME
+    if not cfg_path.is_file():
+        return None
+
+    try:
+        steps = read_config(cfg_path)["steps"]
+    except (ValueError, yaml.YAMLError):
+        return None
+
+    for step in steps:
+        for step_cfg in step.values():
+            if isinstance(step_cfg, dict):
+                return step_cfg
+    return None
+
+
 T = TypeVar("T", bound="FromDict")
 
 
diff --git a/bergson/data.py b/bergson/data.py
index bf5f3418..caa21a58 100644
--- a/bergson/data.py
+++ b/bergson/data.py
@@ -640,40 +640,13 @@ def to_grid(self) -> torch.Tensor:
         return scores[..., 0] if self.num_scores == 1 else scores
 
 
-class ArrayScores:
-    """Dense ``[num_items, num_scores]`` score matrix with the same interface
-    as :class:`Scores`, for scores saved as a plain ``.npy`` (e.g. an ad hoc
-    array from outside the standard scoring pipeline)."""
-
-    def __init__(self, arr: np.ndarray):
-        if arr.ndim == 1:
-            arr = arr[:, None]
-        self.arr = arr
-        self.num_scores = arr.shape[1]
-
-    def __len__(self) -> int:
-        return self.arr.shape[0]
-
-    def __getitem__(self, key: Any) -> NDArray:
-        return self.arr[key]
-
-    def get(self, key: Any, score_idx: int = 0) -> NDArray:
-        return self.arr[key, score_idx]
-
-    def is_written(self) -> bool:
-        return True
-
-
-def load_scores(path: Path) -> Scores | ArrayScores:
+def load_scores(path: Path) -> Scores:
     """Load a score store written by the standard scoring pipeline.
 
-    A plain ``.npy`` array loads as :class:`ArrayScores`. Any other score
-    directory loads its ``scores.bin`` as :class:`Scores`, with ``offsets``
-    set when ``info["attribute_tokens"]`` marks it as a per-token store.
+    A score directory loads its ``scores.bin`` as :class:`Scores`, with
+    ``offsets`` set when ``info["attribute_tokens"]`` marks it as a per-token
+    store.
     """
-    if path.suffix == ".npy":
-        return ArrayScores(np.load(path))
-
     info_path = path / "info.json"
     with open(info_path, "r") as f:
         info = json.load(f)
diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py
index 3bdcc130..47434c78 100644
--- a/bergson/magic/cli.py
+++ b/bergson/magic/cli.py
@@ -28,7 +28,7 @@
 )
 
 from ..config.config import TrainingConfig, ValidationConfig
-from ..config.config_io import save_run_config
+from ..config.config_io import read_first_step_config, save_run_config
 from ..distributed import launch_distributed_run
 from ..utils.load_from_optimizer import (
     save_second_moments_as_optimizer_pt,
@@ -130,9 +130,8 @@ def compute_per_query_magic_scores(
     yields a single aggregate-query score. This scores each query separately:
     for each query document it takes that document's gradient at the final model
     as the backward cotangent and runs ``Trainer.backward`` over the saved
-    trajectory, producing a ``[num_train_docs, num_query_docs]`` matrix (the
-    layout ``validate_scores`` consumes: rows are leave-out docs, columns
-    queries).
+    trajectory, producing ``[num_train_docs, num_query_docs]`` — or
+    ``[num_train_docs, seq_len, num_query_docs]`` when attributing tokens.
 
     The backward is linear in the cotangent, so this is exact; the forward runs
     once and every query reuses its checkpoints (``cleanup=False``). Per-query
@@ -217,7 +216,7 @@ def compute_per_query_magic_scores(
 
         s = bwd_state.weight_grads.detach().cpu()
         if pad_count:
-            s = s[:-weight_pad_count]
+            s = s[:-weight_pad_count] if s.ndim == 1 else s[:-pad_count]
         if main:
             torch.save(s, qpath)
         per_query.append(s)
@@ -241,9 +240,7 @@ def compute_per_query_magic_scores(
     fwd_state.copy_(restored)
     del restored
 
-    # [num_train_docs, num_query_docs] — the layout validate_scores expects
-    # (rows are leave-out docs; columns are queries).
-    return torch.stack(per_query, dim=1)
+    return torch.stack(per_query, dim=-1)
 
 
 def scores_are_per_token(score_path: str) -> bool:
@@ -253,10 +250,14 @@ def scores_are_per_token(score_path: str) -> bool:
             return False
         with open(info_path) as f:
             return bool(json.load(f).get("attribute_tokens", False))
-    if score_path.endswith(".npy"):
-        return False
+    step_cfg = read_first_step_config(score_path)
+    if step_cfg is not None:
+        return bool(step_cfg.get("attribute_tokens") or step_cfg.get("per_token"))
+
     scores = torch.load(score_path, map_location="cpu")
-    return isinstance(scores, torch.Tensor) and scores.ndim == 2 and scores.shape[1] > 1
+    return isinstance(scores, torch.Tensor) and (
+        scores.ndim == 3 or (scores.ndim == 2 and scores.shape[1] > 1)
+    )
 
 
 def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset:
@@ -487,9 +488,8 @@ def worker(
 
     multi_query = False
     if not score_path and run_cfg.query_method == "none":
-        # Per-query MAGIC: one backward per query, sharing the forward. Yields a
-        # [num_query_docs, num_train_docs] score matrix (multi_query), the unit
-        # for a per-query LDS.
+        # Per-query MAGIC: one backward per query, sharing the forward. Yields
+        # the unit for a per-query LDS.
         if not isinstance(run_cfg, MagicConfig):
             raise RuntimeError("run_cfg must be a MagicConfig to compute scores")
         assert query_dataset is not None
@@ -569,7 +569,7 @@ def worker(
             score_path = os.path.join(run_cfg.run_path, "scores.pt")
             torch.save(scores, score_path)
             print(f"Saved attribution scores to {score_path}")
-    elif os.path.isdir(score_path) or score_path.endswith(".npy"):
+    elif os.path.isdir(score_path):
         scores, multi_query = load_attribution_scores(score_path)
     else:
         scores = torch.load(score_path, map_location="cpu")
diff --git a/bergson/score/score.py b/bergson/score/score.py
index 02ee6f8f..dac0c7c4 100644
--- a/bergson/score/score.py
+++ b/bergson/score/score.py
@@ -231,7 +231,7 @@ def create_scorer(
     else:
         num_scores = num_queries_total if num_queries_total is not None else num_queries
     if attribute_tokens:
-        writer = MemmapTokenScoreWriter(
+        writer = MemmapTokenScoreWriter.from_dataset(
             path,
             data,
             num_scores,
diff --git a/bergson/score/score_writer.py b/bergson/score/score_writer.py
index 3f7a06b6..d51dcc82 100644
--- a/bergson/score/score_writer.py
+++ b/bergson/score/score_writer.py
@@ -156,11 +156,12 @@ class MemmapTokenScoreWriter(ScoreWriter):
     def __init__(
         self,
         path: Path,
-        data: Dataset,
+        num_token_grads: np.ndarray,
         num_scores: int,
         *,
         dtype: torch.dtype = torch.float32,
         flush_interval: int = 64,
+        overwrite: bool = False,
     ):
         self.path = path
         self.num_scores = num_scores
@@ -168,10 +169,9 @@ def __init__(
         self.flush_interval = flush_interval
         self.num_batches_since_flush = 0
 
-        num_token_grads = compute_num_token_grads(data)
-        num_items = len(data)
+        num_items = len(num_token_grads)
         self.num_token_grads = num_token_grads
-        self.offsets = np.zeros(len(num_token_grads) + 1, dtype=np.int64)
+        self.offsets = np.zeros(num_items + 1, dtype=np.int64)
         np.cumsum(num_token_grads, out=self.offsets[1:])
         total_tokens = int(self.offsets[-1])
 
@@ -181,7 +181,7 @@ def __init__(
         struct_dtype, struct_dtype_json = _score_struct_dtype(num_scores, np_dtype)
 
         rank = dist.get_rank() if dist.is_initialized() else 0
-        if rank == 0 and not scores_file_path.exists():
+        if rank == 0 and (overwrite or not scores_file_path.exists()):
             print(f"Creating new token scores file: {scores_file_path}")
 
             self.scores = np.memmap(
@@ -216,6 +216,11 @@ def __init__(
             shape=(total_tokens,),
         )
 
+    @classmethod
+    def from_dataset(cls, path: Path, data: Dataset, num_scores: int, **kwargs):
+        """For callers holding the dataset the scores were computed over."""
+        return cls(path, compute_num_token_grads(data), num_scores, **kwargs)
+
     def __call__(self, indices: list[int], scores: torch.Tensor, query_offset: int = 0):
         # scores: [total_valid_in_batch, num_scores]
         scores = scores.to(dtype=self.dtype)
@@ -248,48 +253,24 @@ def save_token_scores(
     *,
     dtype: torch.dtype = torch.float32,
 ) -> None:
-    """One-shot equivalent of :class:`MemmapTokenScoreWriter`'s on-disk
-    layout, for callers that already hold the full flat ``(total_tokens,
-    num_scores)`` array in memory (e.g. summed across upstream per-token
-    score dirs sharing the same ``offsets``), rather than streaming batches
-    through ``__call__``.
+    """One-shot equivalent of :class:`MemmapTokenScoreWriter`, for callers that
+    already hold the full flat ``(total_tokens, num_scores)`` array in memory
+    (e.g. summed across upstream per-token score dirs sharing the same
+    ``offsets``), rather than streaming batches through ``__call__``.
     """
     if scores.ndim == 1:
         scores = scores[:, None]
-    total_tokens, num_scores = scores.shape
-    num_items = len(offsets) - 1
-
-    path.mkdir(parents=True, exist_ok=True)
-    np_dtype = convert_dtype_to_np(dtype)
-    struct_dtype, struct_dtype_json = _score_struct_dtype(num_scores, np_dtype)
-
-    mmap = np.memmap(
-        path / "scores.bin",
-        dtype=np.dtype(struct_dtype),  # type: ignore
-        mode="w+",
-        shape=(total_tokens,),
+    _, num_scores = scores.shape
+    num_token_grads = np.diff(offsets).astype(np.int64)
+
+    writer = MemmapTokenScoreWriter(
+        path, num_token_grads, num_scores, dtype=dtype, overwrite=True
     )
-    # numpy_to_tensor handles bf16, which torch.from_numpy cannot.
-    scores_np = tensor_to_numpy(numpy_to_tensor(np.ascontiguousarray(scores)).to(dtype))
-    for i in range(num_scores):
-        mmap[f"score_{i}"] = scores_np[:, i]
-        mmap[f"written_{i}"] = True
-    mmap.flush()
-
-    np.save(path / "offsets.npy", offsets)
-
-    with (path / "info.json").open("w") as f:
-        json.dump(
-            {
-                "attribute_tokens": True,
-                "num_items": num_items,
-                "num_rows": total_tokens,
-                "num_scores": num_scores,
-                "dtype": struct_dtype_json,
-            },
-            f,
-            indent=2,
-        )
+    writer(
+        list(range(len(num_token_grads))),
+        numpy_to_tensor(np.ascontiguousarray(scores)),
+    )
+    writer.flush()
 
 
 class MemmapSequenceScoreWriter(ScoreWriter):
diff --git a/bergson/validate.py b/bergson/validate.py
index 2d61a575..71c3805f 100644
--- a/bergson/validate.py
+++ b/bergson/validate.py
@@ -18,7 +18,6 @@
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
-import yaml
 from peft import PeftModel
 from scipy.stats import pearsonr, spearmanr
 from tqdm import tqdm
@@ -31,7 +30,7 @@
 )
 
 from .config.config import ScoreConfig, ValidationConfig
-from .config.config_io import CONFIG_FILENAME, load_subconfig, save_run_config
+from .config.config_io import load_subconfig, read_first_step_config, save_run_config
 from .data import Scores, load_scores, pad_and_tensor
 from .magic.data_stream import DataStream, pad_dataset_to_batch_size
 from .magic.trainer import TrainerState, prepare_trainer
@@ -41,35 +40,30 @@
 
 
 def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
-    """Load attribution scores from a score directory, ``.npy``, or ``.pt`` file.
+    """Load attribution scores from a score directory or ``.pt`` file.
 
     Returns ``(scores, multi_query)``. Token score directories are per-token,
     with a query dimension when ``num_scores > 1`` (``[docs, seq_len,
-    queries]``); plain score directories and ``.npy`` arrays are per-document,
-    with one column per query. A 2-D ``.pt`` tensor is ambiguous — per-token
-    MAGIC scores are ``[docs, seq_len]`` and per-query MAGIC scores are
-    ``[docs, queries]`` — so the run config next to it decides: it is
-    per-query iff the run used ``query_method: none``.
+    queries]``); plain score directories are per-document, with one column per
+    query. A 3-D ``.pt`` tensor is per-token per-query ``[docs, seq_len,
+    queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token MAGIC scores are
+    ``[docs, seq_len]`` and per-query MAGIC scores are ``[docs, queries]`` —
+    so the run config next to it decides: it is per-query iff the run used
+    ``query_method: none`` without attributing tokens.
 
     Score directories are negated when their ``score_cfg.higher_is_better`` is
-    set, aligning them with the loss-diff convention. ``.npy`` files carry no
-    ``score_cfg`` and are loaded as-is: they must already be in the loss-diff
-    convention.
+    set, aligning them with the loss-diff convention.
     """
-    if os.path.isdir(score_path) or score_path.endswith(".npy"):
+    if os.path.isdir(score_path):
         loaded = load_scores(Path(score_path))
-        score_cfg = (
-            load_subconfig(score_path, "score_cfg", ScoreConfig)
-            if os.path.isdir(score_path)
-            else None
-        )
+        score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig)
         negate = score_cfg is not None and score_cfg.higher_is_better
 
         if isinstance(loaded, Scores) and loaded.offsets is not None:
             scores = loaded.to_grid()
             if negate:
                 scores = -scores
-            return scores, scores.ndim == 3
+            return scores, loaded.num_scores > 1
 
         arr = np.asarray(loaded[:])
         # Copy: the slice is a read-only view onto the memmap.
@@ -77,35 +71,16 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
         scores = torch.from_numpy(arr.astype(out_dtype, copy=True))
         if negate:
             scores = -scores
-        return scores, scores.ndim == 2 and scores.shape[1] > 1
+        return scores, loaded.num_scores > 1
 
     scores = torch.load(score_path, map_location="cpu")
-    return scores, _pt_scores_are_per_query(score_path, scores)
-
-
-def _pt_scores_are_per_query(score_path: str, scores: torch.Tensor) -> bool:
-    """Whether a 2-D ``.pt`` tensor is per-query ``[docs, queries]`` rather
-    than per-token ``[docs, seq_len]``. The shape alone cannot tell them
-    apart; the run config written next to ``scores.pt`` records which one
-    ``run_magic`` saved."""
-    if not (isinstance(scores, torch.Tensor) and scores.ndim == 2):
-        return False
-    cfg_path = Path(score_path).parent / CONFIG_FILENAME
-    if not cfg_path.is_file():
-        return False
-    with open(cfg_path) as f:
-        doc = yaml.safe_load(f)
-    if not isinstance(doc, dict):
-        return False
-    steps = doc.get("steps")
-    payload = (
-        next(iter(steps[0].values())) if isinstance(steps, list) and steps else doc
-    )
-    return (
-        isinstance(payload, dict)
-        and payload.get("query_method") == "none"
-        and not payload.get("per_token")
-    )
+    if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3):
+        return scores, False
+
+    step_cfg = read_first_step_config(score_path)
+    if step_cfg is None:
+        return scores, scores.ndim == 3
+    return scores, step_cfg.get("query_method") == "none"
 
 
 def bank_loss_cache_key(
diff --git a/tests/test_attribute_tokens.py b/tests/test_attribute_tokens.py
index 916c3954..38f4d2eb 100644
--- a/tests/test_attribute_tokens.py
+++ b/tests/test_attribute_tokens.py
@@ -236,7 +236,7 @@ def test_token_score_writer(tmp_path: Path):
     # lengths [4, 3] → num_token_grads [3, 2]
     ds = Dataset.from_dict({"input_ids": [[1, 2, 3, 4], [5, 6, 7]], "length": [4, 3]})
 
-    writer = MemmapTokenScoreWriter(
+    writer = MemmapTokenScoreWriter.from_dataset(
         tmp_path,
         data=ds,
         num_scores=2,
@@ -457,7 +457,7 @@ def test_token_score_e2e(tmp_path: Path, model, dataset):
     query_grads = {m: torch.randn(1, math.prod(shapes[m])) for m in modules}
 
     score_dtype = get_gradient_dtype(model)
-    writer = MemmapTokenScoreWriter(
+    writer = MemmapTokenScoreWriter.from_dataset(
         tmp_path / "scores",
         data=dataset,
         num_scores=1,
@@ -877,7 +877,7 @@ def test_trackstar_token_scores_sum_to_sequence_scores_on_disk(
 
     # --- Per-token scores via MemmapTokenScoreWriter ---
     tok_path = tmp_path / "tok_scores"
-    tok_writer = MemmapTokenScoreWriter(tok_path, dataset, 1, dtype=dtype)
+    tok_writer = MemmapTokenScoreWriter.from_dataset(tok_path, dataset, 1, dtype=dtype)
     tok_scorer = Scorer(
         query_grads=query_grads,
         modules=sorted_modules,
diff --git a/tests/test_bank_loss_cache.py b/tests/test_bank_loss_cache.py
index f72985a4..a955154f 100644
--- a/tests/test_bank_loss_cache.py
+++ b/tests/test_bank_loss_cache.py
@@ -18,6 +18,7 @@
 import bergson.validate as validate
 from bergson.config.config import DataConfig
 from bergson.magic.config import MagicConfig
+from bergson.score.score_writer import save_sequence_scores
 from bergson.validate import bank_loss_cache_key, evaluate_retrained
 
 MODEL = "trl-internal-testing/tiny-Phi3ForCausalLM"
@@ -111,8 +112,8 @@ def test_evaluate_retrained_reuses_cached_bank_losses(
     cols = n_query if multi_query else 1
     rng = np.random.default_rng(0)
     scores = rng.standard_normal((NUM_DOCS, cols)).astype(np.float32)
-    score_path = tmp_path / "scores.npy"
-    np.save(score_path, scores)
+    score_path = tmp_path / "scores"
+    save_sequence_scores(score_path, scores)
 
     # First run: cold cache, evaluates the bank and writes the cache.
     cfg1 = _run_cfg(tmp_path, "run1", query_path)
@@ -150,8 +151,8 @@ def test_different_query_does_not_reuse_losses(tmp_path, model):
     root = _build_bank(tmp_path, model)
     query_path, n_query = _query_dataset(tmp_path)
     scores = np.random.default_rng(2).standard_normal((NUM_DOCS, n_query)).astype("f4")
-    score_path = tmp_path / "scores.npy"
-    np.save(score_path, scores)
+    score_path = tmp_path / "scores"
+    save_sequence_scores(score_path, scores)
 
     cfg_a = _run_cfg(tmp_path, "run_a", query_path)
     evaluate_retrained(cfg_a, str(root), score_path=str(score_path))
@@ -170,8 +171,8 @@ def test_evaluate_retrained_averages_over_dirs(tmp_path, model):
     query_path, _ = _query_dataset(tmp_path)
 
     scores = np.random.default_rng(0).standard_normal((NUM_DOCS, 1)).astype(np.float32)
-    score_path = tmp_path / "scores.npy"
-    np.save(score_path, scores)
+    score_path = tmp_path / "scores"
+    save_sequence_scores(score_path, scores)
 
     for name, dirs in [
         ("one", str(bank_a)),
diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py
index 7778bfbd..ad81b007 100644
--- a/tests/test_multi_query_validate.py
+++ b/tests/test_multi_query_validate.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 import torch.nn.functional as F
 from datasets import Dataset
@@ -72,13 +71,7 @@ def test_load_attribution_scores_single_column_dir(tmp_path):
     assert scores.shape == (4, 1)
 
 
-def test_load_attribution_scores_npy_and_pt(tmp_path):
-    npy_path = tmp_path / "scores.npy"
-    np.save(npy_path, np.zeros((5, 2), dtype=np.float32))
-    scores, multi_query = load_attribution_scores(str(npy_path))
-    assert multi_query
-    assert scores.shape == (5, 2)
-
+def test_load_attribution_scores_pt_without_config(tmp_path):
     # 2D tensors in .pt files are per-token MAGIC scores, never multi-query.
     pt_path = tmp_path / "scores.pt"
     torch.save(torch.zeros(5, 2), pt_path)
@@ -174,11 +167,12 @@ def test_load_attribution_scores_pt_per_query(tmp_path):
     assert multi_query
     torch.testing.assert_close(scores, values)
 
-    # Per-token runs save [docs, seq_len]; never multi-query.
-    cfg = {"steps": [{"magic": {"query_method": "none", "per_token": True}}]}
+    # Attributing tokens as well makes it 3-D, and still per-query.
+    cfg = {"steps": [{"magic": {"query_method": "none", "attribute_tokens": True}}]}
     (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg))
-    _, multi_query = load_attribution_scores(str(tmp_path / "scores.pt"))
-    assert not multi_query
+    torch.save(torch.zeros(4, 3, 2), tmp_path / "tok.pt")
+    _, multi_query = load_attribution_scores(str(tmp_path / "tok.pt"))
+    assert multi_query
 
     cfg = {"steps": [{"magic": {"query_method": "mean"}}]}
     (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg))
diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py
index 7f05e221..a8910f8e 100644
--- a/tests/test_per_query_magic.py
+++ b/tests/test_per_query_magic.py
@@ -197,3 +197,77 @@ def test_per_query_scores_only_real_queries_when_padded():
         import os
 
         assert not os.path.exists(f"{run_path}/per_query/q3.pt")
+
+
+def _per_query_run(tmp_path, attribute_tokens: bool, num_docs=5, seq_len=8, n_query=2):
+    """Run worker() in per-query mode; return (scores, doc_ids).
+
+    num_docs=5 at batch_size 4 pads by 3 rows, where weight_pad_count is 1 but
+    pad_count is 3 — the gap a rank-blind trim falls into.
+    """
+    from bergson.config.config import DataConfig
+    from bergson.magic.cli import worker
+
+    def ds(n):
+        toks = [[(d * seq_len + t) % 50 + 1 for t in range(seq_len)] for d in range(n)]
+        return Dataset.from_dict(
+            {
+                "input_ids": toks,
+                "labels": toks,
+                "doc_ids": [[d] * seq_len for d in range(n)],
+                "length": [seq_len] * n,
+            }
+        )
+
+    run_path = tmp_path / ("tok" if attribute_tokens else "doc")
+    run_cfg = MagicConfig(
+        run_path=str(run_path),
+        model="EleutherAI/pythia-14m",
+        data=DataConfig(dataset="unused", chunk_length=seq_len),
+        query=DataConfig(dataset="unused", chunk_length=seq_len),
+        batch_size=4,
+        attribute_tokens=attribute_tokens,
+        query_method="none",
+        skip_validation=True,
+    )
+    worker(0, 0, 1, ds(num_docs), ds(n_query), num_docs, n_query, run_cfg)
+
+    doc_ids = run_path / "doc_ids.pt"
+    return (
+        torch.load(run_path / "scores.pt"),
+        torch.load(doc_ids) if doc_ids.exists() else None,
+    )
+
+
+def test_per_query_per_token_aggregates_to_per_doc(tmp_path):
+    """Per-token per-query scores are [docs, seq_len, queries], and summing
+    them over each document's tokens reproduces the per-doc per-query run."""
+    per_tok, doc_ids = _per_query_run(tmp_path, attribute_tokens=True)
+    per_doc, _ = _per_query_run(tmp_path, attribute_tokens=False)
+    num_docs, n_query = per_doc.shape
+
+    assert per_tok.shape == (num_docs, 8, n_query), f"got {tuple(per_tok.shape)}"
+    assert doc_ids is not None, "per-token run must write doc_ids.pt"
+    assert doc_ids.shape == per_tok.shape[:2]
+
+    agg = torch.zeros(num_docs, n_query, dtype=torch.float64)
+    agg.scatter_add_(
+        0,
+        doc_ids.reshape(-1, 1).expand(-1, n_query),
+        per_tok.reshape(-1, n_query).to(torch.float64),
+    )
+
+    torch.testing.assert_close(agg, per_doc.to(torch.float64), atol=1e-5, rtol=1e-4)
+
+
+def test_three_dim_scores_load_as_per_token_multi_query(tmp_path):
+    """A 3-D scores.pt is unambiguous: neither classifier needs the run config."""
+    from bergson.magic.cli import scores_are_per_token
+    from bergson.validate import load_attribution_scores
+
+    path = tmp_path / "scores.pt"
+    torch.save(torch.randn(4, 6, 3), path)
+
+    assert scores_are_per_token(str(path))
+    _, multi_query = load_attribution_scores(str(path))
+    assert multi_query
diff --git a/tests/test_per_token_lds.py b/tests/test_per_token_lds.py
index 188c5bf5..f5fe33ec 100644
--- a/tests/test_per_token_lds.py
+++ b/tests/test_per_token_lds.py
@@ -16,105 +16,6 @@
 from bergson.validate import load_attribution_scores
 
 
-def test_flat_reweight_per_doc_matches_direct_indexing():
-    """1-D (per-doc) weights: ``view(-1)[subset]`` equals direct indexing, so
-    the per-document leave-k-out path is unchanged by the flat generalization."""
-    weights = torch.ones(8)
-    subset = torch.tensor([1, 4, 6])
-
-    flat = weights.clone()
-    flat.view(-1)[subset] = 0.0
-    direct = torch.ones(8)
-    direct[subset] = 0.0
-
-    torch.testing.assert_close(flat, direct)
-
-
-def test_flat_reweight_per_token_hits_grid_positions():
-    """2-D (per-token) weights: a flat subset re-weights exactly the intended
-    ``doc * seq_len + token`` grid positions and nothing else."""
-    n_docs, seq_len = 4, 5
-    weights = torch.ones(n_docs, seq_len)
-    # doc 1 token 2, doc 3 token 0, doc 3 token 4
-    positions = [(1, 2), (3, 0), (3, 4)]
-    subset = torch.tensor([d * seq_len + t for d, t in positions])
-
-    weights.view(-1)[subset] = 2.0
-
-    expected = torch.ones(n_docs, seq_len)
-    for d, t in positions:
-        expected[d, t] = 2.0
-    torch.testing.assert_close(weights, expected)
-
-
-def test_whole_doc_token_subset_reweights_full_row():
-    """Selecting every token position of a document is equivalent to selecting
-    that document in the per-document formulation (its whole row is set)."""
-    n_docs, seq_len = 3, 6
-    weights = torch.ones(n_docs, seq_len)
-    doc = 1
-    subset = torch.arange(doc * seq_len, (doc + 1) * seq_len)
-
-    weights.view(-1)[subset] = 0.0
-
-    expected = torch.ones(n_docs, seq_len)
-    expected[doc] = 0.0
-    torch.testing.assert_close(weights, expected)
-
-
-def test_reweight_sequence_preserves_padding_rows():
-    """The retrain loop's exact sequence -- fill with 1.0, zero the padding
-    rows, then flat-reweight the subset -- leaves padding rows at zero, so
-    batch-size padding docs never re-enter training."""
-    n_docs, seq_len, pad_count = 4, 5, 1
-    weights = torch.ones(n_docs, seq_len)
-    weights.data[-pad_count:] = 0.0
-    subset = torch.tensor([0 * seq_len + 2, 1 * seq_len + 4])
-
-    weights.view(-1)[subset] = 0.0
-
-    assert weights[-pad_count:].eq(0).all()
-    expected_real = torch.ones(n_docs - pad_count, seq_len)
-    expected_real[0, 2] = 0.0
-    expected_real[1, 4] = 0.0
-    torch.testing.assert_close(weights[:-pad_count], expected_real)
-
-
-def test_score_sum_flat_per_token():
-    """``scores.reshape(-1)[subset].sum()`` sums the selected per-token scores
-    regardless of the score tensor's shape (1-D per doc or 2-D per token)."""
-    scores = torch.arange(12, dtype=torch.float32).reshape(3, 4)
-    subset = torch.tensor([0 * 4 + 1, 2 * 4 + 3])  # scores 1 and 11
-
-    got = scores.reshape(-1)[subset].sum()
-
-    torch.testing.assert_close(got, torch.tensor(12.0))
-    # Same expression on a flat (per-doc) score vector.
-    flat = torch.arange(6, dtype=torch.float32)
-    torch.testing.assert_close(
-        flat.reshape(-1)[torch.tensor([0, 5])].sum(), torch.tensor(5.0)
-    )
-
-
-def test_score_sum_flat_multi_query_per_token():
-    """``scores.reshape(-1, num_queries)[subset].sum(dim=0)`` sums each query's
-    scores over the selected token positions of a ``[docs, seq_len, queries]``
-    grid -- the same expression that serves per-doc multi-query scores."""
-    n_docs, seq_len, num_queries = 2, 3, 2
-    scores = torch.arange(12, dtype=torch.float32).reshape(n_docs, seq_len, num_queries)
-    # doc 0 token 1, doc 1 token 2
-    subset = torch.tensor([0 * seq_len + 1, 1 * seq_len + 2])
-
-    got = scores.reshape(-1, num_queries)[subset].sum(dim=0)
-
-    torch.testing.assert_close(got, scores[0, 1] + scores[1, 2])
-    # Per-doc multi-query [docs, queries] is served by the same expression.
-    per_doc = torch.arange(8, dtype=torch.float32).reshape(4, 2)
-    torch.testing.assert_close(
-        per_doc.reshape(-1, 2)[torch.tensor([1, 3])].sum(dim=0), per_doc[1] + per_doc[3]
-    )
-
-
 def test_load_attribution_scores_per_token_pt(tmp_path):
     """A 2-D ``.pt`` tensor is treated as per-token MAGIC scores, never
     multi-query -- so ``validate_scores`` keeps it 2-D for token re-weighting."""
@@ -181,8 +82,8 @@ def test_load_token_dir_scatters_packed_to_grid(tmp_path):
 
 def test_scores_are_per_token_inference(tmp_path):
     """``bergson validate`` infers per-token-ness from the scores themselves:
-    token dirs and 2-D ``.pt`` tensors are per-token; 1-D ``.pt``, ``.npy``,
-    and plain score dirs are per-document."""
+    token dirs and 2-D ``.pt`` tensors are per-token; 1-D ``.pt`` and plain
+    score dirs are per-document."""
     from bergson.magic.cli import scores_are_per_token
 
     token_dir = tmp_path / "token_dir"
@@ -202,12 +103,6 @@ def test_scores_are_per_token_inference(tmp_path):
     torch.save(torch.randn(6, 1), pt_col)
     assert not scores_are_per_token(str(pt_col))
 
-    import numpy as np
-
-    npy = tmp_path / "multi_query.npy"
-    np.save(npy, np.random.randn(6, 5))
-    assert not scores_are_per_token(str(npy))
-
     plain_dir = tmp_path / "plain_dir"
     plain_dir.mkdir()
     assert not scores_are_per_token(str(plain_dir))
@@ -225,3 +120,18 @@ def test_load_token_dir_keeps_query_dim_when_multiscore(tmp_path):
     torch.testing.assert_close(scores[0, 0], torch.tensor([1.0, 2.0]))
     torch.testing.assert_close(scores[1, 0], torch.tensor([5.0, 6.0]))
     torch.testing.assert_close(scores[1, 1], torch.tensor([0.0, 0.0]))
+
+
+def test_scores_are_per_token_from_run_config(tmp_path):
+    """With a run config beside it, the flag decides rather than the shape --
+    including ``attribute_tokens``, not just the deprecated ``per_token``."""
+    import yaml
+
+    from bergson.magic.cli import scores_are_per_token
+
+    # A shape the fallback above would read as per-doc.
+    torch.save(torch.zeros(4, 1), tmp_path / "scores.pt")
+    cfg = {"steps": [{"magic": {"query_method": "mean", "attribute_tokens": True}}]}
+    (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg))
+
+    assert scores_are_per_token(str(tmp_path / "scores.pt"))