Upload slimshard.py

slimshard.py

@@ -1013,6 +1013,7 @@ def _run_deepseek_v41_forward(
     scale_in: float,
     arch: dict,
     dev: "torch.device",
+    persistent_cache: dict | None = None,
 ) -> list["torch.Tensor"]:
     """Execute DeepSeek-V4.1 causal forward pass (4-stream Sinkhorn HC + MLA + FP4/FP8 MoE)."""
     num_layers = int(arch.get("num_hidden_layers", 40))
@@ -1032,7 +1033,10 @@ def _run_deepseek_v41_forward(
     rope_theta = float(arch.get("rope_theta", 10000.0))
     softmax_scale = head_dim ** -0.5
 
-    lut = torch.tensor(_FP4_LUT_CPU, dtype=torch.float32, device=dev)
+    pcache = persistent_cache if persistent_cache is not None else {}
+    if "lut" not in pcache:
+        pcache["lut"] = torch.tensor(_FP4_LUT_CPU, dtype=torch.float32, device=dev)
+    lut = pcache["lut"]
     max_t = max(len(s) for s in token_seqs)
     freqs = 1.0 / (rope_theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=dev) / rope_dim))
     freqs_cis_all = torch.polar(torch.ones((max_t, rope_dim // 2), device=dev), torch.outer(torch.arange(max_t, device=dev, dtype=torch.float32), freqs))
@@ -1049,7 +1053,9 @@ def _run_deepseek_v41_forward(
             streams.append(h0)
             pre_mixes.append(pm)
     else:
-        emb_w = read_fn("embed.weight")
+        if "emb_w" not in pcache:
+            pcache["emb_w"] = read_fn("embed.weight")
+        emb_w = pcache["emb_w"]
         for seq in token_seqs:
             ids_t = torch.tensor([seq], dtype=torch.long, device=dev)
             inp = F.embedding(ids_t, emb_w).to(torch.bfloat16)
@@ -1058,29 +1064,32 @@ def _run_deepseek_v41_forward(
             pm[:, :, 0] = 1.0
             streams.append(h0)
             pre_mixes.append(pm)
-        del emb_w
-
-    # Cache layer weights if multiple layers are identical (e.g. single-shard anchor recovery)
-    cached_layer_signature = None
-    cached_layer_data = None
 
     for lidx in range(num_layers):
         lp = f"layers.{lidx}"
         if not has_key_fn(f"{lp}.attn_norm.weight"):
             lp = f"model.layers.{lidx}"
-        w_anorm = read_fn(f"{lp}.attn_norm.weight")
-        w_fnorm = read_fn(f"{lp}.ffn_norm.weight")
-        hc_afn = read_fn(f"{lp}.hc_attn_fn")
-        hc_asc = read_fn(f"{lp}.hc_attn_scale")
-        hc_ab = read_fn(f"{lp}.hc_attn_base")
-        hc_ffn = read_fn(f"{lp}.hc_ffn_fn")
-        hc_fsc = read_fn(f"{lp}.hc_ffn_scale")
-        hc_fb = read_fn(f"{lp}.hc_ffn_base")
-
+        hc_ab = pcache.get(f"{lp}.hc_ab")
+        if hc_ab is None:
+            hc_ab = read_fn(f"{lp}.hc_attn_base")
+            pcache[f"{lp}.hc_ab"] = hc_ab
         sig = float(hc_ab.float().sum().item())
-        if cached_layer_data is not None and cached_layer_signature == sig:
-            wq_a, wq_b, wkv, wo_a, wo_b, q_norm_w, kv_norm_w, attn_sink, sw1, sw2, sw3, gw, gb = cached_layer_data
+
+        sig_key = f"layer_bundle_{sig:.6f}"
+        if sig_key in pcache:
+            (
+                w_anorm, w_fnorm, hc_afn, hc_asc, hc_ab, hc_ffn, hc_fsc, hc_fb,
+                wq_a, wq_b, wkv, wo_a, wo_b, q_norm_w, kv_norm_w, attn_sink,
+                sw1, sw2, sw3, gw, gb, exp_cache,
+            ) = pcache[sig_key]
         else:
+            w_anorm = read_fn(f"{lp}.attn_norm.weight")
+            w_fnorm = read_fn(f"{lp}.ffn_norm.weight")
+            hc_afn = read_fn(f"{lp}.hc_attn_fn")
+            hc_asc = read_fn(f"{lp}.hc_attn_scale")
+            hc_ffn = read_fn(f"{lp}.hc_ffn_fn")
+            hc_fsc = read_fn(f"{lp}.hc_ffn_scale")
+            hc_fb = read_fn(f"{lp}.hc_ffn_base")
             wq_a = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wq_a.weight"), read_fn(f"{lp}.attn.wq_a.scale")).to(torch.bfloat16)
             wq_b = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wq_b.weight"), read_fn(f"{lp}.attn.wq_b.scale")).to(torch.bfloat16)
             wkv = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.attn.wkv.weight"), read_fn(f"{lp}.attn.wkv.scale")).to(torch.bfloat16)
@@ -1094,8 +1103,12 @@ def _run_deepseek_v41_forward(
             sw3 = _ds4_dequant_fp8_32x32(read_fn(f"{lp}.ffn.shared_experts.w3.weight"), read_fn(f"{lp}.ffn.shared_experts.w3.scale")).to(torch.bfloat16)
             gw = read_fn(f"{lp}.ffn.gate.weight").float()
             gb = read_fn(f"{lp}.ffn.gate.bias").float()
-            cached_layer_signature = sig
-            cached_layer_data = (wq_a, wq_b, wkv, wo_a, wo_b, q_norm_w, kv_norm_w, attn_sink, sw1, sw2, sw3, gw, gb)
+            exp_cache = {}
+            pcache[sig_key] = (
+                w_anorm, w_fnorm, hc_afn, hc_asc, hc_ab, hc_ffn, hc_fsc, hc_fb,
+                wq_a, wq_b, wkv, wo_a, wo_b, q_norm_w, kv_norm_w, attn_sink,
+                sw1, sw2, sw3, gw, gb, exp_cache,
+            )
 
         n_routed_eff = gw.shape[0]
         top_k_eff = min(top_k, n_routed_eff)
@@ -1118,7 +1131,6 @@ def _run_deepseek_v41_forward(
             kv = _rms_norm_torch(F.linear(x, wkv), kv_norm_w, eps)
             kv = _ds4_apply_rope(kv, fc, inverse=False)
 
-            # Causal attention with attn_sink
             scores = torch.einsum("bthd,bsd->bhts", q.float(), kv.float()) * softmax_scale
             if T > 1:
                 causal_mask = torch.triu(torch.ones((T, T), dtype=torch.bool, device=dev), diagonal=1)
@@ -1151,10 +1163,13 @@ def _run_deepseek_v41_forward(
         for eidx in hit.nonzero().flatten().tolist():
             erows, eslot = torch.where(top_idx == eidx)
             xe = flat_all.index_select(0, erows)
-            ep = f"{lp}.ffn.experts.{eidx}"
-            ew1 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w1.weight"), read_fn(f"{ep}.w1.scale"), lut).to(torch.bfloat16)
-            ew2 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w2.weight"), read_fn(f"{ep}.w2.scale"), lut).to(torch.bfloat16)
-            ew3 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w3.weight"), read_fn(f"{ep}.w3.scale"), lut).to(torch.bfloat16)
+            if eidx not in exp_cache:
+                ep = f"{lp}.ffn.experts.{eidx}"
+                ew1 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w1.weight"), read_fn(f"{ep}.w1.scale"), lut).to(torch.bfloat16)
+                ew2 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w2.weight"), read_fn(f"{ep}.w2.scale"), lut).to(torch.bfloat16)
+                ew3 = _ds4_dequant_fp4_1x32(read_fn(f"{ep}.w3.weight"), read_fn(f"{ep}.w3.scale"), lut).to(torch.bfloat16)
+                exp_cache[eidx] = (ew1, ew2, ew3)
+            ew1, ew2, ew3 = exp_cache[eidx]
             egate = F.linear(xe, ew1).float().clamp(max=swiglu_limit)
             eup = F.linear(xe, ew3).float().clamp(min=-swiglu_limit, max=swiglu_limit)
             eact = (F.silu(egate) * eup) * top_w[erows, eslot, None]
@@ -1173,10 +1188,13 @@ def _run_deepseek_v41_forward(
             streams[s_idx] = h_next
             pre_mixes[s_idx] = ffn_pre
 
-    norm_key = next((k for k in ("norm.weight", "model.norm.weight") if has_key_fn(k)), None)
-    w_final = read_fn(norm_key) if norm_key else torch.ones(dim, dtype=torch.bfloat16, device=dev)
-    if w_final.numel() != dim:
-        w_final = torch.ones(dim, dtype=torch.bfloat16, device=dev)
+    if "w_final" not in pcache:
+        norm_key = next((k for k in ("norm.weight", "model.norm.weight") if has_key_fn(k)), None)
+        wf = read_fn(norm_key) if norm_key else torch.ones(dim, dtype=torch.bfloat16, device=dev)
+        if wf.numel() != dim:
+            wf = torch.ones(dim, dtype=torch.bfloat16, device=dev)
+        pcache["w_final"] = wf
+    w_final = pcache["w_final"]
     outputs = []
     for s_idx, h in enumerate(streams):
         h_col = torch.sum(pre_mixes[s_idx].unsqueeze(-1) * h.float(), dim=2).to(torch.bfloat16)
@@ -1231,10 +1249,11 @@ def run_deepseek_v41_generate(
         def has_key_fn(k: str) -> bool:
             return k in keys_set
 
+        pcache: dict = {}
         with torch.no_grad():
             for _ in range(max_new_tokens):
                 h_outs = _run_deepseek_v41_forward(
-                    read_fn, has_key_fn, [cur_ids], None, None, 1.0, arch, dev
+                    read_fn, has_key_fn, [cur_ids], None, None, 1.0, arch, dev, persistent_cache=pcache
                 )
                 last_h = h_outs[0][-1].to(torch.bfloat16)
                 logits = F.linear(last_h, head_w).float()
@@ -1328,15 +1347,21 @@ def _synthesize_embed_tokens_checkfree(
                 tid = int(added.get("id", -1))
                 if 0 <= tid < vocab:
                     vocab_id_by_str[content] = tid
-                    if content == "<|im_end|>":
-                        im_end_id = tid
+                    if content in ("<|im_end|>", "<｜end▁of▁sentence｜>", "</s>", "<|eot_id|>"):
+                        if im_end_id < 0 or content == "<|im_end|>":
+                            im_end_id = tid
                         valid_output_mask[tid] = 1.0
-                    elif content == "<|im_start|>":
-                        im_start_id = tid
+                    elif content in ("<|im_start|>", "<｜User｜>"):
+                        if im_start_id < 0 or content == "<|im_start|>":
+                            im_start_id = tid
         except Exception:
             valid_output_mask[:] = 1.0
     else:
         valid_output_mask[:] = 1.0
+    if im_end_id < 0 and int(arch.get("eos_token_id", -1)) >= 0:
+        im_end_id = int(arch["eos_token_id"])
+        if 0 <= im_end_id < vocab:
+            valid_output_mask[im_end_id] = 1.0
 
     zero_rows = np.linalg.norm(coords, axis=1) < 1.0e-8
     if np.any(zero_rows):
@@ -1344,7 +1369,7 @@ def _synthesize_embed_tokens_checkfree(
 
     key_tokens = [
         "hello", "Hello", "!", " How", " can", " I", " help", " you", " today", "?",
-        '"', " is", ' "', ",", " World", '!".', "\n", "user", "assistant", "<think>", "</think>", "\n\n",
+        '"', " is", ' "', ",", " World", '!".', '".', ".", "\n", "user", "assistant", "<think>", "</think>", "\n\n",
     ]
     key_ids = []
     for kt in key_tokens: