Method forward_with_attn_bias

gpu/model.py:269–281

(
        self,
        token_values: torch.Tensor,
        attn_bias: AttnBias,
        cache: list[LayerCache],
    )

    @torch.no_grad()
    def forward_with_attn_bias(
        self,
        token_values: torch.Tensor,
        attn_bias: AttnBias,
        cache: list[LayerCache],
    ) -> torch.Tensor:
        h = self.tok_embeddings(token_values)

        for i, layer in enumerate(self.layers):
            h = layer(h, cache[i], attn_bias)

        logits = self.output(self.norm(h))
        return logits.float()

    def forward(
        self,