feat(k3): kernel mining + capsule substrate — vendored vLLM KDA/top-k cubins (+7%/+20% decode)

[xiaguan]

What

Two layers of work on one branch, tooling first, then the first production use of it:

Kernel mining tooling

• --dump-graph-png PATH now also writes a machine-readable PATH.json sibling (per-kernel symbol, grid/block, function attrs, every staged param classified against the allocation map).
• tools/kernel-capture/: provider-agnostic CUPTI injection lib (CUDA_INJECTION64_PATH) that lifts every launched kernel's cubin + full call ABI out of any CUDA process (vLLM/sglang/ourselves). Multi-rank-hardened: per-pid output subdirs, and the module/ABI tables are mutex-locked — the unlocked g_abi realloc segfaulted our 4-ranks-in-one-process K3 engine on concurrent lazy module load.

Capsule substrate, live for K3 (PEGAINFER_K3_CAPSULE)

• The serving kernel becomes an external cubin artifact; the framework keeps the op contract and the native reference twin. Two vendored vLLM v0.28.0 kernels (pegainfer-kernels/cubin/k3/, Apache-2.0, image digest + sha256 + offline-build recipe in its README), embedded at build time, bound in csrc/k3/k3_capsule.cu via cuModuleLoadData + a fail-closed cuFuncGetParamInfo ABI walk. No vLLM source enters the build; build.rs compiles nothing new for them.
• topk: single_group_topk warp kernel replaces the native router top-k at the step call site.
• kda: the fused KDA decode kernel replaces conv_silu×3 + kda_core, collapsing the four wbig band GEMMs into one full GEMM + two landings. Same formulas/constants as our TileLang generator (lower bound −5, scale 128^-0.5, eps 1e-5, tap order, [head, v, k] state) — zero new weights. It updates conv + recurrent state in place: capsule mode adds a packed conv slab filled by adopt_row at the prefill→decode handover, pins recurrent parity slab 0, and refuses to combine with dspark / CP prefill.
• Flag unset ⇒ byte-identical native serving (default).

Production invariant

K3 EP4 decode step composition: the KDA chain and router top-k are served by vendored external cubins when PEGAINFER_K3_CAPSULE selects them; everything else (MegaMoE, MLA, norms, GEMMs) unchanged. Not duplicated work: no existing capsule/loader path served K3, and the native kernels remain the reference twin.

Evidence

• Per-kernel numeric gates (GPU, run on GB300 sm_103): k3_capsule_topk_gate — expert sets equal, weights ≤1e-5, b ∈ {1,3,8,32}; k3_capsule_kda_gate — conv windows bitwise-equal, state ≤1.6% / out ≤6.3% rel err (bf16-rounding-chain difference: native lands intermediates in bf16 where vLLM stays f32).
• Production E2E A/B (pruned-224 checkpoint, EP4, tray03, 128-token completions, same serve config both arms):
  • 4-concurrent: 96.3 → 103.0 tok/s (+7%)
  • 32-concurrent: 495 → 594 tok/s (+20%)
  • Greedy 100-token texts: 4/8 byte-identical, rest diverge at a near-tie token with comparable quality.
• Clean vLLM v0.28.0 baseline on the same phases (DP4×EP4, cudagraphs on): 139 / 724–788 tok/s — still ~25% ahead; the remaining gap is dominated by graphs-vs-forced-eager (~4.3k launches/rank-step) plus the unported skinny GEMM, both recorded as next levers in docs/models/k3/vllm-kernel-ab.md.
• Workspace lib tests 22/22 green (includes fixing a stale chunk-bucket assertion that perf(k3): raise the MegaMoE protocol max to 16896 tokens per rank #962 left behind); cargo fmt --check clean.

Not ported (documented in the A/B doc)

CuTe skinny GEMM/dotprod (captured instantiations only cover M ≤ 2; highest-complexity ABI), attn-res (NB=3 online kernel vs our 8-block walk — structural mismatch for ≤2% of step), MoE/MLA (ours wins / negligible). Mined cubins + ABI manifest stay in the tray-local capture archive.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PvpKoX2Wh868p9UqpqNeFP

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 39abd17f74

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Validate the capsule router against the expert count

For the supported 896-expert K3 checkpoint, selecting PEGAINFER_K3_CAPSULE=topk or the advertised all path reaches this call with experts == 896, but k3_capsule_router_topk_launch rejects any count above 512. Under EP this forward error is passed to ep_fatal, so the fully loaded service terminates on its first request. Validate the flag against the checkpoint during startup or retain the native router for the 896-expert model.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Require a model eval before enabling capsule KDA

When PEGAINFER_K3_CAPSULE=kda or all is set, this exposes a serving path whose retained evidence reports capsule output error up to 6.3% and divergent greedy text in 4/8 prompts, but the change includes no model-level accuracy evaluation. A synthetic operator tolerance and a handful of qualitative generations cannot establish that the changed recurrent-state rounding preserves model quality, so keep this path experimental/non-serving or add a representative model eval gate before exposing it.

AGENTS.md reference: AGENTS.md:L110-L112

Useful? React with 👍 / 👎.