roce_ring-sparkring-f16b5f4.patch 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. diff -ruN a/roce_ring/__init__.py b/roce_ring/__init__.py
  2. --- a/roce_ring/__init__.py
  3. +++ b/roce_ring/__init__.py
  4. @@ -1,5 +1,9 @@
  5. -"""``b12x.comm.roce`` (RoCEnante): one-shot all-reduce and all-gather over RoCE for multi-node TP.
  6. +"""``b12x.comm.roce_ring`` (RoCEnante with per-peer paths): one-shot all-reduce and all-gather over RoCE.
  7. +The per-peer path variant of ``b12x.comm.roce`` for a switchless four-node ring whose opposite
  8. +nodes are reached through hardware-forwarded paths (FujitsuPolycom/sparkring); see
  9. +``runtime/b12x/roce_ring-provenance.json`` in the recipe for origin and local changes.
  10. +
  11. Target: clusters of DGX Spark nodes joined by their ConnectX-7 200 GbE ports,
  12. one GPU per node. The GB10's unified memory lets the NIC RDMA-write straight
  13. into pinned host memory that the GPU kernel then reads in place, so no
  14. @@ -19,7 +23,7 @@
  15. from ..._lib.meta import OpMeta, Provenance, install_lazy_api
  16. META = OpMeta(
  17. - name="roce",
  18. + name="roce_ring",
  19. group="comm",
  20. api_style="stateful",
  21. entry_points=(
  22. diff -ruN a/roce_ring/_allgather_cute.py b/roce_ring/_allgather_cute.py
  23. --- a/roce_ring/_allgather_cute.py
  24. +++ b/roce_ring/_allgather_cute.py
  25. @@ -351,7 +351,7 @@
  26. 1,
  27. 1,
  28. current_cuda_stream(),
  29. - compile_spec=KernelCompileSpec.from_key("comm.roce.allgather", 2, cache_key),
  30. + compile_spec=KernelCompileSpec.from_key("comm.roce_ring.allgather", 2, cache_key),
  31. )
  32. def run(
  33. diff -ruN a/roce_ring/_oneshot_cute.py b/roce_ring/_oneshot_cute.py
  34. --- a/roce_ring/_oneshot_cute.py
  35. +++ b/roce_ring/_oneshot_cute.py
  36. @@ -429,7 +429,7 @@
  37. 1,
  38. 1,
  39. current_cuda_stream(),
  40. - compile_spec=KernelCompileSpec.from_key("comm.roce.oneshot", 2, cache_key),
  41. + compile_spec=KernelCompileSpec.from_key("comm.roce_ring.oneshot", 2, cache_key),
  42. )
  43. def run(
  44. diff -ruN a/roce_ring/_path_config.py b/roce_ring/_path_config.py
  45. --- a/roce_ring/_path_config.py
  46. +++ b/roce_ring/_path_config.py
  47. @@ -9,6 +9,7 @@
  48. PATH_COUNT = 2
  49. MAX_PATH_COUNT = 4
  50. PEER_HCA_MAP_ENV = "B12X_ROCE_PEER_HCA_MAP"
  51. +PEER_HCA_MAPS_ENV = "B12X_ROCE_PEER_HCA_MAPS"
  52. OPPOSITE_PATHS_ENV = "B12X_ROCE_OPPOSITE_PATHS"
  53. CANONICAL_TP4_HCAS = (
  54. "rocep1s0f0",
  55. @@ -67,6 +68,8 @@
  56. the opposite-rank entry has four distinct HCA indices. A two-HCA runtime
  57. defaults every non-local peer to ``0/1``; runtimes with more HCAs require
  58. an explicit mapping because local and reciprocal indices can differ.
  59. + ``B12X_ROCE_PEER_HCA_MAPS`` carries one such map per rank, ``;``-separated in rank order, so a
  60. + single launcher environment serves every rank; ``B12X_ROCE_PEER_HCA_MAP`` wins when both are set.
  61. """
  62. if hca_count < PATH_COUNT:
  63. @@ -75,6 +78,14 @@
  64. f"got {hca_count}"
  65. )
  66. raw = os.getenv(PEER_HCA_MAP_ENV, "").strip()
  67. + per_rank = os.getenv(PEER_HCA_MAPS_ENV, "").strip()
  68. + if not raw and per_rank:
  69. + maps = [item.strip() for item in per_rank.split(";")]
  70. + if len(maps) != world_size:
  71. + raise ValueError(
  72. + f"{PEER_HCA_MAPS_ENV} needs {world_size} ';'-separated maps, got {len(maps)}"
  73. + )
  74. + raw = maps[rank]
  75. if not raw:
  76. if hca_count == PATH_COUNT:
  77. return tuple(
  78. @@ -160,6 +171,7 @@
  79. "OPPOSITE_PATHS_ENV",
  80. "PATH_COUNT",
  81. "PEER_HCA_MAP_ENV",
  82. + "PEER_HCA_MAPS_ENV",
  83. "opposite_path_count",
  84. "peer_hca_map",
  85. "peer_path_count",
  86. diff -ruN a/roce_ring/_proxy.py b/roce_ring/_proxy.py
  87. --- a/roce_ring/_proxy.py
  88. +++ b/roce_ring/_proxy.py
  89. @@ -70,7 +70,7 @@
  90. if override:
  91. return Path(override)
  92. root = os.getenv("XDG_CACHE_HOME") or os.path.join(os.path.expanduser("~"), ".cache")
  93. - return Path(root) / "b12x" / "roce"
  94. + return Path(root) / "b12x" / "roce_ring"
  95. def _compiler() -> str:
  96. @@ -79,7 +79,7 @@
  97. if candidate and shutil.which(candidate):
  98. return candidate
  99. raise RuntimeError(
  100. - "b12x.comm.roce needs a C compiler (gcc/cc/clang) and libibverbs headers "
  101. + "b12x.comm.roce_ring needs a C compiler (gcc/cc/clang) and libibverbs headers "
  102. "to build its RDMA proxy"
  103. )
  104. diff -ruN a/roce_ring/_roce_proxy.c b/roce_ring/_roce_proxy.c
  105. --- a/roce_ring/_roce_proxy.c
  106. +++ b/roce_ring/_roce_proxy.c
  107. @@ -44,7 +44,7 @@
  108. #define ROCE_PORT 1
  109. #define ROCE_SEND_DEPTH 256
  110. #define ROCE_ABI_VERSION 4
  111. -#define ROCE_IDLE_SPINS 200000
  112. +#define ROCE_IDLE_SPINS 20000000 /* dsv41: match the recipe's SG17 proxy; 200000 let the proxy nap mid-step */
  113. #define ROCE_DEFAULT_TWO_WAVE_THRESHOLD_BYTES 131072u
  114. #define ROCE_WAVE_MODE_TWO 0u
  115. #define ROCE_WAVE_MODE_MIXED_TWO 1u
  116. diff -ruN a/roce_ring/api.py b/roce_ring/api.py
  117. --- a/roce_ring/api.py
  118. +++ b/roce_ring/api.py
  119. @@ -1,4 +1,4 @@
  120. -"""Public surface for comm.roce (docs in the op ``__init__``)."""
  121. +"""Public surface for comm.roce_ring (docs in the op ``__init__``)."""
  122. from __future__ import annotations