graph TB
    subgraph "Before: Standard Dense Attention"
        A1[Hidden States] --> B1[QKV Projection]
        B1 --> C1[Dense Attention]
        C1 --> D1[Output Projection]
    end
    
    subgraph "After: DeepSeek Sparse Attention"
        A2[Hidden States] --> B2[QKV Projection]
        A2 --> E2[DSA Indexer]
        B2 --> E2
        E2 --> F2[Top-K Selection]
        B2 --> G2[Sparse Attention]
        F2 --> G2
        G2 --> D2[Output Projection]
        E2 --> H2[Indexer Loss]
        H2 --> D2
    end
    
    style E2 fill:#90EE90
    style F2 fill:#90EE90
    style H2 fill:#FFD700

graph TD
    A["Input: hidden_states<br/>[sq, b, hidden]"] --> B["MLA Q Compression<br/>linear_q_proj→linear_q_down_proj"]
    A --> C["MLA KV Compression<br/>linear_kv_down_proj"]
    
    B --> D["q_compressed<br/>[sq, b, q_lora_rank]"]
    C --> E["kv_compressed<br/>[sq, b, kv_lora_rank]"]
    
    D --> F["MLA Q Upsampling<br/>linear_q_up_proj + RoPE"]
    E --> G["MLA KV Upsampling<br/>linear_kv_up_proj + RoPE"]
    
    F --> H["query<br/>[sq, b, np, hn]"]
    G --> I["key<br/>[sk, b, np, hn]"]
    G --> J["value<br/>[sk, b, np, hnv]"]
    
    A --> K["x.detach()"]
    D --> L["q_compressed.detach()"]
    
    K --> M["DSAIndexer"]
    L --> M
    M --> N["Indexer Q Proj<br/>linear_wq_b<br/>[sq, b, index_n_heads, index_head_dim]"]
    K --> O["Indexer K Proj<br/>linear_wk + k_norm<br/>[sk, b, index_head_dim]"]
    K --> P["Indexer Weights<br/>linear_weights_proj<br/>[sq, b, index_n_heads]"]
    
    N --> Q["Apply RoPE"]
    O --> R["Apply RoPE"]
    Q --> S["rotate_activation<br/>&#40;Hadamard transform&#41;"]
    R --> T["rotate_activation<br/>&#40;Hadamard transform&#41;"]
    
    S --> U["Index Scores<br/>q @ k^T → ReLU → weighted sum<br/>[b, sq, sk]"]
    T --> U
    P --> U
    
    U --> V["TopK Selection<br/>[b, sq, index_topk]"]
    
    H --> W["Sparse Attention"]
    I --> W
    J --> W
    V --> W
    
    W --> X["attention_output<br/>[sq, b, hidden]"]
    
    U --> Y["KL Divergence Loss<br/>KL&#40;true_attn || index_scores&#41;"]
    V --> Y
    H --> Y
    I --> Y
    
    Y --> Z["indexer_loss<br/>scalar"]
    
    X --> AA["DSAIndexerLossAutoScaler.apply"]
    Z --> AA
    AA --> AB["Final Output<br/>&#40;with loss attached&#41;"]
    
    style K fill:#FFE4B5
    style L fill:#FFE4B5
    style M fill:#90EE90
    style U fill:#87CEEB
    style V fill:#87CEEB
    style W fill:#FFD700
    style Y fill:#FF6347
    style Z fill:#FF6347
    style AA fill:#DDA0DD

classDiagram
    class TransformerConfig {
        +int num_layers
        +int hidden_size
        +str experimental_attention_variant
    }
    
    class MLATransformerConfig {
        +int q_lora_rank
        +int kv_lora_rank
        +int dsa_indexer_n_heads
        +int dsa_indexer_head_dim
        +int dsa_indexer_topk
        +float dsa_indexer_loss_coeff
    }
    
    class Attention {
        <<abstract>>
        +forward()*
    }
    
    class MultiLatentAttention {
        +get_query_key_value_tensors()
        +forward()
    }
    
    class MLASelfAttention {
        +linear_q_proj
        +linear_kv_down_proj
        +core_attention
        +get_query_key_value_tensors(return_compressed_tensors)
    }
    
    class DSAttention {
        +indexer: DSAIndexer
        +softmax_scale: float
        +forward(query, key, value, x, qr, ...)
    }
    
    class DSAIndexer {
        +linear_wq_b
        +linear_wk
        +k_norm
        +linear_weights_proj
        +rotary_pos_emb
        +forward(x, qr, mask)
        +forward_with_scores(x, qr, mask)
        -_apply_rope()
        -_compute_index_scores()
    }
    
    class DSAIndexerSubmodules {
        +linear_wq_b: ModuleSpec
        +linear_wk: ModuleSpec
        +k_norm: ModuleSpec
        +linear_weights_proj: ModuleSpec
    }
    
    class DSAttentionSubmodules {
        +indexer: ModuleSpec
    }
    
    class MLASelfAttentionSubmodules {
        +core_attention: ModuleSpec
        +linear_q_proj
        +linear_kv_down_proj
        +q_layernorm
        +kv_layernorm
    }
    
    class RotaryEmbedding {
        +forward(seq_len)
    }
    
    class DSAIndexerLossAutoScaler {
        <<autograd.Function>>
        +forward(output, loss)$
        +backward(grad_output)$
        +set_loss_scale(scale)$
        +main_loss_backward_scale$
    }
    
    class DSAIndexerLossLoggingHelper {
        +save_loss_to_tracker()$
        +reduce_loss_in_tracker()$
        +track_indexer_metrics()$
        +tracker: dict$
    }
    
    TransformerConfig <|-- MLATransformerConfig : extends
    Attention <|-- MultiLatentAttention : extends
    MultiLatentAttention <|-- MLASelfAttention : extends
    Attention <|-- DSAttention : implements (core_attention)
    
    MLASelfAttention --> DSAttention : uses as core_attention
    MLASelfAttention --> MLASelfAttentionSubmodules : configured by
    DSAttention --> DSAIndexer : contains
    DSAttention --> DSAttentionSubmodules : configured by
    DSAIndexer --> DSAIndexerSubmodules : configured by
    DSAIndexer --> RotaryEmbedding : uses
    DSAttention ..> DSAIndexerLossAutoScaler : uses
    DSAttention ..> DSAIndexerLossLoggingHelper : logs to
    
    MLASelfAttention ..> MLATransformerConfig : reads config
    DSAttention ..> MLATransformerConfig : reads config
    DSAIndexer ..> MLATransformerConfig : reads config