nano-attention

nano-vllm/nanovllm/layers/attention.py · bb823b3 · L43–75 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

43class Attention(nn.Module):
44
45    def __init__(
46        self,
47        num_heads,
48        head_dim,
49        scale,
50        num_kv_heads,
51    ):
52        super().__init__()
53        self.num_heads = num_heads
54        self.head_dim = head_dim
55        self.scale = scale
56        self.num_kv_heads = num_kv_heads
57        self.k_cache = self.v_cache = torch.tensor([])
58
59    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
60        context = get_context()
61        k_cache, v_cache = self.k_cache, self.v_cache
62        if k_cache.numel() and v_cache.numel():
63            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
64        if context.is_prefill:
65            if context.block_tables is not None:    # prefix cache
66                k, v = k_cache, v_cache
67            o = flash_attn_varlen_func(q, k, v,
68                                       max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
69                                       max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
70                                       softmax_scale=self.scale, causal=True, block_table=context.block_tables)
71        else:    # decode
72            o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
73                                        cache_seqlens=context.context_lens, block_table=context.block_tables, 
74                                        softmax_scale=self.scale, causal=True)
75        return o