Skip to content

vllm.model_executor.models.ricursive_dflash2

DFlash2 drafter for Ricursive Ripple, whose no_rope_layers are NoPE.

_NoRope

Bases: Module

Identity rotation that keeps the attributes the fused context-K/V precompute reads from every layer's RoPE.

Source code in vllm/model_executor/models/ricursive_dflash2.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class _NoRope(nn.Module):
    """Identity rotation that keeps the attributes the fused context-K/V
    precompute reads from every layer's RoPE."""

    def __init__(self, rope: nn.Module) -> None:
        super().__init__()
        self.head_size = rope.head_size
        self.is_neox_style = rope.is_neox_style
        self.cos_sin_cache = rope.cos_sin_cache

    def forward(
        self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return query, key