Cross-Architecture Kernel Porting

LayerNorm implemented across three abstraction layers: CUDA C++, hand-written PTX (SM-120), and Ascend C (CANN 8.x).

---

Why LayerNorm?

LayerNorm is the canonical memory-bandwidth-bound kernel in LLM inference. Every transformer block applies it twice (pre-attention, pre-FFN), and it touches the entire hidden dimension on every token. Optimizing it requires maximizing effective memory bandwidth — not FLOPS.

This makes it the ideal benchmark for comparing how different architectures and abstraction layers handle the fundamental bottleneck: moving data from memory to compute units and back.

Tensor Cores are irrelevant here. mma.sync accelerates matrix multiply-accumulate. LayerNorm is a reduction + element-wise operation. There is no mathematically valid mapping. The optimization lever is LDG.128 (128-bit vectorized loads), not compute throughput.

---

The Three Implementations

1. CUDA C++ (High-Level)

Aspect	Detail
File	cuda/layernorm.cu → layernorm_cuda_kernel
Approach	One block per row, 256 threads, warp reduction via __shfl_down_sync, cross-warp reduction via shared memory
Memory	Scalar float loads (compiler auto-vectorizes to LDG.128 at -O3)
Passes	3 (mean → variance → normalize+scale+shift)
Key insight	Clean, portable, compiler-optimizable. nvcc -O3 already generates vectorized loads and optimal register allocation.

2. PTX Assembly (Low-Level)

Aspect	Detail
File	cuda/layernorm.cu → layernorm_ptx_kernel
Reference	cuda/layernorm.ptx (annotated instruction guide)
Approach	Same algorithm, explicit ld.global.nc.v4.f32 and st.global.wb.v4.f32
Key instructions	ld.global.nc.v4.f32 (non-coherent 128-bit load), st.global.wb.v4.f32 (write-back 128-bit store), shfl.sync.down.b32
Key insight	.nc bypasses L1 cache. This helps streaming kernels but hurts multi-pass kernels like LayerNorm where data is reused across passes.

3. Ascend C (Huawei CANN)

Aspect	Detail
File	ascend/layernorm_ascend.c
Target	Ascend 910B (DaVinci V300)
Approach	Explicit triple-pipe model: MTE (DMA) → Vector → MTE
Key APIs	DataCopy (GM↔UB), ReduceSum, SetFlag/WaitFlag (pipe sync)
Key insight	No implicit memory access. Every byte must be explicitly DMA'd from HBM to the 256 KB Unified Buffer. Adds complexity but enables fine-grained pipe overlap.

---

Architecture Comparison

Feature	NVIDIA SM-120 (Blackwell)	Huawei DaVinci V300 (Ascend 910B)
Execution model	Warp-synchronous (32 threads)	Explicit triple-pipe (MTE/Vector/Cube)
On-chip memory	Shared memory (up to 228 KB)	Unified Buffer (256 KB per AI Core)
Memory access	Implicit (pointer dereference)	Explicit (DataCopy via DMA engine)
Synchronization	__syncthreads(), bar.sync	SetFlag / WaitFlag between pipes
HBM bandwidth	448 GB/s (GDDR7)	1,200 GB/s (HBM2e)
Cache hierarchy	L1 (128 KB) → L2 (96 MB) → GDDR7	No traditional cache; UB is software-managed SRAM
Reduction primitive	__shfl_down_sync (warp interconnect)	ReduceSum (hardware tree in Vector pipe)
Vectorized load	ld.global.nc.v4.f32 (LDG.128)	DataCopy (DMA burst, configurable size)

---

Results

Hardware: NVIDIA GeForce RTX 5060 (SM 12.0, 8.5 GB GDDR7) Environment: WSL2 (Ubuntu 22.04), CUDA 12.8, nvcc 12.8.93 Workload: B=4096, N=4096, float32, 100 iterations per run

Benchmark (3 runs, averaged)

Kernel	Latency (ms)	BW (GB/s)	% of GDDR7 Peak
CUDA C++ (scalar)	0.423 ± 0.020	318.0	71.0%
PTX (LDG.128)	0.441 ± 0.010	304.3	67.9%
Ascend 910B (est.)	~0.11	~864	72% of HBM2e

Raw Data

Run	CUDA C++ (ms)	CUDA C++ (GB/s)	PTX (ms)	PTX (GB/s)
1	0.398	337.6	0.427	314.3
2	0.447	300.0	0.449	298.9
3	0.424	316.4	0.448	299.7
Avg	0.423	318.0	0.441	304.3

Key Finding: Hand-Written PTX is 4.3% Slower

The hand-tuned PTX kernel with explicit ld.global.nc.v4.f32 is consistently slower than the CUDA C++ version across all 3 runs. This is not noise — it is a reproducible, structurally explained result.

Root causes:

1. nvcc -O3 already auto-vectorizes. The compiler transforms scalar float loads into LDG.128 in the generated SASS. Inspect with:

   make ptx && grep "ld.global" cuda/layernorm_extracted.ptx

   The "scalar" kernel already has vectorized loads in the output.

2. .nc (non-coherent) hurts multi-pass kernels. LayerNorm reads the same data 3 times (mean → variance → normalize). The .nc modifier bypasses L1 cache, forcing a full GDDR7 round-trip on every pass. The compiler-generated code uses .ca (cache-all), which keeps the data hot in L1 for passes 2 and 3.

3. asm volatile blocks compiler optimizations. Inline PTX is an opaque barrier to nvcc. It cannot reorder instructions, optimize register allocation, or apply ILP around the asm block.

Staff Engineer takeaway: Knowing when NOT to hand-tune is as important as knowing how. Manual PTX wins for irregular access patterns (scatter/gather, graph traversal). For linear streaming with data reuse, the modern compiler already makes the right trade-offs.

---

Build & Run

Prerequisites

Component	Version
GPU	NVIDIA RTX 50xx (SM-120)
CUDA Toolkit	12.8+ (tested: 12.8.93)
Python	3.10+
PyTorch	2.1+ (for Python benchmarks)

Compile and run

# Build the CUDA benchmark
make cuda

# Run the C++ benchmark (B=4096, N=4096)
./cuda/layernorm_bench 4096 4096

# Run the Python benchmark (includes PyTorch comparison)
make bench

# Extract PTX assembly for inspection
make ptx
grep "ld.global" cuda/layernorm_extracted.ptx

# Run the Ascend emulator (no hardware required)
make bench-ascend

---

Key Takeaways

1. Memory is the bottleneck on both architectures. LayerNorm does ~2 FLOPs per byte. The RTX 5060 has 9+ TFLOPS but only 448 GB/s of bandwidth. The kernel is bandwidth-bound by ~20×. No amount of compute optimization matters.

2. The modern compiler is excellent for memory-bound kernels. nvcc -O3 auto-vectorizes to LDG.128 and achieves 71% of GDDR7 peak bandwidth. Hand-written PTX with .nc is 4.3% slower because it bypasses L1 cache on a multi-pass kernel.

3. .nc is not universally better. Non-coherent loads are optimal for single-pass streaming (e.g., memcpy, element-wise ops). For multi-pass kernels with data reuse (LayerNorm, Softmax, RMSNorm), L1 caching (.ca) wins.

4. Ascend's explicit model is harder but more predictable. NVIDIA's implicit caching makes performance easier to achieve but harder to reason about. Ascend's explicit DataCopy + SetFlag/ WaitFlag model forces you to think about data movement, which produces more predictable latency at the cost of developer effort.

5. Tensor Cores don't apply to LayerNorm. mma.sync is for GEMM. LayerNorm is a reduction. Forcing Tensor Cores into this kernel would be technically dishonest and a red flag in a systems engineering interview.

---

Repository Structure

cross-arch-kernel-porting/
├── README.md
├── Makefile
├── LICENSE
├── cuda/
│   ├── layernorm.cu          # CUDA C++ + inline PTX kernels + benchmark
│   └── layernorm.ptx         # Annotated PTX instruction reference
├── ascend/
│   └── layernorm_ascend.c    # Ascend C (CANN 8.x) implementation
└── benchmarks/
    ├── bench_cuda.py          # PyTorch comparison benchmark
    └── bench_ascend.py        # Ascend 910B cycle-approximate estimator

---

License

MIT

---

Author

João Felipe de Souza — Staff ML Systems Engineer github.com/JohnScheuer