Update class names to MultiHeadLatentAttention

src/__init__.py

@@ -5,7 +5,7 @@ Copyright (c) 2025
 Implementation of the Multi-Latent Attention mechanism from the DeepSeek-V2 paper.
 """
 
-from .mla import MultiLatentAttention, precompute_freqs_cis, reshape_for_broadcast, apply_rotary_emb
+from .mla import MultiHeadLatentAttention, precompute_freqs_cis, reshape_for_broadcast, apply_rotary_emb
 
 __version__ = "0.1.0"
-__all__ = ["MultiLatentAttention", "precompute_freqs_cis", "reshape_for_broadcast","apply_rotary_emb"]
+__all__ = ["MultiHeadLatentAttention", "precompute_freqs_cis", "reshape_for_broadcast","apply_rotary_emb"]

src/mla.py

@@ -58,7 +58,7 @@ def apply_rotary_emb(
 
 
 
-class MultiLatentAttention(nn.Module):
+class MultiHeadLatentAttention(nn.Module):
     """
         Multi-Head Latent Attention(MLA) Module As in DeepSeek_V2 pape
         Key innovation from standard MHA:

src/tests/test_mla.py

@@ -1,6 +1,6 @@
 import unittest
 import torch
-from ..mla import MultiLatentAttention  # Using relative import
+from ..mla import MultiHeadLatentAttention  # Using relative import
 
 class TestMultiLatentAttention(unittest.TestCase):
     def setUp(self):
@@ -15,7 +15,7 @@ class TestMultiLatentAttention(unittest.TestCase):
         self.seq_len = 10
         
         # Initialize MLA
-        self.mla = MultiLatentAttention(
+        self.mla = MultiHeadLatentAttention(
             d_model=self.d_model,
             num_head=self.num_head,
             d_embed=self.d_embed,