[Fix] Prevent multimodal feature offload races

[NewbieOrange]

Motivation

Multimodal features are asynchronously copied from CUDA to CPU after embedding so chunked prefill can retain a fallback without holding GPU memory. The previous implementation released the CUDA source and exposed the CPU destination without tracking completion.

This created two races:

• Queued encoder reads or the device-to-host copy could access a GPU allocation after the allocator had recycled it.
• A later cache miss or prefill chunk could read the CPU backup before the asynchronous copy completed, producing zero-filled or corrupted image and video features.

The DP ViT path made the issue visible because it reads selected feature rows later and across multiple encoder ranks. The fix is generic multimodal scheduling and does not change model math.

Modifications

• Record the consuming CUDA stream before starting non-blocking multimodal feature and precomputed-embedding offloads.
• Maintain one reusable host-offload completion event per worker and record it after the queued copies.
• Synchronize that event before reading CPU precomputed embeddings.
• Make CPU feature uploads wait for pending offload completion, and make deferred encoders wait before reading CPU feature subsets.
• Add test/registered/unit/managers/test_mm_feature_offload_race.py with deterministic pending-offload regressions for precomputed embeddings and encoder feature reads.

Accuracy Tests

The new regression test delays an offload completion while a later consumer reads the same feature:

• Before this change, both tests fail on 367e370. The consumers observe pre-offload zeros instead of the completed value 7.0.
• With this change, both tests pass and observe the completed value after the synchronization.
• Focused suite result: 9 tests passed, including both new race regressions and the existing multimodal item-splitting tests.
• The fix only restores feature-lifetime ordering and introduces no numerical model changes. Full serving accuracy was not rerun.

Test file: test/registered/unit/managers/test_mm_feature_offload_race.py

Speed Tests and Profiling

No throughput regression is expected. The change adds no kernels and does not synchronize during normal feature use. It only waits when a later consumer would otherwise read a pending CPU offload or asks a deferred encoder to read CPU feature data.

No standalone throughput benchmark was run because this is an asynchronous lifetime correctness fix rather than a compute-path change.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

[NewbieOrange]

Just saw its already fixed in #40621, closing. agent mistakenly took the fixed patch and submitted, sorry. :(