Skip to content

Commit d8d909f

Browse files
gongchensugongchensu
committed
Add python interface documentation.
1 parent 61e1fe2 commit d8d909f

File tree

13 files changed

+524
-2
lines changed

13 files changed

+524
-2
lines changed

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
## 文档目录
1010

1111
### InfiniCore
12-
- [`Python APIs`]
12+
- [`Python APIs`](python/README.md)
1313

1414
- [`C++ APIs`]
1515

@@ -22,7 +22,7 @@
2222
- [`InfiniCCL`]：统一集合通信库，提供常用的集合通信功能，包括点对点、广播、聚合等。
2323

2424

25-
[`Python APIs`]:README.md
25+
[`Python APIs`]:python/README.md
2626
[`C++ APIs`]:README.md
2727
[`InfiniRT`]:/infinirt/README.md
2828
[`InfiniOP`]:/infiniop/README.md

‎python/README.md‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
# `infinicore` Python 前端
2+
3+
*InfiniCore* 提供了与 C++ 前端一致的 Python 封装，位于 `python/infinicore/`。该模块通过 `pybind11` 将核心张量、算子与设备上下文暴露给 Python，便于在推理框架或调试脚本中快速集成。
4+
5+
## 模块结构
6+
7+
| 符号 | 说明 |
8+
| --- | --- |
9+
| `device` | 设备句柄类 (`python/infinicore/device.py`)，支持 `"cuda:0"`、`device(\"cpu\", 0)` 等写法或复用已有实例。 |
10+
| `dtype` / `float16` 等 | 数据类型枚举 (`python/infinicore/dtype.py`)。 |
11+
| `Tensor` | 张量包装类 (`python/infinicore/tensor.py`)，内部封装底层 `_infinicore` 对象。 |
12+
| `empty` / `zeros` / `ones` / `empty_like` 等 | 张量构造函数（`python/infinicore/tensor.py`），默认要求显式传入 `dtype` 与 `device`。 |
13+
| 顶层算子（`add`、`matmul`、`rearrange`、`attention`） | 暴露在 `infinicore` 命名空间下，对应实现位于 `python/infinicore/ops/`。 |
14+
| `infinicore.nn` | 神经网络相关模块集合，未来可扩展更多组件。 |
15+
| `infinicore.nn.functional` | 函数式算子集合 (`python/infinicore/nn/functional.py`)。 |
16+
| `use_ntops` / `infinicore.ntops` | 若系统安装 `ntops` 包，将自动置位 `use_ntops=True` 并暴露原始模块。 |
17+
18+
所有符号在包的 `__init__.py` 中进行了显式导出，可直接通过 `import infinicore as ic` 后使用。
19+
20+
相关导出定义见 `InfiniCore/python/infinicore/__init__.py`。
21+
22+
## 张量与构造函数
23+
24+
`Tensor` 是对底层 `_infinicore.Tensor` 的 Python 包装，常用接口包括：
25+
26+
- `shape` / `ndim` / `size(dim)` / `stride(dim)`：获取张量维度与步长信息。
27+
- `dtype` / `device`：返回 `dtype` 与 `device` 包装类。
28+
- `numel()` / `is_contiguous()`：查看张量元素数量与存储布局。
29+
- `copy_(src)` / `to(...)`：执行数据拷贝与跨设备搬运。
30+
- `contiguous()` / `permute(dims)` / `view(shape)` / `as_strided(size, stride)`：布局调整与视图操作。
31+
- `debug(filename=None)`：将张量内容打印或输出到二进制文件。
32+
33+
常用构造函数包括 `empty`、`strided_empty`、`zeros`、`ones`、`from_blob`、`strided_from_blob`、`empty_like` 等：
34+
35+
```python
36+
import infinicore as ic
37+
38+
cpu = ic.device("cpu")
39+
a = ic.empty((4, 8), dtype=ic.float16, device=cpu)
40+
b = ic.ones((4, 8), dtype=ic.float16, device=cpu)
41+
a.copy_(b)
42+
```
43+
44+
> 注意：这些函数要求显式传入 `dtype` 与 `device`，避免隐式从 PyTorch/TensorFlow 对象推断。
45+
46+
## 顶层算子 (`infinicore.*`)
47+
48+
以下函数直接通过 `infinicore` 命名空间导出，全部支持可选的 `out` 关键字参数以复用缓冲区：
49+
50+
详见 [`ops` 文档索引](ops/README.md)。
51+
52+
## 函数式算子 (`infinicore.nn.functional`)
53+
54+
函数式 API 集中在 `infinicore.nn.functional`：
55+
56+
详见 [`nn.functional` 文档](nn/functional/README.md)。
57+
58+
## 运行时上下文
59+
60+
- `_infinicore` 在进程内维护运行时状态；创建张量时请显式传入 `device`，并保持算子的所有输入位于同一设备。
61+
- 如需强制同步，可调用 `infinicore.lib._infinicore.sync_stream()`、`sync_device()` 等底层绑定。
62+
- 在同一执行流内串行调用算子通常无需额外同步。
63+
64+
## 与 `ntops` 的协作
65+
66+
- 导入 `ntops` 成功后，`infinicore.use_ntops` 会被设置为 `True`，并可通过 `infinicore.ntops` 访问原始模块。
67+
- `nn.functional.silu` 在 `use_ntops=True` 且设备类型为 `"cuda"`/`"musa"` 且未传 `out` 时，会委托 `ntops.torch.silu`。
68+
- 若想强制禁用，可直接设置 `infinicore.use_ntops = False`。
69+
70+
## 端到端示例
71+
72+
```python
73+
import infinicore as ic
74+
75+
device = ic.device("cuda:0")
76+
77+
q = ic.empty((8, 1, 128), dtype=ic.float16, device=device)
78+
k = ic.empty((2, 1, 128), dtype=ic.float16, device=device)
79+
v = ic.empty((2, 1, 128), dtype=ic.float16, device=device)
80+
k_cache = ic.empty((2, 128, 128), dtype=ic.float16, device=device)
81+
v_cache = ic.empty((2, 128, 128), dtype=ic.float16, device=device)
82+
83+
out = ic.attention(q, k, v, k_cache, v_cache, pos=0)
84+
85+
if ic.use_ntops:
86+
# 在部分设备上，SiLU 会委托给 ntops 的高性能实现
87+
out = ic.nn.functional.silu(out)
88+
89+
out.debug()
90+
```
91+
92+
## 相关链接
93+
94+
- [`infinicore.ops 顶层算子`](ops/README.md)
95+
- [`nn.functional 函数式文档`](nn/functional/README.md)
96+
- [`InfiniOP` 统一算子库](/infiniop/README.md)

‎python/nn/README.md‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# `infinicore.nn` 模块
2+
3+
`infinicore.nn` 聚合了面向神经网络的辅助模块，目前主要暴露函数式算子集合 `infinicore.nn.functional`，并预留位置扩展其他组件（如模块化层、优化器等）。
4+
5+
## 模块结构
6+
7+
| 子模块 | 说明 |
8+
| --- | --- |
9+
| `functional` | 函数式算子集合，接口文档见 [`functional/README.md`](functional/README.md)。 |
10+
11+
## 使用示例
12+
13+
```python
14+
import infinicore as ic
15+
from infinicore.nn import functional as F
16+
17+
x = ic.ones((4, 1024), dtype=ic.float16, device=ic.device("cuda:0"))
18+
normed = F.rms_norm(x, normalized_shape=[1024], weight=ic.ones((1024,), dtype=x.dtype, device=x.device))
19+
activated = F.silu(normed)
20+
```
21+
22+
## 相关链接
23+
24+
- [`nn.functional 函数式文档`](functional/README.md)
25+
- [`Python API 总览`](../README.md)

‎python/nn/functional/README.md‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# `infinicore.nn.functional` 函数式接口
2+
3+
`infinicore.nn.functional` 集中收录 PyTorch 风格的函数式算子封装。实现位于 `InfiniCore/python/infinicore/nn/functional.py`，依赖 `_infinicore` C++ 绑定并复用运行时上下文。
4+
5+
## 公共约定
6+
7+
- 所有函数都返回 `infinicore.Tensor`；当提供 `out`/`inplace` 等参数时会复用已有缓冲区。
8+
- 输入张量需由 `infinicore` 创建（或至少携带 `_underlying` 指针），否则无法与底层运行时交互。
9+
- 若函数内部调用 `_infinicore.*_` 原位接口，需确保输出张量与输入形状、dtype 一致。
10+
11+
## API 详情
12+
13+
- [`causal_softmax`](causal_softmax/README.md)：因果掩码 Softmax。
14+
- [`rms_norm`](rms_norm/README.md)：Root Mean Square LayerNorm。
15+
- [`silu`](silu/README.md)：SiLU（Sigmoid Linear Unit）激活。
16+
- [`swiglu`](swiglu/README.md)：SwiGLU 前向门控。
17+
18+
## 示例
19+
20+
```python
21+
import infinicore as ic
22+
from infinicore.nn import functional as F
23+
24+
x = ic.empty((4, 1024), dtype=ic.float16, device=ic.device("cuda:0"))
25+
w = ic.empty((1024,), dtype=ic.float16, device=x.device)
26+
27+
normed = F.rms_norm(x, normalized_shape=list(w.shape), weight=w)
28+
activated = F.silu(normed)
29+
gated = F.swiglu(activated, ic.empty_like(activated))
30+
31+
probs = F.causal_softmax(gated, out=ic.empty_like(gated))
32+
```
33+
34+
## 相关链接
35+
36+
- [`Python API 总览`](../../README.md)
37+
- [`ntops` 协作接口说明](../../README.md#与-ntops-的协作)

‎python/nn/functional/causal_softmax/README.md‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# `nn.functional.causal_softmax`
2+
3+
对最后一维应用因果掩码 Softmax，用于自回归注意力场景。函数定义位于 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def causal_softmax(input: Tensor, out: Optional[Tensor] = None) -> Tensor
9+
```
10+
11+
- `input`：任意维度张量，末维视为序列维，将应用因果掩码。
12+
- `out`：可选输出张量，若提供需与 `input` 形状、`dtype`、`device` 完全一致。
13+
14+
默认返回新张量；提供 `out` 时调用 `_infinicore.causal_softmax_` 原位写入。
15+
16+
## 示例
17+
18+
```python
19+
import infinicore as ic
20+
from infinicore.nn import functional as F
21+
22+
device = ic.device("cuda:0")
23+
logits = ic.empty((4, 128), dtype=ic.float16, device=device)
24+
25+
probs = F.causal_softmax(logits)
26+
F.causal_softmax(logits, out=logits) # 原位写回
27+
```
28+
29+
## 相关链接
30+
31+
- [`nn.functional` 文档](../README.md)
32+
- [`infinicore.attention` 算子](../../../ops/attention/README.md)

‎python/nn/functional/rms_norm/README.md‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# `nn.functional.rms_norm`
2+
3+
实现 Root Mean Square LayerNorm。函数定义位于 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def rms_norm(
9+
input: Tensor,
10+
normalized_shape: list[int],
11+
weight: Tensor,
12+
eps: float = 1e-5,
13+
*,
14+
out: Optional[Tensor] = None,
15+
) -> Tensor
16+
```
17+
18+
- `input`：待归一化张量，末维通常为隐藏维度。
19+
- `normalized_shape`：期望归一化的维度大小列表，会与 `weight.shape` 进行严格比较。
20+
- `weight`：缩放系数张量，维度需与 `normalized_shape` 匹配。
21+
- `eps`：数值稳定项，默认 `1e-5`。
22+
- `out`：可选输出张量；若提供需与 `input` 形状、`dtype`、`device` 一致。
23+
24+
函数首先断言 `normalized_shape == weight.shape`，然后调用 `_infinicore.rms_norm` / `_infinicore.rms_norm_` 完成计算。
25+
26+
## 示例
27+
28+
```python
29+
import infinicore as ic
30+
from infinicore.nn import functional as F
31+
32+
device = ic.device("cuda:0")
33+
x = ic.ones((4, 1024), dtype=ic.float16, device=device)
34+
gamma = ic.ones((1024,), dtype=ic.float16, device=device)
35+
36+
y = F.rms_norm(x, normalized_shape=list(gamma.shape), weight=gamma, eps=1e-5)
37+
F.rms_norm(x, normalized_shape=[1024], weight=gamma, out=x) # 原位写回
38+
```
39+
40+
## 相关链接
41+
42+
- [`nn.functional` 文档](../README.md)
43+
- [`Tensor` 构造函数](../../../README.md#张量与构造函数)

‎python/nn/functional/silu/README.md‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# `nn.functional.silu`
2+
3+
Sigmoid Linear Unit (SiLU) 激活函数。实现位于 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def silu(
9+
input: Tensor,
10+
inplace: bool = False,
11+
*,
12+
out: Optional[Tensor] = None,
13+
) -> Tensor
14+
```
15+
16+
- `input`：待激活张量。
17+
- `inplace`：是否原地写回到 `input`。
18+
- `out`：可选输出张量，若提供需与 `input` 形状、`dtype`、`device` 一致。
19+
20+
### 行为说明
21+
22+
- 当 `inplace=True` 时，直接调用 `_infinicore.silu_` 写回 `input` 并返回。
23+
- 当未提供 `out` 且 `infinicore.use_ntops=True` 且设备类型为 `"cuda"` 或 `"musa"` 时，会委托 `ntops.torch.silu` 以复用优化实现。
24+
- 其他情况下调用 `_infinicore.silu` 或 `_infinicore.silu_` 完成计算。
25+
26+
## 示例
27+
28+
```python
29+
import infinicore as ic
30+
from infinicore.nn import functional as F
31+
32+
device = ic.device("cuda:0")
33+
x = ic.ones((4, 8), dtype=ic.float16, device=device)
34+
35+
y = F.silu(x) # 返回新张量
36+
F.silu(x, inplace=True) # 原地更新
37+
```
38+
39+
## 相关链接
40+
41+
- [`nn.functional` 文档](../README.md)
42+
- [`use_ntops` 协作说明](../../../README.md#与-ntops-的协作)

‎python/nn/functional/swiglu/README.md‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# `nn.functional.swiglu`
2+
3+
SwiGLU（Swish-Gated Linear Unit）函数式实现，定义在 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def swiglu(
9+
input: Tensor,
10+
other: Tensor,
11+
*,
12+
out: Optional[Tensor] = None,
13+
) -> Tensor
14+
```
15+
16+
- `input`：激活分支张量。
17+
- `other`：门控分支张量。要求与 `input` 在形状、`dtype`、`device` 完全一致。
18+
- `out`：可选输出张量，若提供需满足上述一致性条件。
19+
20+
内部调用 `_infinicore.swiglu` / `_infinicore.swiglu_` 完成计算；未提供 `out` 时返回新张量，否则原地写入。
21+
22+
## 示例
23+
24+
```python
25+
import infinicore as ic
26+
from infinicore.nn import functional as F
27+
28+
device = ic.device("cuda:0")
29+
a = ic.ones((4, 8), dtype=ic.float16, device=device)
30+
b = ic.ones((4, 8), dtype=ic.float16, device=device)
31+
32+
out = F.swiglu(a, b)
33+
F.swiglu(a, b, out=a) # 原位更新
34+
```
35+
36+
## 相关链接
37+
38+
- [`nn.functional` 文档](../README.md)
39+
- [`infinicore.ops` 索引](../../../ops/README.md)

‎python/ops/README.md‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# `infinicore.ops` 顶层算子
2+
3+
该模块通过 pybind11 将常用算子直接暴露在 `infinicore` 命名空间下，对应源码位于 `InfiniCore/python/infinicore/ops/`。所有函数均支持可选的 `out` 参数以复用输出缓冲区。
4+
5+
## 通用注意事项
6+
7+
- 所有参数必须是 `infinicore.Tensor` 实例（或至少携带 `_underlying` 指针），否则无法传递到底层 `_infinicore`。
8+
- `out` 张量（若提供）需要和输出形状、dtype、设备完全匹配。
9+
- 算子执行依赖底层 `_infinicore` 运行时；请在创建张量时传入期望的 `device`，保持所有输入处于同一设备上。
10+
11+
## 文档索引
12+
13+
- [`add`](add/README.md)
14+
- [`matmul`](matmul/README.md)
15+
- [`rearrange`](rearrange/README.md)
16+
- [`attention`](attention/README.md)
17+
18+
## 示例
19+
20+
```python
21+
import infinicore as ic
22+
23+
device = ic.device("cuda:0")
24+
a = ic.ones((4, 8), dtype=ic.float16, device=device)
25+
b = ic.ones((4, 8), dtype=ic.float16, device=device)
26+
27+
ic.add(a, b, out=a) # 原位累加
28+
c = ic.matmul(a, b.permute([1, 0])) # (4, 8) @ (8, 4)
29+
30+
contiguous = ic.rearrange(c)
31+
32+
attn_out = ic.attention(
33+
q=contiguous,
34+
k=contiguous,
35+
v=contiguous,
36+
k_cache=ic.empty((4, 128, contiguous.shape[-1]), dtype=contiguous.dtype, device=contiguous.device),
37+
v_cache=ic.empty((4, 128, contiguous.shape[-1]), dtype=contiguous.dtype, device=contiguous.device),
38+
pos=0,
39+
)
40+
```
41+
42+
## 相关链接
43+
44+
- [`Python API 总览`](../README.md)
45+
- [`nn.functional 函数式接口`](../nn/functional/README.md)

0 commit comments

Comments
 (0)