diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index b57589fde..efeb37173 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py b/python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py index 563757029..26db5610a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py +++ b/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): diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index f910195d0..5c1034f57 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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( diff --git a/test/registered/unit/models/test_kimi_k3_vision.py b/test/registered/unit/models/test_kimi_k3_vision.py index d934a52ca..22ebdfe9e 100644 --- a/test/registered/unit/models/test_kimi_k3_vision.py +++ b/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]) - return torch.tensor([[float(calls[-1][0]), 0.0]]), torch.tensor([[1, 1, 1]]) + ids = [int(image[0, 0, 0]) for image in images] + 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():