Merge remote-tracking branch 'upstream/main' into feat/gpu-backend-hardening

c/tests/bench_dsa_select.c

@@ -59,6 +59,12 @@ static double bench_ns(void (*fn)(const float*,int,int,int*,int*),
     return ts[N_REPEAT/2];
 }
 
+/* Sort an array of doubles ascending (median-of-medians aggregation below). */
+static void dsort(double *a, int n){
+    for(int s=1;s<n;s++){ double k=a[s]; int b=s-1;
+        while(b>=0 && a[b]>k){ a[b+1]=a[b]; b--; } a[b+1]=k; }
+}
+
 /* the NEW algorithm calls the real partial_select_desc + replicates the production
  * threshold derivation and position scans. */
 static void keep_new(const float *isc, int nk, int keep, int *dst, int *nd_out){
@@ -74,6 +80,7 @@ static void keep_new(const float *isc, int nk, int keep, int *dst, int *nd_out){
 
 /* deterministic score fill for three shapes */
 static uint32_t brng = 0xA5A5A5A5u;
+static void brng_seed(uint32_t s){ brng = s; }
 static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
     return (double)(brng >> 8) * (1.0 / 16777216.0); }
 static void fill_realistic(float *isc, int nk){   /* few hot, long distinct tail */
@@ -87,11 +94,24 @@ static void fill_plateau(float *isc, int nk){     /* tie blocks -> boundary path
     for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7));
 }
 
+/* Multi-seed aggregation. Per @KingIcyCreamProjects (#357 thread): with a single frozen
+ * input per cell, quickselect's deterministic median-of-three pivot means one lucky input
+ * can spike a single nk row (an observed ~75x at nk=8192 on a 9950X3D that was really a
+ * ~13-40x algorithm). Two bugs compounded it: (1) the old bench drew ONE input per cell;
+ * (2) brng was never reset, so each cell's input depended on every prior cell's draws --
+ * reordering nks[] silently shifted all later inputs. Both fixed here: brng is reseeded
+ * per (shape, nk, seed), and we take the MEDIAN of N_SEEDS independent inputs, each itself
+ * a median over N_REPEAT timing reps. A lucky pivot now moves one of the N_SEEDS samples,
+ * not the reported number. */
+
 int main(void){
     int keep = 2048;   /* GLM-5.2 index_topk */
     int nks[] = {2049, 4096, 8192, 16384, 32768, 65536};
+    const int N_SEEDS = 11;   /* odd so the median is a real sample, not interpolated */
     float *isc = malloc((size_t)65536*sizeof(float));
     int *dst   = malloc((size_t)65536*sizeof(int)); int nd;
+    double *old_seeds = malloc((size_t)N_SEEDS*sizeof(double));
+    double *new_seeds = malloc((size_t)N_SEEDS*sizeof(double));
 
     struct { const char *name; void (*fill)(float*,int); } shapes[] = {
         { "realistic", fill_realistic },
@@ -99,27 +119,40 @@ int main(void){
         { "plateau",   fill_plateau   },
     };
 
-    printf("bench_dsa_select: DSA top-keep, old (qsort) vs new (partial-select)  keep=%d\n", keep);
+    printf("bench_dsa_select: DSA top-keep, old (qsort) vs new (partial-select)  keep=%d  (median of %d seeds x %d reps)\n",
+           keep, N_SEEDS, N_REPEAT);
     printf("%-12s %7s %14s %14s %9s\n", "shape", "nk", "old ns/call", "new ns/call", "speedup");
     printf("------------------------------------------------------------------------\n");
 
     for(size_t sh=0; sh<sizeof(shapes)/sizeof(shapes[0]); sh++){
         for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
             int nk=nks[ni];
-            shapes[sh].fill(isc,nk);
-            /* warmup both paths so caches/branch predictors are primed */
-            for(int w=0; w<50; w++){ keep_old(isc,nk,keep,dst,&nd); keep_new(isc,nk,keep,dst,&nd); }
-            /* sanity: both must keep exactly `keep` (correctness is test_dsa_select's
-             * job, but a count divergence here would make the timing meaningless) */
-            keep_old(isc,nk,keep,dst,&nd); int na=nd;
-            keep_new(isc,nk,keep,dst,&nd); int nb=nd;
-            if(na!=keep || nb!=keep){
-                printf("%-12s %7d   (BAD COUNTS: old=%d new=%d, skipped)\n",
-                       shapes[sh].name, nk, na, nb);
+            int bad = 0;
+            for(int sd=0; sd<N_SEEDS; sd++){
+                /* reseed per (shape,nk,seed) so each cell's input is reproducible and
+                 * independent of cell ordering, and so lucky pivots are sampled, not fixed. */
+                brng_seed(0xA5A5A5A5u + (uint32_t)(sd*0x9E3779B9u));
+                shapes[sh].fill(isc,nk);
+                /* warmup both paths so caches/branch predictors are primed */
+                for(int w=0; w<50; w++){ keep_old(isc,nk,keep,dst,&nd); keep_new(isc,nk,keep,dst,&nd); }
+                /* sanity: both must keep exactly `keep` (correctness is test_dsa_select's
+                 * job, but a count divergence here would make the timing meaningless) */
+                keep_old(isc,nk,keep,dst,&nd); int na=nd;
+                keep_new(isc,nk,keep,dst,&nd); int nb=nd;
+                if(na!=keep || nb!=keep){ bad++; continue; }
+                old_seeds[sd] = bench_ns(keep_old,isc,nk,keep);
+                new_seeds[sd] = bench_ns(keep_new,isc,nk,keep);
+            }
+            if(bad == N_SEEDS){
+                printf("%-12s %7d   (BAD COUNTS on all %d seeds, skipped)\n",
+                       shapes[sh].name, nk, bad);
                 continue;
             }
-            double t_old=bench_ns(keep_old,isc,nk,keep);
-            double t_new=bench_ns(keep_new,isc,nk,keep);
+            /* report median-of-seed-medians: robust to a single lucky/unlucky pivot */
+            dsort(old_seeds, N_SEEDS);
+            dsort(new_seeds, N_SEEDS);
+            double t_old = old_seeds[N_SEEDS/2];
+            double t_new = new_seeds[N_SEEDS/2];
             printf("%-12s %7d %14.0f %14.0f %8.2fx\n",
                    shapes[sh].name, nk, t_old, t_new, t_old/t_new);
         }
@@ -127,6 +160,6 @@ int main(void){
     }
 
     printf("bench_dsa_select: done (lower ns is better; speedup = old/new)\n");
-    free(isc); free(dst);
+    free(isc); free(dst); free(old_seeds); free(new_seeds);
     return 0;
 }

c/tests/test_benchmark_cuda_fixture.py

@@ -1,11 +1,12 @@
 import unittest
 
-from tools.benchmark_cuda_fixture import parse_output
+from tools.benchmark_cuda_fixture import parse_output, parse_p0
 
 
 SAMPLE = """
 REPLAY decode: 4 tokens | 12.34 tok/s
 PROFILE: expert-disk 1.25s | expert-matmul 2.50s | attention 0.75s | lm_head 0.10s | other -0.05s
+P0-EXEC: routed CPU 1.200s / 123.40 GB/s (456 row) | routed GPU critical 0.150s | router 0.200s | residual P2P 0.030s / 75 hop | orchestration 0.100s
 """
 
 
@@ -19,6 +20,9 @@ class ParseOutputTest(unittest.TestCase):
         with self.assertRaisesRegex(RuntimeError, "benchmark output missing"):
             parse_output("REPLAY decode: 4 tokens | 12.34 tok/s", "engine failed")
 
+    def test_extracts_p0_profile(self):
+        self.assertEqual(parse_p0(SAMPLE), [1.2, 123.4, 456.0, 0.15, 0.2, 0.03, 75.0, 0.1])
+
 
 if __name__ == "__main__":
     unittest.main()

c/tests/test_logit_nan.c

@@ -0,0 +1,70 @@
+/* Regression for non-finite-logit poisoning of sampling.
+ *
+ * A single NaN or +Inf in the logits (a bad streamed expert tile, or an fp
+ * overflow in the matmul at a low-RAM eviction boundary) used to make softmax
+ * produce an all-NaN g_pbuf; dist_sample then never satisfied cum>=u and fell
+ * through to return token 0 — so the engine silently emitted an unbroken run of
+ * token 0 with no error, on the DEFAULT serve path (TEMP>0, 0<NUCLEUS<1).
+ *
+ * Fix under test: argmax_v() skips NaN (picks the max finite/+Inf entry instead
+ * of being pinned to index 0), and dist_build() detects a non-finite softmax sum
+ * and collapses to a one-hot over the finite argmax (warning once) rather than
+ * dividing every entry into NaN. Degrade + diagnose, never silently corrupt.
+ *
+ * No model file needed: exercises argmax_v / dist_build / dist_sample directly. */
+#include <assert.h>
+#include <math.h>
+#define main coli_glm_main_unused
+#include "../glm.c"
+#undef main
+
+static int approx1(double x){ return x > 0.999 && x < 1.001; }
+
+int main(void){
+    /* --- argmax_v must skip NaN (greedy decode + speculative-verify paths) --- */
+    { float lo[8]={NAN,1.f,5.f,2.f,NAN,-3.f,4.f,0.f};
+      assert(argmax_v(lo,8)==2 && "pick max finite (idx2=5.0), not NaN-pinned idx0"); }
+    { float lo[8]={3.f,INFINITY,1.f,2.f,0.f,-1.f,2.5f,1.5f};
+      assert(argmax_v(lo,8)==1 && "pick the +Inf position"); }
+    { float lo[8]; for(int i=0;i<8;i++) lo[i]=NAN;
+      assert(argmax_v(lo,8)==0 && "all-NaN: no crash, defined fallback"); }
+
+    g_temp=0.7f; g_nuc=0.9f;   /* the default serve/chat sampling path */
+
+    /* --- dist_build: a NaN logit must yield a finite one-hot, not all-NaN --- */
+    { float lo[8]={0.5f,1.f,NAN,8.f,0.2f,-1.f,0.f,0.3f};   /* max finite = idx3 (8.0) */
+      dist_build(lo,8);
+      double sum=0; int nan=0;
+      for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
+      assert(!nan && "g_pbuf must be finite after a NaN logit");
+      assert(approx1(sum) && "g_pbuf must normalize to 1");
+      assert(approx1(g_pbuf[3]) && "mass must land on the max finite logit (idx3)");
+      assert(dist_sample(8,-1)==3 && "sampler emits the finite argmax, not token 0"); }
+
+    /* --- NaN at index 0: poisons the max scan itself (the old mx=lo[0] seed),
+     *     the failure mode that starts before the sum (review note on #369) --- */
+    { float lo[8]={NAN,1.f,0.5f,6.f,0.2f,-1.f,0.f,0.3f};   /* max finite = idx3 (6.0) */
+      dist_build(lo,8);
+      double sum=0; int nan=0;
+      for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
+      assert(!nan && "NaN at lo[0] must not poison via the mx seed");
+      assert(approx1(sum) && "still normalizes to 1");
+      assert(dist_sample(8,-1)==3 && "emits the max finite logit, not token 0"); }
+
+    /* --- all-NaN logits: worst case — must stay finite, no crash --- */
+    { float lo[8]; for(int i=0;i<8;i++) lo[i]=NAN;
+      dist_build(lo,8);
+      for(int i=0;i<8;i++) assert(g_pbuf[i]==g_pbuf[i] && "all-NaN: g_pbuf stays finite"); }
+
+    /* --- regression: clean logits still produce a valid distribution --- */
+    { float lo[8]={0.1f,0.2f,3.0f,0.4f,0.5f,0.6f,0.7f,0.8f};   /* peak = idx2 */
+      dist_build(lo,8);
+      double sum=0; int nan=0;
+      for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
+      assert(!nan && "clean softmax stays finite");
+      assert(approx1(sum) && "clean softmax must sum to 1");
+      assert(g_pbuf[2]>=g_pbuf[0] && "peak token keeps the most mass"); }
+
+    printf("OK test_logit_nan: argmax_v NaN-skip + dist_build finite-collapse\n");
+    return 0;
+}

c/tests/test_openai_tools_e2e.py

@@ -0,0 +1,182 @@
+"""End-to-end tool-calling test for the OpenAI gateway (#401).
+
+Unlike the unit tests in test_openai_server.py (which call parse_tool_calls /
+render_chat directly), this suite runs openai_server.py as a real subprocess
+against a mock engine that speaks the actual SERVE wire protocol
+(READY / SUBMIT / DATA / DONE), then talks to it over real HTTP. It pins down
+the full path a coding client exercises: tool declaration rendering, marker
+suppression in streamed deltas (across chunk boundaries), tool_calls in both
+response shapes, and the <|observation|><tool_response> round trip.
+"""
+import json
+import os
+import socket
+import subprocess
+import sys
+import tempfile
+import unittest
+import urllib.request
+from pathlib import Path
+
+SERVER = Path(__file__).resolve().parent.parent / "openai_server.py"
+MODEL_ID = "glm-5.2-colibri"
+
+# Mock engine: replies are keyed on the prompt so one process covers every case.
+# Prompts received are appended to MOCK_LOG for assertions on the rendering.
+MOCK_ENGINE = r'''#!/usr/bin/env python3
+import sys, os
+out, inp = sys.stdout.buffer, sys.stdin.buffer
+out.write(b"\x01\x01READY\x01\x01\n" + b"STAT 0 0 0 0 0\n"); out.flush()
+
+def reply(rid, text, chunks=1):
+    data = text.encode("utf-8")
+    n = max(1, len(data) // chunks)
+    for i in range(0, len(data), n):
+        part = data[i:i+n]
+        out.write(("DATA %s %d\n" % (rid, len(part))).encode() + part + b"\n"); out.flush()
+    out.write(("DONE %s STAT %d 1.0 50.0 10.0 42 0\n" % (rid, len(text.split()))).encode())
+    out.flush()
+
+while True:
+    line = inp.readline()
+    if not line: break
+    f = line.decode().strip().split()
+    if not f or f[0] != "SUBMIT": continue
+    rid, plen = f[1], int(f[3])
+    prompt = inp.read(plen).decode("utf-8", "replace"); inp.read(1)
+    with open(os.environ["MOCK_LOG"], "a") as log:
+        log.write(prompt + "\n\x00\n")
+    if "<tool_response>" in prompt:
+        reply(rid, "25 degrees and sunny in Rome.")
+    elif "weather in Rome" in prompt:
+        reply(rid, "<tool_call>get_weather<arg_key>city</arg_key>"
+                   "<arg_value>Rome</arg_value></tool_call>")
+    elif "weather in Milan" in prompt:
+        # split across many tiny DATA chunks: streamed marker suppression must
+        # hold even when a marker straddles a chunk boundary
+        reply(rid, "Checking. <tool_call>get_weather<arg_key>city</arg_key>"
+                   "<arg_value>Milan</arg_value></tool_call>", chunks=20)
+    else:
+        reply(rid, "Hello from the mock engine.")
+'''
+
+TOOLS = [{"type": "function", "function": {
+    "name": "get_weather",
+    "description": "Current weather for a city",
+    "parameters": {"type": "object",
+                   "properties": {"city": {"type": "string"}},
+                   "required": ["city"]}}}]
+
+
+@unittest.skipUnless(os.name == "posix",
+                     "the mock engine is a shebang script the gateway execs directly; "
+                     "Windows CreateProcess cannot run it. The gateway logic under test "
+                     "is platform-independent and covered by the POSIX CI jobs.")
+class ToolCallingE2E(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tmp = tempfile.TemporaryDirectory()
+        mock = Path(cls.tmp.name) / "mock_engine.py"
+        mock.write_text(MOCK_ENGINE)
+        mock.chmod(0o755)
+        cls.mock_log = Path(cls.tmp.name) / "prompts.log"
+        cls.mock_log.touch()
+        with socket.socket() as probe:          # free port, then hand it to the server
+            probe.bind(("127.0.0.1", 0))
+            cls.port = probe.getsockname()[1]
+        env = dict(os.environ, MOCK_LOG=str(cls.mock_log))
+        cls.server = subprocess.Popen(
+            [sys.executable, str(SERVER), "--model", cls.tmp.name,
+             "--engine", str(mock), "--port", str(cls.port)],
+            env=env, stderr=subprocess.DEVNULL)
+        cls.base = f"http://127.0.0.1:{cls.port}/v1"
+        for _ in range(100):
+            try:
+                urllib.request.urlopen(cls.base + "/models", timeout=2)
+                return
+            except OSError:
+                if cls.server.poll() is not None:
+                    raise RuntimeError("gateway exited during startup")
+                import time
+                time.sleep(0.1)
+        raise RuntimeError("gateway did not come up")
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.server.terminate()
+        cls.server.wait(timeout=5)
+        cls.tmp.cleanup()
+
+    def post(self, body, stream=False):
+        req = urllib.request.Request(
+            self.base + "/chat/completions", json.dumps(body).encode(),
+            {"Content-Type": "application/json"})
+        resp = urllib.request.urlopen(req, timeout=30)
+        if not stream:
+            return json.loads(resp.read())
+        events = []
+        for raw in resp:
+            line = raw.decode().strip()
+            if line.startswith("data: ") and line != "data: [DONE]":
+                events.append(json.loads(line[6:]))
+        return events
+
+    def test_tool_call_non_stream(self):
+        r = self.post({"model": MODEL_ID, "tools": TOOLS,
+                       "messages": [{"role": "user",
+                                     "content": "What is the weather in Rome?"}]})
+        choice = r["choices"][0]
+        self.assertEqual(choice["finish_reason"], "tool_calls")
+        calls = choice["message"]["tool_calls"]
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0]["function"]["name"], "get_weather")
+        self.assertEqual(json.loads(calls[0]["function"]["arguments"]), {"city": "Rome"})
+        self.assertNotIn("<tool_call>", choice["message"].get("content") or "")
+
+    def test_tool_call_streamed_markers_suppressed(self):
+        events = self.post({"model": MODEL_ID, "tools": TOOLS, "stream": True,
+                            "messages": [{"role": "user",
+                                          "content": "What is the weather in Milan?"}]},
+                           stream=True)
+        deltas = [e["choices"][0]["delta"] for e in events if e["choices"]]
+        text = "".join(d.get("content") or "" for d in deltas)
+        self.assertNotIn("<tool_call>", text)
+        self.assertNotIn("<arg_key>", text)
+        calls = [d["tool_calls"] for d in deltas if d.get("tool_calls")]
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0][0]["function"]["name"], "get_weather")
+        self.assertEqual(json.loads(calls[0][0]["function"]["arguments"]),
+                         {"city": "Milan"})
+        finish = [e["choices"][0]["finish_reason"] for e in events
+                  if e["choices"] and e["choices"][0].get("finish_reason")]
+        self.assertEqual(finish, ["tool_calls"])
+
+    def test_tool_result_round_trip(self):
+        r = self.post({"model": MODEL_ID, "tools": TOOLS, "messages": [
+            {"role": "user", "content": "What is the weather in Rome?"},
+            {"role": "assistant", "content": None, "tool_calls": [
+                {"id": "call_x", "type": "function",
+                 "function": {"name": "get_weather",
+                              "arguments": "{\"city\": \"Rome\"}"}}]},
+            {"role": "tool", "tool_call_id": "call_x",
+             "content": "25 degrees, sunny"}]})
+        choice = r["choices"][0]
+        self.assertEqual(choice["finish_reason"], "stop")
+        self.assertFalse(choice["message"].get("tool_calls"))
+        self.assertIn("25 degrees", choice["message"]["content"])
+        rendered = self.mock_log.read_text().split("\x00")[-2]
+        self.assertIn("<|observation|><tool_response>25 degrees, sunny</tool_response>",
+                      rendered)
+        self.assertIn("# Tools", rendered)
+        self.assertIn('"get_weather"', rendered)
+
+    def test_no_tools_plain_text(self):
+        r = self.post({"model": MODEL_ID,
+                       "messages": [{"role": "user", "content": "Hi!"}]})
+        choice = r["choices"][0]
+        self.assertEqual(choice["finish_reason"], "stop")
+        self.assertIn("mock engine", choice["message"]["content"])
+
+
+if __name__ == "__main__":
+    unittest.main()

c/tests/test_ragged_attention.cu

@@ -0,0 +1,39 @@
+#include "../backend_cuda.h"
+
+#include <cmath>
+#include <cstdio>
+#include <vector>
+
+int main(){
+    int dev=0;if(!coli_cuda_init(&dev,1))return 77;
+    constexpr int S=3,H=2,Q=2,R=1,V=2,K=3,D=H*V,O=3,T=3;
+    std::vector<float> w(H*(Q+V)*K),p(O*D),q(S*H*(Q+R));
+    for(size_t i=0;i<w.size();i++)w[i]=((int)(i%11)-5)*.07f;
+    for(size_t i=0;i<p.size();i++)p[i]=((int)(i%7)-3)*.09f;
+    for(size_t i=0;i<q.size();i++)q[i]=((int)(i%13)-6)*.05f;
+    ColiCudaTensor *tw=nullptr,*tp=nullptr;
+    if(!coli_cuda_tensor_upload(&tw,w.data(),nullptr,0,K,H*(Q+V),dev)||
+       !coli_cuda_tensor_upload(&tp,p.data(),nullptr,0,D,O,dev))return 1;
+    int n[S]={1,2,3};std::vector<std::vector<float>> l(S),r(S);
+    const float *lp[S],*rp[S];
+    const void *keys[S];
+    for(int s=0;s<S;s++){
+        l[s].resize(n[s]*K);r[s].resize(n[s]*R);
+        for(size_t i=0;i<l[s].size();i++)l[s][i]=((int)((i+s*3)%9)-4)*.08f;
+        for(size_t i=0;i<r[s].size();i++)r[s][i]=((int)((i+s)%5)-2)*.06f;
+        lp[s]=l[s].data();rp[s]=r[s].data();keys[s]=&l[s];
+    }
+    float got[S*O],ref[S*O],warm[S*O];
+    int first[S]={1,1,1};
+    if(!coli_cuda_attention_project_ragged(tw,tp,warm,q.data(),keys,lp,rp,first,
+        S,H,Q,R,V,K,1,.2f))return 2;
+    if(!coli_cuda_attention_project_ragged(tw,tp,got,q.data(),keys,lp,rp,n,S,H,Q,R,V,K,T,.2f))return 2;
+    for(int s=0;s<S;s++)if(!coli_cuda_attention_project_batch(tw,tp,ref+s*O,
+        q.data()+s*H*(Q+R),lp[s],rp[s],1,H,Q,R,V,K,n[s],.2f))return 3;
+    double e=0,z=0;for(int i=0;i<S*O;i++){
+        double d=got[i]-ref[i];e+=d*d;z+=(double)ref[i]*ref[i];
+    }
+    double rms=std::sqrt(e/(z+1e-30));std::printf("ragged_relative_rms=%.9g\n",rms);
+    coli_cuda_tensor_free(tw);coli_cuda_tensor_free(tp);coli_cuda_shutdown();
+    return rms<1e-6?0:4;
+}

c/tests/test_resource_plan.py

@@ -10,6 +10,7 @@ from resource_plan import (
     GB,
     analyze_model,
     build_plan,
+    cpu_socket_count,
     environment_for_plan,
     format_plan,
     memory_available,
@@ -66,15 +67,19 @@ class ResourcePlanTest(unittest.TestCase):
         # 0 slots/layer. The value must be a sane positive number of bytes.
         self.assertGreater(memory_available(), 0)
 
+    def test_cpu_socket_count_is_positive(self):
+        self.assertGreaterEqual(cpu_socket_count(), 1)
+
     def test_builds_bounded_three_tier_plan(self):
         gpus = [{"index": 0, "name": "test-gpu", "total_bytes": 12 * GB,
                  "free_bytes": 10 * GB}]
         plan = build_plan(self.model, ram_gb=16, context=32, vram_gb=20,
                           available_memory=32 * GB, available_disk=100 * GB, gpus=gpus,
-                          physical_cpus=24)
+                          physical_cpus=24, cpu_sockets=2)
         self.assertEqual(plan["version"], 2)
         self.assertEqual(plan["policy"]["name"], "quality")
         self.assertEqual(plan["cpu"]["physical_cores"], 24)
+        self.assertEqual(plan["cpu"]["sockets"], 2)
         self.assertTrue(plan["policy"]["preserve_quantization"])
         self.assertFalse(plan["tiers"]["vram"]["requires_host_backing"])
         self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB)
@@ -105,7 +110,7 @@ class ResourcePlanTest(unittest.TestCase):
             {"index": 1, "name": "b", "total_bytes": 12 * GB, "free_bytes": 10 * GB},
         ]
         plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
-                          available_disk=1, gpus=gpus)
+                          available_disk=1, gpus=gpus, cpu_sockets=2)
         env = environment_for_plan(plan, {"RAM_GB": "12", "PIN": "stats.txt",
                                                "COLI_GPUS": "1"})
         self.assertEqual(env["RAM_GB"], "12")
@@ -125,6 +130,16 @@ class ResourcePlanTest(unittest.TestCase):
                                                         "OMP_PROC_BIND": "close"})
         self.assertEqual(explicit_threads["OMP_NUM_THREADS"], "7")
         self.assertEqual(explicit_threads["OMP_PROC_BIND"], "close")
+
+        if sys.platform.startswith("linux"):
+            self.assertEqual(env["COLI_NUMA"], "1")
+            explicit_numa = environment_for_plan(plan, {"COLI_NUMA": "0"})
+            self.assertEqual(explicit_numa["COLI_NUMA"], "0")
+
+    def test_single_socket_plan_does_not_enable_numa(self):
+        plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
+                          gpus=[], physical_cpus=8, cpu_sockets=1)
+        self.assertNotIn("COLI_NUMA", environment_for_plan(plan))
     def test_cpu_binary_does_not_apply_gpu_tier(self):
         plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
                           gpus=[{"index": 0, "name": "a", "total_bytes": 8 * GB,

c/tests/test_sample_nan.c

@@ -0,0 +1,62 @@
+/* Regressione #369: un solo logit NaN o +Inf faceva emettere a dist_build/dist_sample
+ * il token 0 PER SEMPRE, in silenzio (il fallback `g_pbuf[i]>0` e' falso su NaN ovunque).
+ * Ora la distribuzione degenere ripiega sull'argmax dei logit FINITI e avvisa una volta.
+ *
+ * Il test verifica: (a) con logit sani il campionamento resta corretto; (b) con NaN/+Inf
+ * iniettato il token scelto e' l'argmax dei FINITI (mai 0 per default), su ogni posizione
+ * del NaN inclusa lo[0]; (c) nessun NaN sopravvive in g_pbuf. */
+#define main coli_glm_main_unused
+#include "../glm.c"
+#undef main
+#include <stdio.h>
+
+static uint32_t rs=0x1234abcd; static uint32_t xr(){rs^=rs<<13;rs^=rs>>17;rs^=rs<<5;return rs;}
+
+static int pbuf_has_nan(int V){ for(int i=0;i<V;i++) if(!isfinite(g_pbuf[i])) return 1; return 0; }
+
+int main(void){
+    int fail=0, V=2000;
+    float *lo=malloc(V*sizeof(float));
+    g_temp=0.7f; g_nuc=0.90f;                     /* il path serve di default */
+
+    /* (a) logit sani: il campionato deve avere prob > 0 e nessun NaN nel buffer */
+    for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%2000)-1000)/100.f;
+    int known=1337; lo[known]=50.f;               /* picco netto */
+    dist_build(lo,V);
+    if(pbuf_has_nan(V)){ printf("  FAIL: NaN in g_pbuf con logit sani\n"); fail=1; }
+    if(g_pbuf[known]<=0.f){ printf("  FAIL: il picco ha prob 0\n"); fail=1; }
+    if(!fail) printf("  logit sani: distribuzione valida, picco vivo            ok\n");
+
+    /* (b) NaN/+Inf iniettato in varie posizioni; il finito-argmax deve vincere */
+    float bad[3]; bad[0]=NAN; bad[1]=INFINITY; bad[2]=-INFINITY;
+    const char *bn[3]={"NaN","+Inf","-Inf"};
+    for(int b=0;b<3;b++){
+        for(int pos=0;pos<3;pos++){               /* lo[0], meta', ultimo */
+            for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%400)-200)/100.f;
+            int amax=777; lo[amax]=9.0f;           /* massimo FINITO atteso */
+            int at = pos==0?0 : pos==1?V/2 : V-1;
+            if(at==amax) amax=amax+1;              /* non sovrapporre */
+            lo[amax]=9.0f;
+            lo[at]=bad[b];                         /* veleno */
+            dist_build(lo,V);
+            if(pbuf_has_nan(V)){ printf("  FAIL: NaN sopravvive (%s @ %d)\n",bn[b],at); fail=1; continue; }
+            /* con -Inf il fallback puo' non scattare (max finito resta), ma il buffer
+             * deve restare valido e sommare ~1: con +Inf/NaN scatta il delta su amax */
+            int picked=-1; float pv=-1;
+            for(int i=0;i<V;i++) if(g_pbuf[i]>pv){pv=g_pbuf[i];picked=i;}
+            if(b<2 && picked!=amax){                /* NaN e +Inf: delta esatto su amax */
+                printf("  FAIL: %s @ %d -> picked %d, atteso argmax finito %d\n",bn[b],at,picked,amax); fail=1;
+            }
+        }
+    }
+    if(!fail) printf("  NaN/+Inf iniettato: argmax dei finiti vince, mai 0/NaN   ok\n");
+
+    /* (c) caso estremo: TUTTI non finiti -> non deve crashare, buffer valido */
+    for(int i=0;i<V;i++) lo[i]=NAN;
+    dist_build(lo,V);
+    if(pbuf_has_nan(V)){ printf("  FAIL: tutti-NaN lascia NaN nel buffer\n"); fail=1; }
+    else printf("  tutti non-finiti: nessun crash, buffer valido               ok\n");
+
+    printf(fail?"test_sample_nan: FAIL\n":"test_sample_nan: ok\n");
+    return fail;
+}