[AMD] Enable draft-extend CUDA graph and reduce bubble for MTP (#29202)

Co-authored-by: Thomas Wang <thomawan@amd.com>

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

@@ -583,7 +583,12 @@ class DeepseekV4HipRadixBackend(
             is_prefill=True,
         )
         self._attach_unified_kv_prefill_meta(
-            core_attn_metadata, req_pool_indices, seq_lens, extend_seq_lens
+            core_attn_metadata,
+            req_pool_indices,
+            seq_lens,
+            extend_seq_lens,
+            num_tokens,
+            need_compress=need_compress,
         )
         indexer_metadata = (
             self.init_forward_metadata_indexer(core_attn_metadata)
@@ -1176,6 +1181,8 @@ class DeepseekV4HipRadixBackend(
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         extend_seq_lens: torch.Tensor,
+        num_tokens: int,
+        need_compress: bool = True,
     ) -> None:
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
             is_unified_kv_triton,
@@ -1187,10 +1194,20 @@ class DeepseekV4HipRadixBackend(
         bs = req_pool_indices.shape[0]
         seq_lens = seq_lens.to(torch.int64)
         extend_seq_lens = extend_seq_lens.to(torch.int64)
-        # token -> req index (length L = sum(extend_seq_lens))
-        bid = torch.repeat_interleave(
-            torch.arange(bs, device=device, dtype=torch.int64), extend_seq_lens
-        )
+        # token -> req index (length L = sum(extend_seq_lens)).
+        # output_size skips the implicit sum() D2H on draft-extend. dropping it on the
+        # target-extend path triggers a GPU memory access fault.
+        if need_compress:
+            bid = torch.repeat_interleave(
+                torch.arange(bs, device=device, dtype=torch.int64),
+                extend_seq_lens,
+            )
+        else:
+            bid = torch.repeat_interleave(
+                torch.arange(bs, device=device, dtype=torch.int64),
+                extend_seq_lens,
+                output_size=num_tokens,
+            )
         if core.unified is None:
             core.unified = UnifiedKvMetadata()
         core.unified.pf_state_slot = req_pool_indices[bid]

python/sglang/srt/managers/scheduler_components/profiler_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import gc
 import logging
 import os
 import time
@@ -378,7 +379,11 @@ class SchedulerProfilerManager:
             self.torch_profiler_output_dir,
             merge_message,
         )
-        self.torch_profiler = None
+
+        if self.torch_profiler is not None:
+            self.torch_profiler = None
+            gc.collect()
+
         self.profile_in_progress = False
         self.profiler_start_forward_ct = None
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -397,17 +397,22 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             "cuda": EAGLEDraftExtendCudaGraphRunner,
             "musa": EAGLEDraftCudaGraphRunner,
         }
-        supports_hip_aiter_draft_extend_graph = False
+        supports_hip_draft_extend_graph = False
         if _is_hip:
-            # Keep import local so non-HIP environments do not require aiter.
+            # Keep imports local so non-HIP environments do not require these.
+            # aiter packs draft-extend support into the decode (multi-step)
+            # backend; DSV4 exposes it on the draft-extend backend itself.
             from sglang.srt.layers.attention.aiter_backend import (
                 AiterMultiStepDraftBackend,
             )
-
-            supports_hip_aiter_draft_extend_graph = isinstance(
-                self.draft_attn_backend, AiterMultiStepDraftBackend
+            from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
+                DeepseekV4HipRadixBackend,
             )
 
+            supports_hip_draft_extend_graph = isinstance(
+                self.draft_attn_backend, AiterMultiStepDraftBackend
+            ) or isinstance(self.draft_extend_attn_backend, DeepseekV4HipRadixBackend)
+
         graph_supported_backend_types = [
             TritonAttnBackend,
             TRTLLMMLABackend,
@@ -451,7 +456,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                 _is_npu
                 or _is_xpu
                 or supports_cuda_draft_extend_graph
-                or supports_hip_aiter_draft_extend_graph
+                or supports_hip_draft_extend_graph
             )
         ):
             tic = time.perf_counter()