diff --git a/docs/index.md b/docs/index.md index 021ca864a..332dc0bc8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -102,7 +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/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/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. 2026-09-01: five CP8 128k anatomy cuts (vectorized landing, in-place conv window, striped FMHA triangle, vectorized lse_merge + bf16 elementwise family) 6,387 → 5,021 ms. | ## models / deepseek-v2-lite diff --git a/docs/models/k3/cp-lane-design.md b/docs/models/k3/cp-lane-design.md index 33834e9c4..5b79df560 100644 --- a/docs/models/k3/cp-lane-design.md +++ b/docs/models/k3/cp-lane-design.md @@ -26,7 +26,7 @@ CP1 13.9s，CP8/CP16 只差 gang 侧**）→ M1 full@EP16 交叉矩阵 + lane-vs-独占步系统 A/B → M2 agent cache 闭环 → M3 弹性调度。 -Last touched: 2026-08 +Last touched: 2026-09 ## 0. 与 mix-engine-design 的关系 @@ -686,9 +686,51 @@ vLLM 13,155 = **1.49×**，快过 pruned CP16 的 9,316）；64k/128k 与 W-chun 逐 ms 持平（1,717/3,618 vs 1,734/3,615，零回归）；249,891 tok + 48 greedy 生成连贯。验收表 256k 行补齐——**64k 及以上全档反超 vLLM TP16-MNNVL**。 +## CP8 128k 三刀（2026-09-01，`feat/k3-cp8-anatomy-cuts`，账在 +`~/code/bench_results/2026-09-01-k3-cp8-128k-gap-anatomy/`） + +nsys 单 rank 剖析（tray14 深位 rank，128k superstep 6,130 ms）推翻了两个旧假设：间隙不是 host +走时（host-starved 0 ms，全是 launch queue 反压与 doorbell 等待），FMHA kernel 本身在峰 +（1.47–1.53 PF/s），2,132 ms 里 ~1,090 ms 纯粹是深位不均衡。按确定性下刀，每刀一个 commit， +门禁全绿（golden 13/13 逐步精确、paged 3/3、spec 6/6、cp_prefill 2/2）： + +1. **land 向量化**（`e1fdd881`）：TileLang `land_batched` 每线程 1 元素、2 字节存，12288 宽落地 + 只跑 2.3 TB/s；换成 8 列/线程的 CUDA 核（同 cast、同顺序，逐位相同）→ 550→190 µs @16896。 + **GEMM bf16 epilogue 方案被数值否决**：cuBLAS 对 bf16 输出选不同 kernel，累加顺序变， + 4 层 golden 在 3-ULP margin 的步翻转（对 reference 的偏差分布与 main 完全一致——median 1 / + p90 3 / max 12 ULP——但门禁的"margin>2 ULP 逐 token 精确"契约绑定 kernel 选择）。 +2. **conv 原地窗口**（`dad55e4b`）：chunk 路径原本 land → 3 个 tap 的 `cuMemcpy2DAsync` 窗口 + 物化 → 批 conv → carry 拷贝；新核直接读 f32 partial 的 t-3..t 行（段首取 carry），16 行/block + 寄存器滑窗、taps 常驻寄存器；算式逐项复刻（同 -O3 无 fast-math）。CP4 64k：DtoD 拷贝 25,466 + 次/2,593 ms → 4,766/102，conv 877→476 µs。 +3. **FMHA 条带化**（`80f12bb8`）：off-diagonal 段对 (q,k) 奇差由 owner 算、偶差由 key 持有者算 + （它本就为自己的对角展开了 K）；每层 owner 发布一次 Q（搭 latent 窗口），helper 的 FMHA 经 + fabric TMA 直读 Q、结果写双缓冲回传 slab，owner 的 lse_merge 原地读 peer slab；每层 + 1 + ⌊(R−1)/2⌋ 个固定窗口（Stripe kind 带显式 rank 位掩码）。负载 r+½ → ≤(r+3)/2（8 rank： + 深位 7.5→4.5，理想 4）。CP4 64k：深位 FMHA 2,408→1,733（精确等于 2.5/3.5），墙钟 + 4,171→3,888。**fleet CP8 128k：TTFT 6,387 → 5,328 ms（−16.6%），65k 2,810 → 2,394**； + superstep 6,130 → 5,071。slab 每 rank +1.47 GB（Q 623 MB + 2 × 回传 421 MB @seg_cap 16896）。 + +after anatomy（dev2）：fmha 1,516 / mega 1,293 / nvjet 920 / lse_merge 220 / attnres 181 / kda 277 / +gaps 144（全 doorbell）。 + +4. **lse_merge 向量化**（`003b5b5b`）：一 (q,h) 一个 128 线程 block、1 元素/线程 → 16 lane 一行、 + 8 列/线程；逐元素算式不动。CP4 64k 3,888→3,785。 +5. **bf16 逐元素家族向量化**（`e82f072d`）：add2/mul_sigmoid/situ/o_norm_gate 从 TileLang（1 元素/ + 线程，1–2 TB/s）换成 CUDA；TileLang 的 `bfloat16_t` 是 cutlass 的（`+`=`__hadd`、`*`=`__hmul`、 + cast=`cvt.rn`），o_norm_gate 的 128 宽 xor butterfly（64/32 走 smem、16..1 走 shuffle）用 lane + xor 8/4/2/1 + 槽位 j^4/j^2/j^1 逐对复刻——golden 13/13 逐步精确证明 decode 路径也逐位相同。 + 退休 40 个实例化。CP4 64k 3,785→3,639。 + +**fleet CP8 五刀累计：128k 6,387 → 5,021 ms（−21.4%），65k 2,810 → 2,231，16k 1,150 → 946**； +128k 后生成逐字接续原文。剩余（按三刀后 anatomy 推算）：attnres 181（向量化 ~80）、doorbell 144、 +mega 拆账、fmha 残余不均衡（helped 段 dense FMHA 比 causal 慢，4.5 vs 理想 4）。 + ## Next action -PR #970 CI 17/17 全绿（`05da7961`），待 susun review。 -性能杠杆按 ⑥ 实测定序不变：FMHA 条带化（merge 原语已在）→ superstep -图化/融合 → **bucket 细化（8k/16k 档翻盘的主杠杆，验收表的 0.39×/0.78× 就是 -它）** → 协调压缩。KDA 包 prefix-scan 排后。 +#970 已合入 main（`1e3b6f23`）。五刀分支 `feat/k3-cp8-anatomy-cuts`（land 向量化 / conv +原地窗口 / FMHA 条带化 / lse_merge 向量化 / 逐元素家族向量化）待 PR：128k 6,387 → 5,021 ms。 +未决：同进程 `cp_prefill` 两测试连跑的 DeepGEMM grid-sync 超时 flake（分进程必过；main 的 +worktree `~/agent_code/wt-main-k3` 已构建，待跑 3 次定归属）。杠杆顺序：attnres 向量化（~80 ms） +→ mega 拆账（solo 微基准）→ **bucket 细化（8k/16k 档翻盘的主杠杆，验收表 +的 0.39×/0.78× 就是它）** → 协调压缩。KDA 包 prefix-scan 排后。 diff --git a/pegainfer-k3/kernels/README.md b/pegainfer-k3/kernels/README.md index 84e285393..4facc0f2f 100644 --- a/pegainfer-k3/kernels/README.md +++ b/pegainfer-k3/kernels/README.md @@ -31,28 +31,24 @@ every row below is its shape count × 10. | Kernel | Shapes per bucket | Instantiations | Launcher | | --- | --- | --- | --- | | `rms_norm_rbs_batched` | H ∈ {7168, 512, 3584} | 30 | `k3_rms_norm_rbs_batched` | -| `land_batched` | 14 (NT, N, OFF) spans, SK = 1 | 140 | `k3_land_batched` | | `land_rms_norm_rbs_batched` | MLA q_a, SK = 1 | 10 | `k3_land_rms_norm_rbs_batched` | -| `add2_batched` | N = 7168 | 10 | `k3_add2_batched` | -| `mul_sigmoid_batched` | N = 12288 | 10 | `k3_mul_sigmoid_batched` | -| `situ_batched` | N ∈ {6144, 33792} | 20 | `k3_situ_batched` | | `conv_silu_batched` | KP = 12288, W = 4, SK = 1 | 10 | `k3_conv_silu_batched` | | `kda_core_batched` | 96 heads × 128 head_dim | 10 | `k3_kda_core_batched` | | `router_topk_batched` | E ∈ {896, 224}, TOPK = 16 | 20 | `k3_router_topk_batched` | | `attnres_scores_batched` | NB ∈ 1..8, H = 7168 | 80 | `k3_attnres_scores_batched` | | `attnres_mix_batched` | NB ∈ 1..8, H = 7168 | 80 | `k3_attnres_mix_batched` | -**420 instantiations**, about 20 seconds of generation and 7 seconds of nvcc. +**240 instantiations**, about 20 seconds of generation and 7 seconds of nvcc. The pool fans out at *instantiation* granularity, not family granularity — the families differ by more than an order of magnitude in size, so a family-granular -pool would be bound by `land_batched` alone. It defaults to one worker per CPU +pool would be bound by the largest family alone. It defaults to one worker per CPU capped at 32; each worker holds a TileLang lowering, so lower it with `PEGAINFER_K3_TILELANG_JOBS` on memory-tight hosts. One list in `generate.py` is deliberately narrow and is a one-line change: -* `SPLIT_K` — the segment counts the partial consumers (`land`, - `land_rms_norm_rbs`, `conv_silu`) accept. Only `1` — the single partial a +* `SPLIT_K` — the segment counts the partial consumers (`land_rms_norm_rbs`, + `conv_silu`) accept. Only `1` — the single partial a framework GEMM produces — has a launch site; the reference engine's split-K-8 GEMV shapes are not generated. @@ -67,9 +63,10 @@ absorbed paged-KV kernel (`csrc/k3/k3_mla_paged_attn.cu` — a runtime page walk needs no per-capacity instantiation, which is what retired the upstream `mla_attn` family and its `MAX_CTX` list), so the upstream `gemv`, `expert_gemv`, `packed_expert_gemv` and `mla_attn` kernels would be dead -weight. Their consumers are still -here — `land_batched` at `SK = 1` is exactly the bf16 landing of a cuBLASLt -f32 output. +weight. Nor the matmul landing itself: `land_batched` was retired for the +hand-written `csrc/k3/k3_land.cu` (same arithmetic, 8 columns per thread, +batch a runtime value) once the chunked-prefill anatomy showed the +one-element-per-thread kernel at a third of HBM rate. TileLang always names the entry point `main_kernel`, so every instantiation is renamed to a shape-tagged symbol before the sources are concatenated. Each diff --git a/pegainfer-k3/kernels/generate.py b/pegainfer-k3/kernels/generate.py index b2fa7b299..c1315084b 100644 --- a/pegainfer-k3/kernels/generate.py +++ b/pegainfer-k3/kernels/generate.py @@ -5,10 +5,8 @@ decode kernel set: k3_rms_norm_rbs_batched.cu k3_conv_silu_batched.cu - k3_land_batched.cu k3_kda_core_batched.cu - k3_land_rms_norm_rbs_batched.cu k3_attnres_scores_batched.cu - k3_add2_batched.cu k3_attnres_mix_batched.cu - k3_mul_sigmoid_batched.cu k3_situ_batched.cu + k3_land_rms_norm_rbs_batched.cu k3_kda_core_batched.cu + k3_attnres_scores_batched.cu k3_attnres_mix_batched.cu The batch size is a static compile-time dimension, so a single-stream step is served by the `B = 1` instantiation of the same family — its per-row spelling @@ -106,7 +104,7 @@ # the same artifact serves a single-GPU and a 4-way-EP deployment. EXPERTS = [896, 896 // 4] # engine E, and engine Es under 4-way EP -# Segment counts of the partials `land`/`conv_silu` merge. The engine's +# Segment counts of the partials `conv_silu` (and the hand-written landing) merge. The engine's # producers are framework GEMMs (cuBLASLt, DeepGEMM), which emit a single # segment, so only SK=1 is instantiated; the reference engine's SK=8 GEMV # shapes have no launch site here and are not generated. @@ -150,40 +148,16 @@ # the routed-latent norm (LAT). RMS_NORM_N = [HIDDEN, KV_LORA, LATENT] -# land(NT, N, OFF, SK): merge one column span of a (SK, NT) partial and land -# bf16 once. The engine's `lands` list, verbatim, plus the chunked-prefill -# conv-input landing (the sequential engine never lands that projection alone — -# its conv kernel casts in place; the chunk builds windows from the landed rows -# before the conv runs, so it needs the standalone cast). -LAND_CONFIGS = [ - # NT N OFF engine call site - (4 * KDA_DIM, KDA_DIM, 3 * KDA_DIM), # KDA output gate - (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 - (MLA_FUSED, KV_LORA + ROPE_DIM, Q_LORA), # MLA kv_a|k_rope - (MLA_FUSED, KDA_DIM, Q_LORA + KV_LORA + ROPE_DIM), # MLA output gate - (MLA_HEADS * QK_DIM, MLA_HEADS * QK_DIM, 0), # MLA q_b - (MLA_HEADS * 256, MLA_HEADS * 256, 0), # MLA kv_b - (HIDDEN, HIDDEN, 0), # o_proj / routed / shared - (LATENT, LATENT, 0), # routed latent - (2 * SHARED_INTER, SHARED_INTER, 0), # shared gate - (2 * SHARED_INTER, SHARED_INTER, SHARED_INTER), # shared up - (2 * DENSE_INTER, DENSE_INTER, 0), # dense gate - (2 * DENSE_INTER, DENSE_INTER, DENSE_INTER), # dense up - (VOCAB, VOCAB, 0), # logits -] +# land(NT, N, OFF, SK) is no longer generated: the matmul landing is the +# hand-written `csrc/k3/k3_land.cu` (batch a runtime value). # land_rms_norm_rbs(NT, N, OFF, SK, eps): MLA's q_a, the one place a merge and # a round-before-scale norm are fused. LAND_RMS_NORM_CONFIGS = [(MLA_FUSED, Q_LORA, 0)] -# add2 / mul_sigmoid / situ / conv_silu / kda_core take a single width each; -# situ has two (shared, dense) — the routed-expert situ is fused into the -# masked-GEMM chain and the mega kernel. -ADD2_N = [HIDDEN] -MUL_SIGMOID_N = [KDA_DIM] -SITU_N = [SHARED_INTER, DENSE_INTER] +# conv_silu / kda_core take a single width each. The bf16 elementwise family +# (add2, mul_sigmoid, situ, o_norm_gate) is hand-written CUDA now +# (`csrc/k3/k3_elementwise.cu`), shape-agnostic. # --------------------------------------------------------------------------- # @@ -199,19 +173,10 @@ f"(const {BF16}* __restrict__ G, {BF16}* __restrict__ O, " f"const {BF16}* __restrict__ X)" ) -LAND_PARAMS = f"({BF16}* __restrict__ O, const float* __restrict__ P)" LAND_RMS_NORM_PARAMS = ( f"(const {BF16}* __restrict__ G, {BF16}* __restrict__ O, " "const float* __restrict__ P)" ) -BINARY_ABO_PARAMS = ( - f"(const {BF16}* __restrict__ A, const {BF16}* __restrict__ Bt, " - f"{BF16}* __restrict__ O)" -) -SITU_PARAMS = ( - f"(const {BF16}* __restrict__ G, {BF16}* __restrict__ O, " - f"const {BF16}* __restrict__ U)" -) CONV_SILU_PARAMS = ( f"(const {BF16}* __restrict__ Cs, const float* __restrict__ Cw, " f"const float* __restrict__ P, {BF16}* __restrict__ Sn, " @@ -225,10 +190,6 @@ f"const {BF16}* __restrict__ Q, const float* __restrict__ State, " f"float* __restrict__ StateN, const {BF16}* __restrict__ V)" ) -O_NORM_GATE_PARAMS = ( - f"(const {BF16}* __restrict__ G2, const float* __restrict__ Go, " - f"{BF16}* __restrict__ Out, const {BF16}* __restrict__ X)" -) SCORES_PARAMS = ( f"(const {BF16}* __restrict__ Bl, const {BF16}* __restrict__ Ps, " "float* __restrict__ Sc, const float* __restrict__ Sw)" @@ -570,52 +531,6 @@ def plan_rms_norm_rbs() -> Plan: ) -def plan_land() -> Plan: - insts = [] - for nt, n, off in LAND_CONFIGS: - for split_k in SPLIT_K: - for batch in B_CHUNK_BUCKETS: - npad = ceildiv(n, THREADS) * THREADS - insts.append(Inst( - family="land", - order=len(insts), - label=f"land_batched NT={nt} N={n} OFF={off} SK={split_k} B={batch}", - factory="land_batched", - args=(nt, n, off, split_k, batch, THREADS), - num_params=2, - params=LAND_PARAMS, - symbol=f"k3_land_b{batch}_nt{nt}_n{n}_off{off}_sk{split_k}_kernel", - grid=(batch, npad // THREADS), - threads=THREADS, - guard=( - f"b == {batch} && nt == {nt} && n == {n} && " - f"off == {off} && split_k == {split_k}" - ), - call_args=(_bf16("O", False), "P"), - )) - return Plan( - stem=_STEM.format("land"), - signature=( - "k3_land_batched(\n" - " const float* P,\n" - " void* O,\n" - " int b,\n" - " int nt,\n" - " int n,\n" - " int off,\n" - " int split_k,\n" - " cudaStream_t stream)" - ), - doc=( - "// Merge the column span [off, off+n) of each row's (split_k, nt) f32\n" - "// partial and land bf16 once -- the landing of every matmul. split_k = 1\n" - "// is the single-partial case a framework GEMM produces, where the merge\n" - "// degenerates to the slice and the cast." - ), - insts=tuple(insts), - ) - - def plan_land_rms_norm_rbs() -> Plan: insts = [] for nt, n, off in LAND_RMS_NORM_CONFIGS: @@ -659,102 +574,13 @@ def plan_land_rms_norm_rbs() -> Plan: " cudaStream_t stream)" ), doc=( - "// k3_land_batched fused with the round-before-scale norm: MLA's q_a,\n" + "// The matmul landing fused with the round-before-scale norm: MLA's q_a,\n" "// the one place the engine fuses a merge and a norm." ), insts=tuple(insts), ) -def _plan_binary(family: str, widths: list[int], doc: str) -> Plan: - insts = [] - for width in widths: - for batch in B_CHUNK_BUCKETS: - insts.append(Inst( - family=family, - order=len(insts), - label=f"{family}_batched N={width} B={batch}", - factory=f"{family}_batched", - args=(width, batch, THREADS), - num_params=3, - params=BINARY_ABO_PARAMS, - symbol=f"k3_{family}_b{batch}_n{width}_kernel", - grid=(batch, width // THREADS), - threads=THREADS, - guard=f"b == {batch} && n == {width}", - call_args=(_bf16("A"), _bf16("Bt"), _bf16("O", False)), - )) - return Plan( - stem=_STEM.format(family), - signature=( - f"k3_{family}_batched(\n" - " const void* A,\n" - " const void* Bt,\n" - " void* O,\n" - " int b,\n" - " int n,\n" - " cudaStream_t stream)" - ), - doc=doc, - insts=tuple(insts), - ) - - -def plan_add2() -> Plan: - return _plan_binary( - "add2", - ADD2_N, - "// O = A + B in bf16 addition (the residual adds, and routed + shared).\n" - "// One block per (row, column segment).", - ) - - -def plan_mul_sigmoid() -> Plan: - return _plan_binary( - "mul_sigmoid", - MUL_SIGMOID_N, - "// O = A * bf16(sigmoid(B)), the MLA sigmoid output gate. The sigmoid is\n" - "// taken in f32 and lands in bf16 before the product.", - ) - - -def plan_situ() -> Plan: - insts = [] - for width in SITU_N: - for batch in B_CHUNK_BUCKETS: - insts.append(Inst( - family="situ", - order=len(insts), - label=f"situ_batched N={width} B={batch}", - factory="situ_batched", - args=(width, batch, THREADS), - num_params=3, - params=SITU_PARAMS, - symbol=f"k3_situ_b{batch}_n{width}_kernel", - grid=(batch, width // THREADS), - threads=THREADS, - guard=f"b == {batch} && n == {width}", - call_args=(_bf16("G"), _bf16("O", False), _bf16("U")), - )) - return Plan( - stem=_STEM.format("situ"), - signature=( - "k3_situ_batched(\n" - " const void* G,\n" - " const void* U,\n" - " void* O,\n" - " int b,\n" - " int n,\n" - " cudaStream_t stream)" - ), - doc=( - "// 4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25), computed in f32 and landed in\n" - "// bf16 once; the two betas are compiled in." - ), - insts=tuple(insts), - ) - - def plan_conv_silu() -> Plan: insts = [] for split_k in SPLIT_K: @@ -872,51 +698,6 @@ def plan_kda_core() -> Plan: ) -def plan_o_norm_gate() -> Plan: - insts = [] - for batch in B_CHUNK_BUCKETS: - insts.append(Inst( - family="o_norm_gate", - order=len(insts), - label=f"o_norm_gate_batched KH={KDA_HEADS} KD={KDA_HEAD_DIM} B={batch}", - factory="o_norm_gate_batched", - args=(KDA_HEADS, KDA_HEAD_DIM, batch, RMS_EPS), - num_params=4, - params=O_NORM_GATE_PARAMS, - symbol=f"k3_o_norm_gate_b{batch}_kh{KDA_HEADS}_kd{KDA_HEAD_DIM}_kernel", - grid=(batch, KDA_HEADS), - threads=KDA_HEAD_DIM, - guard=( - f"b == {batch} && num_heads == {KDA_HEADS} && " - f"head_dim == {KDA_HEAD_DIM}" - ), - call_args=( - _bf16("G2"), "Go", _bf16("Out", False), _bf16("X"), - ), - )) - return Plan( - stem=_STEM.format("o_norm_gate"), - signature=( - "k3_o_norm_gate_batched(\n" - " const void* X,\n" - " const void* G2,\n" - " const float* Go,\n" - " void* Out,\n" - " int b,\n" - " int num_heads,\n" - " int head_dim,\n" - " cudaStream_t stream)" - ), - doc=( - "// kda_core's tail on its own: per (row, head) the f32 rms_norm of the\n" - "// bf16 attention landing times the o_norm gamma, landed once, times the\n" - "// bf16 sigmoid of the output gate. Chunked prefill computes attention\n" - "// through FlashKDA and finishes rows here; eps is compiled in." - ), - insts=tuple(insts), - ) - - def plan_attnres_scores() -> Plan: insts = [] for blocks in ATTNRES_NB: @@ -1000,14 +781,9 @@ def plan_attnres_mix() -> Plan: PLANNERS = [ plan_rms_norm_rbs, - plan_land, plan_land_rms_norm_rbs, - plan_add2, - plan_mul_sigmoid, - plan_situ, plan_conv_silu, plan_kda_core, - plan_o_norm_gate, plan_attnres_scores, plan_attnres_mix, ] diff --git a/pegainfer-k3/kernels/tilelang_defs.py b/pegainfer-k3/kernels/tilelang_defs.py index 4bbaf7711..47293de31 100644 --- a/pegainfer-k3/kernels/tilelang_defs.py +++ b/pegainfer-k3/kernels/tilelang_defs.py @@ -49,7 +49,6 @@ def _compile(prim): return tilelang.compile(prim) - # --------------------------------------------------------------------------- # # Batched decode kernels, vendored verbatim (including the upstream section # comment below, which names the parity gates that certified them). @@ -168,29 +167,6 @@ def main(X: T.Tensor((B, H), DT), G: T.Tensor((H,), DT), return _compile(main) -@lru_cache(maxsize=None) -def land_batched(NT: int, N: int, OFF: int, SK: int, B: int, threads: int = 256): - """Batched ``land``: block (b, x) merges row b's partials for one column - span. The serial ascending-s merge and the single bf16 landing are the - bs=1 body verbatim.""" - NPAD = ((N + threads - 1) // threads) * threads - - @T.prim_func - def main(P: T.Tensor((B, SK, NT), ACC), O: T.Tensor((B, N), DT)): - with T.Kernel(B, NPAD // threads, threads=threads) as (bb, bx): - acc = T.alloc_fragment((threads,), ACC) - T.clear(acc) - for j in T.Parallel(threads): - for s in T.serial(SK): - acc[j] += P[bb, s, OFF + T.min(bx * threads + j, N - 1)] - for j in T.Parallel(threads): - with T.If(bx * threads + j < N): - with T.Then(): - O[bb, bx * threads + j] = T.Cast(DT, acc[j]) - - return _compile(main) - - @lru_cache(maxsize=None) def land_rms_norm_rbs_batched(NT: int, N: int, OFF: int, SK: int, B: int, eps: float, threads: int = 256): @@ -217,58 +193,6 @@ def main(P: T.Tensor((B, SK, NT), ACC), G: T.Tensor((N,), DT), return _compile(main) -@lru_cache(maxsize=None) -def add2_batched(N: int, B: int, threads: int = 256): - """Batched ``add2``: block (b, x) adds one column segment of row b in bf16. - Requires threads|N.""" - @T.prim_func - def main(A: T.Tensor((B, N), DT), Bt: T.Tensor((B, N), DT), - O: T.Tensor((B, N), DT)): - with T.Kernel(B, N // threads, threads=threads) as (bb, bx): - for j in T.Parallel(threads): - O[bb, bx * threads + j] = (A[bb, bx * threads + j] - + Bt[bb, bx * threads + j]) - - return _compile(main) - - -@lru_cache(maxsize=None) -def mul_sigmoid_batched(N: int, B: int, threads: int = 256): - """Batched ``mul_sigmoid``: block (b, x) gates one column segment of row b. - f32 sigmoid landing back in bf16, then the bf16 product -- the bs=1 body - verbatim. Requires threads|N.""" - @T.prim_func - def main(A: T.Tensor((B, N), DT), Bt: T.Tensor((B, N), DT), - O: T.Tensor((B, N), DT)): - with T.Kernel(B, N // threads, threads=threads) as (bb, bx): - for j in T.Parallel(threads): - n = bx * threads + j - O[bb, n] = A[bb, n] * T.Cast(DT, T.sigmoid(Bt[bb, n].astype(ACC))) - - return _compile(main) - - -@lru_cache(maxsize=None) -def situ_batched(N: int, B: int, threads: int = 256): - """Batched ``situ``: block (b, x) applies the activation to one column - segment of row b; beta=4 / linear_beta=25 stay compiled in and the f32 - chain lands bf16 once, as in the bs=1 body. Requires threads|N.""" - @T.prim_func - def main(G: T.Tensor((B, N), DT), U: T.Tensor((B, N), DT), - O: T.Tensor((B, N), DT)): - with T.Kernel(B, N // threads, threads=threads) as (bb, bx): - for j in T.Parallel(threads): - n = bx * threads + j - g = G[bb, n].astype(ACC) - u = U[bb, n].astype(ACC) - O[bb, n] = T.Cast( - DT, - 4.0 * T.tanh(g / 4.0) * T.sigmoid(g) * (25.0 * T.tanh(u / 25.0)), - ) - - return _compile(main) - - # --- KDA ------------------------------------------------------------------- # @@ -420,32 +344,3 @@ def main( return _compile(main) -def o_norm_gate_batched(KH: int, KD: int, B: int, eps: float): - """``kda_core``'s tail on its own: per (row, head) the f32 rms_norm of the - bf16 attention landing times the o_norm gamma, landed once, times the bf16 - sigmoid of the output-gate projection -- word-for-word the last loop of - ``kda_core_batched``. Chunked prefill computes the attention elsewhere - (FlashKDA) and finishes each row through this identical spelling.""" - KP = KH * KD - - @T.prim_func - def main( - X: T.Tensor((B, KP), DT), - G2: T.Tensor((B, KP), DT), - Go: T.Tensor((KD,), ACC), - Out: T.Tensor((B, KP), DT), - ): - with T.Kernel(B, KH, threads=KD) as (bb, bh): - asq = T.alloc_fragment((KD,), ACC) - atot = T.alloc_fragment((1,), ACC) - for d in T.Parallel(KD): - asq[d] = X[bb, bh * KD + d].astype(ACC) * X[bb, bh * KD + d].astype(ACC) - T.reduce_sum(asq, atot, dim=0) - for d in T.Parallel(KD): - Out[bb, bh * KD + d] = T.Cast( - DT, - X[bb, bh * KD + d].astype(ACC) * T.rsqrt(atot[0] / KD + eps) - * Go[d].astype(ACC), - ) * T.Cast(DT, T.sigmoid(G2[bb, bh * KD + d].astype(ACC))) - - return _compile(main) diff --git a/pegainfer-k3/src/executor/buffers.rs b/pegainfer-k3/src/executor/buffers.rs index dd6bfcf13..112a96f68 100644 --- a/pegainfer-k3/src/executor/buffers.rs +++ b/pegainfer-k3/src/executor/buffers.rs @@ -755,14 +755,6 @@ pub(crate) struct K3Scratch { pub(crate) forget_low: CudaSlice, pub(crate) out_gate: CudaSlice, pub(crate) conv_x: CudaSlice, - /// Prefill-chunk convolution windows, `[rows, K3_CONV_STATE, inner]`: row - /// `t` holds the three landed inputs preceding token `t`, prebuilt from - /// the chunk itself (and the carried window for the first rows), so one - /// batched conv launch serves the whole chunk. - pub(crate) conv_window: CudaSlice, - /// The batched conv launch's successor windows; row `tokens - 1` is the - /// window to carry into the next chunk. - pub(crate) conv_window_next: CudaSlice, pub(crate) conv_q: CudaSlice, pub(crate) conv_k: CudaSlice, pub(crate) conv_v: CudaSlice, @@ -889,8 +881,6 @@ impl K3Scratch { forget_low: wide(K3_HEAD_DIM)?, out_gate: wide(K3_ATTN_INNER)?, conv_x: wide(K3_ATTN_INNER)?, - conv_window: wide(K3_CONV_STATE * K3_ATTN_INNER)?, - conv_window_next: wide(K3_CONV_STATE * K3_ATTN_INNER)?, conv_q: wide(K3_ATTN_INNER)?, conv_k: wide(K3_ATTN_INNER)?, conv_v: wide(K3_ATTN_INNER)?, diff --git a/pegainfer-k3/src/executor/cp.rs b/pegainfer-k3/src/executor/cp.rs index 214b4622b..3203145b5 100644 --- a/pegainfer-k3/src/executor/cp.rs +++ b/pegainfer-k3/src/executor/cp.rs @@ -20,15 +20,19 @@ //! conv kernel then runs unchanged. //! * **MLA**: each rank publishes its segment's post-norm latent and rope //! rows; a rank assembles rows `0..seg_end` straight into the dense-FMHA -//! context scratch (the paged gather is bypassed) and the existing -//! bottom-right-aligned causal FMHA serves its queries at -//! `t_kv = seg_start + seg_len`. +//! context scratch (the paged gather is bypassed). The causal triangle is +//! then *striped* across the gang ([`k3_cp_stripe_kept`]): a rank attends +//! its own segment causally and its odd-distance upstream segments +//! densely, while each even-distance upstream rank attends this rank's +//! published queries over its own keys and returns output + LSE for the +//! owner's log-sum-exp merge — the deep rank's `r + 1` segment-pairs +//! become at most `(r + 3) / 2`. //! //! The exchange transport is in-process ranks and plain peer-access //! device-to-device copies, with every ordering edge expressed **on-device** //! through CUDA events — the host never syncs a stream inside a superstep. -//! Each of the roughly `2 × 69 + 24` windows per superstep runs the same -//! four-beat protocol on every rank: +//! Each of the roughly `2 × 69 + 24 × (1 + stripe slots)` windows per +//! superstep runs the same four-beat protocol on every rank: //! //! 1. record my *publish* event (all my publish writes are enqueued), then //! announce it through my `published` counter; @@ -49,7 +53,6 @@ //! passes the same collective window count, so the slots agree at each //! superstep boundary even as CP ranks rotate. -use std::ops::Range; use std::sync::Arc; use std::sync::Barrier; use std::sync::Mutex; @@ -69,17 +72,20 @@ use cudarc::driver::sys as cu_sys; use half::bf16; use pegainfer_kernels::ops::K3_KDA_HEAD_DIM; use pegainfer_kernels::ops::K3_KDA_HEADS; +use pegainfer_kernels::ops::K3_MLA_HEADS; use pegainfer_kernels::ops::gemm_strided_batched_f32; use pegainfer_kernels::tensor::DeviceContext; use pegainfer_kernels::tensor::active_cu_stream; use super::buffers::K3_CONV_STATE; use super::buffers::K3_KDA_STATE; +use super::buffers::K3_MLA_V_ROW; use super::buffers::copy_rows; use super::whale_gang::K3WhaleGang; use crate::config::K3_ATTN_INNER; use crate::config::K3_HIDDEN; use crate::config::K3_KV_LORA_RANK; +use crate::config::K3_Q_B_OUT; use crate::config::K3_QK_ROPE_HEAD_DIM; /// Raw device base pointers one CP rank publishes for its peers to read. @@ -92,6 +98,12 @@ pub(crate) struct K3CpPeerPtrs { pub(crate) kda_d: u64, pub(crate) mla_latent: u64, pub(crate) mla_rope: u64, + /// This segment's MLA queries for the stripe, `[seg_cap, 96 * 192]` bf16. + pub(crate) mla_q: u64, + /// The two stripe return slabs — attention output `[seg_cap, 96 * 128]` + /// bf16 and LSE `[96, seg_cap]` f32 — one pair per in-flight slot. + pub(crate) mla_ret_o: [u64; 2], + pub(crate) mla_ret_lse: [u64; 2], /// Raw `CUevent` handles for the owning rank's publish/consume events — /// peers wait on these cross-device (the safe wrapper refuses foreign /// contexts); the owning [`K3CpScratch`] keeps the events alive. @@ -101,32 +113,98 @@ pub(crate) struct K3CpPeerPtrs { /// Which ranks a window couples. `Halo` moves the conv carry one hop down /// the chain; `Upstream` fans every upstream rank's publication down to all -/// of its successors (KDA packages, MLA latents). +/// of its successors (KDA packages, MLA latents); `Stripe` couples the +/// explicit rank sets of one MLA stripe window (bitmasks over CP ranks — +/// the queries fan out to a rank's helpers, each return slot comes back from +/// them). #[derive(Clone, Copy)] pub(crate) enum K3CpWindowKind { Halo, Upstream, + Stripe { + /// Ranks whose publication I read this window. + from: u32, + /// Ranks that read mine. + by: u32, + }, +} + +/// The widest gang the stripe bitmasks seat. +pub(crate) const K3_CP_STRIPE_MAX_RANKS: usize = 32; + +fn rank_mask(ranks: impl IntoIterator) -> u32 { + ranks.into_iter().fold(0u32, |mask, rank| { + debug_assert!(rank < K3_CP_STRIPE_MAX_RANKS); + mask | (1u32 << rank) + }) +} + +fn mask_ranks(mask: u32) -> Vec { + (0..K3_CP_STRIPE_MAX_RANKS) + .filter(|&rank| mask & (1u32 << rank) != 0) + .collect() } impl K3CpWindowKind { /// Ranks whose publications `me` reads this window. CP ranks — the fleet /// gang maps them through its member table. - pub(crate) fn reads_from(self, me: usize) -> Range { + pub(crate) fn reads_from(self, me: usize) -> Vec { match self { - Self::Halo => me.saturating_sub(1)..me, - Self::Upstream => 0..me, + Self::Halo => (me.saturating_sub(1)..me).collect(), + Self::Upstream => (0..me).collect(), + Self::Stripe { from, .. } => mask_ranks(from), } } /// Ranks that read `me`'s publications this window. - pub(crate) fn read_by(self, me: usize, cp_size: usize) -> Range { + pub(crate) fn read_by(self, me: usize, cp_size: usize) -> Vec { match self { - Self::Halo => (me + 1).min(cp_size)..(me + 2).min(cp_size), - Self::Upstream => me + 1..cp_size, + Self::Halo => ((me + 1).min(cp_size)..(me + 2).min(cp_size)).collect(), + Self::Upstream => (me + 1..cp_size).collect(), + Self::Stripe { by, .. } => mask_ranks(by), } } } +/// The MLA stripe schedule: the causal triangle's off-diagonal segment pairs +/// `(q, k)`, `k < q`, each attended once — by the owner `q` (its queries over +/// its assembled context) when `q - k` is odd, by the key holder `k` (the +/// owner's published queries over its own expanded keys) when even. With +/// equal segments rank `r` carries `(r + 1) / 2` kept pairs plus +/// `(cp_size - 1 - r) / 2` helped ones on top of its own half-triangle — +/// loads of `cp_size / 2 ± 1/2` all the way down instead of `r + 1/2`, so +/// the deepest rank's attention halves (8 ranks: 7.5 → 4.5 segment-pairs, +/// ideal 4). The keys a helper needs are the ones it already expanded for +/// its own diagonal, so only queries and results move. +/// +/// Every rank derives the identical schedule from `cp_size` alone; the +/// windows it induces are collective and fixed per layer. +pub(crate) fn k3_cp_stripe_kept(cp_size: usize, owner: usize) -> Vec { + debug_assert!(owner < cp_size); + (0..owner).filter(|k| (owner - k) % 2 == 1).collect() +} + +/// Owners whose queries `helper` attends over its own keys, ascending. +pub(crate) fn k3_cp_stripe_helped(cp_size: usize, helper: usize) -> Vec { + debug_assert!(helper < cp_size); + (helper + 1..cp_size) + .filter(|q| (q - helper) % 2 == 0) + .collect() +} + +/// Return slots per MLA layer: the most pairs any one helper carries, which +/// is rank 0's count. Every rank runs exactly this many return windows. +pub(crate) fn k3_cp_stripe_slots(cp_size: usize) -> usize { + if cp_size == 0 { 0 } else { (cp_size - 1) / 2 } +} + +/// Helpers that hand `owner` a result in return slot `slot`. +pub(crate) fn k3_cp_stripe_returns(cp_size: usize, owner: usize, slot: usize) -> Vec { + (0..owner) + .filter(|&h| k3_cp_stripe_helped(cp_size, h).get(slot) == Some(&owner)) + .collect() +} + /// The coordination substrate one CP scratch runs its exchange windows over: /// the in-process gang (peer access + CUDA events) or the fleet whale gang /// (fabric slabs + doorbells). The forward path is agnostic — it snapshots a @@ -568,6 +646,13 @@ pub(crate) struct K3CpScratch { pub(crate) mla_latent_pub: CudaSlice, /// This segment's shared rope halves, `[seg_cap, 64]` bf16. pub(crate) mla_rope_pub: CudaSlice, + /// This segment's MLA queries for the stripe, `[seg_cap, 96 * 192]` bf16. + pub(crate) mla_q_pub: CudaSlice, + /// Stripe return slabs, one pair per in-flight slot: the helped owner's + /// attention output `[seg_cap, 96 * 128]` bf16 and LSE `[96, seg_cap]` + /// f32 over this rank's keys. + pub(crate) mla_ret_o: [CudaSlice; 2], + pub(crate) mla_ret_lse: [CudaSlice; 2], // Local working buffers. /// Received upstream normed tail, `[4, hidden]` (row 3 is bucket padding). pub(crate) halo_normed: CudaSlice, @@ -602,7 +687,7 @@ pub(crate) struct K3CpScratch { gang_ranks: Vec, } -/// The five buffers a CP rank publishes to its peers: pool allocations +/// The buffers a CP rank publishes to its peers: pool allocations /// in-process, fabric-slab carvings on the fleet. struct K3CpPublish { normed_tail: CudaSlice, @@ -610,6 +695,9 @@ struct K3CpPublish { kda_d: CudaSlice, mla_latent: CudaSlice, mla_rope: CudaSlice, + mla_q: CudaSlice, + mla_ret_o: [CudaSlice; 2], + mla_ret_lse: [CudaSlice; 2], } impl K3CpScratch { @@ -626,6 +714,19 @@ impl K3CpScratch { .alloc_zeros(seg_cap * K3_KV_LORA_RANK) .context("alloc K3 CP latent publish buffer")?, mla_rope: stream.alloc_zeros(seg_cap * K3_QK_ROPE_HEAD_DIM)?, + mla_q: stream + .alloc_zeros(seg_cap * K3_Q_B_OUT) + .context("alloc K3 CP stripe query publish buffer")?, + mla_ret_o: [ + stream + .alloc_zeros(seg_cap * K3_MLA_V_ROW) + .context("alloc K3 CP stripe return slab")?, + stream.alloc_zeros(seg_cap * K3_MLA_V_ROW)?, + ], + mla_ret_lse: [ + stream.alloc_zeros(K3_MLA_HEADS * seg_cap)?, + stream.alloc_zeros(K3_MLA_HEADS * seg_cap)?, + ], }; let events = Some(( new_event(ctx).context("create K3 CP publish event")?, @@ -669,6 +770,15 @@ impl K3CpScratch { kda_d: carve_f32(mine.kda_d, K3_KDA_STATE), mla_latent: carve_bf16(mine.mla_latent, seg_cap * K3_KV_LORA_RANK), mla_rope: carve_bf16(mine.mla_rope, seg_cap * K3_QK_ROPE_HEAD_DIM), + mla_q: carve_bf16(mine.mla_q, seg_cap * K3_Q_B_OUT), + mla_ret_o: [ + carve_bf16(mine.mla_ret_o[0], seg_cap * K3_MLA_V_ROW), + carve_bf16(mine.mla_ret_o[1], seg_cap * K3_MLA_V_ROW), + ], + mla_ret_lse: [ + carve_f32(mine.mla_ret_lse[0], K3_MLA_HEADS * seg_cap), + carve_f32(mine.mla_ret_lse[1], K3_MLA_HEADS * seg_cap), + ], }; Self::new_inner( ctx, @@ -706,6 +816,9 @@ impl K3CpScratch { kda_d: publish.kda_d, mla_latent_pub: publish.mla_latent, mla_rope_pub: publish.mla_rope, + mla_q_pub: publish.mla_q, + mla_ret_o: publish.mla_ret_o, + mla_ret_lse: publish.mla_ret_lse, halo_normed: stream.alloc_zeros((K3_CONV_STATE + 1) * K3_HIDDEN)?, halo_partial: stream.alloc_zeros((K3_CONV_STATE + 1) * K3_ATTN_INNER)?, halo_xs: stream.alloc_zeros((K3_CONV_STATE + 1) * K3_ATTN_INNER)?, @@ -769,6 +882,11 @@ impl K3CpScratch { "K3 CP rank {cp_rank} of {}", self.cp_size ); + ensure!( + self.cp_size <= K3_CP_STRIPE_MAX_RANKS, + "K3 CP gang of {} ranks exceeds the {K3_CP_STRIPE_MAX_RANKS}-rank stripe masks", + self.cp_size + ); self.cp_rank = cp_rank; ensure!( segments.len() == self.cp_size, @@ -797,6 +915,9 @@ impl K3CpScratch { kda_d: ptr(&self.kda_d), mla_latent: ptr_bf(&self.mla_latent_pub), mla_rope: ptr_bf(&self.mla_rope_pub), + mla_q: ptr_bf(&self.mla_q_pub), + mla_ret_o: [ptr_bf(&self.mla_ret_o[0]), ptr_bf(&self.mla_ret_o[1])], + mla_ret_lse: [ptr(&self.mla_ret_lse[0]), ptr(&self.mla_ret_lse[1])], publish_event: self .publish_event .as_ref() @@ -834,6 +955,10 @@ impl K3CpScratch { gang.world() ); ensure!(cp_rank < width, "K3 whale CP rank {cp_rank} of {width}"); + ensure!( + width <= K3_CP_STRIPE_MAX_RANKS, + "K3 whale gang of {width} ranks exceeds the {K3_CP_STRIPE_MAX_RANKS}-rank stripe masks" + ); ensure!( gang_ranks[cp_rank] == gang.rank(), "K3 whale gang seats rank {} at CP position {cp_rank}, but this executor is rank {}", @@ -870,6 +995,45 @@ impl K3CpScratch { Ok(()) } + /// Upstream segments this rank attends itself (ascending). + pub(crate) fn stripe_kept(&self) -> Vec { + k3_cp_stripe_kept(self.cp_size, self.cp_rank) + } + + /// Owners whose queries this rank attends over its own keys, by slot. + pub(crate) fn stripe_helped(&self) -> Vec { + k3_cp_stripe_helped(self.cp_size, self.cp_rank) + } + + /// Return windows every rank runs per MLA layer. + pub(crate) fn stripe_slots(&self) -> usize { + k3_cp_stripe_slots(self.cp_size) + } + + /// Helpers handing this rank a result in `slot`. + pub(crate) fn stripe_returns(&self, slot: usize) -> Vec { + k3_cp_stripe_returns(self.cp_size, self.cp_rank, slot) + } + + /// The query window: my queries fan out to my helpers, and I read the + /// queries of every owner I help. + pub(crate) fn stripe_query_kind(&self) -> K3CpWindowKind { + let helpers = (0..self.cp_rank).filter(|h| (self.cp_rank - h) % 2 == 0); + K3CpWindowKind::Stripe { + from: rank_mask(self.stripe_helped()), + by: rank_mask(helpers), + } + } + + /// Return window `slot`: I read the results my slot-`slot` helpers + /// publish, and the owner I serve in that slot reads mine. + pub(crate) fn stripe_return_kind(&self, slot: usize) -> K3CpWindowKind { + K3CpWindowKind::Stripe { + from: rank_mask(self.stripe_returns(slot)), + by: rank_mask(self.stripe_helped().get(slot).copied()), + } + } + /// Snapshot what one exchange window needs — plain copied data, so the /// `consume` closure is free to capture the scratch. On the fleet /// substrate this also claims the window's doorbell value, so every @@ -1071,6 +1235,63 @@ mod tests { ); } + /// Every off-diagonal pair is attended exactly once, by its owner or its + /// key holder, and the window sets are mirror images. + #[test] + fn stripes_cover_the_triangle_once_and_mirror() { + for cp_size in 2..=16usize { + for q in 0..cp_size { + for k in 0..q { + let kept = k3_cp_stripe_kept(cp_size, q).contains(&k); + let helped = k3_cp_stripe_helped(cp_size, k).contains(&q); + assert!( + kept ^ helped, + "pair ({q}, {k}) @ {cp_size}: kept {kept}, helped {helped}" + ); + } + } + let slots = k3_cp_stripe_slots(cp_size); + for h in 0..cp_size { + assert!(k3_cp_stripe_helped(cp_size, h).len() <= slots); + } + for slot in 0..slots { + for q in 0..cp_size { + for h in k3_cp_stripe_returns(cp_size, q, slot) { + assert_eq!(k3_cp_stripe_helped(cp_size, h).get(slot), Some(&q)); + } + } + } + } + } + + /// With equal segments the deepest rank's attention halves and no rank + /// exceeds the ideal by more than half a segment-pair. + #[test] + fn stripes_level_the_causal_triangle() { + for cp_size in [4usize, 8, 16] { + let load = |r: usize| -> f64 { + 0.5 + k3_cp_stripe_kept(cp_size, r).len() as f64 + + k3_cp_stripe_helped(cp_size, r).len() as f64 + }; + let ideal = cp_size as f64 / 2.0; + let worst = (0..cp_size).map(load).fold(0.0, f64::max); + assert!( + worst <= ideal + 0.5, + "{cp_size}: worst {worst} vs ideal {ideal}" + ); + // The deepest rank: r + 1/2 unstriped, at most (r + 3) / 2 striped. + let deep = load(cp_size - 1); + assert!( + deep <= (cp_size as f64 + 2.0) / 2.0, + "{cp_size}: deep {deep}" + ); + assert!( + deep < cp_size as f64 - 0.5, + "{cp_size}: deep {deep} not relieved" + ); + } + } + #[test] fn gang_is_tray_aligned_with_the_poster_last() { assert_eq!(k3_whale_gang(5, 8, 16), vec![4, 6, 7, 8, 9, 10, 11, 5]); diff --git a/pegainfer-k3/src/executor/forward/prefill.rs b/pegainfer-k3/src/executor/forward/prefill.rs index 7138669c9..207a65701 100644 --- a/pegainfer-k3/src/executor/forward/prefill.rs +++ b/pegainfer-k3/src/executor/forward/prefill.rs @@ -6,18 +6,19 @@ use anyhow::Result; use anyhow::ensure; use cudarc::driver::CudaSlice; use half::bf16; -use pegainfer_kernels::ops::K3_CONV_WIDTH; use pegainfer_kernels::ops::K3_MLA_HEADS; use pegainfer_kernels::ops::argmax_bf16_split_into; use pegainfer_kernels::ops::gemm_rows_span_into_checked; -use pegainfer_kernels::ops::k3_conv_silu_batched_launch; +use pegainfer_kernels::ops::k3_conv_silu_chunk_launch; use pegainfer_kernels::ops::k3_flash_kda_fwd_launch; use pegainfer_kernels::ops::k3_flash_mla_prefill_fwd_dense_launch; +use pegainfer_kernels::ops::k3_flash_mla_prefill_fwd_dense_peer_q_launch; use pegainfer_kernels::ops::k3_flash_mla_prefill_fwd_launch; use pegainfer_kernels::ops::k3_land_batched_launch; use pegainfer_kernels::ops::k3_mla_prefill_expand_k_launch; use pegainfer_kernels::ops::k3_mla_prefill_gather_launch; use pegainfer_kernels::ops::k3_mla_prefill_lse_merge_launch; +use pegainfer_kernels::ops::k3_mla_prefill_lse_merge_peer_launch; use pegainfer_kernels::ops::k3_mla_prefill_o_finalize_launch; use pegainfer_kernels::ops::k3_o_norm_gate_batched_launch; use pegainfer_kernels::ops::k3_rms_norm_rbs_batched_launch; @@ -55,6 +56,7 @@ use crate::config::K3_HEADS; use crate::config::K3_HIDDEN; use crate::config::K3_KV_B_OUT; use crate::config::K3_KV_LORA_RANK; +use crate::config::K3_RMS_EPS; use crate::config::K3_VOCAB; use crate::model::K3KdaWeights; use crate::model::K3LayerWeights; @@ -394,9 +396,6 @@ pub(super) fn kda_attention_chunk( 0, &s.normed, &mut s.kda_conv_partial, - &mut s.conv_x, - &mut s.conv_window, - &mut s.conv_window_next, &mut s.conv_q, )?; kda_conv_stream_chunk( @@ -410,9 +409,6 @@ pub(super) fn kda_attention_chunk( 1, &s.normed, &mut s.kda_conv_partial, - &mut s.conv_x, - &mut s.conv_window, - &mut s.conv_window_next, &mut s.conv_k, )?; kda_conv_stream_chunk( @@ -426,9 +422,6 @@ pub(super) fn kda_attention_chunk( 2, &s.normed, &mut s.kda_conv_partial, - &mut s.conv_x, - &mut s.conv_window, - &mut s.conv_window_next, &mut s.conv_v, )?; @@ -632,6 +625,7 @@ pub(super) fn kda_attention_chunk( b, K3_HEADS, K3_HEAD_DIM, + K3_RMS_EPS, &s.kda_attn, &s.out_gate, &w.gamma_o, @@ -654,8 +648,8 @@ pub(super) fn kda_attention_chunk( } /// One q/k/v stream of a prefill chunk or verify step: the batched band -/// projection, the per-group window builds, one batched convolution, and each -/// group's carry into its next segment. +/// projection, then one conv launch per group that walks the segment's +/// partial rows in place and carries the window into its next segment. #[allow(clippy::too_many_arguments)] fn kda_conv_stream_chunk( ctx: &DeviceContext, @@ -668,13 +662,9 @@ fn kda_conv_stream_chunk( stream_index: usize, normed: &CudaSlice, partial: &mut CudaSlice, - xs: &mut CudaSlice, - window: &mut CudaSlice, - window_next: &mut CudaSlice, out: &mut CudaSlice, ) -> Result<()> { let inner = K3_ATTN_INNER; - let window_row = K3_CONV_STATE * inner; k3_gemm_partial( ctx, fused, @@ -685,76 +675,43 @@ fn kda_conv_stream_chunk( partial, K3PartialSpan::whole(inner), )?; - // Land the step's inputs once: the window entries ARE these bf16 rows, - // the same cast the conv kernel itself applies. - k3_land_batched_launch(ctx, b, inner, inner, 0, 1, partial, xs)?; - // Row t of a group's window slot j holds the group's input `t - - // K3_CONV_STATE + j`: from the segment itself once that token exists, - // from the slot's carried window before it. Rows outside every group - // keep stale windows; their conv output is padding and is discarded. - for group in groups { - let tokens = group.commit_rows + group.spec_rows; - let carry = &conv_state[group.parity][stream_index]; - for j in 0..K3_CONV_STATE { - let lead = K3_CONV_STATE - j; - if tokens > lead { - copy_rows_2d( - ctx, - xs, - group.row * inner, - inner, - window, - ((group.row + lead) * K3_CONV_STATE + j) * inner, - window_row, - tokens - lead, - inner, - )?; - } - for t in 0..lead.min(tokens) { - copy_rows_2d( - ctx, - carry, - (group.state_row * K3_CONV_STATE + t + j) * inner, - inner, - window, - ((group.row + t) * K3_CONV_STATE + j) * inner, - inner, - 1, - inner, - )?; - } - } - } - k3_conv_silu_batched_launch( - ctx, - b, - inner, - K3_CONV_WIDTH, - 1, - partial, - taps, - window, - xs, - out, - window_next, - )?; + // Row t of a group's window slot j is the group's input `t - + // K3_CONV_STATE + j`: the segment's own partial row once that token + // exists (the kernel lands it, the same cast the conv applies), the + // slot's carried window before it. Nothing is materialized: the kernel + // reads the neighbouring rows in place. Rows outside every group are + // not touched; their conv output is padding and is discarded. + // // A group's carry into its next segment is the successor window of its // last COMMIT row — for a segment shorter than the window it already - // folds the slots carried in above. It lands in the other parity slab, + // folds the slots carried in. It lands in the other parity slab, // agreeing with the recurrent state's per-segment double buffering. The // speculative tail's successors are never carried: its tokens replay as // the next round's commit rows. - for group in groups.iter().filter(|group| group.commit_rows > 0) { - copy_rows_2d( + let (even, odd) = conv_state.split_at_mut(1); + for group in groups { + let tokens = group.commit_rows + group.spec_rows; + if tokens == 0 { + continue; + } + let (carry, next) = if group.parity == 0 { + (&even[0][stream_index], &mut odd[0][stream_index]) + } else { + (&odd[0][stream_index], &mut even[0][stream_index]) + }; + k3_conv_silu_chunk_launch( ctx, - window_next, - (group.row + group.commit_rows - 1) * window_row, - window_row, - &mut conv_state[group.parity ^ 1][stream_index], - group.state_row * window_row, - window_row, - 1, - window_row, + inner, + tokens, + group.commit_rows, + partial, + group.row, + taps, + carry, + group.state_row, + out, + group.row, + (group.commit_rows > 0).then_some((next, group.state_row)), )?; } Ok(()) @@ -797,12 +754,20 @@ pub(super) fn mla_attention_chunk_fmha( } /// The CP variant of one chunk's MLA attention: the paged gather is replaced -/// by assembly from the gang's published post-norm latents. Each rank -/// publishes its own segment's `kv_norm`/`rope` rows, copies its upstream -/// peers' rows into the context scratch at their global offsets, lands its -/// own rows from local scratch, and runs the same bottom-right-aligned FMHA -/// at `t_kv = seg_start + seg_len` — causality falls out of the alignment -/// exactly as for a local chunk. +/// by assembly from the gang's published post-norm latents, and the causal +/// triangle is striped across the gang (`cp::k3_cp_stripe_kept`). +/// +/// Each rank publishes its own segment's `kv_norm`/`rope` rows and copies +/// its upstream peers' rows into the context scratch at their global +/// offsets. Then, per layer: it publishes its queries; attends its *kept* +/// upstream segments densely and its own segment causally, folding every +/// window into the f32 log-sum-exp accumulator; and runs the return slots — +/// in slot `i` it attends the `i`-th owner it helps (that owner's published +/// queries over this rank's own expanded keys, straight into a return slab) +/// and merges what its own slot-`i` helpers returned, reading their slabs in +/// place. The finalize lands `s.attn`. Merge order is fixed (kept segments +/// ascending, own diagonal, helpers by slot), so the result is deterministic +/// and identical on the in-process and fleet substrates. pub(super) fn mla_attention_chunk_cp( ctx: &DeviceContext, shape: K3StepShape, @@ -820,6 +785,17 @@ pub(super) fn mla_attention_chunk_cp( "K3 CP MLA step shape ({}+{t_q}) disagrees with the rank's segment ({seg_start}+{seg_len})", shape.chunk_start ); + ensure!( + t_kv * K3_KV_LORA_RANK <= s.mla_ctx_latent.data.len(), + "K3 MLA prefill workspace of {} tokens cannot span the {t_kv}-token context", + s.mla_ctx_latent.data.len() / K3_KV_LORA_RANK + ); + let win = s.mla_ctx_k.len() / crate::config::K3_Q_B_OUT; + ensure!( + cp.segments.iter().all(|&(_, len)| len <= win), + "K3 CP segments {:?} exceed the {win}-row expansion window", + cp.segments + ); copy_rows( ctx, &s.kv_norm, @@ -838,6 +814,17 @@ pub(super) fn mla_attention_chunk_cp( t_q, crate::config::K3_QK_ROPE_HEAD_DIM, )?; + // The stripe's query publication rides the same window as the latents: + // both are ready, and one window fewer per layer. + copy_rows( + ctx, + &s.query, + 0, + &mut cp.mla_q_pub, + 0, + t_q, + crate::config::K3_Q_B_OUT, + )?; let group = cp.sync.clone(); let sync = cp.window_sync(K3CpWindowKind::Upstream)?; group.exchange(ctx, &sync, || { @@ -862,6 +849,11 @@ pub(super) fn mla_attention_chunk_cp( } Ok(()) })?; + // The query window proper: nothing is read inside it — every helper's + // read of my queries is an FMHA launched before a return-window publish + // I wait on, which orders my next layer's overwrite behind it. + let sync = cp.window_sync(cp.stripe_query_kind())?; + group.exchange(ctx, &sync, || Ok(()))?; copy_rows( ctx, &s.kv_norm, @@ -894,7 +886,115 @@ pub(super) fn mla_attention_chunk_cp( &s.mla_ctx_rope, )?; } - mla_chunk_attend(ctx, t_q, t_kv, w, s) + + // The decode kernel reads the softmax scale as a bf16 device scalar; feed + // the FMHA the same rounded constant so the two paths agree on it. + let scale = bf16::from_f64(crate::model::k3_mla_scale()).to_f32(); + s.mla_ctx_latent.seq_len = t_kv; + // Kept upstream segments, densely, in key order. + let mut first = true; + for k in cp.stripe_kept() { + let (k_start, k_len) = cp.segments[k]; + mla_expand_window(ctx, w, s, k_start, k_len)?; + k3_flash_mla_prefill_fwd_dense_launch( + ctx, + t_q, + k_len, + K3_MLA_HEADS, + &s.query, + &s.mla_ctx_k, + &s.mla_ctx_nope_v.data, + &mut s.attn, + Some(&mut s.mla_lse_win), + scale, + )?; + k3_mla_prefill_lse_merge_launch( + ctx, + t_q, + K3_MLA_HEADS, + &s.attn, + &s.mla_lse_win, + &mut s.mla_o_acc, + &mut s.mla_lse_acc, + first, + )?; + first = false; + } + // The own diagonal — and the expansion the helper slots below read. + mla_expand_window(ctx, w, s, seg_start, t_q)?; + k3_flash_mla_prefill_fwd_launch( + ctx, + t_q, + t_q, + K3_MLA_HEADS, + &s.query, + &s.mla_ctx_k, + &s.mla_ctx_nope_v.data, + &mut s.attn, + Some(&mut s.mla_lse_win), + scale, + )?; + k3_mla_prefill_lse_merge_launch( + ctx, + t_q, + K3_MLA_HEADS, + &s.attn, + &s.mla_lse_win, + &mut s.mla_o_acc, + &mut s.mla_lse_acc, + first, + )?; + + // Return slots. Slot i's helped FMHA is launched before window i opens + // (slot 0 here, slot i+1 inside window i's consume, so it needs no wait + // on window i's readers); it lands in slab i % 2, whose previous readers + // window i-2 already waited for. The owner merges each returned window + // in place inside window i. + let slots = cp.stripe_slots(); + let helped = cp.stripe_helped(); + let peers = cp.peers.clone(); + let segments = cp.segments.clone(); + let helped_fmha = |cp: &mut K3CpScratch, s: &mut K3Scratch, slot: usize| -> Result<()> { + let Some(&owner) = helped.get(slot) else { + return Ok(()); + }; + let (_, owner_len) = segments[owner]; + let slab = slot % 2; + let (ret_o, ret_lse) = (&mut cp.mla_ret_o[slab], &mut cp.mla_ret_lse[slab]); + k3_flash_mla_prefill_fwd_dense_peer_q_launch( + ctx, + owner_len, + t_q, + K3_MLA_HEADS, + peers[owner].mla_q, + &s.mla_ctx_k, + &s.mla_ctx_nope_v.data, + ret_o, + Some(ret_lse), + scale, + ) + }; + helped_fmha(cp, s, 0)?; + for slot in 0..slots { + let returns = cp.stripe_returns(slot); + let sync = cp.window_sync(cp.stripe_return_kind(slot))?; + group.exchange(ctx, &sync, || { + for helper in returns { + k3_mla_prefill_lse_merge_peer_launch( + ctx, + t_q, + K3_MLA_HEADS, + peers[helper].mla_ret_o[slot % 2], + peers[helper].mla_ret_lse[slot % 2], + &mut s.mla_o_acc, + &mut s.mla_lse_acc, + false, + )?; + } + helped_fmha(cp, s, slot + 1) + })?; + } + k3_mla_prefill_o_finalize_launch(ctx, t_q, K3_MLA_HEADS, &s.mla_o_acc, &mut s.attn) } /// The shared tail of a chunk's MLA attention: kv_b expansion of the diff --git a/pegainfer-k3/src/executor/whale_gang.rs b/pegainfer-k3/src/executor/whale_gang.rs index 98ec9c995..c30fb1531 100644 --- a/pegainfer-k3/src/executor/whale_gang.rs +++ b/pegainfer-k3/src/executor/whale_gang.rs @@ -45,6 +45,7 @@ use anyhow::ensure; use cudarc::driver::sys as cu_sys; use half::bf16; use pegainfer_kernels::ops::K3_MEGA_FABRIC_HANDLE_BYTES; +use pegainfer_kernels::ops::K3_MLA_HEADS; use pegainfer_kernels::ops::k3_chunk_bucket; use pegainfer_kernels::ops::k3_mega_fabric_slab_alloc; use pegainfer_kernels::ops::k3_mega_fabric_slab_import; @@ -55,12 +56,14 @@ use pegainfer_kernels::tensor::active_cu_stream; use super::buffers::K3_CONV_STATE; use super::buffers::K3_KDA_STATE; +use super::buffers::K3_MLA_V_ROW; use super::cp::K3CpPeerPtrs; use super::cp::K3CpScratch; use super::cp::K3CpSyncHandle; use super::cp::K3CpWindowKind; use crate::config::K3_HIDDEN; use crate::config::K3_KV_LORA_RANK; +use crate::config::K3_Q_B_OUT; use crate::config::K3_QK_ROPE_HEAD_DIM; /// Doorbell values per whale: window `w` of whale `seq` rings @@ -85,6 +88,9 @@ pub(crate) struct K3WhaleSlabLayout { kda_d: usize, mla_latent: usize, mla_rope: usize, + mla_q: usize, + mla_ret_o: [usize; 2], + mla_ret_lse: [usize; 2], pub(crate) num_bytes: usize, } @@ -111,6 +117,15 @@ impl K3WhaleSlabLayout { let kda_d = region(K3_KDA_STATE * size_of::()); let mla_latent = region(seg_cap * K3_KV_LORA_RANK * size_of::()); let mla_rope = region(seg_cap * K3_QK_ROPE_HEAD_DIM * size_of::()); + let mla_q = region(seg_cap * K3_Q_B_OUT * size_of::()); + let mla_ret_o = [ + region(seg_cap * K3_MLA_V_ROW * size_of::()), + region(seg_cap * K3_MLA_V_ROW * size_of::()), + ]; + let mla_ret_lse = [ + region(K3_MLA_HEADS * seg_cap * size_of::()), + region(K3_MLA_HEADS * seg_cap * size_of::()), + ]; Self { publish_inbox, consume_inbox, @@ -119,6 +134,9 @@ impl K3WhaleSlabLayout { kda_d, mla_latent, mla_rope, + mla_q, + mla_ret_o, + mla_ret_lse, num_bytes: aligned(offset), } } @@ -266,6 +284,15 @@ impl K3WhaleGang { kda_d: base + self.layout.kda_d as u64, mla_latent: base + self.layout.mla_latent as u64, mla_rope: base + self.layout.mla_rope as u64, + mla_q: base + self.layout.mla_q as u64, + mla_ret_o: [ + base + self.layout.mla_ret_o[0] as u64, + base + self.layout.mla_ret_o[1] as u64, + ], + mla_ret_lse: [ + base + self.layout.mla_ret_lse[0] as u64, + base + self.layout.mla_ret_lse[1] as u64, + ], publish_event: 0, consume_event: 0, }) @@ -304,6 +331,7 @@ impl K3WhaleGang { // launches. let publish_flags: Vec = kind .read_by(cp_rank, cp_size) + .into_iter() .map(|reader| { self.layout .publish_flag(self.bases[gang[reader]], self.rank) @@ -322,6 +350,7 @@ impl K3WhaleGang { consume()?; let consume_flags: Vec = kind .reads_from(cp_rank) + .into_iter() .map(|source| { self.layout .consume_flag(self.bases[gang[source]], self.rank) @@ -475,6 +504,23 @@ mod tests { layout.mla_rope, seg_cap * K3_QK_ROPE_HEAD_DIM * size_of::(), ), + (layout.mla_q, seg_cap * K3_Q_B_OUT * size_of::()), + ( + layout.mla_ret_o[0], + seg_cap * K3_MLA_V_ROW * size_of::(), + ), + ( + layout.mla_ret_o[1], + seg_cap * K3_MLA_V_ROW * size_of::(), + ), + ( + layout.mla_ret_lse[0], + K3_MLA_HEADS * seg_cap * size_of::(), + ), + ( + layout.mla_ret_lse[1], + K3_MLA_HEADS * seg_cap * size_of::(), + ), ]; for (offset, _) in regions { assert_eq!(offset % K3_WHALE_ALIGN, 0, "unaligned region at {offset}"); diff --git a/pegainfer-kernels/KERNELS.md b/pegainfer-kernels/KERNELS.md index 57b207b2d..133fcb6ff 100644 --- a/pegainfer-kernels/KERNELS.md +++ b/pegainfer-kernels/KERNELS.md @@ -131,14 +131,15 @@ f32, and narrowing any of them to bf16 measurably flips routing decisions. | op_id | Runtime owner | Rust wrapper | FFI symbol | Source (OUT_DIR `.cu`) | Backend | Shape / layout notes | | --- | --- | --- | --- | --- | --- | --- | | `k3.norm.rms_rbs` | `pegainfer-k3` | `ops::k3_rms_norm_rbs_batched_launch` | `k3_rms_norm_rbs_batched` | `k3_rms_norm_rbs_batched.cu` | TileLang AOT CUDA | KimiRMSNorm, round-before-scale: the normalized value lands bf16 *before* multiplying gamma. One row per block; `X`/`O` are `[b, h]` bf16 and gamma `[h]` is shared. h ∈ {7168, 512, 3584} (layer norms, MLA kv latent, routed latent) × 15 buckets; eps compiled in. | -| `k3.linear.land` | `pegainfer-k3` | `ops::k3_land_batched_launch` | `k3_land_batched` | `k3_land_batched.cu` | TileLang AOT CUDA | Merge the column span `[off, off+n)` of each row's `[split_k, nt]` f32 partial and land bf16 once — the landing of every matmul. 15 (nt, n, off) spans × 15 buckets (14 from the sequential engine verbatim, plus `(12288, 12288, 0)` for chunked prefill's conv-input landing — the sequential engine casts conv inputs inside the conv kernel, the chunk lands them before its window build), split_k = 1 only — the single partial a framework GEMM (cuBLASLt, DeepGEMM) produces, where the merge degenerates to the slice and the cast. Masked tail store, so `n` need not divide 256. | +| `k3.linear.land` | `pegainfer-k3` | `ops::k3_land_batched_launch` | `k3_land_cuda` | `csrc/k3/k3_land.cu` | CUDA | Merge the column span `[off, off+n)` of each row's `[split_k, nt]` f32 partial and land bf16 once — the landing of every matmul. Replaces the TileLang `land_batched` family (one element per thread, 2.3 TB/s on the chunked-prefill landings): eight columns per thread with 16-byte loads and stores, the same ascending-segment f32 sum and single round-to-nearest-even cast, so bit-identical. Shapes are runtime values (no per-bucket instantiation); rows not 16-byte aligned fall back to a scalar path. split_k = 1 at every launch site — the single partial a framework GEMM (cuBLASLt, DeepGEMM) produces. | | `k3.linear.land_rms_rbs` | `pegainfer-k3` | `ops::k3_land_rms_norm_rbs_batched_launch` | `k3_land_rms_norm_rbs_batched` | `k3_land_rms_norm_rbs_batched.cu` | TileLang AOT CUDA | `k3_land_batched` fused with the round-before-scale norm. One span — MLA q_a, `[0, 1536)` of the `14400`-wide fused projection — × 15 buckets, split_k = 1. | -| `k3.elementwise.add2` | `pegainfer-k3` | `ops::k3_add2_batched_launch` | `k3_add2_batched` | `k3_add2_batched.cu` | TileLang AOT CUDA | `O = A + Bt` in bf16 addition (residual adds, routed + shared). n = 7168 × 15 buckets. | -| `k3.elementwise.mul_sigmoid` | `pegainfer-k3` | `ops::k3_mul_sigmoid_batched_launch` | `k3_mul_sigmoid_batched` | `k3_mul_sigmoid_batched.cu` | TileLang AOT CUDA | `O = A * bf16(sigmoid(Bt))`, the MLA sigmoid output gate; the sigmoid is taken in f32 and lands bf16 before the product. n = 12288 × 15 buckets. | -| `k3.act.situ` | `pegainfer-k3` | `ops::k3_situ_batched_launch` | `k3_situ_batched` | `k3_situ_batched.cu` | TileLang AOT CUDA | `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)` in f32, landing bf16 once; the betas are compiled in. n ∈ {6144 (shared), 33792 (dense)} × 15 buckets; the routed-expert situ is fused into the masked chain and the mega kernel, so no wide routed instantiation exists here. | +| `k3.elementwise.add2` | `pegainfer-k3` | `ops::k3_add2_batched_launch` | `k3_add2_cuda` | `csrc/k3/k3_elementwise.cu` | CUDA | `O = A + Bt` in bf16 addition (residual adds, routed + shared). Eight columns per thread (`add.rn.bf16x2`), bit-identical to the retired TileLang kernel; shapes are runtime values, n a multiple of 8. | +| `k3.elementwise.mul_sigmoid` | `pegainfer-k3` | `ops::k3_mul_sigmoid_batched_launch` | `k3_mul_sigmoid_cuda` | `csrc/k3/k3_elementwise.cu` | CUDA | `O = A * bf16(sigmoid(Bt))`, the MLA sigmoid output gate; the sigmoid is taken in f32 and lands bf16 before the bf16 product. Eight columns per thread, bit-identical to the retired TileLang kernel; shapes are runtime values. | +| `k3.act.situ` | `pegainfer-k3` | `ops::k3_situ_batched_launch` | `k3_situ_cuda` | `csrc/k3/k3_elementwise.cu` | CUDA | `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)` in f32, landing bf16 once; the betas are compiled in. Eight columns per thread, bit-identical to the retired TileLang kernel; shapes are runtime values (shared 6144, dense 33792). The routed-expert situ is fused into the masked chain and the mega kernel. | | `k3.kda.conv_silu` | `pegainfer-k3` | `ops::k3_conv_silu_batched_launch` | `k3_conv_silu_batched` | `k3_conv_silu_batched.cu` | TileLang AOT CUDA | Causal depthwise convolution over the 4-slot window plus silu, one token per row. Consumes the projection's `[b, split_k, 12288]` f32 partial: its bf16 landing is `X`, the newest window slot; `Sn` is the shifted state the caller carries. Conv weights `[4, 12288]` are **f32** and have no batch axis; the window state is `[b, 3, 12288]`, one independent window per row. split_k = 1 × 15 buckets. | | `k3.kda.core` | `pegainfer-k3` | `ops::k3_kda_core_batched_launch` | `k3_kda_core_batched` | `k3_kda_core_batched.cu` | TileLang AOT CUDA | One delta-rule step per row, one (row, head) per block, `threads = head_dim`. State `[b, 96, 128, 128]` f32 laid out `[head, v_dim, k_dim]` per row with decay along k, read from `State` and written to `StateN` (must not alias). `Dt`/`Alog`/`Go` f32 weights with no batch axis, `Bt`/`G2` bf16; gate lower bound and eps compiled in. Gate partial uses split-K 1. 10 buckets — the decode ladder only; prefill chunks cross the recurrence through FlashKDA. | -| `k3.kda.o_norm_gate` | `pegainfer-k3` | `ops::k3_o_norm_gate_batched_launch` | `k3_o_norm_gate_batched` | `k3_o_norm_gate_batched.cu` | TileLang AOT CUDA | `kda_core`'s tail on its own: per (row, head) the f32 rms_norm of the bf16 attention landing `X` times the o_norm gamma `Go [128]` f32, landed once, times the bf16 sigmoid of the output gate `G2` — word-for-word the fused core's last loop. Chunked prefill computes attention through FlashKDA and finishes rows here; eps compiled in. 15 buckets. | +| `k3.kda.conv_silu_chunk` | `pegainfer-k3` | `ops::k3_conv_silu_chunk_launch` | `k3_conv_silu_chunk_cuda` | `csrc/k3/k3_conv_silu_chunk.cu` | CUDA | Chunked-prefill / verify causal conv + silu over one q/k/v stream of one segment of consecutive rows, reading the projection's `[tokens, 12288]` f32 partial rows `t-3..t` in place (the carried `[3, 12288]` bf16 window supplies the positions before the segment) and writing the window after row `commit_rows - 1` back as the carry. Replaces the chunk path's landing + strided window copies + batched conv + carry copy (CP8 128k anatomy: 182 ms conv + 235 ms `cuMemcpy2DAsync` per deep rank). Term-for-term the batched conv kernel's arithmetic — `bf16(0 + p)` landing, ascending-tap f32 products, `bf16(ca)` then `sb * (1 / (1 + expf(-sb)))` — same `-O3`, no fast-math, so bit-identical. Eight columns per thread, 16 rows per block with the taps in registers and the window sliding through them; 889 -> 476 us at bucket 16896. Shapes are runtime values. | +| `k3.kda.o_norm_gate` | `pegainfer-k3` | `ops::k3_o_norm_gate_batched_launch` | `k3_o_norm_gate_cuda` | `csrc/k3/k3_elementwise.cu` | CUDA | `kda_core`'s tail on its own: per (row, head) the f32 rms_norm of the bf16 attention landing `X` times the o_norm gamma `Go [128]` f32, landed once, times the bf16 sigmoid of the output gate `G2` — word-for-word the fused core's last loop. Sixteen lanes carry a head with eight columns each; the 128-wide xor butterfly of the retired TileLang kernel is reproduced pair for pair, so the scale is bit-identical. Chunked prefill computes attention through FlashKDA and finishes rows here; eps is a runtime argument, head_dim must be 128. | | `k3.moe.router_topk` | `pegainfer-k3` | `ops::k3_router_topk_batched_launch` | `k3_router_topk_batched` | `k3_router_topk_batched.cu` | TileLang AOT CUDA | Sigmoid router + biased top-k over already-merged `[b, E]` f32 score rows, one row per block. Serial O(topk*E) scan by thread 0 with lowest-index tie-break; weights gathered from the **un-biased** scores, denominator `+1e-20`, scaled by the bf16 `Rs[0]`. E ∈ {896 (full table), 224 (4-way EP shard)}, TOPK = 16, × 15 buckets. | | `k3.attnres.scores` | `pegainfer-k3` | `ops::k3_attnres_scores_batched_launch` | `k3_attnres_scores_batched` | `k3_attnres_scores_batched.cu` | TileLang AOT CUDA | Attention-residual candidate scoring, one block per (row, candidate): weightless RMS normalization then a dot with the pre-fused f32 scoring vector `[7168]`. Candidate `NB` is that row's prefix sum, below it its own snapshot history `[b, NB, 7168]`. NB ∈ 1..8 (the history grows one entry per 12 layers over 93 layers) × 15 buckets. | | `k3.attnres.mix` | `pegainfer-k3` | `ops::k3_attnres_mix_batched_launch` | `k3_attnres_mix_batched` | `k3_attnres_mix_batched.cu` | TileLang AOT CUDA | Softmax over each row's `NB+1` scores, then a probability-weighted mix of the **un-normalized** candidates landing bf16 once. Grid `(b, 7168/256)`; each block redoes its row's softmax. NB ∈ 1..8 × 15 buckets. | diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index 78ebc5420..27c531df1 100644 --- a/pegainfer-kernels/build.rs +++ b/pegainfer-kernels/build.rs @@ -1524,26 +1524,10 @@ const K3_TILELANG_LAUNCHERS: &[(&str, &str)] = &[ "k3_rms_norm_rbs_batched", "const void*, const void*, void*, int, int", ), - ( - "k3_land_batched", - "const float*, void*, int, int, int, int, int", - ), ( "k3_land_rms_norm_rbs_batched", "const float*, const void*, void*, int, int, int, int, int", ), - ( - "k3_add2_batched", - "const void*, const void*, void*, int, int", - ), - ( - "k3_mul_sigmoid_batched", - "const void*, const void*, void*, int, int", - ), - ( - "k3_situ_batched", - "const void*, const void*, void*, int, int", - ), ( "k3_conv_silu_batched", "const float*, const float*, const void*, void*, void*, void*, int, int, int, int", @@ -1554,10 +1538,6 @@ const K3_TILELANG_LAUNCHERS: &[(&str, &str)] = &[ const void*, const void*, const float*, const float*, float*, void*, \ int, int, int, int", ), - ( - "k3_o_norm_gate_batched", - "const void*, const void*, const float*, void*, int, int, int", - ), ( "k3_attnres_scores_batched", "const void*, const void*, const float*, float*, int, int, int", diff --git a/pegainfer-kernels/csrc/k3/k3_conv_silu_chunk.cu b/pegainfer-kernels/csrc/k3/k3_conv_silu_chunk.cu new file mode 100644 index 000000000..fbadea4e3 --- /dev/null +++ b/pegainfer-kernels/csrc/k3/k3_conv_silu_chunk.cu @@ -0,0 +1,196 @@ +// Kimi-K3 chunked-prefill causal conv + silu over one q/k/v stream, walking +// the segment's rows in place. +// +// The chunk path used to land the projection to bf16, build every row's +// 3-slot window with strided device copies (row t's slot j <- landed row +// t-3+j, or the carried window for the first rows), run the batched decode +// conv over the materialized windows, and copy the last commit row's successor +// window out as the carry. The CP8 128k anatomy charged that to 182 ms of +// conv plus 235 ms of cuMemcpy2DAsync per deep rank. This kernel reads the +// f32 partial rows t-3..t directly (the neighbours hit L2), takes the first +// rows' missing inputs from the carried window, and writes the carry itself. +// +// The arithmetic is the batched conv kernel's spelling, term for term: +// +// xb = bf16(0 + P[r]) (the landing) +// ca = 0 +// ca = ca + (f32(w_j) * Cw[j]) for j = 0, 1, 2 (oldest first) +// ca = ca + (f32(xb_t) * Cw[3]) +// sb = f32(bf16(ca)) +// Y[t] = bf16(sb * (1 / (1 + expf(0 - sb)))) +// next[j] = w at position commit_rows - 3 + j (successor window) +// +// where w at position r is xb_r for r >= 0 and carry[r + 3] before the +// segment. Written as the same expressions so nvcc contracts the same +// products into the same FMAs; compiled with the same -O3 and no fast-math. +// Eight columns per thread with 16-byte loads and stores, a run of rows per +// block with the taps in registers and the window sliding through them; the +// launcher guarantees inner % 8 == 0 and 16-byte-aligned rows. Deterministic +// and CUDA-graph safe. + +#include "../shared/ffi_guard.cuh" + +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kVec = 8; +constexpr int kState = 3; // K3_CONV_WIDTH - 1 +constexpr int kRowsPerBlock = 16; + +struct Vec8 { + float v[kVec]; +}; + +__device__ __forceinline__ Vec8 land8(const float* __restrict__ p) { + const float4 lo = *reinterpret_cast(p); + const float4 hi = *reinterpret_cast(p + 4); + const float raw[kVec] = {lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}; + Vec8 out; +#pragma unroll + for (int i = 0; i < kVec; ++i) { + float xa = 0.0f; + xa = xa + raw[i]; + out.v[i] = __bfloat162float(__float2bfloat16_rn(xa)); + } + return out; +} + +__device__ __forceinline__ Vec8 load8(const __nv_bfloat16* __restrict__ p) { + const uint4 packed = *reinterpret_cast(p); + const __nv_bfloat162* pairs = reinterpret_cast(&packed); + Vec8 out; +#pragma unroll + for (int i = 0; i < kVec / 2; ++i) { + const float2 f = __bfloat1622float2(pairs[i]); + out.v[2 * i] = f.x; + out.v[2 * i + 1] = f.y; + } + return out; +} + +__device__ __forceinline__ void store8(__nv_bfloat16* __restrict__ p, + const Vec8& x) { + __nv_bfloat162 packed[kVec / 2]; +#pragma unroll + for (int i = 0; i < kVec / 2; ++i) { + packed[i] = __floats2bfloat162_rn(x.v[2 * i], x.v[2 * i + 1]); + } + *reinterpret_cast(p) = *reinterpret_cast(packed); +} + +__device__ __forceinline__ Vec8 taps8(const float* __restrict__ p) { + const float4 lo = *reinterpret_cast(p); + const float4 hi = *reinterpret_cast(p + 4); + return Vec8{{lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}}; +} + +// Each block owns one 8-column-per-thread column slice and a run of +// kRowsPerBlock consecutive rows: the taps live in registers and the window +// slides through them, so every partial element is read once and every +// output written once (the f32 read and the bf16 write are the whole HBM +// story; the taps and the carry come from L2 once per block). +__global__ void __launch_bounds__(kThreads) + conv_silu_chunk_kernel(const float* __restrict__ P, + const float* __restrict__ Cw, + const __nv_bfloat16* __restrict__ carry, + __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ next, int tokens, + int commit_rows, int inner) { + const int c = (blockIdx.x * kThreads + threadIdx.x) * kVec; + if (c >= inner) return; + const int t0 = blockIdx.y * kRowsPerBlock; + const int t1 = min(t0 + kRowsPerBlock, tokens); + + Vec8 cw[kState + 1]; +#pragma unroll + for (int j = 0; j <= kState; ++j) { + cw[j] = taps8(Cw + (size_t)j * inner + c); + } + + // The window entering row t0: positions t0-3 .. t0-1, the segment's own + // rows once they exist, the carried window before it. + Vec8 w[kState + 1]; +#pragma unroll + for (int j = 0; j < kState; ++j) { + const int r = t0 - kState + j; + w[j] = r >= 0 ? land8(P + (size_t)r * inner + c) + : load8(carry + (size_t)(r + kState) * inner + c); + } + + for (int t = t0; t < t1; ++t) { + w[kState] = land8(P + (size_t)t * inner + c); + Vec8 y; +#pragma unroll + for (int i = 0; i < kVec; ++i) { + float ca = 0.0f; +#pragma unroll + for (int j = 0; j < kState; ++j) { + ca = ca + (w[j].v[i] * cw[j].v[i]); + } + ca = ca + (w[kState].v[i] * cw[kState].v[i]); + const float sb = __bfloat162float(__float2bfloat16_rn(ca)); + y.v[i] = sb * (1.0f / (1.0f + expf(0.0f - sb))); + } + store8(Y + (size_t)t * inner + c, y); + // The carry into the next segment: the successor window of the last + // commit row, i.e. the inputs at positions commit_rows-3 .. commit_rows-1. + if (next != nullptr && t == commit_rows - 1) { +#pragma unroll + for (int j = 0; j < kState; ++j) { + store8(next + (size_t)j * inner + c, w[j + 1]); + } + } +#pragma unroll + for (int j = 0; j < kState; ++j) { + w[j] = w[j + 1]; + } + } +} + +CUresult map_cuda_error(cudaError_t err) { + if (err == cudaSuccess) return CUDA_SUCCESS; + if (err == cudaErrorInvalidValue) return CUDA_ERROR_INVALID_VALUE; + if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; + if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; + return CUDA_ERROR_LAUNCH_FAILED; +} + +bool aligned16(const void* p) { return (reinterpret_cast(p) & 15) == 0; } + +} // namespace + +extern "C" { + +// One q/k/v stream of one segment of `tokens` consecutive rows: `p [tokens, +// inner]` f32 partial (the segment's rows), `cw [4, inner]` f32 taps, `carry +// [3, inner]` bf16 window preceding the segment, `y [tokens, inner]` bf16 +// conv+silu output. When `next` is non-null (commit_rows >= 1) it receives +// the `[3, inner]` window carrying into the segment after row commit_rows-1. +// Shapes are runtime values. +CUresult k3_conv_silu_chunk_cuda(const float* p, const float* cw, + const void* carry, void* y, void* next, + int tokens, int commit_rows, int inner, + cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + if (p == nullptr || cw == nullptr || carry == nullptr || y == nullptr || + tokens <= 0 || commit_rows < 0 || commit_rows > tokens || inner <= 0 || + inner % kVec != 0 || !aligned16(p) || !aligned16(cw) || + !aligned16(carry) || !aligned16(y) || !aligned16(next) || + (next == nullptr) != (commit_rows == 0)) { + return CUDA_ERROR_INVALID_VALUE; + } + const dim3 grid((inner / kVec + kThreads - 1) / kThreads, + (tokens + kRowsPerBlock - 1) / kRowsPerBlock); + conv_silu_chunk_kernel<<>>( + p, cw, static_cast(carry), + static_cast<__nv_bfloat16*>(y), static_cast<__nv_bfloat16*>(next), + tokens, commit_rows, inner); + return map_cuda_error(cudaGetLastError()); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +} // extern "C" diff --git a/pegainfer-kernels/csrc/k3/k3_elementwise.cu b/pegainfer-kernels/csrc/k3/k3_elementwise.cu new file mode 100644 index 000000000..f70e9df66 --- /dev/null +++ b/pegainfer-kernels/csrc/k3/k3_elementwise.cu @@ -0,0 +1,271 @@ +// Kimi-K3 bf16 elementwise family, eight columns per thread. +// +// Hand-written replacements for the retired TileLang `add2_batched`, +// `mul_sigmoid_batched`, `situ_batched` and `o_norm_gate_batched` kernels. +// Those walked one element per thread with 2-byte loads and stores and ran at +// 1–2 TB/s on the chunked-prefill rows (16896 x 7168..33792); every launch +// here moves 16 bytes per thread per operand. +// +// Each kernel keeps the retired kernel's arithmetic term for term. TileLang's +// `bfloat16_t` is cutlass's: its `+` is `__hadd`, its `*` is `__hmul`, and +// the bf16 cast is `cvt.rn.bf16.f32` — so the bf16x2 forms below (`add.rn`, +// `mul.rn`, `cvt.rn` per lane) round identically. The f32 chains use the +// same `expf` / `tanhf` / `rsqrtf` and IEEE division, and both compile under +// the same `-O3` without fast-math, so every landing is bit-identical. +// +// `o_norm_gate` reduces 128 squares per (row, head). The retired kernel used +// TileLang's xor butterfly (offsets 64 and 32 through shared memory, then +// 16..1 by shuffle), whose value is the same in every lane; sixteen lanes +// holding eight columns each reproduce that tree exactly — lane xor 8/4/2/1 +// pairs columns 64/32/16/8 apart, then slots j^4/j^2/j^1 pair the rest — +// so the norm scale matches bit for bit. + +#include "../shared/ffi_guard.cuh" + +#include +#include +#include + +#include + +namespace { + +constexpr int kVec = 8; +constexpr int kThreads = 256; +constexpr int kHeadDim = 128; +constexpr int kLanes = kHeadDim / kVec; // lanes per (row, head) +constexpr int kRowsPerBlock = kThreads / kLanes; + +__device__ __forceinline__ void unpack8(uint4 v, float f[kVec]) { + const __nv_bfloat162* p = reinterpret_cast(&v); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float2 t = __bfloat1622float2(p[j]); + f[2 * j] = t.x; + f[2 * j + 1] = t.y; + } +} + +__device__ __forceinline__ uint4 pack8(const float f[kVec]) { + uint4 v; + __nv_bfloat162* p = reinterpret_cast<__nv_bfloat162*>(&v); +#pragma unroll + for (int j = 0; j < 4; ++j) { + p[j] = __floats2bfloat162_rn(f[2 * j], f[2 * j + 1]); + } + return v; +} + +__device__ __forceinline__ float sigmoid_f32(float x) { + return 1.0f / (1.0f + expf(0.0f - x)); +} + +// O = A + Bt in bf16 addition. +__global__ void add2_kernel(const uint4* __restrict__ a, const uint4* __restrict__ bt, + uint4* __restrict__ o, long long total_vec) { + long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_vec) { + return; + } + uint4 va = a[idx]; + uint4 vb = bt[idx]; + uint4 vo; + const __nv_bfloat162* pa = reinterpret_cast(&va); + const __nv_bfloat162* pb = reinterpret_cast(&vb); + __nv_bfloat162* po = reinterpret_cast<__nv_bfloat162*>(&vo); +#pragma unroll + for (int j = 0; j < 4; ++j) { + po[j] = __hadd2(pa[j], pb[j]); + } + o[idx] = vo; +} + +// O = A * bf16(sigmoid(Bt)): the sigmoid in f32, landed bf16, then the bf16 +// product. +__global__ void mul_sigmoid_kernel(const uint4* __restrict__ a, const uint4* __restrict__ bt, + uint4* __restrict__ o, long long total_vec) { + long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_vec) { + return; + } + uint4 va = a[idx]; + float b[kVec]; + unpack8(bt[idx], b); + float s[kVec]; +#pragma unroll + for (int j = 0; j < kVec; ++j) { + s[j] = sigmoid_f32(b[j]); + } + uint4 vs = pack8(s); + uint4 vo; + const __nv_bfloat162* pa = reinterpret_cast(&va); + const __nv_bfloat162* ps = reinterpret_cast(&vs); + __nv_bfloat162* po = reinterpret_cast<__nv_bfloat162*>(&vo); +#pragma unroll + for (int j = 0; j < 4; ++j) { + po[j] = __hmul2(pa[j], ps[j]); + } + o[idx] = vo; +} + +// O = bf16(4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)), the f32 chain landed once. +__global__ void situ_kernel(const uint4* __restrict__ g, const uint4* __restrict__ u, + uint4* __restrict__ o, long long total_vec) { + long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_vec) { + return; + } + float gf[kVec]; + float uf[kVec]; + unpack8(g[idx], gf); + unpack8(u[idx], uf); + float r[kVec]; +#pragma unroll + for (int j = 0; j < kVec; ++j) { + r[j] = ((4.0f * tanhf(gf[j] / 4.0f)) * sigmoid_f32(gf[j])) * (25.0f * tanhf(uf[j] / 25.0f)); + } + o[idx] = pack8(r); +} + +// Out = bf16(x * rsqrt(mean(x^2) + eps) * go) * bf16(sigmoid(g2)) per +// (row, head) of 128 columns. `rows` counts (row, head) pairs; a lane past +// the last pair clamps its loads and skips its store so the shuffles stay +// converged. +__global__ void o_norm_gate_kernel(const uint4* __restrict__ x, const uint4* __restrict__ g2, + const float4* __restrict__ go, uint4* __restrict__ out, + long long rows, float eps) { + long long row = (long long)blockIdx.x * kRowsPerBlock + threadIdx.x / kLanes; + int lane = threadIdx.x % kLanes; + bool live = row < rows; + long long vec = (live ? row : rows - 1) * kLanes + lane; + float xf[kVec]; + unpack8(x[vec], xf); + float s[kVec]; +#pragma unroll + for (int j = 0; j < kVec; ++j) { + float sq = xf[j] * xf[j]; + s[j] = 0.0f + sq; + } +#pragma unroll + for (int off = kLanes / 2; off >= 1; off >>= 1) { +#pragma unroll + for (int j = 0; j < kVec; ++j) { + s[j] = s[j] + __shfl_xor_sync(0xffffffffu, s[j], off); + } + } + float t[kVec]; +#pragma unroll + for (int j = 0; j < kVec; ++j) { + t[j] = s[j] + s[j ^ 4]; + } +#pragma unroll + for (int j = 0; j < kVec; ++j) { + s[j] = t[j] + t[j ^ 2]; + } +#pragma unroll + for (int j = 0; j < kVec; ++j) { + t[j] = s[j] + s[j ^ 1]; + } + float atot = t[0]; + float scale = rsqrtf((atot / 128.0f) + eps); + float4 go0 = go[lane * 2]; + float4 go1 = go[lane * 2 + 1]; + float gof[kVec] = {go0.x, go0.y, go0.z, go0.w, go1.x, go1.y, go1.z, go1.w}; + float g2f[kVec]; + unpack8(g2[vec], g2f); + float n[kVec]; + float sg[kVec]; +#pragma unroll + for (int j = 0; j < kVec; ++j) { + n[j] = (xf[j] * scale) * gof[j]; + sg[j] = sigmoid_f32(g2f[j]); + } + uint4 vn = pack8(n); + uint4 vg = pack8(sg); + uint4 vo; + const __nv_bfloat162* pn = reinterpret_cast(&vn); + const __nv_bfloat162* pg = reinterpret_cast(&vg); + __nv_bfloat162* po = reinterpret_cast<__nv_bfloat162*>(&vo); +#pragma unroll + for (int j = 0; j < 4; ++j) { + po[j] = __hmul2(pn[j], pg[j]); + } + if (live) { + out[vec] = vo; + } +} + +CUresult map_cuda_error(cudaError_t err) { + if (err == cudaSuccess) return CUDA_SUCCESS; + if (err == cudaErrorInvalidValue) return CUDA_ERROR_INVALID_VALUE; + if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; + if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; + return CUDA_ERROR_LAUNCH_FAILED; +} + +bool aligned16(const void* p) { return (reinterpret_cast(p) & 15) == 0; } + +template +CUresult launch_binary(Kernel kernel, const void* a, const void* bt, void* o, int b, int n, + cudaStream_t stream) { + if (a == nullptr || bt == nullptr || o == nullptr || b <= 0 || n <= 0 || n % kVec != 0 || + !aligned16(a) || !aligned16(bt) || !aligned16(o)) { + return CUDA_ERROR_INVALID_VALUE; + } + long long total_vec = (long long)b * n / kVec; + long long blocks = (total_vec + kThreads - 1) / kThreads; + kernel<<<(unsigned)blocks, kThreads, 0, stream>>>( + static_cast(a), static_cast(bt), static_cast(o), + total_vec); + return map_cuda_error(cudaGetLastError()); +} + +} // namespace + +extern "C" { + +// O = A + Bt in bf16 addition, all [b, n]. n must be a multiple of 8. +CUresult k3_add2_cuda(const void* a, const void* bt, void* o, int b, int n, + cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + return launch_binary(add2_kernel, a, bt, o, b, n, stream); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +// O = A * bf16(sigmoid(Bt)), all [b, n]. n must be a multiple of 8. +CUresult k3_mul_sigmoid_cuda(const void* a, const void* bt, void* o, int b, int n, + cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + return launch_binary(mul_sigmoid_kernel, a, bt, o, b, n, stream); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +// O = situ(G, U), all [b, n]. n must be a multiple of 8. +CUresult k3_situ_cuda(const void* g, const void* u, void* o, int b, int n, + cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + return launch_binary(situ_kernel, g, u, o, b, n, stream); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +// Out[b, heads * 128] = bf16(rms_norm(X) * Go) * bf16(sigmoid(G2)) per +// (row, head); Go is [128] f32. head_dim must be 128 (the reduction tree is +// the 128-wide one). +CUresult k3_o_norm_gate_cuda(const void* x, const void* g2, const float* go, void* out, int b, + int heads, int head_dim, float eps, cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + if (x == nullptr || g2 == nullptr || go == nullptr || out == nullptr || b <= 0 || + heads <= 0 || head_dim != kHeadDim || !aligned16(x) || !aligned16(g2) || + !aligned16(go) || !aligned16(out)) { + return CUDA_ERROR_INVALID_VALUE; + } + long long rows = (long long)b * heads; + long long blocks = (rows + kRowsPerBlock - 1) / kRowsPerBlock; + o_norm_gate_kernel<<<(unsigned)blocks, kThreads, 0, stream>>>( + static_cast(x), static_cast(g2), + reinterpret_cast(go), static_cast(out), rows, eps); + return map_cuda_error(cudaGetLastError()); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +} // extern "C" diff --git a/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu b/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu index a03a17170..bbb789c51 100644 --- a/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu +++ b/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu @@ -43,6 +43,7 @@ #include #include +#include #ifdef K3_FLASH_MLA_SM100F #include "collective/fmha_fusion.hpp" @@ -172,55 +173,107 @@ extern "C" CUresult k3_mla_prefill_gather( } // Fold one window's FMHA output into the running f32 accumulator via the -// log-sum-exp identity. One block per (q, head) row: only this block touches -// its lse entry, so the read-old/write-new sequence needs no cross-block -// coordination. `reset` starts a fresh accumulation (no -inf seeding). +// log-sum-exp identity. The row walk is bandwidth work (2.1 GB per merge at +// 16896 x 96 x 128), so a thread carries 8 nope columns as one 16-byte bf16 +// load and two float4 accumulator loads/stores: 16 lanes cover one (q, head) +// row and a 256-thread block walks 16 rows. Every lane of a row recomputes +// the two weights from the same lse pair (a broadcast load, identical +// arithmetic), and lane 0 alone writes the merged lse — only this row's +// lanes touch its entry, and the reads precede the divergent store in +// program order, so no cross-lane coordination is needed. `reset` starts a +// fresh accumulation (no -inf seeding). The per-element update keeps the +// exact expression of the scalar original. +static constexpr int K3_FMP_MERGE_LANES = K3_FMP_NOPE / 8; +static constexpr int K3_FMP_MERGE_ROWS = 256 / K3_FMP_MERGE_LANES; + +static __device__ __forceinline__ void k3_fmp_merge8( + float4& a0, float4& a1, uint4 w, float w_acc, float w_win +) { + const __nv_bfloat162* wp = reinterpret_cast(&w); + float2 f0 = __bfloat1622float2(wp[0]); + float2 f1 = __bfloat1622float2(wp[1]); + float2 f2 = __bfloat1622float2(wp[2]); + float2 f3 = __bfloat1622float2(wp[3]); + a0.x = a0.x * w_acc + f0.x * w_win; + a0.y = a0.y * w_acc + f0.y * w_win; + a0.z = a0.z * w_acc + f1.x * w_win; + a0.w = a0.w * w_acc + f1.y * w_win; + a1.x = a1.x * w_acc + f2.x * w_win; + a1.y = a1.y * w_acc + f2.y * w_win; + a1.z = a1.z * w_acc + f3.x * w_win; + a1.w = a1.w * w_acc + f3.y * w_win; +} + static __global__ void k3_mla_prefill_lse_merge_kernel( - const __nv_bfloat16* __restrict__ o_win, + const uint4* __restrict__ o_win, const float* __restrict__ lse_win, - float* __restrict__ o_acc, + float4* __restrict__ o_acc, float* __restrict__ lse_acc, int t_q, int heads, int reset ) { - int q = blockIdx.x; - int h = blockIdx.y; - int d = threadIdx.x; + long long rows = (long long)t_q * heads; + long long row = (long long)blockIdx.x * K3_FMP_MERGE_ROWS + threadIdx.x / K3_FMP_MERGE_LANES; + int lane = threadIdx.x % K3_FMP_MERGE_LANES; + if (row >= rows) { + return; + } + int q = (int)(row / heads); + int h = (int)(row % heads); long long lse_idx = (long long)h * t_q + q; - __shared__ float w_acc, w_win; - if (d == 0) { - float lw = lse_win[lse_idx]; - if (reset) { - w_acc = 0.0f; - w_win = 1.0f; + float w_acc, w_win; + float lw = lse_win[lse_idx]; + if (reset) { + w_acc = 0.0f; + w_win = 1.0f; + if (lane == 0) { lse_acc[lse_idx] = lw; - } else { - float la = lse_acc[lse_idx]; - float m = fmaxf(la, lw); - float merged = m + logf(expf(la - m) + expf(lw - m)); - w_acc = expf(la - merged); - w_win = expf(lw - merged); + } + } else { + float la = lse_acc[lse_idx]; + float m = fmaxf(la, lw); + float merged = m + logf(expf(la - m) + expf(lw - m)); + w_acc = expf(la - merged); + w_win = expf(lw - merged); + if (lane == 0) { lse_acc[lse_idx] = merged; } } - __syncthreads(); - long long o_idx = ((long long)q * heads + h) * K3_FMP_NOPE + d; - float prev = reset ? 0.0f : o_acc[o_idx]; - o_acc[o_idx] = prev * w_acc + __bfloat162float(o_win[o_idx]) * w_win; + long long vec = row * K3_FMP_MERGE_LANES + lane; + uint4 w = o_win[vec]; + float4 a0, a1; + if (reset) { + a0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + a1 = a0; + } else { + a0 = o_acc[vec * 2]; + a1 = o_acc[vec * 2 + 1]; + } + k3_fmp_merge8(a0, a1, w, w_acc, w_win); + o_acc[vec * 2] = a0; + o_acc[vec * 2 + 1] = a1; } -// Leave the f32 accumulator: out[t_q, h, 128] bf16. +// Leave the f32 accumulator: out[t_q, h, 128] bf16, 8 columns per thread. static __global__ void k3_mla_prefill_o_finalize_kernel( - const float* __restrict__ o_acc, - __nv_bfloat16* __restrict__ out, - long long total + const float4* __restrict__ o_acc, + uint4* __restrict__ out, + long long total_vec ) { long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) { + if (idx >= total_vec) { return; } - out[idx] = __float2bfloat16(o_acc[idx]); + float4 a0 = o_acc[idx * 2]; + float4 a1 = o_acc[idx * 2 + 1]; + uint4 packed; + __nv_bfloat162* p = reinterpret_cast<__nv_bfloat162*>(&packed); + p[0] = __floats2bfloat162_rn(a0.x, a0.y); + p[1] = __floats2bfloat162_rn(a0.z, a0.w); + p[2] = __floats2bfloat162_rn(a1.x, a1.y); + p[3] = __floats2bfloat162_rn(a1.z, a1.w); + out[idx] = packed; } extern "C" CUresult k3_mla_prefill_lse_merge( @@ -237,9 +290,13 @@ extern "C" CUresult k3_mla_prefill_lse_merge( || t_q <= 0 || heads <= 0) { return CUDA_ERROR_INVALID_VALUE; } - dim3 grid((unsigned)t_q, (unsigned)heads); - k3_mla_prefill_lse_merge_kernel<<>>( - (const __nv_bfloat16*)o_win, (const float*)lse_win, (float*)o_acc, (float*)lse_acc, + if (((uintptr_t)o_win & 15) != 0 || ((uintptr_t)o_acc & 15) != 0) { + return CUDA_ERROR_INVALID_VALUE; + } + long long rows = (long long)t_q * heads; + long long blocks = (rows + K3_FMP_MERGE_ROWS - 1) / K3_FMP_MERGE_ROWS; + k3_mla_prefill_lse_merge_kernel<<<(unsigned)blocks, 256, 0, (cudaStream_t)stream>>>( + (const uint4*)o_win, (const float*)lse_win, (float4*)o_acc, (float*)lse_acc, t_q, heads, reset); return k3_flash_mla_consume_last_cuda_error(); } @@ -254,11 +311,14 @@ extern "C" CUresult k3_mla_prefill_o_finalize( if (o_acc == nullptr || out == nullptr || t_q <= 0 || heads <= 0) { return CUDA_ERROR_INVALID_VALUE; } - long long total = (long long)t_q * heads * K3_FMP_NOPE; + if (((uintptr_t)o_acc & 15) != 0 || ((uintptr_t)out & 15) != 0) { + return CUDA_ERROR_INVALID_VALUE; + } + long long total_vec = (long long)t_q * heads * (K3_FMP_NOPE / 8); int threads = 256; - long long blocks = (total + threads - 1) / threads; + long long blocks = (total_vec + threads - 1) / threads; k3_mla_prefill_o_finalize_kernel<<<(unsigned)blocks, threads, 0, (cudaStream_t)stream>>>( - (const float*)o_acc, (__nv_bfloat16*)out, total); + (const float4*)o_acc, (uint4*)out, total_vec); return k3_flash_mla_consume_last_cuda_error(); } diff --git a/pegainfer-kernels/csrc/k3/k3_land.cu b/pegainfer-kernels/csrc/k3/k3_land.cu new file mode 100644 index 000000000..790d33a55 --- /dev/null +++ b/pegainfer-kernels/csrc/k3/k3_land.cu @@ -0,0 +1,130 @@ +// Kimi-K3 matmul landing: merge the column span [off, off+n) of each row's +// [split_k, nt] f32 partial and land bf16 once. +// +// Replaces the retired TileLang `land_batched` kernel, which landed one +// element per thread with 2-byte stores and reached 2.3 TB/s on the 12288-wide +// chunked-prefill landings (nsys, CP8 128k anatomy: 385 ms of a 6.1 s deep +// rank). Each thread here lands 8 consecutive columns — two float4 loads per +// segment, one 16-byte store — so the pass streams at HBM rate. +// +// The arithmetic is the retired kernel's spelling: the segments are summed in +// f32 in ascending s order onto a zero accumulator (so a -0 partial lands as +// +0, as it did), then cast once, round-to-nearest-even. At split_k = 1 — the only +// launch site — that is the bare cast, so the landing is bit-identical to the +// retired kernel and the certified single-row spelling. +// +// The vector path needs every row 16-byte aligned in both tensors: nt and off +// multiples of 4, n a multiple of 8. Every K3 span satisfies this; anything +// else takes the scalar path. Batch is a runtime value — no per-bucket +// instantiation. Deterministic and CUDA-graph safe. + +#include "../shared/ffi_guard.cuh" + +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kVec = 8; +constexpr long long kMaxBlocks = 1 << 20; + +__device__ __forceinline__ void add4(float4& a, const float4& b) { + a.x += b.x; + a.y += b.y; + a.z += b.z; + a.w += b.w; +} + +__global__ void land_vec8_kernel(const float* __restrict__ P, + __nv_bfloat16* __restrict__ O, int nt, int n, + int off, int split_k, long long total_vec) { + const int vec_per_row = n / kVec; + for (long long v = (long long)blockIdx.x * blockDim.x + threadIdx.x; + v < total_vec; v += (long long)gridDim.x * blockDim.x) { + const long long row = v / vec_per_row; + const int col = (int)(v - row * vec_per_row) * kVec; + const float* src = P + row * (long long)split_k * nt + off + col; + float4 lo = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 hi = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + add4(lo, *reinterpret_cast(src)); + add4(hi, *reinterpret_cast(src + 4)); + for (int s = 1; s < split_k; ++s) { + src += nt; + add4(lo, *reinterpret_cast(src)); + add4(hi, *reinterpret_cast(src + 4)); + } + __nv_bfloat162 packed[4]; + packed[0] = __floats2bfloat162_rn(lo.x, lo.y); + packed[1] = __floats2bfloat162_rn(lo.z, lo.w); + packed[2] = __floats2bfloat162_rn(hi.x, hi.y); + packed[3] = __floats2bfloat162_rn(hi.z, hi.w); + *reinterpret_cast(O + row * (long long)n + col) = + *reinterpret_cast(packed); + } +} + +__global__ void land_scalar_kernel(const float* __restrict__ P, + __nv_bfloat16* __restrict__ O, int nt, int n, + int off, int split_k, long long total) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + i < total; i += (long long)gridDim.x * blockDim.x) { + const long long row = i / n; + const int col = (int)(i - row * n); + const float* src = P + row * (long long)split_k * nt + off + col; + float acc = 0.0f; + acc = acc + src[0]; + for (int s = 1; s < split_k; ++s) { + acc += src[(long long)s * nt]; + } + O[i] = __float2bfloat16_rn(acc); + } +} + +CUresult map_cuda_error(cudaError_t err) { + if (err == cudaSuccess) return CUDA_SUCCESS; + if (err == cudaErrorInvalidValue) return CUDA_ERROR_INVALID_VALUE; + if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; + if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; + return CUDA_ERROR_LAUNCH_FAILED; +} + +bool aligned16(const void* p) { return (reinterpret_cast(p) & 15) == 0; } + +} // namespace + +extern "C" { + +// Merge the column span [off, off+n) of each row's [split_k, nt] f32 partial +// `p [b, split_k, nt]` and land `o [b, n]` bf16 once. Shapes are runtime +// values; split_k = 1 is the single partial a framework GEMM produces, where +// the merge degenerates to the slice and the cast. +CUresult k3_land_cuda(const float* p, void* o, int b, int nt, int n, int off, + int split_k, cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + if (p == nullptr || o == nullptr || b <= 0 || nt <= 0 || n <= 0 || off < 0 || + split_k <= 0 || off + n > nt) { + return CUDA_ERROR_INVALID_VALUE; + } + __nv_bfloat16* out = static_cast<__nv_bfloat16*>(o); + const bool vec = aligned16(p) && aligned16(o) && nt % 4 == 0 && + off % 4 == 0 && n % kVec == 0; + if (vec) { + const long long total = (long long)b * (n / kVec); + const long long blocks = (total + kThreads - 1) / kThreads; + land_vec8_kernel<<<(unsigned)(blocks < kMaxBlocks ? blocks : kMaxBlocks), + kThreads, 0, stream>>>(p, out, nt, n, off, split_k, + total); + } else { + const long long total = (long long)b * n; + const long long blocks = (total + kThreads - 1) / kThreads; + land_scalar_kernel<<<(unsigned)(blocks < kMaxBlocks ? blocks : kMaxBlocks), + kThreads, 0, stream>>>(p, out, nt, n, off, split_k, + total); + } + return map_cuda_error(cudaGetLastError()); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +} // extern "C" diff --git a/pegainfer-kernels/src/ffi/k3.rs b/pegainfer-kernels/src/ffi/k3.rs index ad463f0e3..6bfdf0ead 100644 --- a/pegainfer-kernels/src/ffi/k3.rs +++ b/pegainfer-kernels/src/ffi/k3.rs @@ -173,6 +173,90 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + /// Matmul landing (`csrc/k3/k3_land.cu`): merge the column span + /// `[off, off + n)` of each row's `p [b, split_k, nt]` f32 partial and + /// land `o [b, n]` bf16 once — ascending-segment f32 sum, one + /// round-to-nearest-even cast, bit-identical to the retired TileLang + /// kernel. Shapes are runtime values (no per-bucket instantiation). + pub fn k3_land_cuda( + p: *const f32, + o: *mut c_void, + b: i32, + nt: i32, + n: i32, + off: i32, + split_k: i32, + stream: CUstream, + ) -> CUresult; + + /// bf16 elementwise family (`csrc/k3/k3_elementwise.cu`), eight columns + /// per thread, bit-identical to the retired TileLang batched kernels. + /// `O = A + Bt` in bf16 addition, all `[b, n]`; n a multiple of 8. + pub fn k3_add2_cuda( + a: *const c_void, + bt: *const c_void, + o: *mut c_void, + b: i32, + n: i32, + stream: CUstream, + ) -> CUresult; + + /// `O = A * bf16(sigmoid(Bt))`, the MLA sigmoid output gate. All `[b, n]`. + pub fn k3_mul_sigmoid_cuda( + a: *const c_void, + bt: *const c_void, + o: *mut c_void, + b: i32, + n: i32, + stream: CUstream, + ) -> CUresult; + + /// K3's situ activation `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)`, computed + /// in f32 and landed bf16 once. All `[b, n]`. + pub fn k3_situ_cuda( + g: *const c_void, + u: *const c_void, + o: *mut c_void, + b: i32, + n: i32, + stream: CUstream, + ) -> CUresult; + + /// `kda_core`'s tail on its own: per (row, head) f32 rms_norm of the bf16 + /// attention landing `X` times the o_norm gamma `Go [128]`, landed once, + /// times the bf16 sigmoid of the output gate `G2`. head_dim must be 128. + pub fn k3_o_norm_gate_cuda( + x: *const c_void, + g2: *const c_void, + go: *const f32, + out: *mut c_void, + b: i32, + num_heads: i32, + head_dim: i32, + eps: f32, + stream: CUstream, + ) -> CUresult; + + /// Chunked-prefill conv + silu over one q/k/v stream of one segment + /// (`csrc/k3/k3_conv_silu_chunk.cu`): `p [tokens, inner]` f32 partial, + /// `cw [4, inner]` f32 taps, `carry [3, inner]` bf16 window preceding + /// the segment, `y [tokens, inner]` bf16 output; `next [3, inner]` bf16 + /// receives the window after row `commit_rows - 1` and is null exactly + /// when `commit_rows == 0`. Reads the partial rows in place — no window + /// is materialized — with the batched conv kernel's arithmetic term for + /// term. Shapes are runtime values. + pub fn k3_conv_silu_chunk_cuda( + p: *const f32, + cw: *const f32, + carry: *const c_void, + y: *mut c_void, + next: *mut c_void, + tokens: i32, + commit_rows: i32, + inner: i32, + 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/ffi/k3_tilelang.rs b/pegainfer-kernels/src/ffi/k3_tilelang.rs index bd9f4f8f8..aba56ad4a 100644 --- a/pegainfer-kernels/src/ffi/k3_tilelang.rs +++ b/pegainfer-kernels/src/ffi/k3_tilelang.rs @@ -33,21 +33,7 @@ unsafe extern "C" { stream: CUstream, ) -> i32; - /// Merge the column span `[off, off + n)` of each row's `P [b, split_k, nt]` - /// f32 partial and land `O [b, n]` bf16 once. `split_k = 1` is the - /// single-partial case a framework GEMM produces. - pub fn k3_land_batched( - p: *const f32, - o: *mut c_void, - b: i32, - nt: i32, - n: i32, - off: i32, - split_k: i32, - stream: CUstream, - ) -> i32; - - /// `k3_land_batched` followed by the round-before-scale norm against the + /// The matmul landing (`k3_land_cuda`) followed by the round-before-scale norm against the /// shared gamma `G [n]`. pub fn k3_land_rms_norm_rbs_batched( p: *const f32, @@ -61,37 +47,6 @@ unsafe extern "C" { stream: CUstream, ) -> i32; - /// `O = A + Bt` in bf16 addition, all `[b, n]`. - pub fn k3_add2_batched( - a: *const c_void, - bt: *const c_void, - o: *mut c_void, - b: i32, - n: i32, - stream: CUstream, - ) -> i32; - - /// `O = A * bf16(sigmoid(Bt))`, the MLA sigmoid output gate. All `[b, n]`. - pub fn k3_mul_sigmoid_batched( - a: *const c_void, - bt: *const c_void, - o: *mut c_void, - b: i32, - n: i32, - stream: CUstream, - ) -> i32; - - /// K3's situ activation `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)`, computed - /// in f32 and landed bf16 once. All `[b, n]`. - pub fn k3_situ_batched( - g: *const c_void, - u: *const c_void, - o: *mut c_void, - b: i32, - n: i32, - stream: CUstream, - ) -> i32; - /// Causal depthwise convolution over the `width`-slot window plus silu. /// `P [b, split_k, kp]` f32 partials land into `X [b, kp]` bf16, which is /// also the newest window slot; `Cs`/`Sn [b, width - 1, kp]` are the carried @@ -137,20 +92,6 @@ unsafe extern "C" { stream: CUstream, ) -> i32; - /// `kda_core`'s tail on its own: per (row, head) f32 rms_norm of the bf16 - /// attention landing `X` times the o_norm gamma `Go [head_dim]`, landed - /// once, times the bf16 sigmoid of the output gate `G2`. eps compiled in. - pub fn k3_o_norm_gate_batched( - x: *const c_void, - g2: *const c_void, - go: *const f32, - out: *mut c_void, - b: i32, - num_heads: i32, - head_dim: i32, - stream: CUstream, - ) -> i32; - /// Attention-residual candidate scoring: weightless RMS normalization then /// a dot with the fused f32 scoring vector `Sw [h]`. `Ps [b, h]` is the /// running prefix sum, `Bl [b, blocks, h]` that row's snapshot history; diff --git a/pegainfer-kernels/src/ops/k3/conv_silu_chunk.rs b/pegainfer-kernels/src/ops/k3/conv_silu_chunk.rs new file mode 100644 index 000000000..ee57dcfd5 --- /dev/null +++ b/pegainfer-kernels/src/ops/k3/conv_silu_chunk.rs @@ -0,0 +1,110 @@ +//! Kimi-K3 chunked-prefill causal conv + silu, walking a segment's rows in +//! place. +//! +//! Replaces the chunk path's landing + strided window copies + batched decode +//! conv + carry copy with one launch per segment that reads the f32 partial +//! rows `t-3..t` directly and takes the first rows' missing inputs from the +//! carried window. Term-for-term the batched conv kernel's arithmetic, so +//! bit-identical — see `csrc/k3/k3_conv_silu_chunk.cu`. + +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::tensor::DeviceContext; + +/// Window slots carried between segments (`K3_CONV_WIDTH - 1`). +const STATE: usize = 3; + +/// Conv + silu over `tokens` consecutive rows of one q/k/v stream. +/// +/// `partial` rows `partial_row..partial_row + tokens` (each `inner` f32) are +/// the segment's projection; `taps` is `[4, inner]` f32; `carry` row +/// `carry_row` (each `STATE * inner` bf16) is the window preceding the +/// segment. `out` rows `out_row..` receive the bf16 conv output. `next`, when +/// given, names the `[STATE, inner]` row that receives the window carrying +/// into the segment after row `commit_rows - 1`; it must be given exactly +/// when `commit_rows > 0`. +#[allow(clippy::too_many_arguments)] +pub fn k3_conv_silu_chunk_launch( + ctx: &DeviceContext, + inner: usize, + tokens: usize, + commit_rows: usize, + partial: &CudaSlice, + partial_row: usize, + taps: &CudaSlice, + carry: &CudaSlice, + carry_row: usize, + out: &mut CudaSlice, + out_row: usize, + next: Option<(&mut CudaSlice, usize)>, +) -> Result<()> { + ensure!(tokens > 0, "K3 conv chunk needs rows"); + ensure!( + commit_rows <= tokens, + "K3 conv chunk commits {commit_rows} of {tokens} rows" + ); + ensure!( + inner % 8 == 0, + "K3 conv chunk needs an 8-column-aligned width, got {inner}" + ); + ensure!( + (commit_rows > 0) == next.is_some(), + "K3 conv chunk carries a window exactly when it commits rows" + ); + ensure!( + partial.len() >= (partial_row + tokens) * inner + && taps.len() >= 4 * inner + && carry.len() >= (carry_row + 1) * STATE * inner + && out.len() >= (out_row + tokens) * inner, + "K3 conv chunk buffers too small for tokens={tokens}, inner={inner}: \ + partial {} (row {partial_row}), taps {}, carry {} (row {carry_row}), out {} (row {out_row})", + partial.len(), + taps.len(), + carry.len(), + out.len() + ); + let (p_ptr, _p_guard) = partial.device_ptr(&ctx.stream); + let (cw_ptr, _cw_guard) = taps.device_ptr(&ctx.stream); + let (carry_ptr, _carry_guard) = carry.device_ptr(&ctx.stream); + let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); + let mut next_guard = None; + let next_ptr = match next { + Some((slab, row)) => { + ensure!( + slab.len() >= (row + 1) * STATE * inner, + "K3 conv chunk carry-out slab too small for row {row}: {}", + slab.len() + ); + let (ptr, guard) = slab.device_ptr_mut(&ctx.stream); + next_guard = Some(guard); + ptr as usize + row * STATE * inner * size_of::() + } + None => 0, + }; + let rc = unsafe { + ffi::k3_conv_silu_chunk_cuda( + (p_ptr as usize + partial_row * inner * size_of::()) as *const f32, + cw_ptr as *const f32, + (carry_ptr as usize + carry_row * STATE * inner * size_of::()) as *const c_void, + (out_ptr as usize + out_row * inner * size_of::()) as *mut c_void, + next_ptr as *mut c_void, + i32::try_from(tokens)?, + i32::try_from(commit_rows)?, + i32::try_from(inner)?, + crate::tensor::active_cu_stream(ctx), + ) + }; + drop(next_guard); + rc.result().map_err(|err| { + anyhow!("K3 conv_silu_chunk (tokens={tokens}, commit={commit_rows}, inner={inner}) launch failed: {err}") + }) +} diff --git a/pegainfer-kernels/src/ops/k3/elementwise.rs b/pegainfer-kernels/src/ops/k3/elementwise.rs new file mode 100644 index 000000000..c1fb5605a --- /dev/null +++ b/pegainfer-kernels/src/ops/k3/elementwise.rs @@ -0,0 +1,161 @@ +//! Kimi-K3 bf16 elementwise family: the residual add, the MLA sigmoid +//! output gate, the situ activation, and the KDA o_norm x output gate. +//! +//! Hand-written replacements for the retired TileLang batched kernels (one +//! element per thread, 2-byte accesses, 1–2 TB/s on the chunked-prefill +//! rows). Eight columns per thread with 16-byte loads and stores and the +//! retired kernels' exact arithmetic — cutlass's bf16 `+`/`*` are `__hadd` / +//! `__hmul`, the casts `cvt.rn`, and `o_norm_gate`'s 128-wide xor butterfly +//! is reproduced pair for pair — so every landing is bit-identical; see +//! `csrc/k3/k3_elementwise.cu`. Batch is a plain launch dimension. + +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::tensor::DeviceContext; + +type BinaryLauncher = unsafe extern "C" fn( + *const c_void, + *const c_void, + *mut c_void, + i32, + i32, + cudarc::driver::sys::CUstream, +) -> cudarc::driver::sys::CUresult; + +fn binary_launch( + ctx: &DeviceContext, + launcher: BinaryLauncher, + what: &str, + b: usize, + n: usize, + a: &CudaSlice, + bt: &CudaSlice, + o: &mut CudaSlice, +) -> Result<()> { + ensure!(b > 0 && n > 0, "K3 {what} needs rows and columns"); + ensure!( + a.len() >= b * n && bt.len() >= b * n && o.len() >= b * n, + "K3 {what} buffers too small for b={b}, n={n}: a {}, bt {}, o {}", + a.len(), + bt.len(), + o.len() + ); + let (a_ptr, _a_guard) = a.device_ptr(&ctx.stream); + let (bt_ptr, _bt_guard) = bt.device_ptr(&ctx.stream); + let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); + unsafe { + launcher( + a_ptr as *const c_void, + bt_ptr as *const c_void, + o_ptr as *mut c_void, + i32::try_from(b)?, + i32::try_from(n)?, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| anyhow!("K3 {what} (B={b}, N={n}) launch failed: {err}")) +} + +/// `o = a + bt`, added in bf16 — the residual adds, and routed + shared. +pub fn k3_add2_batched_launch( + ctx: &DeviceContext, + b: usize, + n: usize, + a: &CudaSlice, + bt: &CudaSlice, + o: &mut CudaSlice, +) -> Result<()> { + binary_launch(ctx, ffi::k3_add2_cuda, "add2", b, n, a, bt, o) +} + +/// `o = a * bf16(sigmoid(bt))`, the MLA sigmoid output gate. The sigmoid is +/// taken in f32 and lands in bf16 before the product. +pub fn k3_mul_sigmoid_batched_launch( + ctx: &DeviceContext, + b: usize, + n: usize, + a: &CudaSlice, + bt: &CudaSlice, + o: &mut CudaSlice, +) -> Result<()> { + binary_launch(ctx, ffi::k3_mul_sigmoid_cuda, "mul_sigmoid", b, n, a, bt, o) +} + +/// K3's situ activation: `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)` in f32, +/// landing bf16 once. The two betas are compiled in. +pub fn k3_situ_batched_launch( + ctx: &DeviceContext, + b: usize, + n: usize, + g: &CudaSlice, + u: &CudaSlice, + o: &mut CudaSlice, +) -> Result<()> { + binary_launch(ctx, ffi::k3_situ_cuda, "situ", b, n, g, u, o) +} + +/// `kda_core`'s tail on its own: per (row, head) the f32 rms_norm of the bf16 +/// attention landing `x` times the o_norm gamma `go [head_dim]`, landed once, +/// times the bf16 sigmoid of the output-gate projection `g2`. `head_dim` must +/// be 128. +#[allow(clippy::too_many_arguments)] +pub fn k3_o_norm_gate_batched_launch( + ctx: &DeviceContext, + b: usize, + num_heads: usize, + head_dim: usize, + eps: f32, + x: &CudaSlice, + g2: &CudaSlice, + go: &CudaSlice, + out: &mut CudaSlice, +) -> Result<()> { + ensure!( + b > 0 && num_heads > 0, + "K3 o_norm_gate needs rows and heads" + ); + ensure!( + head_dim == 128, + "K3 o_norm_gate serves head_dim 128, got {head_dim}" + ); + let kp = num_heads * head_dim; + ensure!( + x.len() >= b * kp && g2.len() >= b * kp && out.len() >= b * kp && go.len() >= head_dim, + "K3 o_norm_gate buffers too small for b={b}, kp={kp}: x {}, g2 {}, go {}, out {}", + x.len(), + g2.len(), + go.len(), + out.len() + ); + let (x_ptr, _x_guard) = x.device_ptr(&ctx.stream); + let (g2_ptr, _g2_guard) = g2.device_ptr(&ctx.stream); + let (go_ptr, _go_guard) = go.device_ptr(&ctx.stream); + let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); + unsafe { + ffi::k3_o_norm_gate_cuda( + x_ptr as *const c_void, + g2_ptr as *const c_void, + go_ptr as *const f32, + out_ptr as *mut c_void, + i32::try_from(b)?, + i32::try_from(num_heads)?, + i32::try_from(head_dim)?, + eps, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| { + anyhow!("K3 o_norm_gate (B={b}, KH={num_heads}, KD={head_dim}) launch failed: {err}") + }) +} diff --git a/pegainfer-kernels/src/ops/k3/flash_mla_prefill.rs b/pegainfer-kernels/src/ops/k3/flash_mla_prefill.rs index 3b97cc88a..16b351540 100644 --- a/pegainfer-kernels/src/ops/k3/flash_mla_prefill.rs +++ b/pegainfer-kernels/src/ops/k3/flash_mla_prefill.rs @@ -138,16 +138,54 @@ pub fn k3_mla_prefill_lse_merge_launch( lse_acc: &mut CudaSlice, reset: bool, ) -> Result<()> { - ensure!(t_q > 0, "K3 MLA prefill LSE merge got an empty span"); ensure!( - o_win.len() >= t_q * heads * K3_MLA_PREFILL_V - && lse_win.len() >= heads * t_q - && o_acc.len() >= t_q * heads * K3_MLA_PREFILL_V - && lse_acc.len() >= heads * t_q, - "K3 MLA prefill LSE merge buffers too small for t_q={t_q}, heads={heads}" + o_win.len() >= t_q * heads * K3_MLA_PREFILL_V && lse_win.len() >= heads * t_q, + "K3 MLA prefill LSE merge window too small for t_q={t_q}, heads={heads}" ); let (ow_ptr, _ow_guard) = o_win.device_ptr(&ctx.stream); let (lw_ptr, _lw_guard) = lse_win.device_ptr(&ctx.stream); + lse_merge_raw(ctx, t_q, heads, ow_ptr, lw_ptr, o_acc, lse_acc, reset) +} + +/// [`k3_mla_prefill_lse_merge_launch`] over a window a *peer* rank computed +/// and published: `o_win`/`lse_win` are raw device pointers into the peer's +/// publish slab, valid inside the exchange window that consumes them (the +/// caller's stream already waits on the peer's publish). The kernel reads +/// them once over the fabric; nothing is staged locally. +#[allow(clippy::too_many_arguments)] +pub fn k3_mla_prefill_lse_merge_peer_launch( + ctx: &DeviceContext, + t_q: usize, + heads: usize, + o_win: u64, + lse_win: u64, + o_acc: &mut CudaSlice, + lse_acc: &mut CudaSlice, + reset: bool, +) -> Result<()> { + ensure!( + o_win != 0 && lse_win != 0, + "K3 MLA prefill LSE merge got a null peer window" + ); + lse_merge_raw(ctx, t_q, heads, o_win, lse_win, o_acc, lse_acc, reset) +} + +#[allow(clippy::too_many_arguments)] +fn lse_merge_raw( + ctx: &DeviceContext, + t_q: usize, + heads: usize, + ow_ptr: u64, + lw_ptr: u64, + o_acc: &mut CudaSlice, + lse_acc: &mut CudaSlice, + reset: bool, +) -> Result<()> { + ensure!(t_q > 0, "K3 MLA prefill LSE merge got an empty span"); + ensure!( + o_acc.len() >= t_q * heads * K3_MLA_PREFILL_V && lse_acc.len() >= heads * t_q, + "K3 MLA prefill LSE merge accumulators too small for t_q={t_q}, heads={heads}" + ); let (oa_ptr, _oa_guard) = o_acc.device_ptr_mut(&ctx.stream); let (la_ptr, _la_guard) = lse_acc.device_ptr_mut(&ctx.stream); let rc = unsafe { @@ -232,19 +270,48 @@ fn k3_flash_mla_prefill_fmha( out: &mut CudaSlice, lse_out: Option<&mut CudaSlice>, scale: f32, +) -> Result<()> { + ensure!( + q.len() >= t_q * heads * K3_MLA_PREFILL_QK, + "K3 FlashMLA prefill query buffer too small for t_q={t_q}, heads={heads}" + ); + let (q_ptr, _q_guard) = q.device_ptr(&ctx.stream); + fmha_raw_q( + ctx, entry, kind, t_q, t_kv, heads, q_ptr, k, nope_v, out, lse_out, scale, + ) +} + +/// The FMHA core over a raw query pointer: `q_ptr` addresses `[t_q, heads, +/// 192]` bf16 rows that are either this rank's own buffer or a peer's +/// published one inside an exchange window (the stream already waits on the +/// publish). The kernel loads each query tile once, so a fabric-resident Q +/// costs its bytes exactly once. +#[allow(clippy::too_many_arguments)] +fn fmha_raw_q( + ctx: &DeviceContext, + entry: FmhaEntry, + kind: &str, + t_q: usize, + t_kv: usize, + heads: usize, + q_ptr: u64, + k: &CudaSlice, + nope_v: &CudaSlice, + out: &mut CudaSlice, + lse_out: Option<&mut CudaSlice>, + scale: f32, ) -> Result<()> { ensure!( t_q > 0 && t_kv > 0, "K3 FlashMLA prefill got an empty span: t_q={t_q}, t_kv={t_kv}" ); + ensure!(q_ptr != 0, "K3 FlashMLA prefill got a null query pointer"); ensure!( - q.len() >= t_q * heads * K3_MLA_PREFILL_QK - && k.len() >= t_kv * heads * K3_MLA_PREFILL_QK + k.len() >= t_kv * heads * K3_MLA_PREFILL_QK && nope_v.len() >= t_kv * heads * K3_MLA_PREFILL_NV && out.len() >= t_q * heads * K3_MLA_PREFILL_V, "K3 FlashMLA prefill buffers too small for t_q={t_q}, t_kv={t_kv}, heads={heads}" ); - let (q_ptr, _q_guard) = q.device_ptr(&ctx.stream); let (k_ptr, _k_guard) = k.device_ptr(&ctx.stream); let (nv_ptr, _nv_guard) = nope_v.device_ptr(&ctx.stream); let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); @@ -358,6 +425,41 @@ pub fn k3_flash_mla_prefill_fwd_dense_launch( ) } +/// [`k3_flash_mla_prefill_fwd_dense_launch`] with the queries read from a +/// *peer* rank's published `[t_q, heads, 192]` rows at raw device pointer +/// `q_peer` — the context-parallel stripe: the rank that holds a key segment +/// attends another rank's queries over it and hands back output + LSE for +/// the owner's merge. Valid only inside the exchange window that published +/// the queries. +#[allow(clippy::too_many_arguments)] +pub fn k3_flash_mla_prefill_fwd_dense_peer_q_launch( + ctx: &DeviceContext, + t_q: usize, + t_kv: usize, + heads: usize, + q_peer: u64, + k: &CudaSlice, + nope_v: &CudaSlice, + out: &mut CudaSlice, + lse_out: Option<&mut CudaSlice>, + scale: f32, +) -> Result<()> { + fmha_raw_q( + ctx, + ffi::k3_flash_mla_prefill_fwd_dense, + "dense (peer queries)", + t_q, + t_kv, + heads, + q_peer, + k, + nope_v, + out, + lse_out, + scale, + ) +} + #[cfg(test)] mod tests { use super::*; diff --git a/pegainfer-kernels/src/ops/k3/land.rs b/pegainfer-kernels/src/ops/k3/land.rs new file mode 100644 index 000000000..8e62c9a85 --- /dev/null +++ b/pegainfer-kernels/src/ops/k3/land.rs @@ -0,0 +1,70 @@ +//! Kimi-K3 matmul landing: merge one column span of a `[split_k, nt]` f32 +//! partial and land bf16 once. +//! +//! Hand-written replacement for the retired TileLang `land_batched` kernel +//! (one element per thread, 2-byte stores, 2.3 TB/s on the chunked-prefill +//! landings). Eight columns per thread with 16-byte loads and stores, the +//! retired kernel's exact arithmetic — ascending-segment f32 sum, one +//! round-to-nearest-even cast — so the landing is bit-identical; see +//! `csrc/k3/k3_land.cu`. Batch is a plain launch dimension. + +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::tensor::DeviceContext; + +/// Merge the column span `[off, off + n)` of each row's `[split_k, nt]` +/// partials and land bf16 once — the landing of every matmul in the certified +/// spelling. +#[allow(clippy::too_many_arguments)] +pub fn k3_land_batched_launch( + ctx: &DeviceContext, + b: usize, + nt: usize, + n: usize, + off: usize, + split_k: usize, + p: &CudaSlice, + o: &mut CudaSlice, +) -> Result<()> { + ensure!( + b > 0 && split_k > 0, + "K3 land needs rows and at least one segment" + ); + ensure!( + off + n <= nt, + "K3 land span [{off}, {off}+{n}) does not fit the partial width {nt}" + ); + ensure!( + p.len() >= b * split_k * nt && o.len() >= b * n, + "K3 land buffers too small for b={b}, nt={nt}, n={n}, split_k={split_k}: p {}, o {}", + p.len(), + o.len() + ); + let (p_ptr, _p_guard) = p.device_ptr(&ctx.stream); + let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); + unsafe { + ffi::k3_land_cuda( + p_ptr as *const f32, + o_ptr as *mut c_void, + i32::try_from(b)?, + i32::try_from(nt)?, + i32::try_from(n)?, + i32::try_from(off)?, + i32::try_from(split_k)?, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| { + anyhow!("K3 land (B={b}, NT={nt}, N={n}, OFF={off}, SK={split_k}) launch failed: {err}") + }) +} diff --git a/pegainfer-kernels/src/ops/k3/mod.rs b/pegainfer-kernels/src/ops/k3/mod.rs index 33804b6e0..34548a8c3 100644 --- a/pegainfer-kernels/src/ops/k3/mod.rs +++ b/pegainfer-kernels/src/ops/k3/mod.rs @@ -1,16 +1,22 @@ //! Kimi-K3 GPU operators. +mod conv_silu_chunk; mod deepgemm; +mod elementwise; mod flash_kda; mod flash_mla_prefill; +mod land; mod mega_moe; mod mla_paged; mod moe_chain; mod router_topk; +pub use conv_silu_chunk::*; pub use deepgemm::*; +pub use elementwise::*; pub use flash_kda::*; pub use flash_mla_prefill::*; +pub use land::*; pub use mega_moe::*; pub use mla_paged::*; pub use moe_chain::*; diff --git a/pegainfer-kernels/src/ops/k3_tilelang.rs b/pegainfer-kernels/src/ops/k3_tilelang.rs index ecd71473d..fe98aba0e 100644 --- a/pegainfer-kernels/src/ops/k3_tilelang.rs +++ b/pegainfer-kernels/src/ops/k3_tilelang.rs @@ -163,52 +163,7 @@ pub fn k3_rms_norm_rbs_batched_launch( check(rc, &format!("K3 rms_norm_rbs_batched (B={b}, H={h})")) } -/// Merge the column span `[off, off + n)` of each row's `[split_k, nt]` -/// partials and land bf16 once — the landing of every matmul in the certified -/// spelling. -#[allow(clippy::too_many_arguments)] -pub fn k3_land_batched_launch( - ctx: &DeviceContext, - b: usize, - nt: usize, - n: usize, - off: usize, - split_k: usize, - p: &CudaSlice, - o: &mut CudaSlice, -) -> Result<()> { - check_bucket(b)?; - ensure!( - off + n <= nt, - "K3 land span [{off}, {off}+{n}) does not fit the partial width {nt}" - ); - ensure!( - p.len() >= b * split_k * nt && o.len() >= b * n, - "K3 land buffers too small for b={b}, nt={nt}, n={n}, split_k={split_k}: p {}, o {}", - p.len(), - o.len() - ); - let (p_ptr, _p_guard) = p.device_ptr(&ctx.stream); - let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); - let rc = unsafe { - ffi::k3_land_batched( - p_ptr as *const f32, - o_ptr as *mut c_void, - b as i32, - nt as i32, - n as i32, - off as i32, - split_k as i32, - ctx.stream.cu_stream(), - ) - }; - check( - rc, - &format!("K3 land_batched (B={b}, NT={nt}, N={n}, OFF={off}, SK={split_k})"), - ) -} - -/// [`k3_land_batched_launch`] fused with the round-before-scale norm — MLA's +/// The matmul landing (`k3_land_batched_launch`) fused with the round-before-scale norm — MLA's /// `q_a`, the one place the engine fuses a merge and a norm. #[allow(clippy::too_many_arguments)] pub fn k3_land_rms_norm_rbs_batched_launch( @@ -256,119 +211,6 @@ pub fn k3_land_rms_norm_rbs_batched_launch( ) } -/// The two `(A, Bt) -> O` elementwise kernels share a signature; only the -/// launcher symbol and the error label differ. -type BinaryLauncher = unsafe extern "C" fn( - *const c_void, - *const c_void, - *mut c_void, - i32, - i32, - cudarc::driver::sys::CUstream, -) -> i32; - -fn k3_binary_batched_launch( - ctx: &DeviceContext, - launcher: BinaryLauncher, - what: &str, - b: usize, - n: usize, - a: &CudaSlice, - bt: &CudaSlice, - o: &mut CudaSlice, -) -> Result<()> { - check_bucket(b)?; - ensure!( - a.len() >= b * n && bt.len() >= b * n && o.len() >= b * n, - "K3 {what} buffers too small for b={b}, n={n}: a {}, bt {}, o {}", - a.len(), - bt.len(), - o.len() - ); - let (a_ptr, _a_guard) = a.device_ptr(&ctx.stream); - let (bt_ptr, _bt_guard) = bt.device_ptr(&ctx.stream); - let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); - let rc = unsafe { - launcher( - a_ptr as *const c_void, - bt_ptr as *const c_void, - o_ptr as *mut c_void, - b as i32, - n as i32, - ctx.stream.cu_stream(), - ) - }; - check(rc, &format!("K3 {what} (B={b}, N={n})")) -} - -/// `o = a + bt`, added in bf16 — the residual adds, and routed + shared. -pub fn k3_add2_batched_launch( - ctx: &DeviceContext, - b: usize, - n: usize, - a: &CudaSlice, - bt: &CudaSlice, - o: &mut CudaSlice, -) -> Result<()> { - k3_binary_batched_launch(ctx, ffi::k3_add2_batched, "add2_batched", b, n, a, bt, o) -} - -/// `o = a * bf16(sigmoid(bt))`, the MLA sigmoid output gate. The sigmoid is -/// taken in f32 and lands in bf16 before the product. -pub fn k3_mul_sigmoid_batched_launch( - ctx: &DeviceContext, - b: usize, - n: usize, - a: &CudaSlice, - bt: &CudaSlice, - o: &mut CudaSlice, -) -> Result<()> { - k3_binary_batched_launch( - ctx, - ffi::k3_mul_sigmoid_batched, - "mul_sigmoid_batched", - b, - n, - a, - bt, - o, - ) -} - -/// K3's situ activation: `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)` in f32, -/// landing bf16 once. The two betas are compiled in. -pub fn k3_situ_batched_launch( - ctx: &DeviceContext, - b: usize, - n: usize, - g: &CudaSlice, - u: &CudaSlice, - o: &mut CudaSlice, -) -> Result<()> { - check_bucket(b)?; - ensure!( - g.len() >= b * n && u.len() >= b * n && o.len() >= b * n, - "K3 situ buffers too small for b={b}, n={n}: g {}, u {}, o {}", - g.len(), - u.len(), - o.len() - ); - let (g_ptr, _g_guard) = g.device_ptr(&ctx.stream); - let (u_ptr, _u_guard) = u.device_ptr(&ctx.stream); - let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); - let rc = unsafe { - ffi::k3_situ_batched( - g_ptr as *const c_void, - u_ptr as *const c_void, - o_ptr as *mut c_void, - b as i32, - n as i32, - ctx.stream.cu_stream(), - ) - }; - check(rc, &format!("K3 situ_batched (B={b}, N={n})")) -} - /// KDA short convolution plus silu, one token per row. /// /// `p` is the projection's f32 partial; its bf16 landing is written to `x` and @@ -539,53 +381,6 @@ pub fn k3_kda_core_batched_launch( ) } -/// `kda_core`'s tail on its own: per (row, head) the f32 rms_norm of the bf16 -/// attention landing `x` times the o_norm gamma `go [head_dim]`, landed once, -/// times the bf16 sigmoid of the output gate `g2` — word-for-word the batched -/// kernel's last loop. Chunked prefill computes the attention through FlashKDA -/// and finishes its rows here. -pub fn k3_o_norm_gate_batched_launch( - ctx: &DeviceContext, - b: usize, - num_heads: usize, - head_dim: usize, - x: &CudaSlice, - g2: &CudaSlice, - go: &CudaSlice, - out: &mut CudaSlice, -) -> Result<()> { - check_bucket(b)?; - let kp = num_heads * head_dim; - ensure!( - x.len() >= b * kp && g2.len() >= b * kp && out.len() >= b * kp && go.len() >= head_dim, - "K3 o_norm_gate buffers too small for b={b}, kp={kp}: x {}, g2 {}, go {}, out {}", - x.len(), - g2.len(), - go.len(), - out.len() - ); - let (x_ptr, _x_guard) = x.device_ptr(&ctx.stream); - let (g2_ptr, _g2_guard) = g2.device_ptr(&ctx.stream); - let (go_ptr, _go_guard) = go.device_ptr(&ctx.stream); - let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); - let rc = unsafe { - ffi::k3_o_norm_gate_batched( - x_ptr as *const c_void, - g2_ptr as *const c_void, - go_ptr as *const f32, - out_ptr as *mut c_void, - b as i32, - num_heads as i32, - head_dim as i32, - ctx.stream.cu_stream(), - ) - }; - check( - rc, - &format!("K3 o_norm_gate_batched (B={b}, KH={num_heads}, KD={head_dim})"), - ) -} - /// Score the `blocks + 1` attention-residual candidates of every row: a /// weightless RMS normalization then a dot with the fused f32 scoring vector. #[allow(clippy::too_many_arguments)]