Skip to content

perf(k3): fold situ/gate epilogues into the Q8_0 their consumer needs - #153

Closed
marktech0813 wants to merge 4 commits into
gittensor-ai-lab:mainfrom
marktech0813:perf/k3-epilogue-q8
Closed

perf(k3): fold situ/gate epilogues into the Q8_0 their consumer needs#153
marktech0813 wants to merge 4 commits into
gittensor-ai-lab:mainfrom
marktech0813:perf/k3-epilogue-q8

Conversation

@marktech0813

Copy link
Copy Markdown
Contributor

Summary

hoist_act already stops the shared attn/ffn norms from being re-quantised. What is left are activations that are written by one kernel and immediately re-read by quantize_q8_0 for a single consumer:

producer consumer /token/rank
kda_gate_out attn_output 69
mla_gate_out attn_output 24
situ shexp / dense ffn_down 93

Each fold does the producer arithmetic and the warp Q8_0 emit in one launch, then the forward calls k3_proj_q8act_f32 on that scratch. On the hot path the float mirror is skipped unless a debug tag needs it.

Rebased onto current main after batched prefill (#144/#136/#148). The fold is wired on the single-token path (a_tok == 1 / n_tok == 1) — decode and any n_tok=1 call — and declines cleanly so the chunk/tile prefill path keeps situ/gate + proj_b. Bit-identical Q8 to the split path. SPARKINFER_K3_EPILOGUE_Q8=0 restores the pair.

Also includes a docs-only fix for docs/technical.md root-relative links that were failing the configs + docs CI job (same breakage exists on main).

Not the MLA slice fill (#133), not phase-major / batch tile prefill (#144/#136/#148).

Replaces #128 — that PR was mergeable (behind_by=0) but carried a stale merge-conflict label contributors cannot clear (bot-owned; only clears on the next eval round).

What kind of change is this?

  • Harness / configs / docs / CI — verified entirely by the ci workflow, no node needed
  • Perf-bearing code (runtime/, kernels/, moe/, server/, CMakeLists.txt) — needs a node run below
  • Pin change (reference.lock, a *.sha256 manifest) — needs the backing bench/results/*.json committed

Node run

Per #139 / #141 / #146: tick only with a bare number in after (this PR). Claimed prefill must beat the current frontier by >2%. Read KIMI_K3_H200X8_IQ1S_SPARKINFER_32K_PP from reference.lock on main today (docs cite 98.80 after #148) — and prefer measuring main in the same session.

No 8× H200 on the authoring box — every claim stays unticked and every after-cell empty rather than inventing a number.

  • Tested on 8× H200 (sm_90)
  • Prefill measured at 32k on 8× H200 — this is the scored metric
  • No 128k decode regression on 8× H200 — within 1% of the frontier
  • Blocked — do not evaluate yet. Something below has to resolve first. Say what, under Known problems.

Prefill tok/s @ 32k — the scored metric.

before (main) after (this PR)
prefill @ 32k

Decode tok/s @ 128k — the guard, not the tier. Must stay within 1% of the frontier.

before (main) after (this PR)
decode @ 128k
# paste same-session main vs this-PR kimi_k3_tp_bench / baseline output here

Known problems

Needs a same-session 8× H200 A/B against post-#148 main before any box can be ticked. The fold targets the single-token path (decode guard + n_tok=1); whether it moves scored prefill now that ingestion is batched is an open measurement question.

Checklist

  • Rebased onto current main (conflicts with batched prefill resolved; fold gated to n_tok/a_tok == 1)
  • CMakeLists keeps main's gpu tests + k3_epilogue_q8_gpu_test
  • Bitwise GPU test vs split path (incl. nullptr float mirror)
  • configs + docs green (technical.md relative links)
  • No sensitive-path edits / author+committer = marktech0813

The decode graph's hottest node is still the activation quantiser. Three
producers — kda_gate_out, mla_gate_out, and situ — exist only to feed one
projection a few microseconds later; folding the producer into that
quantise removes both launches per site and stays bit-identical to the
split path.
The fused Q8 path is the only consumer outside debug tags. Passing
nullptr for the float mirror drops ~186 HBM writes/token/rank without
changing the Q8 bytes the projection reads.
CI Relative links resolve walks paths from the file's directory.
Links in docs/technical.md pointed at repo-root paths, so the
configs + docs job failed with 21 broken links on an otherwise
clean merge.
@github-actions

github-actions Bot commented Aug 7, 2026

Copy link
Copy Markdown

needs-node-run

This PR changes code that can move a measured number:

  • kernels/csrc/cuda/kimi_k3/k3_epilogue_q8.cu
  • kernels/include/sparkinfer/kernels/kimi_k3_fast.h
  • kernels/tests/CMakeLists.txt
  • kernels/tests/k3_epilogue_q8_gpu_test.cu
  • runtime/src/models/kimi_k3.cpp
    CI covers syntax, the no-GPU test suite, configs and the --dry-run plans — but it
    cannot run an 802 GiB model, so it cannot tell whether this changed the numbers.
    To clear this label, edit the PR description and tick the node you measured on,
    plus both metric claims — the evaluation loop skips a PR missing any of them:
  • Tested on 8x H200 (sm_90)
  • Tested on 8x B200 (sm_100)
  • Tested on 4x+ B300 (sm_103)
    then paste the before/after from bench/scripts/kimi_k3_baseline.sh and commit the
    emitted bench/results/*.json. If a pinned baseline changed, the lock CI job will
    require that JSON to back it.
    Not a perf change? A maintainer can add no-node-run or hold.
    Automated by node-attestation CI. This label is advisory — it never closes a PR.

@marktech0813

Copy link
Copy Markdown
Contributor Author

Closing: gittensor-ai-lab/sparkinfer-k3 is not in the current active master_repositories.json emission set (only gittensor-ai-lab/sparkinfer is weighted). Refocusing on scored sparkinfer speedups.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant