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:
Claude
2026-08-11 09:20:14 +00:00
parent ed8bdf8714
commit 2f5be7d635
2 changed files with 22 additions and 17 deletions

View File

@@ -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 && \

View File

@@ -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