fix: GDN NaN clamp (7 sites) + Dockerfile Step 4 tolerance
qwen3_5.py (2642 lines, 12 prebuilt .so, no fallback): - decay_mask: g_diff.clamp(-20,20) before exp() - Neumann row: .clamp(-65504,65504) on iterative update - k_cumdecay: g.clamp(-20,20).exp() - state loop attn_inter: g.clamp(-20,20).exp() - state loop g_exp_term: .clamp(-20,20) - state loop g_last: .clamp(-20,20) - state loop last_state: .clamp(-65504,65504) after update Dockerfile Step 4: each cp gets 2>/dev/null || true (matches tolerance pattern of Steps 1-3, 5-8)
This commit is contained in:
23
Dockerfile
23
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 && \
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user