re-enable avoid torch slice fix when chunked prefill is disabled (ope…

…ndatahub-io#209)
prarit · Sep 26, 2024 · a5d87a1 · a5d87a1
1 parent cc2039c
commit a5d87a1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/attention/backends/rocm_flash_attn.py b/vllm/attention/backends/rocm_flash_attn.py
@@ -573,7 +573,7 @@ def forward(
                 else:
                     out = output
                 ops.paged_attention_rocm(
-                    output[num_prefill_tokens:],
+                    out,
                     exp_sums,
                     max_logits,
                     tmp_output,