[tt-train] Fused SwiGLU forward + single-sender multicast, packer L1 accumulation, batched tile multicast

[mdragulaTT]

Ticket

[Feature Request] SwiGLU forward and backward functions - #14195

Context (follow-up PR)

This is the second PR on fused SwiGLU forward. The first PR introduced the dual-path fused SwiGLU: L1-optimized fast path (full rows of XW1/XW3/M in L1 when they fit), fallback streaming path for larger dimensions, and composite backward. It did not use multicast; every core read weights from DRAM independently, so the fused forward was still slow at scale.

This PR builds on that foundation by adding: single-sender multicast for weights, packer L1 accumulation (XW1, XW3, Y), overlapping phases (dual-NOC: X/weights split), batched tile multicast, W2 double-buffering (row-major layout), and related optimizations—bringing fused SwiGLU forward in line with or ahead of composite on the Tracy profiler while reducing activation memory.

Problem description

Before any fusion, the composite path uses separate matmuls and SiLU with each core reading all weights from DRAM. The first fused PR added an L1-optimized kernel (XW1, XW3, M in L1 when they fit) but still had redundant weight DRAM reads (every core read W1/W2/W3) and no multicast, so it was slow at scale.

This (second) PR adds:

• Single-sender multicast: Core (0,0) reads W1/W2/W3 once and multicasts to all cores → factor-of-num_cores reduction in weight DRAM reads.
• Packer L1 accumulation for XW1, XW3, and Y (no partial-CB staging).
• Batched tile multicast (e.g. 16 tiles per mcast) to cut sync overhead.
• Dual-NOC / overlapping phases: X and Y on one RISC, weights on the other; W2 row-major and double-buffering where applicable.

Backward remains composite (unchanged in this PR). Note on fw+bw measurement: tt-train's autograd does not currently support mixing a custom fused forward with the original composite backward graph. As a workaround, the fused path uses a custom backward that is functionally correct but unoptimized — so fw+bw step times are dominated by the slow backward and are not meaningful for comparing forward performance. All performance data below focuses on forward-only runs. The fw+bw runs are used only to verify that loss decreases correctly with the fused forward.

---

What's changed

Fused SwiGLU forward kernel

• Single compute kernel per core: Phase A (XW1/XW3 with L1 acc), Phase B (M = SiLU(XW1)*XW3), Phase C (Y += M@W2 with L1 acc).
• L1 accumulation for XW1, XW3, and Y: no partial-CB staging; packer L1 acc writes directly into final CBs.
• Padded k-blocks: M rows are padded to block_size for matmul; padding tiles are never read by M@W2.
• Shared compute_utils.hpp helpers: pack_and_push, pack_and_push_block, pack_l1_acc_block.
• Kernel structure comments in compute and dataflow kernels (Phase A/B/C) for alignment and readability.

Single-sender multicast weight distribution

• Core (0,0) reads W1/W2/W3 from DRAM once and multicasts to all cores (loopback: sender is also a receiver).
• W1, W2, W3 share 2 semaphores (sender + receiver); weights are processed sequentially per phase.
• Reusable loopback multicast helpers in dataflow_utils.hpp (mcast_sender_read_batched_rows_and_send_loopback, mcast_receiver_reserve_and_receive).
• Single-core fallback: multicast with 0 receivers is handled transparently.

Batched tile multicast

• W1/W3: batch block_size × block_size tiles per multicast.
• W2: same batching with column-major CB layout.
• Reduces multicast synchronization overhead (e.g. ~4× fewer sync points per phase).

Dual-NOC / X vs weights split

• RISCV_1 (NOC0): X reader + Y writer (this kernel runs on all cores).
• RISCV_0 (NOC1): weight sender (core 0,0) or weight receiver (other cores).
• Aligns with tt-metal matmul 1D mcast: activations on one RISC, weights on the other.

Uniform workload padding for multicast sync

• All cores iterate max_rows_for_sync (max rows assigned to any core).
• Cores with fewer rows run padding iterations: read last valid X row, compute drains CBs but skips writing Y.
• Writer writes only actual rows (no extra outputs).

Weight shape validation

• TT_FATAL checks for fused SwiGLU weight layout: W1[embed, hidden], W3[embed, hidden], W2[hidden, embed].
• Clear error messages when shapes do not match.

Testing

• Existing SwiGLU tests (14 tests passing).
• SwiGLU_RepeatedRuns_NoHang (batch=100, 56 cores, 3 iterations).
• SwiGLU_UnbalancedWorkload_57x1x32x32 (57 rows on 56-core grid).
• Shape validation tests for incorrect weight layouts.

---

Performance

Methodology: All performance comparisons use forward-only runs (backward disabled), isolating the actual change. Training correctness (loss convergence) verified separately with fw+bw runs. Step time and loss from tt-train/scripts/plot_training_comparison.py. Memory from "Memory Usage Summary" FORWARD_PASS block in training logs. Tracy from analyze_tracy_output.py on profiler CSV exports (composite = markers mode; fused = operations mode, op name SwiGLU). SwiGLU calls per block: fixed at 1 (one MLP per transformer block, one SwiGLU per MLP).

---

NanoLlama3 (64×1×256×384, hidden=1024, 6 blocks)

Metric	Composite	Fused (SwiGLU)
Step time (fw only)	Mean 90.68 ms	Mean 84.73 ms (1.07× faster)
Memory (FORWARD_PASS)	Activations 2,068.64 MB, peak 2,128.06 MB	Activations 1,300.61 MB, peak 1,365.05 MB (~37% fewer activations, ~36% lower peak)
Tracy (per block, ms)	4.286 (std 0.009)	3.349 (std 0.001) (21.9% faster)

Loss convergence (fw+bw): Composite final loss 0.122, fused final loss 0.117 — loss decreases correctly with the fused forward. (fw+bw step times not comparable due to unoptimized backward workaround; see note above.)

Forward-only step time comparison (composite vs fused):

Peak memory breakdown (composite vs fused, forward pass):

---

TinyLlama (embed 2048, hidden 5632, 22 blocks)

Metric	Composite	Fused (SwiGLU)
Step time (fw only)	Mean 739.39 ms	Mean 983.56 ms (fused ~1.33× slower)
Memory (FORWARD_PASS)	Activations 193.04 MB, peak 6,568.44 MB	Activations 193.01 MB, peak 6,568.40 MB (~same; see note below)
Tracy (per block, ms)	4.490 (std 0.043)	17.370 (std 0.002) (~3.9× slower)

Loss convergence (fw+bw): Composite final loss 2.135, fused final loss 2.162 — loss decreases correctly.

Why no memory improvement: TinyLlama uses runner_type: memory_efficient (gradient checkpointing). In this mode, block forward runs with gradients disabled — intermediate activations (including the SwiGLU intermediates XW1, XW3, M that fused avoids materializing) are not retained during forward. They are recomputed during backward instead. So the fused path's memory advantage disappears: the intermediates it avoids writing to DRAM were already not being stored. Model (1,863 MB) + optimizer (3,709 MB) dominate the 6,568 MB forward peak; retained activations are only ~193 MB for both paths.

Forward-only step time comparison (composite vs fused):

---

Memory analysis

Numbers are from the FORWARD_PASS phase of the "Memory Usage Summary" (first-iteration fw+bw run). Activations = segment change (net DRAM allocated during forward and held for backward). Peak during forward = cumulative peak at end of FORWARD_PASS. See tt-train/docs/MEMORY_TRACKING.md. Fused reduces activation memory by avoiding materializing full XW1/XW3/M tensors to DRAM.

NanoLlama3 (runner_type: default — all forward activations retained):

Run	Activations (MB)	Peak during forward (MB)
Composite	2,068.64	2,128.06
Fused	1,300.61	1,365.05

Approx. reduction: ~37% activations, ~36% peak during forward.

TinyLlama (runner_type: memory_efficient — gradient checkpointing, intermediates not retained):

Run	Activations (MB)	Peak during forward (MB)
Composite	193.04	6,568.44
Fused	193.01	6,568.40

No activation reduction. With gradient checkpointing enabled, block-internal intermediates (XW1, XW3, M) are already discarded during forward and recomputed during backward — the fused path's memory advantage does not apply.

NanoLlama peak memory breakdown (composite vs fused):

---

Composite SwiGLU (Tracy per-block summary, 1 call per block)

Model	Mean (ms)	Std (ms)	Source
NanoLlama3	4.286	0.009	markers composite_swiglu_begin / _end, 6 blocks/step
TinyLlama	4.490	0.043	markers, 22 blocks/step

---

Fused SwiGLU (Tracy per-block summary, 1 call per block)

Model	Mean (ms)	Std (ms)	vs Composite
NanoLlama3	3.349	0.001	21.9% faster
TinyLlama	17.370	0.002	~3.9× slower

Model-size summary: Fused is faster than composite on NanoLlama3 and slower on TinyLlama. TinyLlama scaling requires further investigation (e.g. fused kernel vs batched matmuls at larger embed/hidden).

---

Summary

• Follow-up to first fused SwiGLU PR: Adds multicast, packer L1 acc, overlapping phases, and W2 multibuffering on top of the initial L1-optimized fused kernel.
• Single-sender multicast: Weight DRAM reads reduced from num_cores to 1; batched tiles reduce sync overhead.
• Forward-only performance: NanoLlama3 — Fused 7% faster end-to-end (84.73 ms vs 90.68 ms), 21.9% faster per-block on Tracy (3.35 ms vs 4.29 ms). TinyLlama — Fused ~1.33× slower end-to-end, ~3.9× slower per-block on Tracy; scaling to larger dimensions requires further investigation.
• Memory (NanoLlama, default runner): Fused ~37% fewer activations, ~36% lower peak during forward (2,069 → 1,301 MB, 2,128 → 1,365 MB). TinyLlama uses gradient checkpointing (memory_efficient runner), so intermediates are not retained and fused shows no memory advantage.
• Backward workaround: tt-train's autograd does not currently support mixing a custom fused forward with the original composite backward. The fused path uses a slow custom backward as a workaround — fw+bw step times are not meaningful for forward comparison. Loss converges correctly in fw+bw runs for both models.

---

Checklist

• 
• 
• 
• New/Existing tests provide coverage for changes

[Copilot]

Pull request overview

This PR optimizes the SwiGLU forward operation by implementing multicast-based weight distribution to reduce DRAM bandwidth consumption. Instead of each core independently reading weight matrices (W1, W2, W3) from DRAM, leftmost column cores (x==0) read weights and multicast them to other cores in their row via NoC. The PR also introduces uniform workload padding to handle imbalanced batch sizes that don't divide evenly across cores.

Changes:

• Split reader kernel into sender (reader_swiglu_fw_sender.cpp) and receiver (reader_swiglu_fw_receiver.cpp) variants
• Implemented multicast synchronization using shared semaphores across W1/W2/W3 (reused since they execute sequentially)
• Added uniform workload padding mechanism where all cores in a grid row loop for max_rows_for_sync iterations to maintain multicast synchronization
• Added reusable multicast helper functions in dataflow_utils.hpp
• Updated test suite with SwiGLU_RepeatedRuns_NoHang test for imbalanced workloads and corrected NanoLlama test naming
• Added const-correctness improvements to compute kernels

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
swiglu_op_test.cpp	Added test for repeated runs with imbalanced batch sizes; updated test naming from NanoGPT to NanoLlama
swiglu_fw_program_factory.hpp	Extended shared variables to track sender/receiver kernels, core counts, and multicast flag
swiglu_fw_program_factory.cpp	Implemented multicast topology setup, kernel splitting, uniform padding logic, and runtime argument assignment
reader_swiglu_fw_sender.cpp	New sender kernel that reads from DRAM and multicasts W1/W2/W3 to receivers in same row
reader_swiglu_fw_receiver.cpp	New receiver kernel that receives W1/W2/W3 via multicast and only reads X from DRAM
reader_swiglu_fw_interleaved_start_id.cpp	Removed (replaced by sender/receiver split)
writer_swiglu_fw_interleaved_start_id.cpp	Minor const-correctness improvements to runtime arguments
swiglu_fw_kernel_m_fits_l1.cpp	Added const qualifiers to function parameters and local register indices
swiglu_fw_kernel.cpp	Added const qualifiers to function parameters and local register indices
dataflow_utils.hpp	Added multicast synchronization primitives and sender/receiver helper functions; improved const-correctness throughout