ci(docker): add DeepSeek-V3 kernels and modernize tooling

Dockerfile

@@ -8,6 +8,7 @@ ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && apt-get install -y --no-install-recommends \
     ca-certificates \
     net-tools \
+    unzip \
     kmod \
     ccache \
     cmake \
@@ -53,18 +54,6 @@ ENV MAX_JOBS=32
 # Set VIRTUAL_ENV so uv pip install targets the venv created below
 ENV VIRTUAL_ENV=/AReaL/.venv
 
-# Install Node.js via nvm and Claude Code
-ENV NVM_DIR=/root/.nvm
-ENV NODE_VERSION=22
-RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash \
-    && . "$NVM_DIR/nvm.sh" \
-    && nvm install $NODE_VERSION \
-    && nvm use $NODE_VERSION \
-    && nvm alias default $NODE_VERSION \
-    && npm install -g npm@latest \
-    && curl -fsSL https://claude.ai/install.sh | bash
-ENV PATH="$NVM_DIR/versions/node/v$NODE_VERSION/bin:/root/.local/bin:$PATH"
-
 ##############################################################
 # STAGE 1: Install base torch FIRST
 # Torch rarely changes and is needed for C++ compilation
@@ -112,6 +101,52 @@ RUN git clone https://github.com/Dao-AILab/flash-attention -b v2.8.3 /flash-atte
     && touch $VIRTUAL_ENV/lib/python3.12/site-packages/flash_attn_3/__init__.py \
     && rm -rf /flash-attention
 
+# FlashMLA (Multi-head Latent Attention for DeepSeek-V3)
+RUN git clone https://github.com/deepseek-ai/FlashMLA.git /flash-mla \
+    && cd /flash-mla \
+    && git submodule update --init --recursive \
+    && uv pip install -v . --no-build-isolation \
+    && rm -rf /flash-mla
+
+# DeepGEMM (FP8 GEMM library for DeepSeek-V3)
+RUN git clone https://github.com/deepseek-ai/DeepGEMM /DeepGEMM \
+    && cd /DeepGEMM \
+    && git submodule update --init --recursive \
+    && uv pip install . --no-build-isolation \
+    && rm -rf /DeepGEMM
+
+# DeepEP (Expert Parallelism communication library for MoE)
+# Note: TORCH_CUDA_ARCH_LIST="9.0" enables SM90 features and aggressive PTX instructions
+# The NVSHMEM path is auto-detected from nvidia.nvshmem module installed above
+RUN git clone https://github.com/deepseek-ai/DeepEP /DeepEP \
+    && cd /DeepEP \
+    && TORCH_CUDA_ARCH_LIST="9.0 9.0a" uv pip install -v . --no-build-isolation \
+    && rm -rf /DeepEP
+
+# flash-linear-attention (pure Triton kernels, no CUDA extensions)
+RUN git clone https://github.com/fla-org/flash-linear-attention /flash-linear-attention \
+    && cd /flash-linear-attention \
+    && uv pip install -v . \
+    && rm -rf /flash-linear-attention
+
+##############################################################
+# STAGE 2.5: Install Node.js and npm-based tools
+##############################################################
+
+# Install Node.js via fnm and Claude Code
+ENV FNM_DIR=/root/.fnm
+ENV NODE_VERSION=24.13.0
+ENV PATH="$FNM_DIR/aliases/default/bin:/root/.local/bin:$PATH"
+RUN curl -fsSL https://fnm.vercel.app/install | bash -s -- --install-dir "$FNM_DIR" --skip-shell \
+    && eval "$($FNM_DIR/fnm env --shell bash)" \
+    && $FNM_DIR/fnm install $NODE_VERSION \
+    && $FNM_DIR/fnm default $NODE_VERSION \
+    && npm install -g npm@latest \
+    && curl -fsSL https://claude.ai/install.sh | bash \
+    && curl -fsSL https://opencode.ai/install | bash \
+    && npm install -g @openai/codex \
+    && npm install -g @google/gemini-cli
+
 ##############################################################
 # STAGE 3: Install project dependencies from pyproject.toml
 # Changes to pyproject.toml/uv.lock will invalidate from here

areal/tools/validate_docker_installation.py

@@ -5,6 +5,12 @@
 This script validates dependencies in the Docker environment, which includes
 additional packages (grouped_gemm, apex, transformer_engine, flash_attn_3)
 and flash-attn version 2.8.1 (installed separately in Dockerfile).
+
+Also validates DeepSeek-V3 related packages:
+- FlashMLA: Multi-head Latent Attention (requires SM90+ GPU)
+- DeepGEMM: FP8 GEMM library (requires SM90+ GPU)
+- DeepEP: Expert Parallelism communication (requires SM80+ GPU)
+- flash-linear-attention (fla): Linear attention with Triton kernels
 """
 
 import sys
@@ -29,6 +35,11 @@ class DockerInstallationValidator(BaseInstallationValidator):
             "megatron.core.parallel_state",
             "megatron.core.tensor_parallel",
         ],
+        # DeepSeek-V3 related packages
+        "flash_mla": ["flash_mla"],
+        "deep_gemm": ["deep_gemm"],
+        "deep_ep": ["deep_ep"],
+        "fla": ["fla.ops", "fla.layers", "fla.modules"],
     }
 
     # Add Docker-specific packages to critical list
@@ -67,6 +78,14 @@ def parse_pyproject(self):
         self.add_additional_package("megatron-core", "==0.13.1", required=True)
         self.add_additional_package("mbridge", "==0.13.0", required=True)
 
+        # Add DeepSeek-V3 related packages (installed in Dockerfile, not pyproject.toml)
+        self.add_additional_package("flash_mla", required=False)  # SM90+ only
+        self.add_additional_package("deep_gemm", required=False)  # SM90+ only
+        self.add_additional_package("deep_ep", required=False)  # SM80+ only
+        self.add_additional_package(
+            "fla", required=True
+        )  # Pure Triton, works everywhere
+
     def test_cuda_functionality(self):
         """Run CUDA functionality tests including Docker-specific packages."""
         super().test_cuda_functionality()
@@ -148,6 +167,69 @@ def test_cuda_functionality(self):
         except Exception as e:
             print(f"⚠ Grouped GEMM test failed: {e}")
 
+        print("\n=== DeepSeek-V3 Package Tests ===")
+
+        # Test FlashMLA (requires SM90+ GPU)
+        try:
+            from flash_mla import flash_mla_with_kvcache, get_mla_metadata  # noqa: F401
+
+            print(
+                "✓ FlashMLA imported successfully "
+                "(functions: get_mla_metadata, flash_mla_with_kvcache)"
+            )
+        except ImportError:
+            print("⚠ FlashMLA not available (requires SM90+ GPU)")
+        except Exception as e:
+            print(f"⚠ FlashMLA import failed: {e}")
+
+        # Test DeepGEMM (requires SM90+ GPU)
+        try:
+            import deep_gemm
+
+            num_sms = deep_gemm.get_num_sms()
+            print(f"✓ DeepGEMM imported successfully (detected {num_sms} SMs)")
+        except ImportError:
+            print("⚠ DeepGEMM not available (requires SM90+ GPU)")
+        except Exception as e:
+            print(f"⚠ DeepGEMM test failed: {e}")
+
+        # Test DeepEP (requires SM80+ GPU and NVSHMEM)
+        try:
+            from deep_ep import Buffer, EventOverlap  # noqa: F401
+
+            print("✓ DeepEP imported successfully (classes: Buffer, EventOverlap)")
+        except ImportError:
+            print("⚠ DeepEP not available (requires SM80+ GPU and NVSHMEM)")
+        except Exception as e:
+            print(f"⚠ DeepEP import failed: {e}")
+
+        # Test flash-linear-attention (fla) with actual layer instantiation
+        try:
+            import fla  # noqa: F401
+            import torch
+            from fla.layers import MultiScaleRetention
+
+            if torch.cuda.is_available():
+                device, dtype = "cuda:0", torch.bfloat16
+                retnet = MultiScaleRetention(hidden_size=1024, num_heads=4).to(
+                    device=device, dtype=dtype
+                )
+                x = torch.randn(1, 64, 1024, device=device, dtype=dtype)
+                y, *_ = retnet(x)
+                assert y.shape == x.shape, f"Shape mismatch: {y.shape} != {x.shape}"
+                print(
+                    "✓ flash-linear-attention (fla) - MultiScaleRetention forward pass OK"
+                )
+            else:
+                print(
+                    "✓ flash-linear-attention (fla) imported successfully "
+                    "(CUDA not available for layer test)"
+                )
+        except ImportError:
+            print("⚠ flash-linear-attention (fla) not available")
+        except Exception as e:
+            print(f"⚠ flash-linear-attention (fla) test failed: {e}")
+
     def get_validation_title(self) -> str:
         """Get the title for validation output."""
         return "AReaL Docker Installation Validation"

areal/tools/validation_base.py

@@ -197,6 +197,8 @@ def get_installed_version(self, pkg_name: str) -> str | None:
             pkg_name.replace("-", "_"),
             pkg_name.replace("_", "-"),
         ]
+        if pkg_name.lower() == "fla":
+            variants.append("flash-linear-attention")
 
         for variant in variants:
             try:

docs/tutorial/installation.md

@@ -132,6 +132,33 @@ fused Adam kernel), install them manually after running `uv sync --extra cuda`:
 the already-installed PyTorch for CUDA compilation. Install PyTorch first via
 `uv sync --extra cuda` before attempting to install these packages.
 
+### DeepSeek-V3 Optimization Packages (Optional)
+
+For running DeepSeek-V3 style MoE models with optimal performance, the Docker image also
+includes the following packages. These packages have complex build requirements and GPU
+architecture constraints. **Refer to the
+[Dockerfile](https://github.com/inclusionAI/AReaL/blob/main/Dockerfile) for the exact
+installation commands and environment variables.**
+
+| Package                | Purpose                                     | GPU Requirement |
+| ---------------------- | ------------------------------------------- | --------------- |
+| FlashMLA               | Multi-head Latent Attention kernels         | SM90+ (Hopper)  |
+| DeepGEMM               | FP8 GEMM library for MoE                    | SM90+ (Hopper)  |
+| DeepEP                 | Expert Parallelism communication library    | SM80+ (Ampere)  |
+| flash-linear-attention | Linear attention with Triton kernels        | Any GPU         |
+| NVSHMEM                | Required for DeepEP internode communication | Any GPU         |
+
+**Notes**:
+
+- FlashMLA and DeepGEMM require Hopper (H100/H800) or newer GPUs.
+- DeepEP requires NVSHMEM for internode and low-latency features. The Docker image
+  installs `nvidia-nvshmem-cu12` via pip, which DeepEP auto-detects.
+- flash-linear-attention uses pure Triton kernels and works on any GPU.
+- These packages are cloned from GitHub and compiled during Docker build. For manual
+  installation, ensure git submodules are initialized
+  (`git submodule update --init --recursive`) for FlashMLA and DeepGEMM as they depend
+  on CUTLASS.
+
 5. Validate your AReaL installation:
 
 We provide a script to validate your AReaL installation. Simply run: