Subgraph-fusion-based Mega-kernel Generation

Project objective

Create a prototype that transforms a provided PyTorch computation subgraph into one or more fused mega-kernels covering forward and backward computation. Reduce memory traffic or memory footprint and improve execution throughput while remaining numerically consistent with the unfused reference.

Identify fusible operators, choose defensible fusion boundaries, and generate or construct both forward and backward kernels. Evaluate on a single GPU with a significant element-wise workload. Measure latency or throughput, peak memory, and memory-traffic indicators across representative inputs. The kernel language, compiler, scheduling, and memory-management approach are open design choices.

Use the supplied RankMixerBlock and DensePFFNv2 implementation as the workload and numerical reference. The student package preserves the full block implementation and provides a small public launch. Full regression matrices and held-out inputs remain with the instructor.

Workload and interfaces

The benchmark entry point is rank_mixer_src/bench_rank_mixer.py. Optimize its RankMixerBlockFwd computation, or replace fusible subgraphs while keeping the full block's behavior:

# Block entry point:
def forward(self, inputs, overlap="none", comp_only=False):
    # inputs and output: [B,T,D]
    ...

# Dense FFN subgraph:
def dense_ffn_v2_compute(inputs, w1, b1, wout, bout, norm_x, norm_gate):
    ...

The default execution is EP=1, overlap=none, forward+backward. It is fixed by the current runner. NPROC=8 launches eight independent GPU workers with EP disabled; it does not turn this into an eight-way expert-parallel workload. Use NPROC=1 for the single-GPU project experiment.

Dense parameters are w1:[T,D,2H], b1:[T,2H], wout:[T,H,D], bout:[T,D]. The two inner LayerNorm modules each have H-element weight and bias. All three outer NezhaLayerNorm modules have D-element gamma and beta. Input and output are [B,T,D]; the runner supplies a same-shape, same-dtype upstream gradient and differentiates the complete block.

Exact computation

out1 = token_mixer(norm0(inputs))
out1 = norm1(out1 + inputs)
out2 = dense_ffn_v2_compute(out1, w1, b1, wout, bout, norm_x, norm_gate)
output = norm2(out1 + out2)

The three outer normalizations are LayerNorm with epsilon=1e-3, gamma initialized to ones, and beta to zeros. Token mixing is reshape(B,T,T,D/T), swap the two T axes, then reshape to [B,T,D]. For example, on [1,2,4], token rows [[0,1,2,3],[4,5,6,7]] become [[0,1,4,5],[2,3,6,7]] under token mixing alone.

The dense subgraph is:

Z[b,t,:] = inputs[b,t,:] @ w1[t,:,:] + b1[t,:]
(U,G) = split Z into two H-dimensional halves
LN(z) = ((z-mean(z))/sqrt(mean((z-mean(z))^2)+epsilon))*gamma + beta
hidden = LN_x(U) * silu(LN_gate(G))
output[b,t,:] = hidden[b,t,:] @ wout[t,:,:] + bout[t,:]

The inner norm_x and norm_gate are PyTorch LayerNorm with epsilon=1e-5, not RMSNorm. Preserve cast placement and arithmetic behavior of the selected backend. EP=1 makes both FFN exchanges identity.

Weights and biases use the included he-uniform initializer: fan_in=product(shape[:-1]), then uniform values in [-sqrt(6/fan_in),sqrt(6/fan_in)]. The process seed is seed+rank; model construction precedes torch.randn input generation. The fixed upstream gradient uses a separate CUDA generator seeded with seed+1000+rank.

Run the benchmark

Extract the package and enter research_topic_1/. Use Linux, Python 3.12, a CUDA 12.6 compatible driver, and a Hopper/SM90 GPU for the documented environment. Install the exact requirements: PyTorch 2.6.0+cu126, NCCL 2.21.5, Triton 3.2.0, and NumPy 1.26.4.

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python preflight.py --target cuda
bash -n scripts/run_minimal_local.sh scripts/run_grouped_local.sh scripts/run_prod_local.sh
NPROC=1 bash scripts/run_prod_local.sh --list-backends
NPROC=1 bash scripts/run_minimal_local.sh
NPROC=1 bash scripts/run_grouped_local.sh

The default scripts use NPROC=8 if unset; explicitly set NPROC=1 for this project's single-GPU runs. For the documented repeated-run configuration:

WARMUP=5 NUM_TESTS=20 NPROC=1 bash scripts/run_grouped_local.sh

For explicit shape and precision controls accepted by the actual parser:

NPROC=1 bash scripts/run_prod_local.sh --batch-size 1024 --tokens 64 --dim 1280 --k 2 --ffn-type dense_v2 --gemm-backend grouped --warmup 1 --num-tests 1
NPROC=1 bash scripts/run_prod_local.sh --dtype float32 --gemm-backend bmm
NPROC=1 bash scripts/run_prod_local.sh --dtype float32 --tf32 --gemm-backend bmm
NPROC=1 bash scripts/run_prod_local.sh --dtype float16 --gemm-backend bmm

The current parser fixes EP=1 and does not accept --ep-group-size, --backward, --check-correctness, or --sleep-ms; passing those options fails. Every iteration already includes backward. Sparse MoE is an available extension using --ffn-type sparse_moe --moe-experts 2 or 4 with grouped GEMM. Equal-size native grouped GEMM uses torch.bmm; variable-size groups use torch.matmul and retain autograd support.

Public checks and instructor coverage

python public_check.py

The public CUDA run uses (B,T,D,k)=(4,16,128,2), seed 7101, dense_v2 with bmm, EP=1, two warmups and three measurements. It executes the full block including backward. The instructor retains the complete dense-bmm, dense-grouped and sparse-MoE-grouped matrix and adds different seeds and shapes. Implement the full documented graph rather than specializing to this public instance.

Held-out dense workloads use 1<=B<=1024, 8<=T<=64, 64<=D<=1280, D%T==0, and k in {2,3}. The default production shape remains a performance target. The interface, initializer and timing behavior are identical across public and private runs. Sparse MoE remains an extension. Full correctness still needs the candidate-versus-reference evidence described below.

Existing checks and numerical validation

The runner performs an untimed forward/backward pass, checks that the output shape matches the input, and then measures repeated forward/backward passes. It does not include a candidate-versus-reference numerical checker, a fixed test-case suite, or prescribed error tolerances. A successful benchmark launch is not numerical acceptance. No replacement 24-case suite or new fixed tolerance is imposed here.

Your project must provide repeatable output and gradient checks against an unchanged copy of this block. Use identical input, state_dict, precision, and upstream gradient for both paths. Compare input gradients and all used parameter gradients. State and justify your proposed tolerances in the technical proposal, report measured errors, and reject non-finite outputs. Keep this additional project validation separate from the checks actually implemented by the runner.

The instructor also retains the determinism diagnostic: repeated gradient runs are compared using torch.equal. This diagnostic does not establish candidate numerical acceptance.

CPU environment check

For development without a GPU, create a Python 3.12 environment, install requirements-cpu.txt, and run python preflight.py --target cpu. This executes the actual full RankMixer block on CPU with EP=1 at (B,T,D,k)=(2,4,16,2), including input and parameter gradients in float32 and bf16. It does not test GPU kernels or timing. See ENVIRONMENT.md for setup and environment verification for current machine results.

Exact timing behavior

Each measured call clears inputs.grad, calls block.zero_grad(set_to_none=True), runs block(inputs, overlap="none"), and calls out.backward(grad_output). The runner fixes eager execution (no_compile=True), L2 flushing on, and prewarm sleep to zero.

The timing utility synchronizes, warms up, flushes L2 and performs an on-stream all-reduce barrier before each start event, then records CUDA events around the complete forward/backward call. The flush and barrier are outside the timed region. It synchronizes before reading event durations, converts milliseconds to seconds, and drops the first measurement when more than one measurement exists. It reports arithmetic mean, minimum, and maximum; printed values use microseconds. Do not describe these outputs as median/p95 or forward-only timings.

The reference output includes B=1024 T=64 D=1280 k=2 H=2560 T_local=64, zero A2A bytes per rank, and mode=fwd+bwd. Measure the fused implementation under the same settings. Report peak memory and memory-traffic indicators separately using a profiler, with explicit methodology and units. Preserve raw measurements if adding further statistics.

Required submission

Submit a proposal explaining fusion eligibility, boundaries, forward/backward construction, memory strategy, tolerances, and evaluation methodology; the prototype with kernel code or generator, environment, configurations, numerical validation, raw timings, profiler summaries, and commands; and a separate report.

The demo must accept the subgraph, produce or invoke the fused implementation, check outputs and gradients, and display throughput and memory results. Explain improvement mechanisms, unsupported cases, and limitations. Evaluation considers correctness, sound and reusable fusion, overhead reduction, measured benefit, robustness, reproducibility, and systems analysis.