flashinfer.cascade.merge_state_in_place¶
- flashinfer.cascade.merge_state_in_place(v: Tensor, s: Tensor, v_other: Tensor, s_other: Tensor, mask: Tensor | None = None) None¶
就地合并自注意力状态
(v, s)与另一个状态(v_other, s_other)。- 参数:
v (torch.Tensor) – 要就地更新的部分注意力输出，形状：
(seq_len, num_heads, head_dim)。s (torch.Tensor) – 要就地更新的部分 logsumexpr 值，预期为 float32 张量，形状：
(seq_len, num_heads)。v_other (torch.Tensor) – 要合并的另一个注意力输出，形状：
(seq_len, num_heads, head_dim)。s_other (torch.Tensor) – 要合并的另一个 logsumexp 值，预期为 float32 张量，形状：
(seq_len, num_heads)。mask (Optional[torch.Tensor]) – 用于确定是否合并相应序列状态的布尔掩码张量。对于 CUDA 图很有用。如果未指定（默认），将合并所有序列的状态。形状：
[seq_len]
示例
>>> import torch >>> import flashinfer >>> seq_len = 2048 >>> num_heads = 32 >>> head_dim = 128 >>> v = torch.randn(seq_len, num_heads, head_dim).half().to("cuda:0") >>> s = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda:0") >>> v_other = torch.randn(seq_len, num_heads, head_dim).half().to("cuda:0") >>> s_other = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda:0") >>> flashinfer.merge_state_in_place(v, s, v_other, s_other)