Separate mm encoder scheduling#1015

Merged

gkumbhat merged 73 commits into

torch-spyre:mainfrom

gkumbhat:separate_mm_encoder_scheduling

Jul 13, 2026

gkumbhat commented Jun 23, 2026 •

edited

Loading

Collaborator

Depends on:

Description

This PR adds a mechanism to run vision encoder model for multi-modal models into a separate python process, allowing decoder model processing, i.e prefil and decode to continue in parallel providing significant performance improvements (numbers shared). This approach builds on top of the approach of running vision encoder on separate hardware than spyre and thus enables bigger performance improvements.

Related Issues

Test Plan

Checklist

I have read the contributing guidelines
My code follows the project's code style (run bash format.sh)
I have added tests for my changes (if applicable)
I have updated the documentation (if applicable)
My commits include a Signed-off-by: line (DCO compliance)

gkumbhat and others added 24 commits

June 4, 2026 10:52


          ✨ Deduplicate MM vision encoder TP ranks

d5ff401

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Add inference mode context around MM encoding invocation

8257f13

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          wip improvements

cf94898

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          Merge branch 'main' into remove_redundant_computation

8afe3d0


          🚧 rebase with main and update runner

f8bf367

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          💡 Remove unnecessary commentary

5e2a761

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🎨 Fix ty checks

7def4ff

Signed-off-by: gkumbhat <Gaurav.Kumbhat@ibm.com>


          ✅ Add unit tests for mm shared memory

f2afc5a

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🎨 Fix ruff check failure

960f572

Signed-off-by: gkumbhat <Gaurav.Kumbhat@ibm.com>


          🚧 Add schedular level MM batch processing

d89fa06

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🚧 Change mistral batching to loop operation

c9dc099

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🚧 WIP for phase 2 - separate process for vision encoder

7caa3ee

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Fix SHM conflict between async encoder

4041d95

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          Merge branch 'main' into separate_mm_encoder_scheduling

5c1dc8f


          🐛 Fix encoder process termination

232505d

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🚚 Rename encode to execute model

510ff1d

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Fix nnpa registration

566de61

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 fix casting of mm_dtype

75cb0c4

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Fix the dtype conversion between vision encoder and spyre runner

31e22da

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Fix mm dtype for NNPA

39859ed

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Fix mm dtype for NNPA

3f83962

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>


          🐛 Fix NNPA slow performance caused by use of fused weights and cleanu…

67b645a

…p dtype casting

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>
Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🔊 Logging improvements

3d9a9b4

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          Merge branch 'main' into separate_mm_encoder_scheduling

5c39e29

github-actions Bot commented Jun 23, 2026

👋 Hi! Thank you for contributing.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, run ./format.sh.
Now you are good to go 🚀.

We also recommend installing prek and configuring it to check your code before every local commit.

gkumbhat and others added 5 commits

June 23, 2026 16:46


          Cleanup old implementation and documentations related to that

e8dd6c2

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🎨 Fix formatting

a888b51

Signed-off-by: gkumbhat <Gaurav.Kumbhat@ibm.com>


          🗑️ Remove architecture diagram section from readme

883d20a

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🎨💡 Cleanup and remove unnecessary code

b0caf82

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🎨 Fix formatting

Signed-off-by: gkumbhat <Gaurav.Kumbhat@ibm.com>

gkumbhat and others added 20 commits

June 29, 2026 21:48


          ⚰️📝🚚 Cleanups

18cccb1

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🧪 add some claude-drafted tests for edge cases

acdc4c7

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          ⚗️ try some e2e tests

27997a8

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          🧵 Handle encoder job cancellation

efce073

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🚸 Handle encoder process crash by killing server

1f85989

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🐛 fixup the e2e tests to run on mac

2e43627

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          🧵 Handle DOS and encoder startup failure scenario

af988ab

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🛡️ Handle some edge case scenario around failures of encode reqs

d352d72

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          📦 Update ibm-fms to 1.12.1 with granite-vision fix

b4f831b

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          Merge branch 'separate_mm_encoder_scheduling' into some-mm-edge-cases

d030638


          🔥 rip out other unit tests

8d3a777

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          🚧 try to add granite vision

63ebf1e

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          ♻️ rename

2e82e36

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          🍱 use nano gc model

819918a

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>


          Merge pull request #3 from joerunde/some-mm-edge-cases

188e3fa

🧪 add some claude-drafted tests for edge cases


          🐛 Fix encoder e2e test failure

f83dfd7

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🎨 Fix formatting and typing errors

9832d6c

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          Merge branch 'main' into separate_mm_encoder_scheduling

dcb3a11


          ✅ Fix schedular MM test to declare request pausing vars

967a841

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🎨 Fix formatting and remove old test file

c412334

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

gkumbhat commented Jul 2, 2026

Collaborator Author

bot:test
export SENDNN_INFERENCE_ASYNC_MM_ENCODER=1

gkumbhat and others added 4 commits

July 6, 2026 09:18


          🚩 Enable MM encoding feature by default

a5926a9

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          Merge branch 'main' into separate_mm_encoder_scheduling

b1eec4f


          🚩 Enable inference async MM encoder by default

6c01349

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>


          🧪 Fix two test failures after vLLM v0.24.0 bump

bdecefd

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

gkumbhat force-pushed the separate_mm_encoder_scheduling branch from 4f0b684 to bdecefd Compare

July 9, 2026 00:00

joerunde reviewed

View reviewed changes

joerunde left a comment

Collaborator

lpgtm- @gkumbhat is working with product teams to ensure this meets their needs

joerunde approved these changes

View reviewed changes


          🚩 Disable separate encoding only for power

cfe1eaf

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

gkumbhat merged commit 4a304fb into torch-spyre:main

16 checks passed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Labels

None yet