fix server 2 server net data error (#51)

Co-authored-by: JesusmiCaH <1010851196jch@gmail.com>

Author: xiongxu1998

benchmarks/benchmark_speculative_decoding.py

@@ -99,7 +99,8 @@ def benchmark_inference(process_idx, args, result_pipe):
 
     result = ""
     start_time = perf_counter()
-    result = model.generate(input_ids=input_ids, drafter=drafter)
+    max_new_tokens = getattr(args, 'seq_len', 128)
+    result = model.generate(input_ids=input_ids, drafter=drafter, max_new_tokens=max_new_tokens)
     time = perf_counter() - start_time
     generated_tokens_nums = []
     for i in range(batch_size):

src/bloombee/client/inference_session.py

@@ -175,13 +175,14 @@ def _infer_batch_dim(value) -> int:
         input_tensors, args_structure = pack_args_kwargs(
             inputs, 
             normalize_arg(keep_indices),
-            normalize_arg(torch.tensor(1 if need_pruning else 0)),
-            prompts, hypo_ids, 
+            normalize_arg(torch.tensor(1 if need_pruning else 0)), 
             normalize_arg(tree_attention_mask),
             normalize_arg(kv_cache_position_ids),
             normalize_arg(draft_tokens),
             normalize_arg(prefill_length),
             normalize_arg(torch.tensor(1 if is_spec_dec else 0)),
+            prompts, 
+            hypo_ids,
         )
         logger.debug(f"_ServerInferenceSession  step id {step_id}")
         request_metadata = dict(session_id=self.session_id, step_id=step_id)
@@ -199,12 +200,10 @@ def _infer_batch_dim(value) -> int:
                 request_metadata["start_from_position"] = self._position
         # Enable server-to-server communication to trigger CROSS_GPU_TRANSFER
         # Speculative decoding keeps strict full-batch semantics; avoid cross-stage push.
-        if self.config.use_server_to_server and not is_spec_dec:
+        if self.config.use_server_to_server:
             next_servers = self._collect_next_servers()
             if next_servers:
                 request_metadata["next_servers"] = next_servers
-        elif is_spec_dec:
-            request_metadata["disable_cross_stage_push"] = 1
 
         request_metadata["args_structure"] = args_structure
 
@@ -495,7 +494,7 @@ def step(   # 执行一次推理步骤，处理输入数据和相应的提示与
                     # 🔍 CLIENT DEBUG: Log server span processing start
                     span_start_time = time.perf_counter()
 
-                    inputs, keep_indices, need_pruning_next = server_session.step( 
+                    inputs, keep_indices, *_ = server_session.step( 
                         inputs,
                         prompts[server_session.span.start : server_session.span.end],
                         hypo_ids,
@@ -516,7 +515,7 @@ def step(   # 执行一次推理步骤，处理输入数据和相应的提示与
                     # 🔍 CLIENT DEBUG: Log server span processing end
                     span_end_time = time.perf_counter()
                     span_duration = (span_end_time - span_start_time) * 1000  # ms
-                    logger.debug(f"[CLIENT_SERVER_END] ServerIdx={server_idx} | Blocks={server_session.span.start}:{server_session.span.end} | Duration={span_duration:.2f}ms")
+                    logger.info(f"[CLIENT_SERVER_END] ServerIdx={server_idx} | Blocks={server_session.span.start}:{server_session.span.end} | Duration={span_duration:.2f}ms")
                     # print('inputs ', inputs)
                     # print('inputs.shape ', inputs.shape)
                     server_idx += 1
@@ -551,7 +550,7 @@ def step(   # 执行一次推理步骤，处理输入数据和相应的提示与
         # 🔍 CLIENT DEBUG: Log inference step end
         inference_step_end = time.perf_counter()
         inference_step_duration = (inference_step_end - inference_step_start) * 1000  # ms
-        logger.debug(f"[CLIENT_INFERENCE_END] Position={self._position} | Duration={inference_step_duration:.2f}ms | Servers={server_idx}")
+        logger.info(f"[CLIENT_INFERENCE_END] Position={self._position} | Duration={inference_step_duration:.2f}ms | Servers={server_idx}")
 
         outputs = outputs.to(device=inputs_device, dtype=inputs_dtype) 
         # print('client inference session outputs ', outputs.shape)

src/bloombee/models/llama/speculative_model.py

@@ -35,11 +35,11 @@ def generate(
         logits_processor: Optional[LogitsProcessorList] = None,
         stopping_criteria: Optional[StoppingCriteriaList] = None,
         streamer: Optional["BaseStreamer"] = None,
-        beam_width: int = 2,
+        beam_width: int = 1,
         max_tree_depth: int = 4,
         use_kv_cache: bool = True,
         kv_cache_window: int = 2048,
-        max_new_tokens: int = 64,
+        max_new_tokens: int = 128,
         session_max_length: Optional[int] = None,
         **model_kwargs,
     ) -> torch.LongTensor:
@@ -132,6 +132,8 @@ def _sample_with_session(
         initial_len = input_ids.shape[1]
         t0 = time.perf_counter()  # 用于记录第一个达标的时间
         has_printed_first_reach = False # 确保只打印一次
+        sample_finish_times = [None] * batch_size
+        sample_finished = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
         while not finished and (seq_lengths.min().item() - initial_len) < max_new_tokens:
             # 1. Build speculative trees using SSM - 传入 seq_lengths
             t1 = time.perf_counter()
@@ -207,16 +209,24 @@ def _sample_with_session(
             finished = unfinished_sequences.max() == 0
             total_time = time.perf_counter() - t1
             logger.info(f"Step {step_idx}: FTotal Time Elapsed={total_time:.4f} seconds")
-            step_idx += 1
             current_generations = seq_lengths - initial_len
-            if not has_printed_first_reach and current_generations.max().item() >= max_new_tokens:
-                first_reach_time = time.perf_counter() - t0
-                logger.info(f"🚀 [First Reach] 第一个样本达到 max_new_tokens，耗时: {first_reach_time:.4f}s")
-                has_printed_first_reach = True
-            
+            for i in range(batch_size):
+                if (current_generations[i] >= max_new_tokens and not sample_finished[i]):
+                    finish_time = time.perf_counter() - t0
+                    sample_finish_times[i] = finish_time
+                    sample_finished[i] = True
+                    logger.info(f"step {step_idx} Sample {i} finished generation ({max_new_tokens} tokens) at {finish_time:.4f}s")
+            step_idx += 1
 
         if streamer is not None:
             streamer.end()
+            
+        logger.info("====== Batch Generation Summary ======")
+        for i, t in enumerate(sample_finish_times):
+            if t is not None:
+                logger.info(f"Sample {i}: finished at {t:.4f}s")
+            else:
+                logger.info(f"Sample {i}: did not reach max_new_tokens")
 
         return current_input_ids
 

src/bloombee/server/backend.py

@@ -369,7 +369,7 @@ def _flag_to_bool(value) -> bool:
                     position_ids = self._position_ids_cache[cache_key] + (cache_len + offset)
                     if self._is_spec_decoding:
                         rotary_position_ids = self._create_tree_position_ids_with_invalid_cache(
-                            width=2,
+                            width=1,
                             depth=4,
                             prefill_length=inference_info.prefill_length - 1,
                             kv_cache_position_ids=kv_cache_position_ids,
@@ -468,14 +468,10 @@ def _flag_to_bool(value) -> bool:
                     self.pruner_manager.train_lm_head(middle_norm_hidden_states, norm_hidden_states)
 
                 if not training_mode and self._is_spec_decoding and self._need_pruning and self._is_last_block:
-                    t6 = time.perf_counter()
                     norm_hidden_states = self.module.rms_norm(output_hidden_states)
                     keep_indices = self.prune_draft_tree(norm_hidden_states, inference_info.draft_tokens, full_mask)
                     keep_indices = keep_indices
 
-                    t7 = time.perf_counter()
-                    logger.info(f"prune_draft_tree spend: {t7 - t6}")
-                    
                 if not training_mode and self._is_spec_decoding and self._is_last_block:
                     original_hidden_states = output_hidden_states
                     batch_size, seq_len, hidden_size = original_hidden_states.shape

src/bloombee/server/block_functions.py

@@ -41,6 +41,7 @@
     MBPIPE_SCHEMA_PREFIX,
 )
 from bloombee.utils.debug import dprint
+import traceback
 
 # [MBPIPE] Cross-stage streaming push support
 _cross_stage_push_callback = None  # Will be set by handler for cross-stage streaming
@@ -424,6 +425,27 @@ def restore_hidden_states(
 
     return restored_hidden_states
 
+def ensure_tensors(flat_tensors):
+    result = []
+    for i, t in enumerate(flat_tensors):
+        if t is None:
+            result.append(torch.tensor(0))
+        elif isinstance(t, torch.Tensor):
+            result.append(t)
+        elif isinstance(t, (list, tuple)):
+            t_clean = [x for x in t if x is not None]
+            if len(t_clean) == 0:
+                result.append(torch.tensor(0))
+            elif isinstance(t_clean[0], torch.Tensor):
+                result.append(torch.stack(t_clean))
+            else:
+                result.append(torch.tensor(t_clean))
+        elif isinstance(t, (int, float, bool)):
+            result.append(torch.tensor(t))
+        else:
+            raise TypeError(f"flat_tensors[{i}] cant trans to tensor: type={type(t)}, value={t}")
+    return tuple(result)
+
 async def iterate_rpc_inference(
     requested_uids: Sequence[ExpertUID],
     requested_backends: Sequence[TransformerBackend],
@@ -1197,7 +1219,7 @@ async def iterate_rpc_inference(
         if args_structure is not None:
             flat_tensors, kwargs = unpack_args_kwargs(flat_tensors, args_structure)
 
-        hidden_states, keep_indices, need_pruning1, prompts, hypo_ids, tree_attention_mask, kv_cache_position_ids, draft_tokens, prefill_length, is_spec_dec1, *_ = flat_tensors
+        hidden_states, keep_indices, need_pruning1, tree_attention_mask, kv_cache_position_ids, draft_tokens, prefill_length, is_spec_dec1, prompts, hypo_ids, *_ = flat_tensors
         draft_tokens = draft_tokens if draft_tokens is not None and not is_dummy(draft_tokens) else None
 
         # Fix for bus error in cross-machine setups: ensure tensors are contiguous
@@ -1229,11 +1251,7 @@ async def iterate_rpc_inference(
             )
         if not need_pruning and _as_python_bool(step_metadata.get("need_pruning", 0)):
             need_pruning = True
-            
-        # logger.info(f"hidden_states: {hidden_states.shape}")
-        # logger.info(f"keep_indices: {keep_indices.shape}")
-        # logger.info(f"draft_tokens: {draft_tokens.shape}")
-            
+
         if is_spec_dec and draft_tokens is not None and draft_tokens.shape[0] != hidden_states.shape[0]:
             hidden_states = restore_hidden_states(hidden_states, keep_indices, draft_tokens.shape[-1])
 
@@ -1845,10 +1863,14 @@ async def process_microbatch(mb_idx: int, mb_start: int, mb_end: int):
 
         serialize_start = perf_counter()
         need_pruning_next = torch.tensor(0)
+        
+        flat_tensors = (hidden_states, keep_indices, need_pruning_next, tree_attention_mask, kv_cache_position_ids, draft_tokens)
+        flat_tensors = ensure_tensors(flat_tensors)
         output_tensors = [
             serialize_torch_tensor(result.to(proto.dtype), proto.compression, allow_inplace=True)
-            for result, proto in zip((hidden_states, keep_indices, need_pruning_next), nested_flatten(requested_backends[-1].outputs_schema))
+            for result, proto in zip(flat_tensors, nested_flatten(requested_backends[-1].outputs_schema))
         ]
+
         serialize_end = perf_counter()
         serialize_time = (serialize_end - serialize_start) * 1000  # ms
         # print('after serialize and send last layer outputs ', )

src/bloombee/server/handler.py

@@ -1480,15 +1480,16 @@ async def _push_outputs(
 
 
             normalized_outputs = self._normalize_serialized_tensors(serialized_outputs)
-            next_tensors = normalized_outputs + list(request.tensors[3:])
+            next_tensors_data = normalized_outputs + list(request.tensors[6:])
+            next_tensors = serialized_outputs + request.tensors[6:]
 
             next_metadata = metadata.copy()
             next_metadata.update(session_id=next_session_id, next_servers=next_servers[1:], pushed=True)
             sender_send_us = self._now_us()
             next_metadata["clock_sync_sender_send_us"] = sender_send_us
 
             stub = self.get_stub(self._p2p, next_peer_id)
-            push_tensor_bytes = sum(len(t.buffer) for t in next_tensors)
+            push_tensor_bytes = sum(len(t.buffer) for t in next_tensors_data)
             serialized_next_metadata = MSGPackSerializer.dumps(next_metadata)
             push_metadata_bytes = len(serialized_next_metadata)
 

src/bloombee/server/server.py

@@ -330,7 +330,7 @@ def __init__(
         # Create configuration
         config = PruningConfig(
             method=PruningMethod.ADAPTIVE_NEURAL,
-            neural_threshold=0.5,
+            neural_threshold=0.9,
             simple_threshold=0.1
         )
 
@@ -716,6 +716,15 @@ def create(
                             dtype=torch.int64,
                             compression=compression
                         ),
+                        BatchTensorDescriptor(
+                            1, 64, 64, dtype=torch.bool
+                        ), # tree_attention_mask
+                        BatchTensorDescriptor(
+                            1, 128, dtype=torch.int64
+                        ), #  kv_cache_position_ids
+                        BatchTensorDescriptor(
+                            1, 128, dtype=torch_dtype
+                        ), # draft_tokens
                     ),
                     min_batch_size=min_batch_size,
                     max_batch_size=max_batch_size,