diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 2671309bfd..36325d3520 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -5,7 +5,6 @@ import torch # this import will also register the custom ops import vllm.model_executor.kernels.mhc as mhc_kernels from vllm.model_executor.custom_op import CustomOp -from vllm.platforms import current_platform # --8<-- [start:mhc_pre] @@ -190,8 +189,6 @@ class HCHeadOp(CustomOp): ) return out.view(*outer_shape, hidden_size) - # This @torch.compile is necessary for accuracy as well as performance. - @torch.compile(backend=current_platform.simple_compile_backend) def forward_hip( self, hidden_states: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py index a647b7d2c6..e7a82a1ee2 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py @@ -302,6 +302,30 @@ def combine_topk_swa_indices_ragged( return combined_ragged, combined_indptr, combined_lens +def _copy_ragged_to_graph_buffers( + ragged_indices: torch.Tensor, + ragged_indptr: torch.Tensor, + ragged_indices_buffer: torch.Tensor, + ragged_indptr_buffer: torch.Tensor, + num_rows: int, + max_entries_per_row: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Copy dynamic ragged metadata into persistent CUDA graph buffers. + + FULL decode graphs capture kernel argument addresses. Keep the returned + tensors backed by stable storage, while indptr continues to bound reads. + """ + indptr_out = ragged_indptr_buffer[: num_rows + 1] + indptr_out.copy_(ragged_indptr, non_blocking=True) + + max_entries = max(num_rows * max_entries_per_row, 1) + ragged_out = ragged_indices_buffer[:max_entries] + nnz = ragged_indices.numel() + if nnz > 0: + ragged_out[:nnz].copy_(ragged_indices, non_blocking=True) + return ragged_out, indptr_out + + @dataclass class DeepseekV4ROCMAiterMLASparseMetadata(FlashMLASparseMetadata): """ROCm-specific DeepSeek V4 metadata carrying ragged decode topk.""" @@ -317,6 +341,23 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None + self.c128a_decode_topk_ragged_indptr_buffer: torch.Tensor | None = None + if self.is_deepseek_v4 and self.compress_ratio == 128: + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + self.c128a_decode_topk_ragged_indices_buffer = torch.empty( + max_tokens * self.c128a_max_compressed, + dtype=torch.int32, + device=self.device, + ) + self.c128a_decode_topk_ragged_indptr_buffer = torch.empty( + max_tokens + 1, + dtype=torch.int32, + device=self.device, + ) + def build( self, common_prefix_len: int, @@ -338,6 +379,16 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder) dense_decode.reshape(dense_decode.shape[0], -1), decode_lens, ) + assert self.c128a_decode_topk_ragged_indices_buffer is not None + assert self.c128a_decode_topk_ragged_indptr_buffer is not None + ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( + ragged_indices, + ragged_indptr, + self.c128a_decode_topk_ragged_indices_buffer, + self.c128a_decode_topk_ragged_indptr_buffer, + dense_decode.shape[0], + self.c128a_max_compressed, + ) return DeepseekV4ROCMAiterMLASparseMetadata( **vars(base), @@ -347,6 +398,20 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder) class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + self.decode_swa_ragged_indices_buffer = torch.empty( + max_tokens * self.window_size, + dtype=torch.int32, + device=self.device, + ) + self.decode_swa_ragged_indptr_buffer = torch.empty( + max_tokens + 1, + dtype=torch.int32, + device=self.device, + ) + def build( self, common_prefix_len: int, @@ -370,6 +435,14 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild base.decode_swa_indices.reshape(base.num_decode_tokens, -1), base.decode_swa_lens, ) + ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( + ragged_indices, + ragged_indptr, + self.decode_swa_ragged_indices_buffer, + self.decode_swa_ragged_indptr_buffer, + base.num_decode_tokens, + self.window_size, + ) return DeepseekV4ROCMAiterSparseSWAMetadata( **vars(base),