compatible with L40

Dockerfile

@@ -7,6 +7,11 @@ RUN apt-get update && \
     git lfs install && \
     rm -rf /var/lib/apt/lists/*
 
+COPY ./vllm_hacked/model_executor/models/llama.py /opt/conda/lib/python3.11/site-packages/vllm/model_executor/models/llama.py
+COPY ./vllm_hacked/v1/sample/sampler.py /opt/conda/lib/python3.11/site-packages/vllm/v1/sample/sampler.py
+COPY ./vllm_hacked/v1/sample/metadata.py /opt/conda/lib/python3.11/site-packages/vllm/v1/sample/metadata.py
+COPY ./vllm_hacked/sampling_params.py /opt/conda/lib/python3.11/site-packages/vllm/sampling_params.py
+
 RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"

vllm_hacked/v1/sample/sampler.py

@@ -187,10 +187,9 @@ class Sampler(nn.Module):
         # Avoid division by zero if there are greedy requests.
         if not all_random:
             temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
-        try:
-            return logits.div_(temp.view(-1, 1))
-        except:
-            return logits.div_(temp.unsqueeze(dim=1))
+        if temp.dim() < logits.dim():
+            temp = temp.view([-1] + [1] * (logits.dim() - 1))
+        return logits / temp
 
     def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
         return logits.argmax(dim=-1).view(-1)