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 2, 11:00 UTC
[ROCm][Perf] DSv3.2 MI355X TP4 decode-step orchestration cleanup (3 micro-opts) (#42982)
Signed-off-by: Frida Andersson <fanderss@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
frida-andersson committed
May 29, 2026 at 13:26 UTC
0b56815a24f4da1af21008355af422891bf654ff
2 files changed
+59
-29
vllm/model_executor/models/deepseek_v2.py
+2
-4
@@ -304,10 +304,8 @@ class DeepseekV2MoE(nn.Module):
304
self.is_rocm_aiter_moe_enabled
305
and self.gate.e_score_correction_bias is not None
306
):
307
- # AITER biased_grouped_topk requires the correction bias dtype to
308
- # match the router logits. Keep DeepSeek's correction bias in fp32
309
- # by requesting fp32 router logits for this routing path.
310
- self.gate.set_out_dtype(torch.float32)
307
+ # Accumulates in fp32; avoids bf16->fp32 cast.
308
+ self.gate.set_out_dtype(self.gate.weight.dtype)
309
310
if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
311
self.shared_experts = None
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+57
-25
@@ -375,7 +375,7 @@ class ROCMAiterMLASparseMetadataBuilder(
375
(1, 1), dtype=torch.int32, device=self.device
376
)
377
378
- self.req_id_per_token_buffer = torch.empty(
378
+ self.req_id_per_token_buffer = torch.zeros(
379
(vllm_config.scheduler_config.max_num_batched_tokens,),
380
dtype=torch.int32,
381
device=device,
@@ -458,6 +458,10 @@ class ROCMAiterMLASparseMetadataBuilder(
458
device=device,
459
)
460
461
+ self._prev_req_extent: int = 0
462
+ self._prev_indices_extent: int = 0
463
+ self._prev_metadata_key: tuple | None = None
464
+
465
def build(
466
self,
467
common_prefix_len: int,
@@ -470,11 +474,23 @@ class ROCMAiterMLASparseMetadataBuilder(
474
req_id_per_token = np.repeat(
475
np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
476
)
473
- # Zero-fill for cudagraphs
474
- self.req_id_per_token_buffer.fill_(0)
475
- self.paged_kv_indices.fill_(0)
476
- self.paged_kv_indptr.fill_(0)
477
- self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
477
+ # Only re-zero the shrink-tail. paged_kv_indptr is fully rewritten
478
+ # by the cumsum below. paged_kv_indices entries past new_indices_extent
479
+ # are never read (the attention kernel only touches the ranges
480
+ # defined by paged_kv_indptr).
481
+ new_req_extent = int(req_id_per_token.shape[0])
482
+ new_indices_extent = num_tokens * self.topk_tokens
483
+ if self._prev_req_extent > new_req_extent:
484
+ self.req_id_per_token_buffer[new_req_extent : self._prev_req_extent].fill_(
485
+ 0
486
+ )
487
+ if self._prev_indices_extent > new_indices_extent:
488
+ self.paged_kv_indices[new_indices_extent : self._prev_indices_extent].fill_(
489
+ 0
490
+ )
491
+ self._prev_req_extent = new_req_extent
492
+ self._prev_indices_extent = new_indices_extent
493
+ self.req_id_per_token_buffer[:new_req_extent].copy_(
494
torch.from_numpy(req_id_per_token), non_blocking=True
495
)
496
query_lens = (
@@ -505,27 +521,43 @@ class ROCMAiterMLASparseMetadataBuilder(
521
# treated as its own batch entry), so persistent metadata can always
522
# be precomputed here. The kernel switches to the persistent
523
# work-stealing path automatically when work_meta_data is non-None.
508
- from aiter import get_mla_metadata_v1
509
-
510
- get_mla_metadata_v1(
511
- qo_indptr,
512
- paged_kv_indptr,
513
- paged_kv_last_page_len,
524
+ # The output is a deterministic function of (num_tokens, max_query_len,
525
+ # num_heads, min(seq_lens, topk_tokens)); fingerprint those CPU-side
526
+ # and skip the launch when nothing changed.
527
+ num_reqs = common_attn_metadata.num_reqs
528
+ clamped_seq_lens = np.minimum(
529
+ common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
530
+ self.topk_tokens,
531
+ )
532
+ metadata_key = (
533
+ num_tokens,
534
+ int(common_attn_metadata.max_query_len),
535
self._num_attention_heads,
515
- 1,
516
- True,
517
- self._mla_work_meta_data,
518
- self._mla_work_info_set,
519
- self._mla_work_indptr,
520
- self._mla_reduce_indptr,
521
- self._mla_reduce_final_map,
522
- self._mla_reduce_partial_map,
523
- page_size=1,
524
- kv_granularity=16,
525
- max_seqlen_qo=1,
526
- uni_seqlen_qo=1,
527
- fast_mode=True,
536
+ clamped_seq_lens.tobytes(),
537
)
538
+ if metadata_key != self._prev_metadata_key:
539
+ from aiter import get_mla_metadata_v1
540
+
541
+ get_mla_metadata_v1(
542
+ qo_indptr,
543
+ paged_kv_indptr,
544
+ paged_kv_last_page_len,
545
+ self._num_attention_heads,
546
+ 1,
547
+ True,
548
+ self._mla_work_meta_data,
549
+ self._mla_work_info_set,
550
+ self._mla_work_indptr,
551
+ self._mla_reduce_indptr,
552
+ self._mla_reduce_final_map,
553
+ self._mla_reduce_partial_map,
554
+ page_size=1,
555
+ kv_granularity=16,
556
+ max_seqlen_qo=1,
557
+ uni_seqlen_qo=1,
558
+ fast_mode=True,
559
+ )
560
+ self._prev_metadata_key = metadata_key
561
562
metadata = ROCMAiterMLASparseMetadata(
563
num_reqs=common_attn_metadata.num_reqs,