[Kimi-K3] Fix "wrong grids" crash in DP-sharded vision preprocessing (#35305)

Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>

python/sglang/kernels/ops/diffusion/__init__.py

@@ -345,6 +345,7 @@ for _op, _backend, _target, _caps, _description in _SPECS:
 # then symbol; a new public kernel belongs here and nowhere else.
 # ---------------------------------------------------------------------------
 _EXPORTS: dict[str, str] = {
+    "load_extension_with_recovery": "ext.loader",
     # Normalization: RMSNorm / LayerNorm / GroupNorm and their fused epilogues
     "FLYDSL_NORM_MIN_ALIGNED_DIM": "norm.fused_residual_norm_flydsl",
     "flydsl_fused_residual_norm_scale_shift": "norm.fused_residual_norm_flydsl",

python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from unittest.mock import patch
 
-from sglang.kernels.ops.diffusion.ext.loader import load_extension_with_recovery
+from sglang.kernels.ops.diffusion import load_extension_with_recovery
 
 
 def test_stale_torch_lock_is_removed_before_loading(tmp_path: Path):

python/sglang/srt/models/kimi_k3.py

@@ -3339,6 +3339,8 @@ class KimiK3ForConditionalGeneration(nn.Module):
                 for backend, indices in deferred_by_backend.items():
                     group_items = [selected_items[index] for index in indices]
                     group_configs = [deferred[index] for index in indices]
+                    # Map backend-group positions through the rank-local shard to global grid rows.
+                    global_indices = [image_indices[index] for index in indices]
                     first_config = group_configs[0]
                     if backend == "gpu":
                         from sglang.srt.multimodal.processors.kimi_k25 import (
@@ -3361,7 +3363,7 @@ class KimiK3ForConditionalGeneration(nn.Module):
                                 x, first_config.transparent_bg_config
                             ),
                         )
-                        expected_grids = grid_thws_host[indices]
+                        expected_grids = grid_thws_host[global_indices]
                         if not torch.equal(produced_grids.cpu(), expected_grids):
                             raise ValueError(
                                 "Kimi-K3 deferred GPU preprocessing produced wrong grids"
@@ -3380,7 +3382,8 @@ class KimiK3ForConditionalGeneration(nn.Module):
                         )
 
                     patch_counts = [
-                        int(grid_thws_host[index].prod().item()) for index in indices
+                        int(grid_thws_host[index].prod().item())
+                        for index in global_indices
                     ]
                     if sum(patch_counts) != pixel_values.shape[0]:
                         raise ValueError(

test/registered/unit/models/test_kimi_k3_vision.py

@@ -503,6 +503,7 @@ def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch):
 
 
 def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
+    """A vision-DP owner uses each assigned image's grid when preprocessing."""
     from unittest.mock import patch as mock_patch
 
     from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
@@ -531,13 +532,15 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
             "pad_height": 0,
         },
     )
+    grids = [[1, 1, 1], [1, 1, 2]]
+    patch_counts = [grid[0] * grid[1] * grid[2] for grid in grids]
     items = [
         MultimodalDataItem(
             modality=Modality.IMAGE,
             offsets=[(index, index)],
             feature=torch.full((3, 2, 2), index, dtype=torch.uint8),
             model_specific_data={
-                "image_grid_thw": torch.tensor([[1, 1, 1]]),
+                "image_grid_thw": torch.tensor([grids[index]]),
                 DEFERRED_PREPROCESSING_KEY: deferred_config,
             },
         )
@@ -546,8 +549,12 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
     calls = []
 
     def fake_preprocess(images, resize_configs, *args, **kwargs):
-        calls.append([int(image[0, 0, 0]) for image in images])
+        ids = [int(image[0, 0, 0]) for image in images]
-        return torch.tensor([[float(calls[-1][0]), 0.0]]), torch.tensor([[1, 1, 1]])
+        calls.append(ids)
+        pixel_values = torch.cat(
+            [torch.full(size=(patch_counts[i], 2), fill_value=float(i)) for i in ids]
+        )
+        return pixel_values, torch.tensor([grids[i] for i in ids])
 
     # Configured TP size (the IPC consumer count) comes from the published
     # bags; the live topology is forced through the context's own override.
@@ -566,7 +573,8 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
 
     assert calls == [[1]]
     assert one.dtype == torch.float32
-    assert one.tolist() == [[1.0, 0.0]]
+    assert one.shape == (2, 2)
+    assert (one == 1.0).all()
 
 
 def test_kimi_k3_scheduler_leaves_feature_placement_to_dp_owner():