Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use random seed for xxhash #752

Merged
merged 2 commits into from
Feb 27, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
4 changes: 0 additions & 4 deletions go.sum
Original file line number Diff line number Diff line change
@@ -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=
Expand Down
17 changes: 14 additions & 3 deletions pkg/plugins/gateway/prefixcacheindexer/hash.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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{
Expand Down
12 changes: 12 additions & 0 deletions pkg/plugins/gateway/prefixcacheindexer/hash_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"}},
Expand Down Expand Up @@ -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
Expand All @@ -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{
Expand All @@ -103,6 +113,8 @@ func Test_MatchPrefix(t *testing.T) {
lastAccessTime: time.Now(),
},
},
hash: xxhash.NewWithSeed(0),
seed: 0,
},
model: "m1",
pods: []*v1.Pod{
Expand Down
6 changes: 5 additions & 1 deletion pkg/utils/util_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Loading