[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>

python/sglang/srt/layers/attention/triton_backend.py

@@ -646,7 +646,13 @@ class TritonAttnBackend(AttentionBackend):
                 )
 
             custom_mask = self.cuda_graph_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
+            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)

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.
+        # Match K shape so RoPE kernel head-count check passes on all backends.
-        dummy_q = k.new_empty((k.shape[0], self.head_dim))
+        dummy_q = k.new_empty(k.shape)
         _, k = self.rotary_emb(positions, dummy_q, k)
         return k
 

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