From c7798c23f80d8480b76a2d745b8fdf9437267beb Mon Sep 17 00:00:00 2001 From: vllmellm Date: Tue, 15 Sep 2026 03:35:25 +0000 Subject: [PATCH 1/2] [Bugfix][ROCm] Fix MXFP4 MoE round-up inflating TP-sharded expert weights on CDNA3 Signed-off-by: vllmellm --- .../layers/fused_moe/oracle/mxfp4.py | 37 +++++++++++++------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index ce19ba39bfc9..d1c151aca247 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -769,18 +769,31 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 128) elif current_platform.is_rocm(): - if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and ( - activation == MoEActivation.SITU or activation == MoEActivation.SILU - ): - # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size - # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a - # no-op for K3's shapes) rather than the generic ROCm 256 round-up, - # which would inflate weights and OOM. - intermediate_size = round_up(intermediate_size, 128) - hidden_size = round_up(hidden_size, 128) - else: - intermediate_size = round_up(intermediate_size, 256) - hidden_size = round_up(hidden_size, 256) + from vllm.platforms.rocm import get_cdna_version + + is_situ_or_silu = activation in ( + MoEActivation.SITU, + MoEActivation.SILU, + ) + + # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size + # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 + # rather than the generic ROCm 256 round-up, which would inflate + # weights and OOM. + aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 + + # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs. + # CDNA4's F16xMXFP4 configuration uses block_k=256. + triton_uses_128 = ( + backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4 + ) + + alignment = ( + 128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256 + ) + + intermediate_size = round_up(intermediate_size, alignment) + hidden_size = round_up(hidden_size, alignment) elif backend == Mxfp4MoeBackend.CPU: # CPU AMX kernel uses BLOCK_N=32, align to 32 intermediate_size = round_up(intermediate_size, 32) From 09fdc5c792da8d2183c507e7f4b2a9993428195f Mon Sep 17 00:00:00 2001 From: vllmellm Date: Thu, 17 Sep 2026 03:32:14 +0000 Subject: [PATCH 2/2] fix format Signed-off-by: vllmellm --- .../ec_transfer/ec_connector/base.py | 9 +++------ .../ec_transfer/ec_connector/metrics.py | 19 ++++++------------- 2 files changed, 9 insertions(+), 19 deletions(-) diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 716f443af244..c56d2cc2129a 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -226,8 +226,7 @@ def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: return None def get_ec_connector_stats(self) -> "ECConnectorStats | None": - """ - Get the EC connector stats collected during the last interval. + """Get the EC connector stats collected during the last interval. Callable on either a worker-side or scheduler-side connector instance: the worker-side instance reports stats gathered during @@ -358,8 +357,7 @@ def has_pending_push_work(self) -> bool: def build_ec_connector_stats( cls, data: dict[str, Any] | None = None ) -> "ECConnectorStats | None": - """ - ECConnectorStats resolution method. This method allows dynamically + """ECConnectorStats resolution method. This method allows dynamically registered connectors to return their own ECConnectorStats object, which can implement custom aggregation logic on the data dict. """ @@ -373,8 +371,7 @@ def build_prom_metrics( labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], ) -> "ECConnectorPromMetrics | None": - """ - Create an ECConnectorPromMetrics subclass which should register + """Create an ECConnectorPromMetrics subclass which should register per-connector Prometheus metrics and implement observe() to expose connector transfer stats via Prometheus. """ diff --git a/vllm/distributed/ec_transfer/ec_connector/metrics.py b/vllm/distributed/ec_transfer/ec_connector/metrics.py index 4a329096cee6..0eb20b9863ae 100644 --- a/vllm/distributed/ec_transfer/ec_connector/metrics.py +++ b/vllm/distributed/ec_transfer/ec_connector/metrics.py @@ -17,8 +17,7 @@ @dataclass class ECConnectorStats: - """ - Base class for EC Connector Stats, a container for transfer performance + """Base class for EC Connector Stats, a container for transfer performance metrics or otherwise important telemetry from the connector. All sub-classes need to be serializable as stats are sent from worker to logger process. @@ -31,14 +30,11 @@ def reset(self): raise NotImplementedError def aggregate(self, other: "ECConnectorStats") -> "ECConnectorStats": - """ - Aggregate stats with another `ECConnectorStats` object. - """ + """Aggregate stats with another `ECConnectorStats` object.""" raise NotImplementedError def reduce(self) -> dict[str, int | float]: - """ - Reduce the observations collected during a time interval to one or + """Reduce the observations collected during a time interval to one or more representative values (eg avg/median/sum of the series). This is meant to be called by the logger to produce a summary of the stats for the last time interval. @@ -107,8 +103,7 @@ def log(self, log_fn=logger.info): class ECConnectorPromMetrics: - """ - A base class for per-connector Prometheus metric registration + """A base class for per-connector Prometheus metric registration and recording. """ @@ -127,8 +122,7 @@ def __init__( self.per_engine_labelvalues = per_engine_labelvalues def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): - """ - Record the supplied transfer statistics to Prometheus metrics. These + """Record the supplied transfer statistics to Prometheus metrics. These statistics are engine-specific, and should be recorded to a metric with the appropriate 'engine' label. These metric instances can be created using the create_metric_per_engine() helper method. @@ -137,8 +131,7 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): class ECConnectorProm: - """ - Support for registering per-connector Prometheus metrics, and + """Support for registering per-connector Prometheus metrics, and recording transfer statistics to those metrics. Uses ECConnectorBase.build_prom_metrics(). """