Skip to content

GLM: an optional residual-direction ablation (carved out of #642) - #901

Open
ClintMoody wants to merge 3 commits into
mw00:mainfrom
ClintMoody:glm-ablate
Open

ClintMoody wants to merge 3 commits into
mw00:mainfrom
ClintMoody:glm-ablate

Conversation

@ClintMoody

Copy link
Copy Markdown

The paged-attention work (#642) grew a setting of its own: a harness that removes chosen directions from the GLM path's residual streams. @catonooka asked for one topic per PR, so it moved here. #642 carries no glm_ablate reference now.

What it does

x <- x - s (x . v) v, applied to the way into every write into the residual streams.

  • src/core/glm_ablate.cuh holds the projection and its three settings. Off unless MAYA_GLM_ABLATE names a file of raw float32 directions.
  • Outputs mode (default): K x n_embd values, K orthonormal directions. Applied to the token embedding and every block output, the mixer's and the FFN's, before the mHC write.
  • Per-layer mode (MAYA_GLM_ABLATE_LAYERS=n_layers): n_layers x K x n_embd values. Layer l's own vectors are applied to the two operands of that layer's FFN write, so the fused read needs no change. A layer whose vectors are all zero is not steered.
  • MAYA_GLM_ABLATE_SCALE (default 1) sets s: 1 removes the component along each direction, above 1 over-projects.

The same operation as exllamav3's EXL3_CVEC hook after each block, and the knob DeepSeek's --ablate gives: find a direction that carries a behaviour, remove it, watch the model.

What it costs

Nothing when unset: each call site checks one bool and returns. The NextN draft block is not touched (a draft only has to be a good guess).

Wiring

src/core/glm_fast_path.cu: the include; the direction's upload on the device before the expert pool is sized; the mixer's and the FFN's writes; the last layer's write; the token embedding; and both operands of every layer's FFN write in per-layer mode. Nine call sites in all.

Tests

tests/glm_ablate_test.cu runs the kernel on a device and compares every value with a double-precision reference:

  • outputs mode: x - (x . v) v for every direction, on every row, and no component along either direction survives;
  • scale 0.5: the component along each direction is halved;
  • per-layer: layer 0 steered, layer 1 (all zero) bit-for-bit untouched.

Run here on a 3090 Ti against this tree: 4, 4 and 5 checks, 0 failures. Three ctest entries (glm_ablate_test, glm_ablate_scale_test, glm_ablate_layers_test). A device that cannot be opened is a SKIP, as the other GPU tests do, so a CPU-only build still runs the suite.

Notes

  • MAYA_* is the alias: include/strata/platform/env_alias.hpp maps every MAYA_X to STRATA_X, and MAYA_X wins when both are set. Each program calls env_aliases() first thing in main().
  • A missing file, a file of the wrong width, or one that is not float32 prints one glm ablate: line and leaves the stock model running. No GPU room turns the setting off.
  • docs/SETTINGS.md has the section, in the shape of the other rows there.

This branch changes no default path, so the CONTRIBUTING.md before/after maya.py --bench pair does not apply to it: with MAYA_GLM_ABLATE unset the work is the same, plus one bool check per call site.

… work

STRATA_GLM_ABLATE=<file> projects directions out of the GLM fast path's
residual streams (research; off unless the variable is set). Outputs mode
removes K orthonormal vectors from the token embedding and every block
output; per-layer mode (STRATA_GLM_ABLATE_LAYERS) gives each layer its own
vectors, the same operation as exllamav3's EXL3_CVEC hook.

This is the harness alone, kept off the paged-attention branch per review.
It is not wired into a translation unit yet; the follow-up wires it.
…cument them

The harness in glm_ablate.cuh was carved out of the paged-attention work,
where it had no business: it is a setting of its own. This adds the wiring
back, on top of the carved header:

- the include, and the direction's upload on the device before the expert
  pool is sized (so the pool sees the bytes it takes);
- the mixer's and the FFN's writes, and the token embedding - x - s (x . v) v
  before each mHC write, in outputs mode;
- both operands of every layer's FFN write, in per-layer mode;
- the docs/SETTINGS.md section and the three MAYA_GLM_ABLATE* rows.

Off unless MAYA_GLM_ABLATE is set: nothing runs, and the answers do not change.
glm_ablate_test runs the rows kernel on a device and compares it with a
double-precision reference: outputs mode, the scale at 0.5, and per-layer
mode (layer 0 steered, layer 1 all zero and untouched). A device that
cannot be opened is a SKIP, as the other GPU tests do.

Adds the target and three ctest entries. The engine does not link it.
@mw00

mw00 commented Nov 9, 2026

Copy link
Copy Markdown
Owner

Thanks for splitting this out. It is off unless the env var names a file, which is good. Which GPUs/model/quant did you test on, and can you confirm with it unset the outputs and speed are identical to main (greedy output match, decode and prefill)? We will look at whether this belongs in Maya.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants