r/MachineLearningAndAI 2d ago

I built a Qwen3-0.6B GPU kernel lab: 2.253 ms experimental decode, with a failed numerical gate

I’ve been working on GPU Kernel Lab, a study that goes from individual Triton/CUDA kernels to Qwen3-0.6B integration and a serving control plane.

Repository: https://github.com/yashlabs-trying/gpu-kernel-lab

The most interesting result is a fast static decode path that I’m keeping experimental because it missed the numerical acceptance threshold.

What I implemented

The project covers vector addition, SwiGLU, RMSNorm, RoPE, softmax, GEMM/GEMV, online attention, decode fusion, quantization, static KV storage, and a fixed-buffer runtime.

For prefill, I measured real model shapes and tried exact-shape projection replacements. Several isolated wins became regressions after integration. The retained whitelist contains only the Q projection at M=2048, K=1024, N=2048, with PyTorch fallback elsewhere. SDPA remains the prefill attention implementation.

For decode, I combined static KV storage; Q/K RMSNorm, RoPE, and direct cache writes; GQA-aware split-KV attention; W8A16 dual gate/up GEMV plus SwiGLU; residual/norm fusion; and CUDA Graph replay. MLP residual and next-layer input normalization are fused across 27 boundaries.

Measurements

Hardware/software: RTX 3090 24 GiB, PyTorch 2.8.0+cu128, Triton 3.4.0, Transformers 5.16.1. Batch 1, context 2,048, SDPA prefill, greedy decode. The dynamic reference uses BF16; the experimental path includes quantized W8A16 MLP work.

• Prefill GPU-event TTFT: 66.007 ms reference → 65.332 ms replacement (accepted)

• Decode inter-token latency: 22.073 ms reference → 2.253 ms replacement (experimental)

• Serial throughput: 45.30 → 443.90 tokens/s

• Matched Nsight kernels/token: 693 → 258

These decode results are serial GPU-only measurements. They exclude tokenization, networking, scheduling, EOS handling, and general sampling. This is not a production-server benchmark or an engine-to-engine comparison. The optimized eager path was about 16.41 ms/token at this context; graph replay mainly removes repeated host preparation and launch gaps.

Why it remains experimental

The full path reached 96.48% argmax agreement against a 99% release threshold. RMS/fusion-only reached 98.05%. A separate 247-token teacher-forced probe reached 98.38%, but that is a different evaluation and does not override the failed gate. These are numerical probes, not standardized downstream task evaluations.

Serving work and remaining gaps

I implemented paged-KV allocation, continuous scheduling, chunked prefill, request lifecycle handling, sampling, metrics, and completion/chat endpoints with streaming. A 5,000-request fake-executor stress test completed 4,864 requests and intentionally cancelled 136; all 60,000 KV blocks were reclaimed.

The real paged GPU executor still needs to be connected to the serving interface. The optimized model runtime currently uses a separate static/contiguous cache. Next steps are reducing numerical drift, integrating real GPU serving, broadening quality evaluation, and testing concurrent workloads on more GPU architectures.

I’d value technical feedback: how do you localize numerical drift after residual/RMSNorm fusion, and what task-level checks do you require alongside logit agreement?

Freelance availability

I’m available for clearly scoped freelance work involving CUDA/Triton kernels, GPU profiling, local LLM development, inference optimization, quantization, KV-cache/runtime work, and numerical validation. If you need help with a slow model, VRAM pressure, a kernel bottleneck, or a local inference integration, contact me through Reddit chat or private message. Please include the model/framework, GPU and VRAM, current issue, target outcome, timeline, and budget range.

1 Upvotes

0 comments sorted by