int8 tilewise - #1699
Open
dxqb wants to merge 14 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>
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>
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
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>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds INT_W8A8_TILEWISE, offered as "int W8A8 tilewise" in the data type dropdown. It keeps one int8 scale per 64x64 weight tile instead of one for the whole tensor, so a few extreme rows or columns no longer set the step size for every other weight in the layer. The tiles are square because the forward reduces over K and the backward over N, and a square tile is the only shape one stored copy of the weight can serve in both directions. The triton mm gains a TILED path for it: the autotuner is pruned to the configs whose BLOCK_SIZE_K is the tile size, so one k-block is exactly one tile, and the block's product is descaled in fp32 with the scale row for that k-group before it is accumulated. A tile-scaled weight therefore needs the kernel - the torch fallback in mm_8bit asserts scale_kn is None, and LinearW8A8 raises when the kernel is unavailable or a layer's dimensions are not multiples of the tile size. The forward and backward helpers in LinearW8A8 lose their implementation suffixes and the unused torch variant, since there is one implementation each again. 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
1 participant
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.
contains other PRs, actual diff: 30d5e9a
Qwen Image:
Flux2:
Summary
Test plan
pre-commit run --all-filespassesAI assistance