fix: Pass kvCache in async stream generate function

jolonf · jolonf · commit b1d5708fbaef · 2025-05-02T23:17:09.000+10:00
diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift
@@ -732,7 +732,7 @@ public func generate(
     input: LMInput, parameters: GenerateParameters, context: ModelContext
 ) throws -> AsyncStream<Generation> {
     let iterator = try TokenIterator(
-        input: input, model: context.model, parameters: parameters)
+        input: input, model: context.model, cache: context.kvCache, parameters: parameters)
     return generate(
         input: input, context: context, iterator: iterator)
 }

Original file line number	Diff line number	Diff line change
`@@ -732,7 +732,7 @@ public func generate(`
`732`	`732`	`input: LMInput, parameters: GenerateParameters, context: ModelContext`
`733`	`733`	`) throws -> AsyncStream<Generation> {`
`734`	`734`	`let iterator = try TokenIterator(`
`735`		`- input: input, model: context.model, parameters: parameters)`
	`735`	`+ input: input, model: context.model, cache: context.kvCache, parameters: parameters)`
`736`	`736`	`return generate(`
`737`	`737`	`input: input, context: context, iterator: iterator)`
`738`	`738`	`}`