How you implement 3-stage pipeline communication with StepMesh?

[Huixxi]

Is there any demos or examples or docs?

[niehao100]

Hi, @Huixxi
You can refer to bmk_comm_latency_multiserver.py, the variable mb represents for 3-stage.

[Huixxi]

Hi, @Huixxi You can refer to bmk_comm_latency_multiserver.py, the variable mb represents for 3-stage.

Thanks a lot! And may I ask how the micro-batch worked I mean given a large batch size when decoding like a 128 total batch, how you split it into multi 3-stage micro batches and how to execute the whole pipeline? Whether it will be split into something like 43 43 42 three small batches(if the micro batch is 3, like shows in the paper (D1,D2,D3)), then pick one token each time from each small batch to compose a micro batch to execute the following pipeline stage. I don't quite understand the implementation detail, and just from the pipeline picture in the paper is a little confused me. Thanks.

[niehao100]

Hi @Huixxi
I think every micro batch size, or every num in (D1,D2,D3) is 128.

[heroes999]

@niehao100 In AF disaggregation, one Attention instance has 8 gpu cards, do we need to all gather 8 tensor to 1 gpu to facilitate NetSend?

[niehao100]

Hi @heroes999
8 GPU at atten node run in DP mode. I'm not sure the final version, in the test version, we process all gather at FFN side before calculation, and reduce scatter at end.

[Huixxi]

@niehao100 have you met this problem when installing # Build and install Fserver （AF's Python SDK）
pip3 install -v -e . ? My torch has already installed.

$ pip3 show torch
Name: torch
Version: 2.7.1
Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration
Home-page: https://pytorch.org/
Author: PyTorch Team
Author-email: packages@pytorch.org
License: BSD-3-Clause
Location: /home/admin/.local/lib/python3.12/site-packages
Requires: filelock, fsspec, jinja2, networkx, nvidia-cublas-cu12, nvidia-cuda-cupti-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-runtime-cu12, nvidia-cudnn-cu12, nvidia-cufft-cu12, nvidia-cufile-cu12, nvidia-curand-cu12, nvidia-cusolver-cu12, nvidia-cusparse-cu12, nvidia-cusparselt-cu12, nvidia-nccl-cu12, nvidia-nvjitlink-cu12, nvidia-nvtx-cu12, setuptools, sympy, triton, typing-extensions
Required-by: compressed-tensors, flashinfer-python, outlines, torchaudio, torchvision, xgrammar

$ python3 -m pip install -v -e .                               
Using pip 25.2 from /home/admin/.local/lib/python3.12/site-packages/pip (python 3.12)
Defaulting to user installation because normal site-packages is not writeable
Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple
Obtaining file:///home/logs/StepMesh
  Running command python setup.py egg_info
  /home/admin/.local/lib/python3.12/site-packages/setuptools/dist.py:332: InformationOnly: Normalizing '0.0.1.dev' to '0.0.1.dev0'
    self.metadata.version = self._normalize_version(self.metadata.version)
  running egg_info
  creating /tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info
  writing /tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info/PKG-INFO
  writing dependency_links to /tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info/dependency_links.txt
  writing top-level names to /tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info/top_level.txt
  writing manifest file '/tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info/SOURCES.txt'
  reading manifest file '/tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info/SOURCES.txt'
  adding license file 'LICENSE'
  writing manifest file '/tmp/pip-pip-egg-info-7l3gtcwt/FServer.egg-info/SOURCES.txt'
  Preparing metadata (setup.py) ... done
Installing collected packages: FServer
  DEPRECATION: Legacy editable install of FServer==0.0.1.dev0 from file:///home/logs/StepMesh (setup.py develop) is deprecated. pip 25.3 will enforce this behaviour change. A possible replacement is to add a pyproject.toml or enable --use-pep517, and use setuptools >= 64. If the resulting installation is not behaving as expected, try using --config-settings editable_mode=compat. Please consult the setuptools documentation for more information. Discussion can be found at https://github.com/pypa/pip/issues/11457
  Running setup.py develop for FServer
    Running command python setup.py develop
    /home/admin/.local/lib/python3.12/site-packages/setuptools/dist.py:332: InformationOnly: Normalizing '0.0.1.dev' to '0.0.1.dev0'
      self.metadata.version = self._normalize_version(self.metadata.version)
    running develop
    /home/admin/.local/lib/python3.12/site-packages/setuptools/_distutils/cmd.py:90: DevelopDeprecationWarning: develop command is deprecated.
    !!

            ********************************************************************************
            Please avoid running ``setup.py`` and ``develop``.
            Instead, use standards-based tools like pip or uv.

            By 2025-Oct-31, you need to update your project and remove deprecated calls
            or your builds will no longer be supported.

            See https://github.com/pypa/setuptools/issues/917 for details.
            ********************************************************************************

    !!
      self.initialize_options()
    Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple
    Obtaining file:///home/logs/StepMesh
      Installing build dependencies: started
      Installing build dependencies: finished with status 'done'
      Checking if build backend supports build_editable: started
      Checking if build backend supports build_editable: finished with status 'done'
      Getting requirements to build editable: started
      Getting requirements to build editable: finished with status 'error'
      error: subprocess-exited-with-error

      × Getting requirements to build editable did not run successfully.
      │ exit code: 1
      ╰─> [23 lines of output]
          Traceback (most recent call last):
            File "/home/admin/.local/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 389, in <module>
              main()
            File "/home/admin/.local/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 373, in main
              json_out["return_val"] = hook(**hook_input["kwargs"])
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            File "/home/admin/.local/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 157, in get_requires_for_build_editable
              return hook(config_settings)
                     ^^^^^^^^^^^^^^^^^^^^^
            File "/tmp/pip-build-env-5qn0xpal/overlay/local/lib/python3.12/dist-packages/setuptools/build_meta.py", line 473, in get_requires_for_build_editable
              return self.get_requires_for_build_wheel(config_settings)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            File "/tmp/pip-build-env-5qn0xpal/overlay/local/lib/python3.12/dist-packages/setuptools/build_meta.py", line 331, in get_requires_for_build_wheel
              return self._get_build_requires(config_settings, requirements=[])
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            File "/tmp/pip-build-env-5qn0xpal/overlay/local/lib/python3.12/dist-packages/setuptools/build_meta.py", line 301, in _get_build_requires
              self.run_setup()
            File "/tmp/pip-build-env-5qn0xpal/overlay/local/lib/python3.12/dist-packages/setuptools/build_meta.py", line 512, in run_setup
              super().run_setup(setup_script=setup_script)
            File "/tmp/pip-build-env-5qn0xpal/overlay/local/lib/python3.12/dist-packages/setuptools/build_meta.py", line 317, in run_setup
              exec(code, locals())
            File "<string>", line 1, in <module>
          ModuleNotFoundError: No module named 'torch'
          [end of output]

      note: This error originates from a subprocess, and is likely not a problem with pip.
    error: subprocess-exited-with-error

    × Getting requirements to build editable did not run successfully.
    │ exit code: 1
    ╰─> See above for output.

    note: This error originates from a subprocess, and is likely not a problem with pip.
    Traceback (most recent call last):
      File "<string>", line 2, in <module>
      File "<pip-setuptools-caller>", line 35, in <module>
      File "/home/logs/StepMesh/setup.py", line 59, in <module>
        setup(
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/__init__.py", line 115, in setup
        return distutils.core.setup(**attrs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/_distutils/core.py", line 186, in setup
        return run_commands(dist)
               ^^^^^^^^^^^^^^^^^^
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/_distutils/core.py", line 202, in run_commands
        dist.run_commands()
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/_distutils/dist.py", line 1002, in run_commands
        self.run_command(cmd)
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/dist.py", line 1102, in run_command
        super().run_command(command)
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/_distutils/dist.py", line 1021, in run_command
        cmd_obj.run()
      File "/home/admin/.local/lib/python3.12/site-packages/setuptools/command/develop.py", line 39, in run
        subprocess.check_call(cmd)
      File "/usr/lib/python3.12/subprocess.py", line 413, in check_call
        raise CalledProcessError(retcode, cmd)
    subprocess.CalledProcessError: Command '['/usr/bin/python3', '-m', 'pip', 'install', '-e', '.', '--use-pep517', '--no-deps', '--user']' returned non-zero exit status 1.
    error: subprocess-exited-with-error
    
    × python setup.py develop did not run successfully.
    │ exit code: 1
    ╰─> See above for output.
    
    note: This error originates from a subprocess, and is likely not a problem with pip.
    full command: /usr/bin/python3 -c '
    exec(compile('"'"''"'"''"'"'
    # This is <pip-setuptools-caller> -- a caller that pip uses to run setup.py
    #
    # - It imports setuptools before invoking setup.py, to enable projects that directly
    #   import from `distutils.core` to work with newer packaging standards.
    # - It provides a clear error message when setuptools is not installed.
    # - It sets `sys.argv[0]` to the underlying `setup.py`, when invoking `setup.py` so
    #   setuptools doesn'"'"'t think the script is `-c`. This avoids the following warning:
    #     manifest_maker: standard file '"'"'-c'"'"' not found".
    # - It generates a shim setup.py, for handling setup.cfg-only projects.
    import os, sys, tokenize, traceback
    
    try:
        import setuptools
    except ImportError:
        print(
            "ERROR: Can not execute `setup.py` since setuptools failed to import in "
            "the build environment with exception:",
            file=sys.stderr,
        )
        traceback.print_exc()
        sys.exit(1)
    
    __file__ = %r
    sys.argv[0] = __file__
    
    if os.path.exists(__file__):
        filename = __file__
        with tokenize.open(__file__) as f:
            setup_py_code = f.read()
    else:
        filename = "<auto-generated setuptools caller>"
        setup_py_code = "from setuptools import setup; setup()"
    
    exec(compile(setup_py_code, filename, "exec"))
    '"'"''"'"''"'"' % ('"'"'/home/logs/StepMesh/setup.py'"'"',), "<pip-setuptools-caller>", "exec"))' develop --no-deps --user --prefix=
    cwd: /home/logs/StepMesh/
error: subprocess-exited-with-error

× python setup.py develop did not run successfully.
│ exit code: 1
╰─> See above for output.

note: This error originates from a subprocess, and is likely not a problem with pip.

[niehao100]

@Huixxi
Please try with python3 setup.py install for detailed output.

It seems you installed torch with none root user, and some ENV_VARIABLE magic is required.

Do everything as root, like reinstall torch as root can be helpful. The easiest way is try with this container #7 (comment)

[heroes999]

Hi @heroes999 8 GPU at atten node run in DP mode. I'm not sure the final version, in the test version, we process all gather at FFN side before calculation, and reduce scatter at end.

Does StepMesh support NetSend/NetRecv threads dynamically binding to different gpus? So that every gpu at atten node(DP) can send and recv their own part of batch？ @niehao100 Otherwise, atten node should add "all gather" to collect all batches from 8 DP.

[Huixxi]

Hi @heroes999 8 GPU at atten node run in DP mode. I'm not sure the final version, in the test version, we process all gather at FFN side before calculation, and reduce scatter at end.

But what if the FFN is EP? Will all-gather and scatter be replaced with something like DeepEP's dispatch and combine?

[niehao100]

Hi @heroes999 8 GPU at atten node run in DP mode. I'm not sure the final version, in the test version, we process all gather at FFN side before calculation, and reduce scatter at end.

Does StepMesh support NetSend/NetRecv threads dynamically binding to different gpus? So that every gpu at atten node(DP) can send and recv their own part of batch？ @niehao100 Otherwise, atten node should add "all gather" to collect all batches from 8 DP.

1. We don't support this feature. Each GPU is pinned with its own send/recv process and RNIC, each GPUs' part of batch is processed seperately.
2. FFN instances in different node will pull different tensors after RS(their own batch).

Hope this can help you.

[heroes999]

Hi @heroes999 8 GPU at atten node run in DP mode. I'm not sure the final version, in the test version, we process all gather at FFN side before calculation, and reduce scatter at end.

Does StepMesh support NetSend/NetRecv threads dynamically binding to different gpus? So that every gpu at atten node(DP) can send and recv their own part of batch？ @niehao100 Otherwise, atten node should add "all gather" to collect all batches from 8 DP.

1. We don't support this feature. Each GPU is pinned with its own send/recv process and RNIC, each GPUs' part of batch is processed seperately.
2. FFN instances in different node will pull different tensors after RS(their own batch).

Hope this can help you.

Got it. We can start several workers(8) in one A node to do send/recv for each gpu, same with F node, right?

[niehao100]

@heroes999 Yes, that's how we use it. We call one FFN node as one server, ATTN as worker, and one process as one instance. Each node has 8 instance, as the DMLC_GROUP_SIZE defined.

[niehao100]

Hi @heroes999 8 GPU at atten node run in DP mode. I'm not sure the final version, in the test version, we process all gather at FFN side before calculation, and reduce scatter at end.

But what if the FFN is EP? Will all-gather and scatter be replaced with something like DeepEP's dispatch and combine?

We use EP in current design. There is combine op.

[niehao100]

Closing the issue,.