Skip to content

[Feature] Add optional FP8 MLA cache for GLM-5.3-Flash - #5045

Merged
RunningLeon merged 6 commits into
InternLM:mainfrom
qescccczmr:feat/glm53-fp8-mla-fp8-cache
Oct 21, 2026
Merged

RunningLeon merged 6 commits into
InternLM:mainfrom
qescccczmr:feat/glm53-fp8-mla-fp8-cache

Conversation

@qescccczmr

@qescccczmr qescccczmr commented Oct 20, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Enable opt-in --quant-policy fp8 for GLM-5.3-Flash target and MTP MLA caches on Hopper, following vLLM #59246. GLM manages KPool itself and sets mla_index_topk=None, so the existing cache-policy conversion did not select its FP8 MLA path. The default BF16 behavior remains unchanged.

The FlashMLA cache row changes from 512 BF16 values (1024 B) to 512 FP8 values, four FP32 scales and 64 BF16 zeros (656 B). Decode pads the 512-wide NoPE query to 576 while preserving the model's attention scale. The writer clears the unused tail for each written token; prefill dequantizes to the logical 512 dimensions. Sparse decode preserves KPool's 2048 selected tokens plus up to three tail tokens and pads the index width to 2112 for the FP8 kernel. This uses FlashMLA 1.0.0+ba89a34 on H200.

Validation

  • 166 H200 tests passed, covering cache writes and reuse, prefill, sparse decode, MTP verification, CUDA Graph and DCP regressions. Ruff and formatting checks passed. FP8 MTP5 generation also passed TP8 eager and TP1/DP8/EP8 CUDA Graph smoke tests.
  • Cache allocation and GSM8K: 8×H200, TP8, MTP5, CUDA Graph, --cache-max-entry-count 0.7, batch size 32, session length 32768, with otherwise matched BF16/FP8 server settings. Cache sizes are actual owning allocations per GPU, excluding model weights and runtime workspaces.
Metric BF16 (--quant-policy none) FP8 (--quant-policy fp8)
Allocated block caches (GiB/GPU) 65.3743 65.3742
Separate KDA/conv/KPool state (GiB/GPU) 3.4480 3.4480
Usable paged-cache capacity (tokens, TP group) 5,319,360 7,994,816
GSM8K 5-shot, full 1319-question test 1209/1319 (91.66%) 1213/1319 (91.96%)

At the same allocated block-cache budget, FP8 increases usable token capacity by 50.30%. The separate state pool is excluded from the block-cache row. GSM8K used temperature 0, seed 42 and a 256-token output limit; the +0.30-point difference is a single paired result, not an established accuracy gain.

  • GPQA Diamond (FP8, MTP5): 1449/1584 correct (91.48%) in one 8×H200 TP8 serving run with --cache-max-entry-count 0.7, prefix caching and 5 MTP draft tokens. The evaluation used temperature 1, top-p 0.95 and thinking enabled. One response reached the output limit. There is no matched BF16 GPQA baseline, so this measures FP8 accuracy rather than an accuracy change.
  • Decode performance: A separate fixed-output sweep used 8×H200, TP8, CUDA Graph, MTP off, cache budget 0.7, 256 output tokens, and 4k/8k/16k/32k/46k/512k input tokens. At concurrency 8, FP8 throughput changed by +7.3%, +5.5%, −0.8%, −2.2%, −3.3% for 4k through 46k respectively; mean TPOT changed by −7.0%, −6.0%, +0.2%, +4.0%, +3.5% (negative is faster). At 512k and concurrency 1, throughput was 5.27 → 5.30 tok/s (+0.6%) and mean TPOT 26.56 → 26.03 ms (2.0% lower); TTFT was about 42 s. These are single paired runs with fixed prompt/output lengths, so small differences are inconclusive. FP8 does not show a consistent long-context speedup here. Fixed-output serving speed with MTP5 remains unmeasured; the GSM8K runs have different generated outputs.

Comment thread tests/pytorch/kernel/test_nope_mla_fp8.py

@RunningLeon RunningLeon left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@RunningLeon
RunningLeon merged commit 1890f50 into InternLM:main Oct 21, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants