nunchaku-style kernel fusion - 20-30% faster training - #1680
Open
dxqb wants to merge 13 commits into
Open
Conversation
Suppress a handful of specific, noisy-but-harmless messages emitted while launching the UI and starting training: - diffusers/transformers logger.warning() lines (Modular Diffusers experimental notice, unexpected-config-attributes, unrecognized loss_type) via filters on the exact emitting loggers - huggingface_hub local_dir_use_symlinks deprecation and the torch.compile inductor performance notes via warnings/logger filters - Qt gnome portal dbus errors via QT_LOGGING_RULES - tensorboard subprocess banner/notices by discarding its stdout/stderr Each filter targets one specific message, so other warnings from the same libraries still come through. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Reworks the Triton 8-bit matmuls around a shared compute core and adds a LoRA epilogue to it, so a LoRA layer over an 8-bit quantized Linear no longer materializes its low-rank product separately. The three mm entry points now share _mm_accumulate (grouped launch order for L2 reuse, compile-time layout specialization, an EVEN_K-specialized main loop) and differ only in their epilogue: the raw product, a per-row dequant scale, or that scale plus a rank-tiled low-rank update. The scaled_lora_mm_8bit and rowcol_scaled_lora_mm_8bit entry points expose the last one. transpose_8bit rewrites the backward pass's B operand to k-major, since the 8-bit tensor-core op wants k-major and Ada has no 8-bit ldmatrix.trans; the wrappers apply it above a token threshold, where the copy pays for itself. On the module side, LoRAFusableLinearMixin declares forward_with_lora, and LinearW8A8, LinearGGUFA8 and LinearSVD implement it with autograd Functions that own the down-projection and the dropout. That lets the LoRA dgrad fold into the backward epilogue as well, so neither direction builds an (M, out_features) intermediate. LoRAModule dispatches on the mixin rather than on BaseLinearSVD, and fused_leaf_forward extends the same path to fused-qkv adapters by narrowing lora_up to the leaf's rows. DoRA opts out, as it recomposes the weight instead of adding a delta. mm_8bit falls back to torch._int_mm / torch._scaled_mm when Triton is not importable, and quantize_axiswise dispatches the two 8-bit dtypes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
M is batch*sequence, so unlike N and K it is data-dependent and unbounded: it moves with resolution, frame count, batch size and, on models that prune prompt padding, the longest caption in the batch. Keying the autotuner on M // 64 makes the number of tuning keys grow linearly with M, so a run keeps paying for fresh autotune passes as sequence lengths vary. QUANTIZED_M now uses M.bit_length(), bucketing per doubling instead, which is logarithmic in M. Proportional resolution is the right shape here because the winning config is decided by block count against SM count, which is linear in M -- equal ratios of M matter equally at every scale, while a fixed stride is too fine at large M and too coarse at small. The LoRA epilogue additionally drops stride_upr from its key. That stride does vary -- it is 1 in the forward, where up is a transposed view of lora_up, and N in the backward, where up is lora_down stored row-major -- but keying on it doubles the key count on every model. The two layouts do load differently, and the slab is only BLOCK_R x BLOCK_N and is reused by every M block in the group: too little traffic next to A and B to move the tile choice. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
# Conflicts: # modules/util/ui/pyside6_util.py
A cold torch.compile cache announces every frame it compiles, which scrolls the progress bar off the screen. There is no knowable total to build a real progress bar from, so the announcement goes into the postfix of the innermost running bar instead, and is cleared again by that bar's next redraw or by its close. tqdm keeps its bars in an unordered WeakSet, so which of the nested bars is the innermost one cannot be recovered from it. modules/util/tqdm_util.py subclasses tqdm to track that itself and adds show_status() next to tqdm.write(); every tqdm import in the repo now comes from there. Bars owned by mgds are outside this and still draw as before. Also gates the warning filters on OT_DEBUG_WARNINGS, so setting it brings every suppressed message back, and silences the diffusers attention-backend experimental notice. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
now contains #1628 to show the autotuning messages as part of the progress bar |
A cold autotune cache benchmarks each kernel once per shape key, and every one of those sweeps announced itself on a line of its own. Route them through tqdm.show_status(), so the message sits in the innermost progress bar while the sweep runs and is cleared once the bar moves on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
as a side effect, this PR halves the size of linear layer activations, because it skips the int32-materialization and directly goes to bf16 batch size 16 peak vram goes down from 11 GB to 9 GB 
|
The legacy fp8 mma issues at half rate with an fp32 accumulator on sm_120, while the block-scaled mxf8f6f4 instruction Blackwell added runs at the full 8-bit tensor core rate. Reaching it through tl.dot_scaled with every scale set to the ue8m0 encoding of 1.0 computes exactly the same product, so the results are bit-identical to what the kernels produce today - it is purely an instruction swap. All five entry points share _mm_accumulate, so the swap lands once and all five benefit. Pre-Blackwell has no such instruction and triton emulates tl.dot_scaled with a bf16 mma, which is slower than the plain tl.dot path, so MXFP8_MMA is decided from the compute capability and sm_89 keeps its current path unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
torch.cuda also serves ROCm devices, where the device capability is the gfx arch number - RDNA4 reports 12 and would take the tl.dot_scaled path, which triton has to emulate there. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…-squashed # Conflicts: # modules/util/triton_mm_8bit.py
|
What about DoRA? |
Collaborator
Author
not possible for DoRA |
The five entry points (mm_8bit, scaled_mm_8bit, rowcol_scaled_mm_8bit and the two lora variants) were the same _mm_accumulate core with a different epilogue bolted on, each with its own Triton kernel, its own custom op and its own autotune cache. They become one mm_8bit(a, b, out_dtype, scale_m=None, scale_n=None, lora_xd=None, lora_up=None) behind a single ot_quant::mm_8bit op. Triton specializes a None argument as a constexpr, so every combination still compiles to the code a separate kernel per epilogue produced. The main loop no longer carries an EVEN_K variant. _prepare_mm zero-pads K up to _K_ALIGN (the largest BLOCK_SIZE_K in the autotune configs) instead, so every block of the loop is full and the loads need no mask; the padded products are zero. Only layers whose width is not a multiple of that alignment pay for the pad copies. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
dxqb
added a commit
to dxqb/OneTrainer
that referenced
this pull request
Aug 19, 2026
Snapshot of the kernel-lora-fusion branch, which this feature builds on and which is under review as its own PR (Nerogar#1680).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
2 participants
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.
Summary
Test plan
pre-commit run --all-filespassesAI assistance