fix: v4.1.2 hotfix - KV cache at batch >= 2, TP/DP checkpoint names, MPS daemon scope

opt_prime/opt_prime/inference.py

@@ -1047,9 +1047,7 @@ def _generate_with_kv_cache(
 
             # Forward single token via decode_step
             if self.tpl.is_first_stage():
-                token_input = next_token.unsqueeze(0) if next_token.dim() == 1 else next_token
-                if token_input.dim() == 1:
-                    token_input = token_input.unsqueeze(0)  # [1, 1]
+                token_input = next_token.reshape(-1, 1)  # [B] -> [B, 1]
                 logits = self._schedule_gen.decode_step(token_input, cur_pos)
             else:
                 logits = self._schedule_gen.decode_step(None, cur_pos)

opt_prime/opt_prime/mps_manager.py

@@ -502,34 +502,12 @@ def cleanup_barrier_files(pipe_dir: str) -> None:
 # ---------------------------------------------------------------------------
 
 def _mps_daemon_running(pipe_dir: str) -> bool:
-    """Best-effort check whether an MPS daemon is already running.
+    """Whether an MPS daemon serves this pipe directory (its control socket exists).
 
-    Looks for the pipe directory's control socket and probes pgrep.
-    Returns True if either signal is present.
+    A host-wide process check would mistake another job's daemon (other pipe
+    directory) for ours, and the stale cleanup would then hit it.
     """
-    # 1) control socket presence (most reliable when MPS daemon started cleanly)
-    if os.path.exists(os.path.join(pipe_dir, "control")):
-        return True
-    # 2) Process check
-    try:
-        out = subprocess.run(
-            ["pgrep", "-x", "nvidia-cuda-mps-server"],
-            capture_output=True, text=True, timeout=2,
-        )
-        if out.returncode == 0 and out.stdout.strip():
-            return True
-    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
-        pass
-    try:
-        out = subprocess.run(
-            ["pgrep", "-x", "nvidia-cuda-mps-control"],
-            capture_output=True, text=True, timeout=2,
-        )
-        if out.returncode == 0 and out.stdout.strip():
-            return True
-    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
-        pass
-    return False
+    return os.path.exists(os.path.join(pipe_dir, "control"))
 
 
 # ---------------------------------------------------------------------------
@@ -786,7 +764,7 @@ def _daemon_responsive(self, timeout_s: float = 2.0) -> bool:
             return False
 
     def _force_cleanup_stale_daemon(self) -> None:
-        """Best-effort cleanup of a stale daemon: try quit, then pkill (own user)."""
+        """Best-effort cleanup of a stale daemon on our pipe dir: quit, drop socket."""
         try:
             subprocess.run(
                 ["nvidia-cuda-mps-control"],
@@ -798,16 +776,8 @@ def _force_cleanup_stale_daemon(self) -> None:
             )
         except (subprocess.TimeoutExpired, OSError):
             pass
-        # Force kill stragglers (only this user's processes)
-        for proc_name in ("nvidia-cuda-mps-server", "nvidia-cuda-mps-control"):
-            try:
-                subprocess.run(
-                    ["pkill", "-u", str(os.geteuid()), "-x", proc_name],
-                    capture_output=True, timeout=5,
-                )
-            except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
-                pass
-        # Remove control socket if present
+        # No pkill: other jobs' daemons (other pipe directories) share the
+        # process names. Remove our control socket so a new daemon can bind.
         ctrl = os.path.join(self.pipe_dir, "control")
         try:
             os.unlink(ctrl)

opt_prime/opt_prime/opti_pri.py

@@ -461,7 +461,7 @@ def __init__(self, module:nn.Module, num_mb, use_gpu=True, pp_size=1, dp_size=1,
         self.tpl = Topology(rank, local_rank, world_size, pp_size, dp_size, tp_size)
 
         self.checkpoint = checkpoint
-        self.ckpt_dir = f"checkpoint_{ckpt_dir_postfix}_{self.tpl.stage}" or f"checkpoint_{self.tpl.stage}"
+        self.ckpt_dir = f"checkpoint_{ckpt_dir_postfix}_{self.tpl.stage}" if ckpt_dir_postfix else f"checkpoint_{self.tpl.stage}"
         if self.checkpoint == True:
             os.makedirs(self.ckpt_dir, exist_ok=True)
 
@@ -1165,6 +1165,10 @@ def save_lora_ckpt(self, step, epoch, lora_dir=None):
             print("[LoRA] No LoRA config set, skipping save.")
             return
 
+        # DP replicas are identical → only DP rank 0 saves
+        if (self.get_rank() // self.tpl.tp_size) % self.tpl.dp_size != 0:
+            return
+
         if lora_dir is None:
             lora_dir = f"lora_checkpoint_stage_{self.tpl.stage}"
         else:
@@ -1234,7 +1238,11 @@ def save_ckpt(self, step, epoch):
         if self.checkpoint != True:
             return
 
-        ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}.pt")
+        # DP replicas are identical → only DP rank 0 saves, one file per TP shard
+        if (self.get_rank() // self.tpl.tp_size) % self.tpl.dp_size != 0:
+            return
+        tp = f"_tp{self.get_rank() % self.tpl.tp_size}" if self.tpl.tp_size > 1 else ""
+        ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}{tp}.pt")
 
         torch.save({
             'epoch': epoch,
@@ -1245,7 +1253,7 @@ def save_ckpt(self, step, epoch):
         print(f"Checkpoint saved: {ckpt_path}")
 
         ckpt_files = sorted(
-                glob.glob(os.path.join(self.ckpt_dir, "ckpt_step_*.pt")),
+                glob.glob(os.path.join(self.ckpt_dir, f"ckpt_step_*{tp}.pt")),
                 key=os.path.getmtime,
                 reverse=True
                 )
@@ -1257,7 +1265,8 @@ def load_ckpt(self, step, epoch):
         if self.checkpoint != True:
             return
 
-        ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}.pt")
+        tp = f"_tp{self.get_rank() % self.tpl.tp_size}" if self.tpl.tp_size > 1 else ""
+        ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}{tp}.pt")
 
         ckpt = torch.load(ckpt_path)
 

opt_prime/opt_prime/schedule_infer.py

@@ -370,7 +370,7 @@ def prefill(
 
         Args:
             input_ids: Input token IDs tensor of shape [batch_size, seq_len]
-            position_ids: Position IDs tensor of shape [1, seq_len] for RoPE.
+            position_ids: Position IDs tensor of shape [batch_size, seq_len] for RoPE.
                           If None, auto-generated as arange(seq_len).
 
         Returns:
@@ -383,9 +383,10 @@ def prefill(
             # Set position_ids for RoPE (critical for LLaMA)
             if position_ids is None and input_ids is not None:
                 seq_len = input_ids.size(1) if input_ids.dim() > 1 else input_ids.size(0)
+                # The export ties position_ids' batch dim to input_ids'
                 position_ids = torch.arange(
                     seq_len, device=input_ids.device
-                ).unsqueeze(0)
+                ).expand(input_ids.size(0) if input_ids.dim() > 1 else 1, -1)
             if position_ids is not None:
                 self.optimus.run_info.env[0]["position_ids"] = \
                     position_ids.to(self.optimus.run_info.device)
@@ -427,8 +428,8 @@ def decode_step(
         if self.optimus.tpl.is_first_stage():
             self.get_input(input_token)
             # Set position_ids for RoPE — must be the absolute position
-            position_ids = torch.tensor(
-                [[position]], device=input_token.device
+            position_ids = torch.full(
+                (input_token.size(0), 1), position, device=input_token.device
             )
             self.optimus.run_info.env[0]["position_ids"] = \
                 position_ids.to(self.optimus.run_info.device)

opt_prime/tests/test_ckpt_names.py

@@ -0,0 +1,76 @@
+#
+# Copyright (c) 2026-present, ETRI, All rights reserved.
+#
+# opt_prime.opti_pri checkpoint naming: TP shards get their own files, only DP rank 0 writes, and the directory
+# name never contains "None". Ranks are laid out as the (pp, dp, tp) device mesh: rank = (pp*dp + dp_i)*tp + tp_i.
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from opt_prime import opti_pri  # noqa: E402
+from opt_prime.opti_pri import Optimus_p  # noqa: E402
+
+
+def _rank(rank, ckpt_dir, tp_size=2, dp_size=2):
+    o = Optimus_p.__new__(Optimus_p)                    # no process group: just the state the save/load paths read
+    o.tpl = SimpleNamespace(rank=rank, stage=0, tp_size=tp_size, dp_size=dp_size)
+    o.checkpoint, o.ckpt_dir, o.lora_config, o.device = True, str(ckpt_dir), "lora", torch.device("cpu")
+    torch.manual_seed(rank)                             # every rank holds different weights
+    submod = nn.ModuleDict({"lora_A": nn.Linear(2, 2)})
+    o.run_info = SimpleNamespace(submod=submod)
+    o.optimizer = torch.optim.SGD(submod.parameters(), lr=0.1)
+    return o
+
+
+def _weight(o):
+    return o.run_info.submod["lora_A"].weight.detach().clone()
+
+
+def test_save_ckpt_one_file_per_tp_shard_from_dp_rank_0(tmp_path):
+    ranks = [_rank(r, tmp_path) for r in range(4)]      # tp_i = r % 2, dp_i = r // 2
+    for o in ranks:
+        o.save_ckpt(1, 0)
+    assert sorted(os.listdir(tmp_path)) == ["ckpt_step_1_epoch_0_tp0.pt", "ckpt_step_1_epoch_0_tp1.pt"]
+    expected = _weight(ranks[1])                        # TP rank 1 of DP rank 0
+    reload = _rank(3, tmp_path)                         # TP rank 1 of DP rank 1 reloads the same shard
+    reload.load_ckpt(1, 0)
+    assert torch.equal(_weight(reload), expected)
+
+
+def test_save_ckpt_pp_only_name_unchanged(tmp_path):
+    _rank(0, tmp_path, tp_size=1, dp_size=1).save_ckpt(1, 0)   # existing PP-only checkpoints keep loading
+    assert os.listdir(tmp_path) == ["ckpt_step_1_epoch_0.pt"]
+
+
+def test_save_lora_ckpt_only_dp_rank_0_writes(tmp_path):
+    ranks = [_rank(r, tmp_path) for r in range(4)]
+    for o in ranks:
+        o.save_lora_ckpt(1, 0, lora_dir=str(tmp_path / "lora"))
+    files = sorted(os.listdir(str(tmp_path / "lora") + "_stage_0"))
+    assert files == ["lora_step_1_epoch_0_tp0.pt", "lora_step_1_epoch_0_tp1.pt"]
+    expected = _weight(ranks[1])
+    reload = _rank(3, tmp_path)
+    reload.load_lora_ckpt(1, 0, lora_dir=str(tmp_path / "lora"))
+    assert torch.equal(_weight(reload), expected)
+
+
+class _Stop(Exception):
+    pass
+
+
+@pytest.mark.parametrize("postfix, name", [(None, "checkpoint_0"), ("run1", "checkpoint_run1_0")])
+def test_ckpt_dir_name(tmp_path, monkeypatch, postfix, name):
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setattr(opti_pri, "Comm", lambda **kw: SimpleNamespace(rank=0, world_size=1, local_rank=0))
+
+    def stop(*a):                                       # __init__ names the directory just before picking the GPU
+        raise _Stop
+    monkeypatch.setattr(torch.cuda, "set_device", stop)
+    with pytest.raises(_Stop):
+        Optimus_p(nn.Linear(2, 2), 1, checkpoint=True, ckpt_dir_postfix=postfix)
+    assert os.listdir(tmp_path) == [name]

opt_prime/tests/test_kv_cache_batch.py

@@ -0,0 +1,37 @@
+#
+# Copyright (c) 2026-present, ETRI, All rights reserved.
+#
+# opt_prime.inference: greedy generation with the KV cache (torch.export path) matches HF at batch size 1 and 2.
+import copy
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Optimus_Inference needs a GPU")
+
+
+@pytest.fixture(scope="module")
+def tiny():
+    from transformers import AutoConfig, AutoModelForCausalLM
+    torch.manual_seed(0)
+    cfg = AutoConfig.for_model("llama", hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+                               num_attention_heads=4, num_key_value_heads=2, vocab_size=128,
+                               attn_implementation="sdpa")
+    return AutoModelForCausalLM.from_config(cfg).eval()
+
+
+@pytest.mark.parametrize("batch", [1, 2])
+def test_kv_cache_generation_matches_hf(tiny, batch):
+    from opt_prime.inference import Optimus_Inference
+    ids = torch.randint(0, 128, (batch, 6), generator=torch.Generator().manual_seed(batch))
+    with torch.no_grad():
+        ref = copy.deepcopy(tiny).cuda().generate(ids.cuda(), max_new_tokens=5, do_sample=False,
+                                                  pad_token_id=0).cpu()
+    eng = Optimus_Inference(copy.deepcopy(tiny), use_kv_cache=True, dynamo_capture=True, dtype=torch.float32)
+    eng.eval()
+    out = eng.generate(ids, max_new_tokens=5, do_sample=False, verbose=False).cpu()
+    assert torch.equal(out, ref)

opt_prime/tests/test_mps_manager.py

@@ -0,0 +1,36 @@
+#
+# Copyright (c) 2026-present, ETRI, All rights reserved.
+#
+# opt_prime.mps_manager: daemon detection and stale cleanup stay scoped to this job's pipe directory.
+import os
+import subprocess
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from opt_prime import mps_manager  # noqa: E402
+
+
+def _fake_run(calls):
+    # Every command "succeeds" as if another job's MPS daemon were running on this host.
+    def run(cmd, *a, **kw):
+        calls.append(cmd)
+        return subprocess.CompletedProcess(cmd, 0, stdout="4242\n", stderr="")
+    return run
+
+
+def test_daemon_running_only_from_own_control_socket(tmp_path, monkeypatch):
+    calls = []
+    monkeypatch.setattr(subprocess, "run", _fake_run(calls))
+    assert not mps_manager._mps_daemon_running(str(tmp_path))      # another job's daemon is not ours
+    (tmp_path / "control").touch()
+    assert mps_manager._mps_daemon_running(str(tmp_path))
+
+
+def test_stale_cleanup_does_not_kill_other_daemons(tmp_path, monkeypatch):
+    calls = []
+    monkeypatch.setattr(subprocess, "run", _fake_run(calls))
+    monkeypatch.setattr(mps_manager.time, "sleep", lambda s: None)
+    (tmp_path / "control").touch()
+    mps_manager.MPSManager(pipe_dir=str(tmp_path), log_dir=str(tmp_path / "log"))._force_cleanup_stale_daemon()
+    assert not any(c[0] in ("pkill", "killall") for c in calls), calls
+    assert not (tmp_path / "control").exists()