alias maca version checker

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore/nn/parameter.hpp

@@ -12,6 +12,12 @@ class Parameter : public Tensor {
               Size tp_rank = 0,
               Size tp_size = 1);
 
+    Parameter(const Tensor &tensor,
+              Size tp_dim,
+              Size tp_rank,
+              Size tp_size,
+              Size num_shards);
+
     Parameter(const Shape &shape,
               const DataType &dtype,
               const Device &device,
@@ -25,8 +31,9 @@ class Parameter : public Tensor {
 
 protected:
     // Tensor parallel configs
-    Size tp_dim_;  // dimension partitioned
-    Size tp_rank_; // rank of this partition among tp group
-    Size tp_size_; // total number of partitions
+    Size tp_dim_;         // dimension partitioned
+    Size tp_rank_;        // rank of this partition among tp group
+    Size tp_size_;        // total number of partitions
+    Size num_shards_ = 0; // logical shards (e.g. KV heads) when tp_size > num_kv_head
 };
 } // namespace infinicore::nn

scripts/install.py

@@ -4,8 +4,7 @@
 import sys
 from set_env import (
     set_env,
-    ensure_aten_torch_compiler_includes,
-    xmake_flags_need_aten_torch_compiler_includes,
+    set_env_by_config,
 )
 
 PROJECT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -16,8 +15,7 @@ def run_cmd(cmd):
 
 
 def install(xmake_config_flags=""):
-    if xmake_flags_need_aten_torch_compiler_includes(xmake_config_flags):
-        ensure_aten_torch_compiler_includes()
+    set_env_by_config(xmake_config_flags)
     run_cmd(f"xmake f {xmake_config_flags} -cv")
     run_cmd("xmake")
     run_cmd("xmake install")

scripts/metax_env.py

@@ -0,0 +1,60 @@
+import os
+
+
+def _first_existing_dir(paths: list[str]) -> str:
+    for p in paths:
+        if p and os.path.isdir(p):
+            return p
+    return ""
+
+
+def _metax_toolkit_root(use_mc: bool) -> str:
+    """Return toolkit root for MetaX builds (MACA when use-mc; otherwise HPCC)."""
+    if use_mc:
+        for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"):
+            v = os.environ.get(key, "").strip()
+            if v:
+                return v
+        return _first_existing_dir(["/opt/maca"])
+    return _first_existing_dir(["/opt/hpcc"])
+
+
+def _prepend_path_var(name: str, prefixes: list[str]) -> None:
+    """Prepend colon-separated *prefixes* to env var *name* (POSIX)."""
+    if not prefixes:
+        return
+    chunk = ":".join(prefixes)
+    cur = os.environ.get(name, "")
+    os.environ[name] = f"{chunk}:{cur}" if cur else chunk
+
+
+def set_env_for_metax_gpu(
+    flags: str,
+    *,
+    parse_xmake_cli_flag_values,
+    truthy_flag_value,
+) -> None:
+    """
+    Prepend compiler include paths needed when building ATen-enabled C++ against torch headers.
+
+    This chooses paths based on xmake backend flags (e.g. --metax-gpu) and toolkit selection
+    (e.g. MetaX HPCC vs MACA when --use-mc=y).
+    """
+    d = parse_xmake_cli_flag_values(flags)
+    if not truthy_flag_value(d.get("aten", "n")):
+        return
+
+    if truthy_flag_value(d.get("metax-gpu", "n")):
+        use_mc = truthy_flag_value(d.get("use-mc", "n"))
+        root = _metax_toolkit_root(use_mc=use_mc)
+        if not root:
+            return
+        dirs = [
+            os.path.join(root, "tools", "cu-bridge", "include"),
+            os.path.join(root, "include", "hcr"),
+            os.path.join(root, "include"),
+        ]
+        for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
+            _prepend_path_var(var, dirs)
+        return
+

scripts/set_env.py

@@ -1,39 +1,7 @@
 import os
 import platform
 
-
-def _hpcc_toolkit_root() -> str:
-    """HPCC/MACA install root (cu-bridge, headers). Env vars first; else common container path."""
-    for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"):
-        v = os.environ.get(key, "").strip()
-        if v:
-            return v
-    if os.path.isdir("/opt/hpcc"):
-        return "/opt/hpcc"
-    return ""
-
-
-def _prepend_path_var(name, prefixes):
-    """Prepend colon-separated *prefixes* to env var *name* (POSIX)."""
-    if not prefixes:
-        return
-    chunk = ":".join(prefixes)
-    cur = os.environ.get(name, "")
-    os.environ[name] = f"{chunk}:{cur}" if cur else chunk
-
-
-def ensure_aten_torch_compiler_includes() -> None:
-    """If HPCC root is known, prepend cu-bridge + HPCC headers for g++ compiling ATen .cc (c10/cuda)."""
-    root = _hpcc_toolkit_root()
-    if not root:
-        return
-    dirs = [
-        os.path.join(root, "tools", "cu-bridge", "include"),
-        os.path.join(root, "include", "hcr"),
-        os.path.join(root, "include"),
-    ]
-    for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
-        _prepend_path_var(var, dirs)
+from metax_env import set_env_for_metax_gpu
 
 
 def _parse_xmake_cli_flag_values(flags: str):
@@ -61,20 +29,17 @@ def _truthy_flag_value(v: str) -> bool:
     return v in ("y", "yes", "true", "1", "on")
 
 
-# xmake.lua GPU / accelerator backends (any of these + aten may compile C++ against torch+cuda-style headers).
-_XMAKE_GPU_BACKEND_KEYS = frozenset(
-    {
-        "metax-gpu",
-    }
-)
-
-
-def xmake_flags_need_aten_torch_compiler_includes(flags: str) -> bool:
-    """True when ATen is enabled with any GPU/accelerator backend (install.py / xmake f ...)."""
+def set_env_by_config(flags: str) -> None:
+    """Set environment variables for InfiniCore builds with xmake config flags."""
     d = _parse_xmake_cli_flag_values(flags)
-    if not _truthy_flag_value(d.get("aten", "n")):
-        return False
-    return any(_truthy_flag_value(d.get(k, "n")) for k in _XMAKE_GPU_BACKEND_KEYS)
+    if _truthy_flag_value(d.get("metax-gpu", "n")):
+        set_env_for_metax_gpu(
+            flags,
+            parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values,
+            truthy_flag_value=_truthy_flag_value,
+        )
+    else:
+        pass
 
 
 def set_env():

src/infinicore/nn/parameter.cc

@@ -30,6 +30,11 @@ Parameter::Parameter(const Tensor &tensor, Size tp_dim, Size tp_rank, Size tp_si
     }
 }
 
+Parameter::Parameter(const Tensor &tensor, Size tp_dim, Size tp_rank, Size tp_size, Size num_shards)
+    : Parameter(tensor, tp_dim, tp_rank, tp_size) {
+    num_shards_ = num_shards;
+}
+
 Parameter::Parameter(
     const Shape &shape,
     const DataType &dtype,

xmake.lua

@@ -167,10 +167,13 @@ option_end()
 
 if has_config("metax-gpu") then
     add_defines("ENABLE_METAX_API")
-    -- Container torch build expects this for ATen headers on hpcc.
-    add_defines("USE_HPCC")
     if has_config("use-mc") then
         add_defines("ENABLE_METAX_MC_API")
+        -- MACA torch build expects USE_MACA for ATen headers (e.g. C10_WARP_SIZE).
+        add_defines("USE_MACA")
+    else
+        -- HPCC torch build expects this for ATen headers on hpcc.
+        add_defines("USE_HPCC")
     end
     includes("xmake/metax.lua")
 end
@@ -481,6 +484,30 @@ target("infinicore_cpp_api")
     -- and `xmake/qy.lua`.
 
     before_build(function (target)
+        -- MetaX + flash-attn: `flash_attn_2_cuda` may use a different `mha_fwd_kvcache` ABI
+        -- depending on the underlying stack version. When building with MACA (`--use-mc=y`),
+        -- the version file is typically `/opt/maca/Version.txt` (HPCC uses `/opt/hpcc/Version.txt`).
+        if has_config("metax-gpu") and get_config("flash-attn") and get_config("flash-attn") ~= "" then
+            local version_txt = "/opt/hpcc/Version.txt"
+            if not os.isfile(version_txt) and has_config("use-mc") then
+                version_txt = "/opt/maca/Version.txt"
+            end
+            if os.isfile(version_txt) then
+                local content = os.iorunv("cat", {version_txt}) or ""
+                content = content:trim()
+                local major_str = content:match("Version:(%d+)") or content:match("^(%d+)")
+                if major_str and major_str ~= "" then
+                    local major = tonumber(major_str)
+                    if major then
+                        local define = "INFINICORE_HPCC_VERSION_MAJOR=" .. tostring(major)
+                        target:add("defines", define)
+                        target:add("cxflags", "-D" .. define)
+                        target:add("cxxflags", "-D" .. define)
+                    end
+                end
+            end
+        end
+
         if has_config("aten") then
             local outdata = os.iorunv("python", {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim()
             local TORCH_DIR = outdata
@@ -531,7 +558,6 @@ target("infinicore_cpp_api")
 target_end()
 
 target("_infinicore")
-    add_packages("boost")
     if is_mode("debug") then
         add_defines("BOOST_STACKTRACE_USE_BACKTRACE")
         add_links("backtrace")

xmake/metax.lua

@@ -11,7 +11,7 @@ end
 
 -- Resolve MetaX flash-attn .so path (used only from this file: `before_link` sandbox cannot see globals from `xmake.lua`).
 local FLASH_ATTN_METAX_CUDA_SO_CONTAINER_DEFAULT =
-    "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so"
+    "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so"
 
 local function metax_flash_attn_cuda_so_path()
     -- Highest priority: override the exact `.so` file to link.
@@ -41,36 +41,9 @@ local function metax_flash_attn_cuda_so_path()
     return container_path
 end
 
--- Set numeric HPCC version macro for flash-attn signature/call compatibility.
--- Must be done before compiling `infinicore_cpp_api` sources.
+-- MetaX flash-attn link flags for pip `flash_attn_2_cuda`.
+-- Version/ABI macros are set in `xmake.lua` for `infinicore_cpp_api` so they apply to all sources.
 target("infinicore_cpp_api")
-    before_build(function (target)
-        if not has_config("metax-gpu") then
-            return
-        end
-        if not (get_config("flash-attn") and get_config("flash-attn") ~= "") then
-            return
-        end
-
-        local version_txt = "/opt/hpcc/Version.txt"
-        if os.isfile(version_txt) then
-            local content = os.iorunv("cat", {version_txt}) or ""
-            content = content:trim()
-            -- Example: `Version:2.32.0.6`
-            local hpcc_major_str = content:match("Version:(%d+)") or content:match("^(%d+)")
-            if hpcc_major_str and hpcc_major_str ~= "" then
-                local hpcc_major = tonumber(hpcc_major_str)
-                if hpcc_major then
-                    local define = "INFINICORE_HPCC_VERSION_MAJOR=" .. tostring(hpcc_major)
-                    -- `defines` is the logical flag list for the target,
-                    -- but we also pass `-D...` directly to ensure it reaches compilation.
-                    target:add("defines", define)
-                    target:add("cxflags", "-D" .. define)
-                    target:add("cxxflags", "-D" .. define)
-                end
-            end
-        end
-    end)
     if get_config("flash-attn") and get_config("flash-attn") ~= "" then
         before_link(function (target)
             local flash_so_metax = metax_flash_attn_cuda_so_path()