diff --git a/docker/Dockerfile.dev b/docker/Dockerfile.dev index c1fd4877d..49a532d96 100644 --- a/docker/Dockerfile.dev +++ b/docker/Dockerfile.dev @@ -25,6 +25,7 @@ RUN apt-get update \ cmake \ curl \ git \ + graphviz \ libclang-dev \ libibverbs-dev \ libnccl-dev \ diff --git a/docs/index.md b/docs/index.md index 63a5950b1..00b370886 100644 --- a/docs/index.md +++ b/docs/index.md @@ -102,6 +102,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/k3/mtp-dspark.md` | K3 speculative decoding via RadixArk's DSpark drafter (DFlash block-diffusion, block 7, not a classic MTP head) — live end-to-end: `--dflash-draft-model-path` arms a per-rank draft lane, a packed verify step with deferred-commit KDA replay commits blocks, full-depth EP4 serve accepts at reference rates (3.3/round cycle, 3.13 prose — matching same-checkpoint sglang) after the Markov row off-by-one fix. Six `spec_verify` gates certify what is exact (verify ≠ plain decode bitwise by construction). | | `models/k3/serving-roadmap.md` | The K3 serving deliverable: a mix engine (no P/D split) — steady TP1×DP×EP with always-on spec decode, plus long prefills as elastic CP-gang lanes inside the fixed EP16 superstep (`cp-lane-design.md`); TP is out entirely. Multi-node EP and spec decode landed 2026-08; side list: varlen prefill packing (priority raised), mega world >4224, full-depth TTFT baseline, real sampling, kv-store reuse. | | `models/k3/mix-engine-design.md` | The mix-engine shape convergence (2026-08-24): span = committed prefix + speculative tail as the one step primitive; two step forms (steady = reserved verify packs + deferred short-prefill filler, whale = 12–16k chunk at a duty cycle); MLA latent-FMHA vs dense folded into the whale bit. Whale parallelism section superseded by `cp-lane-design.md` (transient CP4×TP4 kept only as the A/B control arm); same-backend KDA measurement (KCP4 −2.70% vs TP4 at T100k) recorded here. | +| `models/k3/vllm-kernel-ab.md` | CUPTI-mined decode kernel A/B, PegaInfer EP4 vs vLLM TP4+EP on the pruned checkpoint (bs 1/8/32 ladder, GB300): our MoE/collective structure wins (MegaMoE 3 launches/layer vs ~7, zero collectives vs per-layer TP4 allreduce), but vLLM escapes the cuBLASLt B=1 splitK cliff with its own CuTe-DSL skinny GEMMs (our nvjet dispatch is bit-identical across bs) and fuses KDA decode to one kernel/layer. Cross-referenced with the EP4 decode profile, ranked ports: skinny GEMM → E=224 routing kernel → KDA/attn-res fusion; cubins+ABI captured for all. KDA + top-k now ported as capsule cubins (`PEGAINFER_K3_CAPSULE`): +7%/+20% decode throughput at 4/32-concurrent, gated per-kernel and E2E. | | `models/k3/cp-lane-design.md` | Decided plan (2026-08-24): topology permanently fixed at TP1×attn-DP16×EP16; CP is a per-sequence elastic lane — whale = CP2/CP4 gang (BS=1) coexisting with local lanes in one EP16 superstep, MoE full-width for free; routing splits into extend_len (CP prefill degree) vs total_context_len (distributed-context MLA decode over striped KV); KDA CP = affine-summary KCP, contiguous (uneven) split first, zigzag deferred; no TP baseline (vLLM/sglang serve as external baseline), no EP8 shape; phases: M0 correctness + M0.5 serving integration DONE 2026-08-24 (PR #957: gang = free-running leveling loop; pruned@EP4 16k e2e CP4 1161 ms edges vLLM TP4 1181; full-model 16-GPU round: EP16 CP4 1072 ms @16k, CP4/CP1 2.86x, wins vLLM TP16-MNNVL 1.3-1.5x at 1-2k, loses 0.68x at 8k+ to their 16-way split — the M1 width case) → M1 EP16 crossover matrix + multi-superstep/CP8+ → M2 agent cache loop → M3 elastic scheduler. | ## models / deepseek-v2-lite @@ -195,6 +196,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `subsystems/kernels/kernel-op-reports.md` | Qwen3 kernel/report tooling is feature-gated: `qwen3_kernel_report` covers per-op kernel reports, and `qwen3_model_report` emits runtime-traced eager-DAG decode operator rollups with TensorSpec `KernelCall`s, latency stats, tables, and Graphviz DOT; measured FA2 `CTA_TILE_Q=64` prefill default in place. | | `subsystems/kernels/typed-forward-pipeline.md` | Reusable typed tensor pipeline macro in `pegainfer-kernels` so model crates can express common `typed_ops` chains without model-specific wrapper macros. | | `subsystems/kernels/tvm-ffi-mvp.md` | Optional `tvm-ffi-triton-cubin` bridge in `pegainfer-kernels` plus a packed TVM wrapper for the Qwen3.5 GDR solve Triton AOT CUBIN launcher. | +| `subsystems/kernels/kernel-mining.md` | Kernel mining/capsule substrate: `--dump-graph-png` now emits a machine `.json` (per-kernel ABI + staged params), and `tools/kernel-capture/` is a provider-agnostic CUPTI injection lib that lifts every cubin + call ABI out of vLLM/sglang. JSON-level A/B of Qwen3-4B bs=1 decode located the levers (attention two-pass vs vLLM's fused `fmhaSm100fKernel`; GEMM tile dispatch) and cleared norm/rope/silu/kv-append. Storage/license discipline follows the existing `cubin/glm52/` precedent. | ## playbooks diff --git a/docs/models/k3/vllm-kernel-ab.md b/docs/models/k3/vllm-kernel-ab.md new file mode 100644 index 000000000..5a6bb369c --- /dev/null +++ b/docs/models/k3/vllm-kernel-ab.md @@ -0,0 +1,182 @@ +# K3 decode kernel A/B: PegaInfer EP4 vs vLLM TP4+EP (pruned 224-expert, GB300) + +> **TL;DR:** Mined both engines' decode kernel streams on the same tray/checkpoint +> (`/mnt/shared/weights/kimi-k3-pruned-75pct`, bs ∈ {1,8,32}, CUPTI capture, eager +> both sides), then **ported the two biggest levers as capsule cubins** +> (`PEGAINFER_K3_CAPSULE=all`): vLLM v0.28.0's fused KDA decode kernel and its +> `single_group_topk` router selection, loaded from vendored cubins with a +> fail-closed ABI check — zero vLLM source in our build. E2E on the same serve: +> **+7% decode throughput at 4-concurrent (96.3 → 103.0 tok/s), +20% at +> 32-concurrent (495 → 594 tok/s)**, greedy text 4/8 byte-identical and 4/8 +> diverging at a near-tie token with comparable quality (rounding-chain +> difference, per-kernel gates bound it). Remaining mined-but-not-ported: CuTe +> skinny GEMM (only covers M≤2; CuTe param-staging complexity), attn-res +> (structure mismatch — their NB=3 online kernel vs our 8-block walk — for ≤2% +> of step). MoE stays ours (MegaMoE more fused, zero collectives). +> +> **Last touched:** 2026-08 + +## Setup + +- Same tray (tray03, 4xGB300 sm_103), same checkpoint (224-expert MXFP4 pruned dev, + isomorphic per-rank to full 896-expert @EP16). +- vLLM: `vllm/vllm-openai:kimi-k3` image (0.1.dev19262, native kimi_k3 support), + `--tensor-parallel-size 4 --enable-expert-parallel --enforce-eager`, + FLASHINFER_MLA backend, `VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION=1`. Note the + parallel shape differs by design: vLLM shards attention TP4 (every rank sees all + requests, M = full batch), we run attn-DP (requests round-robined, per-rank M = + bs/4). +- PegaInfer: `--features k3` release, `--k3-ep-size 4`, single process hosting all + 4 ranks; `cuda_graph=false` in this EP configuration, so every launch is eager + and CUPTI-visible, symmetric with the vLLM side. +- Capture: `tools/kernel-capture/` injection lib on both engines (same JSONL + schema). Drive: 1/8/32 concurrent completions x 128 tokens, diverse prose + prompts; phase boundaries recorded as `launches.jsonl` line offsets, census + taken over the 60–95% tail of each phase (steady decode, ~45 steps). +- Artifacts (local disk, not committed): vLLM + `/data/susun/kernel-capture/k3-pruned-ep4/pid*/` (per-rank, 4 processes), + PegaInfer `/data/susun/pegainfer-dev-cache/target/nvidia_cuda_13.2.0-devel-ubuntu24.04/kernel-capture-selfrun/pid1/` + (one process = all ranks), plus `phase_census.json` / `pega_phase_census.json` + under `/data/susun/kernel-capture/`. + +## Headline: launches per rank-step (steady decode) + +| | PegaInfer EP4 | vLLM TP4+EP | +|---|---|---| +| launches per rank-step | ~4.3k | ~2.2k | +| top kernel | `cublasLt::splitKreduce` (17%) | rank-local allreduce (12%) | +| distinct symbols in window | 57–59 | 50–53 | + +## Per-op dispatch table + +| op | PegaInfer | vLLM | verdict | +|---|---|---|---| +| KDA decode (69 layers) | `k3_kda_core` + 3x `k3_conv_silu` + `k3_land_nt256` x2 + rms + nvjet projections per layer | **one** `kda_decode_fusion_many_heads_kernel` per layer (exactly 69/step) | port lever (rank 3) | +| attn-res (24 MLA layers x 8 nb-blocks) | `k3_attnres_scores` + `k3_attnres_mix` per nb-block = **16 launches/layer** | `sm100::fwd_prod_v2::attn_res_fwd_online_v2_kernel` (~7/layer incl. aux) | port lever (rank 4) | +| MLA decode | one `mla_paged_absorbed_attn_kernel`/layer (absorbed, single-pass) | `fusedKimiK3MLADecodeQConcatKVCacheKernel` + CuTe-DSL Blackwell MLA split-kv + occasional reduction | comparable; needs timing (and long-ctx split-kv check) | +| MoE (92 layers) | `mega_quant_x` + `mega_write_routing` + **one fused** `deep_gemm::sm100_fp8_fp4_mega_moe` = 3/layer | MXFP8 quantize + routing (2–3) + 2x MXFP4 `bmm_t128x8x512` + finalize ≈ 7/layer | **we are more fused; keep** | +| collectives | none (attn-DP, MegaMoE pairs ranks over NVLink) | `vllm::cross_device_reduce_1stage` ~278/step (TP4 attention) | structural win, keep | +| dense GEMM | same 5 narrow nvjet tiles (`tss_64x8`, `tss_32x64`, `tss_128x8`) with **bit-identical grid/smem at bs1/8/32**; `splitKreduce` = #1 launched kernel every phase | bs1: own CuTe-DSL skinny GEMM/dotprod family; bs8: splitK + `tst_64x8`; bs32: `tst_64x32`/`tst_128x16` 2-CTA fat tiles, splitK mostly gone | **port lever #1** (skinny GEMM at M=1; M-aware re-pick later) | +| MoE routing top-k | mis-tuned `[1,224]` top-16 (5% of step per profile) | `moe::dev::routing::routingIndices{Block,DynBlock,Cluster}Kernel`, picked by bs, same E=224 shape | **port lever #2** | + +## vLLM's bs-dependent dispatch (the "if bs > N" question, answered) + +Batch-invariant on their side: KDA fused decode, MLA kernels, MXFP4 MoE bmm tile +(`t128x8x512` at every bs — per-expert M stays tiny with 224 experts). What +switches with bs: + +- **Linear-layer GEMM provider**: bs1 = vLLM's own CuTe-DSL skinny/dotprod kernels + (8 variants, latency-tuned for M=1 — the same "vLLM wins bs=1" suspect class as + the Qwen study); bs8 = CuTe splitK + narrow nvjet; bs32 = pure nvjet fat tiles. +- **MoE routing kernel**: `routingIndicesBlockKernel` (bs1) → + `DynBlockKernel` (bs8) → `ClusterKernel` (bs32), `BlockScoresKernel` joins at 8+. +- Triton fused-MoE configs (not active for K3 but shipped) are keyed per + M ∈ {1,2,4,...,4096}; cudagraph capture ladder is `[1,2,4]+range(8,256,8)+...` + — so under graphs a single injected startup enumerates every bs bucket. + +Mining discipline that follows: capture per bs bucket; manifest rows bind per +`(op, bucket)` — the axis the capsule catalog already planned. + +## Port levers, ranked + +Ranking is **timing-informed**: `benchmarks/k3-ep4-decode-profile.md` already +measured the same EP4 decode step (~50 ms) as 52% backbone B=1 dense GEMM at +~50% of the bandwidth floor, 20% MegaMoE, 12% TileLang glue, 7% KDA core, 5% +router top-k — and explicitly *not* launch-bound. The census above tells us what +vLLM does differently at each of those slots; the profile tells us which slots +pay. (License classes per kernel-mining.md: vLLM-tree CUDA and CuTe-DSL / +`moe::dev` TRT-LLM-family kernels are Apache-2.0 and committable; `nvjet_*` is +proprietary — steal the decision, never the cubin.) + +1. **B=1 dense GEMM (52% of step, ~10 ms recoverable).** vLLM solved exactly + this cliff by *leaving cuBLASLt* at M=1: their CuTe-DSL skinny-GEMM/dotprod + family is the mined counterpart of the near-SOL B=1 GEMV the profile calls + for. Cubins + full ABI are in the capture; shapes match our checkpoint. + First candidate for the capsule loader path. +2. **Router top-k (5%, ~2.6 ms).** Our `[1,224]` top-16 call is mis-tuned; + vLLM's `moe::dev::routing::routingIndicesBlockKernel` for the *same E=224 + shape* is captured, per-bs variants included (Block → DynBlock → Cluster). +3. **KDA-layer fusion (attacks the 7% core + a slice of the 12% glue).** + `kda_decode_fusion_many_heads_kernel` collapses our core + 3 convs + + land/rms chain into one launch per layer, 69 of 93 layers. +4. **attn-res fusion (rest of the glue).** `attn_res_fwd_online_v2_kernel` + replaces our 16-launch scores/mix nb-walk. Worth a few ms at most; take it + only if (3) is already being ported from the same source area. +5. Not worth porting now: MLA (0.7% at short ctx — revisit at long context with + split-kv), MoE (we are already more fused than vLLM), M-aware nvjet re-pick + for larger buckets (real, but attn-DP keeps per-rank M small; matters only + when per-rank batch grows). + +## Port log (2026-08-28): capsule cubins for KDA + router top-k + +Shipped as the **capsule substrate**: the serving kernel is an external cubin +artifact (`pegainfer-kernels/cubin/k3/`, provenance + sha256 in its README), +embedded at build time and bound in `csrc/k3/k3_capsule.cu` via +`cuModuleLoadData` + a fail-closed `cuFuncGetParamInfo` walk against the ABI +recorded at capture. No vLLM source enters the build; the native kernels stay +as the reference twin and `PEGAINFER_K3_CAPSULE` (unset ⇒ byte-identical +serving; `all` or csv of `topk,kda`) flips ops per launch site. + +- **Router top-k** (`single_group_topk_warp_kernel`, + offline single-instantiation build, 54 KB): drop-in at the `step.rs` router + call; weights come back in descending-score order vs our selection order — + consumers treat the pairs as unordered. Needs the host-scalar routed scale + (`rs_host` beside the device `rs`). Gate `k3_capsule_topk_gate`: expert sets + equal, weights ≤1e-5, b ∈ {1,3,8,32}. +- **KDA decode** (`kda_decode_fusion_many_heads_kernel`, h96 static-layout + head-grid variant, 42 KB): one launch replaces conv_silu×3 + kda_core, and + the four projection GEMMs collapse to one full `wbig` GEMM + two landings + (`out_gate`, new packed `q|k|v` land config). Formulas are our exact + spellings (`GATE_LOWER_BOUND=-5`, `scale=128^-0.5`, `RMS_EPS=1e-5`, same tap + order, same `[head, v, k]` state layout) — the captured scalars decode to + precisely our generator constants, and **no new weights are needed** + (`cw_*`, `dt_bias`, `a_log`, `gamma_o` are layout-identical). Two structural + deltas: (1) the kernel's conv-tap stride is compiled as `3*12288`, so + capsule mode allocates a packed `[rows, 3 taps, q|k|v, 12288]` slab filled + by `adopt_row` at the prefill→decode handover (the only state-flow boundary; + continuations re-prefill from scratch); (2) conv + recurrent state update + **in place**, so capsule decode pins recurrent parity slab 0 and skips the + ping-pong — dspark and CP refuse to arm with the flag set. Gate + `k3_capsule_kda_gate`: conv windows bitwise-equal, state ≤1.6% / out ≤6.3% + rel err (bf16-rounding chains; native lands intermediates in bf16, vLLM + keeps f32). + +**E2E A/B** (same tray, same serve config, 128-token completions, diverse +prose): 4-concurrent 96.3 → 103.0 tok/s (+7%), 32-concurrent 495 → 594 tok/s +(+20%). Greedy 100-token texts: 4/8 byte-identical, 4/8 diverge at a +near-tie token and continue at comparable quality. Native path with the flag +unset is untouched (same launch sequence, byte-identical). + +vLLM v0.28.0 itself on the same phases (DP4×EP4, production defaults — +cudagraphs ON, `--max-num-seqs 64`, no CUPTI): 4-concurrent 139 tok/s +(TPOT ≈ 28.8 ms), 32-concurrent 724–788 tok/s (TPOT ≈ 41–44 ms). Against +our capsule serve (38.8 / 53.9 ms) it is ~25% ahead at both points — but it +replays graphs while our EP4 path is forced eager at ~4.3k launches per +rank-step, so a large slice of the remaining gap is launch overhead, not +kernel time; graphs over the EP4 fused path (serving-roadmap item) and the +skinny-GEMM lever are the two remaining structural differences. Raw phases in +`/data/susun/kernel-capture/capsule-ab-2026-08-28/`. + +Offline-build recipe: standalone TUs (upstream `.cu` cut above the torch +host-launcher block, `namespace { ... }` reopened around an explicit template +instantiation of the captured flag set) live at +`/data/susun/kernel-capture/offline-build/{kda_tu.cu,topk_tu.cu}` (tray03) +next to the captured ABI manifest `port_abi.json`. + +**Not ported, and why:** skinny GEMM/dotprod — captured cubins only +instantiate M ∈ {1,2} (M is a CuTe compile-time), the A/B buckets (4/32) never +hit them, and staging CuTe's 24-byte tensor-view params is the highest-effort +ABI in the set; attn-res — vLLM's online kernel is templated NB=3 vs our +8-block snapshot walk (structural mismatch, not a slot-for-slot swap) for ≤2% +of step. Both stay mined in the capture with full ABI if the calculus changes. + +## Caveats + +- The census is launch identity/ABI, not time; the profile above supplies time + for our side. A per-op nsys A/B against vLLM's kernels (especially skinny + GEMM vs our nvjet-splitK at M=1) is still the gate before any port lands. +- Contexts here are short (~11+128 tokens); split-kv and long-ctx MLA behavior + unmeasured. +- Side observation, not chased: pruned-checkpoint greedy text from our EP4 serve + degenerated quickly ("the gryl of the gryl of...") on a prose prompt; worth a + spec_verify/accuracy pass someday, unrelated to kernel structure. diff --git a/docs/subsystems/kernels/kernel-mining.md b/docs/subsystems/kernels/kernel-mining.md new file mode 100644 index 000000000..4eb238c61 --- /dev/null +++ b/docs/subsystems/kernels/kernel-mining.md @@ -0,0 +1,74 @@ +# Kernel Mining & Capsule Substrate + +> **TL;DR:** Lift a launched kernel's cubin + full call ABI out of any external CUDA engine (vLLM/sglang) and out of our own decode graph, so a per-model manifest (`qwen34b.toml`, planned) can bind a stolen or self-authored cubin to a logical op and replay it. Two capture sides are built and GB300-verified: our CUDA-graph JSON dumper (`--dump-graph-png` now writes a machine `.json`) and a provider-agnostic CUPTI injection lib (`tools/kernel-capture/`). JSON-level A/B of PegaInfer vs vLLM Qwen3-4B bs=1 decode already located the levers — attention kernel and GEMM tile dispatch — and cleared norm/rope/silu/kv-append as structurally equivalent. The lib is multi-rank-hardened (per-pid output subdirs; module/ABI tables now locked — the unlocked `g_abi` realloc segfaulted PegaInfer's 4-ranks-in-one-process K3 engine on first concurrent lazy module load) and has run a full bs-ladder A/B on K3 EP4 both sides: `docs/models/k3/vllm-kernel-ab.md`. The capsule substrate itself is now **live for K3** (`pegainfer-kernels/cubin/k3/` + `csrc/k3/k3_capsule.cu`): vendored vLLM v0.28.0 cubins for fused KDA decode and router top-k load via `cuModuleLoadData` + fail-closed `cuFuncGetParamInfo` ABI check, behind `PEGAINFER_K3_CAPSULE` with per-op numeric gates — +7%/+20% decode throughput at 4/32-concurrent on the pruned EP4 serve (`docs/models/k3/vllm-kernel-ab.md` port log). **Next: nsys-quantify the attention two-pass vs one-pass gap before porting `fmhaSm100fKernel` (qwen3), and fold the capsule loader into the manifest design.** +> +> **Last touched:** 2026-08 + +## Why + +`docs/roadmap/direction.md` already names the "kernel ledger": for each kernel, its supported shapes/SM/dtype, its measured cost, and where it sits in a model's DAG, maintained machine-readably. This work is the capture layer under that ledger, driven by a concrete question: **when a bench shows vLLM winning bs=1 decode, how do we mine the responsible kernel out fast?** + +Design intent (settled in discussion, not yet built as a runtime): +- The framework does **not** maintain serving-path kernels. It maintains a per-op *contract* (named ports, dtype/layout/axes, attrs — today's `pegainfer-core::ops::call_spec` shapes are exactly this) plus a *reference* twin for the accuracy gate. Kernels are external artifacts. +- A per-`(model, sm)` manifest (`qwen34b.toml`) lists the handful of ops the model touches; each row resolves to `builtin:` (only cuBLASLt + NCCL), `cubin:#entry`, or `native:` during migration. Startup expands the op list, resolves every row, refuses to serve on a miss. +- Loader does three static checks, fail-closed at launch: TOML↔op-schema (named ports typed), TOML↔cubin (`cuFuncGetParamInfo` param count/size/offset vs the declared args), and guard consistency (no bare `const:` that duplicates a guard dim). +- The semantic **oracle stays engine-owned** — never pushed to the kernel provider. Shape can't define `eps` placement, accumulation dtype, or rounding; only an executable reference can. Providers may ship measurements as evidence, but admission is decided by our per-op numeric gate + the model-level `hf_golden_gate`. +- Performance A/B is incumbent-vs-candidate within a `(op, bucket)` catalog slot; the first candidate races the simulator/roofline expectation. + +## Two capture sides (both built, both GB300 sm_103 verified) + +### Our side — CUDA-graph JSON dump + +`pegainfer-core/src/cuda_graph/dump.rs`. `--dump-graph-png PATH` now also writes `PATH.json` (`schema: pegainfer-cuda-graph-dump/v1`) beside the PNG/DOT. Per kernel node: symbol, demangled name, grid/block, dynamic smem, six function attributes (regs, static/const/local bytes, ptx/binary version), and every staged parameter via `cuFuncGetParamInfo` — 8-byte values resolved against the allocation map (`cuPointerGetAttribute`) into device/host ranges. Driver floor moved 12.3 → 12.4. + +Verified: Qwen3-4B bs=1 decode graph = 543 kernel nodes, 2785 params, 1086 device pointers classified. Only cuBLASLt `nvjet` kernels launch with packed `extra` buffers and report `params: null`. + +This is the source for the `qwen34b.toml` first draft (dump our own graph, tag everything `native:`, migrate rows to `cubin:` one at a time — each migration deletes a kernel we maintain). + +### Mining side — CUPTI injection lib + +`tools/kernel-capture/capture.c` (+ `build.sh`). Loaded into any CUDA process via `CUDA_INJECTION64_PATH`; provider-agnostic, no filesystem archaeology: + +```bash +CUDA_INJECTION64_PATH=.../libkernelcapture.so KERNEL_CAPTURE_DIR=out \ + python -m vllm.entrypoints.openai.api_server --model ... --enforce-eager +``` + +- `CUPTI_CBID_RESOURCE_MODULE_LOADED` → dumps each cubin from memory (`module_.cubin`). Triton, cuBLAS/cuBLASLt, CUTLASS, FlashInfer, hand-written CUDA all become module-load events with real ELF bytes. +- `cuLaunchKernel`/`cuLaunchKernelEx` EXIT callback → one `launches.jsonl` record per launch: symbol, grid/block/smem, attributes, staged params + pointer classification. Schema deliberately matches the graph dumper's, so the two sides join by symbol with no format conversion. + +**Load-bearing gotcha — runtime-launched kernels.** Kernels launched through PyTorch's `<<<>>>` (all `vllm::*` and `at::native::*` ops) hand `cuLaunchKernel` a CUfunction that `cuFuncGetParamInfo`/`cuFuncGetAttribute` reject → first cut got launch config but empty params + zero attrs for ~66% of launches (only driver-API kernels like FMHA answered). ENTER-vs-EXIT is not the cause (lazy loading is a red herring here); `CUDA_MODULE_LOADING=EAGER` is worse — it hangs vLLM startup under CUPTI and asks the target to cooperate. **Fix (`cc61e20a`):** at module-load, self-load a private copy of the cubin and `cuModuleEnumerateFunctions` to cache each kernel's layout+attrs from the driver's own parse (reentrancy-guarded against the recursive MODULE_LOADED); `record_launch` resolves from the live handle else the cache. Result on vLLM 0.26.0 Qwen3-4B bs=1: 11418/11418 launches carry attributes, 8374 full param records, 3044 cuBLASLt-nvjet `null` (correct — extra-buffer), 7467 pointers classified. `rms_norm` decodes cleanly: out/input/weight pointers tagged device, `eps=1e-6` recovered as a scalar. + +## JSON-level A/B: PegaInfer vs vLLM Qwen3-4B bs=1 decode (same GB300, bf16) + +Structurally equivalent — each engine writes its own, no order-of-magnitude gap, **not worth stealing**: + +| op | PegaInfer | vLLM | +|---|---|---| +| fused_add_rms / rms | `pegainfer::norm::FusedAddRMSNormRoundKernel` | `vllm::fused_add_rms_norm_kernel` | +| rope | `prefill_qk_norm_rope_warp_kernel` | `vllm::rotary_embedding_kernel` | +| silu | `silu_mul_kernel` | `vllm::act_and_mul_kernel` | +| kv_append | `flashinfer::AppendPagedKVCacheKernel` | `vllm::reshape_and_cache_flash_kernel` | + +Two real levers: + +1. **Attention (primary bs=1 suspect).** PegaInfer runs the FlashInfer **two-pass** split-KV path — `BatchDecodeWithPagedKVCacheKernel` (grid[64,8,1] block[16,4,2] **smem 9KB** regs56) **+ a separate `PersistentVariableLengthMergeStatesKernel`**. vLLM runs a **single fused** `fmhaSm100fKernel_...Q8Kv128PersistentSwapsAbForGen` (grid[1,8,1] block[512,1,1] **smem 143KB** regs128) — a persistent mega-kernel tuned for the generation phase (short Q, long KV, swap-AB), spending Blackwell's large smem to finish in one pass. This is the first port target: it's a C symbol, self-contained, params complete from the first launch. **Same kernel family glm52 already vendors legally** (`pegainfer-kernels/cubin/glm52/fmhaSm100fKernel_*SwapsAbForGen.cubin`). +2. **GEMM tile dispatch (steal the decision, not the cubin).** Both use cuBLASLt `nvjet`; PegaInfer selects 4 tiles (narrow `64x8`/`64x16`), vLLM selects 10, sharing 1 — vLLM's heuristic picks finer tiles per `(M=1,N,K)` (incl. large-N `128x256`/`128x192`), echoing the known cuBLAS narrow-K cliff. cuBLASLt is `builtin:`/proprietary → never a stolen cubin; the value is the shape→tile heuristic for our own selector. + +## Storage & license discipline (answered by existing precedent) + +The repo already commits FMHA cubins: `pegainfer-kernels/cubin/glm52/` holds 7 SM100 cubins, ~110–140 KB each (882 KB total), sha256-pinned in `trtllm_gen/flashInferMetaInfo.h`, loaded by the embedded loader, listed explicitly in `build.rs`, sourced+licensed in a README (FlashInfer 0.6.12, Apache-2.0). That is the pattern for any accepted capsule cubin — small, sha256-pinned, README with provider+version+license, explicit in `build.rs`. No git-lfs (repo has none; per-cubin ~100KB doesn't need it). + +Three-way license split for mined cubins: +- `vllm::*` — Apache-2.0, committable (but these are the equivalent ops, low value). +- `fmhaSm100fKernel_*` — FlashInfer/TRT-LLM Apache-2.0; **same legal path glm52 already uses** — the clean one to steal. +- `nvjet_*` (cuBLASLt) — NVIDIA proprietary, **never redistribute**; this coincides with the design decision to keep GEMM as a `builtin:` special-case rather than a stolen cubin. + +Capture artifacts (the 235 MB of 66 cubins + `launches.jsonl` per run) are mining scratch — never committed. They live on local disk under `/data/susun/`; `tools/kernel-capture/.gitignore` excludes the built `.so` and any output dir. + +## State & next + +- **Built + verified:** graph JSON dumper (`b66b534f`), CUPTI capture lib (`ff647924`), ABI-cache + pointer-classification fix (`cc61e20a`), multi-rank hardening (per-pid subdirs + module/ABI locking; verified against vLLM 4-process EP4 and PegaInfer 4-ranks-one-process K3). Branch `feat/graph-dump-json`. +- **First cross-model application:** K3 pruned-224 EP4 bs-ladder A/B vs vLLM (`docs/models/k3/vllm-kernel-ab.md`) — found KDA-fusion and attn-res-fusion port levers, confirmed our MoE/collective structure wins, and demonstrated the per-bs capture discipline (vLLM switches GEMM provider and routing kernel by batch size). +- **Next action:** nsys A/B the attention span — PegaInfer's `BatchDecode` + `MergeStates` two-pass vs vLLM's single `fmhaSm100fKernel`, on bs=1 mid-context — to size the win before porting. If it pays, archive the Qwen `fmhaSm100fKernel` variant under the glm52 cubin pattern (README + sha256, not yet wired to build), load it in a standalone test via `cuModuleLoadData` + the captured ABI, and numerically A/B against our decode attention. +- **Not yet built:** the `qwen34b.toml` manifest, the loader with the three static checks, the per-op reference/gate tier, the catalog selector. Capture is done; these are the consumers. diff --git a/pegainfer-core/src/cuda_graph/dump.rs b/pegainfer-core/src/cuda_graph/dump.rs index 8b8880d14..0db9d8fac 100644 --- a/pegainfer-core/src/cuda_graph/dump.rs +++ b/pegainfer-core/src/cuda_graph/dump.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::ffi::CStr; +use std::ffi::c_void; use std::fmt::Write as _; use std::io::Write as _; use std::path::Path; @@ -27,6 +28,7 @@ pub struct CudaGraphDumpSummary { pub kernels: usize, pub dot_path: PathBuf, pub png_path: PathBuf, + pub json_path: PathBuf, } struct GraphDescription { @@ -67,12 +69,43 @@ enum GraphNodeKind { grid: [u32; 3], block: [u32; 3], dynamic_shared_mem_bytes: u32, + attributes: KernelAttributes, + /// `None` when the node carries no staged `kernelParams` array (e.g. + /// packed `extra` launch buffers, which stream capture never produces). + params: Option>, }, Other { node_type: String, }, } +/// Per-function facts a replayer needs beyond the launch configuration. +#[derive(Clone, Copy)] +struct KernelAttributes { + num_regs: i32, + static_shared_bytes: i32, + const_bytes: i32, + local_bytes: i32, + ptx_version: i32, + binary_version: i32, +} + +/// One kernel parameter's staged bytes, captured at graph-capture time. +struct KernelParamDump { + offset: usize, + size: usize, + bytes: Vec, + /// Present when an 8-byte value resolves as a live CUDA allocation. + /// Advisory: an integer that collides with an allocation also matches. + pointer: Option, +} + +struct PointerRange { + memory_type: &'static str, + range_start: u64, + range_size: usize, +} + pub fn validate_graph_dump_request(png_path: &Path) -> Result<()> { ensure!( png_path @@ -99,7 +132,8 @@ pub fn validate_graph_dump_request(png_path: &Path) -> Result<()> { } fn require_graph_dump_driver() -> Result<()> { - const MIN_DRIVER_API_VERSION: i32 = 12_030; + // 12.3 for cuFuncGetName, 12.4 for cuFuncGetParamInfo. + const MIN_DRIVER_API_VERSION: i32 = 12_040; let mut version = 0i32; check( @@ -108,7 +142,8 @@ fn require_graph_dump_driver() -> Result<()> { )?; ensure!( version >= MIN_DRIVER_API_VERSION, - "--dump-graph-png requires CUDA driver API 12.3 or newer for kernel names; found {}", + "--dump-graph-png requires CUDA driver API 12.4 or newer for kernel names \ + and parameter layouts; found {}", format_driver_api_version(version) ); Ok(()) @@ -141,6 +176,9 @@ impl CudaGraphState { let dot_path = png_path.with_extension("dot"); std::fs::write(&dot_path, graph.detailed_dot()) .with_context(|| format!("write detailed CUDA Graph DOT to {}", dot_path.display()))?; + let json_path = png_path.with_extension("json"); + std::fs::write(&json_path, graph.machine_json(title)?) + .with_context(|| format!("write CUDA Graph JSON to {}", json_path.display()))?; render_png(&graph.human_dot(title), png_path)?; Ok(CudaGraphDumpSummary { nodes: graph.nodes.len(), @@ -152,6 +190,7 @@ impl CudaGraphState { .count(), dot_path, png_path: png_path.to_path_buf(), + json_path, }) } @@ -184,6 +223,8 @@ impl CudaGraphState { grid, block, dynamic_shared_mem_bytes, + attributes, + params, } => GraphNodeKind::Kernel { raw_symbol, demangled: demangled @@ -192,6 +233,8 @@ impl CudaGraphState { grid, block, dynamic_shared_mem_bytes, + attributes, + params, }, RawNodeKind::Other { node_type } => GraphNodeKind::Other { node_type }, }, @@ -227,6 +270,8 @@ enum RawNodeKind { grid: [u32; 3], block: [u32; 3], dynamic_shared_mem_bytes: u32, + attributes: KernelAttributes, + params: Option>, }, Other { node_type: String, @@ -331,11 +376,149 @@ fn inspect_node(node: CUgraphNode) -> Result { .to_str() .context("CUDA kernel name is not UTF-8")? .to_owned(); + let attributes = kernel_attributes(params.func)?; + let staged_params = kernel_param_dumps(params.func, params.kernelParams) + .with_context(|| format!("read staged kernel parameters of `{raw_symbol}`"))?; Ok(RawNodeKind::Kernel { raw_symbol, grid: [params.gridDimX, params.gridDimY, params.gridDimZ], block: [params.blockDimX, params.blockDimY, params.blockDimZ], dynamic_shared_mem_bytes: params.sharedMemBytes, + attributes, + params: staged_params, + }) +} + +fn kernel_attributes(func: sys::CUfunction) -> Result { + use sys::CUfunction_attribute_enum as Attr; + let attribute = |attribute: sys::CUfunction_attribute, what: &str| -> Result { + let mut value = 0i32; + check( + unsafe { sys::cuFuncGetAttribute(&raw mut value, attribute, func) }, + what, + )?; + Ok(value) + }; + Ok(KernelAttributes { + num_regs: attribute( + Attr::CU_FUNC_ATTRIBUTE_NUM_REGS, + "cuFuncGetAttribute(num_regs)", + )?, + static_shared_bytes: attribute( + Attr::CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, + "cuFuncGetAttribute(shared_size_bytes)", + )?, + const_bytes: attribute( + Attr::CU_FUNC_ATTRIBUTE_CONST_SIZE_BYTES, + "cuFuncGetAttribute(const_size_bytes)", + )?, + local_bytes: attribute( + Attr::CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, + "cuFuncGetAttribute(local_size_bytes)", + )?, + ptx_version: attribute( + Attr::CU_FUNC_ATTRIBUTE_PTX_VERSION, + "cuFuncGetAttribute(ptx_version)", + )?, + binary_version: attribute( + Attr::CU_FUNC_ATTRIBUTE_BINARY_VERSION, + "cuFuncGetAttribute(binary_version)", + )?, + }) +} + +/// Read every staged parameter value of one captured kernel node. +/// +/// `cuFuncGetParamInfo` walks the function's parameter layout; the graph node's +/// `kernelParams` array points at graph-owned host staging for each value. +fn kernel_param_dumps( + func: sys::CUfunction, + kernel_params: *mut *mut c_void, +) -> Result>> { + // The driver caps a kernel's parameter buffer at a few KiB; a runaway index + // here means the walk is broken, not that the kernel is huge. + const MAX_PARAMS: usize = 4_096; + if kernel_params.is_null() { + return Ok(None); + } + let mut dumps = Vec::new(); + for index in 0..MAX_PARAMS { + let mut offset = 0usize; + let mut size = 0usize; + let result = + unsafe { sys::cuFuncGetParamInfo(func, index, &raw mut offset, &raw mut size) }; + if result == sys::CUresult::CUDA_ERROR_INVALID_VALUE { + break; + } + check(result, "cuFuncGetParamInfo")?; + let staged = unsafe { *kernel_params.add(index) }; + ensure!( + !staged.is_null(), + "kernel parameter {index} has no staged value" + ); + let bytes = unsafe { std::slice::from_raw_parts(staged.cast::(), size) }.to_vec(); + let pointer = (size == 8) + .then(|| { + let value = u64::from_le_bytes(bytes[..8].try_into().expect("size == 8")); + classify_pointer(value) + }) + .flatten(); + dumps.push(KernelParamDump { + offset, + size, + bytes, + pointer, + }); + } + Ok(Some(dumps)) +} + +/// Resolve an 8-byte parameter value against the CUDA allocation map. +fn classify_pointer(value: u64) -> Option { + use sys::CUpointer_attribute_enum as Attr; + if value == 0 { + return None; + } + let mut memory_type = 0u32; + let result = unsafe { + sys::cuPointerGetAttribute( + (&raw mut memory_type).cast::(), + Attr::CU_POINTER_ATTRIBUTE_MEMORY_TYPE, + value, + ) + }; + if result != sys::CUresult::CUDA_SUCCESS { + return None; + } + let memory_type = match memory_type { + 1 => "host", + 2 => "device", + 3 => "array", + 4 => "unified", + _ => "unknown", + }; + let mut range_start = 0u64; + let mut range_size = 0usize; + // Best-effort: a pointer with a memory type but no queryable range keeps + // the zero placeholders rather than failing the dump. + let _ = unsafe { + sys::cuPointerGetAttribute( + (&raw mut range_start).cast::(), + Attr::CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, + value, + ) + }; + let _ = unsafe { + sys::cuPointerGetAttribute( + (&raw mut range_size).cast::(), + Attr::CU_POINTER_ATTRIBUTE_RANGE_SIZE, + value, + ) + }; + Some(PointerRange { + memory_type, + range_start, + range_size, }) } @@ -414,6 +597,93 @@ fn communicate(mut child: Child, input: String, program: &'static str) -> Result } impl GraphDescription { + /// Machine-readable dump: everything the DOT carries plus per-kernel + /// function attributes and staged parameter bytes, for downstream tooling + /// (kernel ledger, capsule extraction) rather than human inspection. + fn machine_json(&self, title: &str) -> Result { + let mut version = 0i32; + check( + unsafe { sys::cuDriverGetVersion(&raw mut version) }, + "cuDriverGetVersion", + )?; + let nodes = self + .nodes + .iter() + .enumerate() + .map(|(index, node)| match &node.kind { + GraphNodeKind::Kernel { + raw_symbol, + demangled, + grid, + block, + dynamic_shared_mem_bytes, + attributes, + params, + } => serde_json::json!({ + "id": index, + "type": "kernel", + "symbol": raw_symbol, + "name": demangled, + "grid": grid, + "block": block, + "dynamic_shared_mem_bytes": dynamic_shared_mem_bytes, + "attributes": { + "num_regs": attributes.num_regs, + "static_shared_bytes": attributes.static_shared_bytes, + "const_bytes": attributes.const_bytes, + "local_bytes": attributes.local_bytes, + "ptx_version": attributes.ptx_version, + "binary_version": attributes.binary_version, + }, + "params": params.as_ref().map(|params| { + params + .iter() + .map(|param| { + serde_json::json!({ + "offset": param.offset, + "size": param.size, + "data": hex(¶m.bytes), + "pointer": param.pointer.as_ref().map(|pointer| { + serde_json::json!({ + "memory_type": pointer.memory_type, + "range_start": format!("{:#x}", pointer.range_start), + "range_size": pointer.range_size, + }) + }), + }) + }) + .collect::>() + }), + }), + GraphNodeKind::Other { node_type } => serde_json::json!({ + "id": index, + "type": node_type, + }), + }) + .collect::>(); + let edges = self + .edges + .iter() + .map(|edge| { + serde_json::json!({ + "from": edge.from, + "to": edge.to, + "from_port": edge.from_port, + "to_port": edge.to_port, + "dependency_type": dependency_type_name(edge.dependency_type), + }) + }) + .collect::>(); + let dump = serde_json::json!({ + "schema": "pegainfer-cuda-graph-dump/v1", + "title": title, + "driver_api_version": format_driver_api_version(version), + "nodes": nodes, + "edges": edges, + }); + serde_json::to_string_pretty(&dump).context("serialize CUDA Graph JSON dump") + } + fn detailed_dot(&self) -> String { let mut dot = String::from("digraph cuda_graph_detailed {\n"); dot.push_str(" graph [rankdir=TB];\n node [shape=box];\n"); @@ -425,6 +695,7 @@ impl GraphDescription { grid, block, dynamic_shared_mem_bytes, + .. } => format!( "id={index}\\ntype=kernel\\nname={}\\nraw_symbol={}\\ngrid={}\\nblock={}\\ndynamic_shared_mem_bytes={dynamic_shared_mem_bytes}", dot_escape(demangled), @@ -458,6 +729,14 @@ fn dims(dims: [u32; 3]) -> String { format!("({},{},{})", dims[0], dims[1], dims[2]) } +fn hex(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + let _ = write!(out, "{byte:02x}"); + } + out +} + fn dependency_type_name(dependency_type: u8) -> String { match dependency_type { 0 => "default".to_owned(), diff --git a/pegainfer-core/src/cuda_graph/dump/render.rs b/pegainfer-core/src/cuda_graph/dump/render.rs index 694fed865..5b62d0f95 100644 --- a/pegainfer-core/src/cuda_graph/dump/render.rs +++ b/pegainfer-core/src/cuda_graph/dump/render.rs @@ -452,6 +452,15 @@ mod tests { grid: [1, 1, 1], block: [32, 1, 1], dynamic_shared_mem_bytes: 0, + attributes: super::super::KernelAttributes { + num_regs: 0, + static_shared_bytes: 0, + const_bytes: 0, + local_bytes: 0, + ptx_version: 0, + binary_version: 0, + }, + params: None, }, } } diff --git a/pegainfer-k3/kernels/generate.py b/pegainfer-k3/kernels/generate.py index b2fa7b299..08e13065d 100644 --- a/pegainfer-k3/kernels/generate.py +++ b/pegainfer-k3/kernels/generate.py @@ -158,6 +158,7 @@ LAND_CONFIGS = [ # NT N OFF engine call site (4 * KDA_DIM, KDA_DIM, 3 * KDA_DIM), # KDA output gate + (4 * KDA_DIM, 3 * KDA_DIM, 0), # KDA capsule packed q|k|v (KDA_DIM, KDA_DIM, 0), # chunked-prefill conv inputs (WSM_N, KDA_HEADS, 0), # KDA beta (WSM_N, KDA_HEAD_DIM, KDA_HEADS), # KDA low-rank gate input diff --git a/pegainfer-k3/src/executor/buffers.rs b/pegainfer-k3/src/executor/buffers.rs index d5b38d0a3..76f23236e 100644 --- a/pegainfer-k3/src/executor/buffers.rs +++ b/pegainfer-k3/src/executor/buffers.rs @@ -21,6 +21,7 @@ use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; use half::bf16; use pegainfer_kernels::ops::K3_ATTNRES_MAX_BLOCKS; +use pegainfer_kernels::ops::K3_CAPSULE_CONV_SLOT; use pegainfer_kernels::ops::K3_CONV_WIDTH; use pegainfer_kernels::ops::K3_KDA_HEAD_DIM; use pegainfer_kernels::ops::K3_KDA_HEADS; @@ -39,6 +40,7 @@ use pegainfer_kernels::tensor::DeviceContext; use pegainfer_kernels::tensor::HiddenStates; use super::ep::K3FabricSlab; +use super::forward::capsule::capsule_flags; use super::paged_kv::K3PagedKv; use crate::config::K3_ATTN_INNER; use crate::config::K3_DENSE_INTERMEDIATE; @@ -80,6 +82,12 @@ pub(crate) struct K3KdaState { pub(crate) recurrent: [CudaSlice; 2], /// `[2][3][rows, width - 1, inner]` bf16, one window per q/k/v stream. pub(crate) conv: [[CudaSlice; 3]; 2], + /// Capsule-KDA conv slab, `[rows, width - 1, 3, inner]` bf16 — the + /// packed `q|k|v` tap layout the vendored vLLM decode kernel compiles + /// in. Present only under `PEGAINFER_K3_CAPSULE=kda`; updated in place + /// by the capsule step (no parity pair), filled by [`K3StatePool::adopt_row`] + /// from the native windows at the prefill handover. + pub(crate) conv_packed: Option>, } pub(crate) enum K3LayerState { @@ -162,9 +170,18 @@ impl K3StatePool { let [recurrent_even, recurrent_odd] = recurrent .try_into() .unwrap_or_else(|_| unreachable!("two parities were pushed")); + let conv_packed = capsule_flags() + .kda + .then(|| { + stream + .alloc_zeros::(rows * K3_CAPSULE_CONV_SLOT) + .context("alloc K3 capsule conv slab") + }) + .transpose()?; K3LayerState::Kda(Box::new(K3KdaState { recurrent: [recurrent_even, recurrent_odd], conv: [conv_even, conv_odd], + conv_packed, })) } K3LayerKind::Mla => K3LayerState::Mla, @@ -198,6 +215,9 @@ impl K3StatePool { zero_rows(ctx, stream, row, 1, K3_CONV_STATE * K3_ATTN_INNER)?; } } + if let Some(packed) = kda.conv_packed.as_mut() { + zero_rows(ctx, packed, row, 1, K3_CAPSULE_CONV_SLOT)?; + } } // The paged latent cache is released below; freed pages are // zeroed when next claimed, not here. @@ -248,28 +268,54 @@ impl K3StatePool { for (target, origin) in self.layers.iter_mut().zip(&source.layers) { match (target, origin) { (K3LayerState::Kda(target), K3LayerState::Kda(origin)) => { + // The capsule step reads and updates parity slab 0 in + // place, whatever the step counter says — land there. + let recurrent_parity = if target.conv_packed.is_some() { + 0 + } else { + target_parity + }; copy_rows( ctx, &origin.recurrent[source_parity], source_row, - &mut target.recurrent[target_parity], + &mut target.recurrent[recurrent_parity], row, 1, K3_KDA_STATE, )?; - for (target, origin) in target.conv[target_parity] - .iter_mut() - .zip(&origin.conv[source_parity]) - { - copy_rows( - ctx, - origin, - source_row, - target, - row, - 1, - K3_CONV_STATE * K3_ATTN_INNER, - )?; + if let Some(packed) = target.conv_packed.as_mut() { + // Repack the native per-stream `[width-1, inner]` + // windows into the capsule slab's `[width-1, q|k|v, + // inner]` row — the only place the two layouts meet. + for (stream, origin) in origin.conv[source_parity].iter().enumerate() { + copy_rows_2d( + ctx, + origin, + source_row * K3_CONV_STATE * K3_ATTN_INNER, + K3_ATTN_INNER, + packed, + row * K3_CAPSULE_CONV_SLOT + stream * K3_ATTN_INNER, + 3 * K3_ATTN_INNER, + K3_CONV_STATE, + K3_ATTN_INNER, + )?; + } + } else { + for (target, origin) in target.conv[target_parity] + .iter_mut() + .zip(&origin.conv[source_parity]) + { + copy_rows( + ctx, + origin, + source_row, + target, + row, + 1, + K3_CONV_STATE * K3_ATTN_INNER, + )?; + } } } // The paged latent cache is adopted once for the whole pool, @@ -767,6 +813,9 @@ pub(crate) struct K3Scratch { pub(crate) conv_k: CudaSlice, pub(crate) conv_v: CudaSlice, pub(crate) gated: CudaSlice, + /// Capsule KDA: the packed `q|k|v` pre-conv rows, `[rows, 3 * inner]` + /// bf16, landed in one span from the fused projection partial. + pub(crate) kda_x_packed: CudaSlice, /// Chunked prefill (FlashKDA): the landed pre-activation gate projection, /// `[rows, inner]` — the same `bf16(Σ gp)` landing the fused core takes /// before adding `dt_bias`, which FlashKDA applies in-kernel. @@ -885,6 +934,7 @@ impl K3Scratch { conv_k: wide(K3_ATTN_INNER)?, conv_v: wide(K3_ATTN_INNER)?, gated: wide(K3_ATTN_INNER)?, + kda_x_packed: wide(3 * K3_ATTN_INNER)?, kda_g: wide(K3_ATTN_INNER)?, kda_beta_t: wide(K3_KDA_HEADS)?, kda_attn: wide(K3_ATTN_INNER)?, diff --git a/pegainfer-k3/src/executor/forward/capsule.rs b/pegainfer-k3/src/executor/forward/capsule.rs new file mode 100644 index 000000000..c211231af --- /dev/null +++ b/pegainfer-k3/src/executor/forward/capsule.rs @@ -0,0 +1,43 @@ +//! Runtime selection of the capsule-vendored external decode kernels. +//! +//! `PEGAINFER_K3_CAPSULE` picks which ops run the vendored vLLM v0.28.0 +//! cubins (`pegainfer-kernels/cubin/k3/`) instead of the native kernels: +//! `all`, or a comma list of `topk`, `kda`. Unset or empty = native kernels, +//! byte-identical serving. The flag exists for same-binary A/B; each capsule +//! op has its own numeric gate before it is allowed to default on. + +use std::sync::LazyLock; + +#[derive(Debug, Default, Clone, Copy)] +pub(crate) struct CapsuleFlags { + pub(crate) topk: bool, + pub(crate) kda: bool, +} + +static FLAGS: LazyLock = LazyLock::new(|| { + let raw = std::env::var("PEGAINFER_K3_CAPSULE").unwrap_or_default(); + let mut flags = CapsuleFlags::default(); + for part in raw.split(',').map(str::trim).filter(|p| !p.is_empty()) { + match part { + "all" => { + flags.topk = true; + flags.kda = true; + } + "topk" => flags.topk = true, + "kda" => flags.kda = true, + other => { + // Refuse to start on a typo rather than silently serving the + // wrong kernel set. + panic!("PEGAINFER_K3_CAPSULE: unknown op {other:?} (expected all, topk, kda)"); + } + } + } + if flags.topk || flags.kda { + log::info!("K3 capsule kernels enabled: {flags:?}"); + } + flags +}); + +pub(crate) fn capsule_flags() -> CapsuleFlags { + *FLAGS +} diff --git a/pegainfer-k3/src/executor/forward/decode.rs b/pegainfer-k3/src/executor/forward/decode.rs index 45126b02f..c64b18e8b 100644 --- a/pegainfer-k3/src/executor/forward/decode.rs +++ b/pegainfer-k3/src/executor/forward/decode.rs @@ -5,6 +5,7 @@ use anyhow::Result; use cudarc::driver::CudaSlice; use half::bf16; use pegainfer_kernels::ops::K3_CONV_WIDTH; +use pegainfer_kernels::ops::k3_capsule_kda_decode_launch; use pegainfer_kernels::ops::k3_conv_silu_batched_launch; use pegainfer_kernels::ops::k3_kda_core_batched_launch; use pegainfer_kernels::ops::k3_land_batched_launch; @@ -199,6 +200,113 @@ pub(super) fn kda_attention( ) } +/// The capsule spelling of [`kda_attention`]: one full fused projection GEMM +/// (instead of the gate quarter plus three band GEMMs), two landings out of +/// it (`out_gate` and the packed `q|k|v` rows), and the vendored vLLM fused +/// decode kernel in place of the conv_silu x3 + kda_core chain. The kernel +/// updates the packed conv slab and the recurrent state **in place**, so the +/// caller passes the fixed parity-0 slab and no successor buffers. +#[allow(clippy::too_many_arguments)] +pub(super) fn kda_attention_capsule( + ctx: &DeviceContext, + b: usize, + layer: &K3LayerWeights, + w: &K3KdaWeights, + recurrent: &mut CudaSlice, + conv_packed: &mut CudaSlice, + s: &mut K3Scratch, +) -> Result<()> { + k3_rms_norm_rbs_batched_launch( + ctx, + b, + K3_HIDDEN, + &s.mixed, + &layer.gamma_in.data, + &mut s.normed, + )?; + k3_gemm_full(ctx, &w.wbig, &s.normed, b, &mut s.kda_gate_partial)?; + k3_land_batched_launch( + ctx, + b, + K3_KDA_FUSED, + K3_ATTN_INNER, + 3 * K3_ATTN_INNER, + 1, + &s.kda_gate_partial, + &mut s.out_gate, + )?; + k3_land_batched_launch( + ctx, + b, + K3_KDA_FUSED, + 3 * K3_ATTN_INNER, + 0, + 1, + &s.kda_gate_partial, + &mut s.kda_x_packed, + )?; + k3_gemm_full(ctx, &w.wsm, &s.normed, b, &mut s.kda_wsm_partial)?; + k3_land_batched_launch( + ctx, + b, + K3_KDA_WSM_PADDED, + K3_HEADS, + 0, + 1, + &s.kda_wsm_partial, + &mut s.beta, + )?; + k3_land_batched_launch( + ctx, + b, + K3_KDA_WSM_PADDED, + K3_HEAD_DIM, + K3_HEADS, + 1, + &s.kda_wsm_partial, + &mut s.forget_low, + )?; + k3_gemm_full(ctx, &w.w_f_b, &s.forget_low, b, &mut s.kda_forget_partial)?; + k3_land_batched_launch( + ctx, + b, + K3_ATTN_INNER, + K3_ATTN_INNER, + 0, + 1, + &s.kda_forget_partial, + &mut s.kda_g, + )?; + k3_capsule_kda_decode_launch( + ctx, + b, + &s.kda_x_packed, + &w.cw_q, + &w.cw_k, + &w.cw_v, + conv_packed, + &w.a_log, + &s.kda_g, + &w.dt_bias, + &s.beta, + &s.out_gate, + &w.gamma_o, + recurrent, + &mut s.gated, + )?; + k3_gemm_full(ctx, &w.w_o, &s.gated, b, &mut s.hidden_partial)?; + k3_land_batched_launch( + ctx, + b, + K3_HIDDEN, + K3_HIDDEN, + 0, + 1, + &s.hidden_partial, + &mut s.attn_out, + ) +} + /// One q/k/v stream: its band of the fused projection, then the window. /// /// The band goes to a partial of its own rather than into the fused one: the diff --git a/pegainfer-k3/src/executor/forward/mod.rs b/pegainfer-k3/src/executor/forward/mod.rs index b2a6b90c8..4e6885bd3 100644 --- a/pegainfer-k3/src/executor/forward/mod.rs +++ b/pegainfer-k3/src/executor/forward/mod.rs @@ -37,6 +37,7 @@ //! drafts). Full epilogue; the caller reads the span argmaxes back and //! decides acceptance. +pub(crate) mod capsule; mod decode; mod gemm; mod prefill; diff --git a/pegainfer-k3/src/executor/forward/step.rs b/pegainfer-k3/src/executor/forward/step.rs index 9e146133b..de0ce65a5 100644 --- a/pegainfer-k3/src/executor/forward/step.rs +++ b/pegainfer-k3/src/executor/forward/step.rs @@ -47,6 +47,7 @@ use pegainfer_kernels::ops::extract_hidden_rows_raw_into; use pegainfer_kernels::ops::k3_add2_batched_launch; use pegainfer_kernels::ops::k3_attnres_mix_batched_launch; use pegainfer_kernels::ops::k3_attnres_scores_batched_launch; +use pegainfer_kernels::ops::k3_capsule_router_topk_launch; use pegainfer_kernels::ops::k3_deepgemm_sm100_masked_grouped_fp8_fp4_launch; use pegainfer_kernels::ops::k3_fp8_scale_pack_ue8m0_launch; use pegainfer_kernels::ops::k3_land_batched_launch; @@ -73,7 +74,9 @@ use super::super::cp::K3CpScratch; use super::super::paged_kv::K3_KV_PAGE_TOKENS; use super::super::paged_kv::K3_MLA_LATENT_ROW; use super::super::paged_kv::K3PagedKv; +use super::capsule::capsule_flags; use super::decode::kda_attention; +use super::decode::kda_attention_capsule; use super::gemm::k3_gemm_full; use super::prefill::kda_attention_chunk; use super::prefill::mla_attention_chunk_cp; @@ -276,6 +279,24 @@ pub(super) fn k3_step( match (&layer.attn, layer_state) { (K3LayerAttention::Kda(kda), K3LayerState::Kda(kda_state)) => match mode { + K3StepMode::Decode if capsule_flags().kda => { + // The capsule kernel shifts the packed conv slab and + // updates the recurrent state in place; parity slab 0 is + // the fixed home (`adopt_row` lands there in this mode). + let conv_packed = kda_state + .conv_packed + .as_mut() + .expect("capsule mode allocates the packed conv slab"); + kda_attention_capsule( + ctx, + b, + layer, + kda, + &mut kda_state.recurrent[0], + conv_packed, + scratch, + )?; + } K3StepMode::Decode => { let (recurrent_read, recurrent_write) = parity_pair(&mut kda_state.recurrent, shape.parity); @@ -700,17 +721,31 @@ fn moe_mlp( )?; // Routing reads the pre-down hidden, as in the reference. k3_gemm_full(ctx, &w.w_router, &s.normed, b, &mut s.router_partial)?; - k3_router_topk_batched_launch( - ctx, - b, - experts, - K3_ROUTER_TOPK, - &s.router_partial, - &w.bias, - &w.rs.data, - &mut s.topk_idx, - &mut s.topk_weight, - )?; + if capsule_flags().topk { + k3_capsule_router_topk_launch( + ctx, + b, + experts, + K3_ROUTER_TOPK, + &s.router_partial, + &w.bias, + w.rs_host, + &mut s.topk_idx, + &mut s.topk_weight, + )?; + } else { + k3_router_topk_batched_launch( + ctx, + b, + experts, + K3_ROUTER_TOPK, + &s.router_partial, + &w.bias, + &w.rs.data, + &mut s.topk_idx, + &mut s.topk_weight, + )?; + } k3_gemm_full(ctx, &w.w_lat_down, &s.normed, b, &mut s.latent_partial)?; k3_land_batched_launch( ctx, diff --git a/pegainfer-k3/src/executor/mod.rs b/pegainfer-k3/src/executor/mod.rs index e7670e647..1c4595c82 100644 --- a/pegainfer-k3/src/executor/mod.rs +++ b/pegainfer-k3/src/executor/mod.rs @@ -930,6 +930,12 @@ impl K3Executor { self.dspark.is_none(), "K3 dspark draft lane is already loaded" ); + // The capsule KDA kernel updates recurrent/conv state in place on a + // fixed slab, which the verify lane's parity replay cannot rewind. + ensure!( + !forward::capsule::capsule_flags().kda, + "K3 dspark draft lane is incompatible with PEGAINFER_K3_CAPSULE=kda" + ); // One slot's worst verify round packs its deferred-commit replay // (up to a full accepted block) plus anchor and drafts. ensure!( @@ -1424,6 +1430,10 @@ impl K3Executor { cp_rank: usize, segments: Vec<(usize, usize)>, ) -> Result<()> { + ensure!( + !forward::capsule::capsule_flags().kda, + "K3 CP prefill is incompatible with PEGAINFER_K3_CAPSULE=kda" + ); if let Some(scratch) = self.cp_scratch.as_ref() { // One gang per process, one seg_cap per executor — the scratch // built once serves every superstep. diff --git a/pegainfer-k3/src/model/mod.rs b/pegainfer-k3/src/model/mod.rs index 96e237186..3da0133b6 100644 --- a/pegainfer-k3/src/model/mod.rs +++ b/pegainfer-k3/src/model/mod.rs @@ -125,6 +125,9 @@ pub(crate) struct K3MoeWeights { pub(crate) bias: CudaSlice, /// Routed scaling factor as a device scalar, bf16 `[1]`. pub(crate) rs: DeviceVec, + /// The same routed scaling factor as a host scalar, read back once at + /// build for kernels that take it by value (the capsule router top-k). + pub(crate) rs_host: f32, pub(crate) w_lat_down: DeviceMatrix, pub(crate) w_lat_up: DeviceMatrix, pub(crate) gamma_lat: DeviceVec, @@ -546,10 +549,13 @@ fn build_layer( form, )?; vram.experts += experts.bytes(); + let rs = slots.vector("rs", 1)?; + let rs_host = rs.to_host(ctx)?[0]; K3LayerMlp::Moe(Box::new(K3MoeWeights { w_router: slots.matrix("w_router", routed_experts.count(), K3_HIDDEN)?, bias: slots.f32("bias", routed_experts.count(), 1)?, - rs: slots.vector("rs", 1)?, + rs, + rs_host, w_lat_down: slots.matrix("w_lat_down", K3_ROUTED_EXPERT_HIDDEN, K3_HIDDEN)?, w_lat_up: slots.matrix("w_lat_up", K3_HIDDEN, K3_ROUTED_EXPERT_HIDDEN)?, gamma_lat: slots.vector("gamma_lat", K3_ROUTED_EXPERT_HIDDEN)?, diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index 235614142..4f8fd258f 100644 --- a/pegainfer-kernels/build.rs +++ b/pegainfer-kernels/build.rs @@ -110,6 +110,50 @@ fn generate_glm52_trtllm_fmha_cubins(root: &Path, out_dir: &Path) { }); } +/// K3 vendored decode kernels (cubin/k3/): external serving kernels bound by +/// the capsule loader in csrc/k3/k3_capsule.cu. Provenance in cubin/k3/README. +const K3_CAPSULE_CUBINS: &[(&str, &str)] = &[ + ("kK3CapsuleTopk", "single_group_topk_e512t22_sm103.cubin"), + ("kK3CapsuleKdaDecode", "kda_decode_fusion_h96_sm103.cubin"), +]; + +fn generate_k3_capsule_cubins(root: &Path, out_dir: &Path) { + let cubin_dir = root.join("cubin/k3"); + let mut generated = String::from( + "// Generated by pegainfer-kernels/build.rs from the checked-in K3 capsule cubins.\n\ + #pragma once\n\n", + ); + for &(symbol, file_name) in K3_CAPSULE_CUBINS { + let path = cubin_dir.join(file_name); + println!("cargo:rerun-if-changed={}", path.display()); + let bytes = fs::read(&path).unwrap_or_else(|err| { + panic!("failed to read K3 capsule cubin {}: {err}", path.display()) + }); + writeln!(&mut generated, "static const unsigned char {symbol}[] = {{") + .expect("write cubin array declaration"); + for chunk in bytes.chunks(16) { + generated.push_str(" "); + for byte in chunk { + write!(&mut generated, "0x{byte:02x}, ").expect("write cubin byte"); + } + generated.push('\n'); + } + generated.push_str("};\n"); + writeln!( + &mut generated, + "static constexpr unsigned int {symbol}Size = sizeof({symbol});\n" + ) + .expect("write cubin array size"); + } + let output = out_dir.join("k3_capsule_cubins.inc"); + fs::write(&output, generated).unwrap_or_else(|err| { + panic!( + "failed to write generated K3 capsule cubin header {}: {err}", + output.display() + ) + }); +} + fn workspace_root() -> PathBuf { crate_root().join("..") } @@ -1884,6 +1928,9 @@ fn main() { generate_glm52_trtllm_fmha_cubins(&crate_root(), &out_dir); build_glm52_cutedsl_fp8_dsl(&crate_root(), &out_dir, &cuda_include); } + if k3_enabled { + generate_k3_capsule_cubins(&crate_root(), &out_dir); + } println!( "cargo:warning=Detected CUDA SM targets: {}", sm_targets @@ -1999,6 +2046,10 @@ fn main() { "-I".to_string(), csrc_dir.to_string_lossy().to_string(), ]; + if stem == "k3_capsule" { + // Includes the generated k3_capsule_cubins.inc from OUT_DIR. + nvcc_args.extend(["-I".to_string(), out_dir.to_string_lossy().to_string()]); + } if stem == "glm52_fp8_gemm" { if let Some(args) = glm52_fp8_gemm_arch_args(&nvcc_sm_targets, &nvcc) { nvcc_args.extend(args); diff --git a/pegainfer-kernels/csrc/k3/k3_capsule.cu b/pegainfer-kernels/csrc/k3/k3_capsule.cu new file mode 100644 index 000000000..daaa6d819 --- /dev/null +++ b/pegainfer-kernels/csrc/k3/k3_capsule.cu @@ -0,0 +1,170 @@ +// K3 capsule loader: binds vendored external decode kernels (cubin/k3/, +// embedded at build time via k3_capsule_cubins.inc) to extern "C" launchers +// the Rust side calls like any other kernel. The framework owns the op +// contract and the reference twin; these cubins are replaceable artifacts. +// +// Loading is per-thread (K3 executor threads each own one device/context, +// same discipline as the thread_local cuBLAS handle in shared/linear.cu). +// Each capsule is loaded lazily on first launch and sanity-checked +// fail-closed: cuFuncGetParamInfo must report the exact staged parameter +// layout recorded at capture time, or the launcher refuses to run. + +#include +#include +#include + +#include +#include + +#include "k3_capsule_cubins.inc" + +namespace { + +struct Capsule { + CUmodule module = nullptr; + CUfunction func = nullptr; + CUresult status = CUDA_ERROR_NOT_INITIALIZED; + bool tried = false; +}; + +// Expected staged-parameter layout, from cuFuncGetParamInfo against the +// vendored cubin (recorded in the capture's port_abi manifest). +struct ParamSpec { + int count; + size_t total_bytes; // offset+size of the last parameter +}; + +CUresult capsule_load(Capsule& cap, const unsigned char* image, + const char* entry, const ParamSpec& spec) { + if (cap.tried) { + return cap.status; + } + cap.tried = true; + cap.status = cuModuleLoadData(&cap.module, image); + if (cap.status != CUDA_SUCCESS) { + fprintf(stderr, "[k3-capsule] cuModuleLoadData failed for %s: %d\n", entry, + (int)cap.status); + return cap.status; + } + cap.status = cuModuleGetFunction(&cap.func, cap.module, entry); + if (cap.status != CUDA_SUCCESS) { + fprintf(stderr, "[k3-capsule] entry %s not found: %d\n", entry, + (int)cap.status); + return cap.status; + } + // Static ABI check: parameter count and packed size must match the layout + // this launcher stages. A drifted cubin fails here, not at launch. + size_t offset = 0, size = 0; + int count = 0; + while (count < 64 && + cuFuncGetParamInfo(cap.func, (size_t)count, &offset, &size) == + CUDA_SUCCESS) { + count++; + } + if (count != spec.count || offset + size != spec.total_bytes) { + fprintf(stderr, + "[k3-capsule] ABI mismatch for %s: %d params/%zu bytes, expected " + "%d/%zu\n", + entry, count, offset + size, spec.count, spec.total_bytes); + cap.status = CUDA_ERROR_INVALID_IMAGE; + return cap.status; + } + cap.status = CUDA_SUCCESS; + return cap.status; +} + +thread_local Capsule g_topk; +thread_local Capsule g_kda; + +constexpr char kTopkEntry[] = + "_ZN4vllm3moe17single_group_topk6detail29single_group_topk_warp_kernelIffi" + "LNS0_11ScoringFuncE1ELi512ELi22EEEvPKT_PfPT1_PKT0_lllbfb"; +constexpr char kKdaEntry[] = + "_ZN44_GLOBAL__N__7c05f95b_9_kda_tu_cu_2dd95ecb_4135kda_decode_fusion_" + "many_heads_kernelILb1ELb1ELi96ELi96ELb1ELb0ELb0ELb0ELb0ELb0ELb1ELb1ELb1E" + "Lb1ELb1ELb1EEEvPK13__nv_bfloat16S3_S3_PKfS5_S5_S5_S5_S5_PS1_S6_S6_S5_S3_" + "S5_S3_S3_S5_PKiS8_PfS6_iiifffNS_16KdaDecodeStridesE"; + +// Mirror of the kernel's aggregate parameter (fused_kda_decode_kernel.cu:26). +struct KdaDecodeStrides { + int64_t x_row; + int64_t beta_row; + int64_t onorm_row; + int64_t conv_slot; + int64_t state_slot; +}; + +} // namespace + +// vLLM v0.28.0 single_group_topk warp kernel : sigmoid scoring with bias-corrected selection, +// weights from unbiased scores, renormalized and scaled on device. One warp +// per token, 8 warps per block. +extern "C" CUresult k3_capsule_router_topk_cuda( + const float* scores, const float* bias, int* topk_idx, float* topk_wts, + int b, int num_experts, int topk, float routed_scaling, + cudaStream_t stream) { + // 10 staged params, packed tail (..., bool@56, float@60, bool@64) = 65 + // bytes, from the captured cuFuncGetParamInfo walk. + CUresult rc = capsule_load(g_topk, kK3CapsuleTopk, kTopkEntry, + ParamSpec{10, 65}); + if (rc != CUDA_SUCCESS) { + return rc; + } + int64_t num_tokens = b; + int64_t experts64 = num_experts; + int64_t topk64 = topk; + bool renormalize = true; + bool enable_pdl = false; + void* params[] = { + (void*)&scores, (void*)&topk_wts, (void*)&topk_idx, + (void*)&bias, (void*)&num_tokens, (void*)&experts64, + (void*)&topk64, (void*)&renormalize, (void*)&routed_scaling, + (void*)&enable_pdl, + }; + constexpr unsigned kBlock = 256; // WarpTopKLaunchConfig<512>::BlockDim + constexpr unsigned kWarpsPerBlock = kBlock / 32; + unsigned grid = (unsigned)((b + kWarpsPerBlock - 1) / kWarpsPerBlock); + return cuLaunchKernel(g_topk.func, grid, 1, 1, kBlock, 1, 1, 0, + (CUstream)stream, params, nullptr); +} + +// vLLM v0.28.0 fused KDA decode, 96-head static-layout head-grid variant with +// conv-state update, lower bound and beta sigmoid. One block per (token, +// value head); consumes projected x_q|x_k|x_v rows, updates the bf16 conv +// windows and the fp32 recurrent state in place, applies the gated output +// norm, writes bf16 out rows. +extern "C" CUresult k3_capsule_kda_decode_cuda( + const void* x_q, const void* x_k, const void* x_v, const float* w_q_t, + const float* w_k_t, const float* w_v_t, const float* bias_q, + const float* bias_k, const float* bias_v, void* cs_q, void* cs_k, + void* cs_v, const float* a_log, const void* g, const float* dt_bias, + const void* beta, const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int b, int heads, int value_heads, float lower_bound, + float scale, float onorm_eps, int64_t x_row, int64_t beta_row, + int64_t onorm_row, int64_t conv_slot, int64_t state_slot, + cudaStream_t stream) { + // 29 staged params: 22 pointers, 3 ints, 3 floats, one 40-byte struct at + // offset 200 -> 240 bytes, from the captured cuFuncGetParamInfo walk. + CUresult rc = capsule_load(g_kda, kK3CapsuleKdaDecode, kKdaEntry, + ParamSpec{29, 240}); + if (rc != CUDA_SUCCESS) { + return rc; + } + KdaDecodeStrides strides{x_row, beta_row, onorm_row, conv_slot, state_slot}; + void* params[] = { + (void*)&x_q, (void*)&x_k, (void*)&x_v, + (void*)&w_q_t, (void*)&w_k_t, (void*)&w_v_t, + (void*)&bias_q, (void*)&bias_k, (void*)&bias_v, + (void*)&cs_q, (void*)&cs_k, (void*)&cs_v, + (void*)&a_log, (void*)&g, (void*)&dt_bias, + (void*)&beta, (void*)&onorm_g, (void*)&onorm_weight, + (void*)&ssm_state_indices, (void*)&cu_seqlens, (void*)&state, + (void*)&out, (void*)&b, (void*)&heads, + (void*)&value_heads, (void*)&lower_bound, (void*)&scale, + (void*)&onorm_eps, (void*)&strides, + }; + return cuLaunchKernel(g_kda.func, (unsigned)b, (unsigned)value_heads, 1, 256, + 1, 1, 0, (CUstream)stream, params, nullptr); +} diff --git a/pegainfer-kernels/cubin/k3/README.md b/pegainfer-kernels/cubin/k3/README.md new file mode 100644 index 000000000..e52089669 --- /dev/null +++ b/pegainfer-kernels/cubin/k3/README.md @@ -0,0 +1,37 @@ +# K3 vendored decode kernels (sm_103) + +External serving kernels for the K3 decode path, bound at runtime by the +capsule loader in `csrc/k3/k3_capsule.cu`. Provenance and license per file; +sha256 pins in `SHA256SUMS` are verified by the embedded loader. All sources +are Apache-2.0 (`SPDX-License-Identifier: Apache-2.0`, "Copyright contributors +to the vLLM project"); cuBLAS `nvjet_*` kernels are deliberately absent +(NVIDIA proprietary — never redistributed). + +Both files are **offline single-instantiation builds**: the v0.28.0 wheel +packs these kernels inside multi-MB fatbins (2.3 MB flashkda, 16.7 MB `_C`), +so instead of vendoring those, the exact launched template instantiation was +compiled standalone from the v0.28.0 tagged sources +(`docker.io/vllm/vllm-openai:v0.28.0`, +`sha256:61fc8a896b0a4fbbbdc063bc4b0dbc25ce98e02b5050c24aeb7830ac02039b14`, +linux/arm64, symbols confirmed by CUPTI module-load capture of its Kimi-K3 +DP4×EP4 decode on GB300) with CUDA 13.2 `nvcc -cubin -arch=sm_103a -O3` +(template arguments recovered from the captured launch symbols; TUs are the +upstream `.cu` with the torch host-launcher sections excised — no functional +edits to device code): + +| file | instantiation | upstream source (tag v0.28.0) | +|---|---|---| +| `kda_decode_fusion_h96_sm103.cubin` | `kda_decode_fusion_many_heads_kernel` | `csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu` | +| `single_group_topk_e512t22_sm103.cubin` | `single_group_topk_warp_kernel` | `csrc/libtorch_stable/moe/grouped_topk_kernels.cu` (+ `moeTopKFuncs.cuh`) | + +Rebuild note: the offline TUs are byte-derivable from the upstream tag; the +recipe (cut point + explicit instantiation) is recorded in +`docs/models/k3/vllm-kernel-ab.md`'s port log. A re-vendor against a newer +vLLM must re-capture the launch symbols first — template flag sets are not +stable across versions. + +Mined-but-unwired cubins from the same capture (attn-res online kernel, +CuTe-DSL skinny GEMM M=1 pair, `LLBf16Dotprod` router GEMM) are **not** +vendored here — only kernels a launch site actually binds earn a row. They +live in the capture archive (`/data/susun/kernel-capture/`, tray03) with +their ABI manifest, see the A/B doc's "not ported" rationale. diff --git a/pegainfer-kernels/cubin/k3/SHA256SUMS b/pegainfer-kernels/cubin/k3/SHA256SUMS new file mode 100644 index 000000000..3f08a172f --- /dev/null +++ b/pegainfer-kernels/cubin/k3/SHA256SUMS @@ -0,0 +1,2 @@ +d06fec9e8a5f5e6ea1d1c6873545fc094ff7079e7cd589080e5f74c3c6cd3b12 kda_decode_fusion_h96_sm103.cubin +bbba3c841e7bc8eea505ac6e97a7e2212430697fe67f19c1b01627f694629110 single_group_topk_e512t22_sm103.cubin diff --git a/pegainfer-kernels/cubin/k3/kda_decode_fusion_h96_sm103.cubin b/pegainfer-kernels/cubin/k3/kda_decode_fusion_h96_sm103.cubin new file mode 100644 index 000000000..22ba1eefe Binary files /dev/null and b/pegainfer-kernels/cubin/k3/kda_decode_fusion_h96_sm103.cubin differ diff --git a/pegainfer-kernels/cubin/k3/single_group_topk_e512t22_sm103.cubin b/pegainfer-kernels/cubin/k3/single_group_topk_e512t22_sm103.cubin new file mode 100644 index 000000000..8f1e0f739 Binary files /dev/null and b/pegainfer-kernels/cubin/k3/single_group_topk_e512t22_sm103.cubin differ diff --git a/pegainfer-kernels/src/ffi/k3.rs b/pegainfer-kernels/src/ffi/k3.rs index 5e9ed53ad..dc40fa801 100644 --- a/pegainfer-kernels/src/ffi/k3.rs +++ b/pegainfer-kernels/src/ffi/k3.rs @@ -173,6 +173,64 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + /// Capsule-vendored vLLM v0.28.0 router top-k (`cubin/k3/`, loader in + /// `csrc/k3/k3_capsule.cu`): same contract as [`k3_router_topk_cuda`] + /// except the routed scale is a host scalar and weights are written in + /// the kernel's descending-score order. + pub fn k3_capsule_router_topk_cuda( + scores: *const f32, + bias: *const f32, + topk_idx: *mut i32, + topk_wts: *mut f32, + b: i32, + num_experts: i32, + topk: i32, + routed_scaling: f32, + stream: CUstream, + ) -> CUresult; + + /// Capsule-vendored vLLM v0.28.0 fused KDA decode (96-head variant, see + /// `csrc/k3/k3_capsule.cu`): consumes projected bf16 x_q|x_k|x_v rows, + /// updates the bf16 conv windows and fp32 recurrent state in place, + /// applies the gated output norm, writes bf16 out rows. + #[allow(clippy::too_many_arguments)] + pub fn k3_capsule_kda_decode_cuda( + x_q: *const c_void, + x_k: *const c_void, + x_v: *const c_void, + w_q_t: *const f32, + w_k_t: *const f32, + w_v_t: *const f32, + bias_q: *const f32, + bias_k: *const f32, + bias_v: *const f32, + cs_q: *mut c_void, + cs_k: *mut c_void, + cs_v: *mut c_void, + a_log: *const f32, + g: *const c_void, + dt_bias: *const f32, + beta: *const c_void, + onorm_g: *const c_void, + onorm_weight: *const f32, + ssm_state_indices: *const i32, + cu_seqlens: *const i32, + state: *mut f32, + out: *mut c_void, + b: i32, + heads: i32, + value_heads: i32, + lower_bound: f32, + scale: f32, + onorm_eps: f32, + x_row: i64, + beta_row: i64, + onorm_row: i64, + conv_slot: i64, + state_slot: i64, + stream: CUstream, + ) -> CUresult; + // --- fused MegaMoE (see `csrc/k3/k3_mega_moe_sm100.cu`) --- /// Token-count alignment the MegaMoE API enforces on diff --git a/pegainfer-kernels/src/ops/k3/capsule_kda.rs b/pegainfer-kernels/src/ops/k3/capsule_kda.rs new file mode 100644 index 000000000..06e251287 --- /dev/null +++ b/pegainfer-kernels/src/ops/k3/capsule_kda.rs @@ -0,0 +1,148 @@ +//! Capsule-vendored vLLM fused KDA decode step +//! (`cubin/k3/kda_decode_fusion_h96_sm103.cubin`). +//! +//! One launch replaces the native conv_silu x3 + kda_core chain: short +//! convolution (window update in place), silu, joint q/k L2 norm, lower-bound +//! decay, delta rule (recurrent state update in place), gated output RMS +//! norm. The formulas are the native kernels' exact spellings — same +//! `GATE_LOWER_BOUND=-5`, `scale=head_dim^-0.5`, `RMS_EPS=1e-5`, tap order +//! and `[head, v_dim, k_dim]` state layout — differing only in rounding +//! chains (the native TileLang path lands several intermediates in bf16, the +//! vLLM kernel keeps them f32). +//! +//! Layout contract (compiled into the cubin, tier `heads=96, head_dim=128`): +//! +//! * `x` — packed pre-conv rows `[b, 3 * 12288]` bf16, `q|k|v` bands. +//! * `conv` — packed window slab `[rows, 3 taps, q|k|v, 12288]` bf16; the +//! tap stride `3 * 12288` is compile-time. Updated in place (shift + newest +//! row), unlike the native parity-pair windows. +//! * `state` — `[rows, 96, 128, 128]` f32, updated in place. + +use core::ffi::c_void; + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; + +use crate::ffi; +use crate::ops::K3_CONV_WIDTH; +use crate::ops::K3_KDA_DIM; +use crate::ops::K3_KDA_HEAD_DIM; +use crate::ops::K3_KDA_HEADS; +use crate::tensor::DeviceContext; + +/// Row stride of the packed pre-conv `q|k|v` input, elements. +pub const K3_CAPSULE_X_ROW: usize = 3 * K3_KDA_DIM; +/// Per-slot elements of the packed conv window slab: `taps x (q|k|v) x dim`, +/// with the tap stride `3 * K3_KDA_DIM` compiled into the cubin. +pub const K3_CAPSULE_CONV_SLOT: usize = (K3_CONV_WIDTH - 1) * 3 * K3_KDA_DIM; +/// Per-slot elements of the recurrent state. +pub const K3_CAPSULE_STATE_SLOT: usize = K3_KDA_HEADS * K3_KDA_HEAD_DIM * K3_KDA_HEAD_DIM; + +/// The native kernels' compiled-in constants, restated for the capsule launch +/// (`generate.py`: `GATE_LOWER_BOUND`, `RMS_EPS`). +const LOWER_BOUND: f32 = -5.0; +const ONORM_EPS: f32 = 1e-5; + +/// One fused KDA decode step over `b` rows, states updated in place. +#[allow(clippy::too_many_arguments)] +pub fn k3_capsule_kda_decode_launch( + ctx: &DeviceContext, + b: usize, + x: &CudaSlice, + cw_q: &CudaSlice, + cw_k: &CudaSlice, + cw_v: &CudaSlice, + conv: &mut CudaSlice, + a_log: &CudaSlice, + g: &CudaSlice, + dt_bias: &CudaSlice, + beta: &CudaSlice, + onorm_g: &CudaSlice, + onorm_weight: &CudaSlice, + state: &mut CudaSlice, + out: &mut CudaSlice, +) -> Result<()> { + ensure!(b > 0, "K3 capsule KDA needs rows"); + let taps = K3_CONV_WIDTH * K3_KDA_DIM; + ensure!( + x.len() >= b * K3_CAPSULE_X_ROW + && cw_q.len() >= taps + && cw_k.len() >= taps + && cw_v.len() >= taps + && conv.len() >= b * K3_CAPSULE_CONV_SLOT + && a_log.len() >= K3_KDA_HEADS + && g.len() >= b * K3_KDA_DIM + && dt_bias.len() >= K3_KDA_DIM + && beta.len() >= b * K3_KDA_HEADS + && onorm_g.len() >= b * K3_KDA_DIM + && onorm_weight.len() >= K3_KDA_HEAD_DIM + && state.len() >= b * K3_CAPSULE_STATE_SLOT + && out.len() >= b * K3_KDA_DIM, + "K3 capsule KDA buffers too small for b={b}: x {}, conv {}, state {}, out {}", + x.len(), + conv.len(), + state.len(), + out.len() + ); + let (x_ptr, _x_guard) = x.device_ptr(&ctx.stream); + let (cw_q_ptr, _cwq_guard) = cw_q.device_ptr(&ctx.stream); + let (cw_k_ptr, _cwk_guard) = cw_k.device_ptr(&ctx.stream); + let (cw_v_ptr, _cwv_guard) = cw_v.device_ptr(&ctx.stream); + let (conv_ptr, _conv_guard) = conv.device_ptr_mut(&ctx.stream); + let (a_log_ptr, _a_guard) = a_log.device_ptr(&ctx.stream); + let (g_ptr, _g_guard) = g.device_ptr(&ctx.stream); + let (dt_ptr, _dt_guard) = dt_bias.device_ptr(&ctx.stream); + let (beta_ptr, _bt_guard) = beta.device_ptr(&ctx.stream); + let (og_ptr, _og_guard) = onorm_g.device_ptr(&ctx.stream); + let (ow_ptr, _ow_guard) = onorm_weight.device_ptr(&ctx.stream); + let (state_ptr, _st_guard) = state.device_ptr_mut(&ctx.stream); + let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); + + let band = (K3_KDA_DIM * size_of::()) as u64; + let scale = (K3_KDA_HEAD_DIM as f32).powf(-0.5); + unsafe { + ffi::k3_capsule_kda_decode_cuda( + x_ptr as *const c_void, + (x_ptr + band) as *const c_void, + (x_ptr + 2 * band) as *const c_void, + cw_q_ptr as *const f32, + cw_k_ptr as *const f32, + cw_v_ptr as *const f32, + core::ptr::null(), + core::ptr::null(), + core::ptr::null(), + conv_ptr as *mut c_void, + (conv_ptr + band) as *mut c_void, + (conv_ptr + 2 * band) as *mut c_void, + a_log_ptr as *const f32, + g_ptr as *const c_void, + dt_ptr as *const f32, + beta_ptr as *const c_void, + og_ptr as *const c_void, + ow_ptr as *const f32, + core::ptr::null(), + core::ptr::null(), + state_ptr as *mut f32, + out_ptr as *mut c_void, + i32::try_from(b)?, + i32::try_from(K3_KDA_HEADS)?, + i32::try_from(K3_KDA_HEADS)?, + LOWER_BOUND, + scale, + ONORM_EPS, + i64::try_from(K3_CAPSULE_X_ROW)?, + i64::try_from(K3_KDA_HEADS)?, + i64::try_from(K3_KDA_DIM)?, + i64::try_from(K3_CAPSULE_CONV_SLOT)?, + i64::try_from(K3_CAPSULE_STATE_SLOT)?, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| anyhow!("K3 capsule KDA decode (B={b}) launch failed: {err}")) +} diff --git a/pegainfer-kernels/src/ops/k3/mod.rs b/pegainfer-kernels/src/ops/k3/mod.rs index 33804b6e0..94175bfe8 100644 --- a/pegainfer-kernels/src/ops/k3/mod.rs +++ b/pegainfer-kernels/src/ops/k3/mod.rs @@ -1,5 +1,6 @@ //! Kimi-K3 GPU operators. +mod capsule_kda; mod deepgemm; mod flash_kda; mod flash_mla_prefill; @@ -8,6 +9,7 @@ mod mla_paged; mod moe_chain; mod router_topk; +pub use capsule_kda::*; pub use deepgemm::*; pub use flash_kda::*; pub use flash_mla_prefill::*; diff --git a/pegainfer-kernels/src/ops/k3/router_topk.rs b/pegainfer-kernels/src/ops/k3/router_topk.rs index 17c30001c..5f30572dc 100644 --- a/pegainfer-kernels/src/ops/k3/router_topk.rs +++ b/pegainfer-kernels/src/ops/k3/router_topk.rs @@ -80,3 +80,56 @@ pub fn k3_router_topk_batched_launch( anyhow!("K3 router_topk (B={b}, E={num_experts}, TOPK={topk}) launch failed: {err}") }) } + +/// Capsule-vendored vLLM router top-k (`cubin/k3/single_group_topk_*`): same +/// selection semantics family as [`k3_router_topk_batched_launch`] (sigmoid, +/// biased selection, unbiased renormalized weights, scaled), but the routed +/// scale is a host scalar and the (idx, weight) pairs come back in the +/// kernel's descending-score order rather than selection order. Consumers +/// treat the pairs as unordered. +#[allow(clippy::too_many_arguments)] +pub fn k3_capsule_router_topk_launch( + ctx: &DeviceContext, + b: usize, + num_experts: usize, + topk: usize, + s: &CudaSlice, + bias: &CudaSlice, + routed_scaling: f32, + idx: &mut CudaSlice, + wts: &mut CudaSlice, +) -> Result<()> { + ensure!(b > 0, "K3 capsule router needs rows"); + ensure!( + num_experts <= 512 && topk <= 22 && topk <= num_experts, + "K3 capsule router tier is <=512 experts, <=22 topk; got E={num_experts}, topk={topk}" + ); + ensure!( + s.len() >= b * num_experts + && bias.len() >= num_experts + && idx.len() >= b * topk + && wts.len() >= b * topk, + "K3 capsule router buffers too small for b={b}, experts={num_experts}, topk={topk}" + ); + let (s_ptr, _s_guard) = s.device_ptr(&ctx.stream); + let (bias_ptr, _bias_guard) = bias.device_ptr(&ctx.stream); + let (idx_ptr, _idx_guard) = idx.device_ptr_mut(&ctx.stream); + let (wts_ptr, _wts_guard) = wts.device_ptr_mut(&ctx.stream); + unsafe { + ffi::k3_capsule_router_topk_cuda( + s_ptr as *const f32, + bias_ptr as *const f32, + idx_ptr as *mut i32, + wts_ptr as *mut f32, + i32::try_from(b)?, + i32::try_from(num_experts)?, + i32::try_from(topk)?, + routed_scaling, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| { + anyhow!("K3 capsule router_topk (B={b}, E={num_experts}, TOPK={topk}) launch failed: {err}") + }) +} diff --git a/pegainfer-kernels/src/ops/k3_tilelang.rs b/pegainfer-kernels/src/ops/k3_tilelang.rs index ecd71473d..71ace61dd 100644 --- a/pegainfer-kernels/src/ops/k3_tilelang.rs +++ b/pegainfer-kernels/src/ops/k3_tilelang.rs @@ -727,7 +727,8 @@ mod tests { assert!(check_bucket(bucket).is_ok()); } assert_eq!(k3_chunk_bucket(K3_MAX_BATCH + 1).unwrap(), 256); - assert_eq!(k3_chunk_bucket(4096).unwrap(), K3_MAX_CHUNK); + assert_eq!(k3_chunk_bucket(4096).unwrap(), 4224); + assert_eq!(k3_chunk_bucket(8449).unwrap(), K3_MAX_CHUNK); assert!(k3_chunk_bucket(K3_MAX_CHUNK + 1).is_err()); } diff --git a/pegainfer-kernels/tests/k3_capsule_kda_gate.rs b/pegainfer-kernels/tests/k3_capsule_kda_gate.rs new file mode 100644 index 000000000..1e056b0f9 --- /dev/null +++ b/pegainfer-kernels/tests/k3_capsule_kda_gate.rs @@ -0,0 +1,301 @@ +//! Numeric gate for the capsule-vendored vLLM fused KDA decode +//! (`cubin/k3/kda_decode_fusion_h96_sm103.cubin`) against the native +//! conv_silu x3 + kda_core chain on identical inputs and states. +//! +//! Both sides spell the same math (short conv + silu, joint q/k L2 norm, +//! `exp(-5 * sigmoid(exp(a_log) * (g + dt_bias)))` decay, sigmoid-beta delta +//! rule, gated output RMS norm with eps 1e-5) but round differently: the +//! native TileLang chain lands the conv output, the rsqrt and several other +//! intermediates in bf16, the vLLM kernel keeps them f32. The contract under +//! test is therefore layout + semantics, not bitwise identity: +//! +//! * conv windows must match **bitwise** — the shift is a copy and the newest +//! slot is the same bf16 landing on both sides; +//! * the updated recurrent state and the output row must agree within a +//! rounding-chain tolerance (a layout or indexing mistake produces O(1) +//! garbage, orders of magnitude beyond it). +//! +//! Manual gate: CI compiles this but never runs it. Run on a Blackwell box +//! (set `PEGAINFER_REQUIRE_GPU=1` to turn a missing device into a failure). + +#![cfg(feature = "k3")] + +mod common; + +use half::bf16; +use pegainfer_kernels::ops::K3_CAPSULE_CONV_SLOT; +use pegainfer_kernels::ops::K3_CAPSULE_STATE_SLOT; +use pegainfer_kernels::ops::K3_CAPSULE_X_ROW; +use pegainfer_kernels::ops::K3_CONV_WIDTH; +use pegainfer_kernels::ops::K3_KDA_DIM; +use pegainfer_kernels::ops::K3_KDA_HEAD_DIM; +use pegainfer_kernels::ops::K3_KDA_HEADS; +use pegainfer_kernels::ops::k3_capsule_kda_decode_launch; +use pegainfer_kernels::ops::k3_conv_silu_batched_launch; +use pegainfer_kernels::ops::k3_kda_core_batched_launch; +use pegainfer_kernels::ops::k3_land_batched_launch; +use pegainfer_kernels::tensor::DeviceContext; + +const KP: usize = K3_KDA_DIM; +const WS: usize = K3_CONV_WIDTH - 1; + +struct Lcg(u64); +impl Lcg { + fn next_u32(&mut self) -> u32 { + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (self.0 >> 32) as u32 + } + fn unit_f32(&mut self) -> f32 { + (self.next_u32() as f32 / u32::MAX as f32) * 2.0 - 1.0 + } + fn f32s(&mut self, n: usize, scale: f32) -> Vec { + (0..n).map(|_| self.unit_f32() * scale).collect() + } + fn bf16s(&mut self, n: usize, scale: f32) -> Vec { + (0..n) + .map(|_| bf16::from_f32(self.unit_f32() * scale)) + .collect() + } +} + +/// Stream-ordered d2d copy of `elems` elements between offsets of two bf16 +/// slices (the packed-layout assembly the executor does with the same call). +fn copy_bf16( + ctx: &DeviceContext, + src: &cudarc::driver::CudaSlice, + src_off: usize, + dst: &mut cudarc::driver::CudaSlice, + dst_off: usize, + elems: usize, +) { + use cudarc::driver::DevicePtr; + use cudarc::driver::DevicePtrMut; + let (src_ptr, _sg) = src.device_ptr(&ctx.stream); + let (dst_ptr, _dg) = dst.device_ptr_mut(&ctx.stream); + let e = size_of::(); + unsafe { + cudarc::driver::sys::cuMemcpyDtoDAsync_v2( + dst_ptr + (dst_off * e) as u64, + src_ptr + (src_off * e) as u64, + elems * e, + pegainfer_kernels::tensor::active_cu_stream(ctx), + ) + } + .result() + .expect("d2d copy"); +} + +#[test] +fn capsule_kda_matches_native_chain() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let mut rng = Lcg(0x4b33_4b44_2026_0828); + + // Weights, shared across the batch sweep. + let cw: Vec<_> = (0..3) + .map(|_| { + ctx.stream + .clone_htod(&rng.f32s(K3_CONV_WIDTH * KP, 0.4)) + .expect("cw H2D") + }) + .collect(); + let dt_dev = ctx.stream.clone_htod(&rng.f32s(KP, 0.2)).expect("dt H2D"); + let alog_dev = ctx + .stream + .clone_htod(&rng.f32s(K3_KDA_HEADS, 0.5)) + .expect("alog H2D"); + let go: Vec = rng + .f32s(K3_KDA_HEAD_DIM, 0.3) + .iter() + .map(|v| 1.0 + v) + .collect(); + let go_dev = ctx.stream.clone_htod(&go).expect("go H2D"); + + for &b in &[1usize, 4] { + // Per-stream f32 projection partials and carried conv windows. + let partials: Vec<_> = (0..3) + .map(|_| { + ctx.stream + .clone_htod(&rng.f32s(b * KP, 1.0)) + .expect("p H2D") + }) + .collect(); + let windows: Vec<_> = (0..3) + .map(|_| { + ctx.stream + .clone_htod(&rng.bf16s(b * WS * KP, 1.0)) + .expect("cs H2D") + }) + .collect(); + let forget = ctx.stream.clone_htod(&rng.f32s(b * KP, 1.0)).expect("gp"); + let beta_dev = ctx + .stream + .clone_htod(&rng.bf16s(b * K3_KDA_HEADS, 1.5)) + .expect("beta"); + let g2_dev = ctx.stream.clone_htod(&rng.bf16s(b * KP, 1.0)).expect("g2"); + let state_host = rng.f32s(b * K3_CAPSULE_STATE_SLOT, 0.05); + let state_dev = ctx.stream.clone_htod(&state_host).expect("state H2D"); + + // Native chain: three conv streams, then the core with a distinct + // successor state. + let mut xs = Vec::new(); + let mut ys = Vec::new(); + let mut sns = Vec::new(); + for s in 0..3 { + let mut x = ctx.stream.alloc_zeros::(b * KP).expect("x"); + let mut y = ctx.stream.alloc_zeros::(b * KP).expect("y"); + let mut sn = ctx.stream.alloc_zeros::(b * WS * KP).expect("sn"); + k3_conv_silu_batched_launch( + &ctx, + b, + KP, + K3_CONV_WIDTH, + 1, + &partials[s], + &cw[s], + &windows[s], + &mut x, + &mut y, + &mut sn, + ) + .expect("native conv_silu"); + xs.push(x); + ys.push(y); + sns.push(sn); + } + let mut state_n = ctx + .stream + .alloc_zeros::(b * K3_CAPSULE_STATE_SLOT) + .expect("state_n"); + let mut out_native = ctx.stream.alloc_zeros::(b * KP).expect("out"); + k3_kda_core_batched_launch( + &ctx, + b, + K3_KDA_HEADS, + K3_KDA_HEAD_DIM, + 1, + &ys[0], + &ys[1], + &ys[2], + &forget, + &dt_dev, + &alog_dev, + &beta_dev, + &g2_dev, + &go_dev, + &state_dev, + &mut state_n, + &mut out_native, + ) + .expect("native kda_core"); + + // Capsule side: packed x rows, packed conv slab, bf16 forget landing, + // fresh copy of the same initial state (updated in place). + let mut x_packed = ctx + .stream + .alloc_zeros::(b * K3_CAPSULE_X_ROW) + .expect("x_packed"); + for (s, x) in xs.iter().enumerate() { + for row in 0..b { + copy_bf16( + &ctx, + x, + row * KP, + &mut x_packed, + row * K3_CAPSULE_X_ROW + s * KP, + KP, + ); + } + } + let mut conv_packed = ctx + .stream + .alloc_zeros::(b * K3_CAPSULE_CONV_SLOT) + .expect("conv_packed"); + for (s, w) in windows.iter().enumerate() { + for row in 0..b { + for t in 0..WS { + copy_bf16( + &ctx, + w, + (row * WS + t) * KP, + &mut conv_packed, + row * K3_CAPSULE_CONV_SLOT + (t * 3 + s) * KP, + KP, + ); + } + } + } + let mut g_bf16 = ctx.stream.alloc_zeros::(b * KP).expect("g_bf16"); + k3_land_batched_launch(&ctx, b, KP, KP, 0, 1, &forget, &mut g_bf16).expect("g land"); + let mut state_capsule = ctx.stream.clone_htod(&state_host).expect("state2 H2D"); + let mut out_capsule = ctx.stream.alloc_zeros::(b * KP).expect("out2"); + k3_capsule_kda_decode_launch( + &ctx, + b, + &x_packed, + &cw[0], + &cw[1], + &cw[2], + &mut conv_packed, + &alog_dev, + &g_bf16, + &dt_dev, + &beta_dev, + &g2_dev, + &go_dev, + &mut state_capsule, + &mut out_capsule, + ) + .expect("capsule kda launch"); + + // Conv windows: the shifted slots and the shared bf16 x landing must + // match bitwise, per stream. + let packed = ctx.stream.clone_dtoh(&conv_packed).expect("conv D2H"); + for (s, sn) in sns.iter().enumerate() { + let native = ctx.stream.clone_dtoh(sn).expect("sn D2H"); + for row in 0..b { + for t in 0..WS { + let nat = &native[(row * WS + t) * KP..(row * WS + t + 1) * KP]; + let cap_at = row * K3_CAPSULE_CONV_SLOT + (t * 3 + s) * KP; + let cap = &packed[cap_at..cap_at + KP]; + assert_eq!( + nat, cap, + "b={b} stream {s} row {row} tap {t} window diverges" + ); + } + } + } + + // Recurrent state and output: rounding-chain tolerance. + let sn_native = ctx.stream.clone_dtoh(&state_n).expect("state D2H"); + let sn_capsule = ctx.stream.clone_dtoh(&state_capsule).expect("state2 D2H"); + let mut max_state = 0.0f32; + for (i, (&n, &c)) in sn_native.iter().zip(&sn_capsule).enumerate() { + let err = (n - c).abs() / n.abs().max(0.05); + assert!( + err < 0.08, + "b={b} state[{i}]: native {n} vs capsule {c} (rel {err})" + ); + max_state = max_state.max(err); + } + let out_n = ctx.stream.clone_dtoh(&out_native).expect("out D2H"); + let out_c = ctx.stream.clone_dtoh(&out_capsule).expect("out2 D2H"); + let mut max_out = 0.0f32; + for (i, (&n, &c)) in out_n.iter().zip(&out_c).enumerate() { + let (n, c) = (n.to_f32(), c.to_f32()); + let err = (n - c).abs() / n.abs().max(0.1); + assert!( + err < 0.08, + "b={b} out[{i}]: native {n} vs capsule {c} (rel {err})" + ); + max_out = max_out.max(err); + } + println!( + "b={b}: conv windows bitwise-equal; max rel err state {max_state:.4}, out {max_out:.4}" + ); + } +} diff --git a/pegainfer-kernels/tests/k3_capsule_topk_gate.rs b/pegainfer-kernels/tests/k3_capsule_topk_gate.rs new file mode 100644 index 000000000..592d7ecd7 --- /dev/null +++ b/pegainfer-kernels/tests/k3_capsule_topk_gate.rs @@ -0,0 +1,130 @@ +//! Numeric gate for the capsule-vendored vLLM router top-k +//! (`cubin/k3/single_group_topk_e512t22_sm103.cubin`) against the native +//! `k3_router_topk` kernel on identical logits. +//! +//! Both kernels take raw router logits, apply f32 sigmoid, select top-k over +//! `sigmoid + bias`, and renormalize the *un-biased* sigmoid scores of the +//! picks to the routed scale. They differ in output order (native emits +//! selection order, the capsule emits its own sort order) and in the exact +//! renorm-eps spelling, so the contract under test is order-free: per row, +//! the selected expert *sets* must be equal and each expert's weight must +//! agree to f32 rounding. Downstream (MegaMoE routing) treats the pairs as +//! unordered, so this is exactly the production contract. +//! +//! Manual gate: CI compiles this but never runs it. Run on a Blackwell box +//! (set `PEGAINFER_REQUIRE_GPU=1` to turn a missing device into a failure). + +#![cfg(feature = "k3")] + +mod common; + +use std::collections::HashMap; + +use half::bf16; +use pegainfer_kernels::ops::k3_capsule_router_topk_launch; +use pegainfer_kernels::ops::k3_router_topk_batched_launch; + +/// The pruned dev checkpoint's expert count — the only K3 shape inside the +/// capsule tier (<=512 experts; the 896-expert full model stays native). +const EXPERTS: usize = 224; +const TOPK: usize = 16; +const ROUTED_SCALE: f32 = 2.5; + +struct Lcg(u64); +impl Lcg { + fn next_u32(&mut self) -> u32 { + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (self.0 >> 32) as u32 + } + fn unit_f32(&mut self) -> f32 { + (self.next_u32() as f32 / u32::MAX as f32) * 2.0 - 1.0 + } +} + +#[test] +fn capsule_topk_matches_native_selection() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let mut rng = Lcg(0x4b33_2026_0828); + // Match the native kernel's exact scale spelling: it reads rs as bf16 and + // widens; hand the capsule the identical widened value. + let rs = bf16::from_f32(ROUTED_SCALE); + let rs_dev = ctx.stream.clone_htod(&[rs]).expect("rs H2D"); + let bias: Vec = (0..EXPERTS).map(|_| rng.unit_f32() * 0.02).collect(); + let bias_dev = ctx.stream.clone_htod(&bias).expect("bias H2D"); + + // Production decode buckets plus an odd row count for the capsule's + // 8-rows-per-block grid tail. + for &b in &[1usize, 3, 8, 32] { + let logits: Vec = (0..b * EXPERTS).map(|_| rng.unit_f32() * 4.0).collect(); + let logits_dev = ctx.stream.clone_htod(&logits).expect("logits H2D"); + + let mut native_idx = ctx.stream.alloc_zeros::(b * TOPK).expect("idx alloc"); + let mut native_wts = ctx.stream.alloc_zeros::(b * TOPK).expect("wts alloc"); + k3_router_topk_batched_launch( + &ctx, + b, + EXPERTS, + TOPK, + &logits_dev, + &bias_dev, + &rs_dev, + &mut native_idx, + &mut native_wts, + ) + .expect("native topk launch"); + + let mut cap_idx = ctx.stream.alloc_zeros::(b * TOPK).expect("idx alloc"); + let mut cap_wts = ctx.stream.alloc_zeros::(b * TOPK).expect("wts alloc"); + k3_capsule_router_topk_launch( + &ctx, + b, + EXPERTS, + TOPK, + &logits_dev, + &bias_dev, + rs.to_f32(), + &mut cap_idx, + &mut cap_wts, + ) + .expect("capsule topk launch"); + + let n_idx = ctx.stream.clone_dtoh(&native_idx).expect("D2H"); + let n_wts = ctx.stream.clone_dtoh(&native_wts).expect("D2H"); + let c_idx = ctx.stream.clone_dtoh(&cap_idx).expect("D2H"); + let c_wts = ctx.stream.clone_dtoh(&cap_wts).expect("D2H"); + + for t in 0..b { + let native: HashMap = (0..TOPK) + .map(|r| (n_idx[t * TOPK + r], n_wts[t * TOPK + r])) + .collect(); + let capsule: HashMap = (0..TOPK) + .map(|r| (c_idx[t * TOPK + r], c_wts[t * TOPK + r])) + .collect(); + assert_eq!(native.len(), TOPK, "b={b} row {t}: native emitted a dup"); + assert_eq!(capsule.len(), TOPK, "b={b} row {t}: capsule emitted a dup"); + let mut native_experts: Vec = native.keys().copied().collect(); + let mut capsule_experts: Vec = capsule.keys().copied().collect(); + native_experts.sort_unstable(); + capsule_experts.sort_unstable(); + assert_eq!( + native_experts, capsule_experts, + "b={b} row {t}: expert sets diverge" + ); + for (&e, &nw) in &native { + let cw = capsule[&e]; + // Weights are O(rs / topk); the only legal differences are + // sigmoid/renorm rounding and the 1e-20 eps spelling. + assert!( + (nw - cw).abs() <= 1e-5 * nw.abs().max(1.0), + "b={b} row {t} expert {e}: native weight {nw} vs capsule {cw}" + ); + } + } + } + println!("capsule topk == native topk for b in [1, 3, 8, 32] at E={EXPERTS}, topk={TOPK}"); +} diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index e02125ea5..ca72c5ff5 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -108,12 +108,13 @@ pub(crate) fn start_qwen3( if let Some(path) = dump_graph_png { let summary = executor.dump_decode_graph_png(path)?; info!( - "Qwen3 decode CUDA Graph exported: nodes={}, kernels={}, edges={}, dot={}, png={}", + "Qwen3 decode CUDA Graph exported: nodes={}, kernels={}, edges={}, dot={}, png={}, json={}", summary.nodes, summary.kernels, summary.edges, summary.dot_path.display(), - summary.png_path.display() + summary.png_path.display(), + summary.json_path.display() ); } executor.set_no_prefix_cache(no_prefix_cache); diff --git a/pegainfer-qwen3/src/lib.rs b/pegainfer-qwen3/src/lib.rs index cb16e0b52..bdf59e26a 100644 --- a/pegainfer-qwen3/src/lib.rs +++ b/pegainfer-qwen3/src/lib.rs @@ -234,7 +234,9 @@ pub struct Qwen3LaunchOptions { /// under tensor parallelism every decode graph is pre-captured at startup. pub cuda_graph: bool, /// Export the live rank-0, batch-1 SplitKv decode graph during startup. - /// The requested PNG gets a detailed sibling `.dot` for LLM inspection. + /// The requested PNG gets a detailed sibling `.dot` for LLM inspection and + /// a sibling `.json` with per-kernel attributes and staged parameter bytes + /// for machine consumption. pub dump_graph_png: Option, pub offload: Qwen3OffloadOptions, pub no_prefix_cache: bool, diff --git a/tools/kernel-capture/.gitignore b/tools/kernel-capture/.gitignore new file mode 100644 index 000000000..9ff660eb5 --- /dev/null +++ b/tools/kernel-capture/.gitignore @@ -0,0 +1,2 @@ +libkernelcapture*.so +kernel-capture/ diff --git a/tools/kernel-capture/build.sh b/tools/kernel-capture/build.sh new file mode 100644 index 000000000..9d8bfd775 --- /dev/null +++ b/tools/kernel-capture/build.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Build the CUPTI kernel-capture injection library. +set -euo pipefail + +here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cuda="${CUDA_HOME:-/usr/local/cuda}" +target_dir="$(echo "$cuda"/targets/*-linux)" +cupti_inc="$target_dir/include" +cupti_lib="$target_dir/lib" + +cc -O2 -fPIC -shared \ + -I"$cupti_inc" -I"$cuda/include" \ + "$here/capture.c" \ + -o "$here/libkernelcapture.so" \ + -L"$cupti_lib" -lcupti -lcuda + +echo "built $here/libkernelcapture.so" diff --git a/tools/kernel-capture/capture.c b/tools/kernel-capture/capture.c new file mode 100644 index 000000000..aa253688f --- /dev/null +++ b/tools/kernel-capture/capture.c @@ -0,0 +1,413 @@ +// CUPTI injection library that mines every CUDA module and kernel launch from +// a host process it is loaded into (vLLM, sglang, or PegaInfer itself), so a +// launched kernel's cubin, launch configuration, and staged parameter bytes +// can be lifted out for replay. Provider-agnostic: Triton, cuBLAS/cuBLASLt, +// CUTLASS, FlashInfer, and hand-written CUDA all surface as module-load events +// with real cubin bytes, not just Triton's on-disk cache. +// +// Load it with the CUDA driver's injection hook: +// +// CUDA_INJECTION64_PATH=/path/to/libkernelcapture.so \ +// KERNEL_CAPTURE_DIR=/some/out \ +// python -m vllm.entrypoints.openai.api_server ... +// +// Output under KERNEL_CAPTURE_DIR (default ./kernel-capture), in a pid/ +// subdirectory per injected process so multi-rank engines (TP/EP workers) +// don't clobber each other: +// pid/module_.cubin one file per distinct loaded module +// pid/launches.jsonl one JSON object per captured kernel launch +// +// The driver calls InitializeInjection() once, before the first CUDA call, on +// any library named by CUDA_INJECTION64_PATH. + +#define _GNU_SOURCE +#include +#include +// cuLaunchKernel_params / cuLaunchKernelEx_params come from +// generated_cuda_meta.h, which cupti.h already includes transitively. + +#include +#include +#include +#include +#include +#include +#include +#include + +// A kernel's parameter buffer is a few KiB at most; a walk past this many +// indices means cuFuncGetParamInfo is misbehaving, not that the kernel is huge. +#define MAX_PARAMS 4096 + +#define MAX_ABI_PARAMS 128 +#define MAX_SYMBOL 640 + +static CUpti_SubscriberHandle g_subscriber; +static char g_out_dir[4096]; +static FILE *g_launches; +static pthread_mutex_t g_lock = PTHREAD_MUTEX_INITIALIZER; + +// Bounded set of module ids already written, so a module loaded once but +// launched from thousands of times is dumped a single time. +static uint32_t g_seen_modules[65536]; +static size_t g_seen_count; + +// One kernel's parameter layout, learned from the driver's own parse of the +// cubin at module-load time (see cache_module_abi). Used at launch to size the +// staged kernelParams array for kernels whose launch-time CUfunction does not +// support cuFuncGetParamInfo — every PyTorch/vLLM kernel launched through the +// runtime `<<<>>>` path. +typedef struct { + int num_regs, static_shared, const_bytes, local_bytes; + int ptx_version, binary_version; +} KernelAttrs; + +typedef struct { + char symbol[MAX_SYMBOL]; + int nparams; + size_t offset[MAX_ABI_PARAMS]; + size_t size[MAX_ABI_PARAMS]; + KernelAttrs attrs; +} KernelAbi; + +static KernelAttrs read_attrs(CUfunction func) { + KernelAttrs a = {0}; + cuFuncGetAttribute(&a.num_regs, CU_FUNC_ATTRIBUTE_NUM_REGS, func); + cuFuncGetAttribute(&a.static_shared, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, func); + cuFuncGetAttribute(&a.const_bytes, CU_FUNC_ATTRIBUTE_CONST_SIZE_BYTES, func); + cuFuncGetAttribute(&a.local_bytes, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, func); + cuFuncGetAttribute(&a.ptx_version, CU_FUNC_ATTRIBUTE_PTX_VERSION, func); + cuFuncGetAttribute(&a.binary_version, CU_FUNC_ATTRIBUTE_BINARY_VERSION, func); + return a; +} + +static KernelAbi *g_abi; +static size_t g_abi_count; +static size_t g_abi_cap; +// cuModuleLoadData below triggers another MODULE_LOADED callback on this +// thread; the guard keeps that recursion from re-entering the loader. +static __thread int g_in_self_load; + +static const KernelAbi *abi_lookup(const char *symbol) { + if (!symbol) { + return NULL; + } + for (size_t i = 0; i < g_abi_count; i++) { + if (strcmp(g_abi[i].symbol, symbol) == 0) { + return &g_abi[i]; + } + } + return NULL; +} + +static KernelAbi *abi_new(void) { + if (g_abi_count == g_abi_cap) { + size_t cap = g_abi_cap ? g_abi_cap * 2 : 1024; + KernelAbi *grown = realloc(g_abi, cap * sizeof(*grown)); + if (!grown) { + return NULL; + } + g_abi = grown; + g_abi_cap = cap; + } + return &g_abi[g_abi_count++]; +} + +// Load a private copy of the just-loaded module and record every kernel's +// parameter layout. The driver parses the cubin for us, so this works for +// kernels the runtime registers, which the launch-time handle cannot answer. +static void cache_module_abi(const CUpti_ModuleResourceData *module) { + g_in_self_load = 1; + CUmodule mod = NULL; + if (cuModuleLoadData(&mod, module->pCubin) != CUDA_SUCCESS) { + g_in_self_load = 0; + return; + } + unsigned int count = 0; + if (cuModuleGetFunctionCount(&count, mod) != CUDA_SUCCESS || count == 0) { + cuModuleUnload(mod); + g_in_self_load = 0; + return; + } + CUfunction *funcs = calloc(count, sizeof(*funcs)); + if (funcs && cuModuleEnumerateFunctions(funcs, count, mod) == CUDA_SUCCESS) { + for (unsigned int i = 0; i < count; i++) { + const char *name = NULL; + if (cuFuncGetName(&name, funcs[i]) != CUDA_SUCCESS || !name) { + continue; + } + KernelAbi *abi = abi_new(); + if (!abi) { + break; + } + snprintf(abi->symbol, sizeof(abi->symbol), "%s", name); + abi->attrs = read_attrs(funcs[i]); + abi->nparams = 0; + for (size_t p = 0; p < MAX_ABI_PARAMS; p++) { + size_t offset = 0, size = 0; + if (cuFuncGetParamInfo(funcs[i], p, &offset, &size) != CUDA_SUCCESS) { + break; + } + abi->offset[p] = offset; + abi->size[p] = size; + abi->nparams++; + } + } + } + free(funcs); + cuModuleUnload(mod); + g_in_self_load = 0; +} + +static int module_already_seen(uint32_t module_id) { + for (size_t i = 0; i < g_seen_count; i++) { + if (g_seen_modules[i] == module_id) { + return 1; + } + } + if (g_seen_count < sizeof(g_seen_modules) / sizeof(g_seen_modules[0])) { + g_seen_modules[g_seen_count++] = module_id; + } + return 0; +} + +static void write_cubin(const CUpti_ModuleResourceData *module) { + if (module_already_seen(module->moduleId)) { + return; + } + char path[4200]; + snprintf(path, sizeof(path), "%s/module_%u.cubin", g_out_dir, + module->moduleId); + FILE *f = fopen(path, "wb"); + if (!f) { + return; + } + fwrite(module->pCubin, 1, module->cubinSize, f); + fclose(f); +} + +static void write_hex(FILE *out, const unsigned char *bytes, size_t len) { + static const char digits[] = "0123456789abcdef"; + for (size_t i = 0; i < len; i++) { + fputc(digits[bytes[i] >> 4], out); + fputc(digits[bytes[i] & 0xf], out); + } +} + +// Resolve an 8-byte parameter value against the CUDA allocation map and, if it +// is a live pointer, append a "pointer" object with its memory type and owning +// range — the signal that separates device/kv/weight pointers from scalars for +// downstream binding inference. Advisory: an integer that collides with an +// allocation also matches. +static void write_pointer_field(FILE *out, const unsigned char *bytes) { + uint64_t value; + memcpy(&value, bytes, sizeof(value)); + if (value == 0) { + return; + } + unsigned int memory_type = 0; + if (cuPointerGetAttribute(&memory_type, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, + (CUdeviceptr)value) != CUDA_SUCCESS) { + return; + } + const char *type_name = memory_type == CU_MEMORYTYPE_HOST ? "host" + : memory_type == CU_MEMORYTYPE_DEVICE ? "device" + : memory_type == CU_MEMORYTYPE_ARRAY ? "array" + : memory_type == CU_MEMORYTYPE_UNIFIED ? "unified" + : "unknown"; + uint64_t range_start = 0; + size_t range_size = 0; + cuPointerGetAttribute(&range_start, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, + (CUdeviceptr)value); + cuPointerGetAttribute(&range_size, CU_POINTER_ATTRIBUTE_RANGE_SIZE, + (CUdeviceptr)value); + fprintf(out, + ",\"pointer\":{\"memory_type\":\"%s\",\"range_start\":\"0x%llx\"," + "\"range_size\":%zu}", + type_name, (unsigned long long)range_start, range_size); +} + +// Fill offset[]/size[] for a launch, preferring the live handle and falling +// back to the module-load ABI cache (which the runtime `<<<>>>` launch handle +// cannot answer). Returns the parameter count, or -1 if neither source knows. +static int resolve_layout(CUfunction func, const char *symbol, size_t *offset, + size_t *size) { + for (int index = 0; index < MAX_PARAMS; index++) { + size_t off = 0, sz = 0; + if (cuFuncGetParamInfo(func, index, &off, &sz) != CUDA_SUCCESS) { + if (index > 0) { + return index; + } + break; + } + if ((size_t)index >= MAX_ABI_PARAMS) { + return index; + } + offset[index] = off; + size[index] = sz; + } + const KernelAbi *abi = abi_lookup(symbol); + if (!abi) { + return -1; + } + for (int index = 0; index < abi->nparams; index++) { + offset[index] = abi->offset[index]; + size[index] = abi->size[index]; + } + return abi->nparams; +} + +// Emit one launches.jsonl record: symbol, launch geometry, function +// attributes, and every staged parameter value read through the resolved +// parameter layout. +static void record_launch(const char *symbol, CUfunction func, + unsigned int grid_x, unsigned int grid_y, + unsigned int grid_z, unsigned int block_x, + unsigned int block_y, unsigned int block_z, + unsigned int shared_bytes, void **kernel_params) { + KernelAttrs attrs = read_attrs(func); + + pthread_mutex_lock(&g_lock); + if (attrs.num_regs == 0) { + const KernelAbi *abi = abi_lookup(symbol); + if (abi) { + attrs = abi->attrs; + } + } + fprintf(g_launches, "{\"symbol\":\"%s\",", symbol ? symbol : ""); + fprintf(g_launches, "\"grid\":[%u,%u,%u],\"block\":[%u,%u,%u],", + grid_x, grid_y, grid_z, block_x, block_y, block_z); + fprintf(g_launches, "\"dynamic_shared_mem_bytes\":%u,", shared_bytes); + fprintf(g_launches, + "\"attributes\":{\"num_regs\":%d,\"static_shared_bytes\":%d," + "\"const_bytes\":%d,\"local_bytes\":%d,\"ptx_version\":%d," + "\"binary_version\":%d},", + attrs.num_regs, attrs.static_shared, attrs.const_bytes, + attrs.local_bytes, attrs.ptx_version, attrs.binary_version); + + // A null kernelParams array means the launch passed arguments through the + // packed `extra` buffer (cuBLASLt nvjet kernels do this); the layout walk + // does not apply, so the parameter list is explicitly null. + if (!kernel_params) { + fprintf(g_launches, "\"params\":null}\n"); + fflush(g_launches); + pthread_mutex_unlock(&g_lock); + return; + } + + size_t offset[MAX_ABI_PARAMS], size[MAX_ABI_PARAMS]; + int nparams = resolve_layout(func, symbol, offset, size); + fprintf(g_launches, "\"params\":"); + if (nparams < 0) { + // kernelParams is present but neither the handle nor the cache knows its + // length; emitting a bounded blind walk would risk reading past the array. + fprintf(g_launches, "\"unknown-layout\"}\n"); + fflush(g_launches); + pthread_mutex_unlock(&g_lock); + return; + } + fputc('[', g_launches); + for (int index = 0; index < nparams; index++) { + const unsigned char *staged = (const unsigned char *)kernel_params[index]; + if (index > 0) { + fputc(',', g_launches); + } + fprintf(g_launches, "{\"offset\":%zu,\"size\":%zu,\"data\":\"", offset[index], + size[index]); + if (staged) { + write_hex(g_launches, staged, size[index]); + } + fputc('"', g_launches); + if (staged && size[index] == 8) { + write_pointer_field(g_launches, staged); + } + fputc('}', g_launches); + } + fprintf(g_launches, "]}\n"); + fflush(g_launches); + pthread_mutex_unlock(&g_lock); +} + +static void CUPTIAPI callback(void *userdata, CUpti_CallbackDomain domain, + CUpti_CallbackId cbid, const void *cbdata) { + (void)userdata; + if (domain == CUPTI_CB_DOMAIN_RESOURCE) { + // Our own cuModuleLoadData in cache_module_abi re-enters this callback; + // skip that recursion. + if (cbid == CUPTI_CBID_RESOURCE_MODULE_LOADED && !g_in_self_load) { + const CUpti_ResourceData *resource = (const CUpti_ResourceData *)cbdata; + const CUpti_ModuleResourceData *module = + (const CUpti_ModuleResourceData *)resource->resourceDescriptor; + // Multi-rank engines (one process, one thread per GPU) hit this path + // concurrently on lazy module loads; g_seen_* and the g_abi table are + // shared, so the whole module path holds the same lock as the readers. + // The recursive MODULE_LOADED from our own cuModuleLoadData is filtered + // by g_in_self_load above, before it could reach the lock. + pthread_mutex_lock(&g_lock); + write_cubin(module); + cache_module_abi(module); + pthread_mutex_unlock(&g_lock); + } + return; + } + if (domain != CUPTI_CB_DOMAIN_DRIVER_API) { + return; + } + const CUpti_CallbackData *data = (const CUpti_CallbackData *)cbdata; + // Read on the way *out*: under lazy module loading (the driver default since + // CUDA 12.x) a kernel's module is materialized during the launch call, so on + // the ENTER site cuFuncGetParamInfo/cuFuncGetAttribute return zeros. By EXIT + // the function is loaded and queryable, and the caller's kernelParams array + // — pointed at by functionParams — is still live within this callback. + if (data->callbackSite != CUPTI_API_EXIT) { + return; + } + if (cbid == CUPTI_DRIVER_TRACE_CBID_cuLaunchKernel) { + const cuLaunchKernel_params *p = + (const cuLaunchKernel_params *)data->functionParams; + record_launch(data->symbolName, p->f, p->gridDimX, p->gridDimY, p->gridDimZ, + p->blockDimX, p->blockDimY, p->blockDimZ, p->sharedMemBytes, + p->kernelParams); + } else if (cbid == CUPTI_DRIVER_TRACE_CBID_cuLaunchKernelEx) { + const cuLaunchKernelEx_params *p = + (const cuLaunchKernelEx_params *)data->functionParams; + const CUlaunchConfig *cfg = p->config; + record_launch(data->symbolName, p->f, cfg->gridDimX, cfg->gridDimY, + cfg->gridDimZ, cfg->blockDimX, cfg->blockDimY, cfg->blockDimZ, + cfg->sharedMemBytes, p->kernelParams); + } +} + +// The CUDA driver's injection entry point. Named exactly this; the driver +// dlsym's it out of every library on CUDA_INJECTION64_PATH. +int InitializeInjection(void) { + const char *dir = getenv("KERNEL_CAPTURE_DIR"); + char base_dir[4096]; + snprintf(base_dir, sizeof(base_dir), "%s", + dir && *dir ? dir : "./kernel-capture"); + mkdir(base_dir, 0755); + snprintf(g_out_dir, sizeof(g_out_dir), "%.4080s/pid%d", base_dir, + (int)getpid()); + mkdir(g_out_dir, 0755); + + char path[4200]; + snprintf(path, sizeof(path), "%s/launches.jsonl", g_out_dir); + g_launches = fopen(path, "w"); + if (!g_launches) { + fprintf(stderr, "[kernel-capture] cannot open %s\n", path); + return 0; + } + + if (cuptiSubscribe(&g_subscriber, (CUpti_CallbackFunc)callback, NULL) != + CUPTI_SUCCESS) { + fprintf(stderr, "[kernel-capture] cuptiSubscribe failed\n"); + return 0; + } + cuptiEnableDomain(1, g_subscriber, CUPTI_CB_DOMAIN_RESOURCE); + cuptiEnableCallback(1, g_subscriber, CUPTI_CB_DOMAIN_DRIVER_API, + CUPTI_DRIVER_TRACE_CBID_cuLaunchKernel); + cuptiEnableCallback(1, g_subscriber, CUPTI_CB_DOMAIN_DRIVER_API, + CUPTI_DRIVER_TRACE_CBID_cuLaunchKernelEx); + fprintf(stderr, "[kernel-capture] active, writing to %s\n", g_out_dir); + return 1; +}