[update] MLA stage=2 并启动pipelinev2， 修复pipeline subop 操作类型不一致

examples/mla/triton_decode_attention.py

@@ -37,7 +37,9 @@ import triton.language as tl
 from vllm.platforms import current_platform
 
 is_hip_ = current_platform.is_rocm()
-os.environ["TRITON_HIP_USE_NEW_STREAM_PIPELINE"] = f"0"
+os.environ["TRITON_HIP_USE_NEW_STREAM_PIPELINE"] = f"1"
+os.environ["TRITON_ENABLE_GLOBAL_TO_LOCAL_AND_NUMSTAGE2"]="0"
+os.environ["TRITON_DEFAULT_ENABLE_NUM_VGPRS512"] = "1"
 
 logger = logging.getLogger(__name__)
 
@@ -751,12 +753,12 @@ def decode_attention_v1(
 
 @triton.autotune(
     configs=[
-        triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=2, num_stages=1),
-        triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=4, num_stages=1),
-        triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=2, num_stages=1),
-        triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=4, num_stages=1),
-        triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=2, num_stages=1),
-        triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=4, num_stages=1),
+        triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=2, num_stages=2),
+        triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=2, num_stages=2),
+        triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=2, num_stages=2),
+        triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=4, num_stages=2),
         triton.Config({"BLOCK_N": 128, "BLOCK_DIM":32}, num_warps=2, num_stages=1),
         triton.Config({"BLOCK_N": 128, "BLOCK_DIM":32}, num_warps=4, num_stages=1),
         triton.Config({"BLOCK_N": 256, "BLOCK_DIM":32}, num_warps=2, num_stages=1),

vllm/attention/ops/triton_decode_attention.py

@@ -39,7 +39,9 @@ from vllm import envs
 # from ..backends.triton_config import KERNLE_KINDS
 
 is_hip_ = current_platform.is_rocm()
-os.environ["TRITON_HIP_USE_NEW_STREAM_PIPELINE"] = f"0"
+os.environ["TRITON_HIP_USE_NEW_STREAM_PIPELINE"] = f"1"
+os.environ["TRITON_ENABLE_GLOBAL_TO_LOCAL_AND_NUMSTAGE2"]="0"
+os.environ["TRITON_DEFAULT_ENABLE_NUM_VGPRS512"] = "1"
 
 logger = logging.getLogger(__name__)
 
@@ -1062,12 +1064,12 @@ def decode_attention_v1(
 
 # @triton.autotune(
 #     configs=[ 
-#         triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=2, num_stages=1),
-#         triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=4, num_stages=1),
-#         triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=2, num_stages=1),
-#         triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=4, num_stages=1),
-#         triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=2, num_stages=1),
-#         triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=4, num_stages=1),
+#         triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=2, num_stages=2),
+#         triton.Config({"BLOCK_N": 16, "BLOCK_DIM":64}, num_warps=4, num_stages=2),
+#         triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=2, num_stages=2),
+#         triton.Config({"BLOCK_N": 32, "BLOCK_DIM":64}, num_warps=4, num_stages=2),
+#         triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=2, num_stages=2),
+#         triton.Config({"BLOCK_N": 64, "BLOCK_DIM":32}, num_warps=4, num_stages=2),
 #         triton.Config({"BLOCK_N": 128, "BLOCK_DIM":32}, num_warps=2, num_stages=1),
 #         triton.Config({"BLOCK_N": 128, "BLOCK_DIM":32}, num_warps=4, num_stages=1),
 #         triton.Config({"BLOCK_N": 256, "BLOCK_DIM":32}, num_warps=2, num_stages=1),