A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 14:04 UTC
DeepSeekV4-Pro enable cuda graph full and piecewise mode (#42604)
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
bobofang_amd committed
May 15, 2026 at 16:45 UTC
ccde9540bed09782ee3c14937953acc23319a42b
2 files changed
+73
-3
vllm/model_executor/layers/mhc.py
-3
@@ -5,7 +5,6 @@ import torch
5
# this import will also register the custom ops
6
import vllm.model_executor.kernels.mhc as mhc_kernels
7
from vllm.model_executor.custom_op import CustomOp
8
-from vllm.platforms import current_platform
8
9
10
# --8<-- [start:mhc_pre]
@@ -190,8 +189,6 @@ class HCHeadOp(CustomOp):
189
)
190
return out.view(*outer_shape, hidden_size)
191
193
- # This @torch.compile is necessary for accuracy as well as performance.
194
- @torch.compile(backend=current_platform.simple_compile_backend)
192
def forward_hip(
193
self,
194
hidden_states: torch.Tensor,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py
+73
@@ -302,6 +302,30 @@ def combine_topk_swa_indices_ragged(
302
return combined_ragged, combined_indptr, combined_lens
303
304
305
+def _copy_ragged_to_graph_buffers(
306
+ ragged_indices: torch.Tensor,
307
+ ragged_indptr: torch.Tensor,
308
+ ragged_indices_buffer: torch.Tensor,
309
+ ragged_indptr_buffer: torch.Tensor,
310
+ num_rows: int,
311
+ max_entries_per_row: int,
312
+) -> tuple[torch.Tensor, torch.Tensor]:
313
+ """Copy dynamic ragged metadata into persistent CUDA graph buffers.
314
+
315
+ FULL decode graphs capture kernel argument addresses. Keep the returned
316
+ tensors backed by stable storage, while indptr continues to bound reads.
317
+ """
318
+ indptr_out = ragged_indptr_buffer[: num_rows + 1]
319
+ indptr_out.copy_(ragged_indptr, non_blocking=True)
320
+
321
+ max_entries = max(num_rows * max_entries_per_row, 1)
322
+ ragged_out = ragged_indices_buffer[:max_entries]
323
+ nnz = ragged_indices.numel()
324
+ if nnz > 0:
325
+ ragged_out[:nnz].copy_(ragged_indices, non_blocking=True)
326
+ return ragged_out, indptr_out
327
+
328
+
329
@dataclass
330
class DeepseekV4ROCMAiterMLASparseMetadata(FlashMLASparseMetadata):
331
"""ROCm-specific DeepSeek V4 metadata carrying ragged decode topk."""
@@ -317,6 +341,23 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
341
342
343
class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder):
344
+ def __init__(self, *args, **kwargs):
345
+ super().__init__(*args, **kwargs)
346
+ self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None
347
+ self.c128a_decode_topk_ragged_indptr_buffer: torch.Tensor | None = None
348
+ if self.is_deepseek_v4 and self.compress_ratio == 128:
349
+ max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
350
+ self.c128a_decode_topk_ragged_indices_buffer = torch.empty(
351
+ max_tokens * self.c128a_max_compressed,
352
+ dtype=torch.int32,
353
+ device=self.device,
354
+ )
355
+ self.c128a_decode_topk_ragged_indptr_buffer = torch.empty(
356
+ max_tokens + 1,
357
+ dtype=torch.int32,
358
+ device=self.device,
359
+ )
360
+
361
def build(
362
self,
363
common_prefix_len: int,
@@ -338,6 +379,16 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder)
379
dense_decode.reshape(dense_decode.shape[0], -1),
380
decode_lens,
381
)
382
+ assert self.c128a_decode_topk_ragged_indices_buffer is not None
383
+ assert self.c128a_decode_topk_ragged_indptr_buffer is not None
384
+ ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers(
385
+ ragged_indices,
386
+ ragged_indptr,
387
+ self.c128a_decode_topk_ragged_indices_buffer,
388
+ self.c128a_decode_topk_ragged_indptr_buffer,
389
+ dense_decode.shape[0],
390
+ self.c128a_max_compressed,
391
+ )
392
393
return DeepseekV4ROCMAiterMLASparseMetadata(
394
**vars(base),
@@ -347,6 +398,20 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder)
398
399
400
class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
401
+ def __init__(self, *args, **kwargs):
402
+ super().__init__(*args, **kwargs)
403
+ max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
404
+ self.decode_swa_ragged_indices_buffer = torch.empty(
405
+ max_tokens * self.window_size,
406
+ dtype=torch.int32,
407
+ device=self.device,
408
+ )
409
+ self.decode_swa_ragged_indptr_buffer = torch.empty(
410
+ max_tokens + 1,
411
+ dtype=torch.int32,
412
+ device=self.device,
413
+ )
414
+
415
def build(
416
self,
417
common_prefix_len: int,
@@ -370,6 +435,14 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
435
base.decode_swa_indices.reshape(base.num_decode_tokens, -1),
436
base.decode_swa_lens,
437
)
438
+ ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers(
439
+ ragged_indices,
440
+ ragged_indptr,
441
+ self.decode_swa_ragged_indices_buffer,
442
+ self.decode_swa_ragged_indptr_buffer,
443
+ base.num_decode_tokens,
444
+ self.window_size,
445
+ )
446
447
return DeepseekV4ROCMAiterSparseSWAMetadata(
448
**vars(base),