[DeepSeek v3.2][Bugfix] get_index_k_scale_buffer support cp (#18280)

2026-03-17 09:54:54 -07:00
parent 466ff20e51
commit 17031120b8
4 changed files with 22 additions and 4 deletions
@@ -132,6 +132,10 @@ class MockIndexerMetadata(BaseIndexerMetadata):
        """Return: seq lens for each batch."""
        return torch.tensor(self.seq_lens, dtype=torch.int32, device="cpu")

+    def get_indexer_seq_len(self) -> torch.Tensor:
+        """Return: seq lens for each batch."""
+        return torch.tensor(self.seq_lens, dtype=torch.int32, device=self.device)
+
    def get_nsa_extend_len_cpu(self) -> List[int]:
        """
        Return: extend seq lens for each batch.