[ROCM][AMD][TRITON] Halving warps number for fw_prefill to reduce spilling (vllm-project#12713)

Signed-off-by: Aleksandr Malyshev <maleksan@amd.com>
Co-authored-by: Aleksandr Malyshev <maleksan@amd.com>
Signed-off-by: Louis Ulmer <ulmerlouis@gmail.com>

Author: 2 people

vllm/attention/ops/prefix_prefill.py

@@ -11,7 +11,7 @@
 
 # Static kernels parameters
 BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64
-NUM_WARPS = 8
+NUM_WARPS = 4 if current_platform.is_rocm() else 8
 
 # To check compatibility
 IS_TURING = current_platform.get_device_capability() == (7, 5)