fix mtp static bench (#1009)

fix mtp params of static benchmark model_infer.py.

Author: STwangyingrui

test/benchmark/static_inference/model_infer.py

@@ -66,6 +66,7 @@ def overlap_prefill(
     input_ids,
     mem_indexes,
     b_req_idx,
+    b_mtp_index,
     b_seq_len,
     total_token_num,
     b_ready_cache_len,
@@ -76,16 +77,18 @@ def overlap_prefill(
     _0_input_ids = input_ids[: total_token_num // 2]
     _0_mem_indexes = mem_indexes[: total_token_num // 2]
     _0_b_req_idx = b_req_idx[: batch_size // 2]
+    _0_b_mtp_index = b_mtp_index[: batch_size // 2]
     _0_b_seq_len = b_seq_len[: batch_size // 2]
     _o_b_ready_cache_len = b_ready_cache_len[: batch_size // 2]
     micro_batch1 = ModelInput(
         _0_batch_size,
         _0_total_token_num,
         _0_max_len_in_batch,
         _0_input_ids,
-        _0_mem_indexes,
         _0_b_req_idx,
+        _0_b_mtp_index,
         _0_b_seq_len,
+        _0_mem_indexes,
         True,
         _o_b_ready_cache_len,
         {},
@@ -97,6 +100,7 @@ def overlap_prefill(
     _1_input_ids = input_ids[total_token_num // 2 :]
     _1_mem_indexes = mem_indexes[total_token_num // 2 :]
     _1_b_req_idx = b_req_idx[batch_size // 2 :]
+    _1_b_mtp_index = b_mtp_index[batch_size // 2 :]
     _1_b_seq_len = b_seq_len[batch_size // 2 :]
     _1_b_ready_cache_len = b_ready_cache_len[batch_size // 2 :]
 
@@ -105,9 +109,10 @@ def overlap_prefill(
         _1_total_token_num,
         _1_max_len_in_batch,
         _1_input_ids,
-        _1_mem_indexes,
         _1_b_req_idx,
+        _1_b_mtp_index,
         _1_b_seq_len,
+        _1_mem_indexes,
         True,
         _1_b_ready_cache_len,
         {},
@@ -120,23 +125,25 @@ def overlap_prefill(
 
 
 def overlap_decode(
-    model_part, batch_size, max_len_in_batch, input_ids, mem_indexes, b_req_idx, b_seq_len, total_token_num
+    model_part, batch_size, max_len_in_batch, input_ids, mem_indexes, b_req_idx, b_mtp_index, b_seq_len, total_token_num
 ):
     _0_batch_size = batch_size // 2
     _0_total_token_num = total_token_num // 2
     _0_max_len_in_batch = max_len_in_batch
     _0_input_ids = input_ids[: batch_size // 2]
     _0_mem_indexes = mem_indexes[: batch_size // 2]
     _0_b_req_idx = b_req_idx[: batch_size // 2]
+    _0_b_mtp_index = b_mtp_index[: batch_size // 2]
     _0_b_seq_len = b_seq_len[: batch_size // 2]
     micro_batch1 = ModelInput(
         _0_batch_size,
         _0_total_token_num,
         _0_max_len_in_batch,
         _0_input_ids,
-        _0_mem_indexes,
         _0_b_req_idx,
+        _0_b_mtp_index,
         _0_b_seq_len,
+        _0_mem_indexes,
     )
 
     _1_batch_size = batch_size - batch_size // 2
@@ -145,16 +152,18 @@ def overlap_decode(
     _1_input_ids = input_ids[batch_size // 2 :]
     _1_mem_indexes = mem_indexes[batch_size // 2 :]
     _1_b_req_idx = b_req_idx[batch_size // 2 :]
+    _1_b_mtp_index = b_mtp_index[batch_size // 2 :]
     _1_b_seq_len = b_seq_len[batch_size // 2 :]
 
     micro_batch2 = ModelInput(
         _1_batch_size,
         _1_total_token_num,
         _1_max_len_in_batch,
         _1_input_ids,
-        _1_mem_indexes,
         _1_b_req_idx,
+        _1_b_mtp_index,
         _1_b_seq_len,
+        _1_mem_indexes,
     )
 
     output, output1 = model_part.microbatch_overlap_decode(micro_batch1, micro_batch2)
@@ -170,6 +179,7 @@ def prefill(
     input_ids,
     mem_indexes,
     b_req_idx,
+    b_mtp_index,
     b_seq_len,
     total_token_num,
     b_ready_cache_len,
@@ -179,25 +189,30 @@ def prefill(
         total_token_num,
         max_len_in_batch,
         input_ids,
-        mem_indexes,
         b_req_idx,
+        b_mtp_index,
         b_seq_len,
+        mem_indexes,
         is_prefill=True,
         b_ready_cache_len=b_ready_cache_len,
     )
+
     model_output = model_part.forward(model_input)
     return model_output.logits
 
 
-def decode(model_part, batch_size, max_len_in_batch, input_ids, mem_indexes, b_req_idx, b_seq_len, total_token_num):
+def decode(
+    model_part, batch_size, max_len_in_batch, input_ids, mem_indexes, b_req_idx, b_mtp_index, b_seq_len, total_token_num
+):
     model_input = ModelInput(
         batch_size,
         total_token_num,
         max_len_in_batch,
         input_ids,
-        mem_indexes,
         b_req_idx,
+        b_mtp_index,
         b_seq_len,
+        mem_indexes,
         is_prefill=False,
     )
     model_output = model_part.forward(model_input)
@@ -236,6 +251,7 @@ def run_forward_once(
     b_req_idx = torch.tensor(
         [model_part.req_manager.alloc() for _ in range(batch_size)], dtype=torch.int32, device="cuda"
     )
+    b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
     b_seq_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
     b_ready_cache_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
     for i in range(batch_size):
@@ -260,6 +276,7 @@ def run_forward_once(
         test_data,
         mem_indexes,
         b_req_idx,
+        b_mtp_index,
         b_seq_len,
         total_token_num,
         b_ready_cache_len,  # b_ready_cache_len
@@ -288,6 +305,7 @@ def run_forward_once(
                     test_data,
                     mem_indexes,
                     b_req_idx,
+                    b_mtp_index,
                     b_seq_len,
                     total_token_num,
                     b_ready_cache_len,  # b_ready_cache_len
@@ -302,6 +320,7 @@ def run_forward_once(
         torch.cuda.synchronize()
         step_start = time.time()
         total_token_num += batch_size
+        b_mtp_index += 1
         b_seq_len += 1
         mem_indexes = model_part.req_manager.mem_manager.alloc(predict_ids.shape[0]).cuda()
         max_len_in_batch = input_len + i + 1
@@ -312,6 +331,7 @@ def run_forward_once(
             predict_ids.view(-1),
             mem_indexes,
             b_req_idx,
+            b_mtp_index,
             b_seq_len,
             total_token_num,
         )
@@ -325,6 +345,7 @@ def run_forward_once(
                         predict_ids.view(-1),
                         mem_indexes,
                         b_req_idx,
+                        b_mtp_index,
                         b_seq_len,
                         total_token_num,
                     ),