Split qkv and lm head to low bit

vllm/model_executor/layers/attention/backends/torch_sdpa.py

@@ -91,14 +91,25 @@ def forward(
                 query = query.movedim(1, query.dim() - 2)
                 key = key.movedim(1, key.dim() - 2)
                 value = value.movedim(1, value.dim() - 2)
-                out = torch.nn.functional.scaled_dot_product_attention(
-                    query, 
-                    key, 
-                    value, 
-                    input_metadata.attn_bias,
-                    0.0, 
-                    is_causal=not self.need_mask,
-                    scale=self.scale).movedim(query.dim() - 2, 1).contiguous()
+
+                out = []
+                block_size = 16
+                query_split = torch.split(query, block_size, dim=1)
+                key_split = torch.split(key, block_size, dim=1)
+                value_split = torch.split(value, block_size, dim=1)
+                for q, k, v in zip(query_split, key_split, value_split):
+                    out_split = torch.nn.functional.scaled_dot_product_attention(
+                        q, k, v, input_metadata.attn_bias, 0.0, is_causal=not self.need_mask, scale=self.scale)
+                    out.append(out_split)
+                out = torch.cat(out, dim=1).movedim(query.dim() - 2, 1).contiguous()
+                # out = torch.nn.functional.scaled_dot_product_attention(
+                #     query, 
+                #     key, 
+                #     value, 
+                #     input_metadata.attn_bias,
+                #     0.0, 
+                #     is_causal=not self.need_mask,
+                #     scale=self.scale).movedim(query.dim() - 2, 1).contiguous()
                 # output = out.view_as(query)
                 # FIXME: half input will generate float output, next ipex release will fix this.
                 output = out.view_as(query).to(query.dtype)

vllm/model_executor/layers/sampler.py

@@ -43,7 +43,8 @@ def __init__(self,
     def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor,
                     embedding_bias: Optional[torch.Tensor]) -> torch.Tensor:
         # Get the logits for the next tokens.
-        logits = torch.matmul(hidden_states, embedding.t())
+        # logits = torch.matmul(hidden_states, embedding.t())
+        logits = embedding(hidden_states)
         if embedding_bias is not None:
             logits += embedding_bias
         logits = tensor_model_parallel_gather(logits)

vllm/model_executor/models/llama.py

@@ -343,7 +343,7 @@ def sample(
         hidden_states: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> Optional[SamplerOutput]:
-        next_tokens = self.sampler(self.lm_head.weight, hidden_states,
+        next_tokens = self.sampler(self.lm_head, hidden_states,
                                    sampling_metadata)
         return next_tokens
 

vllm/worker/model_runner.py

@@ -577,6 +577,9 @@ def execute_model(
             model_executable = self.graph_runners[graph_batch_size]
         else:
             model_executable = self.model
+        print(input_tokens.shape)
+        import time
+        start = time.time()
         hidden_states = model_executable(
             input_ids=input_tokens,
             positions=input_positions,
@@ -589,6 +592,8 @@ def execute_model(
             hidden_states=hidden_states,
             sampling_metadata=sampling_metadata,
         )
+        end = time.time()
+        print("Time used: ", (end - start)*1000)
         return output
 
     @torch.inference_mode()