INT4 GEMM Kernel • PTX Intrinsics • Tensor Cores • End-to-End LLM Integration
Custom PTX INT4 GEMM kernel with end-to-end LLM integration on NVIDIA Turing (SM75).
A deep dive into register-level dequantization, WMMA vs raw PTX intrinsics, and the structural challenges of low-bit inference kernels on consumer GPUs.
This project implements a fused CUDA kernel for INT4 weight-only quantization, targeting the Memory Bandwidth Bottleneck in LLM decode-phase inference. The kernel:
- Reads packed INT4 weights from global memory
- Dequantizes on-the-fly in registers (no shared memory round-trip)
- Executes matrix multiplication via Tensor Cores (WMMA + PTX mma.sync.aligned)
- Integrates end-to-end with Qwen2-0.5B via a custom Int4Linear module
Replaced all 24 down_proj MLP layers with the custom INT4 kernel:
| Metric | FP16 Baseline | INT4 Custom | Delta |
|---|---|---|---|
| VRAM | 950.17 MB | 810.23 MB | -139.94 MB (-14.7%) |
| Output Coherence | Reference | Coherent | ✓ |
| Numerical Fidelity | Reference | cos > 0.99 | ✓ |
Prompt: "The capital of France is"
- FP16: "The capital of France is Paris. [multiple choice format]"
- INT4: "The capital of France is Paris. The population of Paris is about 2,000,000 people."
Both outputs are semantically correct. The quantized model successfully identifies Paris as the capital.
Matrix Size: 4096x4096 | Precision: INT4 Weights / FP16 Activations
| Batch (M) | cuBLAS FP16 | Custom PTX INT4 | Notes |
|---|---|---|---|
| 1 (Decode) | 0.12 ms | 0.36 ms | Memory-bound regime |
| 128 (Prefill) | 0.18 ms | 1.36 ms | Compute-bound regime |
The custom kernel does not outperform cuBLAS in isolation. The real value emerges at the model level, where 4x smaller weights reduce total VRAM footprint and enable larger models to fit in constrained memory.
Implemented nvcuda::wmma API with INT4 packing, sign extension, and register-level dequantization. Kernel was correct but limited by WMMA API overhead.
Migrated to mma.sync.aligned.m16n8k8 PTX inline assembly with manual fragment layout control. This exposed critical structural constraints of SM75:
The M=1 Problem
- mma.sync.aligned requires all 32 warp threads to hold valid data
- For M=1 (decode phase), only thread 0 has real input; threads 1-31 hold zeros
- On Turing SM75, this triggers undefined behavior leading to NaN propagation
- Fix: Automatic padding to M=16 in the Python wrapper, transparent to callers
Debugging Methodology
- Progressive isolation testing (M, N, K, group_size individually)
- Timeout-guarded kernel execution to detect GPU hangs
- Bit-identical validation against Python reference dequantization
- Vectorized dequant reference for fast large-shape testing
Built Int4Linear module that transparently replaces nn.Linear layers. Successfully integrated with Qwen2-0.5B, achieving 14.7% VRAM reduction with coherent generation quality.
Turing (SM75) lacks several architectural features that Ampere+ (SM80+) provides:
- No native cp.async — Async weight prefetching must be emulated with manual double-buffering
- No INT4 MMA operands — Ampere supports mma.sync.aligned.*.s4, keeping weights in INT4 through the MMA pipeline. Turing forces dequantization to FP16 before the MMA instruction
- No warp specialization — Cannot dedicate warp groups to loading vs computing
- Higher mma.sync overhead — Instruction dispatch cost is not amortized by native low-bit operations
These limits mean that on Turing, custom INT4 kernels are fighting against fundamental architectural overhead that Ampere/Hopper resolve natively.
- PTX MMA Integration: Raw mma.sync.aligned.m16n8k8 inline assembly
- Register Dequantization: No shared memory round-trip for dequantized weights
- AWQ-Compatible Scales: Per-channel/per-group scaling factors (group_size=128)
- Automatic Padding: Wrapper handles arbitrary M dimensions transparently
- PyTorch Integration: JIT compilation, transparent nn.Linear replacement
- End-to-End Validation: Real Qwen2-0.5B forward pass with coherent output
- src/kernels/int4_gemm_ptx.cu — Main PTX MMA kernel
- src/awq_loader.py — INT4 packing (AWQ-compatible)
- src/kernel_wrapper.py — Python wrapper with automatic M-padding
- benchmarks/qwen2_e2e.py — End-to-end Qwen2 integration
- benchmarks/test_wrapper_final.py — Correctness validation suite
Install dependencies:
pip install torch numpy ninja transformers
Run end-to-end Qwen2 integration:
export PYTHONPATH=.
python benchmarks/qwen2_e2e.py
Run kernel correctness suite:
python benchmarks/test_wrapper_final.py
- quantization-runtime: The mathematical algorithms (AWQ/GPTQ) behind these weights
- mini-llm-inference-engine: Full C++/CUDA inference runtime
- distributed-inference-engine: Multi-GPU parallelism
This project demonstrates:
- ✓ Deep understanding of PTX MMA and Tensor Core programming
- ✓ Ability to debug complex GPU issues (undefined behavior, layout constraints)
- ✓ End-to-end system integration (kernel to wrapper to model to inference)
- ✓ Documentation of architectural constraints
It does NOT demonstrate:
- ✗ Beating cuBLAS in raw kernel latency on Turing
- ✗ Production-grade weight-only quantization performance
For production INT4 inference on Turing, consider:
- Ampere/Hopper GPUs with native INT4 MMA
- Established libraries like CUTLASS or AutoAWQ
- Weight-only quantization primarily for VRAM savings, not latency
MIT — João Felipe De Souza, 2026