diff --git a/docs/technical.md b/docs/technical.md index 9c2200d..89db761 100644 --- a/docs/technical.md +++ b/docs/technical.md @@ -94,12 +94,12 @@ Each of these is a seam, not a rewrite: | knob | what it changes | where | |---|---|---| -| `ShardPolicy` | `ExpertsOnly` (experts banded, attention replicated) vs `Full` | [`kimi_k3.h`](runtime/include/sparkinfer/models/kimi_k3.h) | -| `SPARKINFER_TP_BACKEND` | `nccl` \| `peer` \| `multimem` — all three validated on 8× H200 | [`collective.h`](runtime/include/sparkinfer/tp/collective.h) | -| shard rules | per-tensor Row / Col / Expert / Replicate, 230 CPU tests | [`weight_plan.cpp`](runtime/src/tp/weight_plan.cpp) | -| expert bands | contiguous today; strided would trade load balance for locality | [`shard.cpp`](runtime/src/tp/shard.cpp) | -| reduce points | `K3LayerPhase` — move a collective by moving one call | [`kimi_k3_tp.cpp`](runtime/src/models/kimi_k3_tp.cpp) | -| dtype | f32 today (K3's residual stream is f32 by design); bf16 path exists | [`collective.h`](runtime/include/sparkinfer/tp/collective.h) | +| `ShardPolicy` | `ExpertsOnly` (experts banded, attention replicated) vs `Full` | [`kimi_k3.h`](../runtime/include/sparkinfer/models/kimi_k3.h) | +| `SPARKINFER_TP_BACKEND` | `nccl` \| `peer` \| `multimem` — all three validated on 8× H200 | [`collective.h`](../runtime/include/sparkinfer/tp/collective.h) | +| shard rules | per-tensor Row / Col / Expert / Replicate, 230 CPU tests | [`weight_plan.cpp`](../runtime/src/tp/weight_plan.cpp) | +| expert bands | contiguous today; strided would trade load balance for locality | [`shard.cpp`](../runtime/src/tp/shard.cpp) | +| reduce points | `K3LayerPhase` — move a collective by moving one call | [`kimi_k3_tp.cpp`](../runtime/src/models/kimi_k3_tp.cpp) | +| dtype | f32 today (K3's residual stream is f32 by design); bf16 path exists | [`collective.h`](../runtime/include/sparkinfer/tp/collective.h) | The shard math is **CUDA-free and unit-tested without a GPU** — 4972 checks on `shard.cpp`, 230 on `weight_plan.cpp`, 44 on backend selection. TP bugs do not live in @@ -131,7 +131,7 @@ would shard weights the executor still indexes at full width. That is the next l | Extras | Cross-layer residual attention, `block_size 12` | | Vision | MoonViT-3d — 27 layers, 1024 wide, **non-square fused QKV** (1536 ≠ n_embd), patch 14 | -Three traps that produce silently wrong output rather than an error — all encoded in [`bench/configs/models/kimi_k3.yaml`](bench/configs/models/kimi_k3.yaml): +Three traps that produce silently wrong output rather than an error — all encoded in [`bench/configs/models/kimi_k3.yaml`](../bench/configs/models/kimi_k3.yaml): - **`full_attn_layers` is 1-indexed.** The converter tests `(il + 1) in full_attn_layers`. Off by one and you get garbage, not a crash. - **MLA is stored as MQA.** `head_count_kv = 1`, `key_length = kv_lora + qk_rope = 576`; per-layer `head_count_kv == 0` is what marks a KDA layer. @@ -143,7 +143,7 @@ Three traps that produce silently wrong output rather than an error — all enco Every other model in the SparkInfer family is benchmarked against `ggml-org/llama.cpp` at a pinned commit. Kimi K3 cannot be. -**Upstream llama.cpp cannot load this model at all.** It asserts `n_expert <= LLAMA_MAX_EXPERTS`, and upstream's cap is 512. K3 has 896. There is no upstream number to compare against, so the reference is [`unslothai/llama.cpp`](https://github.com/unslothai/llama.cpp) PR #48, pinned in [`bench/scripts/reference.lock`](bench/scripts/reference.lock). +**Upstream llama.cpp cannot load this model at all.** It asserts `n_expert <= LLAMA_MAX_EXPERTS`, and upstream's cap is 512. K3 has 896. There is no upstream number to compare against, so the reference is [`unslothai/llama.cpp`](https://github.com/unslothai/llama.cpp) PR #48, pinned in [`bench/scripts/reference.lock`](../bench/scripts/reference.lock). Four things in that fork are load-bearing, not cosmetic: @@ -178,7 +178,7 @@ UD-Q2_K_XL is the accuracy knee and remains what the project is ultimately judge But the *default* is what runs when nobody passes a flag, and defaulting to an 802 GiB download that is on no machine means every fresh invocation dies before it does anything. UD-IQ1_S is what is resident, what the llama.cpp reference was measured on, -and what [`bench/refdata/`](bench/refdata)'s reference logits were captured against. +and what [`bench/refdata/`](../bench/refdata)'s reference logits were captured against. `PRIMARY_QUANT=UD-Q2_K_XL` switches to the target. The reference.lock slots carry the quant in their **name** @@ -221,7 +221,7 @@ against it too. One command produces both plus the tier: bench/scripts/kimi_k3_eval.sh --node h200x8 --frontier ``` -It emits the `RESULT_JSON` contract [`bench/scripts/label.py`](bench/scripts/label.py) +It emits the `RESULT_JSON` contract [`bench/scripts/label.py`](../bench/scripts/label.py) already scores for the other models, so K3 needs no second scoring path: - **Correctness gate first.** top-1 ≥ 0.95 and KL ≤ 0.05 against the captured reference, @@ -235,7 +235,7 @@ already scores for the other models, so K3 needs no second scoring path: 2.2×) the frontier binds, so `XL` costs a real 18% over main. A node run posts its verdict to a PR with `/eval RESULT_JSON {...}`; -[`.github/workflows/eval-label.yml`](.github/workflows/eval-label.yml) **re-derives** the +[`.github/workflows/eval-label.yml`](../.github/workflows/eval-label.yml) **re-derives** the tier from the reported measurements rather than trusting the reported label, and honours the command only from maintainers. @@ -245,16 +245,16 @@ the command only from maintainers. | Path | What | |---|---| -| [`bench/`](bench) | **the baseline** — K3 harness, arch/target configs, eval + accuracy scripts | -| [`docs/`](docs) | [`kimi-k3-baseline.md`](docs/kimi-k3-baseline.md) — how to run it, and every trap in the arch | -| [`kernels/`](kernels) | CUDA kernels — flash-decode, decode GEMV, fused MoE FFN, GEMM, RMSNorm, RoPE, GGUF dequant | -| [`runtime/`](runtime) | scheduler, paged KV cache, CUDA-graph decode, native GGUF loading, model forward | -| [`moe/`](moe) | sync-free MoE router + expert dispatch | -| [`server/`](server) | OpenAI-compatible HTTP API (`BUILD_SERVER=ON`) | +| [`bench/`](../bench) | **the baseline** — K3 harness, arch/target configs, eval + accuracy scripts | +| [`docs/`](../docs) | [`kimi-k3-baseline.md`](kimi-k3-baseline.md) — how to run it, and every trap in the arch | +| [`kernels/`](../kernels) | CUDA kernels — flash-decode, decode GEMV, fused MoE FFN, GEMM, RMSNorm, RoPE, GGUF dequant | +| [`runtime/`](../runtime) | scheduler, paged KV cache, CUDA-graph decode, native GGUF loading, model forward | +| [`moe/`](../moe) | sync-free MoE router + expert dispatch | +| [`server/`](../server) | OpenAI-compatible HTTP API (`BUILD_SERVER=ON`) | `kernels/` and `runtime/` now carry a native K3 path (KDA + MLA decode, latent MoE, `situ`, cross-layer residual, expert-parallel dispatch). `moe/` and `server/` are still Qwen-shaped. -**Scoring is speedup-only.** SN74 pays verified marginal speedups labeled **XL / L / M / S / XS**. Sub-2% gains are never aggregated across contexts. See [`.gittensor/weights.json`](.gittensor/weights.json). +**Scoring is speedup-only.** SN74 pays verified marginal speedups labeled **XL / L / M / S / XS**. Sub-2% gains are never aggregated across contexts. See [`.gittensor/weights.json`](../.gittensor/weights.json). --- @@ -308,4 +308,4 @@ Two reference-server flags are mandatory, not tuning: - `--no-context-shift` — K3 is a hybrid recurrent arch; llama.cpp cannot context-shift or restore slots for it, and a long eval dies mid-run without it. - `--no-jinja` — the gate posts raw token ids; a chat template would prepend tokens the candidate never saw. -Details: [`eval/`](eval) · **[EVAL-TRUST.md](EVAL-TRUST.md)** (Polaris TDX receipts, reproducible from source today). +Details: [`eval/`](../eval) · **[EVAL-TRUST.md](../EVAL-TRUST.md)** (Polaris TDX receipts, reproducible from source today). diff --git a/kernels/csrc/cuda/kimi_k3/k3_epilogue_q8.cu b/kernels/csrc/cuda/kimi_k3/k3_epilogue_q8.cu new file mode 100644 index 0000000..d075f9c --- /dev/null +++ b/kernels/csrc/cuda/kimi_k3/k3_epilogue_q8.cu @@ -0,0 +1,225 @@ +// Factor — fold a one-op producer into the Q8_0 quantise its consumer would issue. +// +// =========================================================================== +// THE SHAPE +// =========================================================================== +// The decode graph's most common node is still quantize_q8_0_warp_kernel (557 of +// ~3,449 nodes per rank per token on a recent census). hoist_act already removes the +// re-quantises of the shared attn/ffn norms; what remains are activations that exist +// only to feed ONE projection a few microseconds later: +// +// kda_gate_out -> attn_output 69 layers +// mla_gate_out -> attn_output 24 layers +// situ -> shexp/dense down 1+92 layers +// +// Each pair is an elementwise (or per-head RMS) launch followed by a quantise of the +// vector that launch just wrote. This does both in one kernel and hands the Q8_0 +// straight to k3_proj_q8act_f32, so the standalone producer and the standalone +// quantise both leave the graph. +// +// =========================================================================== +// WHY IT IS BIT-IDENTICAL +// =========================================================================== +// * situ and mla_gate are pure elementwise; the fused path evaluates the same +// expression the standalone kernels use, then runs the shipped warp quantiser +// (amax over magnitudes via shfl_xor, d = amax/127, per-element rn) on those +// values. Order-independent amax, independent qs — same bytes as +// situ/mla_gate followed by k3_quantize_q8_0. +// * kda_gate keeps the frozen 128-lane block_sum over head_dim that +// kda_gate_out_kernel<128> uses; only the apply is then quantised. Idle lanes +// above head_dim are unreachable at K3's head_dim == 128. +// +// The float buffers (gate_out / dense_situ / mla_attn_out) are written only when +// the caller passes them — the forward passes nullptr on the hot path and keeps +// the mirror solely for debug tags. Q8 bytes stay bit-identical either way. +// +// =========================================================================== +// WHAT IT DECLINES +// =========================================================================== +// SPARKINFER_K3_EPILOGUE_Q8=0, a missing pointer, n not a multiple of 32, or +// (for the KDA gate) head_dim != 128. The caller keeps the standalone pair. + +#include "sparkinfer/kernels/kimi_k3_fast.h" +#include "k3_pdl.cuh" + +#include +#include + +#include +#include + +namespace sparkinfer { +namespace kernels { +namespace k3 { +namespace { + +struct BlockQ8_0 { + uint16_t d; + int8_t qs[32]; +}; +static_assert(sizeof(BlockQ8_0) == 34, "bad block_q8_0 layout"); + +__device__ __forceinline__ float sigmoidf_ep(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +// Same tree as k3_kernels.cu — duplicated so this TU stays independent. +template +__device__ __forceinline__ float block_sum(float v, float* shm) { + for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + if (lane == 0) shm[warp] = v; + __syncthreads(); + constexpr int NWARP = BLOCK / 32; + if (threadIdx.x == 0) { + float s = 0.0f; + for (int w = 0; w < NWARP; ++w) s += shm[w]; + shm[NWARP] = s; + } + __syncthreads(); + return shm[NWARP]; +} + +__device__ __forceinline__ void quant_warp_lane(BlockQ8_0* __restrict__ out, + int b, int lane, float v) { + float amax = fabsf(v); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + const float d = amax / 127.0f; + const float id = amax != 0.0f ? 127.0f / amax : 0.0f; + if (lane == 0) out[b].d = __half_as_ushort(__float2half_rn(d)); + out[b].qs[lane] = (int8_t)__float2int_rn(v * id); +} + +__device__ __forceinline__ float situ_one(float g, float u, float beta, float inv_beta, + float lb, float inv_lb, int lb_active) { + const float a = beta * tanhf(g * inv_beta) * sigmoidf_ep(g); + const float ub = lb_active ? (lb * tanhf(u * inv_lb)) : u; + return a * ub; +} + +// One warp per 32-value quant block. Each lane evaluates situ for its element, then +// the warp emits one BlockQ8_0 — same shape as quantize_q8_0_warp_kernel. +__global__ void situ_q8_warp_kernel(BlockQ8_0* __restrict__ q8, + float* situ_out, + const float* __restrict__ gate, + const float* __restrict__ up, + int n_blocks, float beta, float inv_beta, + float lb, float inv_lb, int lb_active) { + k3_pdl_sync(); + const int b = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + if (b >= n_blocks) return; + const int lane = threadIdx.x & 31; + const int i = b * 32 + lane; + const float v = situ_one(gate[i], up[i], beta, inv_beta, lb, inv_lb, lb_active); + if (situ_out) situ_out[i] = v; + quant_warp_lane(q8, b, lane, v); +} + +__global__ void mla_gate_q8_warp_kernel(BlockQ8_0* __restrict__ q8, + float* gated_out, + const float* __restrict__ attn_out, + const float* __restrict__ gate_proj, + int n_blocks) { + k3_pdl_sync(); + const int b = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + if (b >= n_blocks) return; + const int lane = threadIdx.x & 31; + const int i = b * 32 + lane; + const float v = attn_out[i] * sigmoidf_ep(gate_proj[i]); + if (gated_out) gated_out[i] = v; + quant_warp_lane(q8, b, lane, v); +} + +// One block per head, 128 threads. RMS matches kda_gate_out_kernel<128>; each of the +// four warps then emits one Q8_0 block covering its 32 lanes of the head. +template +__global__ void kda_gate_q8_kernel(BlockQ8_0* __restrict__ q8, + float* gate_out, + const float* __restrict__ o, + const float* __restrict__ norm_w, + const float* __restrict__ g2, + int head_dim, float eps) { + k3_pdl_sync(); + const int h = blockIdx.x; + const float* oh = o + (size_t)h * head_dim; + const float* gh = g2 + (size_t)h * head_dim; + float* dst = gate_out ? gate_out + (size_t)h * head_dim : nullptr; + + __shared__ float shm[BLOCK / 32 + 1]; + + float acc = 0.0f; + for (int d = threadIdx.x; d < head_dim; d += BLOCK) { + const float x = oh[d]; + acc += x * x; + } + const float ss = block_sum(acc, shm); + const float inv = rsqrtf(ss / (float)head_dim + eps); + + // head_dim == BLOCK == 128: one element per thread, four warps = four Q8 blocks. + const int d = threadIdx.x; + float v = 0.0f; + if (d < head_dim) + v = (oh[d] * inv * norm_w[d]) * sigmoidf_ep(gh[d]); + if (dst && d < head_dim) dst[d] = v; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int b = h * (head_dim / 32) + warp; + if (d < head_dim) quant_warp_lane(q8, b, lane, v); +} + +bool epilogue_q8_on() { + static const bool on = [] { + const char* e = std::getenv("SPARKINFER_K3_EPILOGUE_Q8"); + return !(e && e[0] == '0'); + }(); + return on; +} + +} // namespace + +bool k3_situ_q8(void* q8_out, float* situ_out, const float* gate, const float* up, + int64_t n, float beta, float linear_beta, cudaStream_t stream) { + if (!epilogue_q8_on()) return false; + if (!q8_out || !gate || !up || n <= 0 || n % 32 != 0) return false; + const int n_blocks = (int)(n / 32); + const int lb_active = linear_beta > 0.0f ? 1 : 0; + const float inv_beta = 1.0f / beta; + const float inv_lb = lb_active ? 1.0f / linear_beta : 1.0f; + constexpr int threads = 256; + const int blocks = (n_blocks * 32 + threads - 1) / threads; + k3_pdl_launch(dim3(blocks), dim3(threads), 0, stream, situ_q8_warp_kernel, + (BlockQ8_0*)q8_out, situ_out, gate, up, n_blocks, beta, inv_beta, + linear_beta, inv_lb, lb_active); + return true; +} + +bool k3_mla_gate_q8(void* q8_out, float* gated_out, const float* attn_out, + const float* gate_proj, int64_t n, cudaStream_t stream) { + if (!epilogue_q8_on()) return false; + if (!q8_out || !attn_out || !gate_proj || n <= 0 || n % 32 != 0) return false; + const int n_blocks = (int)(n / 32); + constexpr int threads = 256; + const int blocks = (n_blocks * 32 + threads - 1) / threads; + k3_pdl_launch(dim3(blocks), dim3(threads), 0, stream, mla_gate_q8_warp_kernel, + (BlockQ8_0*)q8_out, gated_out, attn_out, gate_proj, n_blocks); + return true; +} + +bool k3_kda_gate_q8(void* q8_out, float* gate_out, const float* o, + const float* norm_w, const float* g2, int head_dim, int n_head, + float eps, cudaStream_t stream) { + if (!epilogue_q8_on()) return false; + if (!q8_out || !o || !norm_w || !g2 || head_dim != 128 || n_head <= 0) + return false; + k3_pdl_launch((unsigned)n_head, 128, 0, stream, kda_gate_q8_kernel<128>, + (BlockQ8_0*)q8_out, gate_out, o, norm_w, g2, head_dim, eps); + return true; +} + +} // namespace k3 +} // namespace kernels +} // namespace sparkinfer diff --git a/kernels/include/sparkinfer/kernels/kimi_k3_fast.h b/kernels/include/sparkinfer/kernels/kimi_k3_fast.h index 607bae9..576d2db 100644 --- a/kernels/include/sparkinfer/kernels/kimi_k3_fast.h +++ b/kernels/include/sparkinfer/kernels/kimi_k3_fast.h @@ -179,6 +179,26 @@ bool k3_add3_rows_f32(float* out, float* out_ab, const float* a, const float* b, int64_t b_row_stride, const float* c, int rows, int cols, cudaStream_t stream); +// --------------------------------------------------------------------------- +// Factor — producer + Q8_0 quantise, in one launch +// --------------------------------------------------------------------------- +// Three call sites where an elementwise (or per-head RMS) producer exists only to +// feed a single Q8_0 projection a few microseconds later. Each fold removes the +// producer launch and the standalone quantise; the caller then runs +// k3_proj_q8act_f32 on the emitted scratch. +// +// Bit-identical to the split path (see k3_epilogue_q8.cu). Declines on +// SPARKINFER_K3_EPILOGUE_Q8=0 or a shape the warp quantiser cannot cover. +// Optional float outs (situ_out / gated_out / gate_out) mirror what the standalone +// producer wrote; pass nullptr to skip the store. +bool k3_situ_q8(void* q8_out, float* situ_out, const float* gate, const float* up, + int64_t n, float beta, float linear_beta, cudaStream_t stream); +bool k3_mla_gate_q8(void* q8_out, float* gated_out, const float* attn_out, + const float* gate_proj, int64_t n, cudaStream_t stream); +bool k3_kda_gate_q8(void* q8_out, float* gate_out, const float* o, + const float* norm_w, const float* g2, int head_dim, int n_head, + float eps, cudaStream_t stream); + } // namespace k3 } // namespace kernels } // namespace sparkinfer diff --git a/kernels/tests/CMakeLists.txt b/kernels/tests/CMakeLists.txt index b6aee28..d9fd7b1 100644 --- a/kernels/tests/CMakeLists.txt +++ b/kernels/tests/CMakeLists.txt @@ -103,7 +103,7 @@ add_test(NAME k3_prefill_order_cpu_test COMMAND k3_prefill_order_cpu_test) foreach(k3t k3_kda_step_gpu_test k3_mla_absorb_gpu_test k3_proj_1bar_gpu_test k3_moe_iq1s_mma_gpu_test k3_moe_batched_iq1s_gpu_test k3_kda_chunk_prefill_gpu_test - k3_proj_q8_mma_gpu_test) + k3_proj_q8_mma_gpu_test k3_epilogue_q8_gpu_test) add_executable(${k3t} ${k3t}.cu) target_link_libraries(${k3t} PRIVATE si_k3) target_include_directories(${k3t} PRIVATE ${PROJECT_SOURCE_DIR}/include) diff --git a/kernels/tests/k3_epilogue_q8_gpu_test.cu b/kernels/tests/k3_epilogue_q8_gpu_test.cu new file mode 100644 index 0000000..1155203 --- /dev/null +++ b/kernels/tests/k3_epilogue_q8_gpu_test.cu @@ -0,0 +1,202 @@ +// Factor — producer + Q8_0 in one launch, at K3's real dims. +// Compares the fused Q8 bytes AND the optional float mirror against the +// standalone producer followed by k3_quantize_q8_0 — bitwise, not by tolerance. +#include "sparkinfer/kernels/kimi_k3.h" +#include "sparkinfer/kernels/kimi_k3_fast.h" +#include "k3_fast_test_util.h" + +#include + +using namespace sparkinfer::kernels::k3; +using namespace k3test; +#define CU K3T_CU + +static bool bytes_eq(const void* a, const void* b, size_t n) { + return std::memcmp(a, b, n) == 0; +} + +static void test_situ_q8() { + const int n = 768; // 2-D MoE shexp band at tp=8 + std::printf("situ + Q8_0 (shexp-band width)\n"); + std::mt19937 rng(7); + auto gate = rnd(n, rng); + auto up = rnd(n, rng, -2.f, 2.f); + const float beta = 4.0f, lb = 25.0f; + + float* dg = to_dev(gate); + float* du = to_dev(up); + float* dsitu = nullptr; + float* dsitu2 = nullptr; + void* q8_fused = nullptr; + void* q8_ref = nullptr; + const size_t qb = k3_q8_0_bytes(n); + CU(cudaMalloc(&dsitu, (size_t)n * 4)); + CU(cudaMalloc(&dsitu2, (size_t)n * 4)); + CU(cudaMalloc(&q8_fused, qb)); + CU(cudaMalloc(&q8_ref, qb)); + + const bool took = k3_situ_q8(q8_fused, dsitu, dg, du, n, beta, lb, 0); + ++g_case; + std::printf(" %-46s %s\n", "fast path engaged", took ? "OK" : "FAIL (declined)"); + if (!took) ++g_fail; + + situ_f32(dsitu2, dg, du, n, beta, lb, 0); + k3_quantize_q8_0(q8_ref, dsitu2, n / 32, 0); + CU(cudaDeviceSynchronize()); + CU(cudaGetLastError()); + + auto f_fused = from_dev(dsitu, n); + auto f_ref = from_dev(dsitu2, n); + ++g_case; + const bool f_ok = bytes_eq(f_fused.data(), f_ref.data(), (size_t)n * 4); + std::printf(" %-46s %s\n", "float situ mirror bit-identical", f_ok ? "OK" : "FAIL"); + if (!f_ok) ++g_fail; + + std::vector hq(qb), href(qb); + CU(cudaMemcpy(hq.data(), q8_fused, qb, cudaMemcpyDeviceToHost)); + CU(cudaMemcpy(href.data(), q8_ref, qb, cudaMemcpyDeviceToHost)); + ++g_case; + const bool q_ok = bytes_eq(hq.data(), href.data(), qb); + std::printf(" %-46s %s\n", "Q8_0 bytes bit-identical", q_ok ? "OK" : "FAIL"); + if (!q_ok) ++g_fail; + + // Hot path passes nullptr for the float mirror — Q8 must still match. + void* q8_nomirror = nullptr; + CU(cudaMalloc(&q8_nomirror, qb)); + const bool took2 = k3_situ_q8(q8_nomirror, nullptr, dg, du, n, beta, lb, 0); + CU(cudaDeviceSynchronize()); + ++g_case; + std::printf(" %-46s %s\n", "nullptr float mirror engages", took2 ? "OK" : "FAIL"); + if (!took2) ++g_fail; + std::vector hq2(qb); + CU(cudaMemcpy(hq2.data(), q8_nomirror, qb, cudaMemcpyDeviceToHost)); + ++g_case; + const bool q2_ok = bytes_eq(hq2.data(), href.data(), qb); + std::printf(" %-46s %s\n", "Q8_0 identical without float mirror", q2_ok ? "OK" : "FAIL"); + if (!q2_ok) ++g_fail; + CU(cudaFree(q8_nomirror)); + + CU(cudaFree(dg)); CU(cudaFree(du)); + CU(cudaFree(dsitu)); CU(cudaFree(dsitu2)); + CU(cudaFree(q8_fused)); CU(cudaFree(q8_ref)); +} + +static void test_mla_gate_q8() { + const int n = 1536; // 12 heads * 128 at tp=8 + std::printf("mla_gate + Q8_0\n"); + std::mt19937 rng(11); + auto attn = rnd(n, rng); + auto gate = rnd(n, rng); + + float* da = to_dev(attn); + float* dg = to_dev(gate); + float* dout = nullptr; + float* dout2 = nullptr; + void* q8_fused = nullptr; + void* q8_ref = nullptr; + const size_t qb = k3_q8_0_bytes(n); + CU(cudaMalloc(&dout, (size_t)n * 4)); + CU(cudaMalloc(&dout2, (size_t)n * 4)); + CU(cudaMalloc(&q8_fused, qb)); + CU(cudaMalloc(&q8_ref, qb)); + CU(cudaMemcpy(dout2, da, (size_t)n * 4, cudaMemcpyDeviceToDevice)); + + const bool took = k3_mla_gate_q8(q8_fused, dout, da, dg, n, 0); + ++g_case; + std::printf(" %-46s %s\n", "fast path engaged", took ? "OK" : "FAIL (declined)"); + if (!took) ++g_fail; + + mla_gate_out_f32(dout2, dout2, dg, n, 0); + k3_quantize_q8_0(q8_ref, dout2, n / 32, 0); + CU(cudaDeviceSynchronize()); + + auto f_fused = from_dev(dout, n); + auto f_ref = from_dev(dout2, n); + ++g_case; + const bool f_ok = bytes_eq(f_fused.data(), f_ref.data(), (size_t)n * 4); + std::printf(" %-46s %s\n", "float gate mirror bit-identical", f_ok ? "OK" : "FAIL"); + if (!f_ok) ++g_fail; + + std::vector hq(qb), href(qb); + CU(cudaMemcpy(hq.data(), q8_fused, qb, cudaMemcpyDeviceToHost)); + CU(cudaMemcpy(href.data(), q8_ref, qb, cudaMemcpyDeviceToHost)); + ++g_case; + const bool q_ok = bytes_eq(hq.data(), href.data(), qb); + std::printf(" %-46s %s\n", "Q8_0 bytes bit-identical", q_ok ? "OK" : "FAIL"); + if (!q_ok) ++g_fail; + + CU(cudaFree(da)); CU(cudaFree(dg)); + CU(cudaFree(dout)); CU(cudaFree(dout2)); + CU(cudaFree(q8_fused)); CU(cudaFree(q8_ref)); +} + +static void test_kda_gate_q8() { + const int head_dim = 128, n_head = 12; + const int n = head_dim * n_head; + std::printf("kda_gate + Q8_0 (tp=8 head band)\n"); + std::mt19937 rng(13); + auto o = rnd(n, rng); + auto g2 = rnd(n, rng); + auto w = rnd(head_dim, rng, 0.5f, 1.5f); + const float eps = 1e-5f; + + float* do_ = to_dev(o); + float* dg = to_dev(g2); + float* dw = to_dev(w); + float* dout = nullptr; + float* dout2 = nullptr; + void* q8_fused = nullptr; + void* q8_ref = nullptr; + const size_t qb = k3_q8_0_bytes(n); + CU(cudaMalloc(&dout, (size_t)n * 4)); + CU(cudaMalloc(&dout2, (size_t)n * 4)); + CU(cudaMalloc(&q8_fused, qb)); + CU(cudaMalloc(&q8_ref, qb)); + + const bool took = k3_kda_gate_q8(q8_fused, dout, do_, dw, dg, head_dim, n_head, + eps, 0); + ++g_case; + std::printf(" %-46s %s\n", "fast path engaged", took ? "OK" : "FAIL (declined)"); + if (!took) ++g_fail; + + kda_gate_out_f32(dout2, do_, dw, dg, head_dim, n_head, eps, 0); + k3_quantize_q8_0(q8_ref, dout2, n / 32, 0); + CU(cudaDeviceSynchronize()); + + auto f_fused = from_dev(dout, n); + auto f_ref = from_dev(dout2, n); + ++g_case; + const bool f_ok = bytes_eq(f_fused.data(), f_ref.data(), (size_t)n * 4); + std::printf(" %-46s %s\n", "float gate_out bit-identical", f_ok ? "OK" : "FAIL"); + if (!f_ok) ++g_fail; + + std::vector hq(qb), href(qb); + CU(cudaMemcpy(hq.data(), q8_fused, qb, cudaMemcpyDeviceToHost)); + CU(cudaMemcpy(href.data(), q8_ref, qb, cudaMemcpyDeviceToHost)); + ++g_case; + const bool q_ok = bytes_eq(hq.data(), href.data(), qb); + std::printf(" %-46s %s\n", "Q8_0 bytes bit-identical", q_ok ? "OK" : "FAIL"); + if (!q_ok) ++g_fail; + + // Decline on a non-K3 head width rather than mishandle it. + const bool declined = !k3_kda_gate_q8(q8_fused, dout, do_, dw, dg, 64, n_head, + eps, 0); + ++g_case; + std::printf(" %-46s %s\n", "head_dim=64 declines", declined ? "OK" : "FAIL"); + if (!declined) ++g_fail; + + CU(cudaFree(do_)); CU(cudaFree(dg)); CU(cudaFree(dw)); + CU(cudaFree(dout)); CU(cudaFree(dout2)); + CU(cudaFree(q8_fused)); CU(cudaFree(q8_ref)); +} + +int main() { + if (!have_device()) return 0; + std::printf("K3 epilogue→Q8_0 folds, real dims\n\n"); + test_situ_q8(); + std::printf("\n"); + test_mla_gate_q8(); + std::printf("\n"); + test_kda_gate_q8(); + return report(); +} diff --git a/runtime/src/models/kimi_k3.cpp b/runtime/src/models/kimi_k3.cpp index 69d813f..0f2814d 100644 --- a/runtime/src/models/kimi_k3.cpp +++ b/runtime/src/models/kimi_k3.cpp @@ -2574,6 +2574,9 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha } kda_pre = true; } + // Set inside the a_tok==1 iteration when the gate→Q8 fold takes; consumed + // after the scan loop to skip the re-quantise on attn_output. + bool gated_q8 = false; for (int b = 0; b < a_tok; ++b) { const float* const nrm = s.normed + (int64_t)b * H; // decay_g carries a row axis only when the pre-scan batch filled it; without @@ -2754,8 +2757,20 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha // and `gproj_b` is s.g_proj_out with b == 0. if (!fused_qkvg && !proj_h(gproj_b, nrm, L.ssm_g, qkv, H)) return false; if (!L.ssm_norm.ok()) return false; - k3k::kda_gate_out_f32(gate_b, s.delta_out, (const float*)L.ssm_norm.data, - gproj_b, head_dim, n_head, eps, stream); + // Single-token only: fold gate_out into the Q8_0 attn_output needs. The + // chunk path (a_tok > 1) keeps the split — gate per token, proj_b once. + // Skip the float mirror unless debug is attached — the Q8 scratch is the + // only consumer on the hot path. + if (a_tok == 1 && ggml_qact_proj && L.attn_output.type == 8 && s.proj_q8) { + gated_q8 = k3k::k3_kda_gate_q8( + s.proj_q8, fwd.debug ? gate_b : nullptr, s.delta_out, + (const float*)L.ssm_norm.data, gproj_b, head_dim, n_head, eps, + stream); + } + if (!gated_q8) { + k3k::kda_gate_out_f32(gate_b, s.delta_out, (const float*)L.ssm_norm.data, + gproj_b, head_dim, n_head, eps, stream); + } if (fwd.debug) fwd.debug("dbg_gate_out", layer, gate_b, qkv); } // ---- end of the per-token scan loop ---- @@ -2778,9 +2793,17 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha // collective slot, so a marker hit means s.attn_out is already written for // this token and re-projecting would overwrite it from an aliased gate_out. // A marker miss retains the exact path above. - if (!(pt_hit && pt_hit->kda_out_layer == layer) && - !proj_b(s.attn_out, H, s.gate_out, qkv, L.attn_output, H, qkv)) + // + // gated_q8 (a_tok == 1 only) already wrote proj_q8 — skip the re-quantise. + if (gated_q8) { + if (!(pt_hit && pt_hit->kda_out_layer == layer) && + !k3k::k3_proj_q8act_f32(s.attn_out, s.proj_q8, L.attn_output.data, + L.attn_output.type, H, qkv, stream)) + return false; + } else if (!(pt_hit && pt_hit->kda_out_layer == layer) && + !proj_b(s.attn_out, H, s.gate_out, qkv, L.attn_output, H, qkv)) { return false; + } if (fwd.debug) fwd.debug("kda_out", layer, s.attn_out, H); } else { const int mla_ord = kimi_k3_mla_ordinal(cfg, layer); @@ -2958,6 +2981,7 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha } if (fwd.debug) fwd.debug("dbg_preattn", layer, s.mla_attn_out, qh * cfg.value_length_mla); + bool mla_gated_q8 = false; if (L.has_attn_gate) { // Issued on its own lane so it runs UNDER the attention rather than // after it. Host issue order is irrelevant here -- every launch is @@ -2969,9 +2993,17 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha return false; if (fwd.debug) fwd.debug("dbg_gateproj", layer, s.gate_proj_out, qh * cfg.value_length_mla); if (mla_fork) dag_join(1); - k3k::mla_gate_out_f32(s.mla_attn_out, s.mla_attn_out, s.gate_proj_out, - (int64_t)qh * cfg.value_length_mla, stream, - a_tok, mlaout_row); + const int64_t gate_n = (int64_t)qh * cfg.value_length_mla; + // Single-token fold into Q8; chunk path keeps the batched gate + proj_b. + // Float mirror only when debug needs it. + if (a_tok == 1 && ggml_qact_proj && L.attn_output.type == 8 && s.proj_q8) { + mla_gated_q8 = k3k::k3_mla_gate_q8( + s.proj_q8, fwd.debug ? s.mla_attn_out : nullptr, s.mla_attn_out, + s.gate_proj_out, gate_n, stream); + } + if (!mla_gated_q8) + k3k::mla_gate_out_f32(s.mla_attn_out, s.mla_attn_out, s.gate_proj_out, + gate_n, stream, a_tok, mlaout_row); if (fwd.debug) fwd.debug("dbg_postgate", layer, s.mla_attn_out, qh * cfg.value_length_mla); } @@ -2979,9 +3011,15 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha // the cap, and repointed at the collective's buffer by the driver), and its // per-token pitch is H — the same layout the batched FfnPartial then READS // it at, and the same one nb*hidden reduces. - if (!proj_b(s.attn_out, H, s.mla_attn_out, mlaout_row, L.attn_output, H, - qh * cfg.value_length_mla)) + if (mla_gated_q8) { + if (!k3k::k3_proj_q8act_f32(s.attn_out, s.proj_q8, L.attn_output.data, + L.attn_output.type, H, + qh * cfg.value_length_mla, stream)) + return false; + } else if (!proj_b(s.attn_out, H, s.mla_attn_out, mlaout_row, L.attn_output, H, + qh * cfg.value_length_mla)) { return false; + } if (fwd.debug) fwd.debug("mla_out", layer, s.attn_out, H); } @@ -3128,18 +3166,36 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha cfg.dense_ffn, H)) return false; if (fwd.debug) fwd.debug("dbg_dense_gate", layer, s.dense_gate, cfg.dense_ffn); if (fwd.debug) fwd.debug("dbg_dense_up", layer, s.dense_up, cfg.dense_ffn); - // Elementwise, so the row is purely a grid .y and there is no reduction to - // re-partition. - k3k::situ_f32(s.dense_situ, s.dense_gate, s.dense_up, cfg.dense_ffn, - cfg.situ_beta, cfg.situ_linear_beta, stream, - n_tok, cfg.dense_ffn); + // situ exists only to feed ffn_down's quantise — fold on the single-token + // path when the weight is Q8_0; the chunk path keeps situ + proj_b. + // Float mirror only when debug tags it; otherwise skip the dense_ffn write. + bool dense_situ_q8 = false; + if (n_tok == 1 && ggml_qact_proj && L.ffn_down.type == 8 && s.proj_q8) { + dense_situ_q8 = k3k::k3_situ_q8( + s.proj_q8, fwd.debug ? s.dense_situ : nullptr, s.dense_gate, + s.dense_up, cfg.dense_ffn, cfg.situ_beta, cfg.situ_linear_beta, + stream); + } + if (!dense_situ_q8) { + // Elementwise, so the row is purely a grid .y and there is no reduction to + // re-partition. + k3k::situ_f32(s.dense_situ, s.dense_gate, s.dense_up, cfg.dense_ffn, + cfg.situ_beta, cfg.situ_linear_beta, stream, + n_tok, cfg.dense_ffn); + } if (fwd.debug) fwd.debug("dbg_dense_situ", layer, s.dense_situ, cfg.dense_ffn); // s.ffn_out is read back in FfnFinish, one phase and (in the chunk driver) // one collective later, so it MUST be per-token. It was not, and the leading // dense layer's FFN output was therefore the LAST token of the chunk's for // every token of it -- see the note on this in the FfnFinish residual add. - if (!proj_b(s.ffn_out, H, s.dense_situ, cfg.dense_ffn, L.ffn_down, H, - cfg.dense_ffn)) return false; + if (dense_situ_q8) { + if (!k3k::k3_proj_q8act_f32(s.ffn_out, s.proj_q8, L.ffn_down.data, + L.ffn_down.type, H, cfg.dense_ffn, stream)) + return false; + } else if (!proj_b(s.ffn_out, H, s.dense_situ, cfg.dense_ffn, L.ffn_down, H, + cfg.dense_ffn)) { + return false; + } } else { // --- FORK: routed_down and the whole shared expert ---------------- // @@ -3389,6 +3445,11 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha L.ffn_gate_shexp, shexp_band, H)) return false; if (!proj_hb(s.dense_up, cfg.dense_ffn, normed2_src, H, L.ffn_up_shexp, shexp_band, H)) return false; + k3k::situ_f32(s.dense_situ, s.dense_gate, s.dense_up, shexp_band, + cfg.situ_beta, cfg.situ_linear_beta, l_shx.st, + n_tok, cfg.dense_ffn); + if (!proj_b(s.shexp_out, moe_row, s.dense_situ, cfg.dense_ffn, + L.ffn_down_shexp, H, shexp_band)) return false; } else { if (!proj_h_on(l_shx, s.dense_gate, normed2_src, L.ffn_gate_shexp, shexp_band, H)) @@ -3396,16 +3457,28 @@ bool kimi_k3_forward_layer_phase(KimiK3Forward& fwd, int layer, K3LayerPhase pha if (!proj_h_on(l_shx, s.dense_up, normed2_src, L.ffn_up_shexp, shexp_band, H)) return false; - } - k3k::situ_f32(s.dense_situ, s.dense_gate, s.dense_up, shexp_band, - cfg.situ_beta, cfg.situ_linear_beta, l_shx.st, - n_tok, cfg.dense_ffn); - if (n_tok > 1) { - if (!proj_b(s.shexp_out, moe_row, s.dense_situ, cfg.dense_ffn, - L.ffn_down_shexp, H, shexp_band)) return false; - } else if (!proj_on(l_shx, s.shexp_out, s.dense_situ, - L.ffn_down_shexp, H, shexp_band)) { - return false; + // Same situ+Q8 fold as the leading dense path; nothing reads the + // float situ after this on the q8 path, so skip the mirror. + bool shexp_situ_q8 = false; + if (ggml_qact_proj && L.ffn_down_shexp.type == 8 && l_shx.q8) { + shexp_situ_q8 = k3k::k3_situ_q8( + l_shx.q8, nullptr, s.dense_gate, s.dense_up, shexp_band, + cfg.situ_beta, cfg.situ_linear_beta, l_shx.st); + } + if (!shexp_situ_q8) + k3k::situ_f32(s.dense_situ, s.dense_gate, s.dense_up, + shexp_band, cfg.situ_beta, cfg.situ_linear_beta, + l_shx.st); + if (shexp_situ_q8) { + if (!k3k::k3_proj_q8act_f32(s.shexp_out, l_shx.q8, + L.ffn_down_shexp.data, + L.ffn_down_shexp.type, H, + shexp_band, l_shx.st)) + return false; + } else if (!proj_on(l_shx, s.shexp_out, s.dense_situ, + L.ffn_down_shexp, H, shexp_band)) { + return false; + } } } if (fwd.debug) fwd.debug("dbg_shexp_partial", layer, s.shexp_out, H);