[bugfix] the hccl_buffsize configuration.

Signed-off-by: mojave2 <[email protected]>

Author: kiscad

tests/ut/test_utils.py

@@ -16,7 +16,9 @@
 import math
 import os
 from threading import Lock
+from types import SimpleNamespace
 from unittest import mock
+from unittest.mock import patch
 
 import torch
 from vllm.config import (CompilationConfig, ModelConfig, ParallelConfig,
@@ -170,6 +172,41 @@ def test_find_hccl_library(self):
     def test_current_stream(self):
         with mock.patch("torch.npu.current_stream") as mock_current_stream:
             self.assertEqual(utils.current_stream(), mock_current_stream())
+            
+    @patch.dict(os.environ, {"HCCL_BUFFSIZE": "2048"})
+    @patch("torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options")
+    def test_create_hccl_pg_options_mc2_with_env(self, mock_options):
+        mock_options.return_value = SimpleNamespace(hccl_config=None)
+
+        options = utils.create_hccl_pg_options("mc2")
+
+        mock_options.assert_called_once_with()
+        self.assertIsNotNone(options.hccl_config)
+        self.assertEqual(options.hccl_config["hccl_buffer_size"], 2048)
+
+    @patch.dict(os.environ, {"HCCL_BUFFSIZE": "1024"})
+    @patch("torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options")
+    def test_create_hccl_pg_options_ep_with_env(self, mock_options):
+        mock_options.return_value = SimpleNamespace(hccl_config=None)
+
+        options = utils.create_hccl_pg_options("ep")
+
+        mock_options.assert_called_once_with()
+        self.assertIsNotNone(options.hccl_config)
+        self.assertEqual(options.hccl_config["hccl_buffer_size"], 1024)
+
+    @patch.dict(os.environ, {}, clear=False)
+    @patch("torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options")
+    def test_create_hccl_pg_options_ep_default(self, mock_options):
+        os.environ.pop("HCCL_BUFFSIZE", None)
+        mock_options.return_value = SimpleNamespace(hccl_config=None)
+
+        options = utils.create_hccl_pg_options("ep")
+
+        mock_options.assert_called_once_with()
+        self.assertIsNotNone(options.hccl_config)
+        self.assertEqual(options.hccl_config["hccl_buffer_size"],
+                         utils._DEFAULT_BUFFER_SIZE)
 
     def test_vllm_version_is(self):
         with mock.patch.dict(os.environ, {"VLLM_VERSION": "1.0.0"}):

vllm_ascend/utils.py

@@ -715,12 +715,11 @@ def npu_stream_switch(target_stream: torch.npu.Stream,
 def create_hccl_pg_options(group_name: str):
     options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
     hccl_config = get_hccl_config_for_pg_options(group_name)
-    if hccl_config is not None:
-        options.hccl_config = hccl_config
+    options.hccl_config = hccl_config
     return options
 
 
-def get_hccl_config_for_pg_options(group_name: str) -> Optional[dict]:
+def get_hccl_config_for_pg_options(group_name: str) -> dict:
     """
     Get HCCL process group options for the given communication group name.
 
@@ -730,21 +729,27 @@ def get_hccl_config_for_pg_options(group_name: str) -> Optional[dict]:
     Returns:
         HCCL pg_options or None for mc2 group
     """
-    # FIXME: Current mc2 operators only perform communication space partitioning
-    # based on HCCL_BUFFSIZE configuration. Using pg_options with mc2 group would
-    # result in memory misalignment problems.
-    if group_name and "mc2" in group_name:
-        return None
     hccl_config_map = {
         "dp": {
             "hccl_buffer_size": calculate_dp_buffer_size()
         },
     }
-    return hccl_config_map.get(group_name, get_default_buffer_config())
+    return hccl_config_map.get(group_name, get_buffer_config())
 
 
-def get_default_buffer_config() -> dict:
-    return {"hccl_buffer_size": _DEFAULT_BUFFER_SIZE}
+def get_buffer_config() -> dict:
+    buffer_size = _DEFAULT_BUFFER_SIZE
+    env_buffer_size = os.getenv("HCCL_BUFFSIZE")
+    if env_buffer_size is not None:
+        try:
+            val = int(env_buffer_size)
+            if val > 0:
+                buffer_size = val
+        except ValueError:
+            # Fallback to default if HCCL_BUFFSIZE is not a valid integer
+            pass
+
+    return {"hccl_buffer_size": buffer_size}
 
 
 def calculate_dp_buffer_size() -> int: