diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index f73b79aaa..291672e11 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/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] diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 39eeec2e7..67a74bdbf 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 0b15eb5f1..2c2fc649e 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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()