compatible with L40

levo_inference.py

@@ -47,7 +47,8 @@ class LeVoInference(torch.nn.Module):
             tensor_parallel_size=self.cfg.vllm.device_num,
             enforce_eager=True,
             dtype="bfloat16",  
-            gpu_memory_utilization=self.cfg.vllm.gpu_memory_utilization, 
+            gpu_memory_utilization=self.cfg.vllm.gpu_memory_utilization,
+            max_num_seqs=32,
             tokenizer=None,
             skip_tokenizer_init=True,
             enable_prompt_embeds=True,

vllm_hacked/v1/sample/sampler.py

@@ -103,19 +103,6 @@ class Sampler(nn.Module):
             
         #     logits = scores_processed
         #     # logits = logits_processor_stage1(logits)
-        print("logits.shape:", logits.shape)
-        if logits.dim() == 1 or (logits.size(0) != sampling_metadata.top_k.size(0)):
-            batch_size = sampling_metadata.top_k.size(0)
-            # 1. 构造假的采样结果 (batch_size,)
-            sampled = torch.zeros((batch_size,), device=logits.device, dtype=torch.long)
-            # 2. 构造假的 logprobs (vllm v1 期望 [batch_size, num_logprobs] 或 None)
-            # 这里给 None 通常是最安全的，除非你的配置强制要求返回 logprobs
-            processed_logprobs = None
-            # 3. 返回符合接口要求的 SamplerOutput 对象
-            return SamplerOutput(
-                sampled_token_ids=sampled,
-                logprobs=processed_logprobs
-            )
                 
         '''单条推理CFG'''
         # if logits.shape[0] == 3:
@@ -240,6 +227,15 @@ class Sampler(nn.Module):
         logits = self.apply_temperature(logits, sampling_metadata.temperature,
                                         sampling_metadata.all_random)
 
+        # patch
+        print("logits.shape:", logits.shape)
+        if logits.dim() == 1:
+            logits = logits.unsqueeze(0)
+        # 如果 batch 对不上（比如 logits 变 [1, 4194560] 这种平铺形状）
+        if logits.size(0) != sampling_metadata.top_k.size(0):
+            # 强制 reshape 回 [Batch, -1]
+            logits = logits.view(sampling_metadata.top_k.size(0), -1)
+
         # Apply logits processors that only apply to random sampling
         # (argmax invariant)
         for processor in sampling_metadata.logitsprocs.argmax_invariant: