verl/models/llama/megatron/layers/parallel_attention.py

• LlamaRotaryEmbedding · class · L35-L69 — class LlamaRotaryEmbedding(nn.Module)
• init · method · L37-L49 — def init(self, dim, max_position_embeddings=2048, base=10000, device=None)
• _set_cos_sin_cache · method · L51-L59 — def _set_cos_sin_cache(self, seq_len, device, dtype)
• forward · method · L61-L69 — def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size]
• LlamaLinearScalingRotaryEmbedding · class · L72-L88 — class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding)
• init · method · L75-L77 — def init(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0)
• _set_cos_sin_cache · method · L79-L88 — def _set_cos_sin_cache(self, seq_len, device, dtype)
• LlamaDynamicNTKScalingRotaryEmbedding · class · L91-L113 — class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding)
• init · method · L94-L96 — def init(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0)
• _set_cos_sin_cache · method · L98-L113 — def _set_cos_sin_cache(self, seq_len, device, dtype)
• rotate_half · function · L116-L120 — def rotate_half(x)
• apply_rotary_pos_emb · function · L123-L128 — def apply_rotary_pos_emb(q, k, cos, sin, position_ids)
• repeat_kv · function · L131-L140 — def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor
• ParallelLlamaAttention · class · L143-L280 — class ParallelLlamaAttention(nn.Module)
• init · method · L146-L202 — def init(self, config: LlamaConfig, megatron_config: ModelParallelConfig)
• _init_rope · method · L204-L229 — def _init_rope(self)
• _shape · method · L231-L232 — def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int)
• forward · method · L234-L280 — def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]
• apply_rotary_pos_emb_rmpad · function · L299-L312 — def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length)
• apply_rotary_pos_emb_rmpad_flash · function · L320-L335 — def apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seqlen)
• ParallelLlamaAttentionRmPad · class · L338-L418 — class ParallelLlamaAttentionRmPad(ParallelLlamaAttention)
• forward · method · L340-L418 — def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: torch.Tensor = None, max_seqlen_in_batch: int = None)