issue/340 - patch the code to prevent errors.

Author: wangpengcheng

csrc/engine/compiler/paged_compiler.cpp

@@ -5,16 +5,7 @@
 namespace infinilm::engine {
 PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
     : GraphCompiler(model, barrier) {
-    for (size_t b = 1; b < 64; ++b) {
-        decode_batch_sizes_.push_back(b);
-    }
-    for (size_t b = 64; b < 128; b += 16) {
-        decode_batch_sizes_.push_back(b);
-    }
-    for (size_t b = 128; b < 256; b += 32) {
-        decode_batch_sizes_.push_back(b);
-    }
-    for (size_t b = 256; b <= 512; b += 64) {
+    for (size_t b = 256; b > 0; b--) {
         decode_batch_sizes_.push_back(b);
     }
 }

csrc/engine/rank_worker.cpp

@@ -386,18 +386,23 @@ void RankWorker::thread_loop() {
                             const auto &vocab_size{logits_shape[2]};
                             const auto &total_len{logits_shape[1]};
                             const auto &batch_size{logits_shape[0]};
+                            int32_t seq_length = static_cast<int32_t>(batch_size * total_len);
 
                             auto n_req = local_args.input_offsets.value()->size(0) - 1;
                             int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data();
+                            ASSERT(input_offsets[n_req] == seq_length);
 
                             auto output_ids{infinicore::Tensor::empty({n_req}, infinicore::DataType::I64, rank_info_.device)};
 
                             for (auto i{decltype(n_req)(0)}; i < n_req; ++i) {
-                                auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, size_t(input_offsets[i + 1] - 1), 1}})->view({vocab_size})};
+                                int32_t score_index = input_offsets[i + 1] - 1;
+                                ASSERT(input_offsets[i + 1] > input_offsets[i]);
+                                ASSERT(score_index >= 0 && score_index < seq_length);
+
+                                auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, size_t(score_index), 1}})->view({vocab_size})};
                                 auto out{output_ids->narrow({{0, i, 1}})->view({})};
                                 float random_val = std::uniform_real_distribution<float>(0, 1)(rng_);
-                                infinicore::op::random_sample_(
-                                    out, score, random_val, top_p, top_k, temperature);
+                                infinicore::op::random_sample_(out, score, random_val, top_p, top_k, temperature);
                             }
 
                             output_ids = output_ids->to(infinicore::Device::cpu());

python/infinilm/llm/scheduler.py

@@ -261,6 +261,8 @@ def schedule(self) -> Optional[SchedulerOutput]:
                 scheduled_requests=scheduled_requests,
                 is_prefill=is_prefill,
             )
+            # logger.info("Scheduled decode: %d", len(scheduled_requests))
+
             if self.connector is not None:
                 meta = self.connector.build_connector_meta()
                 scheduler_output.kv_connector_metadata = meta

python/infinilm/processors/basic_llm_processor.py

@@ -4,6 +4,29 @@
 from ..llm.scheduler import SchedulerOutput
 
 
+def extend_to_next_power_of_two(lst):
+    """Pad ``lst`` to the next power-of-two length with ``-1``.
+
+    Padding marks unused slots (same convention as ``block_tables``).
+    Callers must ``narrow`` to the real length before passing data to kernels.
+
+    Args:
+        lst: Input list of numeric offsets or cumulative lengths.
+
+    Returns:
+        A new list. Empty input yields ``[0]``; already power-of-two yields a copy.
+    """
+    if not lst:
+        return [0]
+    n = len(lst)
+    next_pow = 1
+    while next_pow < n:
+        next_pow <<= 1
+    if next_pow == n:
+        return lst[:]
+    return lst + [-1] * (next_pow - n)
+
+
 @register_processor("default")
 class BasicLLMProcessor(InfinilmProcessor):
     def __init__(self, model_dir_path: str):
@@ -35,9 +58,13 @@ def apply_chat_template(
         normalized_conversation = []
         for message in conversation:
             if isinstance(message["content"], list):
-                assert len(message["content"]) == 1, "Only one content item supported in list"
+                assert len(message["content"]) == 1, (
+                    "Only one content item supported in list"
+                )
                 content_item = message["content"][0]
-                assert "type" in content_item and "text" in content_item, "Content dict must have 'type' and 'text' keys"
+                assert "type" in content_item and "text" in content_item, (
+                    "Content dict must have 'type' and 'text' keys"
+                )
                 normalized_conversation.append(
                     {"role": message["role"], "content": content_item["text"]}
                 )
@@ -229,21 +256,44 @@ def _build_model_input_from_batch_scheduler_output(
             block_tables.append(padded_block_table)
             cu_seqlens.append(cu_seqlens[-1] + seq_len)
 
-        return {
-            "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64),
-            "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64),
-            "past_kv_lengths": infinicore.from_list(
-                cached_lens, dtype=infinicore.int32
-            ),
-            "total_kv_lengths": infinicore.from_list(seq_lens, dtype=infinicore.int32),
-            "input_offsets": infinicore.from_list(seq_offsets, dtype=infinicore.int32),
-            "cu_seqlens": infinicore.from_list(cu_seqlens, dtype=infinicore.int32),
-            "block_tables": infinicore.from_list(block_tables, dtype=infinicore.int32),
-            "slot_mapping": infinicore.from_list(slot_mapping, dtype=infinicore.int64),
+        assert seq_offsets[-1] == len(tokens), (
+            f"seq_offsets[-1]={seq_offsets[-1]} != len(tokens)={len(tokens)}"
+        )
+
+        length = len(seq_offsets)
+        seq_offsets = extend_to_next_power_of_two(seq_offsets)
+        cu_seqlens = extend_to_next_power_of_two(cu_seqlens)
+
+        input_ids = infinicore.from_list([tokens], dtype=infinicore.int64)
+        position_ids = infinicore.from_list(position_ids, dtype=infinicore.int64)
+        past_kv_lengths = infinicore.from_list(cached_lens, dtype=infinicore.int32)
+        total_kv_lengths = infinicore.from_list(seq_lens, dtype=infinicore.int32)
+
+        input_offsets = infinicore.from_list(
+            seq_offsets, dtype=infinicore.int32
+        ).narrow(0, 0, length)
+
+        cu_seqlens = infinicore.from_list(cu_seqlens, dtype=infinicore.int32).narrow(
+            0, 0, length
+        )
+
+        block_tables = infinicore.from_list(block_tables, dtype=infinicore.int32)
+        slot_mapping = infinicore.from_list(slot_mapping, dtype=infinicore.int64)
+
+        return_dict = {
+            "input_ids": input_ids,
+            "position_ids": position_ids,
+            "past_kv_lengths": past_kv_lengths,
+            "total_kv_lengths": total_kv_lengths,
+            "input_offsets": input_offsets,
+            "cu_seqlens": cu_seqlens,
+            "block_tables": block_tables,
+            "slot_mapping": slot_mapping,
             "temperature": temperature,
             "top_k": top_k,
             "top_p": top_p,
         }
+        return return_dict
 
     def get_tokenizer(self):
         return self.tokenizer