Add feature of calc ttft by prefill overhead. TODO: kvcache transfer overhead

Signed-off-by: Qifan Deng <[email protected]>

Author: pancak3

pkg/common/config.go

@@ -65,6 +65,11 @@ type Configuration struct {
 	// in milliseconds, optional, default is 0, can't be more than 30% of TimeToFirstToken, will not
 	// cause the actual time to first token to differ by more than 70% from TimeToFirstToken
 	TimeToFirstTokenStdDev int `yaml:"time-to-first-token-std-dev" json:"time-to-first-token-std-dev"`
+
+	// PrefillOverhead time taken to prefill the context, in milliseconds
+	PrefillOverhead           int    `yaml:"prefill-overhead" json:"prefill-overhead"`
+	PrefillOverheadComplexity string `yaml:"prefill-overhead-complexity" json:"prefill-overhead-complexity"`
+
 	// InterTokenLatency time between generated tokens, in milliseconds
 	InterTokenLatency int `yaml:"inter-token-latency" json:"inter-token-latency"`
 	// InterTokenLatencyStdDev standard deviation for time between generated tokens, in milliseconds,
@@ -295,6 +300,16 @@ func (c *Configuration) validate() error {
 	if float32(c.TimeToFirstTokenStdDev) > 0.3*float32(c.TimeToFirstToken) {
 		return errors.New("time to first token standard deviation cannot be more than 30% of time to first token")
 	}
+	if c.PrefillOverhead < 0 {
+		return errors.New("prefill overhead cannot be negative")
+	} else if c.PrefillOverhead == 0 {
+		if c.PrefillOverheadComplexity != "" {
+			return errors.New("prefill overhead complexity is set, but prefill overhead is 0")
+		}
+	}
+	if c.PrefillOverheadComplexity != "" && c.PrefillOverheadComplexity != "n^2" && c.PrefillOverheadComplexity != "nlog(n)" {
+		return errors.New("prefill overhead complexity should be either \"n^2\" or \"nlog(n)\"")
+	}
 	if c.KVCacheTransferLatency < 0 {
 		return errors.New("kv-cache tranfer time cannot be negative")
 	}
@@ -400,6 +415,8 @@ func ParseCommandParamsAndLoadConfig() (*Configuration, error) {
 	f.StringVar(&config.Mode, "mode", config.Mode, "Simulator mode, echo - returns the same text that was sent in the request, for chat completion returns the last message, random - returns random sentence from a bank of pre-defined sentences")
 	f.IntVar(&config.InterTokenLatency, "inter-token-latency", config.InterTokenLatency, "Time to generate one token (in milliseconds)")
 	f.IntVar(&config.TimeToFirstToken, "time-to-first-token", config.TimeToFirstToken, "Time to first token (in milliseconds)")
+	f.IntVar(&config.PrefillOverhead, "prefill-overhead", config.PrefillOverhead, "Time to prefill in milliseconds. This argument is ignored if <time-to-first-token> is not 0.")
+	f.StringVar(&config.PrefillOverheadComplexity, "prefill-overhead-complexity", config.PrefillOverheadComplexity, "Complexity of prefill based on token length. Options are \"n^2\" and \"nlog(n)\". Default is \"n^2\".")
 	f.IntVar(&config.KVCacheTransferLatency, "kv-cache-transfer-latency", config.KVCacheTransferLatency, "Time for KV-cache transfer from a remote vLLM (in milliseconds)")
 	f.IntVar(&config.InterTokenLatencyStdDev, "inter-token-latency-std-dev", config.InterTokenLatencyStdDev, "Standard deviation for time between generated tokens (in milliseconds)")
 	f.IntVar(&config.TimeToFirstTokenStdDev, "time-to-first-token-std-dev", config.TimeToFirstTokenStdDev, "Standard deviation for time before the first token will be returned (in milliseconds)")

pkg/common/config_test.go

@@ -388,6 +388,10 @@ var _ = Describe("Simulator configuration", func() {
 			name: "invalid (negative) zmq-max-connect-attempts for config file",
 			args: []string{"cmd", "--config", "../../manifests/invalid-config.yaml"},
 		},
+		{
+			name: "<prefill-overhead> must be set when <prefill-overhead-complexity> is set",
+			args: []string{"cmd", "--prefill-overhead-complexity", "n^2", "--config", "../../manifests/config.yaml"},
+		},
 	}
 
 	for _, test := range invalidTests {

pkg/llm-d-inference-sim/simulator.go

@@ -22,6 +22,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"net"
 	"os"
 	"strings"
@@ -465,7 +466,7 @@ func (s *VllmSimulator) reqProcessingWorker(ctx context.Context, id int) {
 							model:            displayModel,
 							doRemotePrefill:  req.IsDoRemotePrefill(),
 						},
-						responseTokens, toolCalls, finishReason, usageDataToSend,
+						usageDataToSend.PromptTokens, responseTokens, toolCalls, finishReason, usageDataToSend,
 					)
 				} else {
 					if req.IsDoRemoteDecode() {
@@ -633,8 +634,9 @@ func (s *VllmSimulator) sendResponse(isChatCompletion bool, ctx *fasthttp.Reques
 	}
 
 	// calculate how long to wait before returning the response, time is based on number of tokens
-	numOfTokens := usageData.CompletionTokens
-	totalMillisToWait := s.getTimeToFirstToken(doRemotePrefill) + s.getTotalInterTokenLatency(numOfTokens)
+	nPromptTokens := usageData.PromptTokens
+	nGenTokens := usageData.CompletionTokens
+	totalMillisToWait := s.getTimeToFirstToken(doRemotePrefill, nPromptTokens) + s.getTotalInterTokenLatency(nGenTokens)
 	time.Sleep(time.Duration(totalMillisToWait) * time.Millisecond)
 
 	ctx.Response.Header.SetContentType("application/json")
@@ -652,7 +654,14 @@ func (s *VllmSimulator) sendResponse(isChatCompletion bool, ctx *fasthttp.Reques
 }
 
 // returns time to first token based on the current request's doRemotePrefill
-func (s *VllmSimulator) getTimeToFirstToken(doRemotePrefill bool) int {
+func (s *VllmSimulator) getTimeToFirstToken(doRemotePrefill bool, nPromptTokens int) int {
+	if s.config.TimeToFirstToken == 0 && s.config.PrefillOverhead != 0 {
+		if nPromptTokens <= 1 {
+			return s.config.PrefillOverhead
+		}
+		return s.calcPrefillOverhead(nPromptTokens)
+	}
+
 	mean := float64(s.config.TimeToFirstToken)
 	stddev := float64(s.config.TimeToFirstTokenStdDev)
 	if doRemotePrefill {
@@ -678,6 +687,22 @@ func (s *VllmSimulator) getTotalInterTokenLatency(numOfTokens int) int {
 	return total
 }
 
+// calc the prefill overhead against number of tokens
+func (s *VllmSimulator) calcPrefillOverhead(nPromptTokens int) int {
+	pfOverhead := s.config.PrefillOverhead
+	complexity := s.config.PrefillOverheadComplexity
+	// policies of different complexities of prefill implementation
+	switch complexity {
+	case "n^2", "":
+		// this is simple implementation of n^2
+		return pfOverhead * nPromptTokens * nPromptTokens
+	case "nlog(n)":
+		return int(float64(pfOverhead) * (float64(nPromptTokens) * math.Log2(float64(nPromptTokens))))
+	}
+
+	return 0
+}
+
 // createModelsResponse creates and returns ModelResponse for the current state, returned array of models contains the base model + LoRA adapters if exist
 func (s *VllmSimulator) createModelsResponse() *vllmapi.ModelsResponse {
 	modelsResp := vllmapi.ModelsResponse{Object: "list", Data: []vllmapi.ModelsResponseModelInfo{}}

pkg/llm-d-inference-sim/simulator_test.go

@@ -21,6 +21,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"net"
 	"net/http"
 	"os"
@@ -798,7 +799,7 @@ var _ = Describe("Simulator", func() {
 				simulator.config.TimeToFirstTokenStdDev = timeToFirstTokenStdDev
 				simulator.config.KVCacheTransferLatency = kvCacheLatency
 				simulator.config.KVCacheTransferLatencyStdDev = kvCacheLatencyStdDev
-				timeToFirst := simulator.getTimeToFirstToken(doREmotePrefill)
+				timeToFirst := simulator.getTimeToFirstToken(doREmotePrefill, 1)
 				if doREmotePrefill {
 					Expect(timeToFirst).To(BeNumerically(">=", int(float32(kvCacheLatency)*0.3)))
 					Expect(timeToFirst).To(BeNumerically("<=", int(float32(kvCacheLatency)*1.7)))
@@ -819,6 +820,63 @@ var _ = Describe("Simulator", func() {
 			Entry(nil, 10000, 0, 1000, 0, true),
 			Entry(nil, 10000, 0, 1000, 0, false),
 		)
+
+		It("when <time-to-first-token> is not 0, ignore <prefill-overhead>", func() {
+			timeToFirstToken := 10000
+			prefillOverhead := 100
+			simulator.config.TimeToFirstToken = timeToFirstToken
+			simulator.config.PrefillOverhead = prefillOverhead
+			timeToFirst := simulator.getTimeToFirstToken(false, 1)
+			Expect(timeToFirst).To(BeNumerically(">=", int(float32(timeToFirstToken)*0.3)))
+			Expect(timeToFirst).To(BeNumerically("<=", int(float32(timeToFirstToken)*1.7)))
+		})
+
+		It("when <time-to-first-token> is 0, use <prefill-overhead>", func() {
+			simulator.config.TimeToFirstToken = 0
+			simulator.config.PrefillOverhead = 100
+			timeToFirst := simulator.getTimeToFirstToken(false, 1)
+			Expect(timeToFirst).To(BeNumerically(">=", 100))
+		})
+
+		DescribeTable("time to first token is super linear of prefill against number of prompt tokens",
+			func(prefillOverhead int, tolerance float64, minNTokens int, maxNTokens int) {
+				for nTokens := minNTokens; nTokens <= maxNTokens; nTokens++ {
+					square := prefillOverhead * nTokens * nTokens
+					simulator.config.PrefillOverhead = prefillOverhead
+					timeToFirst := simulator.getTimeToFirstToken(false, nTokens)
+					diffRatio := math.Abs(float64(timeToFirst-square)) / float64(square)
+					Expect(diffRatio).To(BeNumerically("<", tolerance))
+				}
+			},
+			func(prefillOverhead int, tolerance float64, minNTokens int, maxNTokens int) string {
+				return fmt.Sprintf("prefillOverhead: %d tolerance: %f minNTokens: %d maxNTokens: %d",
+					prefillOverhead, tolerance, minNTokens, maxNTokens)
+			},
+			Entry("small numbers", 100, 0.1, 1, 10),
+			Entry("medium numbers, larger range", 200, 0.1, 50, 100),
+			Entry("large numbers", 150, 0.05, 20000, 20010),
+		)
+
+		DescribeTable("time to first token is log-linear of prefill against number of prompt tokens",
+			func(prefillOverhead int, tolerance float64, minNTokens int, maxNTokens int) {
+				simulator.config.PrefillOverheadComplexity = "nlog(n)"
+
+				for nTokens := minNTokens; nTokens <= maxNTokens; nTokens++ {
+					nlogn := int(float64(prefillOverhead) * float64(nTokens) * math.Log2(float64(nTokens)))
+					simulator.config.PrefillOverhead = prefillOverhead
+					timeToFirst := simulator.getTimeToFirstToken(false, nTokens)
+					diffRatio := math.Abs(float64(timeToFirst-nlogn)) / float64(nlogn)
+					Expect(diffRatio).To(BeNumerically("<", tolerance))
+				}
+			},
+			func(prefillOverhead int, tolerance float64, minNTokens int, maxNTokens int) string {
+				return fmt.Sprintf("prefillOverhead: %d tolerance: %f minNTokens: %d maxNTokens: %d",
+					prefillOverhead, tolerance, minNTokens, maxNTokens)
+			},
+			Entry("small numbers", 100, 0.1, 2, 10),
+			Entry("medium numbers, larger range", 200, 0.1, 50, 100),
+			Entry("large numbers", 150, 0.05, 20000, 20010),
+		)
 	})
 
 	Context("fake metrics", func() {

pkg/llm-d-inference-sim/streaming.go

@@ -39,7 +39,7 @@ type streamingContext struct {
 // as defined by isChatCompletion
 // response content is wrapped according SSE format
 // First token is send after timeToFirstToken milliseconds, every other token is sent after interTokenLatency milliseconds
-func (s *VllmSimulator) sendStreamingResponse(context *streamingContext, responseTokens []string, toolCalls []openaiserverapi.ToolCall,
+func (s *VllmSimulator) sendStreamingResponse(context *streamingContext, nPromptTokens int, responseTokens []string, toolCalls []openaiserverapi.ToolCall,
 	finishReason string, usageData *openaiserverapi.Usage) {
 	context.ctx.SetContentType("text/event-stream")
 	context.ctx.SetStatusCode(fasthttp.StatusOK)
@@ -67,11 +67,11 @@ func (s *VllmSimulator) sendStreamingResponse(context *streamingContext, respons
 			if len(toolCalls) > 0 {
 				s.logger.Info("Going to send tools calls")
 				for _, tc := range toolCalls {
-					s.sendTokenChunks(context, w, tc.Function.TokenizedArguments, &tc, finishReason)
+					s.sendTokenChunks(context, w, nPromptTokens, tc.Function.TokenizedArguments, &tc, finishReason)
 				}
 			} else {
 				s.logger.Info("Going to send text", "number of tokens", len(responseTokens))
-				s.sendTokenChunks(context, w, responseTokens, nil, finishReason)
+				s.sendTokenChunks(context, w, nPromptTokens, responseTokens, nil, finishReason)
 			}
 		}
 
@@ -94,11 +94,11 @@ func (s *VllmSimulator) sendStreamingResponse(context *streamingContext, respons
 }
 
 // sendTokenChunks creates and sends response chunks
-func (s *VllmSimulator) sendTokenChunks(context *streamingContext, w *bufio.Writer, tokens []string, tc *openaiserverapi.ToolCall, finishReason string) {
+func (s *VllmSimulator) sendTokenChunks(context *streamingContext, w *bufio.Writer, nPromptTokens int, genTokens []string, tc *openaiserverapi.ToolCall, finishReason string) {
 	// time to first token delay
-	time.Sleep(time.Duration(s.getTimeToFirstToken(context.doRemotePrefill)) * time.Millisecond)
+	time.Sleep(time.Duration(s.getTimeToFirstToken(context.doRemotePrefill, nPromptTokens)) * time.Millisecond)
 
-	for i, token := range tokens {
+	for i, token := range genTokens {
 		if i != 0 {
 			time.Sleep(time.Duration(s.getInterTokenLatency()) * time.Millisecond)
 		}
@@ -119,7 +119,7 @@ func (s *VllmSimulator) sendTokenChunks(context *streamingContext, w *bufio.Writ
 
 		var chunk openaiserverapi.CompletionRespChunk
 		var finishReasonToSend *string
-		if i == len(tokens)-1 && (finishReason == common.LengthFinishReason || finishReason == common.ToolsFinishReason) {
+		if i == len(genTokens)-1 && (finishReason == common.LengthFinishReason || finishReason == common.ToolsFinishReason) {
 			finishReasonToSend = &finishReason
 		}
 		if context.isChatCompletion {