kl-consistency-test

Write, calibrate, and debug prefill-vs-decode logprob KL consistency tests in sglang.

33.0k|8.4k|Updated Jan 8, 2024
One-click install
npx skills add https://github.com/sgl-project/sglang --skill kl-consistency-test
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: kl-consistency-test
Source: https://github.com/sgl-project/sglang/tree/main/.claude/skills/kl-consistency-test
Command: npx skills add https://github.com/sgl-project/sglang --skill kl-consistency-test

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve?

Prefill and decode paths in sglang run different kernels over different shapes, so a corrupted radix-cache prefix, stale conv/mamba checkpoint, or batch-dependent kernel can silently diverge while accuracy benchmarks like gsm8k still pass. This Skill explains how to write, calibrate, and debug the KL divergence tests that catch these state and numerics bugs.

Core Features & Use Cases

  • Two-condition diagnosis: Separates batch-invariance failures (condition 1) from same-function state bugs (condition 2) using the three KLDivergenceMixin helpers, so a helper-specific divergence immediately names the faulty path.
  • Threshold calibration: Guides choosing per-(model, tp) KL thresholds, explains why MoE models amplify numerical differences versus dense models, and why deterministic-inference servers should assert near-zero.
  • Divergence localization: Covers the forward-hook dumper workflow (DUMPER_* environment variables, CUDA graph pitfalls, alignment verification) to pinpoint the first operator whose inputs match but outputs differ.
  • Use Case: A new model shows avg_kl_div of 0.05 only in the prefill_cache_hit helper. Use this Skill to recognize a prefix-restore logic bug rather than float noise, then localize it to the exact layer with the dumper.

Quick Start

Ask the AI to help debug a high KL divergence number in an sglang model test and identify whether it is a batch-invariance or cache-restore problem.

Frequently Asked Questions about kl-consistency-test

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I debug a high KL divergence between prefill and decode in sglang?

First determine which condition fails: run the three KLDivergenceMixin helpers and check whether divergence appears in all helpers (batch-invariance problem) or only in a cache-hit helper (state restore bug). Then use the forward-hook dumper to find the first layer whose inputs are bit-identical but outputs differ.

What KL divergence threshold should I use for sglang consistency tests?

Once every kernel on the path is batch-invariant, prefill and decode agree bit for bit, so assert a stray-ulp floor like 1e-9 rather than a loose tolerance. Thresholds are per (model, tp) pair, and MoE models need separate calibration from dense models because routing amplifies numerical differences.

Why is KL divergence high on MoE models but not dense models?

Top-k routing is a discrete decision over near-tied gate scores, so a tiny numerical difference can flip expert selection and compound across layers. A KL in the hundredths on MoE often reflects the same class of numerical difference that appears as ~1e-4 on a dense model.

Does enabling deterministic inference in sglang fix KL test failures?

The --enable-deterministic-inference flag swaps aten kernels for fixed-reduction versions and pins NCCL settings, but only for kernels it covers. Custom kernels that never reach an aten op remain batch-dependent, so a persistent nonzero KL means some kernel still needs batch-invariance work.

Why does the tensor dumper produce no output during CUDA graph replay?

Python forward hooks do not run during CUDA graph replay because the graph replays kernels, not Python code. Disable both CUDA graph and prefill CUDA graph when dumping, or probe the reused state outside the graph instead.