Optimize exact int8 late-interaction scoring

Author: XiaoConstantine

pkg/conv/actions.go

@@ -328,16 +328,16 @@ func formatContextOutput(result *ContextResult, format string) string {
 	switch format {
 	case "markdown":
 		sb.WriteString("## Previous Session Context\n\n")
-		sb.WriteString(fmt.Sprintf("**Project:** %s\n", result.Project))
-		sb.WriteString(fmt.Sprintf("**Date:** %s (%s)\n", result.Date.Format("2006-01-02"), relativeTime(result.Date)))
-		sb.WriteString(fmt.Sprintf("**Agent:** %s\n", result.Agent))
+		fmt.Fprintf(&sb, "**Project:** %s\n", result.Project)
+		fmt.Fprintf(&sb, "**Date:** %s (%s)\n", result.Date.Format("2006-01-02"), relativeTime(result.Date))
+		fmt.Fprintf(&sb, "**Agent:** %s\n", result.Agent)
 		if result.Topic != "" {
-			sb.WriteString(fmt.Sprintf("**Topic:** %s\n", result.Topic))
+			fmt.Fprintf(&sb, "**Topic:** %s\n", result.Topic)
 		}
 		sb.WriteString("\n")
 
 		if result.Summary != "" {
-			sb.WriteString(fmt.Sprintf("### Summary\n%s\n\n", result.Summary))
+			fmt.Fprintf(&sb, "### Summary\n%s\n\n", result.Summary)
 		}
 
 		if len(result.KeyPoints) > 0 {

pkg/index/index.go

@@ -565,16 +565,6 @@ func (idx *Indexer) prepareFile(ctx context.Context, path string) ([]*store.Docu
 	return docs, nil
 }
 
-// indexFile is a convenience wrapper for single-file indexing (used by Watch).
-func (idx *Indexer) indexFile(ctx context.Context, path string) error {
-	refreshColBERT, err := idx.shouldRefreshColBERT(ctx)
-	if err != nil {
-		return err
-	}
-
-	return idx.syncFile(ctx, path, refreshColBERT, true)
-}
-
 func (idx *Indexer) syncFile(ctx context.Context, path string, refreshColBERT, rebuildColBERTMMap bool) error {
 	relPath, err := filepath.Rel(idx.rootPath, path)
 	if err != nil {

pkg/search/colbert.go

@@ -25,6 +25,11 @@ type ColBERTScorer struct {
 	cache        *segmentCache
 }
 
+type preparedQueryTerm struct {
+	embedding []float32
+	sum       float64
+}
+
 // segmentCache caches PRE-NORMALIZED segment embeddings to avoid recomputation.
 // Storing normalized vectors enables fast similarity via dot product.
 type segmentCache struct {
@@ -228,6 +233,7 @@ func (c *ColBERTScorer) ScoreBatchWithChunkIDs(ctx context.Context, query string
 	if err != nil {
 		return nil, err
 	}
+	preparedTerms := prepareQueryTerms(queryEmbeddings)
 
 	// Batch load pre-computed segment embeddings for all chunks
 	segmentMap, err := c.segmentStore.GetColBERTSegmentsBatch(ctx, chunkIDs)
@@ -274,8 +280,8 @@ func (c *ColBERTScorer) ScoreBatchWithChunkIDs(ctx context.Context, query string
 
 		// Compute MaxSim score using int8 quantized embeddings if available
 		var totalScore float64
-		for _, qEmb := range queryEmbeddings {
-			maxSim := maxSimInt8(qEmb, segments)
+		for _, term := range preparedTerms {
+			maxSim := maxSimPreparedInt8(term, segments)
 			if maxSim > 0 {
 				totalScore += maxSim
 			}
@@ -585,3 +591,42 @@ func maxSimInt8(qEmb []float32, segments []store.ColBERTSegment) float64 {
 	}
 	return maxSim
 }
+
+func prepareQueryTerms(queryEmbeddings [][]float32) []preparedQueryTerm {
+	terms := make([]preparedQueryTerm, len(queryEmbeddings))
+	for i, emb := range queryEmbeddings {
+		terms[i] = preparedQueryTerm{
+			embedding: emb,
+			sum:       sumFloat32(emb),
+		}
+	}
+	return terms
+}
+
+func maxSimPreparedInt8(term preparedQueryTerm, segments []store.ColBERTSegment) float64 {
+	if len(segments) == 0 {
+		return -1
+	}
+
+	maxSim := float64(-1)
+	for _, seg := range segments {
+		var sim float64
+		if seg.EmbeddingInt8 != nil {
+			sim = util.DotProductInt8AffinePrepared(term.embedding, seg.EmbeddingInt8, seg.QuantScale, seg.QuantMin, term.sum)
+		} else if seg.Embedding != nil {
+			sim = util.DotProductUnrolled8(term.embedding, seg.Embedding)
+		}
+		if sim > maxSim {
+			maxSim = sim
+		}
+	}
+	return maxSim
+}
+
+func sumFloat32(values []float32) float64 {
+	var sum float64
+	for _, v := range values {
+		sum += float64(v)
+	}
+	return sum
+}

pkg/search/colbert_test.go

@@ -1,7 +1,11 @@
 package search
 
 import (
+	"math"
 	"testing"
+
+	"github.com/XiaoConstantine/sgrep/pkg/store"
+	"github.com/XiaoConstantine/sgrep/pkg/util"
 )
 
 func TestDecomposeQuery(t *testing.T) {
@@ -29,22 +33,22 @@ func TestDecomposeQuery(t *testing.T) {
 
 func TestDecomposeDocument(t *testing.T) {
 	tests := []struct {
-		name     string
-		content  string
-		minSegs  int
-		maxSegs  int
+		name    string
+		content string
+		minSegs int
+		maxSegs int
 	}{
 		{
-			name:     "empty",
-			content:  "",
-			minSegs:  0,
-			maxSegs:  0,
+			name:    "empty",
+			content: "",
+			minSegs: 0,
+			maxSegs: 0,
 		},
 		{
-			name:     "single line",
-			content:  "func main() {}",
-			minSegs:  1,
-			maxSegs:  1,
+			name:    "single line",
+			content: "func main() {}",
+			minSegs: 1,
+			maxSegs: 1,
 		},
 		{
 			name: "multiple functions",
@@ -221,3 +225,51 @@ func TestSegmentCache(t *testing.T) {
 	// At least one of the original keys should be evicted
 	// Note: eviction is simple and clears half, so behavior may vary
 }
+
+func TestPrepareQueryTerms(t *testing.T) {
+	embeddings := [][]float32{
+		{0.25, -0.5, 0.75},
+		{-0.25, 0.5, -0.75},
+	}
+
+	terms := prepareQueryTerms(embeddings)
+	if len(terms) != len(embeddings) {
+		t.Fatalf("prepareQueryTerms returned %d terms, want %d", len(terms), len(embeddings))
+	}
+
+	for i, term := range terms {
+		if len(term.embedding) != len(embeddings[i]) {
+			t.Fatalf("term %d embedding len = %d, want %d", i, len(term.embedding), len(embeddings[i]))
+		}
+		if diff := math.Abs(term.sum - sumFloat32(embeddings[i])); diff > 1e-9 {
+			t.Fatalf("term %d sum diff = %.12f", i, diff)
+		}
+	}
+}
+
+func TestMaxSimPreparedInt8MatchesCurrent(t *testing.T) {
+	query := util.NormalizeVector([]float32{0.9, -0.3, 0.2, 0.1})
+	segments := make([]store.ColBERTSegment, 3)
+
+	rawSegments := [][]float32{
+		util.NormalizeVector([]float32{0.8, -0.25, 0.25, 0.15}),
+		util.NormalizeVector([]float32{-0.4, 0.9, -0.1, 0.05}),
+		util.NormalizeVector([]float32{0.2, -0.1, 0.95, -0.05}),
+	}
+
+	for i, emb := range rawSegments {
+		quantized, scale, min := util.QuantizeInt8(emb)
+		segments[i] = store.ColBERTSegment{
+			SegmentIdx:    i,
+			EmbeddingInt8: quantized,
+			QuantScale:    scale,
+			QuantMin:      min,
+		}
+	}
+
+	current := maxSimInt8(query, segments)
+	prepared := maxSimPreparedInt8(prepareQueryTerms([][]float32{query})[0], segments)
+	if diff := math.Abs(current - prepared); diff > 1e-7 {
+		t.Fatalf("prepared scorer mismatch: current %.12f prepared %.12f diff %.12f", current, prepared, diff)
+	}
+}

pkg/util/slab.go

@@ -364,6 +364,41 @@ func DotProductInt8Unrolled8(query []float32, doc []int8, scale, min float32) fl
 	return s0 + s1 + s2 + s3 + s4 + s5 + s6 + s7
 }
 
+// DotProductInt8AffinePrepared computes the exact same score as DotProductInt8Unrolled8,
+// but factors out the affine dequantization so query-specific work can be prepared once:
+//
+//	dot(q, ((doc + 128) * scale + min))
+//	= scale*dot(q, doc) + (128*scale + min)*sum(q)
+func DotProductInt8AffinePrepared(query []float32, doc []int8, scale, min float32, querySum float64) float64 {
+	n := len(query)
+	if n != len(doc) {
+		return 0
+	}
+
+	scale64 := float64(scale)
+	bias := float64(128*scale + min)
+
+	var s0, s1, s2, s3, s4, s5, s6, s7 float64
+
+	i := 0
+	for ; i <= n-8; i += 8 {
+		s0 += float64(query[i]) * float64(doc[i])
+		s1 += float64(query[i+1]) * float64(doc[i+1])
+		s2 += float64(query[i+2]) * float64(doc[i+2])
+		s3 += float64(query[i+3]) * float64(doc[i+3])
+		s4 += float64(query[i+4]) * float64(doc[i+4])
+		s5 += float64(query[i+5]) * float64(doc[i+5])
+		s6 += float64(query[i+6]) * float64(doc[i+6])
+		s7 += float64(query[i+7]) * float64(doc[i+7])
+	}
+
+	for ; i < n; i++ {
+		s0 += float64(query[i]) * float64(doc[i])
+	}
+
+	return scale64*(s0+s1+s2+s3+s4+s5+s6+s7) + bias*querySum
+}
+
 // TopKIndices returns indices of k smallest values in distances (partial sort).
 // Uses indices buffer for zero allocation.
 func TopKIndices(distances []float64, indices []int, k int) []int {

pkg/util/slab_test.go

@@ -220,8 +220,8 @@ func TestL2DistanceBatch(t *testing.T) {
 func TestL2DistanceBatch_DimensionMismatch(t *testing.T) {
 	query := []float32{0.0, 0.0, 0.0}
 	vectors := [][]float32{
-		{0.0, 0.0},       // wrong dims
-		{1.0, 0.0, 0.0},  // correct dims
+		{0.0, 0.0},           // wrong dims
+		{1.0, 0.0, 0.0},      // correct dims
 		{0.0, 3.0, 4.0, 5.0}, // wrong dims
 	}
 	distances := make([]float64, len(vectors))
@@ -239,6 +239,38 @@ func TestL2DistanceBatch_DimensionMismatch(t *testing.T) {
 	}
 }
 
+func TestDotProductInt8AffinePreparedMatchesCurrent(t *testing.T) {
+	rng := rand.New(rand.NewSource(42))
+
+	for caseIdx := 0; caseIdx < 128; caseIdx++ {
+		query := make([]float32, 768)
+		doc := make([]float32, 768)
+		for i := range query {
+			query[i] = rng.Float32()*2 - 1
+			doc[i] = rng.Float32()*2 - 1
+		}
+		query = NormalizeVector(query)
+		doc = NormalizeVector(doc)
+
+		docInt8, scale, min := QuantizeInt8(doc)
+		querySum := sumFloat32Values(query)
+
+		got := DotProductInt8AffinePrepared(query, docInt8, scale, min, querySum)
+		want := DotProductInt8Unrolled8(query, docInt8, scale, min)
+		if diff := math.Abs(got - want); diff > 1e-7 {
+			t.Fatalf("case %d: got %.12f want %.12f diff %.12f", caseIdx, got, want, diff)
+		}
+	}
+}
+
+func sumFloat32Values(values []float32) float64 {
+	var sum float64
+	for _, v := range values {
+		sum += float64(v)
+	}
+	return sum
+}
+
 func TestTopKIndices(t *testing.T) {
 	tests := []struct {
 		name      string