recipes/docker-images/glm53-kindling/experimental/compose/megamoe.yaml

# megamoe for decode-sized MoE batches (W4A16, CUTLASS's own tensors), and with
# VLLM_MOE_PREFILL=1 moe_prefill.cu for batches of VLLM_MOE_PREFILL_MIN_TOKENS or
# more; CUTLASS keeps the sizes in between. Needs snapshot.yaml's loader, which
# calls megamoe_vllm.install() after loading. VLLM_MOE_PREFILL_Y8: fc2's per-expert
# rows as e4m3 (scale per 128 columns) for the fused SP finalize: half the bytes.
services:
  glm53:
    volumes:
      - ../experimental/megamoe/megamoe_rm.cu:/opt/megamoe/megamoe_rm.cu:ro
      - ../experimental/megamoe/moe_prefill.cu:/opt/megamoe/moe_prefill.cu:ro
      - ../experimental/megamoe/megamoe_vllm.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/megamoe_vllm.py:ro
    environment:
      - VLLM_MEGAMOE=${VLLM_MEGAMOE:-1}
      - VLLM_MEGAMOE_MAX_TOKENS=${VLLM_MEGAMOE_MAX_TOKENS:-8}
      - VLLM_MEGAMOE_VARIANT=${VLLM_MEGAMOE_VARIANT:-20}
      - VLLM_MOE_PREFILL=${VLLM_MOE_PREFILL:-1}
      - VLLM_MOE_PREFILL_MIN_TOKENS=${VLLM_MOE_PREFILL_MIN_TOKENS:-1024}
      - VLLM_MOE_PREFILL_CHECK=${VLLM_MOE_PREFILL_CHECK:-0}
      - VLLM_MOE_PREFILL_Y8=${VLLM_MOE_PREFILL_Y8:-1}