bugfix && variable regularization

Signed-off-by: Apocalypse990923-qshi <[email protected]>

Author: Apocalypse990923-qshi

vllm_ascend/attention/mla_v1.py

@@ -906,31 +906,6 @@ def _process_weights_for_fused_mlapo(self, act_dtype: torch.dtype):
         self.ctkv_scale = torch.tensor([1], dtype=act_dtype, device=device)
         self.q_nope_scale = torch.tensor([1], dtype=act_dtype, device=device)
 
-    def reorder_by_req(self, num_tokens_per_req_per_rank):
-        num_tokens_per_rank_per_req = [
-            list(i) for i in zip(*num_tokens_per_req_per_rank)
-        ]  # transpose to [rank, req]
-        num_ranks = len(num_tokens_per_rank_per_req)
-        num_reqs = len(num_tokens_per_rank_per_req[0])
-        assert all(len(x) == num_reqs for x in num_tokens_per_rank_per_req)
-
-        # calc each rank's start offset
-        offsets = []
-        offset = 0
-        for rank_tokens in num_tokens_per_rank_per_req:
-            offsets.append(offset)
-            offset += sum(rank_tokens)
-
-        reordered = []
-        for req_idx in range(num_reqs):
-            for rank_idx in range(num_ranks):
-                start = offsets[rank_idx] + sum(
-                    num_tokens_per_rank_per_req[rank_idx][:req_idx])
-                end = start + num_tokens_per_rank_per_req[rank_idx][req_idx]
-                reordered.extend(range(start, end))
-
-        return torch.tensor(reordered, dtype=torch.int32)
-
     def extract_req_dcp_by_chunk_cp(self, lst, chunk_idx, fill_value=0):
         num_reqs = len(lst)
 

vllm_ascend/worker/block_table.py

@@ -356,7 +356,7 @@ def get_split_computed_tokens(
         # else:
         #     assert len(request_ids) == num_requests
         assert request_ids is not None and len(request_ids) == num_requests
-        num_computed_tokens_of_cp_dcp = [[[0] * self.dcp_world_size
+        num_computed_tokens_of_pcp_dcp_for_chunk = [[[0] * self.dcp_world_size
                                           for _ in range(self.pcp_world_size)]
                                          for _ in range(num_requests)]
         total_ranks = self.pcp_world_size * self.dcp_world_size
@@ -382,11 +382,11 @@ def get_split_computed_tokens(
                 else:
                     pcp_idx = start_rank // self.dcp_world_size
                     dcp_idx = start_rank % self.dcp_world_size
-                    num_computed_tokens_of_cp_dcp[req_idx][pcp_idx][
+                    num_computed_tokens_of_pcp_dcp_for_chunk[req_idx][pcp_idx][
                         dcp_idx] += consumed_tokens
                     request_start_rank_dict[req_id] = (start_rank,
                                                        tokens_blank)
-                    return num_computed_tokens_of_cp_dcp
+                    return num_computed_tokens_of_pcp_dcp_for_chunk
 
             virtual_size = total_ranks * cp_kv_cache_interleave_size
             base = int(total_tokens) // virtual_size
@@ -397,7 +397,7 @@ def get_split_computed_tokens(
             for rank_idx in range(total_ranks):
                 pcp_idx = rank_idx // self.dcp_world_size
                 dcp_idx = rank_idx % self.dcp_world_size
-                num_computed_tokens_of_cp_dcp[req_idx][pcp_idx][
+                num_computed_tokens_of_pcp_dcp_for_chunk[req_idx][pcp_idx][
                     dcp_idx] = base * cp_kv_cache_interleave_size
 
             # Distribute remainder tokens starting from start_rank
@@ -406,11 +406,11 @@ def get_split_computed_tokens(
                 pcp_idx = rank // self.dcp_world_size
                 dcp_idx = rank % self.dcp_world_size
                 if i < remain_blocks - 1 or remainder % cp_kv_cache_interleave_size == 0:  # not last block or divisible
-                    num_computed_tokens_of_cp_dcp[req_idx][pcp_idx][
+                    num_computed_tokens_of_pcp_dcp_for_chunk[req_idx][pcp_idx][
                         dcp_idx] += 1 * cp_kv_cache_interleave_size
                     tokens_blank = 0
                 else:  # if last block and undivisible
-                    num_computed_tokens_of_cp_dcp[req_idx][pcp_idx][
+                    num_computed_tokens_of_pcp_dcp_for_chunk[req_idx][pcp_idx][
                         dcp_idx] += remainder % cp_kv_cache_interleave_size
                     tokens_blank = cp_kv_cache_interleave_size - (
                         remainder % cp_kv_cache_interleave_size)
@@ -422,7 +422,7 @@ def get_split_computed_tokens(
             if request_start_rank_dict is not None:
                 request_start_rank_dict[req_id] = (start_rank, tokens_blank)
 
-        return num_computed_tokens_of_cp_dcp
+        return num_computed_tokens_of_pcp_dcp_for_chunk
 
     def clear(self) -> None:
         for block_table in self.block_tables:

vllm_ascend/worker/model_runner_v1.py

@@ -1463,7 +1463,6 @@ def _prepare_inputs(
     ) -> tuple[dict[str, Any], torch.Tensor, np.ndarray, int, torch.Tensor,
                int, torch.Tensor, SpecDecodeMetadata, Optional[torch.Tensor],
                Optional[torch.Tensor], Optional[torch.Tensor], int]:
-        # self.slot_mapping.fill_(0)
         total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         assert total_num_scheduled_tokens > 0
         num_reqs = self.input_batch.num_reqs
@@ -1488,13 +1487,13 @@ def _prepare_inputs(
         self.generate_kv_idx(tokens, scheduler_output)
         self.input_batch.block_table.compute_slot_mapping(
             req_indices, positions_np)
+        self.input_batch.block_table.commit_slot_mapping(
+            total_num_scheduled_tokens)
         tokens, position_pcp, pcp_unpad_mask = self._update_tokens_for_pcp(
             tokens)
         num_scheduled_tokens = np.array(tokens, dtype=np.int32)
         # update total_num_scheduled_tokens
         total_num_scheduled_tokens = sum(num_scheduled_tokens[:num_reqs])
-        self.input_batch.block_table.commit_slot_mapping(
-            total_num_scheduled_tokens)
 
         total_num_pcp_pads = sum(self.num_pcp_pads)
         max_num_scheduled_tokens = max(tokens)