fix(exl3): give the MoE down projection its own Hadamard scratch - #1775
Open
Andybui1012 wants to merge 1 commit into
Open
fix(exl3): give the MoE down projection its own Hadamard scratch#1775Andybui1012 wants to merge 1 commit into
Andybui1012 wants to merge 1 commit into
Conversation
exl3_mgemm reads A and writes the Hadamard-transformed input to A_had. The cooperative-kernel autotuner re-launches the kernel on the same arguments while it times candidates, so when A_had aliases A every launch after the first transforms the previous launch's output. The single-token and small-batch MoE paths passed exl3_small_interm_a for both, so the first call for a given shape returned garbage (normalized error ~1.0 vs a dense reference) and every later call, served from the autotune cache, was correct. In serving this is the first decode step for each expert shape, or for every shape whose autotune record is not in ~/.cache/exllamav3. Allocate a dedicated exl3_small_interm_a_had buffer next to interm_a and pass it as A_had at both call sites. ExLlamaV3 uses a separate buffer here as well (block_sparse_mlp.py, "A_had must not alias A"). The new test builds a small synthetic EXL3 MoE layer through Exl3MoEMethod, uses intermediate sizes no other test autotunes, and checks the first apply() against a per-expert fp32 dense reference for both paths. It fails on the previous code (1.22 and 0.97) and passes with the fix.
AlpinDale
requested changes
Sep 6, 2026
|
|
||
|
|
||
| @pytest.mark.parametrize(("rows", "intermediate"), [(1, 384), (3, 640)]) | ||
| def test_moe_small_batch_first_call_matches_dense_reference(rows: int, intermediate: int, monkeypatch, tmp_path): |
Member
There was a problem hiding this comment.
Suggested change
| def test_moe_small_batch_first_call_matches_dense_reference(rows: int, intermediate: int, monkeypatch, tmp_path): | |
| def test_moe_first_call_matches_dense_reference(rows: int, intermediate: int, monkeypatch, tmp_path): |
Still not happy with it, but this name is a lot better.
| environment override only helps when this test runs first, and is set so | ||
| that running the file on its own never sees a pre-seeded shape. | ||
| """ | ||
| monkeypatch.setenv("EXLLAMAV3_TUNE_CACHE", str(tmp_path / "coop_autotune_v1.bin")) |
Member
There was a problem hiding this comment.
This would make the test pass unconditionally with a warm cache, but I'd say it's mostly a nit.
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.
Problem
exl3_mgemmreadsAand writes the Hadamard-transformed input toA_had. The cooperative-kernel autotuner (CoopKernelAutotuner::launch) re-launches the kernel on the same arguments while it times candidates, so whenA_hadaliasesA, every launch after the first transforms the previous launch's output.Exl3MoEMethod._apply_single_tokenand_apply_small_batchpassedlayer.exl3_small_interm_aas bothAandA_hadfor the down projection. Result: the first call for a given expert shape returns garbage (normalized max error ~1.0 against a dense reference), and every later call, served from the autotune cache, is correct. In serving that is the first decode step for each MoE expert shape, or for every shape whose autotune record is not already in~/.cache/exllamav3/autotune.Fix
Allocate a dedicated
exl3_small_interm_a_hadbuffer next toexl3_small_interm_ainprocess_weights_after_loadingand pass it asA_hadat both call sites. ExLlamaV3 uses a separate buffer for this (block_sparse_mlp.py: "A_had must not alias A (the autotuner relaunches on the first call)"). The gate/up projections already useexl3_small_yh/exl3_small_yh_guas separate scratch.Test
tests/kernels/quantization/test_exl3_moe_small_batch.pybuilds a small synthetic EXL3 MoE layer throughExl3MoEMethod.process_weights_after_loading, uses intermediate sizes no other test autotunes, and checks the firstapply()against a per-expert fp32 dense reference (reconstructed weights with the 128-Hadamard and sign flips folded in), for both the single-token (rows=1) and small-batch (rows=3) paths. It also checks that the second call equals the first.On the previous code:
2 failed(normalized error 1.22 and 0.97). With the fix:2 passed. Verified in a CUDA 13 container on a GB10 (sm_121); the test needs a CUDA device and about 100 MiB of GPU memory.ruff checkandruff format --checkpass on both files.