Repository: ollama/ollama
File: runner/llamarunner/cache.go (Go)

Code before:
	slog.Debug("context limit hit - shifting", "id", slot.Id, "limit", c.numCtx, "input", len(slot.Inputs),
		"keep", numKeep, "discard", discard)

	if c.lc.KvCacheCanShift() {
		if !c.lc.KvCacheSeqRm(slot.Id, numKeep, numKeep+discard) {
			return fmt.Errorf("unable to remove old kv cache entries (id: %v, keep: %v discard: %v)", slot.Id, numKeep, discard)
		}
		c.lc.KvCacheSeqAdd(slot.Id, numKeep+discard, len(slot.Inputs), -discard)

		for i := numKeep + discard; i < len(slot.Inputs); i++ {
			slot.Inputs[i-discard] = slot.Inputs[i]
		}
		slot.Inputs = slot.Inputs[:len(slot.Inputs)-discard]
	} else {
		slog.Debug("kv cache cannot shift, clearing cache and truncating history")

		// Clear the entire KV cache
		if !c.lc.KvCacheSeqRm(slot.Id, 0, -1) {
			return fmt.Errorf("unable to remove kv cache entries (id: %v)", slot.Id)
		}

		// Update the slot.Inputs to match what would happen with a shift operation
		// Keep the first numKeep tokens and the tokens after the discard
		keepInputs := make([]input, numKeep)
		copy(keepInputs, slot.Inputs[:numKeep])

		afterDiscardInputs := make([]input, len(slot.Inputs)-(numKeep+discard))
		copy(afterDiscardInputs, slot.Inputs[numKeep+discard:])

		// Update the inputs to match what would happen after a shift
		newInputs := make([]input, 0, numKeep+len(afterDiscardInputs))
		newInputs = append(newInputs, keepInputs...)
		newInputs = append(newInputs, afterDiscardInputs...)
		slot.Inputs = newInputs
	}

	return nil
}


Review comment (performance):
`c.lc.KvCacheSeqRm` can return false for a different reason than the Deepseek issue but the solution is the same, so I think we can use the same logic and avoid returning an error.

Code after:
	}

	inputLen := len(slot.Inputs)
	discard := c.ShiftDiscard(inputLen, numKeep)

	if discard <= 0 {
		return nil
	}

	slog.Debug("context limit hit - shifting", "id", slot.Id, "limit", c.numCtx, "input", len(slot.Inputs),
		"keep", numKeep, "discard", discard)

	var shiftFailed bool

	if c.lc.KvCacheCanShift() {
		// For models that support shifting, attempt to shift the KV cache
		if !c.lc.KvCacheSeqRm(slot.Id, numKeep, numKeep+discard) {
			shiftFailed = true
			slog.Debug("kv cache removal not supported, clearing cache and returning inputs for reprocessing", "id", slot.Id)
		} else {
			c.lc.KvCacheSeqAdd(slot.Id, numKeep+discard, inputLen, -discard)
		}
	} else {
		// For models that don't support shifting
		shiftFailed = true
		slog.Debug("kv cache cannot shift, clearing cache and returning inputs for reprocessing", "id", slot.Id)
	}

	if shiftFailed {
		// Create new input slice with preserved tokens (numKeep + remaining tokens after discard)
		newInputs := make([]input, numKeep+inputLen-(numKeep+discard))
		copy(newInputs[:numKeep], slot.Inputs[:numKeep])
		copy(newInputs[numKeep:], slot.Inputs[numKeep+discard:])

		// Clear the entire KV cache
		_ = c.lc.KvCacheSeqRm(slot.Id, 0, -1)
		// Reset the slot inputs since we've cleared the cache
		slot.Inputs = []input{}

		// Return error with inputs that need to be reprocessed
		return &ErrReprocessInputs{Inputs: newInputs}
	}

	// Standard shift succeeded - update input array
	for i := numKeep + discard; i < inputLen; i++ {
		slot.Inputs[i-discard] = slot.Inputs[i]
	}
	slot.Inputs = slot.Inputs[:inputLen-discard]

	return nil
}