mla: wire dynamic per-token NVFP4 scaling and cache ABI - #189
Merged
lukealonso merged 4 commits intoJul 29, 2026
Merged
Conversation
📝 WalkthroughWalkthroughAdds a shared NVFP4 MLA cache-format configuration, dynamic per-token scaling for compatible B12x sparse-MLA paths, and ABI-aware KV offload namespaces. Static calibration is validated and fingerprinted, while backend capability checks cover writer and kernel interfaces. ChangesNVFP4 MLA cache format and integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Environment
participant NVFP4_MLA_CACHE_FORMAT
participant B12xMLASparseImpl
participant SparkInfer
participant KVOffload
Environment->>NVFP4_MLA_CACHE_FORMAT: Parse NVFP4 MLA settings
NVFP4_MLA_CACHE_FORMAT->>B12xMLASparseImpl: Provide dynamic and FP8-RoPE configuration
B12xMLASparseImpl->>SparkInfer: Validate and invoke dynamic NVFP4 writer
NVFP4_MLA_CACHE_FORMAT->>KVOffload: Generate cache ABI
KVOffload->>KVOffload: Include ABI and KV geometry in mapper namespace
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…_DYNAMIC_SCALE) The nvfp4_ds_mla writer's outer scale is statically 1.0, parking shallow-layer group scales in E4M3 subnormals (the deep-retrieval regression root cause; PR local-inference-lab#145 calibrates it statically per layer). This mode removes the calibration requirement: the SparkInfer writer derives a per-token second-level scale stored inside the record ([292,296) of the existing pad; record stays 368 bytes) and the readers consume it from the record instead of a per-layer launch scalar. - VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 gate; requires KV_FP8_ROPE=1 368-byte records and a SparkInfer build with per_token_scale / latent_scale_per_token support (fail-closed inspect checks) - mutually exclusive with VLLM_NVFP4_MLA_SCALES_FILE (ValueError) - writer calls (do_kv_cache_update + DCP gather re-quantization) pass per_token_scale only when enabled, keeping older SparkInfer builds working with the mode off - kernel format kwargs pin latent_scale=1.0 and add latent_scale_per_token in mode Draft: not yet compiled/run (no CUDA on authoring host); CN4 validation runbook in glm52-opt design/nvfp4-dynamic-second-level- scale-phaseA-addendum.md. Co-authored-by: Fable <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YHuYXXq9krPdvcgpzY4EfE
yatesdr
force-pushed
the
nvfp4-dynamic-token-scale
branch
from
July 28, 2026 02:32
91dff5a to
ae541fc
Compare
yatesdr
marked this pull request as ready for review
July 28, 2026 05:59
Capture one server-static writer/reader mode, fail closed on incompatible SparkInfer APIs, and include the record ABI in persistent offload namespaces. Add focused tests for invalid modes, both writer sites, and stale-cache separation. Assisted-by: OpenAI Codex
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
vllm/model_executor/layers/mla_cache_format.py (1)
26-65: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the public configuration API in Google style.
from_env()andvalidate()lack docstrings, andrecord_abi()omits its argument, return value, and failure contract.Proposed fix
`@classmethod` def from_env(cls) -> Nvfp4MlaCacheFormat: + """Capture the process-wide NVFP4 MLA cache configuration. + + Returns: + The immutable configuration derived from environment variables. + """ return cls( def validate(self) -> None: + """Reject configurations without a stable NVFP4 MLA record format. + + Raises: + ValueError: If dynamic scaling conflicts with static scales or FP8 RoPE. + """ if self.dynamic_scale and self.scales_file: def record_abi(self, cache_dtype: str) -> str: - """Return an identity suitable for persistent external-cache keys.""" + """Build the persistent cache ABI for a cache dtype. + + Args: + cache_dtype: Cache data type supplied by the backend. + + Returns: + A persistent cache ABI identifier. + + Raises: + ValueError: If NVFP4 MLA configuration is invalid or static scales + cannot be fingerprinted. + """As per coding guidelines, use Google-style docstrings with
Args:/Returns:/Raises:sections.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/model_executor/layers/mla_cache_format.py` around lines 26 - 65, Add Google-style docstrings to Nvfp4MlaCacheFormat.from_env and validate, and expand record_abi’s docstring with Args, Returns, and Raises sections. Document the environment-based configuration, ABI string returned, and ValueError conditions from validation or scale-file fingerprinting without changing behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/v1/kv_offload/test_factory.py`:
- Around line 451-471: Extend the offloading configuration tests around
test_offloading_config_carries_nvfp4_record_abi to cover a non-NVFP4 cache_dtype
through build_offloading_config. Assert that
offloading_config.model.kv_cache_abi remains the "vllm-default-v1" sentinel
expected by FileMapper, while preserving the existing NVFP4 assertion.
In `@tests/v1/kv_offload/test_file_mapper.py`:
- Around line 147-172: Update
test_default_record_abi_preserves_existing_namespace to construct its mapper
through the real build_offloading_config record_abi flow for a non-NVFP4 dtype,
rather than relying on make_mapper_from_offloading_spec’s omitted kv_cache_abi
default. Assert the resulting mapper preserves the existing namespace and
validates the actual sentinel produced by record_abi.
In `@vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py`:
- Around line 150-152: Update build_offloading_config so kv_cache_abi remains
the literal "vllm-default-v1" sentinel for all non-NVFP4 cache dtypes, and call
NVFP4_MLA_CACHE_FORMAT.record_abi only when cache_dtype is "nvfp4_ds_mla" and
the dynamic-scale configuration requires it. Preserve the existing ABI sentinel
for legacy users so FileMapper.__init__ does not add new hash fields unless
NVFP4 dynamic-scale behavior is active.
In `@vllm/v1/kv_offload/file_mapper.py`:
- Around line 63-65: Update the guard in the file-mapping initialization to
recognize the upstream default ABI values produced by build_offloading_config,
rather than checking only the literal "vllm-default-v1". Preserve the legacy
namespace for all default configurations and add kv_cache_abi and
worker_kv_bytes_per_block only for configurations that require the new
namespace.
---
Nitpick comments:
In `@vllm/model_executor/layers/mla_cache_format.py`:
- Around line 26-65: Add Google-style docstrings to Nvfp4MlaCacheFormat.from_env
and validate, and expand record_abi’s docstring with Args, Returns, and Raises
sections. Document the environment-based configuration, ABI string returned, and
ValueError conditions from validation or scale-file fingerprinting without
changing behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 46f46e32-e0be-4458-9bca-e467ae2d9fee
📒 Files selected for processing (10)
tests/model_executor/layers/test_mla_cache_format.pytests/v1/attention/test_b12x_mla_fp8_rope_writer.pytests/v1/kv_offload/test_factory.pytests/v1/kv_offload/test_file_mapper.pyvllm/distributed/kv_transfer/kv_connector/v1/offloading/config.pyvllm/model_executor/layers/mla.pyvllm/model_executor/layers/mla_cache_format.pyvllm/v1/attention/backends/mla/b12x_mla_sparse.pyvllm/v1/kv_offload/config.pyvllm/v1/kv_offload/file_mapper.py
🚧 Files skipped from review as they are similar to previous changes (2)
- vllm/model_executor/layers/mla.py
- vllm/v1/attention/backends/mla/b12x_mla_sparse.py
Recognize the dynamic and static calibration variables so validated configurations do not emit unknown-vLLM-environment warnings. Assisted-by: OpenAI Codex
yatesdr
changed the title
Author
|
CI is currently blocked by repository policy rather than a test failure: |
26 tasks
lukealonso
merged commit Jul 29, 2026
8752e11
into
local-inference-lab:dev/gilded-gnosis
3 of 5 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Wire SparkInfer's per-token NVFP4 MLA scale mode
(local-inference-lab/sparkinfer#86) through vLLM and make the resulting cache
format safe for persistent DRAM/NVMe offload.
VLLM_NVFP4_MLA_DYNAMIC_SCALE=1selects the mode once per process. Bothwriter sites and both sparse-attention readers use the same immutable
configuration. Incompatible SparkInfer APIs, layouts, or simultaneous static
calibration fail during initialization.
Why
The 368-byte
nvfp4_ds_mlarecord uses E2M1 values and E4M3 group scales forthe compressed latent. With an implicit outer scale of 1.0, low-magnitude
layers place many group scales in E4M3's subnormal range. The reconstruction
error changes the hidden-state trajectory consumed by sparse indexing and
caused repeatable deep-context retrieval failures in the tested GLM-5.2
NF3/NVFP4 configuration.
SparkInfer #86 stores one per-token FP32 outer scale in existing record
padding and applies it during decode and MG-prefill dequantization. This PR
connects that mechanism to serving and prevents records with different scale
semantics from sharing a persistent cache namespace.
Changes
Immutable record mode
Nvfp4MlaCacheFormat, captured once at process import.VLLM_NVFP4_MLA_SCALES_FILE.KV_FP8_ROPE=1for the dynamic 368-byte layout.vllm.envs.Writer and reader wiring
per_token_scale=Truethrough:latent_scale_per_token=Truethrough decode and extend.per_token_scale, either readerlacks
latent_scale_per_token, or the record is not the 368-byte FP8-RoPENVFP4 layout.
SparkInfer builds.
Persistent cache ABI
The record bytes do not contain a version marker. The external cache identity
therefore includes:
All writer/readers import the same frozen format object. Dynamic and static
records cannot share a DRAM/NVMe namespace.
Backward compatibility is explicit: non-NVFP4 dtypes and implicit/default
NVFP4 configurations retain the literal
vllm-default-v1sentinel and theirexisting
FileMappernamespace. Only dynamic or calibrated modes opt into thenew ABI fields.
Validation
Focused vLLM tests
67 passed on the integrated CUDA 13.2 image:
bfloat16, auto, and implicit NVFP4;
Pinned Ruff check and format pass across all 11 touched files.
SparkInfer GPU integration
28 production-path tests passed on 4× SM120, covering static/dynamic records,
zero-token handling, multisplit decode, and multitile MG prefill.
End-to-end quality
KLD versus BF16 reference (n=3, lower is better):
Performance
Matched MTP0 decode, 10 fixed 1,024-token samples per arm:
Observed delta: -0.175 tok/s (-0.380%), inside the pre-committed 1% reader
overhead limit. MTP was disabled, so the result is not confounded by draft
acceptance.
Full retrieval, KLD, prefill, and capacity tables are in SparkInfer #86.
Reproduction
Persistent offload automatically derives a distinct dynamic record namespace.
No manual cache-path suffix is required after this PR.
Behavioral evaluation image:
ghcr.io/yatesdr/glm52-serve@sha256:db82fdcb5756d4a547853ba1330538bdd8a3dc0c6443c29bc49ba77b69b51cd1.The later PR-head commits add cache-contract hardening, tests, and env
registration without changing the measured dynamic arithmetic.
Limitations
PCIe system.
Implementation and text were drafted with AI assistance. All changed lines
were independently reviewed, and measurements were executed and validated by
the submitting operators. This is not a duplicate of #145: #145 supplies
checkpoint-specific static calibration; this PR wires calibration-free
per-token scaling and its cache-format contract.