6d8de852ad
CRITICAL DISCOVERY: Qwen3.6-35B-A3B uses head_dim=256 (not 128). text_cfg.head_dim=256, num_heads=24, num_kv_heads=4, GQA=6 This means ALL previous SMEM calculations were wrong: BLOCK=64 + head_dim=256: 64×256×2×2 = 64KB > 48KB → OVERFLOW BLOCK=64 + head_dim=128: 64×128×2×2 = 32KB ≤ 48KB → OK (but wrong model) Fix: head_dim-dependent BLOCK selection in prefix_prefill.py: head_dim ≤ 128: BLOCK=64, NUM_WARPS=4 (32KB SMEM) head_dim = 256: BLOCK=32, NUM_WARPS=4 (32KB SMEM) head_dim > 256: BLOCK=16, NUM_WARPS=2 (16KB SMEM) Also: _Q_CHUNK in _run_sdpa_fallback reduced 256→128 for head_dim=256 to avoid OOM on long sequences (256×100K×24×4=2.3GB vs 128×100K×24×4=1.2GB). Without this patch, Triton prefill CANNOT work for Qwen3.6. patch_enable_triton.py's try/fallback would always fall back to PyTorch.
34 lines
1.5 KiB
Docker
34 lines
1.5 KiB
Docker
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
|
||
|
||
# Run baseline patches (model registration, xformers fallback, tool parser, etc.)
|
||
RUN cd ./qwen3_6_scripts && ./patch_ops.sh
|
||
|
||
# 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
|