github-starred/ollama
2
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-07-17 05:36:16 -05:00
Code Issues 16.9k Packages Projects Releases 122 Wiki Activity

[PR #15814] [MERGED] Model support for batching #62014

New Issue
Closed
opened 2026-04-29 16:58:23 -05:00 by GiteaMirror · 0 comments
GiteaMirror commented 2026-04-29 16:58:23 -05:00
Owner
Copy Link

📋 Pull Request Information

Original PR: https://github.com/ollama/ollama/pull/15814
Author: @jessegross
Created: 4/25/2026
Status: ✅ Merged
Merged: 4/28/2026
Merged by: @jessegross

Base: main ← Head: jessegross/batching

📝 Commits (3)

  • 7133e9d mlxrunner: wrap model forward inputs in a Batch struct
  • 5c75a1e mlxrunner: apply RoPE at per-row positions
  • b170b37 mlxrunner: decouple models from attention cache storage layout

📊 Changes

22 files changed (+2750 additions, -425 deletions)

View changed files

➕ x/mlxrunner/batch/batch.go (+42 -0)
📝 x/mlxrunner/cache/cache.go (+106 -5)
📝 x/mlxrunner/cache/cache_test.go (+22 -10)
📝 x/mlxrunner/cache/recurrent.go (+21 -23)
📝 x/mlxrunner/cache/recurrent_test.go (+201 -4)
➕ x/mlxrunner/cache/rotating_attention_test.go (+237 -0)
📝 x/mlxrunner/cache/rotating_multiturn_test.go (+2 -2)
📝 x/mlxrunner/mlx/fast.go (+10 -35)
📝 x/mlxrunner/mlx/gated_delta.go (+27 -3)
📝 x/mlxrunner/mlx/ops_extra.go (+8 -34)
📝 x/mlxrunner/model/base/base.go (+2 -1)
📝 x/mlxrunner/pipeline.go (+16 -6)
📝 x/models/gemma3/gemma3.go (+19 -18)
📝 x/models/gemma4/gemma4.go (+70 -135)
📝 x/models/glm4_moe_lite/glm4_moe_lite.go (+28 -24)
📝 x/models/llama/llama.go (+19 -18)
➕ x/models/nn/recurrent.go (+259 -0)
➕ x/models/nn/recurrent_test.go (+340 -0)
➕ x/models/nn/sdpa.go (+578 -0)
➕ x/models/nn/sdpa_test.go (+680 -0)

...and 2 more files

📄 Description

Three preparatory changes for multi-sequence batching in mlxrunner. The runner still feeds one sequence at a time after this PR, but the model contract, RoPE, and SDPA/cache abstractions stop assuming B=1 with a single starting position.

  • Wrap model forward inputs in a Batch struct — one extension point for per-row context (positions, sequence IDs, masks) instead of churning every model's Forward signature.
  • Apply RoPE at per-row positions — switch from the scalar-offset RoPE kernel to the array-offset one so each row can start at its own position. The pipeline tracks the position locally and passes it through Batch.SeqOffsets.
  • Decouple models from attention cache storage layout — models describe a logical AttentionMask and receive a per-layer KVHistory carrying K, V, and a MaskApplier. SDPA composes the model mask, query padding, and the cache's storage restrictions, still hitting the kernel's causal/no-mask fast paths when nothing forces a tensor mask. Also unblocks cache variants like paged attention.

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/ollama/ollama/pull/15814 **Author:** [@jessegross](https://github.com/jessegross) **Created:** 4/25/2026 **Status:** ✅ Merged **Merged:** 4/28/2026 **Merged by:** [@jessegross](https://github.com/jessegross) **Base:** `main` ← **Head:** `jessegross/batching` --- ### 📝 Commits (3) - [`7133e9d`](https://github.com/ollama/ollama/commit/7133e9d8b0d971f6029bb9dd5ce8f047a19d06f5) mlxrunner: wrap model forward inputs in a Batch struct - [`5c75a1e`](https://github.com/ollama/ollama/commit/5c75a1e2f500627bca735c54ddb5566b4548cddd) mlxrunner: apply RoPE at per-row positions - [`b170b37`](https://github.com/ollama/ollama/commit/b170b37b035509510cf6cfca7e96e6e3792bfa48) mlxrunner: decouple models from attention cache storage layout ### 📊 Changes **22 files changed** (+2750 additions, -425 deletions) <details> <summary>View changed files</summary> ➕ `x/mlxrunner/batch/batch.go` (+42 -0) 📝 `x/mlxrunner/cache/cache.go` (+106 -5) 📝 `x/mlxrunner/cache/cache_test.go` (+22 -10) 📝 `x/mlxrunner/cache/recurrent.go` (+21 -23) 📝 `x/mlxrunner/cache/recurrent_test.go` (+201 -4) ➕ `x/mlxrunner/cache/rotating_attention_test.go` (+237 -0) 📝 `x/mlxrunner/cache/rotating_multiturn_test.go` (+2 -2) 📝 `x/mlxrunner/mlx/fast.go` (+10 -35) 📝 `x/mlxrunner/mlx/gated_delta.go` (+27 -3) 📝 `x/mlxrunner/mlx/ops_extra.go` (+8 -34) 📝 `x/mlxrunner/model/base/base.go` (+2 -1) 📝 `x/mlxrunner/pipeline.go` (+16 -6) 📝 `x/models/gemma3/gemma3.go` (+19 -18) 📝 `x/models/gemma4/gemma4.go` (+70 -135) 📝 `x/models/glm4_moe_lite/glm4_moe_lite.go` (+28 -24) 📝 `x/models/llama/llama.go` (+19 -18) ➕ `x/models/nn/recurrent.go` (+259 -0) ➕ `x/models/nn/recurrent_test.go` (+340 -0) ➕ `x/models/nn/sdpa.go` (+578 -0) ➕ `x/models/nn/sdpa_test.go` (+680 -0) _...and 2 more files_ </details> ### 📄 Description Three preparatory changes for multi-sequence batching in mlxrunner. The runner still feeds one sequence at a time after this PR, but the model contract, RoPE, and SDPA/cache abstractions stop assuming B=1 with a single starting position. - Wrap model forward inputs in a Batch struct — one extension point for per-row context (positions, sequence IDs, masks) instead of churning every model's Forward signature. - Apply RoPE at per-row positions — switch from the scalar-offset RoPE kernel to the array-offset one so each row can start at its own position. The pipeline tracks the position locally and passes it through Batch.SeqOffsets. - Decouple models from attention cache storage layout — models describe a logical AttentionMask and receive a per-layer KVHistory carrying K, V, and a MaskApplier. SDPA composes the model mask, query padding, and the cache's storage restrictions, still hitting the kernel's causal/no-mask fast paths when nothing forces a tensor mask. Also unblocks cache variants like paged attention. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
GiteaMirror added the pull-request label 2026-04-29 16:58:23 -05:00
Sign in to join this conversation.
No Branch/Tag Specified
Branches Tags
Labels
Clear labels
amd
api
app
bug
build
cli
cloud
compatibility
context-length
create
docker
documentation
embeddings
feature request
feedback wanted
good first issue
gpt-oss
gpu
harmony
help wanted
image
install
intel
js
launch
linux
macos
memory
mlx
model
needs more info
networking
nvidia
ollama.com
performance
pull-request
Mirrored from GitHub Pull Request
 python
question
registry
rendering
thinking
tools
top
vulkan
windows
wsl
No Label pull-request
Milestone
No items
No Milestone
Projects
Clear projects
No project
Assignees
Clear assignees
GiteaMirror ninjasurge
No Assignees
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: github-starred/ollama#62014
Delete Branch "%!s()"

Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?