From 9df6107dca07021a5b6b036771ba0aabb7ebfabb Mon Sep 17 00:00:00 2001
From: andyluo7 <43718156+andyluo7@users.noreply.github.com>
Date: Fri, 17 Apr 2026 13:10:14 -0700
Subject: [PATCH] [AMD] Enable DFLASH speculative decoding on ROCm (#22342)

Signed-off-by: Andy Luo <andyluo7@users.noreply.github.com>
Co-authored-by: Andy Luo <andyluo7@users.noreply.github.com>
---
 .../srt/layers/attention/triton_backend.py    | 16 +++++++++++--
 python/sglang/srt/models/dflash.py            |  4 ++--
 .../sglang/srt/speculative/dflash_worker.py   | 23 ++++++++++++++-----
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py
index c1f1f48fb..d37f9101a 100644
--- a/python/sglang/srt/layers/attention/triton_backend.py
+++ b/python/sglang/srt/layers/attention/triton_backend.py
@@ -646,7 +646,13 @@ class TritonAttnBackend(AttentionBackend):
                 )
 
             custom_mask = self.cuda_graph_custom_mask
-            custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
+            if (
+                spec_info is not None
+                and getattr(spec_info, "custom_mask", None) is not None
+            ):
+                custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
+            else:
+                custom_mask = None
             seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
             mask_indptr = self.mask_indptr[: bs + 1]
             mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
@@ -798,7 +804,13 @@ class TritonAttnBackend(AttentionBackend):
                     )
                 )
             custom_mask = self.cuda_graph_custom_mask
-            custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
+            if (
+                spec_info is not None
+                and getattr(spec_info, "custom_mask", None) is not None
+            ):
+                custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
+            else:
+                custom_mask = None
             seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
             mask_indptr = self.mask_indptr[: bs + 1]
             mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py
index 27f5cdbf5..c7df14c08 100644
--- a/python/sglang/srt/models/dflash.py
+++ b/python/sglang/srt/models/dflash.py
@@ -163,8 +163,8 @@ class DFlashAttention(nn.Module):
         return k_by_head.view_as(k)
 
     def apply_k_rope(self, positions: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
-        # Use a minimal dummy query (1 head) to avoid doing full-Q work.
-        dummy_q = k.new_empty((k.shape[0], self.head_dim))
+        # Match K shape so RoPE kernel head-count check passes on all backends.
+        dummy_q = k.new_empty(k.shape)
         _, k = self.rotary_emb(positions, dummy_q, k)
         return k
 
diff --git a/python/sglang/srt/speculative/dflash_worker.py b/python/sglang/srt/speculative/dflash_worker.py
index 030aa21e5..4a5c4e948 100644
--- a/python/sglang/srt/speculative/dflash_worker.py
+++ b/python/sglang/srt/speculative/dflash_worker.py
@@ -99,26 +99,37 @@ class DFlashWorker:
         draft_server_args = deepcopy(server_args)
         draft_server_args.skip_tokenizer_init = True
         draft_backend = draft_server_args.speculative_draft_attention_backend
-        supported_draft_backends = ("flashinfer", "fa3", "fa4")
+        supported_draft_backends = ("flashinfer", "fa3", "fa4", "triton")
         if draft_backend is None:
             draft_backend, _ = draft_server_args.get_attention_backends()
         if draft_backend is None:
-            draft_backend = "flashinfer"
+            # Use triton on ROCm (no FlashInfer), flashinfer on CUDA
+            import torch as _torch
+
+            draft_backend = "triton" if _torch.version.hip else "flashinfer"
         elif draft_backend == "trtllm_mha":
+            import torch as _torch
+
+            _fb = "triton" if _torch.version.hip else "flashinfer"
             logger.warning(
                 "DFLASH draft worker does not support 'trtllm_mha' because the "
                 "draft path requires non-causal attention. Falling back to "
-                "'flashinfer'."
+                "'%s'.",
+                _fb,
             )
-            draft_backend = "flashinfer"
+            draft_backend = _fb
         elif draft_backend not in supported_draft_backends:
+            import torch as _torch
+
+            _fb = "triton" if _torch.version.hip else "flashinfer"
             logger.warning(
                 "DFLASH draft worker only supports attention_backend in %s for now, "
-                "but got %r. Falling back to 'flashinfer'.",
+                "but got %r. Falling back to '%s'.",
                 supported_draft_backends,
                 draft_backend,
+                _fb,
             )
-            draft_backend = "flashinfer"
+            draft_backend = _fb
         # Make the draft worker backend explicit and self-contained (no further overrides).
         draft_server_args.speculative_draft_attention_backend = None
         draft_server_args.prefill_attention_backend = None