因为转译系统需要通过persistent_kernel.py来完成，所以我们先介绍persistent_kernel.py。
persistent_kernel.py是 Persistent Kernel的Python接口，本质是Python到CUDA持久化内核系统的桥梁，允许用户用python定义复杂的计算图，然后在GPU上高效执行。主要功能包括：
持久化内核管理。提供了 PersistentKernel 作为接口类来管理和执行持久化CUDA内核。
内核编译。将Python定义的计算图编译为CUDA代码并生成共享库。集成了nvcc编译器来编译生成CUDA代码。
内核执行。提供接口来初始化、启动和执行持久化内核。
此外，在 HARD_CODE 定义的C函数是底层入口点，具体如下：
init_func：初始化内核。
launch_func：启动内核执行。会调用到 launch_persistent_kernel。
finalize_func：清理和终止内核。
0x01 流程
persistent_kernel.py的工作流程如下：
初始化：创建 PersistentKernel 类。
定义计算图：使用各种layer方法（如embed_layer、attention_layer等）定义计算图。
编译。调用compile()方法生成和编译CUDA内核。
生成任务图。
创建CUDA代码。
调用nvcc编译器。
创建Python绑定模块
执行：调用call()方法启动内核执行。 self.launch_func()
清理：调用finalize()方法或者自动析构。
具体如下图所示。
3-1
0x02 初始化
初始化函数会创建 PersistentKernel 类。
因为此处只是系统接口，大部分有意义的工作在C++代码中实现，因此此处略过。
class PersistentKernel:
def init(
self,
world_size: int,
mpi_rank: int,
num_workers: int,
num_local_schedulers: int,
num_remote_schedulers: int,
max_seq_length: int,
eos_token_id: int64,
meta_tensors: list[torch.Tensor],
profiler_tensor: torch.Tensor,
spec_decode_config: SpecDecodeConfig
):
self.finalized = False
self._is_compiled = False
self.world_size = world_size
self.mpi_rank = mpi_rank
self.num_workers = num_workers
self.num_local_schedulers = num_local_schedulers
self.num_remote_schedulers = num_remote_schedulers
self.max_seq_length = max_seq_length
self.eos_token_id = eos_token_id
self.kn_graph = KNGraph(CyKNGraph(disable_fingerprint=True))
self.meta_tensors = meta_tensors
self.profiler_tensor = profiler_tensor
self.use_nvshmem = True if world_size > 1 else False
self.spec_decode_config = spec_decode_config
self._spec_decode_handlers = {
"promptlookup": self.prompt_lookup_spec_handler,
}
self._spec_verify_handlers = {
"promptlookup": self.prompt_lookup_verify_handler,
}
0x03 定义计算图
persistent_kernel.py 使用各种layer方法（如embed_layer、attention_layer等）定义计算图。简易流程如下：
MPK-3-2
对应的代码举例如下：
def attach_input(self, torch_tensor: torch.Tensor, name: str = None) -> DTensor:
"""
将PyTorch张量附加到计算图，创建对应的DTensor（分布式张量）。
参数：
torch_tensor: 待附加的PyTorch张量
name: 张量名称（必须指定）
返回：
与输入张量关联的DTensor实例
说明：
仅支持行优先（row-major）内存布局，通过步长校验确保布局正确性
"""
提取张量维度与步长信息
dims = tuple([d for d in torch_tensor.shape])
strides = tuple([s for s in torch_tensor.stride()])
校验是否为行优先布局（高维步长 = 低维步长 × 低维尺寸）
for d in range(len(dims) - 1):
assert strides[d] == strides[d + 1] * dims[d + 1]
转换PyTorch数据类型为框架内部 dtype
dtype = convert_torch_type_to_dtype(torch_tensor.dtype)
创建输入张量节点
t = self.kn_graph.new_input(dims=dims, strides=strides, dtype=dtype)
断言名称非空（当前实现限制）
assert name is not None
将DTensor与PyTorch张量绑定，并注册到计算图
self.kn_graph.attach_torch_tensor(t, torch_tensor, name)
return t
def new_tensor(
self,
dims: tuple,
strides: tuple = None,
dtype: dtype = bfloat16,
name: str = None,
io_category: str = "cuda_tensor",
) -> DTensor:
"""
创建新的DTensor并根据IO类别附加到计算图。
参数：
dims: 张量维度元组
strides: 步长元组（默认自动按行优先计算）
dtype: 数据类型（默认bfloat16）
name: 张量名称（必须指定）
io_category: IO类别（"cuda_tensor"或"nvshmem_tensor"）
返回：
新创建的DTensor实例
说明：
支持CUDA本地张量与NVSHMEM分布式张量两种类型
"""
若指定步长，校验是否为行优先布局
if strides is not None:
for d in range(len(dims) - 1):
assert strides[d] == strides[d + 1] * dims[d + 1]
创建张量节点
t = self.kn_graph.new_input(dims=dims, strides=strides, dtype=dtype)
断言名称非空（当前实现限制）
assert name is not None
根据IO类别绑定张量到计算图
if io_category == "cuda_tensor":
self.kn_graph.attach_cuda_tensor(t, name) # 绑定CUDA张量
elif io_category == "nvshmem_tensor":
self.kn_graph.attach_nvshmem_tensor(t, name) # 绑定NVSHMEM分布式张量
else:
raise RuntimeError(f"Invalid io_category: {io_category}")
return t
def fuse_tensors(
self, inputs: list[DTensor], fused_dim: int, num_groups: int, name: str = None
) -> DTensor:
"""
融合多个张量到单个张量（当前仅支持第0维融合）。
参数：
inputs: 待融合的DTensor列表
fused_dim: 融合维度（必须为0）
num_groups: 分组数量
name: 融合后张量名称
返回：
融合后的DTensor实例
"""
当前仅支持第0维融合
assert fused_dim == 0
调用计算图的张量融合接口
t = self.kn_graph.fuse_tensors(inputs, fused_dim, num_groups, name)
return t
def embed_layer(
self,
input: DTensor, # 输入张量 [batch_size, num_spec_tokens]
weight: DTensor, # 嵌入权重 [vocab_size, hidden_size]
output: DTensor, # 输出张量 [batch_size, hidden_size]
grid_dim: tuple, # CUDA网格维度
block_dim: tuple, # CUDA块维度
input_source: int = 0, # 输入源类型（0: 全 tokens, 1: 输入 token）
):
"""
定义嵌入层计算，将输入张量通过嵌入权重映射到隐藏空间。
参数：
input: 输入张量
weight: 嵌入权重张量
output: 输出张量（用于存储结果）
grid_dim: CUDA kernel的网格维度
block_dim: CUDA kernel的块维度
input_source: 输入源类型标记
说明：
内部创建线程块图（TBGraph），定义输入输出映射关系，并注册为"embedding"任务
"""
创建线程块图（CyTBGraph为底层实现，64为共享内存大小）
tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))
定义输入输出张量的维度映射规则
tb_graph.new_input(input, (-1, 1, -1), -1, True) # 输入张量维度映射
tb_graph.new_input(weight, (1, -1, -1), -1, True) # 权重张量维度映射
tb_graph.new_input(output, (1, 0, -1), -1, True) # 输出张量维度映射
将张量与线程块图关联
self.kn_graph.customized([input, weight, output], tb_graph)
注册嵌入层任务，附加输入源参数
self.kn_graph.register_task(tb_graph, "embedding", [input_source])
def rmsnorm_linear_layer(
self,
input: DTensor, # 输入张量
weight_norm: DTensor, # 归一化权重
weight_linear: DTensor, # 线性层权重
output: DTensor, # 输出张量
grid_dim: tuple, # CUDA网格维度
block_dim: tuple, # CUDA块维度
):
"""
定义RMS归一化+线性变换组合层。
参数：
input: 输入张量（2D）
weight_norm: RMS归一化权重（2D）
weight_linear: 线性层权重（2D）
output: 输出张量（2D）
grid_dim: CUDA kernel的网格维度
block_dim: CUDA kernel的块维度
说明：
先对输入执行RMS归一化，再通过线性层变换，输出结果存储到output
"""
校验输入张量维度（当前仅支持2D张量）
assert input.num_dims == 2
assert weight_linear.num_dims == 2
assert output.num_dims == 2
创建线程块图
tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))