diff --git a/go.mod b/go.mod index 2b1d1ab9..cfeb80d9 100644 --- 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 index 0976d370..71d0f112 100644 --- 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 index cbac85e4..18b33216 100644 --- 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 index 64cbe16f..e9f2f0bd 100644 --- a/pkg/plugins/gateway/prefixcacheindexer/hash_test.go +++ b/pkg/plugins/gateway/prefixcacheindexer/hash_test.go @@ -18,9 +18,11 @@ 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" @@ -28,8 +30,12 @@ import ( ) 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 index 6fde8ac8..8bb7ab51 100644 --- 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) }