Independent from-scratch Qwen-14B step reimplementation (no PyTorch, no cuBLAS, own compiler) — what matched, what it cost, and a LoRA-shape finding

Title: Independent from-scratch Qwen-14B step reimplementation (no PyTorch, no cuBLAS, own compiler) — what matched, what it cost, and a LoRA-shape finding

Not a bug report or a support request. We reimplemented the Qwen-14B forward/training step outside the PyTorch stack — own compiler, own PTX backend, own GEMM kernels, no cuBLAS call — as part of a self-hosted-language project. That project is now shelved, so I'm writing up the parts that seem useful to people working on Qwen rather than letting them sit in a dead repo. Close this if it's off-topic for your tracker.

What we implemented

hxqwen14b: RMSNorm · RoPE · GQA attention · SwiGLU, as hand-written CUDA kernels plus a code emitter in our own language (~142 KB of CUDA, ~349 KB emitter). A 32B variant emitter exists but got less exercise. LoRA training steps ran on top of it.

Everything was gated on numeric equality against a CPU reference, and the own-GEMM path trained hxqwen14b (FP32) cuBLAS-free with correctness PASS (rel-RMS ~1e-6).

Finding 1 — RoPE needed correctly-implemented transcendentals, not library ones

Our step path deliberately carries no libm: every exp/ln/sqrt/sin/cos is a fixed-iteration + − × ÷ routine so the result is bit-identical on any IEEE-754 hardware. RoPE was where this bit hardest — we had to implement sin/cos in raw PTX (Cody-Waite range reduction + odd/even Taylor, C5..C9 terms) before RoPE reached byte-exactness against the reference on silicon.

The reason this may matter to you beyond our project: platform libm transcendentals are not correctly rounded, and glibc and Darwin differ in the last ULP. Any RoPE implementation routing through libm is not bit-reproducible across machines — which is invisible until someone tries to reproduce a fine-tune exactly, on different hardware, later. We verified byte-identical training across 6 environments / 4 architecture-libc combinations (arm64-macOS, x86_64-linux glibc ×2, arm64-linux, x86_64-musl) once every transcendental on the step path was our own.

We are not claiming Qwen should do this — it costs real throughput. Just flagging that "reproducible fine-tune" and "uses libm in RoPE" are in tension, in case it's ever a requirement someone hands you.

Finding 2 — 5 of the 7 GEMMs in a LoRA step are skinny, and one tile size will not do

This is the practical one. Per-GEMM profile of a LoRA step (M=8192, K=N=4096, R=16), our tensor-core kernel vs cuBLAS-TF32:

GEMM	shape	cuBLAS (ms)	ours, big TC tile (ms)	ratio
fwd_y	(4096, 8192) k=16	0.223	0.152	0.7× — ours faster
bwd_dx	(4096, 8192) k=16	0.198	0.151	0.8× — ours faster
fwd_tmp	(16, 8192) k=4096	0.114	0.427	3.8× slower
bwd_u	(16, 8192) k=4096	0.105	0.430	4.1× slower
bwd_dB	(16, 4096) k=8192	0.121	0.689	5.7× slower
bwd_dA	(4096, 16) k=8192	0.120	0.941	7.8× slower

The two square GEMMs were faster than cuBLAS. The entire deficit was the five skinny ones: a 128×64×32 tile launched over a 16-tall output wastes ≥75% of every threadblock plus half the K-tile. R being small is the whole point of LoRA, so every LoRA step has this shape.

Dispatching skinny GEMMs (m < 128 || n < 64) to a right-sized 16×16 tiled kernel instead, and dropping a per-call cudaDeviceSynchronize, moved the full step from 2.24× slower than cuBLAS to 1.67× (46% of the gap, with shape-dispatch alone accounting for +77 of the +85 steps/s). A split-K variant closed part of the residual — cuBLAS uses split-K on tiny-output/large-K and that is exactly the LoRA backward shape.

The transferable version: if you or anyone downstream is writing custom LoRA kernels for Qwen, the backward dA/dB GEMMs are the ones to specialize. They are the worst-shaped and they are where a general tensor-core tile loses the most.

Honest scope

We never beat cuBLAS on the full step (1.59–1.67× off after tuning), and at TF32 the whole stack is 3–8× behind PyTorch. This is not a "faster Qwen" report — it is an independent-reimplementation report. If a from-scratch implementation of your architecture is useful as a cross-check on anything, the kernels are MIT.

Full post-mortem, including three headline numbers we got wrong and retracted: https://github.com/dancinlab/hexa-lang/issues/5035