fix: #552 int8 kv cache quantization cannot be enabled

[BoBoDai]

Summary

Fix --kv-cache-dtype=int8 end-to-end: the flag was silently dropped in the Python layer, so no quantization ever ran; once the chain is wired through, the INT8 path exposes several C++ correctness bugs. This PR makes INT8 KV-cache quantization actually usable.

• csrc/layers/attention/attention.cpp: initialize kv_cache_k_scale/kv_cache_v_scale with Tensor::ones — previously Parameter({1}, F32, device, 0, 0, 1) only sets TP config (tp_dim/tp_rank/tp_size), leaving the storage uninitialized (read as 0.0, so dequantized K/V became all zeros); also create/register the scale parameters before constructing AttentionLayer (it captures them by value, otherwise it keeps stale empty tensors).
• csrc/layers/attention/backends/static_attn.cpp: permute K/V in model precision first, then quantize — quantized dtypes must not pass through permute/rearrange.
• csrc/layers/quantization/kv_quant.cpp: zero_point must be F32 (the infiniop INT8 kernels read it as float*); it was created in the K dtype (bf16).
• examples/test_infer.py, python/infinilm/llm/llm.py, python/infinilm/config/engine_config.py, python/infinilm/llm/model_runner/model_runner.py, python/infinilm/llm/model_runner/speculative_runner.py, python/infinilm/server/inference_server.py, python/infinilm/server/pipeline_worker.py: thread kv_cache_dtype from CLI to the C++ engine (LLM / AsyncLLMEngine / server / worker / draft-runner paths).
• python/infinilm/modeling_utils.py: tolerate missing kv_cache_k_scale / kv_cache_v_scale keys during checkpoint loading — these are registered in C++ only when INT8 is enabled and initialized to 1.0; HF checkpoints never contain them.
• examples/bench.py: KV-cache memory reporting now reflects kv_cache_dtype — the per-case estimate uses the actual cache dtype (int8/fp8 = 1 byte/element) and includes both K and V; the display unit auto-switches (B/KB/MB/GB); a measured value is printed from the live cache tensors ([bench] measured KV cache memory), which matches the estimate exactly.

Motivation

--kv-cache-dtype=int8 has no effect on main: base_config.py parses it but every hop below drops it, so the C++ engine receives None, the quant scheme stays NONE, and the INT8 path never executes. After wiring the chain, the path exposes correctness bugs: uninitialized scale (0.0) → dequantized KV all zero → uniform attention → degenerate output; stale empty scale tensors in the attention backend; bf16 zero_point misread as float*; quantize-before-permute.

Closes #552

Type of Change

• feat — new feature / new model
• fix — bug fix
• perf — performance improvement (no behavioral change)
• refactor — code restructuring without behavior change
• test — adding or fixing tests only
• docs — documentation only
• build / ci — build system or CI configuration
• chore — tooling, formatting, or other non-code changes
• Breaking change

kv_cache_dtype defaults to None; all pre-existing paths are unchanged when the flag is unset.

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Single request test

Model	Platform	Test	Result
TinyLlama-1.1B-Chat (llama)	NVIDIA L40S 46GB	E2E generate, B=1, flag unset (regression)	Output correct, identical to main baseline
TinyLlama-1.1B-Chat	same	E2E generate, --kv-cache-dtype=int8	INT8 path runs without crash; all kv_cache_*_scale params read tensor([1.]); quantize → cache → dequantize round-trip matches element-wise
Qwen2.5-0.5B-Instruct	same	E2E generate, flag unset (regression)	Output correct, identical to main baseline
Qwen2.5-0.5B-Instruct	same	E2E generate, --kv-cache-dtype=int8	INT8 path runs without crash; output degrades to repeated words — expected with the default static scale 1.0 (see Notes: known limitation)

Supplementary: with a calibrated scale (0.05) the same INT8 pipeline produces coherent output for short generations (≤100 tokens) on TinyLlama-1.1B, while 1.0 collapses immediately — the degradation is monotonically tied to the scale step, not the pipeline.

• scale = 1.0
• TinyLlama-1.1B-Chat

python examples/test_infer.py --model /models/TinyLlama-1.1B-Chat-v1.0/ --device=nvidia --max-new-tokens 32

• TinyLlama-1.1B-Chat --kv-cache-dtype=int8

python examples/test_infer.py --model /models/TinyLlama-1.1B-Chat-v1.0/ --device=nvidia --kv-cache-dtype=int8 --max-new-tokens 32

• Qwen2.5-0.5B-Instruct

python examples/test_infer.py --model /models/Qwen2.5-0.5B-Instruct/ --device=nvidia --max-new-tokens 32

• Qwen2.5-0.5B-Instruct --kv-cache-dtype=int8

python examples/test_infer.py --model /models/Qwen2.5-0.5B-Instruct/ --device=nvidia --kv-cache-dtype=int8 --max-new-tokens 32

Scale sensitivity (validation, not part of this PR's code change)

Same INT8 pipeline, only the static scale value differs (TinyLlama-1.1B-Chat, NVIDIA L40S).

The calibrated value was chosen from the measured K/V magnitude (max|K| ≈ 4–5, so max/127 ≈ 0.03–0.04). The degradation is monotonically tied to the scale step, not to the pipeline: 1.0 collapses at once, while 0.05 keeps short generations coherent and only degrades on long contexts — the inherent ceiling of static-scale INT8 (error accumulation). Dynamic per-token scales are the follow-up (see Notes).

• TinyLlama-1.1B-Chat --kv-cache-dtype=int8 (with temporary scale=0.05 patch)

python examples/test_infer.py --model /models/TinyLlama-1.1B-Chat-v1.0/ --device=nvidia --kv-cache-dtype=int8 --max-new-tokens 32

• Qwen2.5-0.5B-Instruct --kv-cache-dtype=int8 (with temporary scale=0.05 patch)

python examples/test_infer.py --model /models/Qwen2.5-0.5B-Instruct/ --device=nvidia --kv-cache-dtype=int8 --max-new-tokens 32

Sanity test

TinyLlama-1.1B-Chat
NVIDIA L40S
Sanity test: test_benchmark.py, MMLU abstract_algebra, 5 samples, default bf16
1/5 = 20%, all 5 predictions identical to main baseline (no regression)

current branch

main branch

service test

Model	Platform	Test	Result
TinyLlama-1.1B-Chat	NVIDIA L40S	Service test (default bf16)	64/64 OK, RPS 1.98, 11.67 tok/s
TinyLlama-1.1B-Chat	NVIDIA L40S	Service test (--kv-cache-dtype=int8)	64/64 OK, RPS 1.86; longer outputs (no EOS) due to known static-scale-1.0 degeneration, see Notes

Benchmark / Performance Impact

N/A for perf. The INT8 KV cache is a memory optimization, so memory is measured
first.

Memory — measured from the live KV cache tensors (get_kv_cache()) and printed by examples/bench.py ([bench] measured KV cache memory), input=32, output=128, B=1:

python examples/bench.py --model <path/to/model> --device nvidia --input-len 32 --output-len 128 --kv-cache-dtype int8
python examples/bench.py --model <path/to/model> --device nvidia --input-len 32 --output-len 128

Model	bfloat16	int8	saving
TinyLlama-1.1B-Chat	3.44 MB	1.72 MB	50%
Qwen2.5-0.5B-Instruct	1.88 MB	0.94 MB (960 KB)	50%

The config-based estimate shown in the case line matches the measured value exactly. At the full cache_len=4096 allocation the same ratio applies (TinyLlama-1.1B-Chat: bf16 92.27 MB vs int8 46.14 MB; SmolLM-135M measured:
90.0 → 45.0 MB).

Throughput (examples/bench.py, NVIDIA L40S, input=32, output=128, B=1, single run; consistent across multiple runs):

Model	Metric	default	--kv-cache-dtype=int8
TinyLlama-1.1B-Chat	total_time	898.4 ms	801.1 ms
	Prefill throughput	116.8 tok/s	254.2 tok/s
	Decode throughput	203.4 tok/s	188.1 tok/s
Qwen2.5-0.5B-Instruct	total_time	751.9 ms	674.1 ms
	Prefill throughput	126.6 tok/s	269.3 tok/s
	Decode throughput	254.5 tok/s	228.7 tok/s

INT8 halves KV-cache memory and roughly doubles prefill throughput (quantized cache writes), at the cost of ~7–10% decode throughput (per-step quantize/dequantize kernels). The default path (flag unset) is unchanged — no regression (CPU timing: main 4084 ms vs PR 3984 ms mean, 64 tokens, ×3).

Notes for Reviewers

• Known limitation, intentionally out of scope: the static per-tensor scale defaults to 1.0 (SGLang convention; SGLang documents the same accuracy caveat). For INT8 the step of 1.0 is coarse, and even a calibrated scale has an inherent precision ceiling — long generations eventually show repetition degeneration as quantization error accumulates. Calibrated scales (checkpoint / JSON) and dynamic per-token scales (stored per cache position) are follow-ups.
• Why the default-1.0 degradation is not a pipeline bug: the degradation is monotonically tied to the static scale step. At 1.0 the output collapses immediately (repeated tokens); with a calibrated scale (0.05) the identical code path produces coherent output for short generations (≤100 tokens) and degrades to repetition only on much longer generations (see Scale sensitivity table). That residual long-context degeneration is the inherent precision ceiling of any static-scale INT8 scheme — quantization error accumulates over the growing context — and is exactly what dynamic per-token scales address (follow-up). The 0.05 value is a validation-only patch, not part of this PR; the shipped default remains 1.0 (SGLang convention).
• deepseek_v2_attention.cpp, deepseek_v2_mla_attention.cpp, ernie4_5_attention.cpp share the same constructor-order pattern (scales created after AttentionLayer captures them); left out of this PR to keep the change minimal — follow-up PR.
• InfiniCore's per_tensor_quant_i8 kernels are NVIDIA/QY-only. A pure-CPU build will fault at the quant op, so the INT8 path is only runnable on NVIDIA/QY; verification was done on NVIDIA (E2E + round-trip) plus allocation-only checks on CPU.
• The pre-existing Chinese comment // 无需反量化 in kv_quant.cpp was not touched by this PR.
• No external dependency introduced; the INT8 path reuses the existing per_tensor_quant_i8 / per_tensor_dequant_i8 ops.

CI / ChatOps

CI will be triggered manually from the Actions tab on this branch after the PR is opened.

---

Checklist

Every contributor must verify every item below before requesting
review. Tick each box only after the check has actually been performed —
do not tick speculatively. If an item truly does not apply, replace the
checkbox with N/A and briefly explain why in an inline comment.

Title, Branch, and Commits

• PR title follows Conventional Commits (e.g. feat(nvidia): …, fix(cuda/gemm): …).
• Branch name follows <type>/xxx-yyyy-zzzz where <type> matches the PR title's Conventional Commits type and words are joined with hyphens (see CONTRIBUTING.md §Branches).
• Each commit message follows Conventional Commits.
• Small PR is a single squashable commit; or, for a large PR, every commit is meaningful, well-formed, and independently reviewable (see CONTRIBUTING.md §Pull Requests).
• No stray merge commits from main — the branch is rebased cleanly on top of the current main.
• No fixup! / squash! / wip commits remain.
• Existing PR/branch/commit that followed the legacy issue format.

Scope and Design

• Changes are minimal — nothing unrelated to the stated motivation was added (CONTRIBUTING.md §Code/General).
• No dead code, commented-out blocks, debug prints, printf/std::cout/print(...) left behind, or TODO without an owner and issue link.
• No unrelated formatting churn that would obscure the diff.
• Public API changes (if any) are intentional, documented, and reflected in affected callers/tests.

General Code Hygiene (applies to all languages)

• The code is self-explanatory; comments were added only where the why is non-obvious (CONTRIBUTING.md §Code/General).
• Every modified or added file ends with a single trailing newline (CONTRIBUTING.md §Code/General).
• No trailing whitespace, tab/space mixing, or stray BOMs.
• Identifiers in comments and error messages are wrapped in backticks (e.g. the `seqlens_k` tensor) (CONTRIBUTING.md §Code/General).
• All comments and error messages are in English (CONTRIBUTING.md §Code/General).
• Comments and error messages are complete sentences — capitalized first letter, terminal punctuation — unless the language/framework convention says otherwise (CONTRIBUTING.md §Code/General; §Python).

C++ Specific (if C++ files changed)

• Code follows the Google C++ Style Guide strictly.
• Error and warning message wording follows the LLVM Coding Standards (CONTRIBUTING.md §C++).
• Constructor initializer list order matches member declaration order (CONTRIBUTING.md §C++).
• No raw new/delete; RAII / smart pointers / existing allocators are used.
• Changed files are formatted by scripts/format.py.
• No changes/reference to csrc/models/llama_legacy/.

Python Specific (if Python files changed)

• Code is PEP 8 compliant.
• Comments are complete English sentences, starting with a capital letter and ending with punctuation; Markdown backticks are used for code references (CONTRIBUTING.md §Python).
• Docstrings (if any) follow PEP 257 (CONTRIBUTING.md §Python).
• Changed files are formatted by scripts/format.py.
• No changes/reference to python/infinilm/auto_config.py.

Testing

• For any platform that could not be tested, an explicit reason is given in the table and a reviewer with access has been tagged.
• Passed single request test (examples/test_infer.py), or specify the reason for skipping.
• Passed offline performance test (examples/bench.py), or specify the reason for skipping.
• Passed sanity test (test/bench/test_benchmark.py), or specify the reason for skipping.
• Passed service test (python/infinilm/server/inference_server.py + scripts/test_perf.py), or specify the reason for skipping.

Build, CI, and Tooling

• The project builds cleanly from a fresh directory on at least one affected platform.
• CI has been triggered manually (Actions → CI on this branch), or /retest was requested.

Documentation

• [N/A] README.md, CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed. — no build flags or developer workflow changed; the --kv-cache-dtype flag already exists in base_config.py (this PR fixes it rather than adding it), and docstrings for the new kv_cache_dtype parameters were added inline (python/infinilm/config/engine_config.py, python/infinilm/server/inference_server.py).
• [N/A] Any user-visible breaking change is called out explicitly under "Motivation" and in the commit/PR title with a ! or BREAKING CHANGE: footer. — no breaking change: kv_cache_dtype defaults to None and all pre-existing paths are unchanged when the flag is unset.

Security and Safety

• No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers have been committed.
• Third-party code is license-compatible and attributed.
• No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.

[BoBoDai]

/retest

[github-actions]

⛔ Only repository members can run retest.