STACK.lock

# cppmega stack pinning — single source of truth for CI image builds.
# Bump a pin → CI rebuilds only the affected wheel(s) thanks to ccache + buildx cache.

base:
  cuda_image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04
  python: "3.13"
  gcc: "15"           # via ppa:ubuntu-toolchain-r/test; needs nvcc --allow-unsupported-compiler
  # CUDA extension wheels must be built and consumed with the same exact
  # stable CUDA wheel. A floating 2.13.* pin can still make wheel builders and
  # runtime images resolve different artifacts, which breaks TE's C++ ABI.
  torch: "2.13.0+cu132"
  torch_index: https://download.pytorch.org/whl/cu132
  torch_arch_list: "9.0;10.0;12.1"   # H100 / B200 / GB10

# Wheels built in matrix on ubuntu-latest (x86_64), uploaded to GH Release as assets.
wheels:
  transformer_engine:
    repo: NVIDIA/TransformerEngine
    ref: v2.16
    version: 2.16
    build: |
      NVTE_FRAMEWORK=pytorch \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - transformer_engine_torch-*.whl
      - transformer_engine_cu13-*.whl
      - transformer_engine-*.whl

  flash_attn:
    # Public upstream flash-attn wheel is FA2 CUDA kernels only for cppmega.
    # Keep flash_attn.cute excluded: FA4 is installed separately from PyPI as
    # flash-attn-4 with its own CUTLASS/quack dependency contract.
    repo: Dao-AILab/flash-attention
    ref: main
    patch: upstream_prs/flash_attn_setup_sm120f.patch
    version: 2.x.x
    build: |
      MAX_JOBS=1 NVCC_THREADS=2 \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - flash_attn-*.whl

  flash_attn_3:
    # FA3 (Hopper) from hopper/ subdir of upstream flash-attention.
    # PACKAGE_NAME=flash_attn_3
    # so it co-installs with flash_attn (FA2/FA4) without conflict. For H100/H200
    # on Modal — useless for B200/GB10 but ships for code parity.
    repo: Dao-AILab/flash-attention
    ref: main
    version: 3.x
    subdir: hopper
    build: |
      MAX_JOBS=1 NVCC_THREADS=2 FLASH_ATTENTION_FORCE_BUILD=TRUE \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - flash_attn_3-*.whl

  mamba_ssm:
    repo: state-spaces/mamba
    ref: 31f3d7b
    version: 2.3.1
    build: |
      MAX_JOBS=1 NVCC_THREADS=2 \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - mamba_ssm-*.whl

  causal_conv1d:
    repo: Dao-AILab/causal-conv1d
    ref: v1.6.1
    version: 1.6.1
    build: |
      MAX_JOBS=1 NVCC_THREADS=2 \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - causal_conv1d-*.whl

  fast_hadamard_transform:
    repo: Dao-AILab/fast-hadamard-transform
    ref: master
    version: 1.1.0
    build: |
      MAX_JOBS=1 NVCC_THREADS=2 \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - fast_hadamard_transform-*.whl

  tilelang:
    # DatasunriseOU/tilelang fork at 5952468a carries:
    #   - upstream tile-ai/tilelang#2071 (046b1bd): removes apache-tvm-ffi<0.1.10 cap
    #     → pyproject pin is "apache-tvm-ffi>=0.1.dev0,<0.2" (compatible with FA4 beta23's >=0.1.12)
    #   - vendored TVM submodule bumped to DatasunriseOU/tvm@9b0a1667, which includes
    #     apache/tvm#18938 (44dbd138d): TVMDerivedObject.__slots__ = ("__dict__", "__weakref__")
    #     → fixes the '_NestedLoopCheckVisitor' object has no attribute '_inst' import crash
    #       under tvm-ffi >=0.1.10 (_ObjectSlotsMeta enforcement).
    # Upstream tile-ai/tilelang HEAD still caps apache-tvm-ffi<0.1.12 (PR #2373/#2736),
    # so it is NOT usable for the FA4 beta23 image. The fork is the clean path (§4 option 1).
    repo: DatasunriseOU/tilelang
    ref: 5952468a
    version: 0.1.9+cuda.git5952468a
    submodules: true
    build: |
      pip wheel . --no-build-isolation -w /wheels
    produces:
      - tilelang-*.whl

  qoptim_cuda:
    repo: NVlabs/COAT
    ref: main
    version: 0.0.0
    subdir: coat/optimizer/kernels
    build: |
      MAX_JOBS=1 NVCC_THREADS=2 \
        pip wheel . --no-build-isolation -w /wheels
    produces:
      - qoptim_cuda-*.whl

runtime_pypi:
  flash_attn_4:
    # FA4 CuTe DSL package. This installs the flash_attn.cute namespace and
    # matching nvidia-cutlass-dsl[cu13] / quack-kernels runtime dependencies.
    # H100/H200 training uses FA3; FA4 is kept for the B200/GB10 CuTe path and
    # for cppmega's cute_dsl_mimo modules.
    package: "flash-attn-4[cu13]==4.0.0b19"
    extra_index_url: https://pypi.nvidia.com

# Repos cloned at runtime in the final image (not built as wheels).
sources:
  megatron_lm:
    repo: NVIDIA/Megatron-LM
    ref: core_v0.18.0
    install_path: /opt/megatron-lm
    install_mode: editable    # `pip install -e .`
    recovered_patches:
      # Archived under cppmega after the old external fork became unavailable.
      # These are not applied automatically; keep them here as provenance for
      # cppmega-specific Megatron work that is not covered by tagged Core.
      - upstream_prs/megatron/0001-merge-PR-3674-DSA-absorbed-MLA-TileLang-fused-sparse.patch
      - upstream_prs/megatron/0002-cherry-pick-PR-4268-delayed-wgrad-overlap-with-P2P-b.patch

# Final runtime image
image:
  registry: ghcr.io/datasunriseou/cppmega
  buildcache_tag: buildcache
  tags:
    - "{{ git_sha }}"
    - latest