[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(