From 4cebfabaa686be2fa46212a5680a6562e56ce77f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 11 Jun 2026 22:21:46 -0700
Subject: [PATCH] fix(spark): whole-token device-name match so GB10 !=
 GB100/GB10X

A loose substring match ("GB10" in name) misdetected a discrete Grace+Blackwell
datacenter GPU (e.g. nvidia-smi name containing "GB100") as a unified-memory DGX
Spark, applying the UMA tuning (pin_memory off, vLLM disabled, allocator capped to
0.80) and regressing that hardware. Match each device-name token with non-alphanumeric
boundaries instead. Found by a platform x device-name gating simulation; the real N1X
(JMJWOA-Generic-GPU) still detects, GB100/B100/GB200/GH200/B200 now correctly reject.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 studio/backend/core/training/worker.py |  8 ++++++--
 unsloth/kernels/flex_attention.py      |  7 ++++++-
 unsloth/models/_utils.py               | 14 ++++++++++++--
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 0d1e51343d..f52dfc48f1 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -749,8 +749,10 @@ def _nvidia_classify_spark_unified_memory(props: Any) -> tuple[str, bool]:
     if getattr(props, "is_integrated", 0):
         return "is_integrated", True
     name_upper = (getattr(props, "name", "") or "").upper()
+    import re
     for token in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK"):
-        if token in name_upper:
+        # Whole-token match so "GB10" does not match a discrete "GB100"/"GB10X".
+        if re.search(r"(?<![A-Z0-9])" + re.escape(token) + r"(?![A-Z0-9])", name_upper):
             return token, True
     return "", False
 
@@ -2455,8 +2457,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
                         timeout = 5,
                     )
                     _names_u = (_smi.stdout or "").upper()
+                    import re as _re
                     _spark_smi = any(
-                        t in _names_u for t in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK")
+                        _re.search(r"(?<![A-Z0-9])" + _re.escape(t) + r"(?![A-Z0-9])", _names_u)
+                        for t in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK")
                     )
                 if _spark_smi and os.environ.get("UNSLOTH_NO_EXPANDABLE_SEGMENTS") != "1":
                     _conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index 8cd1e82141..36bc3a84e8 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -49,7 +49,12 @@ def _flex_is_dgx_spark():
             timeout = 5,
         )
         names = (out.stdout or "").upper()
-        return any(t in names for t in ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110"))
+        # Whole-token match so "GB10" does not match a discrete "GB100"/"GB10X".
+        import re
+        return any(
+            re.search(r"(?<![A-Z0-9])" + re.escape(t) + r"(?![A-Z0-9])", names)
+            for t in ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110")
+        )
     except Exception:
         return False
 
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 7cbe5e1aba..8df76767a8 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -998,6 +998,16 @@ from transformers.modeling_utils import logger as transformers_logger
 _DGX_SPARK_DEVICE_TOKENS = ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110")
 
 
+def _name_has_spark_token(names_upper):
+    # Whole-token match so "GB10" does NOT match "GB100"/"GB10X" -- a discrete
+    # Grace+Blackwell datacenter GPU must not be misread as a unified-memory Spark.
+    import re
+    return any(
+        re.search(r"(?<![A-Z0-9])" + re.escape(tok) + r"(?![A-Z0-9])", names_upper)
+        for tok in _DGX_SPARK_DEVICE_TOKENS
+    )
+
+
 @functools.lru_cache(maxsize = None)
 def is_dgx_spark():
     """True only on DGX Spark / N1X Spark-class machines (gate: aarch64 + NVIDIA
@@ -1017,7 +1027,7 @@ def is_dgx_spark():
         names = " ".join(
             str(torch.cuda.get_device_name(i)).upper() for i in range(torch.cuda.device_count())
         )
-        return any(token in names for token in _DGX_SPARK_DEVICE_TOKENS)
+        return _name_has_spark_token(names)
     except Exception:
         return False
 
@@ -1046,7 +1056,7 @@ def _is_dgx_spark_no_cuda_init():
             timeout = 5,
         )
         names = (out.stdout or "").upper()
-        return any(token in names for token in _DGX_SPARK_DEVICE_TOKENS)
+        return _name_has_spark_token(names)
     except Exception:
         return False