MP config: load protected channel indices - #26
Open
Allenjin123 wants to merge 8 commits into
Open
Conversation
New scmp_kernels/trace.py + one guarded hook at the sc_matmul entry. Records, for every SC matmul (MP and uniform), the effective precision (stoc_len = true post-halving cycle count; rng_levels = resolved enable-grid size) plus shape (rows/d_in/d_out/batch) and caller-tagged identity (operator, decoder block, sub-unit = MoE expert or attention head) — the input an external energy/latency simulator needs. Two output modes (SC_MP_TRACE=<path>, SC_MP_TRACE_MODE=summary|trace): - summary JSON: per-(block, op, unit, stoc_len, rng_levels, smoothed) groups with exact calls/rows/macs/row_cycles; d_in/d_out are representative payload with a dims_vary flag (attention KV growth under decode would otherwise mint one group per step). - trace JSONL: one ordered record per call (seq), spilled to disk every 100k records so long decodes cannot OOM the host. Properties: zero overhead when off (one module-bool read); host-side shape metadata only — no tensor reads, no device sync; thread-local context + locked accumulation (concurrent forwards); atexit safety-net flush (header marked "atexit"); bad SC_MP_TRACE_MODE falls back to summary with a warning instead of breaking package import; enable() refuses mode switches over unflushed records. Header documents coverage (SC matmuls only — FP16 lm_head/embeddings/router are absent). Apps adopt by calling trace.set_context(op, block, unit) before their matmuls; scmp_llm wires SCLinear + all SC attention paths (incl. STE and knock-down-probe helpers) in its own tree. Validated: CPU unit suite (schema, spill, threading, guards); Qwen3-4B MP PPL — per-op rows conserve exactly, trace avg stoc_len matches the runtime MP tracker, macs == rows*d_in*d_out, PPL bit-identical with tracing on; Qwen3-30B-A3B MoE — per-expert rows sum to tokens*top_k, routing load-imbalance captured. Hardened via a 3-lens adversarial review (stale-context, calibration-pollution, unbounded-memory, and schema-ambiguity findings all fixed). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ype) - Record halve as APPLIED, not requested: halving only takes effect in bipolar mode, so unipolar+halve runs no longer claim halved cycles (keeps the field consistent with the effective stoc_len/rng_levels). - Spill-mode flush() now ALWAYS writes the trailer carrying the final n_records (the spilled header has n_records=null); previously a bare flush() without header_extra left the count unrecorded. - header scramble_masks is int-typed (env echo, default 64); a non-numeric env value is echoed raw instead of crashing flush(). Documented as the knob echo, not the per-call resolved min(masks, 2^sc_prec). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
Used by parent PR CrucibleComputingGroup/scmp_llm#9. This submodule change provides the protected-channel config parsing/runtime lookup needed by the |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| global _FH | ||
| out = _PATH | ||
| os.makedirs(os.path.dirname(os.path.abspath(out)) or ".", exist_ok=True) | ||
| _FH = open(out, "w") |
| if _ENABLED and _SEQ: | ||
| try: | ||
| flush(header_extra={"atexit": True}) | ||
| except Exception: |
Comment on lines
+359
to
+361
| print(f"[trace] WARN: SC_MP_TRACE_MODE={_env_mode!r} is not " | ||
| f"'summary'|'trace'; falling back to 'summary'.", | ||
| file=sys.stderr) |
Calibration picks, per operator, the row metric with the best |Spearman rho|
vs true sigma-benefit and exports it as "dispatch_metrics" {op: {metric,
sign}}; negative rho deploys inverted (sign -1 maps to 1 - normalized(raw)
under the min-max normalization). Absent payload => ("amax", +1), byte-
identical to the original dispatch. All metrics are O(D) row reductions —
no runtime cost change.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Autoregressive decode grows the attention inner dim every step, so the per-(D, stoc_len) cum_indicator/k_table entries accumulated without bound (MP's 5 levels x ~50 decode steps x ~270MB at ctx-4096 = tens of GB). Cap both caches at SC_ENABLE_TABLE_CACHE_MAX (default 32) entries with LRU eviction. Fixed-ctx workloads (PPL) never evict — bit-identical results; decode evicts stale-D entries at rebuild-only cost (tables are pure functions of the key). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds support for loading and exposing “protected input channels” in AdaptiveMPConfig so the SC runtime can treat those channels separately, and introduces additional runtime instrumentation/guardrails around SC matmul execution and caching.
Changes:
- Parse and store protected-channel index sets (including block/unit-scoped keys) in
AdaptiveMPConfig, and exposeget_protected_channels(...). - Add a new low-overhead SC matmul tracing facility (
scmp_kernels.trace) and hook it intosc_matmul. - Bound enable-table caching in Triton kernels with an LRU to prevent unbounded growth during autoregressive decode.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| scmp_kernels/trace.py | New trace/summary output module for recording SC matmul shapes + effective precision metadata. |
| scmp_kernels/sc/matmul.py | Hooks trace recording into the unified SC matmul implementation. |
| scmp_kernels/sc/kernels.py | Replaces unbounded enable-table dict caches with a bounded LRU cache controlled by env var. |
| scmp_kernels/mp/config.py | Adds protected-channel parsing/storage and new per-row dispatch metric selection support. |
| scmp_kernels/mp/init.py | Re-exports new MP config utilities/constants. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+600
to
+602
| _ENABLE_TABLE_CACHE_MAX = int(os.environ.get("SC_ENABLE_TABLE_CACHE_MAX", "32")) | ||
| _enable_table_cache: "_OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = _OrderedDict() | ||
| _k_table_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict() |
Comment on lines
+229
to
+251
| if _trace._ENABLED: | ||
| if a.dim() == 3: | ||
| _batch, _rows, _d_in = a.shape | ||
| else: | ||
| _rows, _d_in = a.shape | ||
| _batch = 1 | ||
| _trace.record_matmul( | ||
| rows=int(_rows), d_in=int(_d_in), d_out=int(b.shape[-2]), | ||
| batch=int(_batch), | ||
| stoc_len=int(stoc_len) if stoc_len is not None else 2 ** sc_prec, | ||
| sc_prec=sc_prec, mode=mode, granularity=granularity, | ||
| # halving only takes effect in bipolar mode (guard above) — record | ||
| # whether it was APPLIED, not merely requested, to keep the field | ||
| # consistent with the effective stoc_len/rng_levels. | ||
| halve=bool(halve_bipolar_stoc_len and mode == "bipolar"), | ||
| # RESOLVED enable-grid size (mirrors _resolve_rng_levels), so the | ||
| # field has one meaning across halved and non-halved runs. | ||
| rng_levels=(int(rng_levels) if rng_levels is not None | ||
| else 2 ** sc_prec), | ||
| chunk_d=chunk_d, | ||
| smoothed=smooth_scales is not None, | ||
| ) | ||
|
|
Comment on lines
+265
to
+269
| ``path`` overrides the enable-time path (e.g. one file per sweep | ||
| config) — EXCEPT in trace mode once spilling has begun: the spill file | ||
| (enable-time path) is already on disk, so the remainder is appended | ||
| there and the override is ignored with a warning. ``header_extra`` | ||
| merges caller metadata (model id, eval tag) into the header.""" |
Comment on lines
+142
to
+146
| _ROW_METRIC_EPS = 1e-12 | ||
| # Candidate per-row dispatch metrics (act_global_v2 ρ-selection). All are O(D) | ||
| # reductions over the last dim — same runtime cost class as the original amax. | ||
| ROW_METRIC_NAMES = ("amax", "l2", "crest") | ||
|
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Comments suppressed due to low confidence (2)
scmp_kernels/sc/kernels.py:602
SC_ENABLE_TABLE_CACHE_MAXis parsed withint(...)at import time. A non-integer value will crash module import, and negative values can make_lru_putloop untilpopitemfails. Consider validating and clamping to a safe range (e.g., >= 0) with a default fallback.
_ENABLE_TABLE_CACHE_MAX = int(os.environ.get("SC_ENABLE_TABLE_CACHE_MAX", "32"))
_enable_table_cache: "_OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = _OrderedDict()
scmp_kernels/mp/config.py:325
- The PR title/description focuses on protected-channel indices, but this change set also introduces per-group ladders (
bucket_stoc_len_levels), dispatch-metric selection (dispatch_metrics/compute_row_metric), and an escape-gate path (escape_gate_k), plus separate tracing/caching changes in other modules. Consider updating the PR description/title (or splitting) so reviewers and release notes match the actual scope.
# Per-(op, timestep-bucket, layer-bucket) LADDER, parallel to
# bucket_thresholds. Empty => every bucket uses stoc_len_levels, which is
# the historical behaviour. See get_levels().
bucket_stoc_len_levels: dict[tuple[str, int, int], list[int]] = field(
default_factory=dict)
operator_default_stoc_len_levels: dict[str, list[int]] = field(
default_factory=dict)
protected_channel_stoc_len: Optional[int] = None
protected_channel_indices: dict[tuple[str, int, Optional[int]], list[int]] = field(default_factory=dict)
Comment on lines
+804
to
+813
| levels = config.stoc_len_levels | ||
| esc_sl = int(config.escape_stoc_len) | ||
| esc_idx = levels.index(esc_sl) if esc_sl in levels else len(levels) | ||
| assignment.row_levels[esc_mask] = esc_idx | ||
| for level_idx, sl in enumerate(levels): | ||
| assignment.level_row_indices[sl] = torch.where( | ||
| assignment.row_levels == level_idx)[0] | ||
| if esc_idx == len(levels): | ||
| assignment.level_row_indices[esc_sl] = torch.where(esc_mask)[0] | ||
|
|
| #SBATCH --cpus-per-task=12 | ||
| #SBATCH --mem=180G | ||
| #SBATCH --time=00:30:00 | ||
| #SBATCH --output=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/logs/verify_cum_cache_%j.out |
| export TMPDIR=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/tmp | ||
| export SC_OWEN_MODE=bitrev | ||
| export SC_SCRAMBLE_MASKS=64 | ||
| cd /home/allenjin/Projects/scmp_llm/kernels |
| #SBATCH --cpus-per-task=12 | ||
| #SBATCH --mem=180G | ||
| #SBATCH --time=00:30:00 | ||
| #SBATCH --output=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/logs/sweep_rows_%j.out |
| export TMPDIR=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/tmp | ||
| export SC_OWEN_MODE=bitrev | ||
| export SC_SCRAMBLE_MASKS=64 | ||
| cd /home/allenjin/Projects/scmp_llm/kernels |
The summary accumulator deliberately kept d_in/d_out OUT of the group key so autoregressive decode (KV length grows every step) could not mint one group per step. That also merged *static* shapes whenever they shared a stoc_len — notably an operator's protected-channel slice landing on the same rung as its main slice — yielding groups whose reported dims describe a matmul that was never run, and breaking rows*d_in*d_out == macs. 9/24 ppl/mp_best traces hit it (10-34% of their MACs). Shapes are now part of the key, and unbounded growth is handled by a per-base-key cap instead: the first SC_MP_TRACE_MAX_SHAPES (default 32) distinct (d_in, d_out) get exact groups; anything past it folds into one dims_vary catch-all with representative dims. macs/rows/row_cycles stay exact everywhere, including in the catch-all. reset() clears the shape registry. Tests cover the exact mp_best t96 failure (9144-wide main + 584-wide protected slice at stoc_len 128), the macs identity, cap-bounded KV growth with exact totals, and reset isolation. They load trace.py by path (pure stdlib) so they run on a login node without torch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (6)
scmp_kernels/mp/config.py:808
- _apply_escape_gate() rebuilds indices using config.stoc_len_levels (the global ladder) instead of the per-bucket ladder used to create
assignment. With per-group ladders enabled, this can assign out-of-rangerow_levelsindices and populatelevel_row_indicesunder the wrong stoc_len keys, breaking dispatch.
levels = config.stoc_len_levels
esc_sl = int(config.escape_stoc_len)
esc_idx = levels.index(esc_sl) if esc_sl in levels else len(levels)
assignment.row_levels[esc_mask] = esc_idx
for level_idx, sl in enumerate(levels):
scmp_kernels/sc/matmul.py:246
- The trace metadata sets
halve=...based on the request flag, even when the halving override did not actually apply (e.g., caller passes explicit stoc_len/rng_levels). This can make traces internally inconsistent (halve=true but effective stoc_len/rng_levels are not halved).
# consistent with the effective stoc_len/rng_levels.
halve=bool(halve_bipolar_stoc_len and mode == "bipolar"),
# RESOLVED enable-grid size (mirrors _resolve_rng_levels), so the
# field has one meaning across halved and non-halved runs.
rng_levels=(int(rng_levels) if rng_levels is not None
scmp_kernels/sc/kernels.py:603
- Parsing SC_ENABLE_TABLE_CACHE_MAX with int(...) at import time will raise ValueError on any non-integer env value, preventing the module from importing. For an env knob, it’s safer to fall back to the default (and clamp to >=1) rather than crash on import.
# results are unaffected; on decode, stale-D entries are evicted (rebuild cost
# only). Tables are pure functions of the key, so eviction is always safe.
_ENABLE_TABLE_CACHE_MAX = int(os.environ.get("SC_ENABLE_TABLE_CACHE_MAX", "32"))
_enable_table_cache: "_OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = _OrderedDict()
_k_table_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict()
tests/verify_cum_cache.sbatch:7
- This sbatch script hard-codes user-/cluster-specific settings (account/reservation/partition) and absolute paths (output, HF_HOME, TMPDIR, cd). As committed, it will fail for most users and leaks local filesystem layout into the repo; it should be a portable template that uses relative paths and lets site-specific setup be provided at submission time.
#SBATCH --account=nbleier_owned1
#SBATCH --reservation=rtx6000_arph_nodes
#SBATCH --partition=gpu-rtx6000
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=12
tests/sweep_rows.sbatch:7
- This sbatch script hard-codes user-/cluster-specific settings (account/reservation/partition) and absolute paths (output, HF_HOME, TMPDIR, cd). As committed, it will fail for most users and leaks local filesystem layout into the repo; it should be a portable template that uses relative paths and lets site-specific setup be provided at submission time.
#SBATCH --account=nbleier_owned1
#SBATCH --reservation=rtx6000_arph_nodes
#SBATCH --partition=gpu-rtx6000
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=12
scmp_kernels/trace.py:5
- PR title/description focus on loading protected-channel indices, but this PR also introduces new tracing (scmp_kernels/trace.py + sc_matmul instrumentation), LRU caching changes, per-group ladders, and the escape gate. This scope mismatch makes the change harder to review/revert; please update the PR description to cover these additions or split into separate PRs.
"""SC computation trace — precision/shape log for external HW simulators.
Records, for every ``sc_matmul`` call, the *effective* precision
(``stoc_len`` = SC cycle count, post-halving; ``rng_levels`` = resolved
enable-grid size) together with the matmul shape and the caller-supplied
…okup
Two changes: one a latent-bug fix, the other a new capability.
FIX: _apply_escape_gate rebuilt the row assignment against the GLOBAL ladder
(config.stoc_len_levels) while adaptive_classify_rows had classified against the
per-bucket ladder from get_levels. The rebuild re-keys level_row_indices by
stoc_len VALUE, so a bucket carrying its own ladder had its shared rung keys
overwritten and its non-shared rungs left stale -- 96 assignments for 64 rows,
i.e. those rows dispatched TWICE. Dormant in all 20 deployed mp_best cells (every
one runs escape_gate_k=2.0 with zero per-bucket ladders, verified), so no
archived result is affected -- but it fires the moment per-band ladders exist.
It is invisible downstream too: SCLinear iterates whatever keys it is handed and
the MP tracker prices the assignment rather than measuring the kernel, so a
corrupted cell still lands in band. Third instance of this family after
classify_level_values; any code indexing by rung must resolve via get_levels.
NEW: k_bands section. The residual (non-protected) contraction axis is split into
bands of whole quantization chunks, and band b runs row-rung k at its own length
L[b][k]. Row dispatch is untouched, so the per-row parent sits inside the space
at L[b][k] = L_parent[k] and the refinement cannot lose. Iso-compute is a
per-rung IDENTITY rather than a tolerance:
sum_b (w_b / R) * L[b][k] == L_parent[k]
MACs are linear in the contraction dim, so a column fraction IS the MAC fraction
and the identity holds for any input, any rung occupancy, and any escape-gate
firing pattern. Enforced at table load: overspend is a hard error; underspend is
allowed (a cheaper cell that still wins is a stronger result) but bounded, since
a solver leaving many cycles unspent is a bug, not conservatism.
Band counts are PER OPERATOR -- n_bands is a maximum, clamped by how many chunks
each operator has, because every band needs >= 2 chunks to stay on the chunked
kernel path. Requiring exactly n_bands everywhere silently dropped every narrow
projection (~19 chunks) whenever the count was raised for down_proj (71).
Bands must be whole 128-chunks: the kernel builds its RNG tables over chunk_d
dims and reuses them for every chunk, so relocating a whole chunk is numerically
free while relocating arbitrary channels is not (measured: 1e-7 vs 1.3e-1
relative error).
45 tests cover the schema validation, the rung-index resolver distinction
(get_levels vs classify_level_values, which appends an escape entry), and the
escape-gate case that fails without the fix.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated no new comments.
Suppressed comments (6)
scmp_kernels/sc/matmul.py:243
- The trace metadata currently marks
halveas true wheneverhalve_bipolar_stoc_lenis set, even if the effectivestoc_lenis not actually halved (e.g. when a caller explicitly passes a non-halved stoc_len). Since trace consumers will interprethalveas an effective property, compute it from the resolved stoc_len instead of the request flag.
halve=bool(halve_bipolar_stoc_len and mode == "bipolar"),
scmp_kernels/sc/kernels.py:605
- _ENABLE_TABLE_CACHE_MAX is parsed with int(...) at import time, which will raise ValueError (and fail module import) if SC_ENABLE_TABLE_CACHE_MAX is set to a non-integer. Also, negative values can trigger popitem() on an empty OrderedDict in _lru_put. Parse defensively and clamp to a non-negative maximum to make the cache robust to env misconfiguration.
_ENABLE_TABLE_CACHE_MAX = int(os.environ.get("SC_ENABLE_TABLE_CACHE_MAX", "32"))
_enable_table_cache: "_OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = _OrderedDict()
_k_table_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict()
# Chunked-MLP cum_indicator, same key discipline as _enable_table_cache.
_cum_indicator_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict()
tests/verify_cum_cache.sbatch:14
- This Slurm script hard-codes cluster-specific paths/usernames (e.g. /scratch/.../allenjin,
conda activate annstention) and an absolute output path. Committing these makes the repo harder to use on other clusters and can leak internal filesystem layout. Consider moving this to a local-only artifact or converting it into a template that relies on environment variables / relative paths.
#SBATCH --output=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/logs/verify_cum_cache_%j.out
set +u
source ~/.bashrc
conda activate annstention
set -u
tests/sweep_rows.sbatch:14
- This Slurm script hard-codes cluster-specific paths/usernames and a specific conda env (
annstention), plus an absolute output path under /scratch. If this is intended for general contributors, it should be parameterized (env vars) or turned into a template; otherwise it may be better kept out of the repo.
#SBATCH --output=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/logs/sweep_rows_%j.out
set +u
source ~/.bashrc
conda activate annstention
set -u
scmp_kernels/mp/config.py:549
- The PR title/description focus on protected-channel indices, but this change set also adds substantial new functionality (e.g. K-bands parsing/validation, dispatch_metrics + compute_row_metric, escape gate behavior, and related tests). Please update the PR description/title to reflect the full scope or consider splitting into smaller PRs so reviewers can track intent and risk more easily.
K_BAND_ISO_COST_TOL = 0.25
scmp_kernels/mp/config.py:1082
esc_mask.any().item()introduces a device synchronization on the hot classification path whenever the escape gate is enabled for a bucket. This can add avoidable latency during runtime dispatch. You can avoid the sync by usingtorch.whereand checkingnumel()(which is available without synchronizing).
if not bool(esc_mask.any().item()):
return
# MUST be the SAME ladder the classification used (adaptive_classify_rows
# resolves it via get_levels). Reading config.stoc_len_levels here instead
# was latent-correct only while every bucket shared the global ladder: the
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.
Summary
Testing