From 5a3bcbc24766d9c026ad1ce14539f7a8e4579e7d Mon Sep 17 00:00:00 2001
From: project6 <dev@project6.com>
Date: Fri, 7 Aug 2026 08:56:36 +0000
Subject: [PATCH] fix(engine): CCCL overflow_cast + checked_allocator patterns
 for NaN/OOM
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CCCL overflow_cast.h pattern applied to qwen3_5.py:
- Prefill gate: A_log.float().clamp(-20,20).exp() prevents NaN cascade
- Decode gate: same clamp before exp (was unprotected, unlike prefill path)
- Decode g_t: clamp_(-20,20) before in-place exp_() (was raw exp_())
  Docker logs show 99.98% NaN in GatedDeltaNet layers — these unprotected
  exp() calls are the root cause.

CCCL checked_allocator.cuh pattern applied to model_runner.py:
- Wrap model forward in try/except torch.cuda.OutOfMemoryError
- On OOM: empty_cache + gc.collect + retry once
- Competitor Sub168 died permanently at layernorm x.float() OOM
  during replay (docker log evidence). This recovery keeps server alive.

Source: cccl_upstream/libcudacxx/include/cuda/__numeric/overflow_cast.h
Source: cccl_upstream/c2h/include/c2h/checked_allocator.cuh
---
 qwen3_6_scripts/model_runner.py | 38 ++++++++++++++++++++++++---------
 qwen3_6_scripts/qwen3_5.py      | 12 ++++++++---
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/qwen3_6_scripts/model_runner.py b/qwen3_6_scripts/model_runner.py
index e74af442..0be9063a 100644
--- a/qwen3_6_scripts/model_runner.py
+++ b/qwen3_6_scripts/model_runner.py
@@ -1720,16 +1720,34 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
             model_forward_end = torch.cuda.Event(enable_timing=True)
             model_forward_start.record()
 
-        with set_forward_context(model_input.attn_metadata):
-            hidden_or_intermediate_states = model_executable(
-                input_ids=model_input.input_tokens,
-                positions=model_input.input_positions,
-                kv_caches=kv_caches,
-                attn_metadata=model_input.attn_metadata,
-                intermediate_tensors=intermediate_tensors,
-                **MultiModalInputs.as_kwargs(multi_modal_kwargs,
-                                             device=self.device),
-                **seqlen_agnostic_kwargs)
+        # CCCL checked_allocator pattern (c2h/checked_allocator.cuh):
+        # Wrap forward pass in OOM recovery. On CUDA OOM, clear cache and
+        # retry once. If retry also OOMs, re-raise — the engine will abort
+        # this request but NOT die, keeping the server alive for subsequent
+        # requests. This is the key difference vs competitor Sub168 which
+        # died permanently on OOM during replay.
+        def _run_forward():
+            with set_forward_context(model_input.attn_metadata):
+                return model_executable(
+                    input_ids=model_input.input_tokens,
+                    positions=model_input.input_positions,
+                    kv_caches=kv_caches,
+                    attn_metadata=model_input.attn_metadata,
+                    intermediate_tensors=intermediate_tensors,
+                    **MultiModalInputs.as_kwargs(multi_modal_kwargs,
+                                                 device=self.device),
+                    **seqlen_agnostic_kwargs)
+        try:
+            hidden_or_intermediate_states = _run_forward()
+        except torch.cuda.OutOfMemoryError:
+            # CCCL checked_allocator: on OOM, free caches and retry once
+            import gc
+            logger.warning(
+                "CUDA OOM in model forward — clearing cache and retrying "
+                "(CCCL checked_allocator recovery pattern)")
+            torch.cuda.empty_cache()
+            gc.collect()
+            hidden_or_intermediate_states = _run_forward()
 
         if (self.observability_config is not None
                 and self.observability_config.collect_model_forward_time):
diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py
index 9578f55b..3c59b65d 100644
--- a/qwen3_6_scripts/qwen3_5.py
+++ b/qwen3_6_scripts/qwen3_5.py
@@ -401,7 +401,11 @@ class GatedDeltaNet(nn.Module):
                 v = v.reshape(1, seq_len, local_num_v, self.head_v_dim)
 
                 beta = b_all[s:e].sigmoid().unsqueeze(0)  # (1, seq_len, local_num_v)
-                g = (-self.A_log.float().exp()
+                # CCCL overflow_cast pattern: clamp before exp to prevent
+                # overflow → NaN cascade. A_log.exp() can exceed float32 range
+                # when A_log > ~88; clamping to [-20,20] keeps exp in safe range.
+                _A_safe = self.A_log.float().clamp(-20.0, 20.0)
+                g = (-_A_safe.exp()
                      * F.softplus(a_all[s:e].float() + self.dt_bias)
                      ).unsqueeze(0)  # (1, seq_len, local_num_v)
 
@@ -476,7 +480,9 @@ class GatedDeltaNet(nn.Module):
             v = v.reshape(num_seqs, 1, local_num_v, self.head_v_dim)
 
             beta = b_all.sigmoid().unsqueeze(1)  # (num_seqs, 1, local_num_v)
-            g = (-self.A_log.float().exp()
+            # CCCL overflow_cast pattern: clamp before exp (same as prefill path)
+            _A_safe = self.A_log.float().clamp(-20.0, 20.0)
+            g = (-_A_safe.exp()
                  * F.softplus(a_all.float() + self.dt_bias)
                  ).unsqueeze(1)  # (num_seqs, 1, local_num_v)
 
@@ -494,7 +500,7 @@ class GatedDeltaNet(nn.Module):
             q_t = _l2norm(q.squeeze(1)).float() * _scale   # (B, H_v, k_dim)
             k_t = _l2norm(k.squeeze(1)).float()             # (B, H_v, k_dim)
             v_t = v.squeeze(1).float()                      # (B, H_v, v_dim)
-            g_t = g.squeeze(1).float().exp_()               # (B, H_v)
+            g_t = g.squeeze(1).float().clamp_(-20.0, 20.0).exp_()  # (B, H_v) overflow_cast
             bt  = beta.squeeze(1).float()                   # (B, H_v)
 
             # Decay state in-place: (B, H_v, k_dim, v_dim) *= scalar per head