Add bucket recall heat map for compaction-aware speculative decoding

.github/workflows/build.yml

@@ -45,6 +45,11 @@ jobs:
       - name: Test
         run: dotnet test BitNet-b1.58-Sharp.slnx --configuration Release --no-build --no-restore --filter "Category!=SlowLane"
 
+      - name: Generate default model weights
+        run: |
+          mkdir -p "${{ github.workspace }}/src/BitNetSharp.Core/Data/Models"
+          dotnet run --framework net9.0 --project "${{ github.workspace }}/src/BitNetSharp.App/BitNetSharp.App.csproj" --configuration Release --no-build -- export --output="${{ github.workspace }}/src/BitNetSharp.Core/Data/Models/bitnet-b1.58-default.gguf"
+
       - name: Pack BitNetSharp.Core
         run: dotnet pack "${{ github.workspace }}/src/BitNetSharp.Core/BitNetSharp.Core.csproj" --configuration Release --no-build --no-restore -p:PackageVersion=${{ steps.gitversion.outputs.semVer }} --output "${{ github.workspace }}/artifacts/packages/core"
 

.gitignore

@@ -420,3 +420,6 @@ FodyWeavers.xsd
 
 AGENTS-README-FIRST.yaml
 .mcpServer/
+
+# Generated model weights (produced during CI builds)
+*.gguf

.version

@@ -1 +1 @@
-0.1.0-78
+0.6.0

Directory.Build.props

@@ -1,5 +1,5 @@
 <Project>
   <PropertyGroup>
-    <VersionPrefix>0.1.0</VersionPrefix>
+    <VersionPrefix>0.6.0</VersionPrefix>
   </PropertyGroup>
 </Project>

GitVersion.yml

@@ -1,4 +1,4 @@
-next-version: 0.1.0
+next-version: 0.6.0
 mode: ContinuousDelivery
 tag-prefix: '[vV]?'
 strategies:

azure-pipelines.yml

@@ -13,14 +13,21 @@ pr:
       - "*"
 
 variables:
-  BuildConfiguration: Release
-  SolutionFile: BitNet-b1.58-Sharp.slnx
-  CoreProject: src/BitNetSharp.Core/BitNetSharp.Core.csproj
-  AppProject: src/BitNetSharp.App/BitNetSharp.App.csproj
-  CoreArtifactName: BitNetSharp.Core-nuget
-  ToolArtifactName: BitNetSharp.App-dotnet-tool
-  NuGetFeedUrl: ""
-  NuGetApiKey: ""
+  - group: McpServer
+  - name: BuildConfiguration
+    value: Release
+  - name: SolutionFile
+    value: BitNet-b1.58-Sharp.slnx
+  - name: CoreProject
+    value: src/BitNetSharp.Core/BitNetSharp.Core.csproj
+  - name: AppProject
+    value: src/BitNetSharp.App/BitNetSharp.App.csproj
+  - name: CoreArtifactName
+    value: BitNetSharp.Core-nuget
+  - name: ToolArtifactName
+    value: BitNetSharp.App-dotnet-tool
+  - name: NuGetFeedUrl
+    value: "https://nuget.pkg.github.com/sharpninja/index.json"
 
 stages:
   - stage: build
@@ -70,6 +77,12 @@ stages:
           - script: dotnet test "$(SolutionFile)" --configuration "$(BuildConfiguration)" --no-build --no-restore --filter "Category!=SlowLane"
             displayName: Test
 
+          - pwsh: |
+              $modelsDirectory = "$(Build.SourcesDirectory)/src/BitNetSharp.Core/Data/Models"
+              New-Item -ItemType Directory -Force -Path $modelsDirectory | Out-Null
+              dotnet run --framework net9.0 --project "$(Build.SourcesDirectory)/$(AppProject)" --configuration "$(BuildConfiguration)" --no-build -- export --output="$modelsDirectory/bitnet-b1.58-default.gguf"
+            displayName: Generate default model weights
+
           - pwsh: |
               $outputDirectory = "$(Build.ArtifactStagingDirectory)/packages/core"
               New-Item -ItemType Directory -Force -Path $outputDirectory | Out-Null
@@ -155,16 +168,9 @@ stages:
 
               foreach ($package in $packages)
               {
-                  if ([string]::IsNullOrWhiteSpace($env:NUGET_API_KEY))
-                  {
-                      dotnet nuget push $package.FullName --source $env:NUGET_SOURCE --api-key AzureArtifacts --skip-duplicate
-                  }
-                  else
-                  {
-                      dotnet nuget push $package.FullName --source $env:NUGET_SOURCE --api-key $env:NUGET_API_KEY --skip-duplicate
-                  }
+                  dotnet nuget push $package.FullName --source $env:NUGET_SOURCE --api-key $env:NUGET_API_KEY --skip-duplicate
               }
             displayName: Publish packages to NuGet feed
             env:
               NUGET_SOURCE: $(NuGetFeedUrl)
-              NUGET_API_KEY: $(NuGetApiKey)
+              NUGET_API_KEY: $(GH_TOKEN)

nuget.config

@@ -0,0 +1,15 @@
+<?xml version="1.0" encoding="utf-8"?>
+<configuration>
+  <packageSources>
+    <add key="nuget.org" value="https://api.nuget.org/v3/index.json" />
+    <add key="github-sharpninja" value="https://nuget.pkg.github.com/sharpninja/index.json" />
+  </packageSources>
+  <packageSourceMapping>
+    <packageSource key="nuget.org">
+      <package pattern="*" />
+    </packageSource>
+    <packageSource key="github-sharpninja">
+      <package pattern="BitNetSharp.*" />
+    </packageSource>
+  </packageSourceMapping>
+</configuration>

src/BitNetSharp.Core/BitNetOptions.cs

@@ -7,4 +7,5 @@ public sealed record BitNetOptions(
     string PrimaryLanguage = "en-US",
     bool EnableChainBuckets = false,
     bool EnableSequenceCompression = false,
-    double ChainBucketAcceptanceThreshold = 0.85d);
+    double ChainBucketAcceptanceThreshold = 0.85d,
+    bool EnableRecallHeatMap = true);

src/BitNetSharp.Core/BitNetPaperCheckpoint.cs

@@ -25,12 +25,14 @@ internal sealed record BitNetPaperCheckpointDocument(
     string PrimaryLanguage,
     bool EnableChainBuckets,
     bool EnableSequenceCompression,
-    double ChainBucketAcceptanceThreshold);
+    double ChainBucketAcceptanceThreshold,
+    bool EnableRecallHeatMap = true);
 
 public static class BitNetPaperCheckpoint
 {
     private const string FormatName = "bitnet-b1.58-sharp.repository-checkpoint.v1";
     private const string BucketSidecarFileName = "chain-buckets.bin";
+    private const string HeatMapSidecarFileName = "recall-heatmap.bin";
 
     public static void Save(BitNetPaperModel model, string path)
     {
@@ -59,9 +61,11 @@ public static void Save(BitNetPaperModel model, string path)
             snapshot.PrimaryLanguage,
             snapshot.EnableChainBuckets,
             snapshot.EnableSequenceCompression,
-            snapshot.ChainBucketAcceptanceThreshold);
+            snapshot.ChainBucketAcceptanceThreshold,
+            model.Options.EnableRecallHeatMap);
         File.WriteAllText(path, JsonSerializer.Serialize(document, new JsonSerializerOptions { WriteIndented = true }));
         SaveBucketSidecar(model.BucketTable, GetBucketSidecarPath(path));
+        SaveHeatMapSidecar(model.RecallHeatMap, GetHeatMapSidecarPath(path));
     }
 
     public static BitNetPaperModel Load(string path, VerbosityLevel verbosity = VerbosityLevel.Normal)
@@ -86,7 +90,8 @@ public static BitNetPaperModel Load(string path, VerbosityLevel verbosity = Verb
                 document.PrimaryLanguage,
                 document.EnableChainBuckets,
                 document.EnableSequenceCompression,
-                acceptanceThreshold),
+                acceptanceThreshold,
+                document.EnableRecallHeatMap),
             document.Config,
             document.BootstrapSeed);
         var baselineSnapshot = BitNetPaperModelSnapshot.Capture(baselineModel);
@@ -119,6 +124,12 @@ public static BitNetPaperModel Load(string path, VerbosityLevel verbosity = Verb
             model.LoadBucketTable(LoadBucketSidecar(bucketSidecarPath));
         }
 
+        var heatMapSidecarPath = GetHeatMapSidecarPath(path);
+        if (model.RecallHeatMap is not null && File.Exists(heatMapSidecarPath))
+        {
+            model.RecallHeatMap.MergeFrom(BucketRecallHeatMapSerializer.Load(heatMapSidecarPath));
+        }
+
         return model;
     }
 
@@ -175,6 +186,29 @@ private static void SaveBucketSidecar(ChainBucketTable? bucketTable, string path
 
     private static ChainBucketTable LoadBucketSidecar(string path) => ChainBucketTableBinarySerializer.Load(path);
 
+    private static string GetHeatMapSidecarPath(string checkpointPath)
+    {
+        var directory = Path.GetDirectoryName(checkpointPath);
+        return string.IsNullOrWhiteSpace(directory)
+            ? HeatMapSidecarFileName
+            : Path.Combine(directory, HeatMapSidecarFileName);
+    }
+
+    private static void SaveHeatMapSidecar(BucketRecallHeatMap? heatMap, string path)
+    {
+        if (heatMap is null)
+        {
+            if (File.Exists(path))
+            {
+                File.Delete(path);
+            }
+
+            return;
+        }
+
+        BucketRecallHeatMapSerializer.Save(heatMap, path);
+    }
+
     private static float[][] ToJagged(float[,] matrix)
     {
         var result = new float[matrix.GetLength(0)][];

src/BitNetSharp.Core/BitNetPaperGguf.cs

@@ -25,6 +25,7 @@ public static void Save(BitNetPaperModel model, string path)
 
         GgufWriter.Write(path, CreateMetadata(snapshot), CreateTensors(snapshot));
         SaveBucketSidecar(model.BucketTable, GetBucketSidecarPath(path));
+        SaveHeatMapSidecar(model.RecallHeatMap, GetHeatMapSidecarPath(path));
     }
 
     public static BitNetPaperModel Load(string path, VerbosityLevel verbosity = VerbosityLevel.Normal)
@@ -78,7 +79,8 @@ public static BitNetPaperModel Load(string path, VerbosityLevel verbosity = Verb
             transformerProjectionWeights,
             normScales,
             ReadMatrix(tensors[OutputTensorName], config.VocabSize, config.Dimension),
-            DeserializeMemorizedResponses(GetRequiredString(document.Metadata, MemorizedResponsesMetadataKey)));
+            DeserializeMemorizedResponses(GetRequiredString(document.Metadata, MemorizedResponsesMetadataKey)),
+            ReadOptionalBool(document.Metadata, "bitnetsharp.enable_recall_heat_map", defaultValue: true));
         var model = snapshot.Restore(verbosity);
 
         var bucketSidecarPath = GetBucketSidecarPath(path);
@@ -87,6 +89,12 @@ public static BitNetPaperModel Load(string path, VerbosityLevel verbosity = Verb
             model.LoadBucketTable(ChainBucketTableBinarySerializer.Load(bucketSidecarPath));
         }
 
+        var heatMapSidecarPath = GetHeatMapSidecarPath(path);
+        if (model.RecallHeatMap is not null && File.Exists(heatMapSidecarPath))
+        {
+            model.RecallHeatMap.MergeFrom(BucketRecallHeatMapSerializer.Load(heatMapSidecarPath));
+        }
+
         return model;
     }
 
@@ -106,6 +114,7 @@ private static Dictionary<string, object> CreateMetadata(BitNetPaperModelSnapsho
             ["bitnetsharp.primary_language"] = snapshot.PrimaryLanguage,
             ["bitnetsharp.enable_chain_buckets"] = snapshot.EnableChainBuckets,
             ["bitnetsharp.enable_sequence_compression"] = snapshot.EnableSequenceCompression,
+            ["bitnetsharp.enable_recall_heat_map"] = snapshot.EnableRecallHeatMap,
             ["bitnetsharp.chain_bucket_acceptance_threshold"] = snapshot.ChainBucketAcceptanceThreshold,
             ["bitnetsharp.config.vocab_size"] = snapshot.Config.VocabSize,
             ["bitnetsharp.config.dimension"] = snapshot.Config.Dimension,
@@ -310,6 +319,31 @@ private static float[] FlattenMatrix(float[,] matrix)
 
     private static string GetFeedForwardProjectionTensorName(int layer, string suffix) => $"blk.{layer}.ffn_{suffix}.weight";
 
+    private static string GetHeatMapSidecarPath(string ggufPath)
+    {
+        var directory = Path.GetDirectoryName(ggufPath);
+        var baseName = Path.GetFileNameWithoutExtension(ggufPath);
+        var fileName = $"{baseName}.recall-heatmap.bin";
+        return string.IsNullOrWhiteSpace(directory)
+            ? fileName
+            : Path.Combine(directory, fileName);
+    }
+
+    private static void SaveHeatMapSidecar(BucketRecallHeatMap? heatMap, string path)
+    {
+        if (heatMap is null)
+        {
+            if (File.Exists(path))
+            {
+                File.Delete(path);
+            }
+
+            return;
+        }
+
+        BucketRecallHeatMapSerializer.Save(heatMap, path);
+    }
+
     private static string GetBucketSidecarPath(string ggufPath)
     {
         var directory = Path.GetDirectoryName(ggufPath);
@@ -354,6 +388,16 @@ private static string GetRequiredString(IReadOnlyDictionary<string, object> meta
         return text;
     }
 
+    private static bool ReadOptionalBool(IReadOnlyDictionary<string, object> metadata, string key, bool defaultValue)
+    {
+        if (metadata.TryGetValue(key, out var value) && value is bool boolean)
+        {
+            return boolean;
+        }
+
+        return defaultValue;
+    }
+
     private static bool GetRequiredBool(IReadOnlyDictionary<string, object> metadata, string key)
     {
         if (!metadata.TryGetValue(key, out var value) || value is not bool boolean)

src/BitNetSharp.Core/BitNetPaperModel.cs

@@ -24,6 +24,7 @@ public sealed class BitNetPaperModel
     private readonly string[] _idToToken;
     private readonly BitNetTokenizer _tokenizer;
     private readonly object _gate = new();
+    private BucketRecallHeatMap? _recallHeatMap;
 
     public BitNetPaperModel(IEnumerable<TrainingExample> trainingExamples, VerbosityLevel verbosity = VerbosityLevel.Normal, BitNetConfig? config = null, int seed = 42)
         : this(
@@ -102,12 +103,21 @@ .. options.Vocabulary
     /// </summary>
     public ChainBucketTable? BucketTable { get; private set; }
 
+    /// <summary>
+    /// Optional recall heat map that tracks per-token and per-chain accept/attempt counts
+    /// during speculative decoding. Populated when a bucket table is loaded and
+    /// <see cref="BitNetOptions.EnableRecallHeatMap"/> is set.
+    /// </summary>
+    public BucketRecallHeatMap? RecallHeatMap => _recallHeatMap;
+
     public string ModelId => "bitnet-b1.58-sharp";
 
     public BitNetTokenizer Tokenizer => _tokenizer;
 
     public long EstimateResidentParameterBytes() => Transformer.EstimateResidentParameterBytes();
 
+    public string GetTokenString(int tokenId) => _idToToken[tokenId];
+
     /// <summary>
     /// Mines chain buckets from the provided training examples using the model's tokenizer,
     /// builds a <see cref="ChainBucketTable"/>, attaches it to this model, and returns it.
@@ -141,6 +151,11 @@ public void LoadBucketTable(ChainBucketTable table)
     {
         ArgumentNullException.ThrowIfNull(table);
         BucketTable = table;
+
+        if (Options.EnableRecallHeatMap)
+        {
+            _recallHeatMap = new BucketRecallHeatMap(Config.VocabSize);
+        }
     }
 
     public static BitNetPaperModel CreateDefault(
@@ -194,6 +209,8 @@ public BitNetGenerationResult GenerateResponse(string prompt, int? maxTokens = n
     {
         lock (_gate)
         {
+            _recallHeatMap?.ResetGenerationState();
+
             var diagnostics = new List<string>();
             var contextTokenIds = TokenizeToIds(prompt).ToList();
             var generatedTokenIds = new List<int>();
@@ -293,6 +310,7 @@ public BitNetGenerationResult GenerateResponse(string prompt, int? maxTokens = n
                         if (matchedPrefixLen > 0)
                         {
                             attemptedChains++;
+                            _recallHeatMap?.RecordChainAttempt(chain.ChainId, chain.TokenIds, matchedPrefixLen);
                             var acceptedTokensForChain = 0;
                             for (var ci = matchedPrefixLen; ci < chain.TokenIds.Length && step < maxGeneratedTokens - 1; ci++)
                             {
@@ -322,6 +340,7 @@ public BitNetGenerationResult GenerateResponse(string prompt, int? maxTokens = n
                                 step++;
                                 acceptedTokensForChain++;
                                 acceptedChainTokens++;
+                                _recallHeatMap?.RecordTokenAccepted(chain.ChainId, speculativeId);
 
                                 if (Options.Verbosity == VerbosityLevel.Verbose)
                                 {
@@ -333,6 +352,7 @@ public BitNetGenerationResult GenerateResponse(string prompt, int? maxTokens = n
                             if (acceptedTokensForChain > 0)
                             {
                                 acceptedChains++;
+                                _recallHeatMap?.RecordChainAccepted(chain.ChainId);
                             }
                         }
                     }

src/BitNetSharp.Core/BitNetSharp.Core.csproj

@@ -17,4 +17,12 @@
     <PackageReference Include="System.Numerics.Tensors" Version="10.0.3" />
   </ItemGroup>
 
+  <ItemGroup Condition="Exists('Data/Models/bitnet-b1.58-default.gguf')">
+    <Content Include="Data/Models/bitnet-b1.58-default.gguf"
+             Pack="true"
+             PackagePath="contentFiles/any/any/Models/"
+             CopyToOutputDirectory="PreserveNewest"
+             PackageCopyToOutput="true" />
+  </ItemGroup>
+
 </Project>