diff --git a/.github/workflows/ci-l0-checks.yml b/.github/workflows/ci-l0-checks.yml index ce1e313c4..179deccc1 100644 --- a/.github/workflows/ci-l0-checks.yml +++ b/.github/workflows/ci-l0-checks.yml @@ -89,6 +89,8 @@ jobs: e2e_workflow/scripts/tests/test_parse_profile.py \ e2e_workflow/scripts/tests/test_capture_shapes.py \ e2e_workflow/scripts/tests/test_overlay_setup.py \ + e2e_workflow/scripts/tests/test_seam_trace.py \ + e2e_workflow/scripts/tests/test_kernel_selection.py \ e2e_workflow/scripts/tests/test_attribute_weights_edges.py \ e2e_workflow/scripts/tests/test_op_bench.py \ e2e_workflow/scripts/tests/test_harness_lib.py \ @@ -212,7 +214,7 @@ jobs: /tmp/gitleaks dir . --no-banner --redact --config .gitleaks.toml node-regression: - name: Node regression (expert_skills) + name: Node regression (expert_skills, kernel canonicalization) runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -223,6 +225,12 @@ jobs: # Pure node (fs/path only) — no npm install needed. - name: expert_skills OFF-identical regression run: node e2e_workflow/scripts/test_expert_skills_off_identical.js + # The JS gate and the Python selection verdict canonicalize the same kernel symbols. Only this + # job can run the JS half, so without it the two implementations are pinned on the Python side + # alone and are free to drift — which lets a kernel pass one side and be refused by the other. + # The Python half is TestTheSharedFixtureHoldsOnThisSide in tests/test_kernel_selection.py. + - name: kernel canonicalization JS/Python parity + run: node e2e_workflow/scripts/test_kernel_canonicalization_parity.js dry-run: name: run_e2e dry-run mapping diff --git a/e2e_workflow/e2e_workflow.js b/e2e_workflow/e2e_workflow.js index 3dcf124ec..117dd015c 100644 --- a/e2e_workflow/e2e_workflow.js +++ b/e2e_workflow/e2e_workflow.js @@ -235,10 +235,18 @@ function isImplausibleSpeedup(pct_gpu_time, isolated, integ) { return ((integ && integ.e2e_delta_pct) || 0) > ceilPct * (1 + IMPLAUSIBLE_SPEEDUP_MARGIN) + 1e-9; } // Classify a reject reason into a fix-and-retry class ('' = terminal, not auto-correctable). +// POSTURE IS TESTED FIRST. CORRECTNESS_REJECT_RX contains a bare `mismatch`, which used to swallow +// the `signature_mismatch` that FIXABLE_REJECT_RX names explicitly: a seam whose signature does not +// match the live call site was routed to the correctness corrective ("your kernel computes the wrong +// thing on the live path, look for a data_ptr over-fit") when what it needs is the integration one +// ("find the method the live server actually dispatches and match its call signature"). The posture +// tokens are specific (signature/seam/engagement/capture), so testing them first cannot capture a +// genuine output-correctness reject; and until the candidate is installed at the right seam, any +// statement about its output is meaningless anyway. function rejectClass(reason) { const r = reason || ''; - if (CORRECTNESS_REJECT_RX.test(r)) return 'correctness'; if (FIXABLE_REJECT_RX.test(r)) return 'integration'; + if (CORRECTNESS_REJECT_RX.test(r)) return 'correctness'; return ''; } // A gate 'accept'/'stack' only counts as a REAL win if the measured e2e delta is not an implausible @@ -486,6 +494,9 @@ const EXTRACT_OP_SCHEMA = obj({ target_callable: { type: 'string' }, // module:attr rebind seam for an authored kernel ('' if none) baseline_callable: { type: 'string' }, // module:attr of the FROZEN real online kernel (the speedup denominator) baseline_frozen: { type: 'boolean' }, // true only when baseline_callable resolves outside the task dir + device_kernel: { type: 'string' }, + seam_candidates: arrObj, + selection_validation: { type: 'object', additionalProperties: true }, smoke: { type: 'string' }, notes: { type: 'string' }, }, ['op_kind', 'task_dir', 'smoke']); @@ -514,6 +525,9 @@ const EXTRACT_SCHEMA = obj({ candidate_bind: { type: 'object', additionalProperties: true }, baseline_overlay: { type: 'string' }, baseline_frozen: { type: 'boolean' }, // true only when baseline_overlay/ was seeded AND candidate_bind is declared + device_kernel: { type: 'string' }, + seam_candidates: arrObj, + selection_validation: { type: 'object', additionalProperties: true }, reference_io_sha256: { type: 'string' }, notes: { type: 'string' }, }, ['editable', 'task_dir', 'unittest_smoke']); @@ -729,6 +743,240 @@ async function ensureFlydslGate() { } } +// A profile identity is a device symbol; the replacement target is a live Python callable. Keep the +// two machine fields separate and require runtime evidence before bake-off or authoring. +const CALLABLE_SPEC_RX = /^[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*:[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*$/; +const validCallableSpec = (s) => CALLABLE_SPEC_RX.test(String(s || '').trim()); +// parse_profile.short_name truncates display names at 60 chars; a declared name that hit the limit +// is our own doing, so a prefix match at it is accepted rather than read as a mismatch. +const SHORT_NAME_LIMIT = 60; +// Balanced groups are removed innermost-first until the symbol stops changing. One greedy pass +// spans from the FIRST opener to the LAST closer, so `k(t)` loses the '(' that marks the end +// of the name, and a nested `k>` leaves the leftover delimiter behind. +const stripBalanced = (symbol, opener, closer) => { + const esc = (c) => c.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + const rx = new RegExp(`${esc(opener)}[^${esc(opener)}${esc(closer)}]*${esc(closer)}`, 'g'); + let text = symbol; + for (let previous = null; previous !== text;) { + previous = text; + text = text.replace(rx, ' '); + } + return text; +}; +// Must stay identical to canonical_kernel_name in kernel_selection.py -- the JS gate and the Python +// verdict compare the same two symbols, so any drift lets a kernel pass one side and fail the other. +// Parentheses are stripped as balanced groups rather than by cutting at the first '(': ROCm spells +// the unnamed namespace `(anonymous namespace)`, which puts a parenthesis BEFORE the kernel, so +// cutting there collapsed every such symbol to the return type `void`, and two unrelated kernels +// that both collapsed to `void` then certified each other. Brackets go the same way, which subsumes +// the `[clone .1]` rule: this pipeline appends its own annotations to a kernel name +// (`_fwd_grouped_kernel_stage1 [sliding_attention]`, `main_kernel[prefill]`), and rocprof spells +// memory ops `Memcpy DtoD (Device -> Device)`. The return type is then dropped by name and the FIRST +// identifier taken, as parse_profile.short_name does; taking the last whitespace token also drops +// `void`, but on any of those annotated spellings it answers with the annotation. +function canonicalDeviceKernel(s) { + const stripped = stripBalanced(stripBalanced( + stripBalanced(String(s || ''), '[', ']'), '<', '>'), '(', ')'); + // Whatever opener is left never closes, so the symbol was cut off inside it -- profile artifacts + // store kernel names elided mid-template. The name is what precedes that opener; reading on + // instead lifts a fragment out of the template arguments and states it with the same confidence + // as a real name, which then matches the wrong kernel rather than refusing. + // Removing a group also leaves a gap where it stood, and an unnamed namespace sits mid- + // qualification: `at::native::(anonymous namespace)::CatArrayBatchedCopy` becomes + // `at::native:: ::CatArray...`, where the identifier probe stops at the space. + const cut = stripped.split(/[<([]/)[0].trim().replace(/\s*::\s*/g, '::'); + const match = /^[\w:]+/.exec(cut.replace(/^void\s+/, '')); + return match ? match[0].split('::').pop().toLowerCase().replace(/[^a-z0-9_]+/g, '') : ''; +} +// A name that hit the display limit ends mid-token, so no word boundary can follow it. Tested both +// ways because either side may be the shortened one: heads carry a short_name while traces carry +// the full symbol, and which is which is not fixed. +const truncatedPrefix = (needle, haystack) => + needle.length >= SHORT_NAME_LIMIT && haystack.startsWith(needle); +// [arguments, closed] for the symbol's first `<...>`. Reading only what follows the `<` keeps this +// independent of the return type and namespace, which one side routinely spells and the other does +// not. `closed` is false when the symbol was cut off inside the template -- profile artifacts elide +// long names, and only the visible part of an elided argument list can be held against anything. +// Separators become `_` rather than vanishing, so `<128, 4, ...>` cannot read as a prefix of +// `<128, 48, ...>`. +function templateArguments(value) { + const text = String(value || ''), start = text.indexOf('<'); + const fold = (s) => s.toLowerCase().replace(/[^a-z0-9]+/g, '_'); + if (start < 0) return ['', true]; + let depth = 0; + for (let i = start; i < text.length; i++) { + if (text[i] === '<') depth++; + else if (text[i] === '>' && --depth === 0) return [fold(text.slice(start + 1, i)), true]; + } + return [fold(text.slice(start + 1)), false]; +} +function kernelIdentitiesMatch(a, b) { + const x = canonicalDeviceKernel(a), y = canonicalDeviceKernel(b); + if (!x || !y) return false; + if (x !== y && !truncatedPrefix(x, y) && !truncatedPrefix(y, x)) return false; + // The base token deliberately drops template arguments so a bare declared name can match its + // decorated spelling. When BOTH sides carry them the information is present on both, and ignoring + // it certifies the wrong kernel: one capture here held 20 distinct kernels named + // at::native::vectorized_elementwise_kernel, separated only by their functor. + const [xa, xc] = templateArguments(a), [ya, yc] = templateArguments(b); + if (!xa || !ya) return true; + if (xc && yc) return xa === ya; + // An elided list still has to agree as far as both sides actually spell it out. + return xa.startsWith(ya) || ya.startsWith(xa); +} +function requiredDeviceKernel(h) { + if (!h || h.entity_kind !== 'gpu_kernel') return ''; + // entity_evidence.matched_profiled_kernel is what parse_profile.py --annotate stamped on the row, + // so it is the profiler's own name for this entity and outranks the display short_name. + return String(h.device_kernel || + (h.entity_evidence && h.entity_evidence.matched_profiled_kernel) || + h.short_name || h.name || '').trim(); +} +const candidateSpec = (c) => + String((c && (c.target_callable || c.callable || c.spec)) || '').trim(); +// A declared depth is a STRUCTURAL claim about the call chain. It orders candidates before any +// runtime evidence exists, so a value that is absent or not a finite number is unusable: `Number()` +// turns it into NaN, every comparison against it is silently false, and a genuinely deeper candidate +// stops being able to outrank anything. Refuse the candidate instead of ranking it as depth 0. +const candidateDepth = (c) => { + const depth = Number(c && c.depth); + return Number.isFinite(depth) ? depth : null; +}; +function selectionCandidatesForHead(h, ext) { + const required = requiredDeviceKernel(h); + const fused = !!(h && (h.is_fused_kernel === true || h.op_kind === 'moe')); + const bySpec = new Map(); + for (const candidate of [ + ...((h && h.seam_candidates) || []), + ...((ext && ext.seam_candidates) || []), + ]) { + const spec = candidateSpec(candidate); + // Architect classifications win for an existing spec; Extractor may append missing inner launchers. + if (spec && !bySpec.has(spec)) bySpec.set(spec, candidate); + } + return Array.from(bySpec.values()).filter((candidate) => { + if (!candidate || typeof candidate !== 'object' || !validCallableSpec(candidateSpec(candidate))) + return false; + const role = String(candidate.role || '').toLowerCase(); + if (role === 'kernel_entry') return false; // native/JIT object identity is unsafe to monkeypatch + if (candidateDepth(candidate) === null) return false; + const kernels = Array.isArray(candidate.device_kernels) ? candidate.device_kernels : []; + if (required && !kernels.some((kernel) => kernelIdentitiesMatch(required, kernel))) return false; + return (fused + ? ['outer_wrapper', 'dispatcher', 'op_seam'] + : ['outer_wrapper', 'dispatcher', 'op_seam', 'inner_launcher']).includes(role); + }); +} +function prepareHeadSelection(h) { + if (!h || typeof h !== 'object') return h; + const out = { ...h }; + const required = requiredDeviceKernel(out); + if (required && !out.device_kernel) out.device_kernel = required; + // live_call_seam and malformed target_callable values are prose, never machine targets. + if (!validCallableSpec(out.target_callable)) out.target_callable = ''; + const fused = out.is_fused_kernel === true || out.op_kind === 'moe'; + const deployable = selectionCandidatesForHead(out).filter((candidate) => { + const role = String(candidate.role || '').toLowerCase(); + return fused ? role === 'op_seam' : ['inner_launcher', 'op_seam'].includes(role); + }); + deployable.sort((a, b) => candidateDepth(b) - candidateDepth(a)); + out.target_callable = deployable.length ? candidateSpec(deployable[0]) : ''; + out.selection_status = out.target_callable + ? 'candidate_selected_needs_runtime_verification' + : 'extractor_must_discover_live_launcher'; + return out; +} +function kernelSelectionVerified(h, ext) { + const required = requiredDeviceKernel(h); + if (!required) return { ok: true, why: 'not a profiled GPU-head extraction' }; + const target = String((ext && ext.target_callable) || '').trim(); + if (!validCallableSpec(target)) + return { ok: false, why: `target_callable '${target}' is not an exact module:attr spec` }; + const verdict = ext && ext.selection_validation; + if (!verdict || verdict.contract !== 'kernel_selection') + return { ok: false, why: 'kernel_selection.py verdict is missing' }; + if (verdict.ok !== true) + return { ok: false, + why: `kernel_selection.py failed: ${(verdict.failed || []).join(', ') || 'unknown'}` }; + if (String(verdict.target_callable || '').trim() !== target) + return { ok: false, + why: `selection verdict certifies '${verdict.target_callable || ''}', not '${target}'` }; + if (!kernelIdentitiesMatch(required, verdict.device_kernel || ext.device_kernel)) + return { ok: false, + why: `selection verdict certifies '${verdict.device_kernel || ext.device_kernel || ''}', ` + + `not profiled kernel '${required}'` }; + if (Number(verdict.total_calls_observed || 0) <= 0 || + Number(verdict.target_marker_calls || 0) <= 0 || + Number(verdict.matched_kernel_calls || 0) <= 0) + return { ok: false, why: 'selection verdict lacks positive call/marker/kernel evidence' }; + const candidates = selectionCandidatesForHead(h, ext); + if (!candidates.length) + return { ok: false, why: 'no relevant structured seam_candidates were returned' }; + const tested = new Set((verdict.candidate_targets_tested || []).map((value) => String(value || '').trim())); + const omitted = candidates.map(candidateSpec).filter((spec) => !tested.has(spec)); + if (omitted.length) + return { ok: false, why: `selection probe omitted candidate(s): ${omitted.join(', ')}` }; + const selected = candidates.find((candidate) => candidateSpec(candidate) === target); + const selectedRole = String((selected && selected.role) || '').toLowerCase(); + const fused = h.is_fused_kernel === true || h.op_kind === 'moe'; + if (!selected || (fused ? selectedRole !== 'op_seam' + : !['inner_launcher', 'op_seam'].includes(selectedRole))) + return { ok: false, + why: fused + ? `fused head must select the whole-operation op_seam, got '${selectedRole || 'unknown'}'` + : `selected target role '${selectedRole || 'unknown'}' is not a deployable inner/op seam` }; + // "Deepest" is settled by the profiler's OBSERVED host-span nesting (kernel_selection.py computes + // deeper_live_candidates from it), not by the declared `depth`. The declared integer is a claim by + // the same agent whose choice is under review, and an appended candidate carrying a large depth + // would otherwise outrank a genuinely deeper one. The structural claim is still checked, but only + // as an ADDITIONAL way to fail: it can never buy a pass that the observed nesting did not grant. + const observedDeeper = (verdict.deeper_live_candidates || []) + .map((value) => String(value || '').trim()).filter(Boolean); + if (observedDeeper.length) + return { ok: false, + why: `deeper live candidate(s) observed across calls/ranks: ${observedDeeper.join(', ')}` }; + const live = new Set((verdict.live_candidate_targets || []).map((value) => String(value || '').trim())); + const selectedDepth = candidateDepth(selected); + const deeper = candidates.filter((candidate) => + candidateSpec(candidate) !== target && live.has(candidateSpec(candidate)) && + candidateDepth(candidate) > selectedDepth); + if (deeper.length) + return { ok: false, + why: `declared-deeper live candidate(s): ${deeper.map(candidateSpec).join(', ')}` }; + if (verdict.deepest_verified !== true) + return { ok: false, why: 'selected callable is not machine-verified as deepest' }; + return { ok: true, why: `${target} launches profiled kernel ${required}` }; +} + +const PRE_FLAGGED_HEADS = []; +function admitHeads(queue, stage) { + const admitted = []; + for (const head of (queue || []).filter(Boolean)) { + const label = head.short_name || head.name || '(unnamed)'; + if (head.entity_kind !== 'gpu_kernel') { + log(` ⚠️ FLAG ${label}: entity_kind=${head.entity_kind || 'missing'}; ` + + `the head track requires a profiler-confirmed gpu_kernel (${stage}).`); + PRE_FLAGGED_HEADS.push({ short_name: label, + pct_gpu_time: head.pct_gpu_time, stage, gate: 'wrong_head_granularity', + reason: `entity_kind=${head.entity_kind || 'missing'}; gpu_kernel required` }); + continue; + } + // A gpu_kernel head with NO device identity at all makes requiredDeviceKernel return '', and + // kernelSelectionVerified then passes it as "not a profiled GPU-head extraction". That vacuous + // pass is the whole contract switched off for this head, so refuse it at admission instead. + if (!requiredDeviceKernel(head)) { + log(` ⚠️ FLAG ${label}: entity_kind=gpu_kernel but no device_kernel/short_name/name; ` + + `there is no profiled symbol to verify a seam against (${stage}).`); + PRE_FLAGGED_HEADS.push({ short_name: label, + pct_gpu_time: head.pct_gpu_time, stage, gate: 'missing_kernel_identity', + reason: 'gpu_kernel head carries no device_kernel/short_name/name' }); + continue; + } + admitted.push(prepareHeadSelection(head)); + } + return admitted; +} + // A FROZEN baseline is resolvable when the extractor seeded baseline_overlay/ + declared // meta.candidate_bind (kernel track), or set an importable meta.baseline_callable (op track). // That is the language-independent speedup denominator. @@ -747,22 +995,51 @@ const hasFrozenBaseline = (ext) => // site (deep, opt-A, milestone/head extract_op, and the non-op milestone extract). async function extractWithBaseline(role, phase, intro, inputs, opts) { const smokeOk = (e) => !!(e && e.task_dir && (e.smoke === 'pass' || e.unittest_smoke === 'pass')); + const head = (inputs && inputs.KERNEL) || {}; let ext = await safeAgent(roleAgent(role, phase, intro, inputs), opts); let tries = 0; - while (smokeOk(ext) && !hasFrozenBaseline(ext) && tries < BASELINE_EXTRACT_RETRIES) { + const attemptedTargets = []; + const complete = (e) => hasFrozenBaseline(e) && kernelSelectionVerified(head, e).ok; + while (smokeOk(ext) && !complete(ext) && tries < BASELINE_EXTRACT_RETRIES) { tries++; - log(` ${(opts && opts.label) || role}: extraction froze NO baseline ` + - `(baseline_overlay/ + meta.candidate_bind) — the speedup denominator would fall back to the ` + - `candidate's own scaffold (fake-win). RE-EXTRACTING (retry ${tries}/${BASELINE_EXTRACT_RETRIES}).`); + 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. ' + + '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. 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}).`); ext = await safeAgent( - roleAgent(role, phase, - intro + ' 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.', - inputs), + roleAgent(role, phase, intro + selectionCorrective + baselineCorrective, { + ...(inputs || {}), + 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).`); @@ -1125,8 +1402,7 @@ if (want('setup')) { // OP-IDENTITY GUARD — a fused-MoE / grouped-expert GEMM must be optimized AS the fused op at its live // dispatcher seam, never decomposed into standalone dense GEMMs (a dense candidate has no live call site // → no_rebind_seam). So force op_kind='moe' (the grouped-GEMM branch; gemmSynthFor keys on this to keep - // dense synth OFF) and preserve the live seam as target_callable, so ANY lever (backend-swap / tune / - // author-fused) binds. The head is never SKIPPED — editability is irrelevant, since a non-editable fused + // dense synth OFF). The head is never SKIPPED — editability is irrelevant, since a non-editable fused // kernel is still backend-swapped at its (editable) dispatcher. GENERIC: detects via the Architect's // is_fused_kernel OR the profile class/name; never keys on a backend name. const _isFusedOp = (c) => (c && c.is_fused_kernel === true) || @@ -1136,9 +1412,9 @@ if (want('setup')) { for (const c of headQueue) { if (!_isFusedOp(c)) continue; c.op_kind = 'moe'; // grouped-GEMM branch (gemmSynthFor → no dense synth) - if (!c.target_callable && c.live_call_seam) c.target_callable = c.live_call_seam; // bind at the live seam _fusedTagged++; } + headQueue = admitHeads(headQueue, 'strategize'); if (_fusedTagged) log(`[op-identity] ${_fusedTagged} fused/grouped head(s): op_kind=moe (never dense-GEMM), bound at live seam — optimized as the fused op, never skipped.`); log(`Strategy: ${headQueue.length} head candidates, ${kernelQueue.length} kernel candidates, ${(strategy && strategy.config_directions || []).length} config directions.`); // strategize decided the backends -> if any candidate routed flydsl, provision it now (blocking). @@ -1155,7 +1431,7 @@ if (want('setup')) { profile = { profile_topN_json: ST.profile_topn_json || '' }; strategy = { config_directions: ST.config_directions || [] }; kernelQueue = ST.kernelQueue || []; - headQueue = ST.headQueue || []; + headQueue = admitHeads(ST.headQueue || [], 'resume'); log(`Loaded carried state: EVAL_DIR=${EVAL_DIR}, baseline ${BASELINE_TPUT}, flags='${curFlags}', env='${curEnv}', ${headQueue.length} head + ${kernelQueue.length} kernel candidates.`); } @@ -1194,7 +1470,8 @@ if (want('config') && CONFIG_TUNE_ENABLED && strategy && (strategy.config_direct }), { phase: 'Strategize', label: 'architect:re-strategize', schema: STRATEGY_SCHEMA }); if (restrat && restrat.kernel_candidates) kernelQueue = restrat.kernel_candidates.slice(); - if (restrat && restrat.head_candidates) headQueue = restrat.head_candidates.slice(); + if (restrat && restrat.head_candidates) + headQueue = admitHeads(restrat.head_candidates.slice(), 're-strategize'); // re-strategize may have (re)routed flydsl -> provision it (idempotent; no-op if already done). await ensureFlydslGate(); } else { @@ -1217,7 +1494,7 @@ const acceptedHeads = (ST.accepted_heads || []).slice(); // so Finalize can finish the best one's A/B (Fix C) and so a real isolated win // is surfaced (return.pending_integrations) instead of being silently dropped. const pendingIntegrations = (ST.pending_integrations || []).slice(); -const flaggedHeads = (ST.flagged_heads || []).slice(); // dominant heads that could NOT be optimized (loudly surfaced, never silently skipped) +const flaggedHeads = (ST.flagged_heads || []).concat(PRE_FLAGGED_HEADS); // heads that could NOT be optimized (loudly surfaced, never silently skipped) let headDispatched = 0; const history = ST.history || { insights: [], ledger: [], milestones: [], bottleneck_now: '', suggest_next: '' }; diff --git a/e2e_workflow/roles/kernel_extractor.md b/e2e_workflow/roles/kernel_extractor.md index 7f16edf4b..160c40900 100644 --- a/e2e_workflow/roles/kernel_extractor.md +++ b/e2e_workflow/roles/kernel_extractor.md @@ -122,26 +122,35 @@ If the live regime genuinely cannot be reproduced offline (op only exists fused routing-dependent MoE token counts), say so in `notes` and report `editable:false`/drop rather than freeze an out-of-regime oracle nobody should trust. -1. **Locate the source.** **If `KERNEL.source_hint`/`KERNEL.launcher_hint` is provided (TraceLens +1. **Locate the source and select the live launcher.** `KERNEL.device_kernel` is the profiled GPU + symbol this extraction must reach. `KERNEL.target_callable` is only a hint, and + `KERNEL.live_call_seam` is prose that must never become a machine target. Start from + `KERNEL.seam_candidates[]`. If source/runtime inspection finds a missing inner launcher, append it + to the returned `seam_candidates[]` with its exact role, depth, matching device kernels, and evidence. + **If `KERNEL.source_hint`/`KERNEL.launcher_hint` is provided (TraceLens pre-resolved the file/seam), look there FIRST** — but always CONFIRM by importing the package + grepping the `short_name`/`module:attr` target; never trust the hint blindly (it may point at a launcher/wrapper rather than the true defining file). If no hint, resolve as usual (`python3 -c "import sglang,os;print(os.path.dirname(sglang.__file__))"`, then grep the `short_name` / the `module:attr` target). **OP-IDENTITY IS THE RULE: extract the op the LIVE kernel actually is, at the seam it is actually called - from — never a different op.** Two cases: + from — never a different op.** Three cases: - **Standalone LIBRARY op** (a discrete hipBLASLt/rocBLAS `gemm(...)` / library attention whose only call site is that library call, no editable body) → STOP, report `editable=false`, `target_callable=""`; it belongs to the config/tune-hook track (per-shape DB tune / backend env), not a source rewrite. Do NOT synthesize a standalone-GEMM proxy just to make it look extractable. - **FUSED / monolithic op** (fused-MoE, grouped-expert GEMM, asm/CK fused kernel — `KERNEL` arrives with `op_kind=moe` and `GEMM_SYNTH=false`): **extract the FUSED op** (capture its live I/O oracle), NOT its - constituent standalone GEMMs. Set `target_callable` to the **dispatcher** actually called at runtime — - use `KERNEL.target_callable`/`KERNEL.live_call_seam` if provided (e.g. the vLLM `fused_moe`/ - `fused_experts` dispatcher), which is editable Python EVEN WHEN the underlying kernel is a non-editable + constituent standalone GEMMs. Select the bindable whole-operation **`op_seam`** from + `KERNEL.seam_candidates` (or an exact `KERNEL.target_callable` hint), which is editable Python EVEN + WHEN the underlying kernel is a non-editable library/asm `.so`. That dispatcher seam is what lets a fused op be BACKEND-SWAPPED (aiter/flydsl/triton fused) or AUTHOR-fused-replaced regardless of the underlying kernel's editability. Report `editable=true` (the seam is rebindable). NEVER decompose it into a dense A·Bᵀ GEMM — no live call site. + - **OUTER WRAPPER / dispatcher** → do not stop after rejecting it. Descend to the deepest safe Python + `inner_launcher` or `op_seam` that launches `KERNEL.device_kernel`. A native or Triton + `kernel_entry` remains source evidence, not a monkeypatch target. If no safe callable can be found, + report `editable=false`; never claim selection success from rejection alone. 2. **Capture shapes + oracle** from a live server using `scripts/capture_shapes.py` via a temporary capture overlay, driven by the SAME workload as the profile so shapes match the regime: ```bash @@ -152,17 +161,37 @@ freeze an out-of-regime oracle nobody should trust. # PRISTINE install instead of the server the accepted kernels actually built. python3 "$SKILL_DIR/scripts/overlay_setup.py" add-capture \ --overlay "$TASK/_capture_overlay" --from "$CURRENT_OVERLAY" \ - --target "" --out "$TASK" --max 5 \ + --target "" --out "$TASK" --max 5 \ --capture-file "$SKILL_DIR/scripts/capture_shapes.py" + # Repeat add-marker for every relevant safe Python candidate, deepest first. The shim always + # installs captures BEFORE markers, so a marker can never freeze an alias of a function that the + # capture hook had not wrapped yet. + python3 "$SKILL_DIR/scripts/overlay_setup.py" add-marker \ + --overlay "$TASK/_capture_overlay" --target "" \ + --marker-file "$SKILL_DIR/scripts/seam_trace.py" cp -r "$CURRENT_OVERLAY"/. "$TASK/baseline_overlay"/ 2>/dev/null || \ python3 -c "import sys;sys.path.insert(0,'$SKILL_DIR/scripts');import overlay_setup as o;o._ensure_overlay('$TASK/baseline_overlay')" BACKEND="" OUT_DIR="$TASK/_capture" GPU="$GPU_ID" MODEL="$MODEL_PATH" \ ISL= OSL= CONC= REPEATS=0 PROFILE=0 \ OVERLAY_PYTHONPATH="$TASK/_capture_overlay" \ - EXTRA_ENV="CAPTURE_TARGET= CAPTURE_OUT=$TASK CAPTURE_MAX=5" \ + EXTRA_ENV="CAPTURE_TARGET= CAPTURE_OUT=$TASK CAPTURE_MAX=5 GEAK_SELECTION_TRACE=$TASK/selection_trace.json" \ bash "$EVAL_DIR/bench_e2e.sh" 2>&1 | tee "$EVAL_DIR/logs/capture_.log" + python3 "$SKILL_DIR/scripts/kernel_selection.py" \ + --target "" --device-kernel "" \ + --capture-meta "$TASK"/capture.pid-*.rank-*/meta.json \ + --torch-trace "$TASK"/selection_trace.pid-*.rank-*.call-*.json \ + --candidate-target "" \ + --out "$TASK/selection_validation.json" ``` + Repeat `--candidate-target` for every relevant candidate. `seam_trace` writes an atomic trace for + each PID/rank/root call, and `capture_shapes` writes atomic process-local artifacts. The verifier + merges all calls for each PID, isolates External ids between trace files, and requires every capture + PID to pass. Installation proof is distinct from execution: a marked mutually exclusive branch may + stay inactive, while a selected outer target fails if a deeper marked candidate launches the same + device kernel in any call/rank. Require `deepest_verified:true`, then copy the selected process's + `meta.json` and `reference_io.pt` into the task root. + 🔴 **Capture on the CURRENT stack, not the install.** The oracle you freeze is the truth source the candidate is judged against, and the baseline you time against is `baseline_overlay/`. Both must be the server as it runs RIGHT NOW (config + every accepted kernel). Capturing on the pristine install @@ -451,7 +480,14 @@ Return JSON: "editable": true, "task_dir": "/kernels/_task", "source_path_in_sglang": "", + "device_kernel": "", "target_callable": "", + "seam_candidates": [ + {"target_callable": "", + "role": "outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry", + "device_kernels": [""], "depth": 0, "evidence": "source/runtime evidence"} + ], + "selection_validation": {"contract": "kernel_selection", "ok": true, "deepest_verified": true}, "candidate_bind": {"kind": "module|rebind", "module": "", "file": "kernel_src/.py"}, "baseline_overlay": "/baseline_overlay", "baseline_frozen": true, @@ -807,7 +843,14 @@ Return JSON: "regimes_captured": ["prefill"], "candidate_backends": ["aiter","hipblaslt","triton","ck"], "reference_io_sha256": "", + "device_kernel": "", "target_callable": "", + "seam_candidates": [ + {"target_callable": "", + "role": "outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry", + "device_kernels": [""], "depth": 0, "evidence": "source/runtime evidence"} + ], + "selection_validation": {"contract": "kernel_selection", "ok": true, "deepest_verified": true}, "baseline_callable": "", "smoke": "pass|fail", "notes": "transpose/bias inference, regime, whether oracle was synthesized vs captured" diff --git a/e2e_workflow/roles/profiler.md b/e2e_workflow/roles/profiler.md index ba1f2a887..364282077 100644 --- a/e2e_workflow/roles/profiler.md +++ b/e2e_workflow/roles/profiler.md @@ -100,8 +100,21 @@ An upstream orchestrator may already have profiled the SAME baseline workload wi `
` args, `classification`←map from `kernel_category`/`bound_type` (MoE/grouped-GEMM→library_gemm or triton per `kernel_kind`; attention→library_attn; etc.), `editable`←`op_to_source_patchable`. Carry `source_file`/`kernel_path` into each entry's `notes` (the Architect/Extractor reuse them). Write - `profile_topN.json` + `.md` via your own Write (you may shell out to `parse_profile.py` only if you - also have a trace; otherwise assemble the JSON yourself) and set `source:"tracelens"`. + `profile_topN.json` + `.md` via your own Write and set `source:"tracelens"`. + **Then annotate the assembled rows from profiler evidence; never hand-write `entity_kind`:** + ```bash + TLT=$(ls -1 "$TRACELENS_TRACE_FILE"/*rank0*.pt.trace.json.gz 2>/dev/null | head -1) + [ -z "$TLT" ] && TLT=$(ls -1 "$TRACELENS_TRACE_FILE"/*.pt.trace.json.gz \ + "$TRACELENS_TRACE_FILE"/*.json.gz "$TRACELENS_TRACE_FILE"/*.json 2>/dev/null | head -1) + python3 "$EVAL_DIR/parse_profile.py" \ + --annotate "$EVAL_DIR/profile/round_${ROUND}/profile_topN.json" \ + --torch-trace "$TLT" \ + --annotate-out "$EVAL_DIR/profile/round_${ROUND}/profile_topN.json" + ``` + If no trace is available, fall back to the normal collection below instead of guessing a row's + entity kind. Annotation expands an outer dispatcher/custom-op row through torch-profiler External-id + edges into its concrete device children. Preserve the resulting `device_kernel`, `profile_parent`, + and split GPU percentages: rejecting a dispatcher without discovering its children is not success. - **If `TRACELENS_TRACE_FILE` is also a non-empty path that EXISTS → run an ADDITIONAL trace-analysis pass on top of analysis.md to sharpen the picture** (this is required by contract when the trace is present). `TRACELENS_TRACE_FILE` is a `torch_trace` **directory** that holds one steady-state serving diff --git a/e2e_workflow/roles/system_architect.md b/e2e_workflow/roles/system_architect.md index 8eb429597..d1ee3df50 100644 --- a/e2e_workflow/roles/system_architect.md +++ b/e2e_workflow/roles/system_architect.md @@ -244,6 +244,20 @@ OPTIONAL upstream TraceLens prior (may be empty strings — treat empty/missing e2e by MORE than the noise band. Otherwise drop it — say so. 5. Write `EVAL_DIR/strategy.md` (human-readable plan) and return the routing. +> **Every head candidate must identify a device kernel and a structured callable chain.** +> - Copy `entity_kind` and `device_kernel` from the profiled row. Route only `gpu_kernel` rows; if a +> dispatcher was expanded, route its device children rather than recreating the outer aggregate. +> - `live_call_seam` is prose context only. Never copy arrows, signatures, paths, or prose into +> `target_callable`. +> - Build `seam_candidates[]` from source and the baseline server log. Each entry has an exact importable +> `target_callable` (`module:attr`), `role` (`outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry`), +> matching `device_kernels`, `depth`, and evidence. Include every plausible callable on the live path. +> - Keep native/JIT `kernel_entry` objects as source evidence only; replacing them can break their +> `.run`, `.warmup`, or cache protocols. For non-fused heads prefer the deepest safe +> `inner_launcher`/`op_seam`; fused heads must select the whole-operation `op_seam`. +> - `target_callable` is only an initial hint. The Extractor may add a missing inner launcher and must +> prove the final choice with runtime markers; merely rejecting an outer wrapper is not discovery. + Return JSON: ```json { @@ -256,6 +270,15 @@ Return JSON: "head_candidates": [ {"id": "h0", "short_name": "...", "op_kind": "gemm|attn", "pct_gpu_time": 0.0, "shapes": "[[1024,5120],[5120,34816]]", "dtype": "bf16", "regime": "prefill|decode|both", + "entity_kind": "gpu_kernel", + "device_kernel": "", + "target_callable": "", + "seam_candidates": [ + {"target_callable": "", + "role": "outer_wrapper|dispatcher|op_seam|inner_launcher|kernel_entry", + "device_kernels": [""], "depth": 0, + "runtime_verified": false, "evidence": "source/log evidence"} + ], "transpose_b": true, "bias": false, "candidate_backends": ["aiter","hipblaslt","triton","ck"], "is_fused_kernel": false, diff --git a/e2e_workflow/scripts/capture_shapes.py b/e2e_workflow/scripts/capture_shapes.py index 5464de7ea..9261f7f1d 100755 --- a/e2e_workflow/scripts/capture_shapes.py +++ b/e2e_workflow/scripts/capture_shapes.py @@ -54,6 +54,22 @@ } +def _rank(): + return next((str(os.environ[key]) for key in + ("RANK", "LOCAL_RANK", "TP_RANK", "SLURM_PROCID") + if key in os.environ), "unknown") + + +def _process_out_dir(out_dir): + """Isolate selection-capture artifacts so TP workers cannot corrupt each other.""" + unique = (os.environ.get("CAPTURE_PROCESS_UNIQUE") == "1" + or bool(os.environ.get("GEAK_SELECTION_TRACE"))) + if not unique: + return out_dir + return os.path.join( + out_dir, f"capture.pid-{os.getpid()}.rank-{_rank()}") + + def _shapes_dtypes(args, kwargs): """Light shape/dtype walk (no clone) so we can catalog EVERY distinct shape cheaply, independent of the memory-bounded oracle capture.""" @@ -153,7 +169,20 @@ def _sig(args, kwargs): def _wrapper(*args, **kwargs): s = _STATE - out = s["orig"](*args, **kwargs) + # This marker is consumed by kernel_selection.py from the capture run's torch trace. It turns + # "the hook saw calls" into stronger evidence: the GPU kernel selected from the baseline profile + # must actually execute while THIS callable is active. record_function is effectively a no-op + # when no profiler is collecting, so existing non-profiled capture users keep the same behavior. + marker = f"GEAK_TARGET::{s['target']}" + try: + record_function = _torch().profiler.record_function + except Exception: + record_function = None + if record_function: + with record_function(marker): + out = s["orig"](*args, **kwargs) + else: + out = s["orig"](*args, **kwargs) s["calls"] += 1 in_graph = _capturing() try: @@ -239,7 +268,9 @@ def _flush(write_oracle=True): # workload capture (< max_cases distinct shapes) and a late regime-coverage case (appended past # max_cases) land on disk; records is bounded, so this rewrites only a handful of times. if write_oracle and records and len(records) > s["oracle_records"]: - torch.save({"target": s["target"], "records": records}, io_path) + tmp_io = f"{io_path}.tmp-{os.getpid()}-{threading.get_ident()}" + torch.save({"target": s["target"], "records": records}, tmp_io) + os.replace(tmp_io, io_path) import hashlib h = hashlib.sha256() with open(io_path, "rb") as fh: @@ -274,6 +305,8 @@ def walk(o): # not just single-shape h.check_correct_multi. meta = { "target": s["target"], + "process_id": os.getpid(), + "rank": _rank(), "module": s["mod"].__name__ if s["mod"] else None, "attr": s["attr"], "num_cases": len(records), @@ -291,8 +324,11 @@ def walk(o): "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.", } - with open(os.path.join(out_dir, "meta.json"), "w") as fh: + meta_path = os.path.join(out_dir, "meta.json") + tmp_meta = f"{meta_path}.tmp-{os.getpid()}-{threading.get_ident()}" + with open(tmp_meta, "w") as fh: json.dump(meta, fh, indent=2) + os.replace(tmp_meta, meta_path) sys.stderr.write(f"[capture_shapes] flushed {len(records)} case(s) " f"(regimes={sorted(s['regime_seen'])}), " f"oracle_complete={s['oracle_written']} -> {out_dir}\n") @@ -365,8 +401,15 @@ def install(target, out_dir, max_cases=5): 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.") + out_dir = _process_out_dir(out_dir) s.update(target=target, out_dir=out_dir, max_cases=int(max_cases), orig=orig, mod=mod, attr=attr, installed=True) + if os.environ.get("GEAK_SELECTION_TRACE"): + # Selection runs are intentionally short and server teardown may use SIGTERM, which skips + # Python atexit. Persist the first eager capture immediately. + s["flush_every"] = 1 + elif os.environ.get("CAPTURE_FLUSH_EVERY"): + s["flush_every"] = max(1, int(os.environ["CAPTURE_FLUSH_EVERY"])) setattr(mod, attr, _make_wrapper(orig)) atexit.register(_flush) sys.stderr.write(f"[capture_shapes] hooked {target}; recording up to {max_cases} cases -> {out_dir}\n") diff --git a/e2e_workflow/scripts/kernel_selection.py b/e2e_workflow/scripts/kernel_selection.py new file mode 100644 index 000000000..89a6fa2d6 --- /dev/null +++ b/e2e_workflow/scripts/kernel_selection.py @@ -0,0 +1,529 @@ +#!/usr/bin/env python3 +"""Verify that an extracted callable is the live seam for a profiled GPU kernel. + +The extractor marks every structured candidate with ``seam_trace`` and captures +the selected callable with ``capture_shapes``. ``seam_trace`` profiles a bounded +warmup/call window and emits install proof plus nested live-call markers. This tool +verifies, from machine evidence, that: + +* the target is an exact import spec rather than a prose call chain; +* the target was called by the live server; +* the expected profiled GPU kernel ran while the target marker was active. +* no deeper marked live candidate launches the same kernel. + +The JSON verdict is consumed by the workflow. A rejected outer wrapper is not +selection success; the extractor must descend and repeat this probe until both +this contract and the binding contract pass. +""" + +import argparse +import gzip +import json +import os +import re +import sys + + +MARKER_PREFIX = "GEAK_TARGET::" +INSTALL_PREFIX = "GEAK_INSTALLED::" +_SPEC_RE = re.compile( + r"^[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*:" + r"[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*$" +) + + +def valid_callable_spec(value): + return bool(_SPEC_RE.fullmatch(str(value or "").strip())) + + +def _open(path): + return gzip.open(path, "rt") if str(path).endswith(".gz") else open(path, "rt") + + +def load_events(path): + with _open(path) as fh: + doc = json.load(fh) + return doc.get("traceEvents", doc if isinstance(doc, list) else []) + + +def merge_process_traces(paths): + """Merge call traces without allowing pid or External-id collisions across files.""" + merged = [] + for trace_index, path in enumerate(paths): + prefix = f"trace-{trace_index}:" + for original in load_events(path): + if not isinstance(original, dict): + continue + event = dict(original) + event["pid"] = prefix + str(original.get("pid")) + event["tid"] = prefix + str(original.get("tid")) + args = dict(original.get("args") or {}) + if args.get("External id") is not None: + args["External id"] = prefix + str(args["External id"]) + if args.get("correlation") is not None: + args["correlation"] = prefix + str(args["correlation"]) + event["args"] = args + merged.append(event) + return merged + + +# parse_profile.short_name truncates display names here. A declared name that hit the limit is our +# own doing, so the matcher accepts a prefix at it rather than reading the truncation as a mismatch. +SHORT_NAME_LIMIT = 60 + + +def _strip_balanced(text, opener, closer): + """Remove balanced ``opener...closer`` groups innermost-first until the symbol stops changing. + + A single greedy pass spans from the FIRST opener to the LAST closer, so ``k(t)`` loses the + ``(`` that ends the name and a nested ``k>`` leaves a stray delimiter behind. Removing + the innermost group repeatedly is the only way to survive both. + """ + pattern = re.compile( + re.escape(opener) + "[^" + re.escape(opener + closer) + "]*" + re.escape(closer)) + previous = None + while previous != text: + previous = text + text = pattern.sub(" ", text) + return text + + +def canonical_kernel_name(value): + """Return a stable token for matching mangled/demangled kernel symbols. + + This must stay identical to ``canonicalDeviceKernel`` in e2e_workflow.js: the JS gate and this + verdict compare the same two symbols, and a canonicalization that drifted between them would let + a kernel pass one side and be refused by the other. + + Parentheses are stripped as balanced groups rather than by cutting at the first ``(``. ROCm names + the unnamed namespace ``(anonymous namespace)``, which puts a parenthesis *before* the kernel, so + cutting there collapsed every such symbol to the return type ``void`` -- and two unrelated + kernels that both collapse to ``void`` then certified each other. + + Brackets go the same way, which subsumes the ``[clone .1]`` rule: this pipeline appends its own + annotations to a kernel name (``_fwd_grouped_kernel_stage1 [sliding_attention]``, + ``main_kernel[prefill]``, ``hgemm_..._SPK1 [qkv_proj]``), and rocprof spells memory ops + ``Memcpy DtoD (Device -> Device)``. + + The return type is then dropped by name and the FIRST identifier is taken, exactly as + ``parse_profile.short_name`` does. Taking the last whitespace token instead also drops ``void``, + but on any of the annotated spellings above it answers with the annotation -- ``DtoD``, + ``sliding_attention``, ``qkv_proj`` -- and a head then fails to match its own kernel. + """ + text = _strip_balanced(str(value or ""), "[", "]") + text = _strip_balanced(text, "<", ">") + text = _strip_balanced(text, "(", ")") + # Whatever opener is left never closes, so the symbol was cut off inside it -- profile artifacts + # store kernel names elided mid-template. The name is what precedes that opener; reading on + # instead lifts a fragment out of the template arguments and states it with the same confidence + # as a real name ('...elementwise_kernel_manual_unroll<128, 4, at::native::gpu_k' answered + # 'gpu_k'), which then matches the wrong kernel rather than refusing to answer. + text = re.split(r"[<(\[]", text, maxsplit=1)[0] + # Removing a group leaves a gap where it stood, and an unnamed namespace sits mid-qualification: + # `at::native::(anonymous namespace)::CatArrayBatchedCopy` becomes `at::native:: ::CatArray...`, + # where the identifier probe below stops at the space and the last `::` segment is empty. + text = re.sub(r"\s*::\s*", "::", text.strip()) + match = re.match(r"[\w:]+", re.sub(r"^void\s+", "", text)) + return re.sub(r"[^a-z0-9_]+", "", match.group(0).split("::")[-1].lower()) if match else "" + + +def _truncated_prefix(needle, haystack): + """True when ``needle`` is ``haystack`` cut at the display limit. + + A name that hit the limit ends mid-token, so no word boundary can follow it. Tested both ways + because either side may be the shortened one: heads carry a short_name while traces carry the + full symbol, and which is which is not fixed. + """ + return len(needle) >= SHORT_NAME_LIMIT and haystack.startswith(needle) + + +def _template_arguments(value): + """``(arguments, closed)`` for the symbol's first ``<...>``. + + Reading only what follows the ``<`` keeps this independent of the return type and namespace, + which one side routinely spells and the other does not. ``closed`` is False when the symbol was + cut off inside the template -- profile artifacts elide long names, and only the visible part of + an elided argument list can be held against anything. Separators become ``_`` rather than + vanishing, so ``<128, 4, ...>`` cannot read as a prefix of ``<128, 48, ...>``. + """ + text = str(value or "") + start = text.find("<") + if start < 0: + return "", True + depth = 0 + for index in range(start, len(text)): + if text[index] == "<": + depth += 1 + elif text[index] == ">": + depth -= 1 + if depth == 0: + return re.sub(r"[^a-z0-9]+", "_", text[start + 1:index].lower()), True + return re.sub(r"[^a-z0-9]+", "_", text[start + 1:].lower()), False + + +def kernel_matches(expected, observed): + want = canonical_kernel_name(expected) + got = canonical_kernel_name(observed) + if not want or not got: + return False + if want != got and not _truncated_prefix(want, got) and not _truncated_prefix(got, want): + return False + # The base token deliberately drops template arguments so a bare declared name can match its + # decorated spelling. When BOTH sides carry them the information is present on both, and ignoring + # it certifies the wrong kernel: one capture here held 20 distinct kernels named + # at::native::vectorized_elementwise_kernel, separated only by their functor. + want_args, want_closed = _template_arguments(expected) + got_args, got_closed = _template_arguments(observed) + if not want_args or not got_args: + return True + if want_closed and got_closed: + return want_args == got_args + # An elided list still has to agree as far as both sides actually spell it out. + return want_args.startswith(got_args) or got_args.startswith(want_args) + + +def _device_projection(event): + """True for the GPU-timeline copy of a host annotation (``gpu_user_annotation``). + + The profiler re-emits every ``record_function`` marker on the device timeline, where the span + covers the kernels it launched rather than the Python call. Those projections do not preserve + host call nesting -- a short device span for an OUTER seam routinely lands inside the device span + of an INNER one -- so only host-side spans may establish nesting or launch causality. + """ + return str(event.get("cat") or "").startswith("gpu_") + + +def _complete_spans(events, name): + spans = [] + stacks = {} + for event in events: + if not isinstance(event, dict) or event.get("name") != name: + continue + if _device_projection(event): + continue + phase = event.get("ph", "X") + if phase == "X" and event.get("ts") is not None: + start = float(event["ts"]) + spans.append((start, start + float(event.get("dur") or 0), event)) + elif phase == "B": + stacks.setdefault((event.get("pid"), event.get("tid")), []).append(event) + elif phase == "E": + key = (event.get("pid"), event.get("tid")) + if stacks.get(key): + begin = stacks[key].pop() + start = float(begin.get("ts") or 0) + spans.append((start, float(event.get("ts") or start), begin)) + return spans + + +def _outermost_spans(spans): + """Return one logical call for same-thread spans nested inside an identical marker. + + ``capture_shapes`` and ``seam_trace`` can both wrap the selected callable with + ``GEAK_TARGET::``. One real invocation then produces the same marker nested inside itself. + Keep every raw span for launch-causality analysis, but count only the outermost copy when reporting + calls. Different threads and non-nested calls remain distinct. + """ + outermost = [] + for index, span in sorted( + enumerate(spans), + key=lambda item: (item[1][0], -item[1][1], item[0])): + start, end, event = span + contained = any( + outer_start <= start and end <= outer_end + and event.get("pid") == outer_event.get("pid") + and event.get("tid") == outer_event.get("tid") + for outer_start, outer_end, outer_event in outermost + ) + if not contained: + outermost.append(span) + return outermost + + +def _within_any_span(event, spans, same_thread=False): + if event.get("ts") is None: + return False + ts = float(event["ts"]) + return any( + start <= ts <= end + and (not same_thread or ( + event.get("pid") == marker.get("pid") and event.get("tid") == marker.get("tid"))) + for start, end, marker in spans + ) + + +def _marker_kernel_evidence(trace_events, target_callable, device_kernel): + marker = MARKER_PREFIX + target_callable + spans = _complete_spans(trace_events or [], marker) + # Only CPU events on the marker's own thread can establish launch causality. Global timestamp + # overlap is unsafe under concurrent serving. The resulting External ids bridge to async GPU events. + related_external_ids = set() + # Triton (and any raw hipModuleLaunchKernel) device rows carry only `correlation`, never an + # `External id`, so the External-id bridge alone silently misses them. The host-side launch + # runtime event inside the marker span carries BOTH ids, so its `correlation` is an equally + # strict launch-causality bridge: same thread, inside the marker span, same launch. + related_correlations = set() + for event in trace_events or []: + if not isinstance(event, dict) or not _within_any_span(event, spans, same_thread=True): + continue + if event.get("cat") == "kernel": + continue + args = event.get("args") or {} + ext = args.get("External id") + if ext is not None: + related_external_ids.add(ext) + corr = args.get("correlation") + if corr is not None: + related_correlations.add(corr) + matched = [] + for event in trace_events or []: + if not isinstance(event, dict) or event.get("cat") != "kernel": + continue + args = event.get("args") or {} + ext = args.get("External id") + corr = args.get("correlation") + linked = (ext is not None and ext in related_external_ids) or ( + corr is not None and corr in related_correlations) + if linked and kernel_matches(device_kernel, event.get("name")): + matched.append(str(event.get("name") or "")) + return { + "target": target_callable, + "marker": marker, + "spans": spans, + "marker_calls": len(_outermost_spans(spans)), + "related_external_ids": related_external_ids, + "related_correlations": related_correlations, + "matched": matched, + } + + +def _is_nested(inner_spans, outer_spans): + for inner_start, inner_end, inner_event in inner_spans: + for outer_start, outer_end, outer_event in outer_spans: + same_thread = ( + inner_event.get("pid") == outer_event.get("pid") + and inner_event.get("tid") == outer_event.get("tid") + ) + if same_thread and outer_start <= inner_start and inner_end <= outer_end and ( + outer_start < inner_start or inner_end < outer_end): + return True + return False + + +def verify(target_callable, device_kernel, capture_meta, trace_events, candidate_targets=None): + target_callable = str(target_callable or "").strip() + device_kernel = str(device_kernel or "").strip() + failed = [] + + if not valid_callable_spec(target_callable): + failed.append("invalid_target_callable") + if not device_kernel: + failed.append("missing_device_kernel") + + meta_target = "" + observed_calls = 0 + if isinstance(capture_meta, dict): + module = str(capture_meta.get("module") or "").strip() + attr = str(capture_meta.get("attr") or "").strip() + meta_target = f"{module}:{attr}" if module and attr else "" + observed_calls = int(capture_meta.get("total_calls_observed") or 0) + if meta_target != target_callable: + failed.append("capture_target_mismatch") + if observed_calls <= 0: + failed.append("target_not_observed") + + candidates = [] + for candidate in list(candidate_targets or []) + [target_callable]: + candidate = str(candidate or "").strip() + if candidate and candidate not in candidates: + candidates.append(candidate) + invalid_candidates = [candidate for candidate in candidates if not valid_callable_spec(candidate)] + if invalid_candidates: + failed.append("invalid_candidate_target") + evidence = { + candidate: _marker_kernel_evidence(trace_events, candidate, device_kernel) + for candidate in candidates if valid_callable_spec(candidate) + } + installed_candidates = [ + candidate for candidate in candidates + if _complete_spans(trace_events, INSTALL_PREFIX + candidate) + ] + missing_candidate_markers = sorted(set(candidates) - set(installed_candidates)) + if missing_candidate_markers: + failed.append("candidate_marker_not_installed") + selected_evidence = evidence.get(target_callable) or { + "marker": MARKER_PREFIX + target_callable, "spans": [], "marker_calls": 0, + "related_external_ids": set(), "related_correlations": set(), "matched": [], + } + marker = selected_evidence["marker"] + spans = selected_evidence["spans"] + if not spans: + failed.append("target_marker_missing") + + related_external_ids = selected_evidence["related_external_ids"] + matched = selected_evidence["matched"] + if not matched: + failed.append("device_kernel_not_under_target") + deeper = [ + candidate for candidate, candidate_evidence in evidence.items() + if candidate != target_callable and candidate_evidence["matched"] + and _is_nested(candidate_evidence["spans"], spans) + ] + if deeper: + failed.append("deeper_live_candidate_exists") + # A candidate whose marker was installed and never fired is NOT evidence that the callable is + # off the live path. Installation rebinds one module attribute; a callable reached through the + # torch dispatcher (`torch.ops.*`) or through an alias a caller imported before installation + # still runs with the wrapper bypassed, and is therefore invisible here. `deeper` can only see + # what the markers saw, so report the unobserved candidates separately instead of letting them + # be read as "not deeper": whoever holds the declared depths decides what an unobserved + # declared-deeper candidate means. + unobserved_candidates = sorted( + candidate for candidate in installed_candidates + if candidate != target_callable and not (evidence.get(candidate) or {}).get("spans")) + deepest_verified = bool( + spans and matched and not deeper and not invalid_candidates + and not missing_candidate_markers) + + return { + "contract": "kernel_selection", + "ok": not failed, + "target_callable": target_callable, + "device_kernel": device_kernel, + "capture_target": meta_target, + "total_calls_observed": observed_calls, + "target_marker": marker, + "target_marker_calls": selected_evidence["marker_calls"], + "matched_kernel_calls": len(matched), + "matched_kernel_names": sorted(set(matched)), + "correlated_external_ids": len(related_external_ids), + "correlated_launch_correlations": len(selected_evidence.get("related_correlations") or ()), + "candidate_targets_tested": installed_candidates, + "live_candidate_targets": sorted( + candidate for candidate, candidate_evidence in evidence.items() + if candidate_evidence["matched"]), + "deeper_live_candidates": sorted(deeper), + "missing_candidate_markers": missing_candidate_markers, + "installed_but_never_live_candidates": unobserved_candidates, + "deepest_verified": deepest_verified, + "evidence": "installed+live_nested_candidate_markers+torch_profiler_external_id_or_launch_correlation", + "failed": failed, + } + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--target", required=True, help="exact module:attr selected seam") + parser.add_argument("--device-kernel", required=True, help="GPU kernel selected from profile") + parser.add_argument("--capture-meta", required=True, nargs="+", + help="one or more process-local capture_shapes meta.json files") + parser.add_argument("--torch-trace", required=True, nargs="+", + help="one or more process-local capture traces (.json or .json.gz)") + parser.add_argument("--candidate-target", action="append", default=[], + help="exact module:attr candidate marked in the same trace; repeat for all candidates") + parser.add_argument("--out", default="", help="optional verdict JSON path") + args = parser.parse_args(argv) + + trace_paths = list(args.torch_trace) + attempts = [] + for meta_path in args.capture_meta: + with open(meta_path) as fh: + meta = json.load(fh) + capture_pid = str(meta.get("process_id") or "") + matching_traces = trace_paths + if capture_pid: + matching_traces = [ + path for path in trace_paths + if re.search(rf"\.pid-{re.escape(capture_pid)}(?:\.|$)", + os.path.basename(path)) + ] + if not matching_traces: + verdict = verify( + args.target, args.device_kernel, meta, [], args.candidate_target) + verdict["failed"].append("capture_process_trace_missing") + verdict["ok"] = False + verdict["deepest_verified"] = False + attempts.append((meta_path, [], verdict)) + continue + attempts.append(( + meta_path, + matching_traces, + verify(args.target, args.device_kernel, meta, + merge_process_traces(matching_traces), args.candidate_target), + )) + # The strongest single process seeds the verdict shape, but every field that survives below is + # recomputed across ALL attempts. Only the meta path stays process-specific, so it is named for + # the one process it describes rather than presented as the run's selection. + best_process_meta_file, _, verdict = max( + attempts, + key=lambda item: ( + item[2]["matched_kernel_calls"], + item[2]["target_marker_calls"], + -len(item[2]["failed"]), + ), + ) + verdict = dict(verdict) + all_verdicts = [item[2] for item in attempts] + verdict["ok"] = all(item["ok"] for item in all_verdicts) + verdict["deepest_verified"] = all( + item["deepest_verified"] for item in all_verdicts) + verdict["failed"] = sorted({ + failure for item in all_verdicts for failure in item["failed"] + }) + verdict["live_candidate_targets"] = sorted({ + target for item in all_verdicts + for target in item["live_candidate_targets"] + }) + verdict["deeper_live_candidates"] = sorted({ + target for item in all_verdicts + for target in item["deeper_live_candidates"] + }) + verdict["missing_candidate_markers"] = sorted({ + target for item in all_verdicts + for target in item["missing_candidate_markers"] + }) + tested_sets = [ + set(item["candidate_targets_tested"]) for item in all_verdicts + ] + verdict["candidate_targets_tested"] = sorted( + set.intersection(*tested_sets) if tested_sets else set()) + # Never-live has to hold on every capture process: one rank observing the candidate is enough + # to make it live, so this intersects rather than unions. + unobserved_sets = [ + set(item["installed_but_never_live_candidates"]) for item in all_verdicts + ] + verdict["installed_but_never_live_candidates"] = sorted( + set.intersection(*unobserved_sets) if unobserved_sets else set()) + for field in ( + "total_calls_observed", "target_marker_calls", "matched_kernel_calls", + "correlated_external_ids", "correlated_launch_correlations"): + verdict[field] = sum(int(item.get(field) or 0) for item in all_verdicts) + verdict["matched_kernel_names"] = sorted({ + name for item in all_verdicts for name in item["matched_kernel_names"] + }) + verdict["process_verdicts"] = [ + { + "capture_meta_file": meta_path, + "trace_files": paths, + "ok": item["ok"], + "failed": item["failed"], + "live_candidate_targets": item["live_candidate_targets"], + "deeper_live_candidates": item["deeper_live_candidates"], + } + for meta_path, paths, item in attempts + ] + all_paths = [path for _, paths, _ in attempts for path in paths] + verdict["best_process_meta_file"] = best_process_meta_file + verdict["trace_file"] = all_paths[0] if len(all_paths) == 1 else "" + verdict["trace_files"] = all_paths + verdict["trace_files_considered"] = trace_paths + payload = json.dumps(verdict, indent=2) + if args.out: + with open(args.out, "w") as fh: + fh.write(payload + "\n") + print(payload) + return 0 if verdict["ok"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/e2e_workflow/scripts/overlay_setup.py b/e2e_workflow/scripts/overlay_setup.py index e7cfd55dd..7d0a4e34e 100755 --- a/e2e_workflow/scripts/overlay_setup.py +++ b/e2e_workflow/scripts/overlay_setup.py @@ -25,6 +25,8 @@ --impl-module fast_act --impl-attr fast_silu_and_mul [--impl-file fast_act.py] add-capture install a shape/IO capture hook on module:attr (uses capture_shapes.py) --overlay O --target sglang...:fn --out [--max 5] [--capture-file capture_shapes.py] + add-marker install a marker-only hook on one candidate seam (uses seam_trace.py) + --overlay O --target sglang...:fn [--marker-file seam_trace.py] check print where a module resolves from (run with the overlay on PYTHONPATH) --module sglang.srt.layers.activation [--path-only] Every add-* takes --from BASE to SEED a new overlay from an existing one, so a candidate overlay is @@ -45,7 +47,7 @@ with open(_MAN) as _fh: _m = json.load(_fh) except Exception as _e: - _m = {"modules": [], "rebinds": [], "captures": []} + _m = {"modules": [], "rebinds": [], "markers": [], "captures": []} # (a) inject patched submodules under their dotted names BEFORE anything imports them. for _e in _m.get("modules", []): @@ -77,13 +79,23 @@ except Exception as _ex: sys.stderr.write("[overlay] rebind FAILED %r: %r\n" % (_e, _ex)) -# (c) capture hooks (shape/IO oracle recording). +# (c) capture hooks (shape/IO oracle recording) go on FIRST, so the capture wrapper is the innermost +# stand-in and is already bound before any marker install imports a module that does +# `from import ` (which would otherwise alias the un-captured function). for _e in _m.get("captures", []): try: import capture_shapes capture_shapes.install(_e["target"], _e["out"], int(_e.get("max", 5))) except Exception as _ex: sys.stderr.write("[overlay] capture install FAILED %r: %r\n" % (_e, _ex)) + +# (d) marker-only hooks used to compare every candidate seam in one trace. +for _e in _m.get("markers", []): + try: + import seam_trace + seam_trace.install(_e["target"]) + except Exception as _ex: + sys.stderr.write("[overlay] seam marker install FAILED %r: %r\n" % (_e, _ex)) ''' @@ -120,7 +132,7 @@ def _ensure_overlay(overlay, base=""): man = os.path.join(overlay, "_overlay_manifest.json") if not os.path.exists(man): with open(man, "w") as fh: - json.dump({"modules": [], "rebinds": [], "captures": []}, fh, indent=2) + json.dump({"modules": [], "rebinds": [], "markers": [], "captures": []}, fh, indent=2) return man @@ -199,6 +211,19 @@ def cmd_add_capture(a): print(f"launch with: PYTHONPATH={a.overlay}:$PYTHONPATH") +def cmd_add_marker(a): + man = _ensure_overlay(a.overlay, getattr(a, "base", "")) + marker = a.marker_file or os.path.join(os.path.dirname(os.path.abspath(__file__)), "seam_trace.py") + shutil.copy2(marker, os.path.join(a.overlay, "seam_trace.py")) + m = _load_man(man) + m["markers"] = [e for e in m.get("markers", []) if e["target"] != a.target] + m["markers"].append({"target": a.target}) + _save_man(man, m) + print(f"OVERLAY_DIR={a.overlay}") + print(f"add-marker {a.target}") + print(f"launch with: PYTHONPATH={a.overlay}:$PYTHONPATH") + + def cmd_check(a): f = module_file(a.module) if getattr(a, "path_only", False): @@ -245,6 +270,13 @@ def main(): p.add_argument("--from", dest="base", default="", help="seed the overlay from this existing overlay dir") p.set_defaults(func=cmd_add_capture) + p = sub.add_parser("add-marker") + p.add_argument("--overlay", required=True) + p.add_argument("--target", required=True, help="module:attr to mark without capturing I/O") + p.add_argument("--marker-file", default="", dest="marker_file") + p.add_argument("--from", dest="base", default="", help="seed the overlay from this existing overlay dir") + p.set_defaults(func=cmd_add_marker) + p = sub.add_parser("check") p.add_argument("--module", required=True) p.add_argument("--path-only", action="store_true", dest="path_only") diff --git a/e2e_workflow/scripts/parse_profile.py b/e2e_workflow/scripts/parse_profile.py index d6e7146cc..ab8545c22 100644 --- a/e2e_workflow/scripts/parse_profile.py +++ b/e2e_workflow/scripts/parse_profile.py @@ -87,16 +87,26 @@ def classify(name): return "other", "unknown", True, "Unclassified — inspect source to route." +SHORT_NAME_LIMIT = 60 + + def short_name(name): """Best-effort readable short name from a mangled C++/triton symbol.""" n = name # drop leading 'void ' and template/return noise n = re.sub(r"^void\s+", "", n) + # A symbol from an unnamed namespace opens with '(' where the identifier should be, so the + # '[\w:]+' probe below finds nothing and the whole signature falls through as the "short" name + # -- ROCm's '(anonymous namespace)::kda_packed_decode_kernel<...>(...)' used to shorten to the + # trailing parameter type rather than the kernel. + n = re.sub(r"\(anonymous namespace\)", "", n).lstrip(": ") # take the first identifier-ish token before '(' or '<' m = re.match(r"[\w:]+", n) base = m.group(0) if m else n base = base.split("::")[-1] - return base[:60] + # Truncation is for display only. kernel_selection.canonical_kernel_name knows about this limit + # and tolerates a prefix match at it, because mangled and Tensile symbols routinely exceed it. + return base[:SHORT_NAME_LIMIT] # --------------------------------------------------------------------------- # @@ -259,9 +269,12 @@ def _phase_of(ts): total_us += dur launches += 1 d = agg.setdefault(name, {"calls": 0, "total_us": 0.0, "shapes": set(), - "dtypes": set(), "by_case": {}, "by_phase": {}}) + "dtypes": set(), "by_case": {}, "by_phase": {}, + "cat_counts": {}}) d["calls"] += 1 d["total_us"] += dur + cat = e.get("cat", "") + d["cat_counts"][cat] = d["cat_counts"].get(cat, 0) + 1 # attribute this launch to its serving phase (measured from the step span it falls in) phase, stepM = _phase_of(e.get("ts")) if phase: @@ -341,7 +354,8 @@ def parse_rocprof_dir(d): us = ns / 1000.0 total_us += us launches += calls - e = agg.setdefault(name, {"calls": 0, "total_us": 0.0, "shapes": set(), "dtypes": set()}) + e = agg.setdefault(name, {"calls": 0, "total_us": 0.0, "shapes": set(), + "dtypes": set(), "src": "rocprof_kernel_stats"}) e["calls"] += calls e["total_us"] += us break # one stats file is the authoritative aggregate @@ -353,6 +367,299 @@ def norm_key(name): return re.sub(r"[^a-z0-9]", "", short_name(name).lower()) +# --------------------------------------------------------------------------- # +# INV-6: what KIND of thing is this row? +# +# The head track rewrites GPU kernels. A row that is actually a dispatcher op or a Python launcher can +# top a table by wall time while owning no device code to rewrite, and the downstream stages have no way +# to tell — a name is not evidence. So every Top-N row now carries `entity_kind` plus the EVIDENCE that +# produced it, derived from the profiler event category, never from a name pattern or a blacklist. +# +# gpu_kernel a dispatched compute kernel (rocprofv3 kernel-stats row, or a torch trace event +# with cat="kernel"). The only kind the head track accepts. +# memory_op a device memcpy/memset. Real device time, but no kernel source to rewrite. +# dispatcher_op a host-side operator span (torch cat="cpu_op"/"user_annotation"): its duration +# includes everything it dispatched, so it double-counts its own children. +# python_launcher a host span with no device work attributed to it at all. +# unresolved claimed by an upstream tool but not found among the dispatched kernels. +# --------------------------------------------------------------------------- # +ENTITY_KINDS = ("gpu_kernel", "memory_op", "dispatcher_op", "python_launcher", "unresolved") + +# torch profiler event category -> entity kind. Categories are emitted by the profiler itself, so this +# is a fact about how the work was observed, not a guess about what the symbol is called. +_CAT_ENTITY = { + "kernel": "gpu_kernel", + "gpu_memcpy": "memory_op", + "gpu_memset": "memory_op", + "cpu_op": "dispatcher_op", + "user_annotation": "dispatcher_op", + "cuda_runtime": "python_launcher", + "hip_runtime": "python_launcher", +} + + +def classify_entity(d, source): + """-> (entity_kind, evidence). `d` is an agg entry; `source` the profile source string.""" + cats = d.get("cat_counts") or {} + evidence_sources = set(d.get("evidence_sources") or []) + if d.get("src"): + evidence_sources.add(d["src"]) + observed_kinds = {_CAT_ENTITY.get(cat, "unresolved") for cat, n in cats.items() if n} + if "rocprof_kernel_stats" in evidence_sources: + observed_kinds.add("gpu_kernel") + device_kinds = observed_kinds & {"gpu_kernel", "memory_op"} + host_kinds = observed_kinds & {"dispatcher_op", "python_launcher"} + if device_kinds and host_kinds: + return "unresolved", { + "basis": "exact_name_multiple_entity_kinds", + "categories": dict(sorted(cats.items())), + "evidence_sources": sorted(evidence_sources), + "why": "the exact name denotes both host and device events; refusing to guess which entity " + "the external Top-N row represents", + } + if cats: + # A name can appear under several categories; the device categories decide. + kinds = {} + for cat, n in cats.items(): + kinds[_CAT_ENTITY.get(cat, "unresolved")] = kinds.get(_CAT_ENTITY.get(cat, "unresolved"), 0) + n + for k in ("gpu_kernel", "memory_op", "dispatcher_op", "python_launcher"): + if kinds.get(k): + return k, {"basis": "torch_profiler_event_category", + "categories": dict(sorted(cats.items())), "device_events": kinds[k]} + if d.get("src") == "rocprof_kernel_stats": + return "gpu_kernel", {"basis": "rocprofv3_kernel_stats_row", + "why": "rocprofv3 kernel-stats aggregates dispatched kernels only", + "calls": d.get("calls", 0)} + return "unresolved", {"basis": "no_category_evidence", + "why": f"source={source!r} carried no per-event category for this row"} + + +def index_host_entities(path): + """Name -> agg-shaped stub for HOST-side spans (cpu_op / annotations / runtime launch stubs). + + These are deliberately excluded from the timing aggregate — a cpu_op's duration subsumes the + kernels it dispatched, so counting it would double-count device time. But for --annotate we still + want to be able to say "this claimed head IS a dispatcher op" rather than the weaker "not found". + Used only by the annotate path, so the normal parse is unchanged. + """ + idx = {} + try: + with _open(path) as fh: + data = json.load(fh) + except Exception: + return idx + events = data.get("traceEvents", data if isinstance(data, list) else []) + host = {"cpu_op", "user_annotation", "cuda_runtime", "hip_runtime"} + for e in events: + if not isinstance(e, dict): + continue + cat = e.get("cat") + if cat not in host: + continue + d = idx.setdefault(e.get("name", "?"), {"calls": 0, "total_us": 0.0, "cat_counts": {}}) + d["calls"] += 1 + d["total_us"] += float(e.get("dur", 0.0) or 0.0) + d["cat_counts"][cat] = d["cat_counts"].get(cat, 0) + 1 + return idx + + +def index_dispatch_edges(path): + """Return host-op -> dispatched GPU kernels using profiler External ids. + + TraceLens may report one aggregate host/custom-op row whose percentage is the sum of several + kernels. Such a row is useful context but it is not a kernel-selection candidate. This index + provides the machine-observed children needed to split it before strategizing. + """ + try: + with _open(path) as fh: + data = json.load(fh) + except Exception: + return {} + events = data.get("traceEvents", data if isinstance(data, list) else []) + # Map every CPU External id to both its own op and enclosing dispatcher annotations. Kernels often + # carry the id of an inner launch op rather than the outer custom op that TraceLens aggregated. + # Thread-local interval stacks recover that ancestry without an O(events^2) containment scan. + host_by_ext = {} + spans_by_thread = {} + for event in events: + if not isinstance(event, dict) or event.get("cat") not in ("cpu_op", "user_annotation"): + continue + ext = (event.get("args") or {}).get("External id") + if ext is not None: + host_by_ext.setdefault(ext, set()).add(event.get("name", "?")) + if ext is not None and event.get("ts") is not None: + start = float(event["ts"]) + end = start + float(event.get("dur") or 0.0) + spans_by_thread.setdefault((event.get("pid"), event.get("tid")), []).append( + (start, end, event.get("name", "?"), ext)) + for spans in spans_by_thread.values(): + stack = [] + for start, end, name, ext in sorted(spans, key=lambda item: (item[0], -item[1])): + while stack and stack[-1][0] <= start: + stack.pop() + names = host_by_ext.setdefault(ext, set()) + names.add(name) + names.update(parent_name for _, parent_name in stack) + stack.append((end, name)) + edges = {} + for event in events: + if not isinstance(event, dict) or event.get("cat") != "kernel": + continue + ext = (event.get("args") or {}).get("External id") + hosts = host_by_ext.get(ext) or () + for host in hosts: + child = edges.setdefault(host, {}).setdefault( + event.get("name", "?"), {"calls": 0, "total_us": 0.0}) + child["calls"] += 1 + child["total_us"] += float(event.get("dur") or 0.0) + return edges + + +def expand_dispatcher_rows(rows, dispatch_edges, entity_agg=None, source="torch-trace"): + """Replace aggregate dispatcher rows with their observed device-kernel children. + + Rejection is not discovery. If the trace proves which kernels a dispatcher launched, preserve + its Amdahl mass but route it onto those concrete children. Unexpanded rows remain unchanged and + will later be classified/flagged normally. + """ + dispatch_edges = dispatch_edges or {} + concrete_keys = { + norm_key(row.get("name") or row.get("short_name") or "") + for row in (rows or []) + if (row.get("name") or row.get("short_name") or "") not in dispatch_edges + } + + expanded, records = [], [] + for row in rows or []: + rname = row.get("name") or row.get("short_name") or "" + # Expansion requires an exact observed host name. A normalized match is too lossy here: a host + # op and its GPU child routinely normalize to the same token. + hit = dispatch_edges.get(rname) + if hit and entity_agg is not None: + evidence = entity_agg.get(rname) + kind = classify_entity(evidence or {}, source)[0] + if kind != "dispatcher_op": + hit = None + if not hit: + expanded.append(row) + continue + total = sum(float(entry.get("total_us") or 0.0) for entry in hit.values()) + if total <= 0: + expanded.append(row) + continue + child_names = [] + for kernel_name, entry in sorted( + hit.items(), key=lambda item: float(item[1].get("total_us") or 0.0), reverse=True): + share = float(entry.get("total_us") or 0.0) / total + if norm_key(kernel_name) in concrete_keys: + # The upstream Top-N already emitted this device kernel. Keeping an expanded copy would + # double-count its Amdahl mass. + child_names.append(kernel_name) + continue + child = dict(row) + cls, backend, editable, hint = classify(kernel_name) + child.update({ + "name": kernel_name, + "short_name": short_name(kernel_name), + "device_kernel": kernel_name, + "profile_parent": rname, + "profile_parent_entity_kind": "dispatcher_op", + "profile_parent_pct_gpu_time": row.get("pct_gpu_time"), + "pct_gpu_time": round(float(row.get("pct_gpu_time") or 0.0) * share, 6), + "calls": int(entry.get("calls") or 0), + "total_ms": round(float(row.get("total_ms") or 0.0) * share, 6), + "classification": cls, + "backend_guess": backend, + "editable": editable, + "opt_hint": hint, + }) + child["avg_us"] = ( + child["total_ms"] * 1000.0 / child["calls"] if child["calls"] else 0.0) + child["notes"] = ( + f"Expanded from dispatcher '{rname}' using torch-profiler External-id edges. " + + str(row.get("notes") or "")) + expanded.append(child) + child_names.append(kernel_name) + records.append({"dispatcher": rname, "device_kernels": child_names}) + + expanded.sort(key=lambda row: float(row.get("pct_gpu_time") or 0.0), reverse=True) + for rank, row in enumerate(expanded, 1): + row["rank"] = rank + return expanded, records + + +def merge_entity_evidence(*aggregates): + """Merge profiler evidence without letting a same-name source overwrite another entity kind.""" + merged = {} + for agg in aggregates: + for name, raw in (agg or {}).items(): + d = dict(raw) + d["cat_counts"] = dict(raw.get("cat_counts") or {}) + sources = set(raw.get("evidence_sources") or []) + if raw.get("src"): + sources.add(raw["src"]) + d["evidence_sources"] = sorted(sources) + if name not in merged: + merged[name] = d + continue + cur = merged[name] + cur["calls"] = int(cur.get("calls") or 0) + int(d.get("calls") or 0) + cur["total_us"] = float(cur.get("total_us") or 0.0) + float(d.get("total_us") or 0.0) + for cat, count in d["cat_counts"].items(): + cur.setdefault("cat_counts", {})[cat] = cur.setdefault("cat_counts", {}).get(cat, 0) + count + cur_sources = set(cur.get("evidence_sources") or []) + cur_sources.update(d.get("evidence_sources") or []) + cur["evidence_sources"] = sorted(cur_sources) + return merged + + +def annotate_rows(rows, agg, source): + """Stamp entity_kind onto Top-N rows that were assembled OUTSIDE this script (the TraceLens + fast path builds them by hand from a third-party report). A row survives as a gpu_kernel only if its + name matches something this profiler actually observed being dispatched; otherwise it is + `unresolved` and the head track will refuse it. Returns (rows, stats).""" + agg = agg or {} + # C9: resolve by EXACT name first. norm_key() is a lossy fold (short_name + strip non-alnum), so a + # host dispatcher span 'aten::mul' and a device kernel 'mul' collapse to the SAME key; the old + # `setdefault` let whichever was aggregated first stamp its kind + matched_profiled_kernel onto the + # other row, silently misclassifying it and defeating the INV-6 gate in EITHER direction (a + # dispatcher_op admitted as a gpu_kernel, or a real kernel downgraded). The normalized fold is kept + # only as a fallback, and only when the key maps to exactly ONE observed entity. + by_exact = dict(agg) + by_norm, norm_collisions = {}, set() + for name in agg: + k = norm_key(name) + if k in by_norm and by_norm[k][0] != name: + norm_collisions.add(k) + else: + by_norm.setdefault(k, (name, agg[name])) + stats = {k: 0 for k in ENTITY_KINDS} + for r in rows: + rname = r.get("name") or r.get("short_name") or "" + k = norm_key(rname) + hit = None + if rname in by_exact: + hit = (rname, by_exact[rname]) + elif k in by_norm and k not in norm_collisions: + hit = by_norm[k] + if hit: + kind, ev = classify_entity(hit[1], source) + ev["matched_profiled_kernel"] = hit[0] + elif k in norm_collisions: + kind = "unresolved" # fail closed: refuse to guess which colliding entity + ev = {"basis": "ambiguous_name_match", + "why": f"'{rname}' normalizes to a key shared by multiple entities observed in " + f"{source}; refusing to guess its kind rather than misclassify it"} + else: + kind = "unresolved" + ev = {"basis": "name_not_found_in_profile", + "why": f"not among the {len(by_exact)} kernels observed in {source}"} + r["entity_kind"] = kind + r["entity_evidence"] = ev + stats[kind] = stats.get(kind, 0) + 1 + return rows, stats + + def build_summary(agg, total_us, launches, source, top_n, enrich=None, conc=0, isl=0, osl=0, chunk=None, capture_sizes=None, phase_meta=None): items = [] @@ -414,6 +721,11 @@ def build_summary(agg, total_us, launches, source, top_n, enrich=None, "editable": editable, "opt_hint": hint, } + entity_kind, entity_evidence = classify_entity(d, source) + entry["entity_kind"] = entity_kind + entry["entity_evidence"] = entity_evidence + if entity_kind == "gpu_kernel": + entry["device_kernel"] = name # SERVING-PHASE annotation (only when the trace exposed step spans). Enrich from the # HW-name-matched torch agg when this agg (e.g. rocprof) has no by_phase of its own. pd = d if d.get("by_phase") else (enrich_by_key.get(norm_key(name)) if enrich else None) @@ -561,6 +873,10 @@ def main(): "(for the kernel_workflow harness; needs --torch-trace for shapes)") ap.add_argument("--target", default="", help="optional kernel-name substring filter for --workload-out") + ap.add_argument("--annotate", default="", + help="annotate an existing Top-N JSON with profiler-backed entity kinds") + ap.add_argument("--annotate-out", default="", + help="output path for --annotate (defaults to replacing the input)") # serving-phase accounting (optional; pass the SAME ISL/OSL/concurrency as the bench). These only # EXPOSE the analytic per-phase call model (est_calls == serving_weight_model.analytic_calls) and # let est_shape snap decode M to the capture size; they never rescale the profile `weight`. @@ -589,6 +905,38 @@ def main(): if args.rocprof_dir: rp_agg, rp_total, rp_launch = parse_rocprof_dir(args.rocprof_dir) + if args.annotate: + with open(args.annotate) as fh: + doc = json.load(fh) + # Preserve all same-name evidence. A host span and a dispatched kernel can legitimately share + # an exact name; overwriting either one fabricates certainty, so classify_entity marks that row + # unresolved and the head gate refuses it. + host_agg = index_host_entities(args.torch_trace) if args.torch_trace else {} + ev_agg = merge_entity_evidence(rp_agg, host_agg, torch_agg) + ev_src = doc.get("source") or ("merged" if (rp_agg and torch_agg) else + "rocprofv3" if rp_agg else "torch-trace") + # A third-party Top-N may aggregate an outer custom op and its device children into one row. + # Turn that row into the actual kernels first; merely labelling/rejecting the dispatcher would + # repeat the 0802 failure mode instead of finding the launcher/kernel that 0720 optimized. + dispatch_edges = index_dispatch_edges(args.torch_trace) if args.torch_trace else {} + rows, expansions = expand_dispatcher_rows( + doc.get("top_kernels") or [], dispatch_edges, ev_agg, ev_src) + rows, stats = annotate_rows(rows, ev_agg, ev_src) + doc["top_kernels"] = rows + doc["entity_kind_contract"] = "v1" + doc["entity_kind_counts"] = {k: v for k, v in stats.items() if v} + doc["dispatcher_expansions"] = expansions + dest = args.annotate_out or args.annotate + with open(dest, "w") as fh: + fh.write(json.dumps(doc, indent=2)) + sys.stderr.write(f"annotated {len(rows)} row(s) -> {dest}: " + + ", ".join(f"{k}={v}" for k, v in stats.items() if v) + "\n") + bad = stats.get("unresolved", 0) + print(json.dumps({"annotated": len(rows), "entity_kind_counts": + {k: v for k, v in stats.items() if v}}, indent=2)) + # Non-zero exit when a claimed head is not a kernel this profiler ever saw dispatched. + sys.exit(1 if bad else 0) + if rp_agg and torch_agg: summ = build_summary(rp_agg, rp_total, rp_launch, "merged", args.top, enrich=torch_agg, phase_meta=torch_pmeta, **phk) diff --git a/e2e_workflow/scripts/seam_trace.py b/e2e_workflow/scripts/seam_trace.py new file mode 100644 index 000000000..5dd1acb9d --- /dev/null +++ b/e2e_workflow/scripts/seam_trace.py @@ -0,0 +1,214 @@ +"""Install safe profiler markers on candidate live call seams.""" + +import atexit +import functools +import importlib +import inspect +import os +import sys +import threading +import types + + +MARKER_PREFIX = "GEAK_TARGET::" +INSTALL_PREFIX = "GEAK_INSTALLED::" +_INSTALLED = {} +_TLS = threading.local() +_PROFILE = { + "lock": threading.Lock(), + "active": False, + "done": False, + "owner": None, + "profiler": None, + "active_calls": 0, + "root_calls": 0, + "out": "", + "trace_index": 0, + "atexit_registered": False, +} + + +def _rank(): + for key in ("RANK", "LOCAL_RANK", "TP_RANK", "SLURM_PROCID"): + value = os.environ.get(key) + if value is not None: + return str(value) + return "unknown" + + +def _trace_path(trace_index): + template = os.environ.get("GEAK_SELECTION_TRACE", "").strip() + if not template: + return "" + pid = os.getpid() + if "{pid}" in template or "{rank}" in template: + path = template.format(pid=pid, rank=_rank()) + else: + root, ext = os.path.splitext(template) + path = f"{root}.pid-{pid}.rank-{_rank()}{ext or '.json'}" + if os.environ.get("GEAK_SELECTION_TRACE_UNIQUE", "1") == "0": + # UNIQUE=0 opts out of the per-CALL suffix only. The per-PROCESS part has to stay: returning + # the raw template leaves a literal '{pid}' in the name, and kernel_selection pairs a capture + # to its trace by matching '.pid-' -- without it every rank races on one path and the + # verdict degrades to capture_process_trace_missing. + return path + root, ext = os.path.splitext(path) + return f"{root}.call-{trace_index}{ext or '.json'}" + + +def _record_install_markers(): + """Put proof of successful installation in every process-local trace.""" + import torch + for target in sorted(_INSTALLED): + with torch.profiler.record_function(INSTALL_PREFIX + target): + pass + + +def _start_profile(): + """Start the next bounded process-local root-call profile.""" + with _PROFILE["lock"]: + if _PROFILE["active"] or _PROFILE["done"]: + return False + budget = max(1, int(os.environ.get("GEAK_SELECTION_PROFILE_CALLS", "32"))) + if _PROFILE["trace_index"] >= budget: + _PROFILE["done"] = True + return False + out = _trace_path(_PROFILE["trace_index"] + 1) + if not out: + return False + ident = threading.get_ident() + profiler = None + try: + import torch + activities = [torch.profiler.ProfilerActivity.CPU] + if hasattr(torch.profiler.ProfilerActivity, "CUDA"): + activities.append(torch.profiler.ProfilerActivity.CUDA) + profiler = torch.profiler.profile(activities=activities) + profiler.__enter__() + _record_install_markers() + _PROFILE.update(active=True, owner=ident, profiler=profiler, out=out, + active_calls=0, root_calls=0) + if not _PROFILE["atexit_registered"]: + atexit.register(_finish_profile) + _PROFILE["atexit_registered"] = True + return True + except Exception as exc: + # Anything that fails after __enter__ has to close the profiler here. _finish_profile + # returns early once done is set, so a profiler left entered on this path would stay + # entered -- and collecting -- for the life of the process. + _PROFILE.update(active=False, done=True, owner=None, profiler=None) + sys.stderr.write(f"[seam_trace] profiler start failed: {exc!r}\n") + if profiler is not None: + try: + profiler.__exit__(None, None, None) + except Exception as close_exc: + # It stays entered, and collecting, for the life of the server. Nothing here can + # undo that, but an unexplained slowdown later is worth one line now. + sys.stderr.write(f"[seam_trace] profiler close failed: {close_exc!r}\n") + return False + + +def _finish_profile(): + """Stop and atomically export this process's trace once.""" + with _PROFILE["lock"]: + if not _PROFILE["active"] or _PROFILE["done"]: + return + profiler = _PROFILE.get("profiler") + out = _PROFILE.get("out") + _PROFILE["trace_index"] += 1 + budget = max(1, int(os.environ.get("GEAK_SELECTION_PROFILE_CALLS", "32"))) + _PROFILE.update(active=False, done=_PROFILE["trace_index"] >= budget, + owner=None, profiler=None) + try: + profiler.__exit__(None, None, None) + os.makedirs(os.path.dirname(os.path.abspath(out)), exist_ok=True) + root, ext = os.path.splitext(out) + tmp = f"{root}.tmp-{os.getpid()}{ext}" + profiler.export_chrome_trace(tmp) + os.replace(tmp, out) + sys.stderr.write(f"[seam_trace] selection trace -> {out}\n") + except Exception as exc: + sys.stderr.write(f"[seam_trace] profiler export failed: {exc!r}\n") + + +def _wrappable(value): + """Only replace callables whose identity/protocol a Python wrapper preserves.""" + if isinstance(value, (types.FunctionType, types.MethodType, functools.partial)): + return True + if isinstance(value, (types.BuiltinFunctionType, types.BuiltinMethodType)): + return False + if any(hasattr(value, attr) + for attr in ("fn", "cache", "warmup", "run", "__torch_dispatch__")): + return False + return callable(value) and type(value).__module__ != "builtins" + + +def _enter_call(): + _start_profile() + depth = getattr(_TLS, "depth", 0) + _TLS.depth = depth + 1 + with _PROFILE["lock"]: + if _PROFILE["active"]: + _PROFILE["active_calls"] += 1 + return depth == 0 + + +def _leave_call(root_call): + _TLS.depth = max(0, getattr(_TLS, "depth", 1) - 1) + finish = False + with _PROFILE["lock"]: + if _PROFILE["active"]: + _PROFILE["active_calls"] = max(0, _PROFILE["active_calls"] - 1) + if root_call: + _PROFILE["root_calls"] += 1 + finish = root_call and _PROFILE["active_calls"] == 0 + if finish: + _finish_profile() + + +def _resolve_owner(target): + """Resolve ``module:attr`` (attr may be dotted, e.g. ``Class.method``) to (owner, leaf). + + Class-method seams such as ``pkg.mod:RunnerCore.run`` are declared candidates just as often as + module-level functions; resolving only the flat case silently dropped them from the probe. + """ + module_name, attr = target.split(":", 1) + owner = importlib.import_module(module_name) + parts = attr.split(".") + for part in parts[:-1]: + owner = getattr(owner, part) + return owner, parts[-1] + + +def install(target): + """Wrap one module:attr with a record_function marker; idempotent per target.""" + if target in _INSTALLED: + return + module, attr = _resolve_owner(target) + original = getattr(module, attr) + if not _wrappable(original): + raise RuntimeError(f"cannot safely mark non-Python callable {target}") + + @functools.wraps(original) + def marked(*args, **kwargs): + root_call = _enter_call() + try: + import torch + record_function = torch.profiler.record_function + except Exception: + record_function = None + try: + if record_function is None: + return original(*args, **kwargs) + with record_function(MARKER_PREFIX + target): + return original(*args, **kwargs) + finally: + _leave_call(root_call) + + try: + marked.__signature__ = inspect.signature(original) + except (TypeError, ValueError): + pass + _INSTALLED[target] = original + setattr(module, attr, marked) + sys.stderr.write(f"[seam_trace] marked {target}\n") diff --git a/e2e_workflow/scripts/test_kernel_canonicalization_parity.js b/e2e_workflow/scripts/test_kernel_canonicalization_parity.js new file mode 100644 index 000000000..89455d6ac --- /dev/null +++ b/e2e_workflow/scripts/test_kernel_canonicalization_parity.js @@ -0,0 +1,60 @@ +#!/usr/bin/env node +// Regression guard for kernel symbol canonicalization (no GPU, no model needed). +// +// Invariant under test: canonicalDeviceKernel / kernelIdentitiesMatch in e2e_workflow.js agree with +// canonical_kernel_name / kernel_matches in scripts/kernel_selection.py. The JS gate and the Python +// verdict compare the SAME two symbols, so any drift between them lets a kernel pass one side and be +// refused by the other. Both sides are pinned to tests/kernel_symbols.json; the Python half of this +// is TestTheSharedFixtureHoldsOnThisSide in tests/test_kernel_selection.py. +// +// The functions are extracted from the real workflow source rather than reimplemented here -- a copy +// would pass this test while the shipped code drifted, which is the failure it exists to catch. +// +// Run: node e2e_workflow/scripts/test_kernel_canonicalization_parity.js +'use strict'; +const fs = require('fs'); +const path = require('path'); + +const ROOT = path.resolve(__dirname, '..', '..'); // .../GEAK +const WORKFLOW = path.join(ROOT, 'e2e_workflow', 'e2e_workflow.js'); +const FIXTURE = path.join(ROOT, 'e2e_workflow', 'scripts', 'tests', 'kernel_symbols.json'); + +let failures = 0; +const ok = (cond, msg) => { if (!cond) { console.error(' FAIL:', msg); failures++; } else console.log(' ok:', msg); }; + +const src = fs.readFileSync(WORKFLOW, 'utf8'); +const start = src.indexOf('const SHORT_NAME_LIMIT'); +const end = src.indexOf('function requiredDeviceKernel'); +ok(start !== -1 && end !== -1 && start < end, 'canonicalization block located in e2e_workflow.js'); +if (failures) process.exit(1); + +const build = new Function( + `${src.slice(start, end)}\nreturn { canonicalDeviceKernel, kernelIdentitiesMatch, SHORT_NAME_LIMIT };`); +const { canonicalDeviceKernel, kernelIdentitiesMatch, SHORT_NAME_LIMIT } = build(); + +const fixture = JSON.parse(fs.readFileSync(FIXTURE, 'utf8')); + +console.log('\n# canonical tokens'); +for (const c of fixture.canonical) { + const got = canonicalDeviceKernel(c.symbol); + ok(got === c.canonical, + `${JSON.stringify(c.symbol.slice(0, 56))} -> ${JSON.stringify(c.canonical)}` + + (got === c.canonical ? '' : ` (got ${JSON.stringify(got)})`)); +} + +console.log('\n# match verdicts'); +for (const c of fixture.matches) { + const forward = kernelIdentitiesMatch(c.a, c.b), back = kernelIdentitiesMatch(c.b, c.a); + ok(forward === c.match && back === c.match, + `${c.why} (got ${forward}/${back}, want ${c.match})`); +} + +console.log('\n# constants shared with the Python side'); +ok(SHORT_NAME_LIMIT === 60, 'SHORT_NAME_LIMIT matches parse_profile.SHORT_NAME_LIMIT'); + +// The fixture is only worth anything while it still carries symbols that real captures produced. +const real = fixture.canonical.filter((c) => c.real).length; +ok(real >= 5, `fixture keeps ${real} symbols taken verbatim from ROCm captures`); + +console.log(failures ? `\nFAILED (${failures})` : '\nPASS'); +process.exit(failures ? 1 : 0); diff --git a/e2e_workflow/scripts/tests/kernel_symbols.json b/e2e_workflow/scripts/tests/kernel_symbols.json new file mode 100644 index 000000000..4cae4657f --- /dev/null +++ b/e2e_workflow/scripts/tests/kernel_symbols.json @@ -0,0 +1,178 @@ +{ + "_comment": [ + "Shared fixture for kernel symbol canonicalization. Two implementations must agree on it:", + " - canonical_kernel_name / kernel_matches in e2e_workflow/scripts/kernel_selection.py", + " - canonicalDeviceKernel / kernelIdentitiesMatch in e2e_workflow/e2e_workflow.js", + "The JS gate and the Python verdict compare the SAME two symbols, so any drift between them lets", + "a kernel pass one side and be refused by the other. Asserted from Python by", + "tests/test_kernel_selection.py and from node by scripts/test_kernel_canonicalization_parity.js.", + "", + "Symbols marked 'real' were lifted verbatim from ROCm captures. Hand-written C++ never produced", + "the shapes that broke this: an unnamed namespace putting a '(' in front of the kernel, a bare", + "kernel with no namespace to separate it from the 'void' return type, and the annotations this", + "pipeline appends to a name after the fact -- a phase tag, a cluster label, a memcpy direction." + ], + "canonical": [ + { + "real": true, + "symbol": "void (anonymous namespace)::kda_packed_decode_kernel<8, false>((anonymous namespace)::KdaPackedDecodeParams)", + "canonical": "kda_packed_decode_kernel" + }, + { + "real": true, + "symbol": "void aiter::greedy_sample_kernel(float const*, int*, int, int)", + "canonical": "greedy_sample_kernel" + }, + { + "real": true, + "symbol": "void wvSplitKrc_<__hip_bfloat16, 8, 4>(void*, int)", + "canonical": "wvsplitkrc_" + }, + { + "real": true, + "symbol": "void paged_attention_ll4mi_QKV_mfma16_kernel<0, __hip_bfloat16, (vllm::Fp8KVCacheDataType)0, 256>(float*, int*)", + "canonical": "paged_attention_ll4mi_qkv_mfma16_kernel" + }, + { + "real": true, + "symbol": "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bLi256EEEvPT_i", + "canonical": "_zn5aiter24add_rmsnorm_quant_kernelidf16bli256eeevpt_i" + }, + { + "real": true, + "symbol": "reshape_and_cache_shuffle_5d", + "canonical": "reshape_and_cache_shuffle_5d" + }, + { + "symbol": "at::native::vectorized_elementwise_kernel<4, at::native::AddFunctor, at::detail::Array >(int, float)", + "canonical": "vectorized_elementwise_kernel" + }, + { + "symbol": "vectorized_elementwise_kernel [clone .isra.0]", + "canonical": "vectorized_elementwise_kernel" + }, + { + "real": true, + "_note": "profile artifacts store long kernel names elided mid-template, leaving an opener that never closes", + "symbol": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_k...", + "canonical": "elementwise_kernel_manual_unroll" + }, + { + "real": true, + "symbol": "_ZN7ck_tile6kentryILi2ENS_15MoeFlatmmKernelINS_33GemmSpatiallyLocalTilePart...", + "canonical": "_zn7ck_tile6kentryili2ens_15moeflatmmkernelins_33gemmspatiallylocaltilepart" + }, + { + "real": true, + "_note": "the annotations below are appended by this pipeline, not by the compiler. Reading the name from the END of the symbol answers with the annotation and the head stops matching its own kernel.", + "symbol": "_fwd_grouped_kernel_stage1 [sliding_attention]", + "canonical": "_fwd_grouped_kernel_stage1" + }, + { "real": true, "symbol": "main_kernel[prefill]", "canonical": "main_kernel" }, + { "real": true, "symbol": "main_kernel [cluster B (long, ~34us)]", "canonical": "main_kernel" }, + { "real": true, "symbol": "hgemm_bf16_32x64x256x3_SPK1 [qkv_proj]", "canonical": "hgemm_bf16_32x64x256x3_spk1" }, + { "real": true, "symbol": "allreduce_prototype_twoshot@prefill", "canonical": "allreduce_prototype_twoshot" }, + { + "real": true, + "_note": "rocprof spells memory ops with the direction after the operation", + "symbol": "Memcpy DtoD (Device -> Device)", + "canonical": "memcpy" + }, + { + "real": true, + "_note": "an unnamed namespace can also sit mid-qualification, leaving 'at::native:: ::CatArray...' once the group is removed", + "symbol": "void at::native::(anonymous namespace)::CatArrayBatchedCopy", "canonical": "" }, + { "symbol": "", "canonical": "" } + ], + "matches": [ + { + "why": "the bare name and its decorated spelling are one kernel", + "a": "paged_attention_ll4mi_QKV_mfma4_kernel", + "b": "void paged_attention_ll4mi_QKV_mfma4_kernel<__hip_bfloat16, 128, 256>(float*, int)", + "match": true + }, + { + "why": "a bare declared name has no template arguments to compare, so it matches the decorated spelling", + "a": "vectorized_elementwise_kernel", + "b": "void at::native::vectorized_elementwise_kernel<4, AddFunctor >(int)", + "match": true + }, + { + "why": "two instantiations of one template are two kernels once both sides carry the arguments", + "a": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor >(int)", + "b": "void at::native::vectorized_elementwise_kernel<4, at::native::MulFunctor >(int)", + "match": false + }, + { + "why": "real head row: a short_name cut inside the ARGUMENT list keeps a whole template, and only one side spells the return type and namespace", + "a": "clamp_position_kernel(long*, long const*, unsigned lon", + "b": "void (anonymous namespace)::clamp_position_kernel(long*, long const*, unsigned long, int)", + "match": true + }, + { + "why": "one build leaves the namespace unnamed and another declares it; same kernel, same instantiation", + "a": "void (anonymous namespace)::store_kvcache<512l, 512l, 1, false, long>(int)", + "b": "void sglang::store_kvcache<512l, 512l, 1, false, long>(sglang::StoreKVCacheParams)", + "match": true + }, + { + "why": "both argument lists are elided, so they must agree as far as they are spelled: 128, 4 against 128, 8", + "a": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_k...", + "b": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_k...", + "match": false + }, + { + "why": "folding separators to _ stops '128, 4,' from reading as a prefix of '128, 48,'", + "a": "k<128, 4, at::native::gpu_k", "b": "k<128, 48, at::native::gpu_k", "match": false + }, + { + "why": "real sglang neighbours: a name appearing inside another is a different kernel", + "a": "_fwd_kernel", "b": "_fwd_kernel_stage2", "match": false + }, + { + "why": "a name embedded in another is not evidence of the same kernel", + "a": "gemm_kernel", "b": "fused_gemm_kernel_v2", "match": false + }, + { + "why": "two unnamed-namespace kernels used to both reduce to 'void' and certify each other", + "a": "void (anonymous namespace)::clamp_position_kernel(long*, long)", + "b": "void (anonymous namespace)::kda_packed_decode_kernel<8, false>(int)", + "match": false + }, + { + "why": "a token this short appears inside unrelated kernels and is not evidence", + "a": "gemm", "b": "gemm_kernel_v2", "match": false + }, + { + "why": "a prefix that is not the display-limit truncation is a different kernel", + "a": "attention_kernel", "b": "attention_kernelbackward", "match": false + }, + { + "why": "a name cut at the 60-char display limit ends mid-token, so only a prefix rule can match it", + "a": "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bDF16bLi256ELi16ELb1", + "b": "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bDF16bLi256ELi16ELb1ELb0ELb1ELi1EEEvPT0_i", + "match": true + }, + { + "why": "one character short of the display limit is just an unrelated fragment", + "a": "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bDF16bLi256ELi16ELb", + "b": "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bDF16bLi256ELi16ELb1ELb0ELb1ELi1EEEvPT0_i", + "match": false + }, + { + "why": "real head row: the phase tag is ours, so the row still has to match the kernel it names", + "a": "_fwd_grouped_kernel_stage1 [sliding_attention]", + "b": "_fwd_grouped_kernel_stage1", + "match": true + }, + { + "why": "real head row: a memcpy row and its own rocprof spelling", + "a": "Memcpy", "b": "Memcpy DtoD (Device -> Device)", "match": true + }, + { "why": "an empty side is never evidence", "a": "", "b": "gemm_kernel", "match": false } + ] +} diff --git a/e2e_workflow/scripts/tests/test_capture_shapes.py b/e2e_workflow/scripts/tests/test_capture_shapes.py index f46f7fc9b..c0ef7e4ea 100644 --- a/e2e_workflow/scripts/tests/test_capture_shapes.py +++ b/e2e_workflow/scripts/tests/test_capture_shapes.py @@ -576,7 +576,7 @@ def test_oracle_file_is_written_and_hashed(self): io_path = os.path.join(self.out_dir, "reference_io.pt") self.assertTrue(os.path.exists(io_path)) payload, path = self.torch.saved[-1] - self.assertEqual(path, io_path) + self.assertTrue(path.startswith(io_path + ".tmp-")) self.assertEqual(payload["target"], "fake_serving_layer:op") self.assertEqual([r["regime"] for r in payload["records"]], ["decode", "prefill"]) sha = self._meta()["reference_io_sha256"] @@ -857,6 +857,15 @@ def test_install_from_env_reads_target_out_and_max(self): self.assertEqual(cs._STATE["out_dir"], self.out_dir) self.assertIsNot(mod.op, cs._STATE["orig"]) + def test_selection_capture_uses_a_process_local_artifact_directory(self): + self._target_module() + with _env(GEAK_SELECTION_TRACE=os.path.join(self.out_dir, "selection.json")): + with _stderr(): + cs.install("fake_serving_layer:op", self.out_dir) + expected_prefix = os.path.join( + self.out_dir, f"capture.pid-{os.getpid()}.rank-") + self.assertTrue(cs._STATE["out_dir"].startswith(expected_prefix)) + # --------------------------------------------------------------------------- # # import-time self-install (the overlay PYTHONPATH / sitecustomize entry point) diff --git a/e2e_workflow/scripts/tests/test_kernel_selection.py b/e2e_workflow/scripts/tests/test_kernel_selection.py new file mode 100644 index 000000000..9679e964e --- /dev/null +++ b/e2e_workflow/scripts/tests/test_kernel_selection.py @@ -0,0 +1,741 @@ +#!/usr/bin/env python3 +"""Regression tests for machine-verified live kernel selection.""" + +import importlib.util +import json +import os +import tempfile +import unittest + + +SCRIPTS = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SPEC = importlib.util.spec_from_file_location( + "kernel_selection", os.path.join(SCRIPTS, "kernel_selection.py")) +ks = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(ks) + +# parse_profile owns the display-name limit that canonical_kernel_name's prefix rule depends on. +PP_SPEC = importlib.util.spec_from_file_location( + "parse_profile", os.path.join(SCRIPTS, "parse_profile.py")) +pp = importlib.util.module_from_spec(PP_SPEC) +PP_SPEC.loader.exec_module(pp) + +TARGET = ( + "vllm.v1.attention.ops.chunked_prefill_paged_decode:" + "chunked_prefill_paged_decode" +) +KERNEL = "kernel_paged_attention_2d" + + +def trace(kernel=KERNEL, under_target=True): + marker_start = 100 + kernel_start = 120 if under_target else 250 + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": marker_start, "dur": 100}, + {"cat": "kernel", "name": f"void vllm::{kernel}(int)", + "ph": "X", "ts": kernel_start, "dur": 25, "args": {"External id": 7}}, + ] + if under_target: + events.insert(1, {"cat": "cpu_op", "name": "launch", "ph": "X", + "pid": 1, "tid": 2, "ts": 110, "dur": 5, + "args": {"External id": 7}}) + return events + + +class TestCallableSpec(unittest.TestCase): + def test_only_exact_machine_specs_are_accepted(self): + self.assertTrue(ks.valid_callable_spec(TARGET)) + self.assertFalse(ks.valid_callable_spec(TARGET + " -> inner kernel")) + self.assertFalse(ks.valid_callable_spec("vllm/path.py:launcher")) + + +class TestKernelMatching(unittest.TestCase): + def test_demangled_kernel_matches_profile_identity(self): + self.assertTrue(ks.kernel_matches( + KERNEL, "void vllm::kernel_paged_attention_2d(int)")) + self.assertFalse(ks.kernel_matches(KERNEL, "unrelated_attention_kernel")) + + def test_nested_template_arguments_reduce_to_the_bare_kernel_name(self): + """Demangled C++ symbols nest their template arguments, and the argument list can carry both + '::' and '(' -- the two delimiters the rest of the canonicalization splits on. Leaving any of + it behind produces a token that never matches the profile's own name for the same kernel.""" + for decorated in ( + "at::native::vectorized_elementwise_kernel<4, at::native::AddFunctor, " + "at::detail::Array >(int, float)", + "void ns::vectorized_elementwise_kernel>(void*)", + "vectorized_elementwise_kernel [clone .isra.0]", + ): + with self.subTest(decorated=decorated): + self.assertEqual(ks.canonical_kernel_name(decorated), + "vectorized_elementwise_kernel") + self.assertTrue(ks.kernel_matches("vectorized_elementwise_kernel", decorated)) + + def test_a_symbol_elided_mid_template_answers_with_the_name_not_a_fragment(self): + """Profile artifacts store long kernel names elided mid-template, so the opener never closes + and no amount of balanced stripping removes it. Reading the last whitespace token then lifts + a fragment out of the template arguments -- this real elision answered 'gpu_k' -- and states + it as confidently as a real name, which matches the wrong kernel instead of refusing.""" + self.assertEqual( + ks.canonical_kernel_name( + "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_k..."), + "elementwise_kernel_manual_unroll") + self.assertEqual( + ks.canonical_kernel_name("void aiter::grouped_topk_kernel(c10::BFloat1"), + "grouped_topk_kernel") + self.assertEqual(ks.canonical_kernel_name("gemm_kernel"), "") + + def test_an_elided_name_still_matches_the_full_symbol_it_was_cut_from(self): + elided = "_ZN7ck_tile6kentryILi2ENS_15MoeFlatmmKernelINS_33GemmSpatiallyLocalTilePart..." + full = ("_ZN7ck_tile6kentryILi2ENS_15MoeFlatmmKernelINS_33GemmSpatiallyLocalTilePartitioner" + "INS_13TileGemmShapeINS_8sequenceIJLi16") + self.assertTrue(ks.kernel_matches(elided, full)) + + +# Symbols lifted verbatim from ROCm captures under /shared_nfs/hyperloom-claw. They are here because +# hand-written C++ never produced the shapes that broke this: an unnamed namespace putting a '(' in +# front of the kernel, and a bare kernel with no namespace at all to separate it from 'void'. +REAL_ROCM_SYMBOLS = { + "void (anonymous namespace)::kda_packed_decode_kernel<8, false>" + "((anonymous namespace)::KdaPackedDecodeParams)": "kda_packed_decode_kernel", + "void aiter::greedy_sample_kernel(float const*, int*, int, int)": + "greedy_sample_kernel", + "void paged_attention_ll4mi_QKV_mfma16_kernel<0, __hip_bfloat16, " + "(vllm::Fp8KVCacheDataType)0, 256>(float*, int*)": "paged_attention_ll4mi_qkv_mfma16_kernel", + "void wvSplitKrc_<__hip_bfloat16, 8, 4>(void*, int)": "wvsplitkrc_", + "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bLi256EEEvPT_i": "_zn5aiter24add_rmsnorm_quant_kernelidf16bli256eeevpt_i", + "reshape_and_cache_shuffle_5d": "reshape_and_cache_shuffle_5d", +} + + +class TestRealRocmSymbolsCanonicalizeToTheirKernel(unittest.TestCase): + def test_each_symbol_reduces_to_the_kernel_it_names(self): + for symbol, token in REAL_ROCM_SYMBOLS.items(): + with self.subTest(symbol=symbol[:48]): + self.assertEqual(ks.canonical_kernel_name(symbol), token) + + def test_an_unnamed_namespace_does_not_collapse_the_symbol_to_its_return_type(self): + """ROCm spells the unnamed namespace '(anonymous namespace)', which puts a parenthesis BEFORE + the kernel. Cutting the symbol at its first '(' left the return type as the whole token, so + every such kernel reduced to 'void' -- and two unrelated ones then certified each other, + which is exactly the unearned credit this verdict exists to refuse.""" + one = "void (anonymous namespace)::clamp_position_kernel(long*, long)" + two = "void (anonymous namespace)::kda_packed_decode_kernel<8, false>(int)" + self.assertEqual(ks.canonical_kernel_name(one), "clamp_position_kernel") + self.assertEqual(ks.canonical_kernel_name(two), "kda_packed_decode_kernel") + self.assertFalse(ks.kernel_matches(one, two)) + self.assertFalse(ks.kernel_matches(two, one)) + + def test_a_kernel_with_no_namespace_is_not_fused_with_its_return_type(self): + """Stripping separators rather than splitting on them glued 'void' onto any kernel that had + no '::' to fall back on, so the bare name never matched its own decorated spelling.""" + bare = "paged_attention_ll4mi_QKV_mfma4_kernel" + decorated = f"void {bare}<__hip_bfloat16, 128, 256>(float*, int)" + self.assertEqual(ks.canonical_kernel_name(decorated), bare.lower()) + self.assertTrue(ks.kernel_matches(bare, decorated)) + + def test_a_name_truncated_by_the_display_limit_still_matches_its_full_symbol(self): + """parse_profile.short_name caps display names, and mangled and Tensile symbols run well past + the cap. The truncated name ends mid-token, so no word boundary can follow it -- without an + explicit prefix rule our own shortening reads as a different kernel.""" + symbol = "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bDF16bLi256ELi16ELb1ELb0ELb1ELi1EEEvPT0_i" + declared = symbol[:ks.SHORT_NAME_LIMIT] + self.assertGreater(len(symbol), ks.SHORT_NAME_LIMIT) + self.assertTrue(ks.kernel_matches(declared, symbol)) + self.assertFalse(ks.kernel_matches(declared[:ks.SHORT_NAME_LIMIT - 1], symbol)) + + def test_a_name_embedded_inside_another_is_a_different_kernel(self): + """The first two pairs are real neighbours inside one sglang capture. Accepting a name + because it appears inside the other certified _fwd_kernel as _fwd_kernel_stage2 -- a + different kernel, which is precisely the credit this verdict exists to withhold.""" + for a, b in ( + ("_fwd_kernel", "_fwd_kernel_stage2"), + ("reshape_and_cache_kernel", "reshape_and_cache_kernel_flash"), + ("gemm_kernel", "fused_gemm_kernel_v2"), + ("gemm", "gemm_kernel_v2"), + ("attention_kernel", "attention_kernelbackward"), + ): + with self.subTest(a=a, b=b): + self.assertFalse(ks.kernel_matches(a, b)) + self.assertFalse(ks.kernel_matches(b, a)) + + def test_two_instantiations_of_one_template_are_not_one_kernel(self): + """The base token drops template arguments so a bare declared name can match its decorated + spelling. When both sides carry those arguments the information is present on both, and one + capture here held 20 distinct kernels whose only difference was the functor.""" + fill = ("void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor, " + "std::array >(int, at::native::FillFunctor)") + power = ("void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous " + "namespace)::pow_tensor_scalar_kernel_impl, std::array >(int)") + self.assertEqual(ks.canonical_kernel_name(fill), ks.canonical_kernel_name(power)) + self.assertFalse(ks.kernel_matches(fill, power)) + self.assertTrue(ks.kernel_matches(fill, fill)) + + def test_a_bare_declared_name_still_matches_a_templated_symbol(self): + """The tightening above applies only when BOTH sides carry template arguments. A head that + declares the bare name has none to compare, and refusing it would reject every real head.""" + self.assertTrue(ks.kernel_matches( + "vectorized_elementwise_kernel", + "void at::native::vectorized_elementwise_kernel<4, AddFunctor >(int)")) + + def test_arguments_are_read_past_the_return_type_and_namespace(self): + """Only one side spells the return type and namespace. This pair is a real head row whose + stored short_name had been cut inside the ARGUMENT list while its template stayed whole -- + comparing the symbols entire refused a kernel against itself.""" + self.assertTrue(ks.kernel_matches( + "clamp_position_kernel(long*, long const*, unsigned lon", + "void (anonymous namespace)::clamp_position_kernel(long*, long const*, unsigned " + "long, int)")) + + def test_one_kernel_matches_across_builds_that_name_its_namespace_differently(self): + """Both spellings occur across the captures: one build leaves the namespace unnamed, another + declares it. Same kernel, same instantiation.""" + self.assertTrue(ks.kernel_matches( + "void (anonymous namespace)::store_kvcache<512l, 512l, 1, false, long>(int)", + "void sglang::store_kvcache<512l, 512l, 1, false, long>(sglang::StoreKVCacheParams)")) + + def test_an_elided_argument_list_still_has_to_agree_as_far_as_it_is_spelled(self): + """Artifacts elide long names mid-template, so neither list is complete. Skipping the check + entirely would let every elided instantiation of one template certify the others -- these two + differ only in the visible '128, 4' against '128, 8'.""" + four = "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_k..." + eight = "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_k..." + self.assertEqual(ks.canonical_kernel_name(four), ks.canonical_kernel_name(eight)) + self.assertFalse(ks.kernel_matches(four, eight)) + self.assertTrue(ks.kernel_matches(four, four)) + + def test_a_separator_cannot_make_one_argument_list_a_prefix_of_another(self): + """Deleting separators instead of folding them read '<128, 4,' as a prefix of '<128, 48,'.""" + self.assertFalse(ks.kernel_matches("k<128, 4, at::native::gpu_k", + "k<128, 48, at::native::gpu_k")) + + def test_the_display_limit_agrees_with_the_module_that_applies_it(self): + """The prefix rule above is only sound while both sides mean the same number of characters.""" + self.assertEqual(ks.SHORT_NAME_LIMIT, pp.SHORT_NAME_LIMIT) + + +class TestTheSharedFixtureHoldsOnThisSide(unittest.TestCase): + """kernel_symbols.json is the one place both canonicalizers are pinned. The node half of this is + scripts/test_kernel_canonicalization_parity.js; a rule added here without adding it there is + exactly the drift that lets a kernel pass the JS gate and be refused by this verdict.""" + + @classmethod + def setUpClass(cls): + with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), + "kernel_symbols.json")) as handle: + cls.fixture = json.load(handle) + + def test_every_symbol_reduces_to_its_recorded_token(self): + for case in self.fixture["canonical"]: + with self.subTest(symbol=case["symbol"][:48]): + self.assertEqual(ks.canonical_kernel_name(case["symbol"]), case["canonical"]) + + def test_every_recorded_pair_gets_the_recorded_verdict(self): + for case in self.fixture["matches"]: + with self.subTest(why=case["why"]): + self.assertEqual(ks.kernel_matches(case["a"], case["b"]), case["match"]) + self.assertEqual(ks.kernel_matches(case["b"], case["a"]), case["match"]) + + def test_the_fixture_still_covers_symbols_taken_from_real_captures(self): + """Hand-written C++ is what let the unnamed-namespace and bare-kernel defects through.""" + self.assertGreaterEqual(sum(1 for c in self.fixture["canonical"] if c.get("real")), 5) + + +class TestSelectionVerdict(unittest.TestCase): + def meta(self, calls=7, target=TARGET): + module, attr = target.split(":", 1) + return {"module": module, "attr": attr, "total_calls_observed": calls} + + def test_0720_live_inner_launcher_is_selection_success(self): + verdict = ks.verify(TARGET, KERNEL, self.meta(), trace()) + self.assertTrue(verdict["ok"]) + self.assertEqual(verdict["matched_kernel_calls"], 1) + self.assertEqual(verdict["failed"], []) + + def test_capture_and_seam_markers_count_one_logical_call(self): + """capture_shapes and seam_trace both emit GEAK_TARGET for the selected callable. The trace + therefore nests the same marker inside itself, but the report must count the invocation once.""" + events = trace() + events.insert(2, { + "cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 105, "dur": 80, + }) + verdict = ks.verify(TARGET, KERNEL, self.meta(calls=1), events) + self.assertTrue(verdict["ok"], verdict) + self.assertEqual(verdict["target_marker_calls"], 1) + self.assertEqual(verdict["matched_kernel_calls"], 1) + + def test_nested_markers_on_other_threads_and_later_calls_stay_distinct(self): + marker = ks.MARKER_PREFIX + TARGET + spans = ks._outermost_spans([ + (100.0, 200.0, {"pid": 1, "tid": 2, "name": marker}), + (110.0, 190.0, {"pid": 1, "tid": 2, "name": marker}), + (120.0, 180.0, {"pid": 1, "tid": 3, "name": marker}), + (300.0, 320.0, {"pid": 1, "tid": 2, "name": marker}), + ]) + self.assertEqual([(start, end) for start, end, _ in spans], + [(100.0, 200.0), (120.0, 180.0), (300.0, 320.0)]) + + def test_kernel_seen_elsewhere_is_not_enough(self): + verdict = ks.verify(TARGET, KERNEL, self.meta(), trace(under_target=False)) + self.assertFalse(verdict["ok"]) + self.assertIn("device_kernel_not_under_target", verdict["failed"]) + + def test_async_kernel_after_marker_is_linked_by_external_id(self): + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 40}, + {"cat": "cpu_op", "name": "launch", "ph": "X", "pid": 1, "tid": 2, + "ts": 110, "dur": 5, + "args": {"External id": 42}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "ts": 300, "dur": 20, + "args": {"External id": 42}}, + ] + verdict = ks.verify(TARGET, KERNEL, self.meta(), events) + self.assertTrue(verdict["ok"], verdict) + self.assertEqual(verdict["correlated_external_ids"], 1) + + def test_triton_kernel_without_external_id_is_linked_by_launch_correlation(self): + # ROCm/kineto emits Triton device rows (hipModuleLaunchKernel) with only `correlation`; + # `External id` is absent, so the External-id bridge alone reports a false negative. + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 40}, + {"cat": "cuda_runtime", "name": "hipModuleLaunchKernel", "ph": "X", + "pid": 1, "tid": 2, "ts": 110, "dur": 5, "args": {"correlation": 25}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "pid": 2, "tid": 0, + "ts": 300, "dur": 20, "args": {"correlation": 25, "stream": 0}}, + ] + verdict = ks.verify(TARGET, KERNEL, self.meta(), events) + self.assertTrue(verdict["ok"], verdict) + self.assertEqual(verdict["matched_kernel_calls"], 1) + self.assertEqual(verdict["correlated_launch_correlations"], 1) + + def test_launch_correlation_outside_the_marker_span_is_not_enough(self): + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 40}, + {"cat": "cuda_runtime", "name": "hipModuleLaunchKernel", "ph": "X", + "pid": 1, "tid": 2, "ts": 500, "dur": 5, "args": {"correlation": 25}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "pid": 2, "tid": 0, + "ts": 600, "dur": 20, "args": {"correlation": 25, "stream": 0}}, + ] + verdict = ks.verify(TARGET, KERNEL, self.meta(), events) + self.assertFalse(verdict["ok"]) + self.assertIn("device_kernel_not_under_target", verdict["failed"]) + + def test_correlation_ids_do_not_collide_across_merged_call_traces(self): + # Every per-call trace restarts `correlation` at 1, so an unprefixed merge would let call-2's + # kernel be attributed to call-1's in-span launch. + def one(marked, corr, kernel_ts): + evs = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "pid": 2, "tid": 0, + "ts": kernel_ts, "dur": 20, "args": {"correlation": corr}}, + ] + if marked: + evs += [ + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 40}, + {"cat": "cuda_runtime", "name": "hipModuleLaunchKernel", "ph": "X", + "pid": 1, "tid": 2, "ts": 110, "dur": 5, "args": {"correlation": corr}}, + ] + return {"traceEvents": evs} + + with tempfile.TemporaryDirectory() as tmp: + paths = [] + # call-1: marked launch, correlation 1. call-2: NO marker at all, correlation 1 again. + for i, doc in enumerate((one(True, 1, 300), one(False, 1, 900))): + p = os.path.join(tmp, "t%d.json" % i) + json.dump(doc, open(p, "w")) + paths.append(p) + merged = ks.merge_process_traces(paths) + verdict = ks.verify(TARGET, KERNEL, self.meta(), merged) + self.assertEqual(verdict["matched_kernel_calls"], 1, verdict) + + def test_0802_outer_wrapper_fails_when_0720_launcher_is_marked(self): + outer = "vllm.v1.attention.layer:unified_attention_with_output" + inner = TARGET + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + outer, + "ph": "X", "pid": 1, "tid": 2, "ts": 80, "dur": 1}, + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + inner, + "ph": "X", "pid": 1, "tid": 2, "ts": 82, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + outer, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 100}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + inner, + "ph": "X", "pid": 1, "tid": 2, "ts": 120, "dur": 50}, + {"cat": "cpu_op", "name": "launch", "ph": "X", "pid": 1, "tid": 2, + "ts": 130, "dur": 5, "args": {"External id": 9}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "ts": 300, "dur": 20, + "args": {"External id": 9}}, + ] + outer_meta = {"module": "vllm.attention", "attr": "outer", "total_calls_observed": 2} + outer_verdict = ks.verify(outer, KERNEL, outer_meta, events, [outer, inner]) + self.assertFalse(outer_verdict["ok"]) + self.assertIn("deeper_live_candidate_exists", outer_verdict["failed"]) + self.assertEqual(outer_verdict["deeper_live_candidates"], [inner]) + + inner_verdict = ks.verify(inner, KERNEL, self.meta(), events, [outer, inner]) + self.assertTrue(inner_verdict["ok"], inner_verdict) + self.assertTrue(inner_verdict["deepest_verified"]) + + def test_every_declared_probe_candidate_must_have_an_installed_marker(self): + missing = "vllm.attention:missing_candidate" + verdict = ks.verify(TARGET, KERNEL, self.meta(), trace(), [TARGET, missing]) + self.assertFalse(verdict["ok"]) + self.assertIn("candidate_marker_not_installed", verdict["failed"]) + self.assertEqual(verdict["missing_candidate_markers"], [missing]) + + def test_installed_but_inactive_alternative_branch_does_not_fail(self): + alternative = "vllm.attention:prefill_only" + events = trace() + events.insert(1, { + "cat": "cpu_op", "name": ks.INSTALL_PREFIX + alternative, + "ph": "X", "pid": 1, "tid": 2, "ts": 92, "dur": 1, + }) + verdict = ks.verify( + TARGET, KERNEL, self.meta(), events, [TARGET, alternative]) + self.assertTrue(verdict["ok"], verdict) + self.assertEqual( + sorted(verdict["candidate_targets_tested"]), + sorted([TARGET, alternative]), + ) + + def test_a_probed_candidate_that_never_fired_is_reported_not_silently_dropped(self): + # A marker rebinds one module attribute. A callable dispatched through `torch.ops.*`, or + # reached through an alias a caller imported before installation, runs with the wrapper + # bypassed and produces zero spans while being fully live. Reading that as "not deeper" + # would let a shallower seam collect `deepest_verified`, so the verdict has to say which + # candidates were probed and never observed, apart from the ones never probed at all. + invisible = "vllm.model_executor.layers.attention.attention:unified_attention_with_output" + unprobed = "vllm.attention:never_installed" + events = trace() + events.insert(1, { + "cat": "cpu_op", "name": ks.INSTALL_PREFIX + invisible, + "ph": "X", "pid": 1, "tid": 2, "ts": 92, "dur": 1, + }) + verdict = ks.verify( + TARGET, KERNEL, self.meta(), events, [TARGET, invisible, unprobed]) + self.assertEqual(verdict["installed_but_never_live_candidates"], [invisible]) + self.assertEqual(verdict["missing_candidate_markers"], [unprobed]) + self.assertNotIn(invisible, verdict["deeper_live_candidates"]) + + def test_device_projected_annotation_does_not_invert_call_nesting(self): + # The profiler re-emits each marker on the GPU timeline as `gpu_user_annotation`, where an + # OUTER seam's short device span can land inside the INNER launcher's long one. Only the + # host spans describe the real call nesting; the selected inner launcher must still pass. + outer = "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe:_fused_moe_kernel_sequence" + events = trace() + events.insert(1, {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + outer, + "ph": "X", "pid": 1, "tid": 2, "ts": 92, "dur": 1}) + events.insert(2, {"cat": "cpu_op", "name": ks.MARKER_PREFIX + outer, + "ph": "X", "pid": 1, "tid": 2, "ts": 95, "dur": 200}) + # device projections: outer's is a tiny span inside the target's long one, same pid/tid + events += [ + {"cat": "gpu_user_annotation", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 9, "tid": 9, "ts": 500, "dur": 400}, + {"cat": "gpu_user_annotation", "name": ks.MARKER_PREFIX + outer, + "ph": "X", "pid": 9, "tid": 9, "ts": 600, "dur": 5}, + ] + verdict = ks.verify(TARGET, KERNEL, self.meta(), events, [TARGET, outer]) + self.assertEqual(verdict["deeper_live_candidates"], []) + self.assertTrue(verdict["ok"], verdict) + self.assertTrue(verdict["deepest_verified"]) + + def test_capture_of_a_different_callable_is_rejected(self): + wrong = "vllm.model_executor.layers.attention.attention:outer_wrapper" + verdict = ks.verify(TARGET, KERNEL, self.meta(target=wrong), trace()) + self.assertFalse(verdict["ok"]) + self.assertIn("capture_target_mismatch", verdict["failed"]) + + def test_zero_live_calls_is_rejected(self): + verdict = ks.verify(TARGET, KERNEL, self.meta(calls=0), trace()) + self.assertFalse(verdict["ok"]) + self.assertIn("target_not_observed", verdict["failed"]) + + def test_cli_writes_the_same_machine_verdict(self): + with tempfile.TemporaryDirectory() as root: + meta_path = os.path.join(root, "meta.json") + trace_path = os.path.join(root, "trace.json") + out_path = os.path.join(root, "selection.json") + with open(meta_path, "w") as fh: + json.dump(self.meta(), fh) + with open(trace_path, "w") as fh: + json.dump({"traceEvents": trace()}, fh) + rc = ks.main([ + "--target", TARGET, + "--device-kernel", KERNEL, + "--capture-meta", meta_path, + "--torch-trace", trace_path, + "--out", out_path, + ]) + self.assertEqual(rc, 0) + with open(out_path) as fh: + self.assertTrue(json.load(fh)["ok"]) + + def test_cli_fails_closed_when_capture_pid_has_no_trace(self): + with tempfile.TemporaryDirectory() as root: + meta_path = os.path.join(root, "capture.pid-111.rank-0", "meta.json") + os.makedirs(os.path.dirname(meta_path)) + trace_path = os.path.join(root, "selection.pid-999.rank-0.json") + out_path = os.path.join(root, "selection.json") + meta = self.meta() + meta["process_id"] = 111 + with open(meta_path, "w") as fh: + json.dump(meta, fh) + with open(trace_path, "w") as fh: + json.dump({"traceEvents": trace()}, fh) + rc = ks.main([ + "--target", TARGET, + "--device-kernel", KERNEL, + "--capture-meta", meta_path, + "--torch-trace", trace_path, + "--out", out_path, + ]) + self.assertEqual(rc, 1) + with open(out_path) as fh: + verdict = json.load(fh) + self.assertFalse(verdict["ok"]) + self.assertIn("capture_process_trace_missing", verdict["failed"]) + self.assertEqual(verdict["trace_file"], "") + + def test_cli_merges_calls_before_deciding_deepest_candidate(self): + mid = "vllm.attention:mid" + inner = TARGET + with tempfile.TemporaryDirectory() as root: + meta_path = os.path.join(root, "capture.pid-111.rank-0", "meta.json") + os.makedirs(os.path.dirname(meta_path)) + meta = self.meta(target=mid) + meta["process_id"] = 111 + with open(meta_path, "w") as fh: + json.dump(meta, fh) + + common_installs = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + mid, + "ph": "X", "pid": 1, "tid": 2, "ts": 80, "dur": 1}, + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + inner, + "ph": "X", "pid": 1, "tid": 2, "ts": 82, "dur": 1}, + ] + call_one = common_installs + [ + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + mid, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 50}, + {"cat": "cpu_op", "name": "launch", "ph": "X", + "pid": 1, "tid": 2, "ts": 110, "dur": 5, + "args": {"External id": 7}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "ts": 180, "dur": 10, + "args": {"External id": 7}}, + ] + call_two = common_installs + [ + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + mid, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 80}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + inner, + "ph": "X", "pid": 1, "tid": 2, "ts": 120, "dur": 40}, + {"cat": "cpu_op", "name": "launch", "ph": "X", + "pid": 1, "tid": 2, "ts": 130, "dur": 5, + "args": {"External id": 9}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "ts": 200, "dur": 10, + "args": {"External id": 9}}, + ] + trace_paths = [] + for index, events in enumerate((call_one, call_two), 1): + path = os.path.join( + root, f"selection.pid-111.rank-0.call-{index}.json") + with open(path, "w") as fh: + json.dump({"traceEvents": events}, fh) + trace_paths.append(path) + out_path = os.path.join(root, "selection.json") + rc = ks.main([ + "--target", mid, + "--device-kernel", KERNEL, + "--capture-meta", meta_path, + "--torch-trace", *trace_paths, + "--candidate-target", mid, + "--candidate-target", inner, + "--out", out_path, + ]) + self.assertEqual(rc, 1) + with open(out_path) as fh: + verdict = json.load(fh) + self.assertIn("deeper_live_candidate_exists", verdict["failed"]) + self.assertEqual(verdict["deeper_live_candidates"], [inner]) + + def test_cli_requires_deepest_selection_on_every_capture_process(self): + mid = "vllm.attention:mid" + inner = TARGET + + def process_events(with_inner): + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + mid, + "ph": "X", "pid": 1, "tid": 2, "ts": 80, "dur": 1}, + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + inner, + "ph": "X", "pid": 1, "tid": 2, "ts": 82, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + mid, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 80}, + ] + if with_inner: + events.append({ + "cat": "cpu_op", "name": ks.MARKER_PREFIX + inner, + "ph": "X", "pid": 1, "tid": 2, "ts": 120, "dur": 40, + }) + events.extend([ + {"cat": "cpu_op", "name": "launch", "ph": "X", + "pid": 1, "tid": 2, "ts": 130, "dur": 5, + "args": {"External id": 9}}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "ts": 200, "dur": 10, + "args": {"External id": 9}}, + ]) + return events + + with tempfile.TemporaryDirectory() as root: + meta_paths, trace_paths = [], [] + for pid, with_inner in ((111, False), (222, True)): + meta_path = os.path.join( + root, f"capture.pid-{pid}.rank-0", "meta.json") + os.makedirs(os.path.dirname(meta_path)) + meta = self.meta(target=mid) + meta["process_id"] = pid + with open(meta_path, "w") as fh: + json.dump(meta, fh) + meta_paths.append(meta_path) + trace_path = os.path.join( + root, f"selection.pid-{pid}.rank-0.call-1.json") + with open(trace_path, "w") as fh: + json.dump({"traceEvents": process_events(with_inner)}, fh) + trace_paths.append(trace_path) + out_path = os.path.join(root, "selection.json") + rc = ks.main([ + "--target", mid, + "--device-kernel", KERNEL, + "--capture-meta", *meta_paths, + "--torch-trace", *trace_paths, + "--candidate-target", mid, + "--candidate-target", inner, + "--out", out_path, + ]) + self.assertEqual(rc, 1) + with open(out_path) as fh: + verdict = json.load(fh) + self.assertFalse(verdict["ok"]) + self.assertEqual(verdict["deeper_live_candidates"], [inner]) + self.assertEqual( + sorted(verdict["live_candidate_targets"]), sorted([mid, inner])) + + +class TestTheVerdictRefusesMalformedInput(unittest.TestCase): + """These are the fail-closed edges. Each one is a way an extraction could arrive incomplete and + still be read as "nothing to check here", which is the vacuous pass the contract exists to stop.""" + + def meta(self, calls=7, target=TARGET): + module, attr = target.split(":", 1) + return {"module": module, "attr": attr, "total_calls_observed": calls} + + def test_a_prose_target_is_named_as_such_rather_than_probed(self): + verdict = ks.verify("the attention wrapper", KERNEL, self.meta(), trace()) + self.assertFalse(verdict["ok"]) + self.assertIn("invalid_target_callable", verdict["failed"]) + + def test_a_head_with_no_device_symbol_cannot_certify_anything(self): + verdict = ks.verify(TARGET, "", self.meta(), trace()) + self.assertFalse(verdict["ok"]) + self.assertIn("missing_device_kernel", verdict["failed"]) + + def test_one_malformed_candidate_sinks_the_whole_probe(self): + """The coverage check compares the declared candidate set against what was probed. A member + that cannot be probed at all must fail rather than quietly shrink the set being compared.""" + verdict = ks.verify(TARGET, KERNEL, self.meta(), trace(), + candidate_targets=["live_call_seam (see notes)"]) + self.assertFalse(verdict["ok"]) + self.assertIn("invalid_candidate_target", verdict["failed"]) + + def test_an_empty_kernel_name_never_matches_by_accident(self): + self.assertFalse(ks.kernel_matches("", KERNEL)) + self.assertFalse(ks.kernel_matches(KERNEL, "")) + self.assertFalse(ks.kernel_matches("<>", KERNEL)) + + +class TestSpansAreReadFromEitherTraceShape(unittest.TestCase): + def test_begin_end_pairs_bound_a_span_the_way_a_complete_event_does(self): + """Kineto writes a marker as a complete `X` event, but a trace that was cut short (or came + from a different exporter) carries the same span as a B/E pair. Reading only `X` would report + a live seam as never having run.""" + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "B", "pid": 1, "tid": 2, "ts": 100}, + {"cat": "cpu_op", "name": "launch", "ph": "X", "pid": 1, "tid": 2, + "ts": 110, "dur": 5, "args": {"External id": 7}}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "E", "pid": 1, "tid": 2, "ts": 200}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "ts": 300, "dur": 10, + "args": {"External id": 7}}, + ] + module, attr = TARGET.split(":", 1) + verdict = ks.verify( + TARGET, KERNEL, + {"module": module, "attr": attr, "total_calls_observed": 1}, events) + self.assertTrue(verdict["ok"], verdict) + self.assertEqual(verdict["matched_kernel_calls"], 1) + + def test_an_end_with_no_begin_is_dropped_instead_of_inventing_a_span(self): + spans = ks._complete_spans([ + {"cat": "cpu_op", "name": "m", "ph": "E", "pid": 1, "tid": 2, "ts": 200}, + ], "m") + self.assertEqual(spans, []) + + def test_an_event_with_no_timestamp_is_not_inside_any_span(self): + self.assertFalse(ks._within_any_span({"cat": "cpu_op"}, [(0.0, 10.0, {})])) + + def test_a_device_row_does_not_donate_its_own_id_to_the_launch_set(self): + """Kernel rows sit inside the marker span on the device timeline too. Harvesting ids from + them would make a kernel vouch for itself, so only host events may establish causality.""" + events = [ + {"cat": "cpu_op", "name": ks.INSTALL_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 90, "dur": 1}, + {"cat": "cpu_op", "name": ks.MARKER_PREFIX + TARGET, + "ph": "X", "pid": 1, "tid": 2, "ts": 100, "dur": 100}, + {"cat": "kernel", "name": KERNEL, "ph": "X", "pid": 1, "tid": 2, + "ts": 120, "dur": 10, "args": {"External id": 7}}, + ] + module, attr = TARGET.split(":", 1) + verdict = ks.verify( + TARGET, KERNEL, + {"module": module, "attr": attr, "total_calls_observed": 1}, events) + self.assertFalse(verdict["ok"]) + self.assertIn("device_kernel_not_under_target", verdict["failed"]) + + +class TestMergingProcessTraces(unittest.TestCase): + def test_a_trace_file_carrying_non_event_entries_is_merged_without_raising(self): + """`traceEvents` from a truncated export can contain nulls/strings. The verifier reads other + people's files, so a malformed entry must be skipped rather than abort every rank's merge.""" + with tempfile.TemporaryDirectory() as root: + path = os.path.join(root, "selection.pid-1.rank-0.call-1.json") + with open(path, "w") as fh: + json.dump({"traceEvents": [ + None, + "not an event", + {"cat": "cpu_op", "name": "m", "ph": "X", "pid": 1, "tid": 2, + "ts": 1, "dur": 1, "args": {"External id": 5, "correlation": 6}}, + ]}, fh) + merged = ks.merge_process_traces([path]) + self.assertEqual(len(merged), 1) + self.assertEqual(merged[0]["pid"], "trace-0:1") + self.assertEqual(merged[0]["args"]["External id"], "trace-0:5") + self.assertEqual(merged[0]["args"]["correlation"], "trace-0:6") + + +if __name__ == "__main__": + unittest.main() diff --git a/e2e_workflow/scripts/tests/test_overlay_setup.py b/e2e_workflow/scripts/tests/test_overlay_setup.py index 6824b27bd..0eb0de9b5 100644 --- a/e2e_workflow/scripts/tests/test_overlay_setup.py +++ b/e2e_workflow/scripts/tests/test_overlay_setup.py @@ -59,7 +59,7 @@ def _load(mod_name, filename): ov = _load("overlay_setup", "overlay_setup.py") -EMPTY_MANIFEST = {"modules": [], "rebinds": [], "captures": []} +EMPTY_MANIFEST = {"modules": [], "rebinds": [], "markers": [], "captures": []} class _RecordingRun: @@ -210,6 +210,13 @@ def test_shim_calls_capture_shapes_install_as_capture_shapes_defines_it(self): self.assertIn("def install(target, out_dir, max_cases=5):", self._read(os.path.join(SCRIPTS_DIR, "capture_shapes.py"))) + def test_shim_installs_captures_before_markers(self): + # A marker on any module that does `from import ` imports that module + # and freezes the alias. If markers ran first, the later capture hook would rebind only the + # defining module and the live call (through the alias) would never be recorded. + shim = ov.SITECUSTOMIZE + self.assertLess(shim.index('_m.get("captures", [])'), shim.index('_m.get("markers", [])')) + def test_rerun_preserves_an_edited_shim_and_an_existing_manifest(self): # Re-running any add-* must not reset an overlay that already carries accepted kernels. ov._ensure_overlay(self.overlay) @@ -596,6 +603,54 @@ def test_distinct_targets_compound(self): self.assertEqual([e["target"] for e in self._manifest()["captures"]], ["m:a", "m:b"]) +class TestAddMarker(_OverlayCase): + def test_marker_targets_compound_and_copy_the_probe(self): + marker = self._write("custom_marker.py", "def install(target):\n pass\n") + self._run(ov.cmd_add_marker, self._ns(target="m:outer", marker_file=marker)) + self._run(ov.cmd_add_marker, self._ns(target="m:inner", marker_file=marker)) + self.assertEqual( + [entry["target"] for entry in self._manifest()["markers"]], + ["m:outer", "m:inner"], + ) + self.assertEqual( + self._read(os.path.join(self.overlay, "seam_trace.py")), + "def install(target):\n pass\n", + ) + + def test_readding_a_marker_is_idempotent(self): + marker = self._write("custom_marker.py", "def install(target):\n pass\n") + self._run(ov.cmd_add_marker, self._ns(target="m:inner", marker_file=marker)) + self._run(ov.cmd_add_marker, self._ns(target="m:inner", marker_file=marker)) + self.assertEqual(self._manifest()["markers"], [{"target": "m:inner"}]) + + def test_a_marker_overlay_seeds_from_the_base_like_every_other_add(self): + """Two overlay dirs on PYTHONPATH do not compound, so an overlay that was not seeded from the + live stack IS the pristine install. A probe-only overlay that skipped --from would move the + seam onto a different build of the code than the one the profile came from.""" + marker = self._write("custom_marker.py", "def install(target):\n pass\n") + base = os.path.join(self.tmp, "current_overlay") + ov._ensure_overlay(base) + self._run(ov.cmd_add_rebind, self._ns( + overlay=base, target="m:accepted", impl_module="fast", impl_attr="go", impl_file="")) + + probe = os.path.join(self.tmp, "probe_overlay") + self._run(ov.cmd_add_marker, + self._ns(overlay=probe, target="m:inner", marker_file=marker, base=base)) + + self.assertEqual([e["target"] for e in self._manifest(probe)["rebinds"]], ["m:accepted"], + "the accepted kernel from the base overlay was dropped") + self.assertEqual([e["target"] for e in self._manifest(probe)["markers"]], ["m:inner"]) + + def test_add_marker_accepts_from_on_the_command_line(self): + marker = self._write("custom_marker.py", "def install(target):\n pass\n") + base = os.path.join(self.tmp, "current_overlay") + ov._ensure_overlay(base) + probe = os.path.join(self.tmp, "probe_overlay") + self._main(["add-marker", "--overlay", probe, "--target", "m:inner", + "--marker-file", marker, "--from", base]) + self.assertEqual([e["target"] for e in self._manifest(probe)["markers"]], ["m:inner"]) + + # --------------------------------------------------------------------------- # # check -- "is the overlay actually shadowing this module?" # --------------------------------------------------------------------------- # diff --git a/e2e_workflow/scripts/tests/test_parse_profile.py b/e2e_workflow/scripts/tests/test_parse_profile.py index 153902292..f7641ec14 100644 --- a/e2e_workflow/scripts/tests/test_parse_profile.py +++ b/e2e_workflow/scripts/tests/test_parse_profile.py @@ -241,6 +241,19 @@ def test_strips_void_template_and_namespace(self): def test_truncates_to_60_chars(self): self.assertEqual(pp.short_name("k" * 70), "k" * 60) + self.assertEqual(pp.SHORT_NAME_LIMIT, 60) + + def test_an_unnamed_namespace_shortens_to_the_kernel_not_its_parameter(self): + """ROCm spells the unnamed namespace '(anonymous namespace)', so the symbol opens with '(' + where the identifier should be. The '[\\w:]+' probe found nothing, the whole signature fell + through as the "short" name, and splitting it on '::' handed back the trailing parameter + type -- this real symbol used to shorten to 'KdaPackedDecodeParams)'.""" + self.assertEqual( + pp.short_name("void (anonymous namespace)::kda_packed_decode_kernel<8, false>" + "((anonymous namespace)::KdaPackedDecodeParams)"), + "kda_packed_decode_kernel") + self.assertEqual(pp.norm_key("void (anonymous namespace)::clamp_position_kernel(int)"), + "clamppositionkernel") def test_non_identifier_name_is_returned_as_is(self): self.assertEqual(pp.short_name(""), "") @@ -253,6 +266,209 @@ def test_norm_key_joins_hw_name_to_torch_name(self): self.assertEqual(pp.norm_key("at::native::Fill_Kernel"), "fillkernel") +class TestEntityEvidence(unittest.TestCase): + def test_exact_name_host_device_collision_is_unresolved(self): + host = {"shared_name": {"calls": 2, "total_us": 5.0, + "cat_counts": {"cpu_op": 2}}} + device = {"shared_name": {"calls": 3, "total_us": 7.0, + "cat_counts": {"kernel": 3}}} + merged = pp.merge_entity_evidence(host, device) + rows, stats = pp.annotate_rows([{"name": "shared_name"}], merged, "torch-trace") + self.assertEqual(rows[0]["entity_kind"], "unresolved") + self.assertEqual(rows[0]["entity_evidence"]["basis"], + "exact_name_multiple_entity_kinds") + self.assertEqual(stats["unresolved"], 1) + + def test_same_name_device_sources_remain_a_gpu_kernel(self): + rocprof = {"kernel_name": {"calls": 2, "total_us": 5.0, + "src": "rocprof_kernel_stats"}} + torch = {"kernel_name": {"calls": 3, "total_us": 7.0, + "cat_counts": {"kernel": 3}}} + merged = pp.merge_entity_evidence(rocprof, torch) + rows, _ = pp.annotate_rows([{"name": "kernel_name"}], merged, "merged") + self.assertEqual(rows[0]["entity_kind"], "gpu_kernel") + + def test_c9_norm_key_collision_resolves_each_row_by_exact_name(self): + # C9 (reviewer's LITERAL scenario): a host dispatcher span 'aten::mul' (cpu_op) and a device + # kernel 'mul' (cat="kernel") are DIFFERENT exact names that norm_key() folds to the SAME key + # 'mul' (short_name drops the aten:: namespace, then non-alnum is stripped). The old + # setdefault(first-collision-wins) fold let whichever was aggregated first stamp its kind onto + # the other row. Exact-name resolution must win for BOTH, and a third row that only matches via + # the lossy fold must fail closed as ambiguous rather than borrow either neighbor's kind. + self.assertEqual(pp.norm_key("aten::mul"), pp.norm_key("mul")) # the collision premise + self.assertEqual(pp.norm_key("aten::Mul"), pp.norm_key("mul")) + host = {"aten::mul": {"calls": 2, "total_us": 5.0, "cat_counts": {"cpu_op": 2}}} + device = {"mul": {"calls": 3, "total_us": 7.0, "cat_counts": {"kernel": 3}}} + merged = pp.merge_entity_evidence(host, device) + rows, stats = pp.annotate_rows( + [{"name": "mul"}, {"name": "aten::mul"}, {"name": "aten::Mul"}], merged, "torch-trace") + by_name = {r["name"]: r for r in rows} + # device kernel keeps gpu_kernel — EXACT name wins over the shared fold + self.assertEqual(by_name["mul"]["entity_kind"], "gpu_kernel") + self.assertEqual(by_name["mul"]["entity_evidence"]["matched_profiled_kernel"], "mul") + # host dispatcher keeps dispatcher_op — EXACT name wins (not stamped as a kernel) + self.assertEqual(by_name["aten::mul"]["entity_kind"], "dispatcher_op") + self.assertEqual(by_name["aten::mul"]["entity_evidence"]["matched_profiled_kernel"], + "aten::mul") + # a NON-exact row whose norm_key hits the shared collision is stamped unresolved, refusing to + # guess which colliding entity it is — never silently misclassified as either kind. + self.assertEqual(by_name["aten::Mul"]["entity_kind"], "unresolved") + self.assertEqual(by_name["aten::Mul"]["entity_evidence"]["basis"], "ambiguous_name_match") + self.assertNotIn("matched_profiled_kernel", by_name["aten::Mul"]["entity_evidence"]) + self.assertEqual(stats["gpu_kernel"], 1) + self.assertEqual(stats["dispatcher_op"], 1) + self.assertEqual(stats["unresolved"], 1) + + def test_a_row_naming_nothing_the_profiler_saw_is_unresolved_not_assumed(self): + """The head track routes only gpu_kernel rows. A name the profile never contains has no + evidence at all, and defaulting it to "kernel" is exactly how an unverifiable head got in.""" + rows, stats = pp.annotate_rows( + [{"name": "a_kernel_nobody_profiled"}], + pp.merge_entity_evidence({"real_kernel": {"calls": 1, "total_us": 5.0, + "cat_counts": {"kernel": 1}}}), + "torch-trace") + self.assertEqual(rows[0]["entity_kind"], "unresolved") + self.assertEqual(rows[0]["entity_evidence"]["basis"], "name_not_found_in_profile") + self.assertEqual(stats["unresolved"], 1) + + def test_a_decorated_row_name_still_resolves_through_the_normalized_key(self): + """Top-N rows arrive with demangled/decorated names ('void ns::k(int)'). When the fold + is unambiguous that row IS the profiled kernel, and refusing it would drop a real head.""" + device = {"void ns::gemm_kernel(int)": {"calls": 3, "total_us": 9.0, + "cat_counts": {"kernel": 3}}} + rows, stats = pp.annotate_rows( + [{"name": "gemm_kernel"}], pp.merge_entity_evidence(device), "torch-trace") + self.assertEqual(rows[0]["entity_kind"], "gpu_kernel") + self.assertEqual(rows[0]["entity_evidence"]["matched_profiled_kernel"], + "void ns::gemm_kernel(int)") + self.assertEqual(stats["gpu_kernel"], 1) + + def test_c9_collision_resolution_is_merge_order_independent(self): + # The setdefault bug made the collision winner depend on which aggregate was merged first. + # Reversing the merge input order must yield IDENTICAL kinds for every row: neither the + # dispatcher nor the kernel may stamp its kind onto the other in EITHER direction. + host = {"aten::mul": {"calls": 2, "total_us": 5.0, "cat_counts": {"cpu_op": 2}}} + device = {"mul": {"calls": 3, "total_us": 7.0, "cat_counts": {"kernel": 3}}} + rowspec = [{"name": "mul"}, {"name": "aten::mul"}, {"name": "aten::Mul"}] + + def kinds(*aggs): + merged = pp.merge_entity_evidence(*aggs) + rows, _ = pp.annotate_rows([dict(r) for r in rowspec], merged, "torch-trace") + return {r["name"]: r["entity_kind"] for r in rows} + + forward = kinds(host, device) + reverse = kinds(device, host) # merge the device kernel FIRST this time + self.assertEqual(forward, reverse) + self.assertEqual(forward, {"mul": "gpu_kernel", "aten::mul": "dispatcher_op", + "aten::Mul": "unresolved"}) + + +class TestDispatcherExpansion(_TmpMixin, unittest.TestCase): + def test_external_id_edges_expand_wrapper_into_device_kernels(self): + events = [ + {"cat": "cpu_op", "name": "vllm::attention", + "args": {"External id": 17}}, + {"cat": "kernel", "name": "kernel_paged_attention_2d", + "ts": 10, "dur": 80, "args": {"External id": 17}}, + {"cat": "kernel", "name": "_fwd_kernel", + "ts": 20, "dur": 20, "args": {"External id": 17}}, + ] + edges = pp.index_dispatch_edges(self._trace_file(events)) + rows, records = pp.expand_dispatcher_rows([{ + "rank": 1, + "name": "vllm::attention", + "short_name": "attention (paged + prefill)", + "pct_gpu_time": 25.0, + "total_ms": 10.0, + "calls": 3, + "notes": "aggregate", + }], edges) + + self.assertEqual([r["short_name"] for r in rows], + ["kernel_paged_attention_2d", "_fwd_kernel"]) + self.assertEqual([r["pct_gpu_time"] for r in rows], [20.0, 5.0]) + self.assertTrue(all(r["profile_parent"] == "vllm::attention" for r in rows)) + self.assertEqual(records[0]["dispatcher"], "vllm::attention") + + def test_nested_launch_external_id_is_attributed_to_outer_dispatcher(self): + events = [ + {"cat": "cpu_op", "name": "vllm::attention", "pid": 1, "tid": 2, + "ts": 100, "dur": 100, "args": {"External id": 17}}, + {"cat": "cpu_op", "name": "triton::launch", "pid": 1, "tid": 2, + "ts": 120, "dur": 10, "args": {"External id": 18}}, + {"cat": "kernel", "name": "kernel_paged_attention_2d", + "ts": 250, "dur": 80, "args": {"External id": 18}}, + ] + edges = pp.index_dispatch_edges(self._trace_file(events)) + self.assertIn("kernel_paged_attention_2d", edges["vllm::attention"]) + self.assertIn("kernel_paged_attention_2d", edges["triton::launch"]) + + def test_unrelated_dispatcher_is_preserved_for_fail_closed_classification(self): + original = {"name": "vllm::unknown", "pct_gpu_time": 9.0} + rows, records = pp.expand_dispatcher_rows([original], {}) + self.assertEqual(rows, [original]) + self.assertEqual(records, []) + + def test_existing_device_row_is_not_duplicated_by_dispatcher_expansion(self): + rows, _ = pp.expand_dispatcher_rows([ + {"name": "vllm::attention", "pct_gpu_time": 20.0, "total_ms": 10.0}, + {"name": "kernel_paged_attention_2d", "pct_gpu_time": 20.0, "total_ms": 10.0}, + ], { + "vllm::attention": { + "kernel_paged_attention_2d": {"calls": 4, "total_us": 100.0}, + }, + }) + self.assertEqual( + [row["name"] for row in rows].count("kernel_paged_attention_2d"), 1) + self.assertEqual(sum(row["pct_gpu_time"] for row in rows), 20.0) + + def test_a_dispatcher_whose_children_cost_nothing_is_left_alone(self): + """Expansion divides the row's Amdahl mass by the children's measured time. With no measured + time there is nothing to divide by, and inventing a split would be worse than not expanding.""" + row = {"name": "vllm::attention", "pct_gpu_time": 20.0, "total_ms": 10.0} + rows, records = pp.expand_dispatcher_rows([row], { + "vllm::attention": {"kernel_paged_attention_2d": {"calls": 4, "total_us": 0.0}}, + }) + self.assertEqual(rows, [row]) + self.assertEqual(records, []) + + def test_a_trace_that_cannot_be_read_yields_no_edges_rather_than_raising(self): + """The trace path is supplied by the caller and may be absent or truncated. parse_profile is + the step that TRIAGES a run; it must degrade to "no evidence", never take the run down.""" + self.assertEqual(pp.index_dispatch_edges(self._trace_file(raw="{not json")), {}) + self.assertEqual(pp.index_dispatch_edges("/nonexistent/trace.json"), {}) + + def test_sibling_spans_do_not_inherit_each_others_kernels(self): + """Ancestry is what attributes a kernel to an outer dispatcher. A span that already CLOSED is + not an ancestor, so a launch under the second sibling must not be credited to the first.""" + events = [ + {"cat": "cpu_op", "name": "outer", "pid": 1, "tid": 2, + "ts": 100, "dur": 200, "args": {"External id": 1}}, + {"cat": "cpu_op", "name": "first_child", "pid": 1, "tid": 2, + "ts": 110, "dur": 10, "args": {"External id": 2}}, + {"cat": "cpu_op", "name": "second_child", "pid": 1, "tid": 2, + "ts": 150, "dur": 10, "args": {"External id": 3}}, + {"cat": "kernel", "name": "second_kernel", "ts": 400, "dur": 5, + "args": {"External id": 3}}, + ] + edges = pp.index_dispatch_edges(self._trace_file(events)) + self.assertIn("second_kernel", edges["second_child"]) + self.assertIn("second_kernel", edges["outer"]) + self.assertNotIn("first_child", edges) + + def test_exact_host_device_collision_is_not_expanded(self): + row = {"name": "shared", "pct_gpu_time": 20.0} + evidence = pp.merge_entity_evidence( + {"shared": {"calls": 1, "total_us": 5.0, "cat_counts": {"cpu_op": 1}}}, + {"shared": {"calls": 1, "total_us": 7.0, "cat_counts": {"kernel": 1}}}, + ) + rows, records = pp.expand_dispatcher_rows( + [row], {"shared": {"child_kernel": {"calls": 1, "total_us": 7.0}}}, + evidence, "torch-trace") + self.assertEqual(rows, [row]) + self.assertEqual(records, []) + + # --------------------------------------------------------------------------- # # serving-phase step windows # --------------------------------------------------------------------------- # @@ -627,7 +843,8 @@ def test_blank_cells_default_to_zero_duration_and_one_call(self): d = self._rocprof_dir("Name,Calls,TotalDurationNs\nqux_kernel,,\n") agg, total_us, launches = pp.parse_rocprof_dir(d) self.assertEqual(agg["qux_kernel"], {"calls": 1, "total_us": 0.0, - "shapes": set(), "dtypes": set()}) + "shapes": set(), "dtypes": set(), + "src": "rocprof_kernel_stats"}) self.assertEqual((total_us, launches), (0.0, 1)) def test_repeated_kernel_rows_accumulate(self): @@ -1061,5 +1278,130 @@ def test_module_entrypoint_invokes_main(self): self.assertIn("# Profile Top-1", out.getvalue()) +# --------------------------------------------------------------------------- # +# host-side entity index + the --annotate CLI the architect's head rows go through +# --------------------------------------------------------------------------- # +class TestHostEntityIndex(_TmpMixin, unittest.TestCase): + """Host spans are deliberately kept OUT of the timing aggregate (a cpu_op's duration subsumes the + kernels it dispatched). They are indexed separately so --annotate can say "this claimed head is a + dispatcher" instead of the much weaker "not found".""" + + def test_host_spans_are_counted_by_category_and_device_rows_are_not(self): + idx = pp.index_host_entities(self._trace_file([ + {"cat": "cpu_op", "name": "vllm::attention", "dur": 10.0}, + {"cat": "cpu_op", "name": "vllm::attention", "dur": 5.0}, + {"cat": "user_annotation", "name": "step", "dur": 100.0}, + {"cat": "hip_runtime", "name": "hipLaunchKernel", "dur": 1.0}, + {"cat": "kernel", "name": "gemm_kernel", "dur": 80.0}, + "not an event", + ])) + self.assertEqual(idx["vllm::attention"], { + "calls": 2, "total_us": 15.0, "cat_counts": {"cpu_op": 2}}) + self.assertEqual(idx["step"]["cat_counts"], {"user_annotation": 1}) + self.assertIn("hipLaunchKernel", idx) + self.assertNotIn("gemm_kernel", idx) + + def test_an_unreadable_trace_indexes_nothing_rather_than_raising(self): + self.assertEqual(pp.index_host_entities("/nonexistent/trace.json"), {}) + self.assertEqual(pp.index_host_entities(self._trace_file(raw="{not json")), {}) + + +class TestAnnotateCli(_TmpMixin, unittest.TestCase): + """`--annotate` is how a Top-N doc becomes routable head candidates. It is the step that must + refuse a row it cannot prove is a GPU kernel, and its EXIT CODE is what the workflow reads.""" + + def _main(self, argv): + out, err = io.StringIO(), io.StringIO() + saved = sys.argv + sys.argv = ["parse_profile.py"] + argv + try: + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + with self.assertRaises(SystemExit) as caught: + pp.main() + finally: + sys.argv = saved + return caught.exception.code, out.getvalue(), err.getvalue() + + def _doc(self, top_kernels, name="profile_topN.json"): + return self._write(name, json.dumps({"top_kernels": top_kernels})) + + def test_a_dispatcher_row_is_replaced_by_the_kernels_it_launched(self): + trace = self._trace_file([ + {"cat": "cpu_op", "name": "vllm::attention", "pid": 1, "tid": 2, + "ts": 100, "dur": 100, "args": {"External id": 17}}, + {"cat": "kernel", "name": "kernel_paged_attention_2d", "pid": 1, "tid": 7, + "ts": 110, "dur": 80, "args": {"External id": 17}}, + ]) + doc = self._doc([{"name": "vllm::attention", "short_name": "attention", + "pct_gpu_time": 20.0, "total_ms": 10.0, "calls": 2}]) + out_path = os.path.join(self._dir(), "annotated.json") + code, stdout, stderr = self._main( + ["--torch-trace", trace, "--annotate", doc, "--annotate-out", out_path]) + + self.assertEqual(code, 0) + with open(out_path) as fh: + annotated = json.load(fh) + self.assertEqual(annotated["entity_kind_contract"], "v1") + self.assertEqual([r["name"] for r in annotated["top_kernels"]], + ["kernel_paged_attention_2d"]) + self.assertEqual([r["entity_kind"] for r in annotated["top_kernels"]], ["gpu_kernel"]) + self.assertEqual(annotated["entity_kind_counts"], {"gpu_kernel": 1}) + self.assertEqual(annotated["dispatcher_expansions"][0]["dispatcher"], "vllm::attention") + self.assertIn("gpu_kernel=1", stderr) + self.assertEqual(json.loads(stdout)["annotated"], 1) + with open(doc) as fh: # --annotate-out leaves the input untouched + self.assertNotIn("entity_kind_contract", json.load(fh)) + + def test_a_row_that_is_not_a_profiled_kernel_exits_non_zero(self): + """The workflow branches on this exit code: a claimed head the profiler never dispatched must + stop the head track rather than be routed as if it had been measured.""" + trace = self._trace_file([ + {"cat": "kernel", "name": "gemm_kernel", "ts": 10, "dur": 80, + "args": {"External id": 1}}, + ]) + doc = self._doc([{"name": "gemm_kernel", "pct_gpu_time": 20.0}, + {"name": "live_call_seam", "pct_gpu_time": 5.0}]) + code, stdout, stderr = self._main(["--torch-trace", trace, "--annotate", doc]) + + self.assertEqual(code, 1) + with open(doc) as fh: # no --annotate-out: annotated in place + annotated = json.load(fh) + kinds = {r["name"]: r["entity_kind"] for r in annotated["top_kernels"]} + self.assertEqual(kinds, {"gemm_kernel": "gpu_kernel", "live_call_seam": "unresolved"}) + self.assertEqual(annotated["entity_kind_counts"], {"gpu_kernel": 1, "unresolved": 1}) + self.assertIn("unresolved=1", stderr) + self.assertEqual(json.loads(stdout)["entity_kind_counts"]["unresolved"], 1) + + def test_a_host_only_row_is_labelled_a_dispatcher_rather_than_reported_missing(self): + """A dispatcher that launched nothing this trace can attribute is still IDENTIFIED, which is + a successful classification -- the exit code marks rows with no evidence at all. Refusing a + dispatcher from the head track is admitHeads' job, and it needs this label to do it.""" + trace = self._trace_file([ + {"cat": "cpu_op", "name": "vllm::opaque", "pid": 1, "tid": 2, + "ts": 100, "dur": 100}, + {"cat": "kernel", "name": "gemm_kernel", "ts": 10, "dur": 80}, + ]) + doc = self._doc([{"name": "vllm::opaque", "pct_gpu_time": 20.0}]) + code, _, stderr = self._main(["--torch-trace", trace, "--annotate", doc]) + + self.assertEqual(code, 0) + with open(doc) as fh: + annotated = json.load(fh) + self.assertEqual(annotated["top_kernels"][0]["entity_kind"], "dispatcher_op") + self.assertEqual(annotated["entity_kind_counts"], {"dispatcher_op": 1}) + self.assertIn("dispatcher_op=1", stderr) + + def test_annotating_with_rocprof_evidence_needs_no_torch_trace(self): + doc = self._doc([{"name": "gemm_kernel", "pct_gpu_time": 20.0}]) + code, _, _ = self._main( + ["--rocprof-dir", self._rocprof_dir(), "--annotate", doc]) + + self.assertEqual(code, 0) + with open(doc) as fh: + annotated = json.load(fh) + self.assertEqual(annotated["top_kernels"][0]["entity_kind"], "gpu_kernel") + self.assertEqual(annotated["dispatcher_expansions"], []) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/e2e_workflow/scripts/tests/test_seam_trace.py b/e2e_workflow/scripts/tests/test_seam_trace.py new file mode 100644 index 000000000..b1a2e4201 --- /dev/null +++ b/e2e_workflow/scripts/tests/test_seam_trace.py @@ -0,0 +1,365 @@ +#!/usr/bin/env python3 +"""Tests for marker-only candidate tracing.""" + +import importlib.util +import os +import sys +import tempfile +import types +import unittest + + +SCRIPTS = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SPEC = importlib.util.spec_from_file_location("seam_trace", os.path.join(SCRIPTS, "seam_trace.py")) +st = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(st) + + +class _Context: + def __init__(self, events, name): + self.events = events + self.name = name + + def __enter__(self): + self.events.append(("enter", self.name)) + + def __exit__(self, *unused): + self.events.append(("exit", self.name)) + + +class _Profiler: + def __init__(self, events): + self.events = events + + def __enter__(self): + self.events.append(("profile", "start")) + return self + + def __exit__(self, *unused): + self.events.append(("profile", "stop")) + + def export_chrome_trace(self, path): + with open(path, "w") as fh: + fh.write("{}") + + +class TestSeamTrace(unittest.TestCase): + def setUp(self): + self.events = [] + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + # UNIQUE=0 drops only the per-call suffix, so an explicit {pid} keeps the path predictable. + os.environ["GEAK_SELECTION_TRACE"] = os.path.join(self.tmp.name, "selection.{pid}.json") + self.trace = os.path.join(self.tmp.name, f"selection.{os.getpid()}.json") + os.environ["GEAK_SELECTION_TRACE_UNIQUE"] = "0" + os.environ["GEAK_SELECTION_PROFILE_CALLS"] = "1" + self.addCleanup(os.environ.pop, "GEAK_SELECTION_TRACE", None) + self.addCleanup(os.environ.pop, "GEAK_SELECTION_TRACE_UNIQUE", None) + self.addCleanup(os.environ.pop, "GEAK_SELECTION_PROFILE_CALLS", None) + + torch = types.ModuleType("torch") + torch.profiler = types.SimpleNamespace( + ProfilerActivity=types.SimpleNamespace(CPU="cpu", CUDA="cuda"), + profile=lambda activities: _Profiler(self.events), + record_function=lambda name: _Context(self.events, name), + ) + self.saved_torch = sys.modules.get("torch") + sys.modules["torch"] = torch + self.addCleanup(self._restore_torch) + + module = types.ModuleType("_seam_trace_fixture") + module.inner = lambda value: value + 1 + module.outer = lambda value: module.inner(value) * 2 + sys.modules[module.__name__] = module + self.module = module + self.addCleanup(sys.modules.pop, module.__name__, None) + + st._INSTALLED.clear() + st._PROFILE.update(lock=st._PROFILE["lock"], active=False, done=False, + owner=None, profiler=None, active_calls=0, root_calls=0, out="", + trace_index=0, atexit_registered=False) + + def _restore_torch(self): + if self.saved_torch is None: + sys.modules.pop("torch", None) + else: + sys.modules["torch"] = self.saved_torch + + def test_first_outer_call_profiles_nested_candidate_markers_once(self): + st.install("_seam_trace_fixture:outer") + st.install("_seam_trace_fixture:inner") + self.assertEqual(self.module.outer(3), 8) + self.assertTrue(os.path.isfile(self.trace)) + names = [value for action, value in self.events if action == "enter"] + self.assertEqual(names, [ + st.INSTALL_PREFIX + "_seam_trace_fixture:inner", + st.INSTALL_PREFIX + "_seam_trace_fixture:outer", + st.MARKER_PREFIX + "_seam_trace_fixture:outer", + st.MARKER_PREFIX + "_seam_trace_fixture:inner", + ]) + self.assertEqual(self.events.count(("profile", "start")), 1) + self.assertEqual(self.events.count(("profile", "stop")), 1) + + def test_jit_protocol_callable_is_rejected_before_replacement(self): + class JitLike: + def __call__(self, value): + return value + + def run(self, value): + return value + + original = JitLike() + self.module.jit_entry = original + with self.assertRaisesRegex(RuntimeError, "cannot safely mark"): + st.install("_seam_trace_fixture:jit_entry") + self.assertIs(self.module.jit_entry, original) + + def test_each_root_call_exports_before_process_exit(self): + os.environ["GEAK_SELECTION_PROFILE_CALLS"] = "32" + st.install("_seam_trace_fixture:outer") + self.assertEqual(self.module.outer(3), 8) + self.assertTrue(os.path.isfile(self.trace)) + self.assertEqual(self.events.count(("profile", "stop")), 1) + + def test_class_method_candidate_is_marked_and_proven_installed(self): + class Runner: + def run(self, value): + return value + 1 + + self.module.Runner = Runner + st.install("_seam_trace_fixture:Runner.run") + self.assertEqual(Runner().run(3), 4) + names = [value for action, value in self.events if action == "enter"] + self.assertIn(st.INSTALL_PREFIX + "_seam_trace_fixture:Runner.run", names) + self.assertIn(st.MARKER_PREFIX + "_seam_trace_fixture:Runner.run", names) + + def test_process_local_call_traces_do_not_overwrite(self): + os.environ.pop("GEAK_SELECTION_TRACE_UNIQUE", None) + os.environ["GEAK_SELECTION_PROFILE_CALLS"] = "2" + st.install("_seam_trace_fixture:outer") + self.module.outer(1) + self.module.outer(2) + traces = sorted( + name for name in os.listdir(self.tmp.name) + if name.startswith(f"selection.{os.getpid()}.call-") and name.endswith(".json")) + self.assertEqual(len(traces), 2) + self.assertIn(".call-1.json", traces[0]) + self.assertIn(".call-2.json", traces[1]) + + def test_installing_the_same_target_twice_keeps_the_first_wrapper(self): + """A candidate may be named twice (once by the architect, once by the extractor). Re-wrapping + would nest a marker inside itself and double-count the seam's calls.""" + st.install("_seam_trace_fixture:outer") + wrapper = self.module.outer + st.install("_seam_trace_fixture:outer") + self.assertIs(self.module.outer, wrapper) + self.assertEqual(self.module.outer(3), 8) + markers = [value for action, value in self.events + if action == "enter" and value.startswith(st.MARKER_PREFIX)] + self.assertEqual(markers, [st.MARKER_PREFIX + "_seam_trace_fixture:outer"]) + + def test_a_callable_whose_signature_cannot_be_read_is_still_marked(self): + """`inspect.signature` refuses some C callables reached through a partial. The seam is still a + pure-Python object to replace, so the probe must install rather than abort the whole capture.""" + import functools + + self.module.opaque = functools.partial(range) + st.install("_seam_trace_fixture:opaque") + self.assertEqual(list(self.module.opaque(3)), [0, 1, 2]) + names = [value for action, value in self.events if action == "enter"] + self.assertIn(st.MARKER_PREFIX + "_seam_trace_fixture:opaque", names) + + def test_the_seam_still_runs_when_the_process_has_no_torch(self): + """Markers are installed from sitecustomize, which runs before the server imports torch. A + seam that raised (or silently skipped the call) there would break the server being profiled.""" + st.install("_seam_trace_fixture:outer") + sys.modules["torch"] = None # makes `import torch` raise, as it does before torch is built + self.assertEqual(self.module.outer(3), 8) + self.assertEqual([value for action, value in self.events if action == "enter"], []) + + +class TestTracePathIsProcessLocal(unittest.TestCase): + def setUp(self): + for key in ("GEAK_SELECTION_TRACE", "GEAK_SELECTION_TRACE_UNIQUE", "RANK"): + self.addCleanup(os.environ.pop, key, None) + os.environ.pop(key, None) + + def test_no_trace_env_means_no_path_and_no_profile(self): + """The marker overlay is also loaded by processes that are not the capture (a tuner, a + one-shot import check). With no destination they must not start a profiler at all.""" + self.assertEqual(st._trace_path(1), "") + self.assertFalse(st._start_profile()) + + def test_an_explicit_template_is_filled_with_pid_and_rank(self): + """A TP deployment runs one marked process per rank. Whether the operator places {pid}/{rank} + by hand or leaves it to the default suffix, two ranks must never be handed one filename.""" + os.environ["GEAK_SELECTION_TRACE"] = "/tmp/sel-{pid}-{rank}.json" + os.environ["RANK"] = "3" + self.assertEqual(st._trace_path(7), f"/tmp/sel-{os.getpid()}-3.call-7.json") + + def test_a_path_with_no_placeholders_gets_the_pid_and_rank_appended(self): + os.environ["GEAK_SELECTION_TRACE"] = "/tmp/sel.json" + os.environ["RANK"] = "3" + self.assertEqual(st._trace_path(7), f"/tmp/sel.pid-{os.getpid()}.rank-3.call-7.json") + + def test_the_unique_opt_out_drops_the_call_suffix_but_not_the_process(self): + """The opt-out is for operators who want one file per process instead of one per call. It + cannot also drop the pid: kernel_selection pairs a capture to its trace by matching + '.pid-', so a shared filename loses the pairing and every rank races on one path.""" + os.environ["GEAK_SELECTION_TRACE"] = "/tmp/sel.json" + os.environ["GEAK_SELECTION_TRACE_UNIQUE"] = "0" + os.environ["RANK"] = "3" + self.assertEqual(st._trace_path(7), f"/tmp/sel.pid-{os.getpid()}.rank-3.json") + + def test_the_unique_opt_out_still_substitutes_an_explicit_placeholder(self): + """Returning the raw template here used to leave a literal '{pid}' in the filename.""" + os.environ["GEAK_SELECTION_TRACE"] = "/tmp/sel-{pid}.json" + os.environ["GEAK_SELECTION_TRACE_UNIQUE"] = "0" + self.assertEqual(st._trace_path(7), f"/tmp/sel-{os.getpid()}.json") + + def test_the_rank_comes_from_the_launcher_env_when_it_declares_one(self): + os.environ["RANK"] = "2" + self.assertEqual(st._rank(), "2") + + def test_an_unranked_process_is_still_named_rather_than_dropped(self): + self.assertEqual(st._rank(), "unknown") + + +def _raise_record_function(name): + raise RuntimeError("record_function unavailable on this build") + + +class TestProfileLifecycleFailsSoft(unittest.TestCase): + """Every failure here happens inside the SERVER under measurement. A probe that raises would + take the capture down with it, so each path degrades to "no trace" plus a stderr line.""" + + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + os.environ["GEAK_SELECTION_TRACE"] = os.path.join(self.tmp.name, "selection.json") + os.environ["GEAK_SELECTION_TRACE_UNIQUE"] = "0" + self.addCleanup(os.environ.pop, "GEAK_SELECTION_TRACE", None) + self.addCleanup(os.environ.pop, "GEAK_SELECTION_TRACE_UNIQUE", None) + self.addCleanup(os.environ.pop, "GEAK_SELECTION_PROFILE_CALLS", None) + self.saved_torch = sys.modules.get("torch") + self.addCleanup(self._restore_torch) + st._PROFILE.update(active=False, done=False, owner=None, profiler=None, + active_calls=0, root_calls=0, out="", trace_index=0, + atexit_registered=False) + + def _restore_torch(self): + if self.saved_torch is None: + sys.modules.pop("torch", None) + else: + sys.modules["torch"] = self.saved_torch + + def _install_torch(self, profile): + torch = types.ModuleType("torch") + torch.profiler = types.SimpleNamespace( + ProfilerActivity=types.SimpleNamespace(CPU="cpu"), + profile=profile, + record_function=lambda name: _Context([], name), + ) + sys.modules["torch"] = torch + + def test_a_profiler_that_refuses_to_start_disables_further_attempts(self): + def refuse(activities): + raise RuntimeError("no profiling on this device") + + self._install_torch(refuse) + self.assertFalse(st._start_profile()) + self.assertTrue(st._PROFILE["done"]) + self.assertFalse(st._start_profile()) + + def test_a_failed_export_does_not_propagate_into_the_served_call(self): + events = [] + + class _Unexportable(_Profiler): + def export_chrome_trace(self, path): + raise OSError("disk full") + + self._install_torch(lambda activities: _Unexportable(events)) + self.assertTrue(st._start_profile()) + st._finish_profile() + self.assertFalse(st._PROFILE["active"]) + + def test_a_failure_after_entering_the_profiler_still_exits_it(self): + """The profiler is entered before the install markers are written. When that write raised, + the handler set done=True while active was still True -- and _finish_profile returns early on + done, so the profiler stayed entered, and collecting, for the life of the server.""" + exited = [] + + class _EnteredProfiler(_Profiler): + def __exit__(self, *exc): + exited.append(True) + return False + + self._install_torch(lambda activities: _EnteredProfiler([])) + sys.modules["torch"].profiler.record_function = _raise_record_function + st._INSTALLED["mod:fn"] = None + self.addCleanup(st._INSTALLED.pop, "mod:fn", None) + self.assertFalse(st._start_profile()) + self.assertEqual(exited, [True]) + self.assertFalse(st._PROFILE["active"]) + self.assertIsNone(st._PROFILE["profiler"]) + + def test_a_profiler_that_also_refuses_to_exit_does_not_take_the_call_down(self): + """Everything on this path runs inside the served call, including the cleanup.""" + class _UnstoppableProfiler(_Profiler): + def __exit__(self, *exc): + raise RuntimeError("profiler already torn down") + + self._install_torch(lambda activities: _UnstoppableProfiler([])) + sys.modules["torch"].profiler.record_function = _raise_record_function + st._INSTALLED["mod:fn"] = None + self.addCleanup(st._INSTALLED.pop, "mod:fn", None) + self.assertFalse(st._start_profile()) + self.assertFalse(st._PROFILE["active"]) + + def test_a_process_with_no_trace_destination_never_enters_a_profiler(self): + """The marker overlay is loaded by processes that are not the capture. _start_profile has to + stop at the missing destination rather than profile a run nothing will read.""" + os.environ.pop("GEAK_SELECTION_TRACE", None) + self._install_torch(lambda activities: _Profiler([])) + self.assertFalse(st._start_profile()) + self.assertFalse(st._PROFILE["done"]) + + def test_the_budget_is_rechecked_before_each_start_not_only_after_a_finish(self): + os.environ["GEAK_SELECTION_PROFILE_CALLS"] = "1" + st._PROFILE["trace_index"] = 1 + self._install_torch(lambda activities: _Profiler([])) + self.assertFalse(st._start_profile()) + self.assertTrue(st._PROFILE["done"]) + + def test_finishing_a_profile_that_never_started_is_a_no_op(self): + st._finish_profile() + self.assertEqual(st._PROFILE["trace_index"], 0) + + def test_the_call_budget_stops_the_probe_rather_than_tracing_the_whole_run(self): + """An unbounded probe on a serving process writes a trace per root call for the life of the + server. The budget is what keeps a capture window bounded.""" + os.environ["GEAK_SELECTION_PROFILE_CALLS"] = "1" + self._install_torch(lambda activities: _Profiler([])) + self.assertTrue(st._start_profile()) + st._finish_profile() + self.assertTrue(st._PROFILE["done"]) + self.assertFalse(st._start_profile()) + + +class TestWrappableRefusesWhatItCannotPreserve(unittest.TestCase): + def test_a_builtin_is_refused_because_a_python_wrapper_changes_its_identity(self): + self.assertFalse(st._wrappable(len)) + + def test_a_plain_callable_object_is_accepted(self): + class Seam: + def __call__(self): + return 1 + + self.assertTrue(st._wrappable(Seam())) + + def test_a_non_callable_attribute_is_refused(self): + self.assertFalse(st._wrappable(object())) + + +if __name__ == "__main__": + unittest.main()