compatible with L40

vllm_hacked/v1/sample/sampler.py

@@ -223,18 +223,16 @@ class Sampler(nn.Module):
 
         assert sampling_metadata.temperature is not None
 
+        print("logits.shape:", logits.shape)
         # Apply temperature.
         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)
+            target_batch = sampling_metadata.top_k.size(0)
+            logits = logits.expand(target_batch, -1).contiguous()
 
         # Apply logits processors that only apply to random sampling
         # (argmax invariant)