Docker Build for DeepSeek-v4 (Don't Merge)

.github/workflows/deploy-projection.yml

@@ -12,7 +12,7 @@ on:
   workflow_dispatch:
   push:
     branches:
-      - dev/wenx/deepseek-v4
+      - dev/tas/deepseek-v4
     paths:
       - "deepseek-v4/projection/site/**"
       - ".github/workflows/deploy-projection.yml"

.github/workflows/docker/Dockerfile

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=docker.io/rocm/primus:v26.2
+ARG BASE_IMAGE=docker.io/rocm/primus:v26.3
 FROM ${BASE_IMAGE}
 
 ARG PRIMUS_TURBO_COMMIT
@@ -23,7 +23,7 @@ RUN rm -rf /var/lib/apt/lists/*
 ENV ROCSHMEM_HOME=/opt/rocshmem
 ENV UCX_HOME=/opt/ucx
 # ENV MPI_HOME=/opt/ompi
-# Use the system OpenMPI prefix from the v26.2 base image.
+# Use the system OpenMPI prefix from the v26.3 base image.
 ENV MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi
 ENV ROCM_HOME=/opt/rocm
 ENV PRIMUS_TURBO_FRAMEWORK=${PRIMUS_TURBO_FRAMEWORK}
@@ -55,9 +55,10 @@ RUN rm -rf /opt/rocSHMEM
 
 RUN pip3 uninstall aiter amd-aiter -y || true && \
     cd /opt && \
-    git clone --recursive https://github.com/ROCm/aiter.git && \
+    git clone https://github.com/ROCm/aiter.git && \
     cd aiter && \
-    git checkout ${PRIMUS_TURBO_AITER_COMMIT} && \
+    git checkout -f ${PRIMUS_TURBO_AITER_COMMIT} && \
+    git submodule sync && \
     git submodule update --init --recursive && \
     PREBUILD_KERNELS=3 pip install --no-cache-dir --use-pep517 .
 
@@ -71,6 +72,18 @@ RUN cd /opt && \
 
 RUN rm -rf /opt/Primus-Turbo
 
+# ---------------------------------------------------------------------------
+# Install Triton
+# ---------------------------------------------------------------------------
+RUN cd /opt && \
+    git clone https://github.com/triton-lang/triton.git && \
+    cd triton && \
+    git checkout 09500db9 && \
+    pip3 install -r python/requirements.txt && \
+    MAX_JOBS=96 pip3 install --no-build-isolation --force-reinstall --no-deps .
+
+RUN rm -rf /opt/triton
+
 # ---------------------------------------------------------------------------
 # Install UCCL-EP (skip for JAX framework)
 # ---------------------------------------------------------------------------
@@ -86,6 +99,24 @@ RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \
     rm -rf /opt/uccl; \
     fi
 
+# ---------------------------------------------------------------------------
+# Install fixed origami (rocm-libraries@223648a) over the base image's bundled
+# 0.1.0. The bundled origami's rank_configs() raises
+# `ValueError: vector::reserve` during MoE grouped-gemm kernel selection and
+# crashes training (turbo's _safe_rank_configs only catches RuntimeError).
+# Skipped for JAX. TODO: drop once a base image ships origami with the fix.
+RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \
+    rm -rf /tmp/rocm-libraries && \
+    git clone --filter=blob:none --no-checkout https://github.com/ROCm/rocm-libraries.git /tmp/rocm-libraries && \
+    cd /tmp/rocm-libraries && \
+    git sparse-checkout init --cone && \
+    git sparse-checkout set shared/origami && \
+    git checkout 223648a26928ebed7f3dd0ccdc044c09f1dccf9b && \
+    (pip uninstall -y origami || true) && \
+    pip install --no-cache-dir ./shared/origami/python && \
+    rm -rf /tmp/rocm-libraries; \
+    fi
+
 # Set the default working directory
 WORKDIR /opt
 

.github/workflows/docker/Dockerfile.ainic

@@ -19,9 +19,9 @@ RUN apt-get update && \
 # ---------------------------------------------------------------------------
 # Enviroment variables
 # ---------------------------------------------------------------------------
-ENV WORKDIR=/opt
+ENV WORKDIR=/workspace
 ENV ROCM_PATH=/opt/rocm
-ENV MPI_PATH=/opt/ompi
+ENV MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi
 
 # =============================== Build AINIC Driver ===============================
 # WARNING: Please ensure the following environment variables are correctly set:
@@ -30,16 +30,24 @@ ENV MPI_PATH=/opt/ompi
 # WARNING: If these paths are missing, tools and libraries may not function correctly.
 # INFO: Installation completed successfully
 
-COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.5-a-56.tar.gz ${WORKDIR}
+COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.5-a-77.tar.gz ${WORKDIR}
 RUN cd ${WORKDIR} && \
+    rm -rf rccl amd-anp && \
     echo "Building ainic bundle... current directory: ${WORKDIR}" && \
-    tar zxf ainic_bundle_1.117.5-a-56.tar.gz && \
-    cd ainic_bundle_1.117.5-a-56 && \
+    tar zxf ainic_bundle_1.117.5-a-77.tar.gz && \
+    cd ainic_bundle_1.117.5-a-77 && \
     tar zxf host_sw_pkg.tar.gz && \
     cd host_sw_pkg && \
-    ./install.sh --domain=user -y 2>&1 | tee log_install.txt && \
-    cd ${WORKDIR} && \
-    apt-get install -y ./amd/ainic/deb-repo/libionic*.deb
+    printf '%s\n' \
+        'Package: libionic1 libionic-dev' \
+        'Pin: version 54.0-187-1' \
+        'Pin-Priority: 1001' \
+        '' \
+        'Package: perftest' \
+        'Pin: version 1:25.04.0.0.84-128-1' \
+        'Pin-Priority: 1001' \
+        > /etc/apt/preferences.d/ainic-a77 && \
+    ./install.sh --domain=user -y 2>&1 | tee log_install.txt
 
 # ---------------------------------------------------------------------------
 # Build rccl

.gitignore

@@ -17,3 +17,8 @@ pp_simulation_result
 # GitHub Pages) despite the generic `data` ignore above.
 !deepseek-v4/projection/site/data/
 !deepseek-v4/projection/site/data/*.json
+
+*.log
+*.nohup
+*.zip
+.triton_cache_shared/

check_hc_expand.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""Correctness + microbench for the fused HC expand Triton kernel vs eager."""
+import importlib.util
+import os
+import time
+
+import torch
+
+os.environ.setdefault("PRIMUS_HC_EXPAND_TRITON", "1")
+# Load the kernel module directly by path to avoid the heavy primus/__init__ chain.
+_MOD = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    "primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/hc_expand.py",
+)
+_spec = importlib.util.spec_from_file_location("hc_expand", _MOD)
+_hc = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_hc)
+hc_expand_triton = _hc.hc_expand_triton
+
+
+def eager_expand(x, out, post, comb):
+    write = post.unsqueeze(-1) * out.unsqueeze(-2)
+    mix = torch.matmul(comb, x)
+    return write + mix
+
+
+def run(leading, K, D, dtype):
+    dev = "cuda"
+    torch.manual_seed(0)
+    mk = lambda *s: torch.randn(*s, device=dev, dtype=dtype)  # noqa: E731
+    x = mk(*leading, K, D)
+    out = mk(*leading, D)
+    post = mk(*leading, K)
+    comb = torch.softmax(mk(*leading, K, K).float(), dim=-1).to(dtype)
+
+    # fp32 reference (ground truth), plus eager and triton at the test dtype.
+    x32, o32, p32, c32 = (t.float().clone().requires_grad_(True) for t in (x, out, post, comb))
+    xa, oa, pa, ca = (t.clone().requires_grad_(True) for t in (x, out, post, comb))
+    xb, ob, pb, cb = (t.clone().requires_grad_(True) for t in (x, out, post, comb))
+
+    ref = eager_expand(x32, o32, p32, c32)
+    eag = eager_expand(xa, oa, pa, ca)
+    tri = hc_expand_triton(xb, ob, pb, cb)
+
+    g32 = torch.randn_like(ref)
+    g = g32.to(dtype)
+    ref.backward(g32)
+    eag.backward(g)
+    tri.backward(g)
+
+    def maxerr(a, b):
+        return (a.float() - b.float()).abs().max().item()
+
+    tag = f"{tuple(leading)} K={K} D={D} {str(dtype).split('.')[-1]}"
+    ok = True
+    for name, r, e, t in [
+        ("fwd", ref, eag, tri),
+        ("dx", x32.grad, xa.grad, xb.grad),
+        ("dout", o32.grad, oa.grad, ob.grad),
+        ("dpost", p32.grad, pa.grad, pb.grad),
+        ("dcomb", c32.grad, ca.grad, cb.grad),
+    ]:
+        e_err = maxerr(r, e)  # eager-vs-fp32 (the dtype's intrinsic noise floor)
+        t_err = maxerr(r, t)  # triton-vs-fp32
+        # Pass if triton is no worse than eager, with an fp32 reduction-order
+        # slack relative to the tensor magnitude (length-D sums reorder).
+        slack = 1e-4 + 5e-5 * r.abs().max().item()
+        passed = t_err <= max(1.5 * e_err, slack)
+        ok = ok and passed
+        print(f"  {name:6s} triton_err={t_err:.3e} eager_err={e_err:.3e} {'OK' if passed else 'FAIL'}")
+    print(f"[{tag}] {'PASS' if ok else 'FAIL'}")
+    return ok
+
+
+def bench(leading, K, D, dtype, iters=50):
+    dev = "cuda"
+    x = torch.randn(*leading, K, D, device=dev, dtype=dtype, requires_grad=True)
+    out = torch.randn(*leading, D, device=dev, dtype=dtype, requires_grad=True)
+    post = torch.randn(*leading, K, device=dev, dtype=dtype, requires_grad=True)
+    comb = (
+        torch.softmax(torch.randn(*leading, K, K, device=dev).float(), dim=-1).to(dtype).requires_grad_(True)
+    )
+    g = torch.randn(*leading, K, D, device=dev, dtype=dtype)
+
+    def step(fn):
+        for t in (x, out, post, comb):
+            t.grad = None
+        y = fn(x, out, post, comb)
+        y.backward(g)
+
+    for fn, name in [(eager_expand, "eager"), (hc_expand_triton, "triton")]:
+        step(fn)
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        for _ in range(iters):
+            step(fn)
+        torch.cuda.synchronize()
+        ms = (time.perf_counter() - t0) / iters * 1e3
+        print(f"  {name:7s} fwd+bwd {ms:8.3f} ms/iter")
+
+
+if __name__ == "__main__":
+    print("=== correctness ===")
+    all_ok = True
+    all_ok &= run((4096,), 4, 7168, torch.bfloat16)
+    all_ok &= run((4096,), 4, 7168, torch.float32)
+    all_ok &= run((2, 17), 4, 320, torch.bfloat16)
+    all_ok &= run((3,), 2, 128, torch.float32)
+    all_ok &= run((5,), 8, 64, torch.bfloat16)
+    print("\n=== bench (production shape B*S=4096, K=4, D=7168, bf16) ===")
+    bench((4096,), 4, 7168, torch.bfloat16)
+    print("\nALL", "PASS" if all_ok else "FAIL")