Commit 3d9dd733 authored by laibao's avatar laibao
Browse files

docs: 精简 lmcustomop README,移除 Python 接口/参数说明/约束注意事项

parent ce097b06
......@@ -54,45 +54,6 @@ python setup_lmcustomop.py bdist_wheel
| --- | --- |
| `rms_rotary_embedding_fuse` | 对 `query/key` 执行 RMSNorm 与 RoPE 融合计算(in-place) |
### Python 接口
```python
from lmcustomop import rms_rotary_embedding_fuse
query, key = rms_rotary_embedding_fuse(
positions,
query,
key,
head_size,
cos_sin_cache,
is_neox,
weight_q,
weight_k,
residual_q,
residual_k,
epsilon=1e-5,
)
```
### 参数说明
- `positions`: `int64`,形状 `[num_tokens]``[batch_size, seq_len]`
- `query`: 浮点张量,形状 `[num_tokens, num_heads, head_size]``[batch_size, seq_len, num_heads, head_size]`
- `key`: 浮点张量,形状 `[num_tokens, num_kv_heads, head_size]``[batch_size, seq_len, num_kv_heads, head_size]`
- `head_size`: 每个 head 的维度
- `cos_sin_cache`: RoPE cache,第二维为 `rot_dim`(要求 `rot_dim <= 512`
- `is_neox`: 是否使用 GPT-NeoX 风格旋转
- `weight_q` / `weight_k`: RMSNorm 的权重
- `residual_q` / `residual_k`: 残差输入(需同时提供,或同时不提供)
- `epsilon`: RMSNorm 数值稳定项,默认 `1e-5`
### 约束与注意事项
- 算子会 **原地修改** `query``key`
- `query/key``positions` 的 token 维度必须匹配。
- `num_heads` 必须能被 `num_kv_heads` 整除。
- 当前 kernel 分支覆盖 `head_size``64/128/256/512` 的常见场景。
---
## 安装验证
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment