Move context tensor padding into dedicated method

Author: tdene

megatron/core/inference/contexts/dynamic_context.py

@@ -1006,6 +1006,21 @@ def build_active_slices(self, batch_size: int):
         graph_scratch_space = torch.cumsum(self.active_request_query_lengths[:batch_size], dim=0)
         self.active_request_last_token_idxs[:batch_size].copy_(graph_scratch_space - 1)
 
+    def pad_active_slices(self):
+        """Pad the active slices of specific tensors."""
+        # Some tensors need to be padded at the token level.
+        padding_token_slice = slice(self.active_token_count, self.padded_active_token_count)
+
+        self.token_to_block_idx[padding_token_slice] = self.kv_block_allocator.dummy_block_idx
+        self.token_to_local_position_within_kv_block[padding_token_slice] = 0
+        self.token_to_position_in_request[padding_token_slice] = 0
+
+        # Other tensors need to be padded at the request level.
+        padding_request_slice = slice(
+            self.total_request_count - self.paused_request_count,
+            self.padded_active_request_count,
+        )
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
@@ -1620,23 +1635,6 @@ def initialize_attention_state(
                 prefill_req_count=padded_prefill_req_count,
                 decode_req_count=padded_decode_req_count,
             )
-        self.padded_active_token_count = self.padded_batch_dimensions.token_count
-        self.padded_active_request_count = self.padded_batch_dimensions.req_count
-        self.padding_slice = slice(self.active_token_count, self.padded_active_token_count)
-
-        self.build_active_slices(self.padded_active_request_count)
-        batch_size = self.total_request_count - self.paused_request_count
-
-        # Update token position indexes.
-        self.token_to_block_idx[self.active_token_count : self.padded_active_token_count] = (
-            self.kv_block_allocator.dummy_block_idx
-        )
-        self.token_to_local_position_within_kv_block[
-            self.active_token_count : self.padded_active_token_count
-        ] = 0
-        self.token_to_position_in_request[
-            self.active_token_count : self.padded_active_token_count
-        ] = 0
 
         self.active_attn_metadata = (
             self.graph_attn_metadata  # type: ignore[assignment]
@@ -1657,6 +1655,14 @@ def initialize_attention_state(
                     decode_req_count=adjusted_decode_req_count,
                 )
 
+        self.padded_active_token_count = self.padded_batch_dimensions.token_count
+        self.padded_active_request_count = self.padded_batch_dimensions.req_count
+        self.padding_slice = slice(self.active_token_count, self.padded_active_token_count)
+
+        self.build_active_slices(self.padded_active_request_count)
+        self.pad_active_slices()
+
+        batch_size = self.total_request_count - self.paused_request_count
         assert self.active_attn_metadata is not None
         self.active_attn_metadata["mha_metadata"].update(
             request_query_lengths=self.active_request_query_lengths[:batch_size],