FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3

RUN mkdir /workspace
WORKDIR /workspace/

# Copy all scripts and the V2 module
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
COPY ./paged_attention_v2_pytorch.py /workspace/paged_attention_v2_pytorch.py
COPY ./paged_attention_v2_triton.py /workspace/paged_attention_v2_triton.py

# Run baseline patches (model registration, xformers fallback, tool parser, etc.)
RUN cd ./qwen3_6_scripts && ./patch_ops.sh

# CRITICAL: Enable ixformer native V1/V2 paged attention kernels.
# Fixes: V1 head_mapping int→Tensor, V2 NotImplementedError → native kernel,
# Triton path mismatch.
RUN python3 /workspace/qwen3_6_scripts/patch_ixformer_native.py

# 1. PagedAttention V2 — fills the NotImplementedError hole
#    Enables partitioned attention for long sequences (>8192 tokens)
RUN python3 /workspace/qwen3_6_scripts/patch_paged_attention_v2.py

# 2. Triton kernel tuning: BLOCK=64, NUM_WARPS=4
#    SMEM: BLOCK_N=64 × head_dim=128 × 2B × 2(K+V) = 32KB ≤ 48KB
#    Occupancy: 4 warps allows 2 blocks/SM vs 1 at 8 warps
RUN python3 /workspace/qwen3_6_scripts/patch_triton_tuning.py

# 3. Enable Triton kernels with automatic fallback to PyTorch if they hang
#    Triton Flash Attention is 10-50x faster than PyTorch for-loop fallback
RUN python3 /workspace/qwen3_6_scripts/patch_enable_triton.py

# 5. head_dim=256 support: Qwen3.6 uses head_dim=256
#    BLOCK=64 overflows SMEM (64×256×2×2=64KB > 48KB)
#    → BLOCK=32 for head_dim=256 (32×256×2×2=32KB ≤ 48KB)
RUN python3 /workspace/qwen3_6_scripts/patch_head256_triton.py

# 4. Raise decode threshold: compiled paged_attention_v1 up to 65536
#    instead of falling back to Python at 32768
RUN python3 /workspace/qwen3_6_scripts/patch_vectorized_decode.py