From 7a8edf77eeebb0a11a3a1d9eefd00d64eb9aece5 Mon Sep 17 00:00:00 2001
From: MarkBase Admin <admin@markbase.local>
Date: Sun, 5 Jul 2026 22:41:49 +0800
Subject: [PATCH] v2: remove remaining logit scaling hacks from batch/optimized
 paths

---
 Sources/MarkBase/BatchGeneration.swift     |  6 ------
 Sources/MarkBase/BatchGenerationTrue.swift | 20 --------------------
 Sources/MarkBase/ModelOptimized.swift      |  6 ------
 3 files changed, 32 deletions(-)

diff --git a/Sources/MarkBase/BatchGeneration.swift b/Sources/MarkBase/BatchGeneration.swift
index c960eb1..260e0f5 100644
--- a/Sources/MarkBase/BatchGeneration.swift
+++ b/Sources/MarkBase/BatchGeneration.swift
@@ -161,12 +161,6 @@ extension E4BModel {
             cmdBuf: cmdBuf
         )
         
-        // Logits scaling
-        if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize {
-            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
-            try scaleBufferOptimized(logitsBuffer, scale: logitsScale, count: vocabSize, cmdBuf: cmdBuf)
-        }
-        
         // Softcapping
         if let cap = finalLogitSoftcapping {
             try applyLogitSoftcappingOptimized(
diff --git a/Sources/MarkBase/BatchGenerationTrue.swift b/Sources/MarkBase/BatchGenerationTrue.swift
index 0699b7a..0cd016c 100644
--- a/Sources/MarkBase/BatchGenerationTrue.swift
+++ b/Sources/MarkBase/BatchGenerationTrue.swift
@@ -160,26 +160,6 @@ embedCmdBuf.waitUntilCompleted()
         encLM.dispatchThreads(gridLM, threadsPerThreadgroup: tgLM)
         encLM.endEncoding()
         
-        // Logits scaling and softcapping (batch)
-        if embedWeight.groupSize == 32 {
-            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
-            // Use eltwise_scale for batch scaling
-            let pso = try engine.pipeline(named: "eltwise_scale")
-            let enc = layerCmdBuf.makeComputeCommandEncoder()!
-            enc.setComputePipelineState(pso)
-            
-            enc.setBuffer(context.batchOutputBuffer, offset: 0, index: 0)
-            var ls = logitsScale
-            enc.setBytes(&ls, length: 4, index: 1)
-            var total = UInt32(batchSize * vocabSize)
-            enc.setBytes(&total, length: 4, index: 2)
-            
-            let tg = MTLSize(width: 256, height: 1, depth: 1)
-            let grid = MTLSize(width: batchSize * vocabSize, height: 1, depth: 1)
-            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
-            enc.endEncoding()
-        }
-        
         // Softcapping (skip if kernel not found)
         if let cap = finalLogitSoftcapping {
             // Try to use tanh_scale kernel
diff --git a/Sources/MarkBase/ModelOptimized.swift b/Sources/MarkBase/ModelOptimized.swift
index 3ce4239..a8843e7 100644
--- a/Sources/MarkBase/ModelOptimized.swift
+++ b/Sources/MarkBase/ModelOptimized.swift
@@ -110,12 +110,6 @@ extension E4BModel {
         try quantizedMatmulOptimized(input: lmInput, weights: embedWeight, 
                                      output: logitsBuffer, cmdBuf: cmdBuf3)
         
-        // Logits scaling (if needed)
-        if embedWeight.groupSize == 32 && embedWeight.inDim == hiddenSize {
-            let logitsScale = Float(30.0 / 116.23 / sqrt(Float(hiddenSize)))
-            try scaleBufferOptimized(logitsBuffer, scale: logitsScale, count: vocabSize, cmdBuf: cmdBuf3)
-        }
-        
         // Logit softcapping
         if let cap = finalLogitSoftcapping {
             try applyLogitSoftcappingOptimized(buffer: logitsBuffer, cap: cap,