diff --git a/e2e_workflow/e2e_workflow.js b/e2e_workflow/e2e_workflow.js index a30cf3c88..2b5aeea7f 100644 --- a/e2e_workflow/e2e_workflow.js +++ b/e2e_workflow/e2e_workflow.js @@ -666,6 +666,23 @@ const CAPTURE_CASE_BYTE_LIMIT = String(A.capture_case_byte_limit != null ? A.cap const CAPTURE_PERSIST_POLICY = String(A.capture_persist_policy != null ? A.capture_persist_policy : 'share_large'); const CAPTURE_WORKSPACE_BUDGET = String(A.capture_workspace_budget != null ? A.capture_workspace_budget : '32GiB'); const CAPTURE_STORAGE_ENV = `CAPTURE_BYTE_BUDGET=${CAPTURE_BYTE_BUDGET} CAPTURE_CASE_BYTE_LIMIT=${CAPTURE_CASE_BYTE_LIMIT} CAPTURE_PERSIST_POLICY=${CAPTURE_PERSIST_POLICY}`; +// ---- PER-HEAD EXTRACTION WALL-CLOCK BUDGET ---- +// A retry COUNT does not bound extraction cost. The real envelope is the product of three limits that +// were never multiplied out: BASELINE_EXTRACT_RETRIES (3, so 4 invocations) x safeAgent's internal +// tries (3) x AGENT_TIMEOUT_MS (2h when no global budget is set) = up to 24h on ONE head. And the +// global ELAPSED clock only exists when time_budget_s is passed, so a budget-less run has no guard at +// all. In the 20260907 session that envelope was not hypothetical: one gemm head spent 13h34m over 7 +// captures and one moe head 2h45m over 6, and NEITHER ever reached the optimization lane — ~16h that +// produced no kernel work. The only extraction that did reach it took 63min over 3 attempts. +// +// So bound the thing that actually costs: wall-clock per head, spanning every attempt. The default is +// ~1.9x the one observed success, which still cuts ~12h off that session. A head that exhausts it is +// NOT silently dropped -- it lands on the SAME failure path as any other failed extraction, so a +// dominant head is flagged and surfaced rather than quietly skipped. +const CAPTURE_BUDGET_MS = parseInt(A.capture_budget_s != null ? A.capture_budget_s : 7200, 10) * 1000; // 120min +// Below one server boot there is no attempt worth starting -- boots measured at 15-30min on this model. +// Starting one anyway buys a guaranteed-useless invocation and, worse, an overrun past the budget. +const CAPTURE_MIN_ATTEMPT_MS = parseInt(A.capture_min_attempt_s != null ? A.capture_min_attempt_s : 900, 10) * 1000; const TASK = A.task || ''; const APPLY_TO_ORIGINAL = String(A.apply_to_original != null ? A.apply_to_original : 'false'); const EVAL_DIR_OVERRIDE = A.eval_dir || ''; @@ -1213,9 +1230,38 @@ function agentTimeoutFor() { return Math.max(120000, Math.min(AGENT_TIMEOUT_MS, remainingMs() - FINAL_RESERVE_MS)); } +// A relative stopwatch for one bounded stretch of work. Same construction as the global ELAPSED clock +// above and for the same reason: rungs armed all at once at ABSOLUTE offsets, never as a self-rearming +// chain, so a late rung delays only itself instead of compounding into a drifting under-count. Date.now() +// is unavailable in Workflow scripts, so this is the only way to read elapsed time here. +// With no setTimeout the budget is unenforceable, and the stopwatch reports Infinity -- every caller +// then behaves exactly as it did before this existed, rather than aborting work it cannot time. +function stopwatch(budgetMs) { + const armed = typeof setTimeout === 'function' && budgetMs > 0; + let elapsed = 0; + const timers = []; + if (armed) { + const step = Math.max(30000, Math.ceil(budgetMs / 240)); // >=30s granularity, <=241 rungs + for (let at = step; at <= budgetMs + step; at += step) { + const mark = at; + const t = setTimeout(() => { if (mark > elapsed) elapsed = mark; }, mark); // max: stays monotonic + if (t && t.unref) t.unref(); + timers.push(t); + } + } + return { + remainingMs: () => (armed ? Math.max(0, budgetMs - elapsed) : Infinity), + spentMin: () => Math.round(elapsed / 60000), + stop: () => { for (const t of timers) clearTimeout(t); timers.length = 0; }, + }; +} + function agentBounded(rawPrompt, opts) { const prompt = withProcessSafety(rawPrompt); - const timeoutMs = agentTimeoutFor(); + // opts.timeoutCapMs lets a CALLER that owns its own budget cap this attempt below the global hung-guard. + // Absent (every existing call site) the cap is Infinity and the timeout is byte-identical to before. + const cap = (opts && Number.isFinite(opts.timeoutCapMs)) ? Math.max(60000, opts.timeoutCapMs) : Infinity; + const timeoutMs = Math.min(agentTimeoutFor(), cap); if (typeof setTimeout !== 'function' || !(timeoutMs > 0)) return agent(prompt, opts); let to; const guard = new Promise((resolve) => { @@ -1233,6 +1279,13 @@ function agentBounded(rawPrompt, opts) { async function safeAgent(prompt, opts, tries = 3) { let lastErr = 'unknown'; for (let i = 0; i < tries; i++) { + // opts.abortIf lets a caller with its own budget stop RETRYING inside this funnel. Without it the + // three internal tries are invisible to every caller, and each one is armed with the full hung-guard + // -- which is how a "3 retries" loop turns into a 24h envelope. Absent, nothing changes. + if (opts && typeof opts.abortIf === 'function' && opts.abortIf()) { + log(`agent[${(opts && opts.label) || '?'}] not retrying (${i}/${tries} used): caller's budget is spent.`); + return null; + } try { const r = await agentBounded(prompt, opts); if (r) return r; @@ -1437,7 +1490,9 @@ function kernelSelectionVerified(h, ext) { if (!verdict || verdict.contract !== 'kernel_selection') return { ok: false, why: 'kernel_selection.py verdict is missing' }; if (verdict.ok !== true) - return { ok: false, + // `codes`/`verdict` travel with the failure so the retry can tell "wrong seam" from "right seam, + // wrong NAME". Only the first is repaired by descending to another callable. + return { ok: false, codes: (verdict.failed || []).slice(), verdict, why: `kernel_selection.py failed: ${(verdict.failed || []).join(', ') || 'unknown'}` }; if (String(verdict.target_callable || '').trim() !== target) return { ok: false, @@ -1541,85 +1596,143 @@ async function extractWithBaseline(role, phase, intro, inputs, opts) { `\`${CAPTURE_STORAGE_ENV}\`. Use kernel_selection.py with --task-dir "$TASK" so the selected oracle is ` + `promoted and all capture.pid-* dirs are reclaimed. Unittests for large MoE oracles MUST use ` + `h.iter_eager_cases_from_oracle / h.check_correct_multi_lazy.`; - let ext = await safeAgent(roleAgent(role, phase, captureIntro, { - ...(inputs || {}), - CAPTURE_STORAGE_ENV, - CAPTURE_BYTE_BUDGET, - CAPTURE_CASE_BYTE_LIMIT, - CAPTURE_PERSIST_POLICY, - }), opts); - let tries = 0; - const attemptedTargets = []; - const complete = (e) => hasFrozenBaseline(e) && kernelSelectionVerified(head, e).ok; - while (smokeOk(ext) && !complete(ext) && tries < BASELINE_EXTRACT_RETRIES) { - tries++; - const selection = kernelSelectionVerified(head, ext); - const needBaseline = !hasFrozenBaseline(ext); - const priorTarget = String((ext && ext.target_callable) || '').trim(); - if (priorTarget && !attemptedTargets.includes(priorTarget)) attemptedTargets.push(priorTarget); - const selectionCorrective = selection.ok ? '' : - ` PRIOR ATTEMPT DID NOT SELECT THE PROFILED GPU KERNEL: ${selection.why}. ` + - 'Treat KERNEL.live_call_seam as prose only. Merge KERNEL.seam_candidates with any missing inner ' + - 'launcher found from source/runtime inspection; preserve existing candidate classifications. ' + - 'BEFORE re-running capture, reclaim prior process-local artifacts: ' + - '`python3 "$SKILL_DIR/scripts/capture_shapes.py" --cleanup-task-dir "$TASK" --no-promote` ' + - '(issue #429 — never accumulate capture.pid-* oracles across retries). ' + - 'Install safe markers for every relevant candidate, never native/JIT kernel_entry objects. Run ' + - 'kernel_selection.py over every process-local capture and all root-call traces with --task-dir "$TASK". ' + - 'Return its JSON verbatim as selection_validation. Select the deepest live inner_launcher/op_seam across all ' + - 'calls/ranks; a fused head must select the whole-op op_seam. Rejecting the previous outer wrapper ' + - 'is not success. Do not return any ATTEMPTED_TARGET_CALLABLES value again.'; - const baselineCorrective = needBaseline ? - ' PRIOR ATTEMPT DID NOT FREEZE A BASELINE. You MUST seed baseline_overlay/ from ' + - 'CURRENT_OVERLAY (the live serving stack = the speedup denominator), declare meta.candidate_bind ' + - '(the ONE overlay entry built from kernel_src/), prove both legs differ via h.assert_legs_differ, ' + - 'then return baseline_frozen:true. An extraction with no frozen baseline is INVALID and will be discarded.' : ''; - log(` ${(opts && opts.label) || role}: extraction contract incomplete ` + - `(${selection.ok ? 'kernel selected' : selection.why}; ` + - `${needBaseline ? 'baseline missing (baseline_overlay/ + meta.candidate_bind)' : 'baseline frozen'}). ` + - `RE-EXTRACTING (retry ${tries}/${BASELINE_EXTRACT_RETRIES}).`); - // Best-effort reclaim before the agent retries (issue #429 storage amplification). - if (ext && ext.task_dir) { - try { - const { execFileSync } = require('child_process'); - execFileSync('python3', [ - `${WORKFLOW_DIR}/scripts/capture_shapes.py`, - '--cleanup-task-dir', String(ext.task_dir), - '--no-promote', - ], { stdio: 'pipe', timeout: 120000 }); - log(` ${(opts && opts.label) || role}: reclaimed capture artifacts under ${ext.task_dir}`); - } catch (cleanupErr) { - log(` ${(opts && opts.label) || role}: capture reclaim skipped (${cleanupErr && cleanupErr.message ? cleanupErr.message : cleanupErr})`); + // One budget for this head's WHOLE extraction, armed before the first attempt and spanning every + // retry -- including safeAgent's internal ones, which no caller could see before. Capping each + // attempt at what is LEFT is the half that matters: a single invocation could otherwise run the + // full 2h hung-guard and blow a 2h budget by itself, no matter what the loop below checks. + const budget = stopwatch(CAPTURE_BUDGET_MS); + const spent = () => budget.remainingMs() < CAPTURE_MIN_ATTEMPT_MS; + const bounded = () => ({ ...(opts || {}), timeoutCapMs: budget.remainingMs(), abortIf: spent }); + try { + let ext = await safeAgent(roleAgent(role, phase, captureIntro, { + ...(inputs || {}), + CAPTURE_STORAGE_ENV, + CAPTURE_BYTE_BUDGET, + CAPTURE_CASE_BYTE_LIMIT, + CAPTURE_PERSIST_POLICY, + }), bounded()); + let tries = 0; + let budgetSpent = false; + const attemptedTargets = []; + const complete = (e) => hasFrozenBaseline(e) && kernelSelectionVerified(head, e).ok; + while (smokeOk(ext) && !complete(ext) && tries < BASELINE_EXTRACT_RETRIES) { + if (spent()) { + // Stop BEFORE paying for an attempt, not after. The result so far still flows into the checks + // below, so an extraction that is merely incomplete is reported as incomplete for its own reason + // -- the budget only decides that there will be no further attempt to fix it. + budgetSpent = true; + log(` ${(opts && opts.label) || role}: capture budget spent (${budget.spentMin()}min of ` + + `${Math.round(CAPTURE_BUDGET_MS / 60000)}min over ${tries + 1} attempt(s)) — no further re-extraction. ` + + `Raise with args.capture_budget_s if this head is worth more.`); + break; } + tries++; + const selection = kernelSelectionVerified(head, ext); + const needBaseline = !hasFrozenBaseline(ext); + const priorTarget = String((ext && ext.target_callable) || '').trim(); + // A seam whose marker launched kernels IS live; what failed is the hand-transcribed device_kernel + // NAME. Banning that seam (below) would forbid the one correct answer, and "descend deeper" is a + // repair for a defect that is not there — a search with no terminating condition. One gemm head + // burned 13.6h and seven captures that way, re-hunting a seam that was right on attempt 1. + const nameCodes = (selection.codes || []); + const nameOnly = !selection.ok && + (nameCodes.includes('device_kernel_name_mismatch') || + nameCodes.includes('device_kernel_not_in_profile')) && + !nameCodes.includes('device_kernel_not_under_target'); + if (priorTarget && !nameOnly && !attemptedTargets.includes(priorTarget)) + attemptedTargets.push(priorTarget); + const observedNames = (((selection.verdict || {}).kernels_under_target) || []) + .map((entry) => `${entry && entry.name} (${entry && entry.launches} launches)`); + const profileNames = (((selection.verdict || {}).profile_kernel_candidates) || []).slice(0, 20); + const selectionCorrective = selection.ok ? '' : nameOnly ? + ` PRIOR ATTEMPT SELECTED A LIVE SEAM BUT DECLARED THE WRONG GPU KERNEL NAME: ${selection.why}. ` + + `The seam '${priorTarget}' demonstrably launched GPU work, so DO NOT descend to another callable ` + + 'and DO NOT re-run capture to hunt a different seam — keep this target_callable. Fix KERNEL.device_kernel ' + + 'by COPYING a name verbatim from the profiler, never by re-typing or abbreviating it. ' + + (observedNames.length ? `Kernels this seam actually launched: ${observedNames.join('; ')}. ` : '') + + (profileNames.length ? `Names recorded in PROFILE_TOPN: ${profileNames.join('; ')}. ` : '') + + 'Then re-run kernel_selection.py with --profile-top-n "$PROFILE_TOPN" and return its JSON verbatim ' + + 'as selection_validation.' : + ` PRIOR ATTEMPT DID NOT SELECT THE PROFILED GPU KERNEL: ${selection.why}. ` + + 'Treat KERNEL.live_call_seam as prose only. Merge KERNEL.seam_candidates with any missing inner ' + + 'launcher found from source/runtime inspection; preserve existing candidate classifications. ' + + 'BEFORE re-running capture, reclaim prior process-local artifacts: ' + + '`python3 "$SKILL_DIR/scripts/capture_shapes.py" --cleanup-task-dir "$TASK" --no-promote` ' + + '(issue #429 — never accumulate capture.pid-* oracles across retries). ' + + 'Install safe markers for every relevant candidate, never native/JIT kernel_entry objects. Run ' + + 'kernel_selection.py over every process-local capture and all root-call traces with --task-dir "$TASK". ' + + 'Return its JSON verbatim as selection_validation. Select the deepest live inner_launcher/op_seam across all ' + + 'calls/ranks; a fused head must select the whole-op op_seam. Rejecting the previous outer wrapper ' + + 'is not success. Do not return any ATTEMPTED_TARGET_CALLABLES value again.'; + const baselineCorrective = needBaseline ? + ' PRIOR ATTEMPT DID NOT FREEZE A BASELINE. You MUST seed baseline_overlay/ from ' + + 'CURRENT_OVERLAY (the live serving stack = the speedup denominator), declare meta.candidate_bind ' + + '(the ONE overlay entry built from kernel_src/), prove both legs differ via h.assert_legs_differ, ' + + 'then return baseline_frozen:true. An extraction with no frozen baseline is INVALID and will be discarded.' : ''; + log(` ${(opts && opts.label) || role}: extraction contract incomplete ` + + `(${selection.ok ? 'kernel selected' : selection.why}; ` + + `${needBaseline ? 'baseline missing (baseline_overlay/ + meta.candidate_bind)' : 'baseline frozen'}). ` + + `RE-EXTRACTING (retry ${tries}/${BASELINE_EXTRACT_RETRIES}).`); + // Best-effort reclaim before the agent retries (issue #429 storage amplification). + if (ext && ext.task_dir) { + try { + const { execFileSync } = require('child_process'); + execFileSync('python3', [ + `${WORKFLOW_DIR}/scripts/capture_shapes.py`, + '--cleanup-task-dir', String(ext.task_dir), + '--no-promote', + ], { stdio: 'pipe', timeout: 120000 }); + log(` ${(opts && opts.label) || role}: reclaimed capture artifacts under ${ext.task_dir}`); + } catch (cleanupErr) { + log(` ${(opts && opts.label) || role}: capture reclaim skipped (${cleanupErr && cleanupErr.message ? cleanupErr.message : cleanupErr})`); + } + } + ext = await safeAgent( + roleAgent(role, phase, captureIntro + selectionCorrective + baselineCorrective, { + ...(inputs || {}), + CAPTURE_STORAGE_ENV, + CAPTURE_BYTE_BUDGET, + CAPTURE_CASE_BYTE_LIMIT, + CAPTURE_PERSIST_POLICY, + PRIOR_TARGET_CALLABLE: priorTarget, + PRIOR_SELECTION_VALIDATION: (ext && ext.selection_validation) || {}, + ATTEMPTED_TARGET_CALLABLES: attemptedTargets.slice(), + }), + bounded()); } - ext = await safeAgent( - roleAgent(role, phase, captureIntro + selectionCorrective + baselineCorrective, { - ...(inputs || {}), - CAPTURE_STORAGE_ENV, - CAPTURE_BYTE_BUDGET, - CAPTURE_CASE_BYTE_LIMIT, - CAPTURE_PERSIST_POLICY, - PRIOR_TARGET_CALLABLE: priorTarget, - PRIOR_SELECTION_VALIDATION: (ext && ext.selection_validation) || {}, - ATTEMPTED_TARGET_CALLABLES: attemptedTargets.slice(), - }), - opts); - } - const finalSelection = kernelSelectionVerified(head, ext); - if (smokeOk(ext) && !finalSelection.ok) { - log(` ${(opts && opts.label) || role}: kernel selection still unverified after ` + - `${BASELINE_EXTRACT_RETRIES} re-extractions — ABORTING (${finalSelection.why}).`); - return { ...ext, smoke: 'fail', unittest_smoke: 'fail', selection_failed: true, - notes: `kernel selection failed: ${finalSelection.why} — ${ext.notes || ''}` }; - } - if (smokeOk(ext) && !hasFrozenBaseline(ext)) { - log(` ${(opts && opts.label) || role}: STILL no frozen baseline after ${BASELINE_EXTRACT_RETRIES} ` + - `re-extractions — ABORTING this extraction (refusing a fake speedup vs the candidate's own scaffold).`); - return { ...ext, smoke: 'fail', unittest_smoke: 'fail', - notes: `no frozen baseline after ${BASELINE_EXTRACT_RETRIES} re-extractions ` + - `(baseline_overlay/ + meta.candidate_bind required as the speedup denominator) — ${ext.notes || ''}` }; + // Why the attempts stopped. Reported on every failure below so a budget cut is never mistaken for + // a head that is genuinely unextractable -- the first is raised with a knob, the second is not. + const why = budgetSpent + ? `the ${Math.round(CAPTURE_BUDGET_MS / 60000)}min capture budget (args.capture_budget_s) was spent ` + + `after ${tries + 1} attempt(s)` + : `${BASELINE_EXTRACT_RETRIES} re-extractions`; + const finalSelection = kernelSelectionVerified(head, ext); + if (smokeOk(ext) && !finalSelection.ok) { + log(` ${(opts && opts.label) || role}: kernel selection still unverified after ` + + `${why} — ABORTING (${finalSelection.why}).`); + return { ...ext, smoke: 'fail', unittest_smoke: 'fail', selection_failed: true, + capture_budget_spent: budgetSpent, + notes: `kernel selection failed after ${why}: ${finalSelection.why} — ${ext.notes || ''}` }; + } + if (smokeOk(ext) && !hasFrozenBaseline(ext)) { + log(` ${(opts && opts.label) || role}: STILL no frozen baseline after ${why} ` + + `— ABORTING this extraction (refusing a fake speedup vs the candidate's own scaffold).`); + return { ...ext, smoke: 'fail', unittest_smoke: 'fail', capture_budget_spent: budgetSpent, + notes: `no frozen baseline after ${why} ` + + `(baseline_overlay/ + meta.candidate_bind required as the speedup denominator) — ${ext.notes || ''}` }; + } + // A head that never smoke-passed at all, cut by the budget, would otherwise return null/unmarked and + // read downstream as an ordinary extraction failure. Say which it was. + if (budgetSpent && !smokeOk(ext)) { + return { ...(ext || {}), smoke: 'fail', unittest_smoke: 'fail', capture_budget_spent: true, + notes: `extraction never smoke-passed and ${why} — ${(ext && ext.notes) || ''}` }; + } + return ext; + } finally { + // Timers are unref'd, so a leaked stopwatch cannot hold the process open — but it would keep + // firing for a head that already finished, and there is one per extraction. + budget.stop(); } - return ext; } // abDone == the integrator measured BOTH legs (ref + cand) and emitted a real @@ -3209,6 +3322,31 @@ const history = ST.history || { insights: [], ledger: [], milestones: [], bottle // never decomposed into a standalone dense GEMM — so dense-GEMM synth is off for it. function gemmSynthFor(h) { return (h && h.op_kind === 'moe') ? 'false' : GEMM_SYNTH; } +// Every kernel_extractor invocation needs the same evaluation context; only the head and its GPU +// differ. PROFILE_TOPN is the load-bearing one: kernel_selection.py --check-device-kernel reads it +// to reject a mis-declared kernel name BEFORE a capture is paid for, so a call site that forgot it +// silently lost that check. Assembled once here instead of retyped at each site. +function extractorInputs(kernel, gpuId, extra) { + return { + EVAL_DIR, MODEL_PATH, GPU_ID: gpuId, WORKLOAD, KERNEL: kernel, + CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, + ...(profile && profile.profile_workload_json ? { PROFILE_WORKLOAD_JSON: profile.profile_workload_json } : {}), + PROFILE_TOPN: profile ? profile.profile_topN_json : '', + ...(extra || {}), + }; +} + +// The op track adds dense-GEMM synth plus the shape regimes. The unittest MUST span BOTH: steady-state +// serving is decode/TPOT-bound, so a head GEMM tuned only on the GPU-time-dominant prefill M regresses +// decode and loses e2e. The decode M (= running batch ≈ conc) is passed explicitly so it is never +// dropped, plus a per-step M=1. See kernel_extractor.md "Shapes must span BOTH regimes". +const extractOpInputs = (kernel, gpuId) => extractorInputs(kernel, gpuId, { + GEMM_SYNTH: gemmSynthFor(kernel), + REQUIRE_DECODE_BUCKET: true, + DECODE_M_BUCKETS: [1, CONC], + PREFILL_M_NOTE: 'also include the profiled large prefill M (chunk size, ~thousands) per (N,K)', +}); + // =========================================================================== // PHASE: TuningSkillset — the VENDORED tuning skillset, run WHOLE and STANDALONE, BEFORE HeadKernel. // @@ -3676,13 +3814,8 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { const GLOBAL_KB = `${EVAL_DIR}/deep_head/GLOBAL_KB.md`; const prepHead = async (h) => { const ext = await extractWithBaseline( - 'kernel_extractor', 'extract_op', 'Build a standalone op unittest for a head kernel.', { - EVAL_DIR, MODEL_PATH, GPU_ID: GPU_LIST[0], WORKLOAD, KERNEL: h, GEMM_SYNTH: gemmSynthFor(h), - ...(profile && profile.profile_workload_json ? { PROFILE_WORKLOAD_JSON: profile.profile_workload_json } : {}), - CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, - REQUIRE_DECODE_BUCKET: true, DECODE_M_BUCKETS: [1, CONC], - PREFILL_M_NOTE: 'also include the profiled large prefill M (chunk size, ~thousands) per (N,K)', - }, + 'kernel_extractor', 'extract_op', 'Build a standalone op unittest for a head kernel.', + extractOpInputs(h, GPU_LIST[0]), { phase: 'HeadKernel', label: `extract_op ${h.short_name}`, schema: EXTRACT_OP_SCHEMA }); const isDominant = (h.pct_gpu_time || 0) >= HEAD_PROTECT_PCT; if (!ext || ext.smoke !== 'pass' || !ext.task_dir) { @@ -4042,13 +4175,8 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { return ISO.with(1, async (g) => { const gpu = g[0]; const ext = await extractWithBaseline( - 'kernel_extractor', 'extract_op', 'Build a standalone op unittest for a head kernel.', { - EVAL_DIR, MODEL_PATH, GPU_ID: gpu, WORKLOAD, KERNEL: h, GEMM_SYNTH: gemmSynthFor(h), - ...(profile && profile.profile_workload_json ? { PROFILE_WORKLOAD_JSON: profile.profile_workload_json } : {}), - CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, - REQUIRE_DECODE_BUCKET: true, DECODE_M_BUCKETS: [1, CONC], - PREFILL_M_NOTE: 'also include the profiled large prefill M (chunk size, ~thousands) per (N,K)', - }, + 'kernel_extractor', 'extract_op', 'Build a standalone op unittest for a head kernel.', + extractOpInputs(h, gpu), { phase: 'HeadKernel', label: `extract_op ${h.short_name}`, schema: EXTRACT_OP_SCHEMA }); if (!ext || ext.smoke !== 'pass' || !ext.task_dir) return { h, gpu, ext, dead: 'extract' }; const bake = await safeAgent( @@ -4244,18 +4372,8 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { // fused/monolithic head to op_kind=moe with GEMM_SYNTH off (gemmSynthFor) so it is extracted as the // fused op bound at its live seam — never decomposed into a standalone dense GEMM. Nothing is skipped. const ext = await extractWithBaseline( - 'kernel_extractor', 'extract_op', 'Build a standalone op unittest for a head kernel.', { - EVAL_DIR, MODEL_PATH, GPU_ID: h.gpu_id, WORKLOAD, KERNEL: h, GEMM_SYNTH: gemmSynthFor(h), - ...(profile && profile.profile_workload_json ? { PROFILE_WORKLOAD_JSON: profile.profile_workload_json } : {}), - CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, - // The unittest MUST span BOTH regimes. Steady-state serving is decode/TPOT-bound, so a - // head GEMM tuned only on GPU-time-dominant prefill M regresses decode and loses e2e. - // Pass the decode M explicitly (= running batch ≈ conc) so it is never dropped, plus a - // per-step M=1. See kernel_extractor.md "Shapes must span BOTH regimes". - REQUIRE_DECODE_BUCKET: true, - DECODE_M_BUCKETS: [1, CONC], - PREFILL_M_NOTE: 'also include the profiled large prefill M (chunk size, ~thousands) per (N,K)', - }, + 'kernel_extractor', 'extract_op', 'Build a standalone op unittest for a head kernel.', + extractOpInputs(h, h.gpu_id), { phase: 'HeadKernel', label: `extract_op ${h.short_name}`, schema: EXTRACT_OP_SCHEMA }); const isDominant = (h.pct_gpu_time || 0) >= HEAD_PROTECT_PCT; if (!ext || ext.smoke !== 'pass' || !ext.task_dir) { @@ -4570,11 +4688,8 @@ while (want('kernel') && !TIME_DEADLINE_HIT && dispatched < BUDGET && (dispatche // once (no timing conflict) and accepted overlays carry forward in order. const optimized = await parallel(cands.map((c) => async () => { const ext = await extractWithBaseline( - 'kernel_extractor', 'extract', 'Capture shapes + oracle; emit an immutable unittest task dir.', { - EVAL_DIR, MODEL_PATH, GPU_ID: c.gpu_id, WORKLOAD, KERNEL: c, - CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, - ...(profile && profile.profile_workload_json ? { PROFILE_WORKLOAD_JSON: profile.profile_workload_json } : {}), - }, + 'kernel_extractor', 'extract', 'Capture shapes + oracle; emit an immutable unittest task dir.', + extractorInputs(c, c.gpu_id), { phase: 'Milestone', label: `extract ${c.short_name}`, schema: EXTRACT_SCHEMA }); if (!ext || ext.editable === false || ext.unittest_smoke !== 'pass' || !ext.task_dir) { return { c, skip: true, reason: `extraction failed/non-editable (${ext ? ext.notes || ext.unittest_smoke : 'none'})` }; diff --git a/e2e_workflow/roles/kernel_extractor.md b/e2e_workflow/roles/kernel_extractor.md index 860b68283..ec793275d 100644 --- a/e2e_workflow/roles/kernel_extractor.md +++ b/e2e_workflow/roles/kernel_extractor.md @@ -27,15 +27,36 @@ You are invoked once per kernel candidate. Read first: harness_lib.py # VENDORED scripts/harness_lib.py — the SHARED timing/correctness lib; IMMUTABLE leg_runner.py # VENDORED scripts/leg_runner.py — runs ONE leg under the ambient overlay; IMMUTABLE overlay_setup.py # VENDORED scripts/overlay_setup.py — builds the candidate overlay; IMMUTABLE + kernel_selection.py # VENDORED scripts/kernel_selection.py — supplies the kernel-name matcher the + # baseline dispatch gate uses; stdlib-only; IMMUTABLE unittest.py # driver: h.measure_legs + h.run_correctness, prints the metric; IMMUTABLE meta.json # name, source path, target_callable, candidate_bind, shapes, dtypes, backend, # regime, served_regimes, build, random_draws (default 3), checksum ``` -**Vendor the three shared scripts into the task dir** -(`for f in harness_lib.py leg_runner.py overlay_setup.py; do cp "$SKILL_DIR/scripts/$f" "$TASK/"; done`). +**Vendor the four shared scripts into the task dir** +(`for f in harness_lib.py leg_runner.py overlay_setup.py kernel_selection.py; do cp "$SKILL_DIR/scripts/$f" "$TASK/"; done`). `unittest.py` imports `harness_lib` for ALL timing + correctness — never hand-roll a timing loop or an allclose check. This is what makes every task measure the same way; it also keeps the task self-contained + immutable (the validator sha-checks them alongside `reference_io.pt`). +`kernel_selection.py` supplies the kernel-name matcher only (stdlib-only, no torch); without it +`h.assert_baseline_dispatch` degrades to a no-op and the baseline is never proven to be deployment. + +🔴 **Rehydrate the oracle with `h.reconstruct_captured` — never hand-roll the walk.** `capture_shapes` +records more than data+dtype+shape: loader-set Python attributes (`attrs`) carry BACKEND DISPATCH +DECISIONS — aiter's fused-MoE gate reads `getattr(w1, "is_shuffled", False)` to choose FlyDSL vs CK. +`torch.save` does not persist them and `.to(device)` drops them, so a hand-written rehydrator silently +replays the op on a DIFFERENT kernel than the captured server ran — and the golden, frozen from that +same wrong baseline, agrees with itself. `h.reconstruct_captured` (and `h.eager_cases_from_oracle`, +which calls it) re-apply `attrs` after the device move. Need an extra step (a uint8 raw view for a +packed fp4 operand)? Wrap it — do not replace it. + +When rehydration genuinely cannot preserve the attribute (`.view(dt)`/`.set_()` return fresh tensors), +or the oracle predates `attrs` and the flag is simply not in the file, declare it in +`meta.live_tensor_attrs` (`{operand: {attr: value}}`, `"pos[]"` for a positional) and call +`h.apply_declared_attrs(args, META)` ONCE, after rehydration, where every case set draws its operands — +and pass `META` to `h.eager_cases_from_oracle` so correctness gets it too. Record WHY the declared value +is deployment's (a server.log line, a profile): `assert_baseline_dispatch` proves the result reaches the +right kernel, not that the value is right. ### 🔴 THE TWO LEGS ARE THE SAME CODE UNDER TWO PYTHONPATHS — read this before writing anything There is no `baseline_callable`, and no second copy of the source to time against. Both legs run the @@ -80,7 +101,8 @@ Inputs: `EVAL_DIR`, `MODEL_PATH`, `GPU_ID`, `WORKLOAD`, `KERNEL` (the Architect' short_name, classification, extract_hint = the `module:attr` callable to hook, candidate_backends, regime, and — when an upstream TraceLens prior was available — OPTIONAL `source_hint` (resolved source file), `launcher_hint` (launcher seam), `bound_type`), `CURRENT_OVERLAY` (the accepted-kernel stack -carried forward — may be empty on the first milestone), `CURRENT_FLAGS`/`CURRENT_ENV`, `SKILL_DIR`. +carried forward — may be empty on the first milestone), `CURRENT_FLAGS`/`CURRENT_ENV`, `SKILL_DIR`, +and `PROFILE_TOPN` (the profiler's own topN JSON — the ONLY authority on how a device kernel is spelled). ### Resolve + HONOR the ONLINE REGIME first (same contract as PHASE=extract_op) The #1 cause of "isolated win, e2e loss/crash" is a unittest that SYNTHESIZES its inputs with OFFLINE @@ -155,6 +177,20 @@ freeze an out-of-regime oracle nobody should trust. capture overlay, driven by the SAME workload as the profile so shapes match the regime: ```bash TASK="$EVAL_DIR/kernels/_task"; mkdir -p "$TASK" + # FIRST, before any capture: prove KERNEL.device_kernel is a name the PROFILER actually recorded. + # It is hand-transcribed prose until this passes. One `_HA_S_` mistyped as `_AS_` is invisible to + # every later check except this one, and reads downstream as "the seam is wrong" -- which sends the + # extractor descending through callables it can never fix. This costs no GPU, no server, no capture. + # On failure: DO NOT proceed. COPY the right name verbatim out of profile_kernel_candidates. + # PROFILE_TOPN can be empty on a resumed run whose state predates the profile. Then the check + # CANNOT run -- say so in `notes`; its absence is not a pass, and must not abort the extraction. + if [ -s "$PROFILE_TOPN" ]; then + python3 "$SKILL_DIR/scripts/kernel_selection.py" \ + --target "" --device-kernel "" \ + --profile-top-n "$PROFILE_TOPN" --check-device-kernel || exit 1 + else + echo "WARN: PROFILE_TOPN is empty - pre-capture device_kernel name check SKIPPED (report in notes)" + fi # FREEZE the live serving stack as this task's baseline env, then hang the capture hook off a COPY # of it. --from is what stacks them: two overlay dirs on PYTHONPATH do NOT compound (only the first # sitecustomize is imported), so capturing on a bare hook overlay would silently capture the @@ -182,6 +218,7 @@ freeze an out-of-regime oracle nobody should trust. --capture-meta "$TASK"/capture.pid-*.rank-*/meta.json \ --torch-trace "$TASK"/selection_trace.pid-*.rank-*.call-*.json \ --candidate-target "" \ + --profile-top-n "$PROFILE_TOPN" \ --task-dir "$TASK" \ --out "$TASK/selection_validation.json" ``` @@ -195,6 +232,13 @@ freeze an out-of-regime oracle nobody should trust. every `capture.pid-*` directory (issue #429 — do NOT leave per-rank oracles around). On failure or before a capture retry, reclaim without promote: + 🔴 **`device_kernel_name_mismatch` is NOT `device_kernel_not_under_target`.** The first means the + seam launched GPU work but nothing spelled like `--device-kernel`; the verdict lists what it *did* + launch in `kernels_under_target`. The seam is LIVE — keep it, keep the capture, and fix the NAME by + copying one of those strings verbatim. Descending to another callable there repairs a defect that is + not present and starts a search with no terminating condition. Only `device_kernel_not_under_target` + (the marker launched nothing at all) means the seam is wrong and you should go deeper. + ```bash python3 "$SKILL_DIR/scripts/capture_shapes.py" --cleanup-task-dir "$TASK" --no-promote ``` @@ -317,8 +361,15 @@ freeze an out-of-regime oracle nobody should trust. case sigs). The BASELINE leg records its outputs for those draws in its own process: ```python base_out = h.baseline_random_outputs(TASK, meta, draws=meta.get("random_draws", 3)) + floor = h.baseline_noise_floor(TASK, meta, tol, draws=meta.get("random_draws", 3), + baseline_outputs=base_out) # pass to run_correctness(noise_floor=) ``` - and the candidate is compared against them (same seed ⇒ same inputs). **🔴 Do NOT randomize SHAPES + and the candidate is compared against them (same seed ⇒ same inputs). **🔴 Always pass + `noise_floor=`.** Records the baseline a second time at the same seed and scores it against the + first: an op that reduces with atomics or split-k (FlyDSL MoE, `persist_cu*`) does not reproduce + itself bit-for-bit, and `correct`'s `atol = tol*RMS(ref)` inflates a 1e-05 wobble on a near-zero + element into a ~0.5 "relative error". Without the floor the honest candidate is FAILED for the + baseline's own nondeterminism. Costs one extra oracle leg (~15-25s). **🔴 Do NOT randomize SHAPES — dims stay online-aligned; only the input VALUES vary.** Fold its correctness verdict into the overall PASS/FAIL (a delta vs baseline on ANY draw FAILS the unittest); print its per-draw `speedup` as a SECONDARY robustness @@ -466,22 +517,27 @@ freeze an out-of-regime oracle nobody should trust. is still a byte-copy of the baseline file, so the legs differ by PATH while `speedup≈1.0` — that is the expected smoke result. If the target cannot be resolved on the live stack at all, do NOT fall back to a `kernel_src/` strawman: return `editable:false` with a clear reason. - > **Exit-code contract — a missing replay leg is a UT DEFECT, not a kernel/smoke failure.** The UT - > routes correctness through `h.run_correctness(...)`, which for a graph-deploy kernel (`cuda_graph=true`) - > RAISES `h.HarnessIncompleteError` when no ≥2-shape replay bundle was wired — and it has ALREADY - > printed the `UT_HARNESS_INCOMPLETE: …` sentinel line itself (so the smoke sees it even if `main()` - > forgets to catch). The generated `main()` MUST translate the exception to a DEDICATED exit code; do - > NOT re-print the sentinel (it is already on stdout — a second print is just noise): + > **Exit-code contract — a missing replay leg or a wrong baseline is a UT DEFECT, not a kernel/smoke + > failure.** TWO harness calls raise `h.HarnessIncompleteError`: `h.run_correctness(...)`, when a + > graph-deploy kernel (`cuda_graph=true`) was wired no ≥2-shape replay bundle, and `h.measure_legs(...)`, + > when the BASELINE leg never launches `meta.device_kernel`. Both have ALREADY printed the + > `UT_HARNESS_INCOMPLETE: …` sentinel themselves. The generated `main()` MUST wrap **both** — not just + > correctness, or a dispatch mismatch escapes as an uncaught traceback and scores exit 1, the code + > reserved for a genuine kernel failure. Do NOT re-print the sentinel (already on stdout): > ```python > try: + > per_case = h.measure_legs(TASK, META) # raises if the baseline is not deployment > ok, report = h.run_correctness(META["regime"], ...) # eager+random+replay legs > except h.HarnessIncompleteError: > sys.exit(3) # 3 = regenerate UT (sentinel already printed) > sys.exit(0 if ok else 1) # 1 = real correctness FAIL, 2 = env > ``` - > On smoke **exit 3 OR a `UT_HARNESS_INCOMPLETE` line on stdout: REGENERATE the UT** — add the replay - > bundle (build ≥2 boundary cases via `h.boundary_decode_seq_lens`/`h.shuffled_block_table` for attn, or - > the family×M-buckets for gemm; wire `fill/run/read_out`) and re-run the smoke. Retry up to 3 times. + > On smoke **exit 3 OR a `UT_HARNESS_INCOMPLETE` line on stdout: REGENERATE the UT** — read WHICH + > defect the sentinel names. Missing replay bundle: add it (build ≥2 boundary cases via + > `h.boundary_decode_seq_lens`/`h.shuffled_block_table` for attn, or the family×M-buckets for gemm; + > wire `fill/run/read_out`). Baseline dispatch mismatch: recapture the oracle, declare the lost + > attribute in `meta.live_tensor_attrs`, or re-select the seam — NEVER by exporting a tuned config + > the captured server did not have. Then re-run the smoke. Retry up to 3 times. > Do **NOT** record `unittest_smoke:"fail"` or drop the head for exit 3 — that status is reserved for a > genuine baseline-bind / correctness failure (exit 1). Only after 3 failed regenerations set > `unittest_smoke:"fail"` with `reason="harness_incomplete_unrecoverable"`. @@ -662,7 +718,8 @@ needs an op task dir the **Op Benchmarker** can bake-off across backends. `edit= Inputs: `EVAL_DIR`, `MODEL_PATH`, `GPU_ID`, `WORKLOAD`, `KERNEL` (Architect head candidate: short_name, op_kind=gemm|attn, the profiled `shapes`, dtype, regime, `target_callable` for attn, and OPTIONAL TraceLens `source_hint`/`launcher_hint`/`bound_type`), `GEMM_SYNTH` (bool, default true), -`CURRENT_FLAGS`/`CURRENT_ENV`, `SKILL_DIR`, and OPTIONAL `PROFILE_WORKLOAD_JSON` (the profiler's +`CURRENT_FLAGS`/`CURRENT_ENV`, `SKILL_DIR`, `PROFILE_TOPN` (the profiler's own topN JSON — the ONLY +authority on how a device kernel is spelled), and OPTIONAL `PROFILE_WORKLOAD_JSON` (the profiler's per-(shape,dtype) weighted workload model — slice this kernel's cases into `workload_path`, see below). > **TraceLens shape double-check (mandatory when the shapes came from TraceLens).** If `KERNEL.shapes` diff --git a/e2e_workflow/scripts/capture_shapes.py b/e2e_workflow/scripts/capture_shapes.py index cf8482420..34f096f8c 100755 --- a/e2e_workflow/scripts/capture_shapes.py +++ b/e2e_workflow/scripts/capture_shapes.py @@ -7,8 +7,10 @@ for the first few DISTINCT input-shape signatures seen during a short bench window, and writes a torch-loadable `reference_io.pt` + `meta.json`. -This module is meant to be imported at server startup through an overlay PYTHONPATH (it registers the -hook on import), OR called as a function from a custom preimport. It does NOT launch the server +This module is meant to be imported at server startup through an overlay PYTHONPATH (it arms the +hook there and binds it when the server itself imports the target — importing the target from +sitecustomize would reorder the application's imports; see install()), OR called as a function from a +custom preimport. It does NOT launch the server itself — pair it with scripts/bench_e2e.sh (drive the same workload as the profile so shapes match the regime). @@ -38,11 +40,23 @@ # Heuristic names for expert/static parameter tensors (used by moe_slim / share_large). _WEIGHT_KEY_RE = re.compile( r"(^w[123]$|^weight$|expert_w|gate_up|down_proj|up_proj|_weight$)", re.I) +# Names already recorded as first-class snapshot keys: a __dict__ duplicate would override the real +# one on restore. (Plain torch.Tensor keeps these as descriptors, but subclasses/wrappers do not.) +_ATTR_SKIP = frozenset(("shape", "dtype", "device", "data", "contiguous", "attrs")) +_ATTR_MAX_COUNT = 32 +_ATTR_MAX_STR = 200 +# Env that steers BACKEND DISPATCH (tuned-config tables, backend enables, arch pins) — recorded for +# comparison only, never re-exported: replaying with a table the captured server did not have is its +# own infidelity. Broad by design, so credential-looking names are redacted: this runs in the SERVER +# process (GEAK_KB_STORE_TOKEN matches ^GEAK_) and meta.json is published with the task dir. +_ENV_CAPTURE_RE = re.compile( + r"^(AITER|GEAK|SGLANG|VLLM|TORCH|TORCHINDUCTOR|PYTORCH|TRITON|HIP|ROCM|HSA|CK|GPU)_|FLYDSL") +_ENV_REDACT_RE = re.compile(r"(KEY|TOKEN|SECRET|PASSWORD|PASSWD|CRED|AUTH|COOKIE|SESSION)", re.I) _STATE = { "target": None, "out_dir": None, "max_cases": 5, "num_steps": 0, "records": [], "seen": set(), "lock": threading.Lock(), "orig": None, - "mod": None, "attr": None, "installed": False, "calls": 0, + "mod": None, "attr": None, "installed": False, "bind_pending": False, "calls": 0, # regime coverage for the oracle: the classic failure is a single-case oracle (only ONE shape recorded, # e.g. one decode step), which under-tests correctness. We guarantee at least one case per regime # (decode vs prefill) even if that overshoots max_cases, so the immutable oracle exercises BOTH the q=1 @@ -371,6 +385,21 @@ def reclaim_workspace_captures(eval_dir, workspace_budget=0): return telemetry +def _env_snapshot(): + """Dispatch-steering env of the capturing process (see ``_ENV_CAPTURE_RE``), secrets redacted. + + Lets a UT that reproduces the wrong kernel be diagnosed from the task dir instead of server logs: + an empty ``AITER_CONFIG_FMOE`` here means the captured server ran the heuristic path, so a UT that + exports a tuned table is measuring a third code path deployment never took. + """ + env = {} + for key in sorted(os.environ): + if _ENV_CAPTURE_RE.search(key): + env[key] = ("" if _ENV_REDACT_RE.search(key) + else str(os.environ[key])[:_ATTR_MAX_STR]) + return env + + def _write_capture_manifest(out_dir, extra=None): """Lightweight size/shape manifest written when the heavy oracle is skipped or reclaiming.""" s = _STATE @@ -453,14 +482,76 @@ def _torch(): return torch +def _attr_value(value): + """``(value, keep)`` — keep only attributes that survive ``torch.save`` as plain data. + + Anything else (tensors, modules, callables) is DROPPED rather than repr'd: a restored repr string + would still read truthy to a ``getattr(w, "...", False)`` dispatch gate. + """ + if value is None or isinstance(value, (bool, int, float)): + return value, True + if isinstance(value, str): + return value[:_ATTR_MAX_STR], True + if (isinstance(value, (list, tuple)) and len(value) <= 8 + and all(v is None or isinstance(v, (bool, int, float, str)) for v in value)): + return type(value)(value), True + return None, False + + +def _tensor_attrs(x): + """Loader-set Python attributes from the tensor's ``__dict__`` — ``torch.save`` drops these. + + Several backends carry their DISPATCH DECISION there rather than in the data: aiter's fused-MoE + gate reads ``getattr(w1, "is_shuffled", False)`` to choose FlyDSL vs CK, so an oracle replayed + without the label runs a different kernel than the captured server did. + """ + try: + items = list((getattr(x, "__dict__", None) or {}).items()) + except Exception: + return {} + attrs = {} + for key, value in items: + if not isinstance(key, str) or key.startswith("_") or key in _ATTR_SKIP: + continue + attr_value, keep = _attr_value(value) + if not keep: + continue + attrs[key] = attr_value + if len(attrs) >= _ATTR_MAX_COUNT: + break + return attrs + + +def _to_cpu_clone(x): + """``x.detach().to("cpu").clone()``, with a byte-view fallback for sub-byte dtypes. + + ROCm torch 2.9 has no ``copy_kernel`` for ``float4_e2m1fn_x2``, so the plain D2H raises + ``NotImplementedError`` — and the hook's blanket except then records ZERO cases, silently, for + exactly the MXFP4 seams we most want to capture. The same storage copied through a ``uint8`` + view needs no per-dtype kernel and is bitwise-equal. Mirrored by harness_lib._to_device. + """ + torch = _torch() + try: + return x.detach().to("cpu").clone() + except (NotImplementedError, RuntimeError) as exc: + try: + return x.detach().view(torch.uint8).to("cpu").clone().view(x.dtype) + except Exception: + raise exc # fallback inapplicable (non-contiguous, OOM, ...) — report the real cause + + def _snapshot(x): """Detach+clone tensors to CPU so later in-place ops can't corrupt the oracle. Pass scalars/None through; summarize unsupported objects by repr so the record stays loadable.""" torch = _torch() if torch.is_tensor(x): - return {"__tensor__": True, "data": x.detach().to("cpu").clone(), + snap = {"__tensor__": True, "data": _to_cpu_clone(x), "dtype": str(x.dtype), "device": str(x.device), "shape": list(x.shape), "contiguous": bool(x.is_contiguous())} + attrs = _tensor_attrs(x) + if attrs: + snap["attrs"] = attrs + return snap if isinstance(x, (list, tuple)): return type(x)(_snapshot(v) for v in x) if isinstance(x, dict): @@ -771,6 +862,7 @@ def walk(o): "budget_exceeded": bool(s.get("budget_exceeded")), "budget_skip_count": int(s.get("budget_skip_count") or 0), "oracle_save_count": int(s.get("oracle_save_count") or 0), + "capture_env": _env_snapshot(), "build": False, # default: pure-python/triton; Extractor flips to True for HIP/CK/asm tasks "note": "Oracle captured from baseline. Do NOT edit unittest.py or reference_io.pt during opt.", } @@ -840,32 +932,65 @@ def _w(*args, **kwargs): return _w +class _BindOnImport: + """Meta-path shim that runs ``callback`` right after ``mod_name`` is first executed. + + It never imports anything itself: it defers to the normal finders for the spec and only decorates + that spec's loader. See install() for why the overlay must not pull the target in early. + """ + + def __init__(self, mod_name, callback): + self.mod_name, self.callback, self._busy = mod_name, callback, False + + def find_spec(self, name, path=None, target=None): + if name != self.mod_name or self._busy: + return None + self._busy = True # our own find_spec re-enters the meta path; don't recurse + try: + spec = importlib.util.find_spec(name) + except Exception: + return None + finally: + self._busy = False + if spec is None or spec.loader is None or not hasattr(spec.loader, "exec_module"): + return None + inner, cb = spec.loader.exec_module, self.callback + + def exec_module(module): + inner(module) + try: + sys.meta_path.remove(self) + except ValueError: + pass + cb() + + spec.loader.exec_module = exec_module # loader instance is created per-spec, so this is local + return spec + + def install(target, out_dir, max_cases=5): """Wrap module:attr to record I/O. Registers an atexit flush. Idempotent. - Fails FAST at install (server startup) if the target is a native/non-Python callable that a plain - Python wrapper cannot safely stand in for — converting the old unpredictable mid-run SIGSEGV (which - took the whole server down and lost the run) into a clear, actionable startup error so the Extractor - picks a Python-level seam. Override with CAPTURE_WRAP_UNSAFE=1 to force (e.g. when the caller only - reads shapes, never the JIT internals).""" + If the target module is not imported yet — the normal case, since the overlay's sitecustomize runs + during interpreter startup — the wrap is DEFERRED until the application imports it. Importing the + target from sitecustomize instead would reorder the whole application's imports, and that is not + hypothetical: pulling ``aiter.fused_moe`` in at startup makes FlyDSL's JIT abort the process with + ``LLVM ERROR: Do not know how to expand this operator's operand!`` when it later compiles + flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w2, while the identical run without the overlay compiles the + same kernel fine. An instrumentation hook must not perturb what it observes. + + Fails FAST at install (server startup) on things that need no import — a bad byte budget or persist + policy, or a module root that does not exist. The non-Python-callable check necessarily waits for + the bind: a plain Python wrapper cannot safely stand in for a native/triton-JIT callable (it + SIGSEGVs the server mid-run), so we raise a clear error the moment the target resolves rather than + letting the Extractor find out from a corefile. Override with CAPTURE_WRAP_UNSAFE=1 to force.""" s = _STATE - if s["installed"]: + if s["installed"] or s.get("bind_pending"): return mod_name, attr = target.split(":", 1) - mod = importlib.import_module(mod_name) - # attr may be dotted (e.g. Class.method): resolve the binding owner + leaf, but keep the full - # module path + dotted attr in meta so kernel_selection's f"{module}:{attr}" == target check holds. - owner = mod - for part in attr.split(".")[:-1]: - owner = getattr(owner, part) - leaf = attr.split(".")[-1] - orig = getattr(owner, leaf) - if not _wrappable(orig) and os.environ.get("CAPTURE_WRAP_UNSAFE", "0") != "1": - raise RuntimeError( - f"[capture_shapes] refusing to wrap non-Python callable {target} " - f"({type(orig).__module__}.{type(orig).__name__}): a plain-function stand-in for a native/" - f"triton-JIT callable SIGSEGVs the server (e.g. mxfp4 matmul_ogs). Hook a Python-level seam " - f"(its caller) instead, or set CAPTURE_WRAP_UNSAFE=1 to force.") + root = mod_name.split(".")[0] + if root not in sys.modules and importlib.util.find_spec(root) is None: + raise ModuleNotFoundError(f"[capture_shapes] no module named {root!r} (target {target})") out_dir = _process_out_dir(out_dir) try: byte_budget = parse_byte_budget(os.environ.get("CAPTURE_BYTE_BUDGET", _DEFAULT_BYTE_BUDGET)) @@ -882,8 +1007,7 @@ def install(target, out_dir, max_cases=5): f"[capture_shapes] invalid CAPTURE_PERSIST_POLICY={persist_policy!r} " f"(expected full|share_large|moe_slim)") s.update(target=target, out_dir=out_dir, max_cases=int(max_cases), - orig=orig, mod=mod, attr=attr, installed=True, - byte_budget=byte_budget, case_byte_limit=case_byte_limit, + attr=attr, byte_budget=byte_budget, case_byte_limit=case_byte_limit, persist_policy=persist_policy, share_min_bytes=share_min_bytes or (16 << 20), shared_tensors={}, shared_bytes_est=0, oracle_bytes_est=0, budget_exceeded=False, @@ -894,13 +1018,43 @@ def install(target, out_dir, max_cases=5): s["flush_every"] = 1 elif os.environ.get("CAPTURE_FLUSH_EVERY"): s["flush_every"] = max(1, int(os.environ["CAPTURE_FLUSH_EVERY"])) + if mod_name in sys.modules: + _bind(mod_name, attr) + return + s["bind_pending"] = True + sys.meta_path.insert(0, _BindOnImport(mod_name, lambda: _bind(mod_name, attr))) + sys.stderr.write( + f"[capture_shapes] armed {target}; will hook when {mod_name} is imported\n") + + +def _bind(mod_name, attr): + """Resolve the target and swap in the wrapper. Runs either inline (module already imported) or + from the _BindOnImport shim, i.e. on the application's own import of the module.""" + s = _STATE + if s["installed"]: + return + mod = importlib.import_module(mod_name) + # attr may be dotted (e.g. Class.method): resolve the binding owner + leaf, but keep the full + # module path + dotted attr in meta so kernel_selection's f"{module}:{attr}" == target check holds. + owner = mod + for part in attr.split(".")[:-1]: + owner = getattr(owner, part) + leaf = attr.split(".")[-1] + orig = getattr(owner, leaf) + if not _wrappable(orig) and os.environ.get("CAPTURE_WRAP_UNSAFE", "0") != "1": + raise RuntimeError( + f"[capture_shapes] refusing to wrap non-Python callable {s['target']} " + f"({type(orig).__module__}.{type(orig).__name__}): a plain-function stand-in for a native/" + f"triton-JIT callable SIGSEGVs the server (e.g. mxfp4 matmul_ogs). Hook a Python-level seam " + f"(its caller) instead, or set CAPTURE_WRAP_UNSAFE=1 to force.") + s.update(orig=orig, mod=mod, installed=True, bind_pending=False) setattr(owner, leaf, _make_wrapper(orig)) atexit.register(_flush) sys.stderr.write( - f"[capture_shapes] hooked {target}; recording up to {max_cases} cases -> {out_dir}" - f" (byte_budget={byte_budget or 'unlimited'}" - f" case_limit={case_byte_limit or 'unlimited'}" - f" policy={persist_policy})\n") + f"[capture_shapes] hooked {s['target']}; recording up to {s['max_cases']} cases -> {s['out_dir']}" + f" (byte_budget={s['byte_budget'] or 'unlimited'}" + f" case_limit={s['case_byte_limit'] or 'unlimited'}" + f" policy={s['persist_policy']})\n") # Allow configuration purely via env (so a generic overlay sitecustomize can call install()): diff --git a/e2e_workflow/scripts/harness_lib.py b/e2e_workflow/scripts/harness_lib.py index 3144b52a0..ff967f096 100644 --- a/e2e_workflow/scripts/harness_lib.py +++ b/e2e_workflow/scripts/harness_lib.py @@ -38,13 +38,16 @@ form (observed-vs-ceiling) available to any downstream e2e comparison; an observed delta far above the ceiling is box drift / measurement error, not the kernel. """ +import importlib.util import json import math import os +import re import shutil import signal import subprocess import sys +import tempfile import time @@ -600,11 +603,44 @@ def to_device_like(ref, dev): return ref.to(dev) if hasattr(ref, "to") else ref +def apply_captured_attrs(t, attrs): + """Re-attach the loader-set attributes ``capture_shapes._tensor_attrs`` recorded. + + MUST run AFTER any ``.to(device)``: ``.to()`` returns a fresh tensor with an empty ``__dict__``, + so attributes applied before the move are silently dropped — exactly how a replayed MoE oracle + loses ``w1.is_shuffled`` and falls into a different dispatch branch than deployment. + """ + for key, value in (attrs or {}).items(): + try: + setattr(t, key, value) + except (AttributeError, RuntimeError, TypeError): + pass # tensor subclasses may refuse arbitrary attributes; a missing label beats a crash + return t + + +def _to_device(t, device): + """``t.to(device)`` with a byte-view fallback for sub-byte dtypes. + + Reverse of capture_shapes._to_cpu_clone: ROCm torch 2.9 has no ``copy_kernel`` for + ``float4_e2m1fn_x2``, so rehydrating an MXFP4 oracle onto the GPU raises. Same storage, + same bits, through a ``uint8`` view. + """ + try: + return t.to(device) + except (NotImplementedError, RuntimeError) as exc: + try: + import torch + return t.view(torch.uint8).to(device).view(t.dtype) + except Exception: + raise exc + + def reconstruct_captured(obj, device="cpu"): """Inverse of capture_shapes._snapshot (shared refs must already be resolved).""" if isinstance(obj, dict) and obj.get("__tensor__"): t = obj["data"] - return t.to(device) if hasattr(t, "to") else t + t = _to_device(t, device) if hasattr(t, "to") else t + return apply_captured_attrs(t, obj.get("attrs")) if isinstance(obj, dict) and set(obj.keys()) == {"__repr__"}: return obj["__repr__"] if isinstance(obj, dict): @@ -614,6 +650,54 @@ def reconstruct_captured(obj, device="cpu"): return obj +def apply_declared_attrs(args, meta): + """Apply ``meta.live_tensor_attrs`` = {operand: {attr: value}} to a rehydrated ``args`` bundle. + + The RETROFIT path for an oracle captured before ``capture_shapes`` recorded ``attrs``: the blob + carries the loader's preshuffled BYTES with its FLAG stripped, unrecoverable from the file, so the + declaration is the only repair short of recapturing. A capture that HAS ``attrs`` needs nothing — + ``reconstruct_captured`` replays those, and an explicit declaration overrides them. + + Operands are addressed by kwarg name, or ``"pos[]"`` for a positional, in any of the bundle + shapes in use: the ``{"pos", "kw"}`` split, a flat kwargs mapping (``fn(**args)``, which is what + ``iter_eager_cases_from_oracle`` yields), or a bare positional sequence. + + RAISES on a declaration that matches no operand. The per-task helper this replaces returned early + on an empty spec, so a typo'd or stale name read exactly like "nothing to restore" and the leg + went on quietly measuring the wrong backend — the failure this mechanism exists to prevent. + """ + spec = (meta or {}).get("live_tensor_attrs") or {} + if not spec: + return args + if isinstance(args, dict) and ("pos" in args or "kw" in args): + pos, kw = list(args.get("pos") or ()), (args.get("kw") or {}) + elif isinstance(args, dict): + pos, kw = [], args + elif isinstance(args, (list, tuple)): + pos, kw = list(args), {} + else: + pos, kw = [], {} + missing = [] + for name, attrs in spec.items(): + m = re.match(r"^pos\[(\d+)\]$", str(name)) + if m: + i = int(m.group(1)) + target = pos[i] if i < len(pos) else None + else: + target = kw.get(name) + if target is None or not hasattr(target, "shape"): + missing.append(name) + continue + apply_captured_attrs(target, attrs) + if missing: + raise HarnessIncompleteError( + "meta.live_tensor_attrs declares %s, which %s not a tensor operand of this call. The " + "declaration exists to restore a dispatch-steering attribute the capture dropped; a name " + "that does not land restores nothing and the leg silently runs the wrong backend." + % (", ".join(repr(x) for x in missing), "is" if len(missing) == 1 else "are")) + return args + + def load_reference_io(path, map_location="cpu"): """Load a capture_shapes oracle blob (supports optional ``shared`` weight pool).""" torch = _torch() @@ -640,8 +724,12 @@ def resolve_oracle_shared(obj, shared): return obj -def iter_eager_cases_from_oracle(path, device="cpu"): - """Yield ``{args, ref, sig, regime}`` one record at a time (memory-friendly for multi-GiB MoE).""" +def iter_eager_cases_from_oracle(path, device="cpu", meta=None): + """Yield ``{args, ref, sig, regime}`` one record at a time (memory-friendly for multi-GiB MoE). + + Pass ``meta`` when the task declares ``live_tensor_attrs``, else the correctness cases miss the + dispatch flag the timing legs were retrofitted with and the two gates grade different backends. + """ blob = load_reference_io(path, map_location="cpu") shared = blob.get("shared") or {} for record in blob.get("records") or []: @@ -657,16 +745,16 @@ def iter_eager_cases_from_oracle(path, device="cpu"): for name, value in zip(names, args_pos): args.setdefault(name, value) yield { - "args": args, + "args": apply_declared_attrs(args, meta), "ref": ref, "sig": record.get("sig", ""), "regime": record.get("regime", ""), } -def eager_cases_from_oracle(path, device="cpu"): +def eager_cases_from_oracle(path, device="cpu", meta=None): """Materialize all eager cases; prefer ``iter_eager_cases_from_oracle`` for large oracles.""" - return list(iter_eager_cases_from_oracle(path, device=device)) + return list(iter_eager_cases_from_oracle(path, device=device, meta=meta)) def check_correct_multi_lazy(call, case_iter, tol, max_keep_live=2): @@ -871,7 +959,7 @@ def check_graph_replay(fill, run, read_out, cases, tol, capture_idx=0, warmup=3) # --------------------------------------------------------------------------- (b) random-value parity vs live baseline def check_random_vs_baseline(baseline_call, current_call, shapes, tol, draws=3, warmup=10, repeats=50, inner=1, graph=False, seed=0, - baseline_outputs=None): + baseline_outputs=None, noise_floor=None, noise_margin=2.0): """Validate the candidate against the LIVE frozen baseline on MANY RANDOM INPUT VALUE DRAWS at the SAME online-aligned shapes (NOT random shapes — dims are fixed per `sig`, only values vary). The frozen oracle (`reference_io.pt`) pins ONE recorded input+golden; this catches value-dependent bugs @@ -892,6 +980,10 @@ def check_random_vs_baseline(baseline_call, current_call, shapes, tol, `"|"`. Same seed => same inputs, so the two legs never have to be co-resident. When it is given, `baseline_call` is ignored and `speedup` is None here (timing comes from `measure_legs`). + `noise_floor` (from `baseline_noise_floor`) = per-key error the BASELINE shows against ITSELF at the + same seed. A case that misses `tol` but stays within `noise_margin` x that floor is passed and + labelled — the deviation is the op's own launch-to-launch reduction order, not the candidate's. + Absent (None) the gate is exactly as strict as before. `baseline_call(args) -> out` is the LEGACY in-process form, kept for op_bench / single-process tasks. `current_call(args) -> out` invokes the candidate in kernel_src/. `shapes` is a list of {"sig":