feat(mlkem_top): KeyGen stage 2d (forward NTT of s/e)

Add ST_N stage: forward NTT (ntt_core mode=0, no scaling) of s[0],s[1],
e[0],e[1] in place (slots S0,S1,E0,E1). Per slot: stream 256 coeffs into
ntt_core during LOAD, collect 256 outputs back to same slot. n_pending
waits for core IDLE between slots.

Verified vs ml-kem-r golden: 1024/1024 shat/ehat coeffs exact (17318 cyc).

sync_rtl/top/TB/tb_mlkem_kg_2d_xsim.v

@@ -0,0 +1,57 @@
+// tb_mlkem_kg_2d_xsim.v - Stage 2d: verify forward NTT of s/e in mlkem_top.
+// After ST_N, slots S0,S1,E0,E1 must hold shat_0,shat_1,ehat_0,ehat_1.
+// Golden: kg_c000_sehat.hex (4 polys x 256 = 1024 lines, mod-q).
+`timescale 1ns/1ps
+module tb_mlkem_kg_2d_xsim;
+    reg clk=0, rst_n=0, start_i=0;
+    reg [255:0] d_i, z_i=0;
+    wire busy_o, done_o;
+    reg [3:0] dbg_slot_i=0; reg [7:0] dbg_idx_i=0; wire [11:0] dbg_coeff_o;
+    wire [255:0] dbg_rho_o, dbg_sigma_o;
+
+    mlkem_top #(.K(2)) dut (
+        .clk(clk), .rst_n(rst_n), .d_i(d_i), .z_i(z_i), .start_i(start_i),
+        .busy_o(busy_o), .done_o(done_o),
+        .dbg_slot_i(dbg_slot_i), .dbg_idx_i(dbg_idx_i), .dbg_coeff_o(dbg_coeff_o),
+        .dbg_rho_o(dbg_rho_o), .dbg_sigma_o(dbg_sigma_o)
+    );
+    always #5 clk = ~clk;
+
+    localparam [255:0] D_LIT = 256'h2426f1941779574d3f1b163bd57f7e173e229e630ec7f7073bdf365137c4bb6d;
+
+    reg [11:0] gold [0:1023];
+    reg [3:0]  slot_of [0:3];   // S0,S1,E0,E1
+    integer c, p, idx, errors, gi;
+
+    initial begin
+        $readmemh("sync_rtl/top/TB/vectors/kg_c000_sehat.hex", gold);
+        slot_of[0]=4'd4; slot_of[1]=4'd5; slot_of[2]=4'd6; slot_of[3]=4'd7;
+
+        d_i = D_LIT;
+        rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
+        start_i=1; @(posedge clk); start_i=0;
+        c=0; while(!done_o && c<300000) begin @(posedge clk); c=c+1; end
+        if(!done_o) begin $display("FAIL: timeout after %0d cyc", c); $finish; end
+        $display("=== Stage 2d: NTT(s/e) -> shat/ehat ===  done in %0d cyc", c);
+
+        errors = 0;
+        for (p = 0; p < 4; p = p + 1) begin
+            for (idx = 0; idx < 256; idx = idx + 1) begin
+                dbg_slot_i = slot_of[p];
+                dbg_idx_i  = idx[7:0];
+                @(posedge clk); @(posedge clk);
+                gi = p*256 + idx;
+                if (dbg_coeff_o !== gold[gi]) begin
+                    if (errors < 8)
+                        $display("  MISMATCH slot%0d[%0d]: got=%03x exp=%03x",
+                                 slot_of[p], idx, dbg_coeff_o, gold[gi]);
+                    errors = errors + 1;
+                end
+            end
+        end
+        if (errors == 0) $display("ALL TESTS PASSED (1024/1024 shat/ehat coeffs)");
+        else             $display("TESTS FAILED: %0d mismatches", errors);
+        $finish;
+    end
+    initial begin #10000000; $display("FAIL: global timeout"); $finish; end
+endmodule

sync_rtl/top/mlkem_top.v

@@ -72,6 +72,7 @@ module mlkem_top #(
     localparam ST_G      = 4'd1;   // run G(d||K), capture rho/sigma
     localparam ST_A      = 4'd2;   // generate A_hat[i][j] via SampleNTT
     localparam ST_C      = 4'd3;   // generate s[i],e[i] via CBD
+    localparam ST_N      = 4'd4;   // forward NTT of s[i],e[i] in place
     localparam ST_DONE   = 4'd15;
 
     reg [3:0]   st, st_next;
@@ -167,13 +168,41 @@ module mlkem_top #(
     // signed (two's complement) -> [0,Q): add Q when negative
     wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'(Q)) : cbd_coeff;
 
+    // ---- ntt_core: forward NTT (mode=0, no scaling) of s[i],e[i] in place ----
+    // N-stage bookkeeping: process slots S0,S1,E0,E1 (= SLOT_S0 + n_slot).
+    reg  [2:0]  n_slot;        // 0..2K (4 polys)
+    reg  [8:0]  n_ridx;        // load read index 0..256
+    reg  [7:0]  n_widx;        // output write index 0..255
+    reg         n_valid;       // feeding coeffs to ntt_core
+    reg         n_pending;     // waiting for ntt_core IDLE to start next slot
+    wire [3:0]  n_slot_addr = SLOT_S0 + {1'b0, n_slot};
+
+    wire        ntt_ready;
+    wire [11:0] ntt_coeff;
+    wire        ntt_vo;
+    wire        ntt_done;
+    wire [11:0] ntt_in = polymem[n_slot_addr*256 + n_ridx[7:0]];
+
+    ntt_core u_ntt (
+        .clk(clk), .rst_n(rst_n),
+        .coeff_in(ntt_in),
+        .valid_i(n_valid),
+        .ready_o(ntt_ready),
+        .mode(1'b0),                  // forward NTT, no scaling
+        .coeff_out(ntt_coeff),
+        .valid_o(ntt_vo),
+        .ready_i(1'b1),               // always accept output
+        .done_o(ntt_done)
+    );
+
     always @(*) begin
         st_next = st;
         case (st)
             ST_IDLE: if (start_i) st_next = ST_G;
             ST_G:    if (sha3_vo) st_next = ST_A;
             ST_A:    if (a_pair >= K*K) st_next = ST_C;
-            ST_C:    if (c_poly >= 2*K) st_next = ST_DONE;
+            ST_C:    if (c_poly >= 2*K) st_next = ST_N;
+            ST_N:    if (n_slot >= 2*K) st_next = ST_DONE;
             ST_DONE: st_next = ST_IDLE;
             default: st_next = ST_IDLE;
         endcase
@@ -196,6 +225,11 @@ module mlkem_top #(
             c_poly     <= 3'd0;
             c_widx     <= 8'd0;
             c_busy     <= 1'b0;
+            n_slot     <= 3'd0;
+            n_ridx     <= 9'd0;
+            n_widx     <= 8'd0;
+            n_valid    <= 1'b0;
+            n_pending  <= 1'b0;
         end else begin
             st <= st_next;
 
@@ -271,6 +305,52 @@ module mlkem_top #(
                     end
                 end
             end
+
+            // Arm N stage when C finishes: start NTT on slot S0
+            if (st == ST_C && st_next == ST_N) begin
+                n_slot    <= 3'd0;
+                n_ridx    <= 9'd0;
+                n_widx    <= 8'd0;
+                n_valid   <= 1'b1;       // begin loading first poly
+                n_pending <= 1'b0;
+            end
+
+            // ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ----
+            if (st == ST_N) begin
+                // LOAD phase: stream 256 coeffs into ntt_core
+                if (n_valid && ntt_ready) begin
+                    if (n_ridx == 9'd255) begin
+                        n_valid <= 1'b0;       // last coeff presented this cycle
+                        n_ridx  <= 9'd0;
+                    end else begin
+                        n_ridx <= n_ridx + 9'd1;
+                    end
+                end
+
+                // OUTPUT phase: collect 256 results, write back to same slot
+                if (ntt_vo) begin
+                    polymem[n_slot_addr*256 + n_widx] <= ntt_coeff;
+                    n_widx <= n_widx + 8'd1;   // wraps 255->0 after last
+                end
+
+                // Slot complete when ntt_core returns to DONE
+                if (ntt_done) begin
+                    if (n_slot + 3'd1 < 2*K) begin
+                        n_slot    <= n_slot + 3'd1;
+                        n_widx    <= 8'd0;
+                        n_pending <= 1'b1;     // wait one cycle for core IDLE
+                    end else begin
+                        n_slot <= n_slot + 3'd1;  // == 2K -> ST_DONE
+                    end
+                end
+
+                // Kick next slot's load once core is back IDLE
+                if (n_pending && ntt_ready && !ntt_done) begin
+                    n_valid   <= 1'b1;
+                    n_ridx    <= 9'd0;
+                    n_pending <= 1'b0;
+                end
+            end
         end
     end