From 8acc47129b442c6452a4268a7ef140a036563902 Mon Sep 17 00:00:00 2001
From: project6-dev <dev@project6.ai>
Date: Thu, 13 Aug 2026 02:12:16 +0000
Subject: [PATCH] fix(precision): guard all corex .so outputs with nan_to_num +
 reduce max-model-len
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

MoE kernels:
- topk_softmax: add .contiguous() + nan_to_num + re-normalize weights
- direct_routed: nan_to_num on w2_reduce output
- exact_reduce: nan_to_num on serial_float output

GDN kernels:
- packed_decode: nan_to_num on core_out

BI-V100 CUB may produce non-finite values in fp16 softmax/reduce.
These guards prevent garbage propagation without disabling the kernels.

max-model-len: 256000 → 131072 (4x32GB BI-V100 OOM prevention)
Dockerfile: unchanged (no force push needed)
---
 computility-run.yaml       |  2 +-
 qwen3_6_scripts/qwen3_5.py | 15 ++++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/computility-run.yaml b/computility-run.yaml
index 9fd7ac4d..dcb8603d 100644
--- a/computility-run.yaml
+++ b/computility-run.yaml
@@ -8,7 +8,7 @@ command:
     - --served-model-name
     - llm
     - --max-model-len
-    - '256000'
+    - '131072'
     - --gpu-memory-utilization
     - '0.95'
     - --trust-remote-code
diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py
index aee361aa..d08d945d 100644
--- a/qwen3_6_scripts/qwen3_5.py
+++ b/qwen3_6_scripts/qwen3_5.py
@@ -1187,6 +1187,8 @@ class GatedDeltaNet(nn.Module):
                     core_out = _corex_gdn_packed_decode.packed_decode(
                         temporal_state, packed_mixed_qkv, b_all, a_all,
                         self.A_log, self.dt_bias)
+                    core_out = torch.nan_to_num(
+                        core_out, nan=0.0, posinf=0.0, neginf=0.0)
             else:
                 q, k, v = torch.split(
                     mixed_qkv_conv,
@@ -1611,9 +1613,13 @@ class Qwen3_5MoeSparseBlock(nn.Module):
         # Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh
         if _USE_COREX_MOE_TOPK_SOFTMAX:
             topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax(
-                router_logits.float(), self.top_k, True)
+                router_logits.float().contiguous(), self.top_k, True)
             topk_ids = topk_ids.to(torch.int64)
-            topk_weights = topk_weights.to(hidden_states.dtype)
+            # BI-V100 CUB softmax may produce non-finite → clamp before cast
+            topk_weights = torch.nan_to_num(
+                topk_weights, nan=0.0, posinf=1.0, neginf=0.0)
+            denom = topk_weights.sum(dim=-1, keepdim=True).clamp(min=1e-6)
+            topk_weights = (topk_weights / denom).to(hidden_states.dtype)
         else:
             topk_logits, topk_ids = torch.topk(
                 router_logits.float(), self.top_k, dim=-1)     # (T, top_k)
@@ -1651,8 +1657,10 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                 gate_up = _corex_moe_direct_routed.w13(
                     hidden_states, w13, eids)
                 act = self.act_fn(gate_up)
-                return _corex_moe_direct_routed.w2_reduce(
+                out = _corex_moe_direct_routed.w2_reduce(
                     act, w2, eids, ws)
+                return torch.nan_to_num(
+                    out, nan=0.0, posinf=0.0, neginf=0.0)
 
             use_corex_gather = (
                 _USE_COREX_MOE_WEIGHT_GATHER
@@ -1697,6 +1705,7 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                     and ws.dtype == torch.float16
                     and expert_out.shape[0] == 8):
                 out = _corex_moe_exact_reduce.serial_float(expert_out, ws)
+                out = torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0)
             else:
                 out = (expert_out * ws.unsqueeze(-1)).sum(
                     0, keepdim=True).to(hidden_states.dtype)   # (1, H)