diff --git a/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py
index 877db9935..7ee33e4a9 100644
--- a/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py
+++ b/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py
@@ -94,7 +94,10 @@ def _patch_cache_dit_similarity():
     cache_manager.CachedContextManager.similarity = patched_similarity
 
 
-def _build_parallelism_config(sp_group, tp_group):
+def _build_parallelism_config(
+    sp_group: Optional[torch.distributed.ProcessGroup],
+    tp_group: Optional[torch.distributed.ProcessGroup],
+):
     if sp_group is None and tp_group is None:
         return None
 
@@ -106,7 +109,7 @@ def _build_parallelism_config(sp_group, tp_group):
 
     tp_size = None
     if tp_group is not None:
-        tp_size = dist.get_world_size(tp_group.device_group)
+        tp_size = dist.get_world_size(tp_group)
 
     return ParallelismConfig(
         backend=ParallelismBackend.NATIVE_PYTORCH,