ROCm: Flex Attention Enablement with custom backends (#4178)

Co-authored-by: linsun12 <linsun12@amd.com>

python/sglang/srt/server_args.py

@@ -710,13 +710,23 @@ class ServerArgs:
         )
 
         # Kernel backend
-        parser.add_argument(
-            "--attention-backend",
-            type=str,
-            choices=["flashinfer", "triton", "torch_native"],
-            default=ServerArgs.attention_backend,
-            help="Choose the kernels for attention layers.",
-        )
+        if is_hip():
+            parser.add_argument(
+                "--attention-backend",
+                type=str,
+                choices=["triton", "torch_native", "aiter", "aiter_decode"],
+                default=ServerArgs.attention_backend,
+                help="Choose the kernels for attention layers.",
+            )
+        else:
+            parser.add_argument(
+                "--attention-backend",
+                type=str,
+                choices=["flashinfer", "triton", "torch_native"],
+                default=ServerArgs.attention_backend,
+                help="Choose the kernels for attention layers.",
+            )
+
         parser.add_argument(
             "--sampling-backend",
             type=str,