Merge pull request taco-project#31 from nvidia-china-sae/linhu/quickfix-mla-tp

quick fix for mla + tp

Author: zhuofan1123

csrc/tp_transfer_thread_group.cpp

@@ -70,12 +70,12 @@ void TPTransferThreadGroup::tp_group_transfer(
                     dst_layer_ptrs = static_cast<void**>(gpu_blocks_ + i * num_layers + layer_id);
                     src_layer_ptrs = static_cast<void**>(cpu_blocks_ + layer_id);
                     dst_startoff_inside_chunks = 0;
-                    src_startoff_inside_chunks = i * dst_chunk_size_in_bytes;
+                    src_startoff_inside_chunks = is_mla ? 0 : i * dst_chunk_size_in_bytes;
                     copy_size_in_bytes = dst_chunk_size_in_bytes;
                 } else {
                     dst_layer_ptrs = static_cast<void**>(cpu_blocks_ + layer_id);
                     src_layer_ptrs = static_cast<void**>(gpu_blocks_ + i * num_layers + layer_id);
-                    dst_startoff_inside_chunks = i * src_chunk_size_in_bytes;
+                    dst_startoff_inside_chunks = is_mla ? 0 : i * src_chunk_size_in_bytes;
                     src_startoff_inside_chunks = 0;
                     copy_size_in_bytes = src_chunk_size_in_bytes;
                 }

tests/test_kvmanager.py

@@ -123,7 +123,7 @@ def generate_gpu_blocks(model_config, cache_config, test_config):
         num_layer=num_layers,
         num_block=num_gpu_blocks,
         tokens_per_block=tokens_per_block,
-        num_head=num_kv_heads//tp_size,
+        num_head=num_kv_heads//tp_size if not use_mla else num_kv_heads,
         head_size=head_size,
         is_mla=model_config.use_mla
     )