Fix sampling returning a constant wrong token for vocabularies larger than 65,536 (#384)

* Fix multinomial sampling for vocabularies larger than 65,536 tokens

The Bool mask pipeline in selectNextTokenUsingSampling (cumsum < rnd
combined with arithmetic ops feeding argmin) truncates at 2^16 elements,
so any model with a vocabulary above 65,536 tokens sampled a constant
wrong token id on every step, and the same op sequence can crash in
libBNNS (issue #365). Replace it with an inverse-CDF binary search on
the CPU: read back the cumulative probabilities once per token and find
the first index reaching the drawn value, which is correct for any
vocabulary size.

* Simplify comment

---------

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

Author: shoemoney

Sources/Generation/Decoders.swift

@@ -20,36 +20,44 @@ func selectNextTokenUsingGreedyDecoding(from scores: MLTensor) -> MLTensor {
 /// - Parameter scores: Processed logits tensor [batch_size, vocab_size]
 /// - Returns: Sampled token ID tensor [batch_size, 1]
 @available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, watchOS 11.0, *)
-func selectNextTokenUsingSampling(from scores: MLTensor) -> MLTensor {
+func selectNextTokenUsingSampling(from scores: MLTensor) async -> MLTensor {
     // Convert logits to probabilities
     let probs = scores.softmax(alongAxis: -1)
 
-    // Multinomial sampling using cumulative sum method:
-    // 1. Generate random number in [0, 1)
-    // 2. Compute cumulative sum of probabilities
-    // 3. Find first index where cumsum >= random_number
+    // Multinomial sampling via inverse CDF, searched on the CPU.
     //
-    // This is equivalent to torch.multinomial() but using available MLTensor ops
-
-    let batchSize = scores.shape[0]
-    let rndTensor = MLTensor(randomUniform: [batchSize, 1], in: 0..<1, scalarType: Float.self)
+    // CPU inverse-CDF search, replacing optimized tensor argmin path that breaks for vocabs > 2^16 (#365).
     let cumulativeProbs = probs.cumulativeSum(alongAxis: -1)
-
-    // Ensure random tensor matches the type of cumulativeProbs
-    let rnd = cumulativeProbs.scalarType == Float.self ? rndTensor : rndTensor.cast(to: cumulativeProbs.scalarType)
-
-    // Create mask where cumsum >= rnd (these are candidates)
-    // We want the FIRST position where this is true
-    // Strategy: Set all positions where cumsum < rnd to a large value (1000.0)
-    // Set all positions where cumsum >= rnd to their index value
-    // Then argmin will give us the first qualifying index
-
-    let mask = cumulativeProbs .< rnd
-    let penalized = mask * 1000.0 // Large value for positions to skip
-    let indexed = penalized + cumulativeProbs // Positions >= rnd will have small values
-
-    let sampledIndex = indexed.argmin(alongAxis: -1).reshaped(to: [1, 1])
-    // Ensure indices are Int32 for concatenation with input tokens
-    return sampledIndex.scalarType == Int32.self ? sampledIndex : sampledIndex.cast(to: Int32.self)
+    let floatCumulativeProbs = cumulativeProbs.scalarType == Float.self ? cumulativeProbs : cumulativeProbs.cast(to: Float.self)
+    let cdf = await floatCumulativeProbs.shapedArray(of: Float.self).scalars
+
+    let vocabSize = scores.shape.last ?? 1
+    let batchSize = max(cdf.count / vocabSize, 1)
+
+    var sampledTokens = [Int32]()
+    sampledTokens.reserveCapacity(batchSize)
+    for batch in 0..<batchSize {
+        let row = cdf[(batch * vocabSize)..<((batch + 1) * vocabSize)]
+        // Scale the draw by the total mass to be robust to floating-point
+        // rounding in the last CDF entry.
+        let rnd = Float.random(in: 0..<1) * (row.last ?? 1)
+
+        // Binary search for the first index whose cumulative probability
+        // reaches the drawn value.
+        var low = row.startIndex
+        var high = row.endIndex - 1
+        while low < high {
+            let mid = (low + high) / 2
+            if row[mid] < rnd {
+                low = mid + 1
+            } else {
+                high = mid
+            }
+        }
+        sampledTokens.append(Int32(low - row.startIndex))
+    }
+
+    // Int32 indices, shaped [batch_size, 1] for concatenation with input tokens
+    return MLTensor(shape: [batchSize, 1], scalars: sampledTokens)
 }
 #endif // canImport(CoreML)

Sources/Generation/Generation.swift

@@ -90,7 +90,7 @@ extension Generation {
                 case .greedy:
                     selectNextTokenUsingGreedyDecoding(from: processedScores)
                 case .sample:
-                    selectNextTokenUsingSampling(from: processedScores)
+                    await selectNextTokenUsingSampling(from: processedScores)
                 default:
                     fatalError("Generation mode \(config.generationMode) not implemented yet")
                 }

Tests/GenerationTests/GenerationIntegrationTests.swift

@@ -340,6 +340,22 @@ final class GenerationIntegrationTests: XCTestCase {
         // Model should be called fewer times due to early stopping
         XCTAssertLessThan(model.callCount, 10, "Model should be called fewer times due to EOS")
     }
+
+    func testSamplingWithLargeVocabulary() async throws {
+        // Regression test for #365: Bool-mask tensor ops truncate at 2^16 elements,
+        // so sampling from a vocabulary larger than 65,536 tokens (e.g. Qwen's
+        // 151,936) deterministically returned token 65536 regardless of the scores.
+        let vocabSize = 151_936
+        var logits = [Float](repeating: 0.0, count: vocabSize)
+        logits[vocabSize - 1] = 100.0 // Effectively one-hot after softmax
+
+        let scores = MLTensor(shape: [1, 1, vocabSize], scalars: logits)
+        let sampled = await selectNextTokenUsingSampling(from: scores)
+
+        XCTAssertEqual(sampled.shape, [1, 1], "Sampled token should have shape [1, 1]")
+        let tokenId = await sampled.shapedArray(of: Int32.self).scalars.first
+        XCTAssertEqual(tokenId, Int32(vocabSize - 1), "Sampling must select the only token with probability mass, never token 65536")
+    }
 }
 
 // MARK: - Test Helper