flashinfer.xqa.xqa_mla

flashinfer.xqa.xqa_mla(q: Tensor, k_cache: Tensor, v_cache: Tensor, page_table: Tensor, seq_lens: Tensor, output: Tensor, workspace_buffer: Tensor, semaphores: Tensor, page_size: int, q_scale: float | Tensor = 1.0, kv_scale: float | Tensor = 1.0, sm_count: int | None = None, enable_pdl: bool | None = None) → None

使用 XQA MLA（多头潜在注意力）内核应用注意力，并使用分页 KV 缓存。 :param q: 查询张量，形状为 [batch_size, beam_width, num_q_heads, head_dim]。

数据类型应为 torch.float8_e4m3fn。目前仅支持 beam_width 1。

参数:

  • k_cache (torch.Tensor) – 分页 K 缓存张量，形状为 [total_num_cache_heads, head_dim]。数据类型应为 torch.float8_e4m3fn

  • v_cache (torch.Tensor) – 分页 V 缓存张量，形状为 [total_num_cache_heads, head_dim]。数据类型应为 torch.float8_e4m3fn

  • page_table (torch.Tensor) – 页表张量，形状为 batch_size, nb_pages_per_seq。数据类型应为 torch.int32。K 和 V 共享相同的表。

  • seq_lens (torch.Tensor) – 序列长度张量，形状为 [batch_size, beam_width]。数据类型应为 torch.uint32。

  • output (torch.Tensor) – 输出张量，形状为 [batch_size, beam_width, num_q_heads, head_dim]。数据类型应为 torch.bfloat16。此张量将在原地修改。

  • workspace_buffer (torch.Tensor) – 用于临时计算的工作空间缓冲区。数据类型应为 torch.uint8。

  • semaphores (torch.Tensor) – 用于同步的信号量缓冲区。数据类型应为 torch.uint32。

  • page_size (int) – 分页 KV 缓存中每个页面的大小。必须是 [16, 32, 64, 128] 中的一个。

  • q_scale (Union[float, torch.Tensor], default=1.0) – 查询张量的缩放因子。

  • kv_scale (Union[float, torch.Tensor], default=1.0) – KV 缓存的缩放因子。

  • sm_count (Optional[int], default=None) – 要使用的流式多处理器数量。如果为 None，将从设备推断。

  • enable_pdl (Optional[bool], default=None) – 是否启用 PDL（持久数据加载器）优化。如果为 None，如果硬件支持，则设置为 True。

注意

该函数会自动从张量形状推断几个参数：- batch_size 来自 q.shape[0] - head_dim 来自 q.shape[-1] - head_group_ratio 对于 MLA 固定为 128 - max_seq_len 来自 page_table.shape[-1] * page_size