Kv cache support without KV events (#107)

* first steps of kv-cache support: update build process, add KVCacheHelper

Signed-off-by: Maya Barnea <[email protected]>

* fix build and lint problems after merge

Signed-off-by: Maya Barnea <[email protected]>

* Add command line and yaml config file parameter 'enable-kvcache'

Signed-off-by: Maya Barnea <[email protected]>

* - continue kv-cache support implementation
- add block cache and initial version of test for it

Signed-off-by: Maya Barnea <[email protected]>

* move kv cache related code to a separate package

Signed-off-by: Maya Barnea <[email protected]>

* update kv cache test to be more modular

Signed-off-by: Maya Barnea <[email protected]>

* fixes in kv-cache-manager + thread save test

Signed-off-by: Maya Barnea <[email protected]>

* fix make test

Signed-off-by: Maya Barnea <[email protected]>

* fix typo

Signed-off-by: Maya Barnea <[email protected]>

* remove redundant field initialization

Signed-off-by: Maya Barnea <[email protected]>

* rename cache + fixes according to PR comments

Signed-off-by: Maya Barnea <[email protected]>

* fixes according to PR comments

Signed-off-by: Maya Barnea <[email protected]>

* fixes by PR comments

Signed-off-by: Maya Barnea <[email protected]>

---------

Signed-off-by: Maya Barnea <[email protected]>

Author: mayabar

.gitignore

@@ -1,4 +1,6 @@
 bin
+lib
+vendor
 .vscode
 # MacOSX
 .DS_Store

Dockerfile

@@ -3,6 +3,12 @@ FROM quay.io/projectquay/golang:1.24 AS builder
 ARG TARGETOS
 ARG TARGETARCH
 
+# Install build tools
+# The builder is based on UBI8, so we need epel-release-8.
+RUN dnf install -y 'https://dl.fedoraproject.org/pub/epel/epel-release-latest-8.noarch.rpm' && \
+    dnf install -y gcc-c++ libstdc++ libstdc++-devel clang zeromq-devel pkgconfig && \
+    dnf clean all
+
 WORKDIR /workspace
 # Copy the Go Modules manifests
 COPY go.mod go.mod
@@ -15,18 +21,29 @@ RUN go mod download
 COPY cmd/llm-d-inference-sim/main.go cmd/cmd.go
 COPY . .
 
+# HuggingFace tokenizer bindings
+RUN mkdir -p lib
+RUN curl -L https://github.com/daulet/tokenizers/releases/download/v1.20.2/libtokenizers.${TARGETOS}-${TARGETARCH}.tar.gz | tar -xz -C lib
+RUN ranlib lib/*.a
+
 # Build
 # the GOARCH has not a default value to allow the binary be built according to the host where the command
 # was called. For example, if we call make image-build in a local env which has the Apple Silicon M1 SO
 # the docker BUILDPLATFORM arg will be linux/arm64 when for Apple x86 it will be linux/amd64. Therefore,
 # by leaving it empty we can ensure that the container and binary shipped on it will have the same platform.
-RUN CGO_ENABLED=0 GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH} go build -a -o bin/llm-d-inference-sim cmd/cmd.go
+ENV CGO_ENABLED=1
+ENV GOOS=${TARGETOS:-linux}
+ENV GOARCH=${TARGETARCH}
+RUN go build -a -o bin/llm-d-inference-sim -ldflags="-extldflags '-L$(pwd)/lib'" cmd/cmd.go
+
+# Use ubi9 as a minimal base image to package the manager binary
+# Refer to https://catalog.redhat.com/software/containers/ubi9/ubi-minimal/615bd9b4075b022acc111bf5 for more details
+FROM registry.access.redhat.com/ubi9/ubi-minimal:latest
 
-# Use ubi9 as a micro base image to package the manager binary
-# Refer to https://catalog.redhat.com/software/containers/ubi9/ubi-micro/615bdf943f6014fa45ae1b58 for more details
-FROM registry.access.redhat.com/ubi9/ubi-micro:latest
 WORKDIR /
 COPY --from=builder /workspace/bin/llm-d-inference-sim /app/llm-d-inference-sim
-USER 65532:65532
+
+# USER 65532:65532
+USER root
 
 ENTRYPOINT ["/app/llm-d-inference-sim"]

Makefile

@@ -14,16 +14,17 @@
 
 # Makefile for the llm-d-inference-sim project
 
-CONTAINER_RUNTIME ?= docker
-
 SHELL := /usr/bin/env bash
 
 # Defaults
+TARGETOS ?= $(shell go env GOOS)
+TARGETARCH ?= $(shell go env GOARCH)
 PROJECT_NAME ?= llm-d-inference-sim
-REGISTRY ?= ghcr.io/llm-d
-IMAGE_TAG_BASE ?= $(REGISTRY)/$(PROJECT_NAME)
+IMAGE_REGISTRY ?= ghcr.io/llm-d
+IMAGE_TAG_BASE ?= $(IMAGE_REGISTRY)/$(PROJECT_NAME)
 SIM_TAG ?= dev
 IMG = $(IMAGE_TAG_BASE):$(SIM_TAG)
+
 CONTAINER_TOOL := $(shell { command -v docker >/dev/null 2>&1 && echo docker; } || { command -v podman >/dev/null 2>&1 && echo podman; } || echo "")
 BUILDER := $(shell command -v buildah >/dev/null 2>&1 && echo buildah || echo $(CONTAINER_TOOL))
 PLATFORMS ?= linux/amd64 # linux/arm64 # linux/s390x,linux/ppc64le
@@ -35,17 +36,36 @@ SRC = $(shell find . -type f -name '*.go')
 help: ## Print help
 	@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n  make \033[36m<target>\033[0m\n"} /^[a-zA-Z_0-9-]+:.*?##/ { printf "  \033[36m%-15s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)
 
+LDFLAGS ?= -extldflags '-L$(shell pwd)/lib'
+CGO_ENABLED=1
+TOKENIZER_LIB = lib/libtokenizers.a
+
+.PHONY: download-tokenizer
+download-tokenizer: $(TOKENIZER_LIB)
+$(TOKENIZER_LIB):
+	## Download the HuggingFace tokenizer bindings.
+	@echo "Downloading HuggingFace tokenizer bindings..."
+	mkdir -p lib
+	curl -L https://github.com/daulet/tokenizers/releases/download/v1.20.2/libtokenizers.$(TARGETOS)-$(TARGETARCH).tar.gz | tar -xz -C lib
+	ranlib lib/*.a
+
 ##@ Development
 
+.PHONY: clean
+clean:
+	go clean -testcache -cache
+	rm -f $(TOKENIZER_LIB)
+	rmdir lib
+
 .PHONY: format
 format: ## Format Go source files
 	@printf "\033[33;1m==== Running gofmt ====\033[0m\n"
 	@gofmt -l -w $(SRC)
 
 .PHONY: test
-test: check-ginkgo ## Run tests
+test: check-ginkgo download-tokenizer ## Run tests
 	@printf "\033[33;1m==== Running tests ====\033[0m\n"
-	ginkgo -r -v
+	CGO_ENABLED=1 ginkgo -ldflags="$(LDFLAGS)" -v -r
 
 .PHONY: post-deploy-test
 post-deploy-test: ## Run post deployment tests
@@ -60,25 +80,29 @@ lint: check-golangci-lint ## Run lint
 ##@ Build
 
 .PHONY: build
-build: check-go ##
+build: check-go download-tokenizer ##
 	@printf "\033[33;1m==== Building ====\033[0m\n"
-	go build -o bin/$(PROJECT_NAME) cmd/$(PROJECT_NAME)/main.go
+	go build -ldflags="$(LDFLAGS)" -o bin/$(PROJECT_NAME) cmd/$(PROJECT_NAME)/main.go
 
 ##@ Container Build/Push
 
-.PHONY: image-build-and-push
-image-build-and-push: image-build image-push ## Build and push Docker image $(IMG) to registry
-
 .PHONY:	image-build
 image-build: check-container-tool ## Build Docker image ## Build Docker image using $(CONTAINER_TOOL)
 	@printf "\033[33;1m==== Building Docker image $(IMG) ====\033[0m\n"
-	$(CONTAINER_TOOL) build --build-arg TARGETOS=$(TARGETOS) --build-arg TARGETARCH=$(TARGETARCH) -t $(IMG) .
+	$(CONTAINER_TOOL) build \
+		--platform $(TARGETOS)/$(TARGETARCH) \
+		--build-arg TARGETOS=$(TARGETOS)\
+		--build-arg TARGETARCH=$(TARGETARCH)\
+		-t $(IMG) .
 
 .PHONY: image-push
 image-push: check-container-tool ## Push Docker image $(IMG) to registry
 	@printf "\033[33;1m==== Pushing Docker image $(IMG) ====\033[0m\n"
 	$(CONTAINER_TOOL) push $(IMG)
 
+.PHONY: image-build-and-push
+image-build-and-push: image-build image-push ## Build and push Docker image $(IMG) to registry
+
 ##@ Install/Uninstall Targets
 
 # Default install/uninstall (Docker)

README.md

@@ -112,6 +112,7 @@ For more details see the <a href="https://docs.vllm.ai/en/stable/getting_started
 - `min-tool-call-array-param-length`: the minimum possible length of array parameters in a tool call, optional, defaults to 1
 - `tool-call-not-required-param-probability`: the probability to add a parameter, that is not required, in a tool call, optional, defaults to 50
 - `object-tool-call-not-required-field-probability`: the probability to add a field, that is not required, in an object in a tool call, optional, defaults to 50
+- `enable-kvcache`: if true, the KV cache support will be enabled in the simulator. In this case, the KV cache will be simulated, and ZQM events will be published when a KV cache block is added or evicted.
 
 In addition, as we are using klog, the following parameters are available:
 - `add_dir_header`: if true, adds the file directory to the header of the log messages

go.mod

@@ -1,17 +1,18 @@
 module github.com/llm-d/llm-d-inference-sim
 
-go 1.23.4
+go 1.24.1
 
-toolchain go1.24.0
+toolchain go1.24.2
 
 require (
 	github.com/buaazp/fasthttprouter v0.1.1
 	github.com/go-logr/logr v1.4.2
 	github.com/google/uuid v1.6.0
+	github.com/llm-d/llm-d-kv-cache-manager v0.2.0
 	github.com/onsi/ginkgo/v2 v2.23.4
 	github.com/onsi/gomega v1.37.0
 	github.com/openai/openai-go v0.1.0-beta.10
-	github.com/prometheus/client_golang v1.21.1
+	github.com/prometheus/client_golang v1.22.0
 	github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
 	github.com/spf13/pflag v1.0.6
 	github.com/valyala/fasthttp v1.59.0
@@ -23,25 +24,57 @@ require (
 	github.com/andybalholm/brotli v1.1.1 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/daulet/tokenizers v1.20.2 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+	github.com/emicklei/go-restful/v3 v3.11.0 // indirect
+	github.com/fxamacker/cbor/v2 v2.7.0 // indirect
+	github.com/go-openapi/jsonpointer v0.21.0 // indirect
+	github.com/go-openapi/jsonreference v0.20.2 // indirect
+	github.com/go-openapi/swag v0.23.0 // indirect
 	github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/google/gnostic-models v0.6.9 // indirect
 	github.com/google/go-cmp v0.7.0 // indirect
 	github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 // indirect
+	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
-	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.62.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/redis/go-redis/v9 v9.7.3 // indirect
 	github.com/tidwall/gjson v1.18.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect
 	github.com/tidwall/sjson v1.2.5 // indirect
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
+	github.com/x448/float16 v0.8.4 // indirect
 	go.uber.org/automaxprocs v1.6.0 // indirect
 	golang.org/x/net v0.38.0 // indirect
+	golang.org/x/oauth2 v0.27.0 // indirect
 	golang.org/x/sys v0.32.0 // indirect
+	golang.org/x/term v0.30.0 // indirect
 	golang.org/x/text v0.23.0 // indirect
+	golang.org/x/time v0.9.0 // indirect
 	golang.org/x/tools v0.31.0 // indirect
 	google.golang.org/protobuf v1.36.5 // indirect
+	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
+	k8s.io/api v0.33.0 // indirect
+	k8s.io/apimachinery v0.33.0 // indirect
+	k8s.io/client-go v0.33.0 // indirect
+	k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff // indirect
+	k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738 // indirect
+	sigs.k8s.io/controller-runtime v0.21.0 // indirect
+	sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 // indirect
+	sigs.k8s.io/randfill v1.0.0 // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.6.0 // indirect
+	sigs.k8s.io/yaml v1.4.0 // indirect
 )