Support piecewise cuda graph for MLA (#11812)

2025-11-10 09:13:48 +08:00
parent 885cfca273
commit db24d34603
10 changed files with 196 additions and 89 deletions
@@ -823,7 +823,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
            query_pass = query[..., self.rotary_dim :]
            key_pass = key[..., self.rotary_dim :]

-        self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device)
        cos_sin = self.cos_sin_cache[
            torch.add(positions, offsets) if offsets is not None else positions
        ]