fix comments

Author: jiayyu

atom/model_engine/model_runner.py

@@ -101,7 +101,7 @@ def clean(self):
 
         self.prev_batch: Optional[ScheduledBatch] = None
 
-        self.pre_num_decode_token_pre_seq = 1
+        self.pre_num_decode_token_per_seq = 1
         self.draft_token_ids: Optional[torch.Tensor] = None
 
     def prepare_sampled_ids(
@@ -152,7 +152,7 @@ def get_prev_alive_locations(self, batch: ScheduledBatch) -> tuple[list[int], in
             i for i, seq_id in enumerate(self.prev_batch.req_ids)
             if seq_id in batch.req_ids
         ]
-        num_deferred_tokens = len(alive_seq_indices) * self.pre_num_decode_token_pre_seq
+        num_deferred_tokens = len(alive_seq_indices) * self.pre_num_decode_token_per_seq
         is_all_alive = len(alive_seq_indices) == len(self.prev_batch.req_ids)
         return alive_seq_indices, num_deferred_tokens, is_all_alive
 
@@ -212,7 +212,7 @@ def prepare_input_ids(
                 self.input_ids.np[:num_norm_tokens] = token_ids
                 self.input_ids.copy_to_gpu(num_norm_tokens)
             # no new requests added and old requests finished
-            if self.draft_token_ids is not None and self.pre_num_decode_token_pre_seq > 1:
+            if self.draft_token_ids is not None and self.pre_num_decode_token_per_seq > 1:
                 alive_prev = self.prev_token_ids[alive_seq_indices]
                 alive_draft = self.draft_token_ids[alive_seq_indices]
                 combined = torch.cat([
@@ -236,7 +236,7 @@ def prepare_input_ids(
             #     self.input_ids_loc.gpu[:num_deferred_tokens],
             #     out=self.input_ids.gpu[:num_deferred_tokens],
             # )
-            if self.draft_token_ids is not None and self.pre_num_decode_token_pre_seq > 1:
+            if self.draft_token_ids is not None and self.pre_num_decode_token_per_seq > 1:
                 alive_prev = self.prev_token_ids[alive_seq_indices]  # (num_alive_seqs,)
                 alive_draft = self.draft_token_ids[alive_seq_indices]  # (num_alive_seqs, mtp_n_grams-1)
                 combined = torch.cat([
@@ -1084,7 +1084,7 @@ def propose_draft_token_ids(
         self.forward_vars["draft_tokens"].gpu[:bs, :self.drafter.mtp_k] = draft_token
         self.forward_vars["draft_tokens"].copy_to_cpu()
         self.tokenID_processor.draft_token_ids = draft_token
-        self.tokenID_processor.pre_num_decode_token_pre_seq = 2
+        self.tokenID_processor.pre_num_decode_token_per_seq = 2
 
         return None
 

atom/model_engine/scheduler.py

@@ -194,19 +194,23 @@ def postprocess(
         # update token_ids with the actual sampled token ids
         finished_seqs = []
         stream_outputs = []
+        num_placeholder = (
+            2 * self.mtp_k if is_deferred_out and self.use_spec else
+            1              if is_deferred_out else
+            self.mtp_k     if self.use_spec else
+            0
+        )
 
         for seq in self.running:
             if seq.id not in prev_token_ids:
                 continue
             token_ids = prev_token_ids[seq.id]
             new_tokens = []
             if is_deferred_out:
-                idx = seq.token_ids.index(self.eos_token_id)
-                seq.token_ids[idx:] = token_ids
+                seq.token_ids[-num_placeholder:] = token_ids
 
                 if seq.output_tokens:
-                    idx = seq.output_tokens.index(self.eos_token_id)
-                    seq.output_tokens[idx:] = token_ids
+                    seq.output_tokens[-num_placeholder:] = token_ids
 
                 else:
                     seq.output_tokens.extend(token_ids)
@@ -256,12 +260,7 @@ def postprocess(
         if stream_output_queue is not None and stream_outputs:
             stream_output_queue.put_nowait(stream_outputs)
 
-        num_placeholder = (
-            2 * self.mtp_k if is_deferred_out and self.use_spec else
-            1              if is_deferred_out else
-            self.mtp_k     if self.use_spec else
-            0
-        )
+
         if num_placeholder > 0:
             # placeholder for the each decode step
             for seq in seqs:

atom/model_ops/attentions/aiter_mla.py

@@ -263,7 +263,7 @@ def prepare_kv_indices():
     def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData:
         var = self.model_runner.forward_vars
         sparse_kv_indptr = var["sparse_kv_indptr"].gpu if self.is_sparse else None
-        max_q_len= 1 if not hasattr(self, "drafter") else var["mtp_k"] + 1
+        max_q_len = 1 if not hasattr(self, "drafter") else var["mtp_k"] + 1
         ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len)
         attn_matadata = AttentionMetaData(
             slot_mapping=var["slot_mapping"].gpu[: bs * max_q_len],