update key and value for the None condition

Signed-off-by: apinge <Tong.Qiu2@amd.com>

Author: apinge

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -729,8 +729,11 @@ def forward(
 
         # decode:extend:prefill
         query = query[:num_actual_tokens]
-        key = key[:num_actual_tokens] if key is not None else key_cache[:num_actual_tokens]
-        value = value[:num_actual_tokens] if value is not None else value_cache[:num_actual_tokens]
+        if key is not None:
+            key = key[:num_actual_tokens]
+        if value is not None:
+            value = value[:num_actual_tokens]
+
 
         output_actual_tokens = output[:num_actual_tokens]