apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: jangwoong-mla-rope-bench
  namespace: kbm-g-np-motif
spec:
  managedBy: trainer.kubeflow.org/trainjob-controller
  podTemplateOverrides:
    - spec:
        containers:
          - name: node
            volumeMounts:
              - mountPath: /dev/shm
                name: shm
              - mountPath: /mair
                name: mair
        volumes:
          - emptyDir:
              medium: Memory
              sizeLimit: 64Gi
            name: shm
          - name: mair
            persistentVolumeClaim:
              claimName: mair
      targetJobs:
        - name: node
  runtimeRef:
    apiGroup: trainer.kubeflow.org
    kind: ClusterTrainingRuntime
    name: torch-distributed
  suspend: false
  trainer:
    args:
      - /bin/bash
      - '-c'
      - |
        set -e
        ACTIVATIONPATH=/mair/team-sys/jangwoong/activation
        DATESTAMP=$(date +'%y_%m_%d_%H_%M')
        SAVE_PATH=$ACTIVATIONPATH/benchmarks/results/mla_rope/${DATESTAMP}
        mkdir -p $SAVE_PATH

        pip install triton pandas

        # Build activation from local source (copy to /tmp to avoid NFS race)
        mkdir -p /tmp/activation_src && rm -rf /tmp/activation_src/* && \
          rsync -a --exclude=build $ACTIVATIONPATH/ /tmp/activation_src/ && \
          pip install --no-build-isolation /tmp/activation_src 2>&1 | tail -50

        python -c "import activation; print('fused_q_rope_inplace:', activation.fused_q_rope_inplace); print('fused_kv_split_rope_cat:', activation.fused_kv_split_rope_cat)"

        nvidia-smi | tee $SAVE_PATH/nvidia_smi.txt

        echo "=== MLA RoPE benchmark ==="
        cd $ACTIVATIONPATH/benchmarks
        CUDA_VISIBLE_DEVICES=0 python run_cases.py --case mla_rope --dtype bf16 \
          --save-path $SAVE_PATH/bench 2>&1 | tee $SAVE_PATH/bench.log

        echo "=== Done. Results at: $SAVE_PATH ==="
        exit 0
    env:
      - name: PYTHONUNBUFFERED
        value: '1'
      - name: PYTORCH_ALLOC_CONF
        value: expandable_segments:True
      - name: CUDA_LAUNCH_BLOCKING
        value: '0'
      - name: OMP_NUM_THREADS
        value: '1'
    image: ghcr.io/motiftechnologies/llm-training:v0.1.8
    numNodes: 1
    numProcPerNode: 1
    resourcesPerNode:
      limits:
        cpu: '96'
        memory: 1024Gi
        nvidia.com/gpu: '8'
      requests:
        cpu: '96'
        memory: 1024Gi
        nvidia.com/gpu: '8'