/* verilator lint_off BLKLOOPINIT */
// Transformer Inference Engine — Prefill + Decode phases
module transformer_infer #(parameter DW=16, HEADS=8, DK=64, SEQ=128, FFN_DIM=256)(
    input clk, rst_n, input en, input start,
    input is_decode, // 0=prefill, 1=decode (single token)
    // Token embedding input
    input signed [DW-1:0] token_embed [0:HEADS*DK-1], input embed_valid,
    // KV-Cache interface
    output reg [9:0] kv_cache_wr_addr, output reg kv_cache_wr_en,
    output reg signed [DW-1:0] kv_cache_wr_data,
    input signed [DW-1:0] kv_cache_rd_data [0:HEADS*DK*2-1],
    // Output logits
    output reg signed [DW-1:0] out_logits [0:HEADS*DK-1],
    output reg out_valid, output reg done
);
    reg [3:0] state;
    reg [9:0] seq_pos, head_cnt;
    reg signed [DW*2-1:0] qk_score, attn_weight;
    reg signed [DW-1:0] attn_out [0:DK-1];
    integer hi, di;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin state<=0; seq_pos<=0; head_cnt<=0; qk_score<='0; attn_weight<='0;
            kv_cache_wr_addr<=0; kv_cache_wr_en<=0; kv_cache_wr_data<='0; out_valid<=0; done<=0;
            for(di=0;di<DK;di++) attn_out[di]<='0; for(di=0;di<HEADS*DK;di++) out_logits[di]<='0; end
        else if(en) case(state)
          0: if(start) begin state<=1; seq_pos<=0; head_cnt<=0; end
          1: if(embed_valid) begin // MHA: compute Q*K^T for each head
                // Simplified: single dot product per head
                qk_score <= token_embed[head_cnt*DK+0] * kv_cache_rd_data[head_cnt*DK*2+0];
                kv_cache_wr_addr <= seq_pos * HEADS * DK + head_cnt * DK;
                kv_cache_wr_en <= 1;
                kv_cache_wr_data <= token_embed[head_cnt*DK];
                head_cnt <= head_cnt + 1;
                if(head_cnt >= HEADS-1) begin head_cnt<=0; state<=2; end
            end
          2: begin // Softmax + V weighting (simplified)
                attn_weight <= qk_score >>> 3; // /sqrt(DK)
                for(di=0;di<DK;di++) attn_out[di] <= attn_out[di] + (attn_weight[DW*2-1:DW] * kv_cache_rd_data[head_cnt*DK*2+DK+di]);
                state <= 3;
            end
          3: begin // FFN: two linear layers with GeLU (simplified)
                for(di=0;di<HEADS*DK;di++) out_logits[di] <= attn_out[di%DK]; // Simplified passthrough
                seq_pos <= seq_pos + 1;
                if(is_decode || seq_pos >= SEQ-1) state <= 4;
                else state <= 1;
            end
          4: begin out_valid <= 1; done <= 1; end
        endcase
    end
endmodule