Conversation
The fused kernel additionally caches x in shared memory, which exceeds the default 48 KiB limit for HeadDim=128 (the Qwen3 case); set cudaFuncAttributeMaxDynamicSharedMemorySize like the non-fused launcher already does. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Qwen3 models specify head_dim explicitly (128), which does not
equal HiddenSize / NumQueryHeads for the 0.6B/4B/32B variants.
Add an optional HeadDim to TransformerConfig (0 = derived as
before, so existing models are unchanged) and route all users
through head_size() / the new attn_channels():
- o_proj weight becomes {C, attn_channels}
- the o_proj matmuls take Hq*Hs input channels
- att activations and d_att_y are sized (B, T, attn_channels);
d_att_y can no longer alias d_lnf when the sizes differ
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Qwen3-0.6B ships an explicit lm_head.weight despite tie_word_embeddings=true; skip it instead of rejecting the file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Also makes the distinction to attn_channels() (Nq*Hd, the o_proj input) obvious at the allocation sites. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Qwen3 (QK-norm) support across the training/runtime stack by extending model configuration, weights, run-state shapes, and introducing fused QK-norm+RoPE CUDA kernels with Python bindings and expanded kernel test coverage.
Changes:
- Add Qwen3 architecture support (config parsing/saving, head-dim overrides, QK-norm enablement).
- Extend model weights/run-state to support QK-norm parameters and nonstandard attention channel sizes.
- Implement and bind fused
qk_norm_and_rope_{forward,backward}kernels; split kernel tests into a propertest/suite.
Reviewed changes
Copilot reviewed 28 out of 28 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| test/test_kernels_swiglu.py | Adds dedicated SwiGLU forward/backward + FP8 quant tests. |
| test/test_kernels_sr.py | Adds stochastic-rounding vector add/reduce tests. |
| test/test_kernels_rope.py | Adds RoPE forward/backward correctness tests vs Python reference. |
| test/test_kernels_rmsnorm.py | Adds RMSNorm forward/backward + fused residual RMSNorm tests. |
| test/test_kernels_quant.py | Adds abs-max and quantization kernel tests (BF16/FP8, transpose). |
| test/test_kernels_qk_norm.py | Adds QK-norm and fused QK-norm+RoPE forward/backward tests. |
| test/test_kernels_other.py | Adds tests for misc kernels (fillers, transpose, norms, classifier, AdamW). |
| test/test_kernels_matmul.py | Adds cublasLt-backed matmul tests (modes, bias, accumulation). |
| test/test_kernels_encoder.py | Adds encoder forward/backward tests incl. determinism/accumulation. |
| test/conftest.py | Centralizes dtype list and numeric tolerances for tests. |
| src/training/transformer_config.h | Adds QWEN3 arch, optional HeadDim, and QK-norm feature flag. |
| src/training/transformer_config.cpp | Loads/saves QWEN3, head_dim, and enables QK-norm for QWEN3. |
| src/models/llama_weights.h | Adds optional Q/K norm weights to block weights container. |
| src/models/llama_weights.cpp | Allocates/shards QK-norm weights; adjusts attention projection shapes and scale storage. |
| src/models/llama_run_state.h | Extends activations/gradients to handle attention channel size + QK-norm buffers. |
| src/models/llama_run_state.cpp | Allocates QK-norm scratch/QKV rope temps; adjusts tensor shapes for nonstandard attn channels. |
| src/models/llama_optimizer.cpp | Includes Q/K norm weights in optimizer tensor iteration. |
| src/models/llama_model.h | Adds recompute_qk_rope() helper to control rematerialization behavior. |
| src/models/llama_model.cpp | Integrates fused QK-norm+RoPE forward/backward paths; updates matmul dims for attention channels. |
| src/models/llama_gradients.cpp | Zeros Q/K norm gradients on first micro-step / shard acquisition. |
| src/kernels/qk_norm.cu | Implements fused QK-norm+RoPE forward/backward kernels and updated reductions. |
| src/kernels/kernels.h | Declares fused QK-norm+RoPE APIs and scratch sizing. |
| src/kernels/kernels.cpp | Adds Tensor-dispatch wrappers for fused QK-norm+RoPE forward/backward. |
| src/binding/python/kernels.py | Exposes fused QK-norm+RoPE ops + scratch sizing to Python. |
| src/binding/py_train.cpp | Exposes Q/K norm gradients to Python trainer when present. |
| src/binding/kernel_binding.cpp | Adds nanobind bindings for fused QK-norm+RoPE forward/backward + scratch sizing. |
| scripts/test_kernels.py | Removes monolithic kernel test script (migrated into test/). |
| scripts/bench_qk.py | Adds a simple benchmark harness for QK-norm kernels. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
AKA QK-norm