address comments

Signed-off-by: Congcong Chen <congcongchen@microsoft.com>

Author: congcongchen123

vllm/attention/backends/differential_flash_attn.py

@@ -673,9 +673,7 @@ def __init__(
             differential_flash_attention_config = {}
         self.differential_flash_attention_config = \
             differential_flash_attention_config
-        self.used_shared_kv_cache = \
-            self.differential_flash_attention_config.get(
-                "used_shared_kv_cache", False)
+        self.used_shared_kv_cache = kv_sharing_target_layer_name is not None
         self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
         if blocksparse_params is not None:
             raise ValueError(

vllm/model_executor/models/phi4flash.py

@@ -147,7 +147,6 @@ def __init__(self,
 
         params = {
             'differential_flash_attention_config': {
-                'used_shared_kv_cache': self.yoco_cross,
                 'lambda_init': self.lambda_init,
                 'lambda_q1': self.lambda_q1,
                 'lambda_k1': self.lambda_k1,
@@ -661,6 +660,8 @@ def forward(
         mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
 
         attn_metadata = get_forward_context().attn_metadata
+        # input_ids and hidden_states isn't a one-to-one mapping in prefill
+        # stage due to YOCO optimization.
         hidden_states = self.model(input_ids, positions, attn_metadata,
                                    mamba_cache_params, intermediate_tensors,
                                    inputs_embeds)