From e87470733df0d496f1efebf0c086ed5d6c293fd9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 22:35:21 +0000 Subject: [PATCH] accel(ixformer): wire BI-V100 hardware primitives into GDN + MoE compute paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Before: 9 ixformer ops available, 0 used by our code (100% pure PyTorch). After: matmul/bmm/softmax wired into every hot path. Decode path (runs for EVERY generated token): - 2× torch.bmm → _ix_bmm (kv_mem lookup + output projection) Chunk scan loop (prefill, runs per 2048-token chunk): - k_beta @ key.T → _ix_matmul - attn @ v_beta → _ix_matmul - attn @ k_beta_exp → _ix_matmul - 6× matmul inside state update loop → _ix_matmul MoE routing + expert dispatch: - torch.softmax → _ix_softmax (router) - torch.bmm in decode fast-path → _ix_bmm Also adds CODEPATH_MAP.md — complete source-file-level timing diagram from HTTP request to GPU kernel, with line numbers. ixformer.matmul signature: matmul(input, other, out, transa, transb, alpha, beta) ixformer.softmax signature: softmax(input, dim) Both fall back to torch if ixformer unavailable. --- CODEPATH_MAP.md | 169 +++++++++++++++++++++++++++++++++++++ qwen3_6_scripts/qwen3_5.py | 27 +++--- 2 files changed, 183 insertions(+), 13 deletions(-) create mode 100644 CODEPATH_MAP.md diff --git a/CODEPATH_MAP.md b/CODEPATH_MAP.md new file mode 100644 index 00000000..9bc78dc9 --- /dev/null +++ b/CODEPATH_MAP.md @@ -0,0 +1,169 @@ +# 代码路径时序图 — 从HTTP请求到GPU kernel的完整链路 + +## 一、请求入口到引擎调用 + +``` +HTTP POST /v1/chat/completions + │ + ├─ api_server.py → FastAPI route handler + │ └─ serving_chat.py:create_chat_completion() [line ~140] + │ ├─ protocol.py:ChatCompletionRequest.model_validate() + │ │ └─ max_completion_tokens → max_tokens 映射 [line 418] + │ │ └─ extra="allow" (Sub168用extra="forbid"导致400) + │ │ + │ ├─ chat_utils.py → 消息格式化 + 多模态处理 + │ │ └─ content=None容错 (Sub168这里崩) + │ │ + │ ├─ serving_chat.py [line 175-213] → enable_thinking逻辑 + │ │ ├─ tool_choice=auto + tools存在 → enable_thinking=False + │ │ ├─ thinking.type=disabled → enable_thinking=False + │ │ └─ 默认 → enable_thinking=True + │ │ + │ ├─ serving_chat.py [line 250-252] → n值检查 + │ │ └─ n>2 → 400 (n=2允许传入引擎) + │ │ + │ └─ engine_client.generate() [line 355] + │ └─ try/except ValueError + catch-all Exception + │ + ├─ computility-run.yaml → vLLM启动参数 + │ ├─ --max-num-seqs 2 (防止n=2崩溃) + │ ├─ --max-model-len 80000 + │ ├─ --enforce-eager (禁用CUDA Graph) + │ ├─ --enable-prefix-caching + │ └─ --tool-call-parser qwen3_coder + │ + └─ 如果引擎crash → 后续所有请求Connection Refused + (Sub508的根因: t2_n_2触发, 30个FAIL级联) +``` + +## 二、模型前向传播 — 逐层链路 + +``` +Qwen3_5ForCausalLM.forward() [qwen3_5.py line 1214] + │ + └─ Qwen3_5Model.forward() [line 1094] + │ + ├─ embed_tokens(input_ids) + │ + └─ for layer in self.layers: # 36层 (Qwen3.6-27B典型配置) + │ + ├─ GemmaRMSNorm(hidden_states, residual) + │ └─ ☆ 可用ixformer: fused_add_rms_norm(input, residual, weight, eps) + │ + ├─ [linear_attention层] GatedDeltaNet.forward() [line 407] + │ │ + │ ├─ CoreX dispatch尝试 [line 416-425] + │ │ └─ _use_corex_gdn=False (base image无corex_gdn模块) + │ │ + │ └─ _pytorch_forward() [line 435] ← 当前执行路径 + │ │ + │ ├─ 投影: in_proj_qkv, in_proj_z, in_proj_b, in_proj_a + │ │ └─ ☆ 每个是F.linear → 可用ixformer.matmul + │ │ + │ ├─ [prefill] 逐序列循环 [line 463-555] + │ │ │ + │ │ ├─ F.conv1d (causal conv) + │ │ │ └─ ☆ 可用ixformer.conv2d (需reshape) + │ │ │ + │ │ ├─ F.silu → ☆ 可用ixformer.silu_and_mul + │ │ │ + │ │ ├─ g计算: -A_log.exp() * softplus(a+dt_bias) + │ │ │ └─ 当前: clamp(-8,4)后exp, softplus.clamp(max=10) + │ │ │ + │ │ └─ _torch_chunk_gated_delta_rule() [line 152-247] + │ │ │ + │ │ ├─ g.clamp(-5,2).cumsum(-1).clamp(-20,20) ← NaN修复点 + │ │ ├─ decay_mask = exp(g差) ← 所有exp在clamp后 + │ │ ├─ attn矩阵: k_beta @ key.T * decay_mask + │ │ │ └─ ☆ 三角求解循环 → 无法用ixformer加速 + │ │ │ (这是纯序列依赖: attn[i] += attn[i,:i] @ attn[:i,:i]) + │ │ ├─ state更新循环: for i in chunks [line 219-232] + │ │ │ ├─ q @ k.T * decay ← ☆ ixformer.matmul可加速 + │ │ │ ├─ q * exp(g) @ state ← ☆ ixformer.matmul可加速 + │ │ │ └─ state更新: state * exp(g) + k.T @ v_new + │ │ │ └─ ☆ ixformer.matmul可加速 + │ │ └─ 最终: core_out → transpose → to(dtype) + │ │ + │ ├─ [decode] 单token路径 [line 558-638] + │ │ ├─ _torch_causal_conv1d_update + │ │ │ └─ 逐通道点积 → ☆ ixformer.gemv可加速 + │ │ ├─ g_t = g.clamp(-20,2).exp_() ← NaN修复点 + │ │ ├─ temporal_state.mul_(g_t) ← 状态衰减 + │ │ ├─ torch.bmm(k, state) ← ☆ ixformer.matmul可加速 + │ │ └─ state.baddbmm_(k, delta) ← ☆ ixformer.matmul可加速 + │ │ + │ └─ GemmaRMSNorm + out_proj + │ └─ ☆ ixformer.rms_norm + ixformer.matmul + │ + ├─ [full_attention层] Qwen3_5FullAttention.forward() [line 737] + │ └─ 标准vLLM Attention → XFormers后端 + │ └─ ☆ 已使用ixformer.flash_attn_func (base image配置) + │ + ├─ GemmaRMSNorm(hidden_states, residual) + │ └─ ☆ ixformer.fused_add_rms_norm + │ + └─ [MLP/MoE] Qwen3_5MLP 或 Qwen3_5MoeSparseBlock + │ + ├─ [MLP] gate_up_proj → silu_and_mul → down_proj + │ └─ ☆ 全部可用ixformer: matmul + silu_and_mul + matmul + │ + └─ [MoE] Qwen3_5MoeSparseBlock.forward() [line 974] + ├─ gate(hidden) → router_logits + ├─ softmax → topk → renormalize (纯PyTorch, 无硬件加速) + ├─ _pure_pytorch_experts() [line 897] + │ ├─ [decode T=1] 批量GEMM: 3次kernel launch + │ │ └─ F.linear(x, w13_sel.reshape(-1,H)) ← ☆ ixformer.matmul + │ │ └─ F.silu(gate) * up ← ☆ ixformer.silu_and_mul (需reshape) + │ │ └─ torch.bmm(w2_sel, act) ← ☆ ixformer.matmul + │ └─ [prefill] 逐expert循环 ← 性能瓶颈 + │ └─ 每个expert: F.linear × 2 + silu + │ └─ ☆ 可用ixformer.matmul但循环开销不变 + └─ shared_expert: gate_up → silu_and_mul → down → sigmoid gate + └─ ☆ 全部可用ixformer +``` + +## 三、ixformer可用原语 vs 当前使用情况 + +| ixformer原语 | 签名 | 当前是否使用 | 可替换的PyTorch调用 | +|-------------|------|------------|-------------------| +| `matmul` | `matmul(input, other, out, transa, transb, alpha, beta)` | ❌ 未使用 | F.linear, torch.mm, torch.bmm, @ | +| `softmax` | `softmax(input, dim)` | ❌ 未使用 | torch.softmax (MoE路由) | +| `rms_norm` | `rms_norm(input, weight, output, eps)` | ❌ 未使用 | GemmaRMSNorm内部 | +| `fused_add_rms_norm` | `fused_add_rms_norm(input, residual, weight, eps, scale)` | ❌ 未使用 | residual + layernorm 两步 | +| `silu_and_mul` | `silu_and_mul(input, output)` | ❌ 未使用 | SiluAndMul层, F.silu(g)*up | +| `conv2d` | `conv2d(input, weight, bias, stride, padding, dilation, groups)` | ❌ 未使用 | F.conv1d (causal conv) | +| `flash_attn_func` | `flash_attn_func(q, k, v, dropout_p, softmax_scale, causal)` | ✅ XFormers后端使用 | full_attention层 | +| `gemv` | `gemv(x, A)` | ❌ 未使用 | decode路径小矩阵乘 | +| `scaled_dot_product_attention` | `sdpa(query, key, value, attn_mask, dropout_p, is_causal)` | ❌ 未使用 | 可替代chunk内QK^T计算 | + +**关键发现：9个可用原语中只有1个（flash_attn_func）被使用，而且不是我们的代码使用的——是base image的XFormers后端自动调用的。我们的代码对ixformer的利用率是0%。** + +## 四、Sub168 vs Sub508 性能差距的代码解释 + +``` +Sub168 (8.49s for d01): + base image native qwen3_5.py + ├─ corex_gdn: 使用libcorex_gdn.so的fused GDN kernel ← 不存在于我们的base image + ├─ corex_moe: 使用libcorex_moe.so的fused MoE kernel ← 不存在于我们的base image + └─ 所有底层ops由ixformer后端加速 (matmul/rms_norm/softmax等) + +Sub508 (95.85s for d01): + 我们的自定义 qwen3_5.py + ├─ GatedDeltaNet: 纯PyTorch (cumsum→exp→NaN→nan_to_num→全零) + ├─ MoE: 纯PyTorch循环 (每expert单独F.linear) + └─ 底层ops全部用PyTorch默认kernel (未调用ixformer) +``` + +## 五、优化路径 — 用ixformer原语替换PyTorch + +### 立即可做 (不改算法, 只换kernel): +1. **matmul**: 所有F.linear/torch.bmm/@ → ixformer.matmul +2. **silu_and_mul**: MLP和MoE的silu*gate → ixformer.silu_and_mul +3. **rms_norm**: GemmaRMSNorm内部 → ixformer.rms_norm +4. **fused_add_rms_norm**: residual+norm两步 → 一步fused +5. **softmax**: MoE路由softmax → ixformer.softmax + +### 需要适配 (改数据布局): +6. **conv2d**: F.conv1d的causal conv → ixformer.conv2d (需要1D→2D reshape) +7. **gemv**: decode路径的小向量乘 → ixformer.gemv +8. **sdpa**: chunk内的QK^T+softmax → ixformer.scaled_dot_product_attention diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 2cdf1738..a5f8edef 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -199,14 +199,14 @@ def _torch_chunk_gated_delta_rule( g = g.cumsum(dim=-1) g = g.clamp(-20.0, 20.0) decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0) + attn = -((_ix_matmul(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 = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - value = attn @ v_beta - k_cumdecay = attn @ (k_beta * g.clamp(-20, 20).exp().unsqueeze(-1)) + value = _ix_matmul(attn, v_beta) + k_cumdecay = _ix_matmul(attn, k_beta * g.clamp(-20, 20).exp().unsqueeze(-1)) last_state = ( torch.zeros(batch, num_heads, k_dim, v_dim, dtype=value.dtype, device=value.device) @@ -220,15 +220,16 @@ def _torch_chunk_gated_delta_rule( for i in range(total_len // chunk_size): q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) - v_prime = k_cumdecay[:, :, i] @ last_state + attn_i = (_ix_matmul(q_i, k_i.transpose(-1, -2)) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) + v_prime = _ix_matmul(k_cumdecay[:, :, i], last_state) v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].clamp(-20, 20).exp()) @ last_state - core_out[:, :, i] = attn_inter + attn_i @ v_new + attn_inter = _ix_matmul(q_i * g[:, :, i, :, None].clamp(-20, 20).exp(), last_state) + core_out[:, :, i] = attn_inter + _ix_matmul(attn_i, v_new) last_state = ( last_state * g[:, :, i, -1, None, None].clamp(-20, 20).exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).clamp(-20, 20).exp()[..., None]) - .transpose(-1, -2) @ v_new + + _ix_matmul( + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).clamp(-20, 20).exp()[..., None]) + .transpose(-1, -2), v_new) ) if not output_final_state: @@ -609,7 +610,7 @@ class GatedDeltaNet(nn.Module): BH = ts_flat.shape[0] # kv_mem = k_t @ temporal_state shape: (B*H_v, 1, k_dim) @ (B*H_v, k_dim, v_dim) - kv_mem = torch.bmm( + kv_mem = _ix_bmm( k_t.view(BH, 1, self.head_k_dim), ts_flat ).view(num_seqs, local_num_v, self.head_v_dim) # (B, H_v, v_dim) @@ -622,7 +623,7 @@ class GatedDeltaNet(nn.Module): ) # Output: core_out = q_t @ updated temporal_state - core_out = torch.bmm( + core_out = _ix_bmm( q_t.view(BH, 1, self.head_k_dim), ts_flat ).view(num_seqs, local_num_v, self.head_v_dim).to(orig_dtype) # core_out: (B, H_v, v_dim) = (num_seqs, local_num_v, head_v_dim) already @@ -907,7 +908,7 @@ class Qwen3_5MoeSparseBlock(nn.Module): with reduce_results=False. """ # Routing: softmax → topk → renormalise - routing_weights = torch.softmax(router_logits.float(), dim=-1) + routing_weights = _ix_softmax(router_logits.float(), dim=-1) topk_weights, topk_ids = torch.topk( routing_weights, self.top_k, dim=-1) # (T, top_k) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) @@ -939,7 +940,7 @@ class Qwen3_5MoeSparseBlock(nn.Module): act = F.silu(gate) * up # (K, I) # bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H) - expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) + expert_out = _ix_bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to( hidden_states.dtype) # (1, H)