done

tests/models/test_models.py

@@ -26,7 +26,7 @@
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["float"])
 @pytest.mark.parametrize("max_tokens", [128])
-@pytest.mark.parametrize("enforce_eager", [False])
+@pytest.mark.parametrize("enforce_eager", [False, True])
 def test_models(
     hf_runner,
     vllm_runner,

vllm/engine/llm_engine.py

@@ -798,7 +798,6 @@ def _process_model_outputs(
         # Log stats.
         if self.log_stats:
             self.stat_logger.log(self._get_stats(scheduler_outputs))
-        # breakpoint()
         return request_outputs
 
     def step(self) -> List[RequestOutput]:

vllm/model_executor/layers/attention.py

@@ -141,29 +141,17 @@ def forward(
         # If key_cache and value_cache are not provided, the new key and value
         # vectors will not be cached. This happens during the initial memory
         # profiling run.
-        num_valid_tokens = input_metadata.num_valid_tokens
-        if (num_valid_tokens > 0 and key_cache is not None
-                and value_cache is not None):
-            key_to_cache = key[:num_valid_tokens]
-            value_to_cache = value[:num_valid_tokens]
+        if (key_cache is not None and value_cache is not None):
             cache_ops.reshape_and_cache(
-                key_to_cache,
-                value_to_cache,
+                key,
+                value,
                 key_cache,
                 value_cache,
                 input_metadata.slot_mapping.flatten(),
                 input_metadata.kv_cache_dtype,
             )
 
-        num_prompt_tokens = input_metadata.num_prompt_tokens
-        num_generation_tokens = input_metadata.num_generation_tokens
-        print(num_generation_tokens)
-
-        if num_prompt_tokens > 0:
-            assert num_generation_tokens == 0
-            query = query[:num_prompt_tokens]
-            key = key[:num_prompt_tokens]
-            value = value[:num_prompt_tokens]
+        if input_metadata.is_prompt:
             # normal attention
             if (key_cache is None or value_cache is None
                     or input_metadata.block_tables.numel() == 0):
@@ -202,7 +190,7 @@ def forward(
                             input_metadata)
 
                 if self.use_ref_attention:
-                    output[:num_prompt_tokens] = self.ref_masked_attention(
+                    output = self.ref_masked_attention(
                         query,
                         key,
                         value,
@@ -222,18 +210,17 @@ def forward(
                     key = key.unflatten(0, (num_tokens))
                     value = value.unflatten(0, (num_tokens))
 
-                output[:
-                       num_prompt_tokens] = xops.memory_efficient_attention_forward(
-                           query,
-                           key,
-                           value,
-                           attn_bias=input_metadata.attn_bias,
-                           p=0.0,
-                           scale=self.scale,
-                           op=xops.fmha.
-                           MemoryEfficientAttentionFlashAttentionOp[0] if
-                           (is_hip()) else None,
-                       ).view_as(query)
+                out = xops.memory_efficient_attention_forward(
+                    query,
+                    key,
+                    value,
+                    attn_bias=input_metadata.attn_bias,
+                    p=0.0,
+                    scale=self.scale,
+                    op=xops.fmha.MemoryEfficientAttentionFlashAttentionOp[0] if
+                    (is_hip()) else None,
+                )
+                output = out.view_as(query)
             else:
                 # prefix-enabled attention
                 output = torch.empty_like(query)
@@ -252,13 +239,11 @@ def forward(
                     getattr(self, "alibi_slopes", None),
                 )
 
-        if num_generation_tokens > 0:
-            breakpoint()
-            assert num_prompt_tokens == 0
+        else:
             # Decoding run.
             output = _paged_attention(
-                output[num_prompt_tokens:num_valid_tokens],
-                query[num_prompt_tokens:num_valid_tokens],
+                output,
+                query,
                 key_cache,
                 value_cache,
                 input_metadata,

vllm/worker/model_runner.py

@@ -611,7 +611,6 @@ def execute_model(
             kv_caches=kv_caches,
             input_metadata=input_metadata,
         )
-        breakpoint()
 
         # Sample the next token.
         output = self.model.sample(
@@ -878,7 +877,6 @@ def forward(
                                                  non_blocking=True)
         self.input_buffers["block_tables"].copy_(input_metadata.block_tables,
                                                  non_blocking=True)
-        breakpoint()
         # Run the graph.
         self.graph.replay()