diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 114094e7a..0cb937df8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -65,8 +65,8 @@ class GlmImageLayerKVCache: self.k_cache = k self.v_cache = v else: - self.k_cache = torch.cat([self.k_cache, k], dim=2) - self.v_cache = torch.cat([self.v_cache, v], dim=2) + self.k_cache = torch.cat([self.k_cache, k], dim=1) + self.v_cache = torch.cat([self.v_cache, v], dim=1) def get(self): return self.k_cache, self.v_cache