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 7, 16:30 UTC
[ROCm] Add extra step in config initialization to populate custom ops before compilation config init (#34848)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Feb 26, 2026 at 02:05 UTC
6042e66cd5304fc043d96aaa0c22d56f939af320
3 files changed
+62
-40
vllm/config/vllm.py
+2
@@ -809,6 +809,8 @@ class VllmConfig:
809
if "-quant_fp8" not in custom_ops:
810
custom_ops.append("+quant_fp8")
811
812
+ current_platform.apply_config_platform_defaults(self)
813
+
814
if self.compilation_config.mode is None:
815
if self.optimization_level > OptimizationLevel.O0:
816
self.compilation_config.mode = CompilationMode.VLLM_COMPILE
vllm/platforms/interface.py
+14
@@ -393,6 +393,20 @@ class Platform:
393
"""
394
pass
395
396
+ @classmethod
397
+ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
398
+ """
399
+ Apply the platform-specific default values to the config.
400
+
401
+ This function is called during the initialization of global VllmConfig, after
402
+ parsing cli arguments.
403
+ It can modify the defaults of the config according to the platform. For example,
404
+ it can enable custom_ops based on the enabled features.
405
+
406
+ The config is passed by reference, so it can be modified in place.
407
+ """
408
+ pass
409
+
410
@classmethod
411
def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
412
"""
vllm/platforms/rocm.py
+46
-40
@@ -482,19 +482,61 @@ class RocmPlatform(Platform):
482
return device_props.total_memory
483
484
@classmethod
485
- def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
485
+ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
486
from vllm._aiter_ops import rocm_aiter_ops
487
from vllm.config.compilation import CUDAGraphMode
488
489
- cache_config = vllm_config.cache_config
489
compilation_config = vllm_config.compilation_config
491
- parallel_config = vllm_config.parallel_config
492
- is_eager_execution = compilation_config == CUDAGraphMode.NONE
490
+ is_eager_execution = compilation_config.cudagraph_mode == CUDAGraphMode.NONE
491
use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled()
492
use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()
493
use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled()
494
use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
495
use_aiter_triton_rope = rocm_aiter_ops.is_triton_rotary_embed_enabled()
496
+ # Aiter rms norm perform best when CUDA Graph capture is enabled.
497
+ if (
498
+ use_aiter_rms_norm
499
+ and not is_eager_execution
500
+ and "-rms_norm" not in compilation_config.custom_ops
501
+ ):
502
+ compilation_config.custom_ops.append("+rms_norm")
503
+
504
+ if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops:
505
+ compilation_config.custom_ops.append("+quant_fp8")
506
+
507
+ if use_aiter_fused_se and "-grouped_topk" in compilation_config.custom_ops:
508
+ logger.warning_once(
509
+ "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled, which "
510
+ "requires the 'grouped_topk' custom op. Overriding the "
511
+ "user-provided '-grouped_topk'."
512
+ )
513
+ compilation_config.custom_ops.remove("-grouped_topk")
514
+ # Ensure grouped_topk is always enabled when using AITER if
515
+ # its not disabled by user
516
+ if (
517
+ use_aiter_fused_moe
518
+ and "+grouped_topk" not in compilation_config.custom_ops
519
+ and "-grouped_topk" not in compilation_config.custom_ops
520
+ ):
521
+ compilation_config.custom_ops.append("+grouped_topk")
522
+ # Enable rotary embedding when using AITER if its not disabled by user
523
+ if (
524
+ use_aiter_triton_rope
525
+ and "+rotary_embedding" not in compilation_config.custom_ops
526
+ and "-rotary_embedding" not in compilation_config.custom_ops
527
+ ):
528
+ compilation_config.custom_ops.append("+rotary_embedding")
529
+
530
+ # Default dispatch to rocm's sparse_attn_indexer implementation
531
+ compilation_config.custom_ops.append("+sparse_attn_indexer")
532
+
533
+ @classmethod
534
+ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
535
+ from vllm.config.compilation import CUDAGraphMode
536
+
537
+ cache_config = vllm_config.cache_config
538
+ compilation_config = vllm_config.compilation_config
539
+ parallel_config = vllm_config.parallel_config
540
541
if compilation_config.cudagraph_mode.has_full_cudagraphs():
542
# decode context parallel does not support full cudagraphs
@@ -533,42 +575,6 @@ class RocmPlatform(Platform):
575
576
if parallel_config.worker_cls == "auto":
577
parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
536
- # Aiter rms norm perform best when CUDA Graph capture is enabled.
537
- if (
538
- use_aiter_rms_norm
539
- and not is_eager_execution
540
- and "-rms_norm" not in compilation_config.custom_ops
541
- ):
542
- compilation_config.custom_ops.append("+rms_norm")
543
-
544
- if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops:
545
- compilation_config.custom_ops.append("+quant_fp8")
546
-
547
- if use_aiter_fused_se and "-grouped_topk" in compilation_config.custom_ops:
548
- logger.warning_once(
549
- "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled, which "
550
- "requires the 'grouped_topk' custom op. Overriding the "
551
- "user-provided '-grouped_topk'."
552
- )
553
- compilation_config.custom_ops.remove("-grouped_topk")
554
- # Ensure grouped_topk is always enabled when using AITER if
555
- # its not disabled by user
556
- if (
557
- use_aiter_fused_moe
558
- and "+grouped_topk" not in compilation_config.custom_ops
559
- and "-grouped_topk" not in compilation_config.custom_ops
560
- ):
561
- compilation_config.custom_ops.append("+grouped_topk")
562
- # Enable rotary embedding when using AITER if its not disabled by user
563
- if (
564
- use_aiter_triton_rope
565
- and "+rotary_embedding" not in compilation_config.custom_ops
566
- and "-rotary_embedding" not in compilation_config.custom_ops
567
- ):
568
- compilation_config.custom_ops.append("+rotary_embedding")
569
-
570
- # Default dispatch to rocm's sparse_attn_indexer implementation
571
- compilation_config.custom_ops.append("+sparse_attn_indexer")
578
579
@classmethod
580
def verify_model_arch(cls, model_arch: str) -> None: