Repair C++ logit validation

CMakeLists.txt

@@ -53,6 +53,14 @@ endif()
 # ── Tests ───────────────────────────────────────────────────────
 if(DIFFUSE_BUILD_TESTS)
     enable_testing()
+    find_package(Python3 COMPONENTS Interpreter QUIET)
+    if(Python3_Interpreter_FOUND)
+        add_test(
+            NAME test-validate-logits
+            COMMAND ${Python3_EXECUTABLE}
+                    ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-validate-logits.py)
+    endif()
+
     add_executable(test-forward tests/test-forward.cpp)
     target_link_libraries(test-forward PRIVATE diffuse)
     add_test(NAME test-forward COMMAND test-forward)

tests/test-e2e.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """End-to-end test: create tiny model, convert, load in C++, run forward pass."""
 
+import importlib.util
 import json
 import os
 import sys
@@ -102,6 +103,26 @@ def main():
             print(f"FAIL: test-forward exited with code {result.returncode}")
             sys.exit(1)
 
+        # Exercise the real dump-logits executable through the same Python
+        # wrapper used for PyTorch/C++ validation.
+        validator_path = os.path.join(project_dir, "tools", "validate-logits.py")
+        validator_spec = importlib.util.spec_from_file_location(
+            "diffuse_validate_logits", validator_path)
+        validator = importlib.util.module_from_spec(validator_spec)
+        validator_spec.loader.exec_module(validator)
+
+        dump_binary = os.path.join(build_dir, "dump-logits")
+        dump_tokens = [1, 2, 3, 4]
+        dumped_logits = validator.run_cpp(
+            gguf_path, dump_binary, dump_tokens, n_threads=4)
+        if dumped_logits.shape != (len(dump_tokens), VOCAB):
+            print(f"FAIL: unexpected dumped logit shape {dumped_logits.shape}")
+            sys.exit(1)
+        if not np.all(np.isfinite(dumped_logits)):
+            print("FAIL: dump-logits returned non-finite values")
+            sys.exit(1)
+        print("dump-logits wrapper OK")
+
         print("\nEnd-to-end test PASSED!")
 
 

tests/test-validate-logits.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Regression tests for the dump-logits Python wrapper."""
+
+import importlib.util
+import os
+import stat
+import tempfile
+import unittest
+
+import numpy as np
+
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+VALIDATOR_PATH = os.path.join(PROJECT_DIR, "tools", "validate-logits.py")
+VALIDATOR_SPEC = importlib.util.spec_from_file_location(
+    "diffuse_validate_logits", VALIDATOR_PATH)
+VALIDATOR = importlib.util.module_from_spec(VALIDATOR_SPEC)
+VALIDATOR_SPEC.loader.exec_module(VALIDATOR)
+
+
+class ValidateLogitsTest(unittest.TestCase):
+    def test_run_cpp_reads_shape_and_logits(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            fake_dump = os.path.join(tmpdir, "fake-dump-logits")
+            with open(fake_dump, "w", encoding="utf-8") as script:
+                script.write(
+                    "#!/usr/bin/env python3\n"
+                    "import struct, sys\n"
+                    "output = sys.argv[sys.argv.index('-o') + 1]\n"
+                    "with open(output, 'wb') as f:\n"
+                    "    f.write(struct.pack('<ii6f', 2, 3, 1, 2, 3, 4, 5, 6))\n"
+                )
+            os.chmod(fake_dump, os.stat(fake_dump).st_mode | stat.S_IXUSR)
+
+            logits = VALIDATOR.run_cpp(
+                "unused.gguf", fake_dump, [10, 11], n_threads=2)
+
+        np.testing.assert_array_equal(
+            logits,
+            np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tools/dump-logits.cpp

@@ -6,6 +6,7 @@
 
 #include "diffuse.h"
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>

tools/validate-logits.py

@@ -8,23 +8,20 @@
     python validate-logits.py \
         --model /path/to/LLaDA-8B-Instruct \
         --gguf llada-8b-f16.gguf \
-        --cpp-bin ./build/diffuse-cli \
+        --cpp-bin ./build/dump-logits \
         --tokens "1,2,3,4,5,6,7,8"
 
     # Or with prompt (requires tokenizer):
     python validate-logits.py \
         --model /path/to/LLaDA-8B-Instruct \
         --gguf llada-8b-f16.gguf \
-        --cpp-bin ./build/diffuse-cli \
+        --cpp-bin ./build/dump-logits \
         --prompt "Hello world"
 """
 
 import argparse
-import json
 import os
-import struct
 import subprocess
-import sys
 import tempfile
 import time
 
@@ -87,34 +84,44 @@ def run_pytorch(model_dir, token_ids):
 
 
 def run_cpp(gguf_path, cpp_bin, token_ids, n_threads=4):
-    """Run forward pass with diffuse-cpp. Returns logits [n_tokens, vocab_size].
-
-    Since the CLI doesn't output raw logits yet, we use a helper binary.
-    For now, write a small C++ helper that dumps logits to a binary file.
-    """
-    # Build tokens string
+    """Run dump-logits and return float32 logits [n_tokens, vocab_size]."""
     tokens_str = ",".join(map(str, token_ids))
-
-    # We need a way to get raw logits from C++. For validation, we write
-    # a temporary Python script that calls the C library via ctypes.
-    # Actually, the simplest approach: modify the test to dump logits.
-    # For now, use the test-forward binary approach:
-
-    # Write a small Python script that uses the GGUF file to get metadata,
-    # then we'll compare at the tensor level instead.
-    print("NOTE: Raw logits comparison requires the C++ binary to dump logits.")
-    print("      For now, validating conversion correctness at tensor level.")
-    print("      Full logit comparison will be available after adding --dump-logits to CLI.")
-
-    return None
+    with tempfile.TemporaryDirectory() as tmpdir:
+        output_path = os.path.join(tmpdir, "logits.bin")
+        command = [
+            cpp_bin,
+            "-m", gguf_path,
+            "--tokens", tokens_str,
+            "-o", output_path,
+            "-t", str(n_threads),
+        ]
+        result = subprocess.run(command, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"dump-logits exited with code {result.returncode}:\n{result.stderr}"
+            )
+
+        with open(output_path, "rb") as logits_file:
+            header = np.fromfile(logits_file, dtype=np.int32, count=2)
+            logits = np.fromfile(logits_file, dtype=np.float32)
+
+    if header.size != 2:
+        raise RuntimeError("dump-logits output is missing its shape header")
+    n_tokens, n_vocab = map(int, header)
+    if n_tokens != len(token_ids):
+        raise RuntimeError(
+            f"dump-logits returned {n_tokens} rows for {len(token_ids)} tokens"
+        )
+    if logits.size != n_tokens * n_vocab:
+        raise RuntimeError(
+            f"dump-logits wrote {logits.size} values for shape "
+            f"{n_tokens} x {n_vocab}"
+        )
+    return logits.reshape(n_tokens, n_vocab)
 
 
 def compare_logits(pytorch_logits, cpp_logits, top_k=10):
     """Compare logit distributions between PyTorch and C++."""
-    if cpp_logits is None:
-        print("\nSkipping logit comparison (C++ logits not available)")
-        return
-
     n_tokens, vocab_size = pytorch_logits.shape
     assert cpp_logits.shape == pytorch_logits.shape, \
         f"Shape mismatch: PyTorch {pytorch_logits.shape} vs C++ {cpp_logits.shape}"
@@ -155,7 +162,8 @@ def main():
     parser = argparse.ArgumentParser(description="Validate forward pass logits")
     parser.add_argument("--model", "-m", required=True, help="HF model directory")
     parser.add_argument("--gguf", "-g", required=True, help="GGUF file path")
-    parser.add_argument("--cpp-bin", default="./build/diffuse-cli", help="C++ CLI binary")
+    parser.add_argument("--cpp-bin", default="./build/dump-logits",
+                        help="Path to the dump-logits binary")
     parser.add_argument("--tokens", help="Comma-separated token IDs")
     parser.add_argument("--prompt", "-p", help="Text prompt (tokenized automatically)")
     parser.add_argument("--threads", "-t", type=int, default=4, help="C++ threads")