feat(quantization): support compressed-tensors W8A8 and mixed modules

[sam0336]

Summary

• Parse Hugging Face compressed-tensors metadata, including explicit groups,
  presets, ignored modules, quantization status, KV-cache metadata, and nested
  text_config.quantization_config.
• Resolve quantization per module with deterministic target precedence:
  exact module name > re: regex > module type. Reject ambiguous matches and
  unsupported schemes instead of silently selecting an incompatible kernel.
• Enable module-level W8A8/BF16 mixed execution for Qwen3 attention and MLP
  projections while enforcing the uniform schemes required by fused QKV and
  fused gate/up projections.
• Load weight_scale tensors as FP32 and shard output-channel scales and biases
  consistently under tensor parallelism.
• Add focused C++ tests, a Python scale-loading regression test, a real Qwen3
  W8A8 configuration fixture, and standalone xmake test targets.

Motivation

InfiniLM previously selected one quantization method for the entire model. A
compressed-tensors checkpoint can instead define several config_groups and
an ignore list, so global selection can quantize a module that the checkpoint
stores in BF16 or choose a kernel for an unsupported scheme.

This PR adds a fail-closed, module-aware resolution path and validates it with
the public
RedHatAI/Qwen3-4B-Instruct-2507-quantized.w8a8
checkpoint and an independently generated hybrid checkpoint containing both
W8A8 and BF16 projections.

Related issue: N/A — this is a Project #4 quantization contribution.

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

Design and Supported Execution Scheme

The parser retains the original JSON in raw_config for diagnostics and
forward-compatible metadata inspection. Execution is intentionally narrower
than parsing: a module is dispatched to the existing W8A8 kernel only when its
resolved group has all of the following properties:

• format = "int-quantized";
• static symmetric INT8 per-channel weights;
• dynamic symmetric INT8 per-token input activations;
• no output-activation quantization.

W8A8/INT8 and UNQUANTIZED preset groups are supported. Other parsed
strategies remain visible in metadata but produce a path-qualified error at
resolution time. This prevents a checkpoint from appearing supported while
being executed with the wrong quantization semantics.

For fused projections, q_proj, k_proj, and v_proj must resolve to the same
scheme; the same invariant applies to gate_proj and up_proj. o_proj and
down_proj are independently selectable, which enables practical mixed
W8A8/BF16 checkpoints.

Test Results of Involved Models on Supported Platforms

Platform matrix

Platform	Build / config tests	Real W8A8 inference	Status
NVIDIA, 2 x RTX 4090	PASS	PASS, TP1 and TP2	Fully validated for this PR
Linux x86_64 host	PASS for host-side parser/matcher tests	N/A	This is not a claim of CPU-backend W8A8 support
QY / MetaX / Iluvatar / Cambricon / Hygon / Ascend / Kunlun	Not tested	Not tested	Hardware was not available; maintainer validation requested

Focused tests

The following commands pass on the current branch:

xmake run compressed_tensors_config_test
xmake run module_target_matcher_test
xmake run quant_config_test
xmake run compressed_tensors_test

Coverage includes:

• explicit groups and W8A8/unquantized presets;
• target matching and exact/regex/type precedence;
• ignored modules and equal-specificity ambiguity detection;
• unsupported formats and schemes failing closed;
• root and nested quantization_config extraction;
• module-aware factory selection;
• W8A8 parameter layout, FP32 scales, and TP sharding;
• a real Qwen3-4B-Instruct-2507 W8A8 config.json shape.

End-to-end model validation

Checkpoint	TP	Result
RedHatAI/Qwen3-4B-Instruct-2507-quantized.w8a8	1	PASS: load + text generation
RedHatAI/Qwen3-4B-Instruct-2507-quantized.w8a8	2	PASS: load + text generation
Hybrid W8A8/BF16 checkpoint	1	PASS: load + text generation
Hybrid W8A8/BF16 checkpoint	2	PASS: load + text generation

The hybrid checkpoint uses W8A8 for Transformer linear layers while keeping
lm_head and model.layers.0.self_attn.o_proj in BF16. Passing both TP1 and
TP2 demonstrates that ignore changes actual parameter layout and execution,
not only parser output.

Benchmark / Performance Impact

Methodology

• Revision: InfiniLM 29ddfb9bafa1805fb8685e45e09debac9c638cec, rebased on
  InfiniLM-v0.2.9c at 270feb3e14976ca14c3b012df1d1740f939a941e.
• Date and hardware: 2026-09-18, 2 x NVIDIA RTX 4090 (24 GiB each), NVIDIA
  driver 570.124.06. TP1 results below use GPU 0 only.
• Runtime: Python 3.12 and the same local InfiniCore installation for both
  checkpoints.
• Models: Qwen/Qwen3-4B-Instruct-2507
  BF16 and the corresponding public Red Hat AI compressed-tensors W8A8
  checkpoint linked above.
• Harness: examples/bench.py, greedy decoding, 32 generated tokens, one
  warmup, and three measured repetitions per shape.
• Reported numbers are arithmetic means of the three repetitions.
• BF16 and W8A8 use the same InfiniLM revision, prompt, device, and benchmark
  configuration. Only the checkpoint changes.

Offline performance, TP1

Batch	Input tokens	TTFT BF16 / W8A8 (ms)	TTFT reduction	Prefill BF16 / W8A8 (tok/s)	Prefill gain	Decode BF16 / W8A8 (tok/s)	Decode gain
1	128	18.31 / 14.66	19.9%	6,996 / 8,745	25.0%	75.30 / 91.86	22.0%
1	512	48.75 / 32.59	33.2%	10,505 / 15,730	49.7%	74.82 / 91.66	22.5%
1	2,048	217.26 / 145.83	32.9%	9,426 / 14,045	49.0%	72.98 / 89.43	22.5%
4	128	49.68 / 33.40	32.8%	10,311 / 15,347	48.8%	287.24 / 373.79	30.1%
4	512	198.48 / 127.84	35.6%	10,318 / 16,019	55.2%	287.45 / 369.81	28.7%
8	128	89.74 / 62.10	30.8%	11,410 / 16,493	44.5%	573.04 / 734.81	28.2%
8	512	406.64 / 247.93	39.0%	10,073 / 16,520	64.0%	564.38 / 721.94	27.9%

Across these shapes, W8A8 improves prefill throughput by 25.0%-64.0% and
decode throughput by 22.0%-30.1%. In the batch-1 benchmark process, checkpoint
loading decreased from 2,443.8 ms to 1,677.8 ms (31.3%).

Concurrent service performance, TP1

Both services used 64 requests, concurrency 20, greedy sampling, 64 maximum
new tokens, max_batch_size=20, a static KV cache of 512 tokens, and the same
GPU. Each model used a freshly started service and one measured workload, with
no separate service warmup. The client used a fixed random seed of 0 and
consumed the OpenAI-compatible streaming endpoint; default prefix caching
remained enabled for both runs. Each run completed 64/64 requests with no
client or server errors.

Model	Total time (s)	Request rate (RPS)	Output chunks	Average latency (s)	Average TTFT (s)	Resident GPU memory (MiB)
BF16	59.82	1.07	4,019	16.10	15.24	11,306
W8A8	46.16	1.39	4,011	12.27	11.56	6,712
W8A8 change	-22.8%	+29.9%	-0.2%	-23.8%	-24.1%	-40.6%

Resident GPU memory was sampled with nvidia-smi after each workload while the
service remained loaded; GPU 1 stayed idle. W8A8 used 4,594 MiB less resident
memory. Service checkpoint loading also decreased from 2,672.2 ms to 1,743.0 ms
(34.8%).

Accuracy and evaluation throughput

Benchmark	BF16 accuracy	W8A8 accuracy	Accuracy delta	BF16 / W8A8 throughput	Throughput gain
MMLU validation, 1,531 samples	68.84% (1,054/1,531)	67.80% (1,038/1,531)	-1.04 pp	1,965 / 2,415 tok/s	22.9%
C-Eval validation, 1,346 samples	53.86% (725/1,346)	55.05% (741/1,346)	+1.19 pp	1,691 / 2,037 tok/s	20.5%

The C-Eval increase is not presented as evidence that quantization improves
model quality; small benchmark deltas can result from changed answer choices
near decision boundaries. The relevant observation is that no material
accuracy collapse was observed on either suite.

Weight footprint

The checked checkpoint index metadata reports 8.05 GB of BF16 safetensors and
5.19 GB of W8A8 safetensors, a 35.5% reduction. This is an end-to-end checkpoint
footprint, not an idealized 2x figure: embeddings, normalization weights, and
other non-linear tensors remain at higher precision. The service measurement
above separately captures the resident runtime allocation after the workload;
it is not presented as a sampled peak-memory trace.

Reproduction

Build the extension and focused tests using the same InfiniCore installation:

git checkout 29ddfb9bafa1805fb8685e45e09debac9c638cec
git submodule update --init --recursive
xmake f -m release
xmake build _infinilm
xmake build compressed_tensors_config_test module_target_matcher_test \
  quant_config_test compressed_tensors_test
xmake run compressed_tensors_config_test
xmake run module_target_matcher_test
xmake run quant_config_test
xmake run compressed_tensors_test

Run the public W8A8 checkpoint through the official single-request entry point
on both validated tensor-parallel configurations:

W8A8=/path/to/Qwen3-4B-Instruct-2507-quantized.w8a8

for TP in 1 2; do
  python examples/test_infer.py --device nvidia --backend cpp \
    --model "$W8A8" --tp "$TP" --max-new-tokens 64 \
    --prompt "Explain in one sentence why mixed quantization is useful."
done

The evidence archive includes make_mixed_checkpoint.py. It dequantizes one
selected projection to BF16, removes its scale tensor, updates the checkpoint
index, and adds the module to the compressed-tensors ignore list without
modifying the source checkpoint:

HYBRID=/path/to/Qwen3-4B-Instruct-2507-Hybrid-W8A8-BF16

python make_mixed_checkpoint.py \
  --source "$W8A8" --output "$HYBRID" \
  --module model.layers.0.self_attn.o_proj

for TP in 1 2; do
  python examples/test_infer.py --device nvidia --backend cpp \
    --model "$HYBRID" --tp "$TP" --max-new-tokens 64 \
    --prompt "Explain in one sentence why mixed quantization is useful."
done

Run the TP1 performance comparison. Repeated input lengths produce three
measurements for each shape:

BF16=/path/to/Qwen3-4B-Instruct-2507
W8A8=/path/to/Qwen3-4B-Instruct-2507-quantized.w8a8

for MODEL in "$BF16" "$W8A8"; do
  python examples/bench.py --device nvidia --backend cpp --model "$MODEL" \
    --tp 1 --batch-size 1 \
    --input-len 128,128,128,512,512,512,2048,2048,2048 \
    --output-len 32 --warmup
  python examples/bench.py --device nvidia --backend cpp --model "$MODEL" \
    --tp 1 --batch-size 4 \
    --input-len 128,128,128,512,512,512 \
    --output-len 32 --warmup
  python examples/bench.py --device nvidia --backend cpp --model "$MODEL" \
    --tp 1 --batch-size 8 \
    --input-len 128,128,128,512,512,512 \
    --output-len 32 --warmup
done

Run the evaluation suites from a populated Hugging Face dataset cache:

python test/bench/test_benchmark.py --device nvidia --backend cpp \
  --model "$W8A8" --tp 1 --bench mmlu --subject all --split val \
  --max-new-tokens 5 --cache-dir /path/to/hf-cache \
  --output-csv w8a8_mmlu_val.csv

python test/bench/test_benchmark.py --device nvidia --backend cpp \
  --model "$W8A8" --tp 1 --bench ceval --subject all --split val \
  --max-new-tokens 5 --cache-dir /path/to/hf-cache \
  --output-csv w8a8_ceval_val.csv

Repeat the same two commands with --model "$BF16" for the baseline.

Run the concurrent service comparison once for each model:

python python/infinilm/server/inference_server.py \
  --device nvidia --model "$MODEL" --tp 1 \
  --max-new-tokens 64 --max-batch-size 20 --max-cache-len 512 \
  --temperature 1.0 --top-p 0.8 --top-k 1 \
  --host 127.0.0.1 --port 8000

# In a second terminal. The evidence harness uses the same streaming API and
# workload as scripts/test_perf.py while explicitly closing its HTTP client.
python run_service_benchmark_clean.py \
  --model "$MODEL" --num-requests 64 --concurrency 20

nvidia-smi --query-gpu=index,name,memory.used,memory.total --format=csv,noheader

The installed OpenAI 3.13.0 plus httpx2 2.12.0 client completed 64/64
requests with scripts/test_perf.py, but emitted asynchronous-generator
warnings during interpreter teardown. The attached evidence harness uses
standard httpx 0.28.1 against the same streaming endpoint and workload so
that transport cleanup is deterministic. No InfiniLM source change is needed.

Evidence

The external
infinilm-compressed-tensors-w8a8-29ddfb9-rtx4090-evidence.tar.gz archive
contains the raw BF16/W8A8 offline, accuracy, single-request, and service logs;
CSV summaries; the hybrid checkpoint metadata and construction script; and the
explicit-cleanup service harness. The archive and six representative NVIDIA result screenshots are attached
below. The archive SHA-256 is provided next to its download link.
infinilm-compressed-tensors-w8a8-29ddfb9-rtx4090-evidence.tar.gz
SHA-256：2620c839a0d42bda8315ff87e4a4ba03bd95b7ea266ba383310e6f09efb705e3

Notes for Reviewers

• Review the fail-closed boundary between metadata parsing and executable W8A8
  scheme selection. Parsing a field does not imply kernel support.
• Target precedence and ambiguity errors are intentional; object iteration
  order must not determine the selected quantization group.
• Fused QKV and gate/up projections cannot currently mix schemes internally.
  They fail with an explicit error, while output/down projections remain
  independently selectable.
• KV-cache metadata is parsed but this PR does not add a new KV-cache
  quantization kernel.
• The service comparison measures the complete HTTP streaming path. Its raw
  client and server logs, together with the explicit-cleanup harness, are
  included in the external evidence archive described above.

CI / ChatOps

This branch is rebased onto the requested InfiniLM-v0.2.9c target. A fork
workflow dispatch passed Check Format and matrix generation for commit
29ddfb9, but its hardware jobs could not acquire the upstream self-hosted
runners. A maintainer has been asked to run /retest on this PR.

Checklist

Title, Branch, and Commits

• PR title follows Conventional Commits (feat(quantization): support compressed-tensors W8A8 and mixed modules).
• Branch name follows <type>/xxx-yyyy-zzzz: feat/compressed-tensors-w8a8.
• Each commit message follows Conventional Commits.
• This is a large PR whose eight commits are meaningful, well-formed, and independently reviewable.
• No stray merge commits; the branch is rebased cleanly on the requested InfiniLM-v0.2.9c target.
• No fixup! / squash! / wip commits remain.
• N/A — No existing PR, branch, or commit uses the legacy issue format.

Scope and Design

• Changes are minimal and limited to compressed-tensors parsing, module-aware W8A8 selection, required loading/sharding fixes, and tests.
• No dead code, commented-out blocks, debug prints, or unowned TODO remains.
• No unrelated formatting churn obscures the diff.
• N/A — No stable user-facing API is changed; checkpoint selection remains automatic through model metadata.

General Code Hygiene (applies to all languages)

• Comments are limited to non-obvious behavior and invariants.
• Every modified or added file ends with a trailing newline.
• git diff --check reports no trailing whitespace or whitespace errors.
• Identifiers in new comments and diagnostics follow the repository's Markdown convention.
• New comments and diagnostics are in English.
• New comments and diagnostics follow the applicable language conventions.

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

• Code follows the repository's Google C++ style requirements.
• Error and warning wording follows the LLVM diagnostic convention.
• Constructor initializer order matches member declaration order.
• No raw new / delete was introduced.
• Changed C++ files pass scripts/format.py --check.
• No changes or references to csrc/models/llama_legacy/ were added.

Python Specific (if Python files changed)

• Changed Python code is PEP 8 compliant.
• New Python comments are complete English sentences.
• New docstrings follow PEP 257.
• Changed Python files pass scripts/format.py --check.
• No changes or references to python/infinilm/auto_config.py were added.

Testing

• Untested platforms are identified above and maintainer validation is requested; tag an appropriate reviewer after opening the PR.
• Passed examples/test_infer.py with the public W8A8 and generated hybrid checkpoints on NVIDIA TP1 and TP2.
• Passed examples/bench.py for the BF16/W8A8 comparison.
• Passed test/bench/test_benchmark.py on MMLU and C-Eval validation splits.
• The official scripts/test_perf.py workload completed 64/64 requests; the reported comparison uses the equivalent explicit-cleanup harness because the installed OpenAI/httpx2 client emitted teardown warnings.

Build, CI, and Tooling

• A clean target rebuild passes on NVIDIA; confirm the fresh-directory build through CI.
• Requested an upstream /retest from a maintainer after the fork workflow passed format and matrix generation.

Documentation

• N/A — No new flag or manual workflow is required; model support is selected automatically from checkpoint metadata and is documented in this PR.
• N/A — This PR has no user-visible breaking change.

Security and Safety

• No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers are committed.
• N/A — No third-party source code is introduced.
• No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks are introduced.

[sam0336]

Could a maintainer please run /retest for this PR? The fork workflow dispatch passed Check Format and matrix generation for commit 29ddfb9, but its NVIDIA, MetaX, Cambricon, Hygon, and Moore jobs cannot acquire the upstream self-hosted runners. The implementation has also been rebuilt and validated on 2 x RTX 4090; reproducible logs and screenshots are attached in the PR description.