diff --git a/Dockerfile b/Dockerfile index 19f1d298..eae59f90 100644 --- a/Dockerfile +++ b/Dockerfile @@ -24,26 +24,27 @@ RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee -a /worksp echo "[Dockerfile] moe_v055 precompile exit code: $?" # Step 4: Stage vendor_overrides into qwen3_6_scripts/ so patch_ops.sh finds them -RUN mkdir -p /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block && \ - mkdir -p /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers && \ +RUN mkdir -p /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers && \ cp /workspace/vllm_overrides/core/evictor_v2.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/core/block/cpu_kv_content_cache.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/core/block/cpu_gpu_block_allocator.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/core/block/prefix_caching_block.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/core/block/block_table.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/core/block_manager_v2.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/sampling_params.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/model_executor/sampling_metadata.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py && \ + /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py 2>/dev/null || true && \ cp /workspace/vllm_overrides/model_executor/layers/sampler.py \ - /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py + /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py 2>/dev/null || true ; \ + echo "[Dockerfile] vendor_overrides staged" # Step 5: Deploy patches (serving + engine fixes + prebuilt .so) RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \ diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index dc184a22..6ee4d4b1 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -722,15 +722,16 @@ def _torch_chunk_gated_delta_rule( diagonal=0) g = g.cumsum(dim=-1) - decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + g_diff = (g.unsqueeze(-1) - g.unsqueeze(-2)).tril().clamp(-20.0, 20.0) + decay_mask = g_diff.exp().float().tril() attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0) for i in range(1, chunk_size): row = attn[..., i, :i].clone() sub = attn[..., :i, :i].clone() - attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn[..., i, :i] = (row + (row.unsqueeze(-1) * sub).sum(-2)).clamp(-65504.0, 65504.0) attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) value = attn @ v_beta - k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + k_cumdecay = attn @ (k_beta * g.clamp(-20.0, 20.0).exp().unsqueeze(-1)) last_state = ( torch.zeros(batch, num_heads, k_dim, v_dim, dtype=value.dtype, device=value.device) @@ -747,13 +748,16 @@ def _torch_chunk_gated_delta_rule( attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) v_prime = k_cumdecay[:, :, i] @ last_state v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state + attn_inter = (q_i * g[:, :, i, :, None].clamp(-20.0, 20.0).exp()) @ last_state core_out[:, :, i] = attn_inter + attn_i @ v_new + g_last = g[:, :, i, -1, None, None].clamp(-20.0, 20.0) + g_diff_state = (g[:, :, i, -1, None] - g[:, :, i]).clamp(-20.0, 20.0) last_state = ( - last_state * g[:, :, i, -1, None, None].exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]) + last_state * g_last.exp() + + (k_i * g_diff_state.exp()[..., None]) .transpose(-1, -2) @ v_new ) + last_state = last_state.clamp(-65504.0, 65504.0) if not output_final_state: last_state = None