Skip to content

Latest commit

 

History

History
69 lines (56 loc) · 2.23 KB

File metadata and controls

69 lines (56 loc) · 2.23 KB
sequenceDiagram
    participant Caller
    participant DLayer as DeepseekV2DecoderLayer.forward

    Caller->>+DLayer: (positions, hidden_states, ...)

    DLayer->>DLayer: 1. Comm.prepare_attn(hs, res)
    
    box "2. self_attn.forward (DeepseekV2AttentionMLA)"
        participant Attn as self_attn.forward
    end
    DLayer->>+Attn: forward(positions, hs, ...)
    
    Attn->>Attn: dispatch_attn_forward_method(forward_batch)
    Note right of Attn: Determines path (MHA, MLA, etc.)
    
    alt AttnForwardMethod is MLA
        Attn->>Attn: forward_absorb_prepare(...)
        Attn->>Attn: forward_absorb_core(...)
        Note right of Attn: Uses weight absorption
    else AttnForwardMethod is MHA
        Attn->>Attn: forward_normal_prepare(...)
        Attn->>Attn: forward_normal_core(...)
        Note right of Attn: Standard Multi-Head Attention
    end
    
    Attn-->>-DLayer: hidden_states

    DLayer->>DLayer: 3. Comm.prepare_mlp(hs, res)
    
    box "4. mlp.forward"
        participant MLP as mlp.forward
    end
    DLayer->>+MLP: forward(hidden_states, ...)
    
    alt is_dense (uses DeepseekV2MLP)
        MLP->>MLP: gate_up_proj(hs) -> act_fn -> down_proj
    
    else is_layer_sparse (uses DeepseekV2MoE)
        box "Within DeepseekV2MoE.forward"
            participant MoE as MoE.forward
        end
        MLP->>+MoE: (hidden_states, ...)
        
        alt _enable_deepep_moe
            box "path: forward_deepep"
                participant DeepEP as forward_deepep
            end
            MoE->>+DeepEP: (hs, ...)
            DeepEP->>DeepEP: select_experts -> dispatch -> experts -> combine
            Note right of DeepEP: Expert Parallelism via All-to-All
            DeepEP-->>-MoE: result
        else not _enable_deepep_moe
            box "path: forward_normal"
                participant Normal as forward_normal
            end
            MoE->>+Normal: (hidden_states)
            Normal->>Normal: gate -> experts -> all_reduce
            Note right of Normal: Standard MoE execution
            Normal-->>-MoE: result
        end
        MoE-->>-MLP: hidden_states
    end
    MLP-->>-DLayer: hidden_states

    DLayer->>DLayer: 5. Comm.postprocess_layer(hs, res)

    DLayer-->>-Caller: final hidden_states
Loading