Use random seed for xxhash (#752)

Signed-off-by: Varun Gupta <[email protected]>
vllm-project · Feb 27, 2025 · fe9f91c · fe9f91c
1 parent 1b2fe77
commit fe9f91c
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 10 deletions.
diff --git a/go.mod b/go.mod
@@ -5,7 +5,7 @@ go 1.22.5
 toolchain go1.22.6
 
 require (
-	github.com/cespare/xxhash v1.1.0
+	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/envoyproxy/go-control-plane v0.12.0
 	github.com/go-playground/validator/v10 v10.22.1
 	github.com/google/uuid v1.6.0
@@ -39,7 +39,6 @@ require (
 
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
-	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cncf/xds/go v0.0.0-20240423153145-555b57ec207b // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect

diff --git a/go.sum b/go.sum
@@ -1,15 +1,11 @@
 dario.cat/mergo v0.3.16 h1:wrt7QIfeqlABnUvmf9WpFwB0mGBwtySAJKTgCpnsbOE=
 dario.cat/mergo v0.3.16/go.mod h1:WBLT9ZmE3lPoWsEzCh9LPo3TiwVN+ZKEjmz+hD27ysY=
-github.com/OneOfOne/xxhash v1.2.2 h1:KMrpdQIwFcEqXDklaen+P1axHaj9BSKzvpUUfnHldSE=
-github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
 github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
 github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
 github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
-github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
-github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cncf/xds/go v0.0.0-20240423153145-555b57ec207b h1:ga8SEFjZ60pxLcmhnThWgvH2wg8376yUJmPhEH4H3kw=

diff --git a/pkg/plugins/gateway/prefixcacheindexer/hash.go b/pkg/plugins/gateway/prefixcacheindexer/hash.go
@@ -19,11 +19,12 @@ package prefixcacheindexer
 import (
 	"bytes"
 	"encoding/binary"
+	"math/rand"
 	"strconv"
 	"sync"
 	"time"
 
-	"github.com/cespare/xxhash"
+	"github.com/cespare/xxhash/v2"
 	"github.com/vllm-project/aibrix/pkg/utils"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/klog/v2"
@@ -90,6 +91,8 @@ func getPrefixCacheEvictionDuration() time.Duration {
 type PrefixHashTable struct {
 	mu     sync.RWMutex
 	blocks map[uint64]Block
+	hash   *xxhash.Digest
+	seed   uint64
 }
 
 type Block struct {
@@ -98,8 +101,12 @@ type Block struct {
 }
 
 func NewPrefixHashTable() PrefixCacheIndexer {
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+	seed := r.Uint64()
 	instance := &PrefixHashTable{
 		blocks: map[uint64]Block{},
+		hash:   xxhash.NewWithSeed(seed),
+		seed:   seed,
 	}
 
 	ticker := time.NewTicker(prefixCacheEvictionInterval)
@@ -128,7 +135,9 @@ func (c *PrefixHashTable) MatchPrefix(tokens []int, model string, pods []*v1.Pod
 		}
 
 		chunk := tokens[i:end]
-		prefixHash := xxhash.Sum64(IntArrayToByteArray(chunk))
+		_, _ = c.hash.Write(IntArrayToByteArray(chunk))
+		prefixHash := c.hash.Sum64()
+		c.hash.ResetWithSeed(c.seed)
 		block, ok = c.blocks[prefixHash]
 		if !ok || len(block.modelToPods[model]) == 0 {
 			lastTokenMatchIndex = i
@@ -166,7 +175,9 @@ func (c *PrefixHashTable) AddPrefix(unMatchedTokens []int, model, pod string) {
 		}
 
 		chunk := unMatchedTokens[i:end]
-		prefixHash := xxhash.Sum64(IntArrayToByteArray(chunk))
+		_, _ = c.hash.Write(IntArrayToByteArray(chunk))
+		prefixHash := c.hash.Sum64()
+		c.hash.ResetWithSeed(c.seed)
 		block, ok := c.blocks[prefixHash]
 		if !ok {
 			block = Block{

diff --git a/pkg/plugins/gateway/prefixcacheindexer/hash_test.go b/pkg/plugins/gateway/prefixcacheindexer/hash_test.go
@@ -18,18 +18,24 @@ package prefixcacheindexer
 
 import (
 	"fmt"
+	"math/rand"
 	"testing"
 	"time"
 
+	"github.com/cespare/xxhash/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/vllm-project/aibrix/pkg/utils"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 func Test_PrefixHashTableE2E(t *testing.T) {
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+	seed := r.Uint64()
 	cache := PrefixHashTable{
 		blocks: map[uint64]Block{},
+		hash:   xxhash.NewWithSeed(seed),
+		seed:   seed,
 	}
 	pods := []*v1.Pod{
 		{ObjectMeta: metav1.ObjectMeta{Name: "p1"}},
@@ -64,6 +70,8 @@ func Test_PrefixHashTableE2E(t *testing.T) {
 }
 
 func Test_MatchPrefix(t *testing.T) {
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+	seed := r.Uint64()
 	tests := []*struct {
 		name          string
 		inputText     string
@@ -79,6 +87,8 @@ func Test_MatchPrefix(t *testing.T) {
 			inputText: "Hello World! What a Good Day! 你好世界！多么美好的一天啊！",
 			cache: PrefixHashTable{
 				blocks: map[uint64]Block{},
+				hash:   xxhash.NewWithSeed(seed),
+				seed:   seed,
 			},
 			model: "m1",
 			pods: []*v1.Pod{
@@ -103,6 +113,8 @@ func Test_MatchPrefix(t *testing.T) {
 						lastAccessTime: time.Now(),
 					},
 				},
+				hash: xxhash.NewWithSeed(0),
+				seed: 0,
 			},
 			model: "m1",
 			pods: []*v1.Pod{

diff --git a/pkg/utils/util_test.go b/pkg/utils/util_test.go
@@ -29,8 +29,12 @@ func TestTokenizeInputText(t *testing.T) {
 	tokens, err := TokenizeInputText(inputStr)
 	assert.Equal(t, nil, err)
 
+	outputStr, err := DetokenizeText(tokens)
+	assert.NoError(t, err)
+	assert.Equal(t, inputStr, outputStr)
+
 	tiktoken.SetBpeLoader(tiktoken_loader.NewOfflineLoader())
 	tke, _ := tiktoken.GetEncoding(encoding)
-	outputStr := tke.Decode(tokens)
+	outputStr = tke.Decode(tokens)
 	assert.Equal(t, inputStr, outputStr)
 }