accel(ixformer): wire BI-V100 hardware primitives into GDN + MoE compute paths
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.
This commit is contained in:
Claude
169
CODEPATH_MAP.md
Normal file
169
CODEPATH_MAP.md
Normal file
@@ -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
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user