From d4ef3c75d2476310bcfcffb63bbfd1d44592ea4a Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Wed, 24 Jun 2026 01:32:00 +0200 Subject: [PATCH] fix(merge): serialize GPU compression SVD to stop async-race abort MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The compression SVD (torch.svd_lowrank) can race asynchronously on some stacks (Blackwell/cu130), corrupting CUDA state and aborting a later kernel — a C++ abort the Python try/except can't catch. This is what crashed Krea 2 merges (two different-form LoRAs route through dense merge + compression, the only path that runs GPU compression SVD here). Make the per-group cuda.synchronize() permanent, but gate it tightly on (GPU compression active) so it only fires on compression merges and fast no-compression merges keep their fully-pipelined speed. Previously this sync only ran under the PROFILE_MERGE sentinel. No version bump. Co-Authored-By: Claude Opus 4.8 --- lora_optimizer.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/lora_optimizer.py b/lora_optimizer.py index 1f81285..efd7896 100644 --- a/lora_optimizer.py +++ b/lora_optimizer.py @@ -7505,6 +7505,14 @@ class LoRAOptimizer(_LoRAMergeBase): _merge_prof_lock = threading.Lock() if _profiling_on else None _prof_cuda = (use_gpu and compute_device is not None and getattr(compute_device, "type", None) == "cuda") + # A GPU compression SVD (torch.svd_lowrank) can race asynchronously on + # some stacks (Blackwell/cu130), corrupting state and aborting a later + # kernel — a C++ abort the Python try/except can't catch. When the + # compression SVD runs on GPU, serialize each group so the SVD finishes + # before the next group launches. Gated tightly (compression ON + GPU + # SVD) so fast no-compression merges keep their pipelined speed. + _compress_sync = (_prof_cuda and compress_rank > 0 + and resolved_svd_device is not None) def _prof_t(): """Synced timestamp — without the sync, async CUDA kernels would @@ -7963,7 +7971,7 @@ class LoRAOptimizer(_LoRAMergeBase): logging.info(f"[merge-debug] processing {idx + 1}/{len(group_items)}: " f"{_dbg_tk} (n_loras={prefix_stats.get(label_prefix, {}).get('n_loras', '?')})") result = _merge_one_group(label_prefix, target_group) - if _merge_prof is not None and _prof_cuda: + if _compress_sync or (_merge_prof is not None and _prof_cuda): torch.cuda.synchronize(compute_device) if _sl_key is not None: _sl_store(_sl_key, result)