matt/ollama37
1
0
Fork 0
mirror of https://github.com/dogkeeper886/ollama37.git synced 2025-12-11 08:17:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Disable causal attention based on batch index

Browse Source 
Currently we are using positions, which are relative to a
sequence and may not be unique.
This commit is contained in:
Jesse Gross
2025-03-10 17:17:19 -07:00
committed by Michael Yang
parent 475005504e
commit a8e83a7654
2 changed files with 10 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
kvcache/causal.go
View File

@@ -144,6 +144,7 @@ func (c *Causal) StartForward(ctx ml.Context, opts input.Options) error {
c.curBatchSize = len(opts.Positions)
c.curSequences = opts.Sequences
c.curPositions = opts.Positions
c.opts.Except = nil
var err error
c.curLoc, err = c.findStartLoc()
@@ -234,7 +235,7 @@ func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) {
mask := make([]float32, batchSize*length)
for i := range c.curBatchSize {
enabled := !slices.Contains(c.opts.Except, c.curPositions[i])
enabled := !slices.Contains(c.opts.Except, i)
for j := c.curCellRange.min; j <= c.curCellRange.max; j++ {
if !slices.Contains(c.cells[j].sequences, c.curSequences[i]) ||
(enabled && c.cells[j].pos > c.curPositions[i]) ||
@@ -405,15 +406,12 @@ func (c *Causal) SetLayer(layer int) {
}
type CausalOptions struct {
// Enabled controls whether the causal mask is generated for a particular position.
Except []int32
// Enabled controls whether the causal mask is generated for a particular index in a batch
Except []int
}
// SetCausal enables or disables causal mask generation for subsequent calls to Get.
// This state carries over to future forward passes. The default value is true.
//
// ctx may be set to nil if this is called from outside of a forward pass, for
// example, when initializing the cache.
// SetCausal disables causal mask generation for a particular range of indicies in
// the current batch for subsequent calls to Get. The state resets for the next forward pass.
func (c *Causal) SetCausal(ctx ml.Context, opts CausalOptions) {
if !slices.Equal(c.opts.Except, opts.Except) {
c.opts = opts