iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
b89e85e99f7ae832877effa4eab43adebcca4813
vLLM-4x-RTX-3090/patches/0002-qwen4_exp-FP8-E4M3-KV-reader-for-QSA-with-a-static-s.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

2002 lines
83 KiB
Diff
Raw Blame History

From 7de1b4ee5f23d9ac5698b27c744f100577080e41 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Tue, 22 Sep 2026 04:09:54 -0400
Subject: [PATCH 002/103] qwen4_exp: FP8 E4M3 KV reader for QSA with a static
scale sidecar
FP8 E4M3 KV read path for the Qwen4Exp sparse attention (QSA), indexer and
pre-indexer kernels on sm_86, dequantised with a static per-layer scale loaded
from a JSON sidecar, validated against the model's QSA layer set and applied
per pipeline stage. Includes the offline absmax collector that produces the
sidecar and a CUDA-graph-safe clipping counter.
---
vllm/models/qwen4_exp/nvidia/model.py | 113 ++-
vllm/models/qwen4_exp/nvidia/ops/qsa.py | 213 +++++-
vllm/models/qwen4_exp/nvidia/ple_layer.py | 882 ++++++++++++++++++++++
vllm/models/qwen4_exp/nvidia/qsa.py | 487 +++++++++++-
4 files changed, 1657 insertions(+), 38 deletions(-)
diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py
index 54ee6f0..3f884cd 100644
--- a/vllm/models/qwen4_exp/nvidia/model.py
+++ b/vllm/models/qwen4_exp/nvidia/model.py
@@ -5,6 +5,8 @@
from collections.abc import Iterable
from itertools import islice
+import os
+
import torch
from torch import nn
@@ -79,7 +81,10 @@ from ..config import Qwen4ExpConfig
from .hyperconnection import GatedResidual, HyperConnectionConfig
from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
from .ple_layer import Qwen4ExpPLELayer
-from .qsa import Qwen4ExpQSAAttention
+from .qsa import Qwen4ExpQSAAttention, load_qsa_static_kv_scales
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
def without_modelopt_fp4(
@@ -158,6 +163,104 @@ _EXTRA_WEIGHTS_MAPPER = WeightsMapper(
)
+_QSA_KV_SCALES_ENV = "VLLM_QSA_KV_SCALES"
+
+
+def _maybe_load_qsa_static_kv_scales(model: nn.Module) -> None:
+ """Apply a static QSA K/V scale sidecar when VLLM_QSA_KV_SCALES names one.
+
+ Placed where the loader's own docstring requires: after checkpoint weights
+ are loaded, before profiling, warmup or cudagraph capture. Unset env var ->
+ no-op (the attention layer logs the default 1.0 per layer so the
+ uncalibrated state stays visible). Set but unusable -> the exception
+ propagates: a calibration that was requested and silently not applied is
+ worse than one never requested.
+ """
+ path = os.environ.get(_QSA_KV_SCALES_ENV, "").strip()
+ if not path:
+ return
+ # The PLE offload worker builds the whole target on the meta device and loads
+ # only PLE weights through load_weights; a fill_ on its meta scale buffers
+ # would log success without calibrating any GPU worker.
+ local_qsa: list[Qwen4ExpQSAAttention] = []
+ inner: Qwen4ExpModel | None = None
+ for module in model.modules():
+ if isinstance(module, Qwen4ExpModel) and inner is None:
+ inner = module
+ if isinstance(module, Qwen4ExpQSAAttention):
+ local_qsa.append(module)
+ for layer in local_qsa:
+ k_scale = getattr(layer, "_k_scale", None)
+ if isinstance(k_scale, torch.Tensor) and k_scale.device.type == "meta":
+ logger.info(
+ "Qwen4Exp QSA: model is on meta (PLE offload worker); "
+ "static K/V scales are applied by the GPU workers only"
+ )
+ return
+ if inner is None:
+ raise ValueError("Qwen4Exp QSA: no Qwen4ExpModel found for the sidecar")
+ if not inner._qsa_layer_ids:
+ # A requested calibration for a model without QSA layers is a
+ # configuration error, never a silent no-op.
+ raise ValueError(
+ f"Qwen4Exp QSA: {_QSA_KV_SCALES_ENV} is set but the model has no "
+ "QSA (indexer) attention layers"
+ )
+ # Whole-model QSA layer set (all pipeline stages) = "<stack>.<idx>.self_attn.attn"
+ # for every configured QSA index. The stack prefix is read from a LOCAL QSA
+ # layer's registered layer_name whenever one exists: this hook also runs
+ # inside the nested Qwen4ExpForCausalLM.load_weights, where a module path
+ # relative to that submodule ("model.layers") would not match the names
+ # registered under the outer prefix ("language_model.model.layers") - the
+ # earlier boot failure. Only a stage with no QSA layer
+ # falls back to the module path (then no local name can disagree).
+ if local_qsa:
+ stack, sep, tail = local_qsa[0].layer_name.rpartition(".layers.")
+ if not sep or not tail.split(".")[0].isdigit():
+ raise ValueError(
+ f"Qwen4Exp QSA: unexpected layer name {local_qsa[0].layer_name!r}"
+ )
+ layers_path = f"{stack}.layers"
+ else:
+ layers_path = next(
+ (name for name, module in model.named_modules() if module is inner.layers),
+ None,
+ )
+ if layers_path is None:
+ raise ValueError("Qwen4Exp QSA: cannot locate the decoder layer stack")
+ expected = {
+ f"{layers_path}.{idx}.self_attn.attn" for idx in inner._qsa_layer_ids
+ }
+ if not local_qsa:
+ # A pipeline stage without a full-attention layer (only for a very
+ # uneven split): validate the complete sidecar here too, apply nothing.
+ logger.warning(
+ "Qwen4Exp QSA: this pipeline stage holds no QSA layer "
+ "(pipeline layers %d-%d); sidecar %s validated, nothing to apply",
+ inner.start_layer,
+ inner.end_layer,
+ path,
+ )
+ # strict: every QSA layer of the whole model must be in the sidecar. An empty
+ # or partial file silently leaving layers at scale 1.0 is exactly the failure
+ # this call site exists to prevent.
+ applied = load_qsa_static_kv_scales(
+ model, path, strict=True, expected_layer_names=expected
+ )
+ shown = ", ".join(sorted(applied)[:4])
+ logger.info(
+ "Qwen4Exp QSA: applied static K/V scales from %s to %d of %d layer(s) "
+ "(pipeline layers %d-%d): %s%s",
+ path,
+ len(applied),
+ len(expected),
+ inner.start_layer,
+ inner.end_layer,
+ shown,
+ " ..." if len(applied) > 4 else "",
+ )
+
+
class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock):
"""Qwen3Next MoE with Qwen4Exp HC validation."""
@@ -855,7 +958,9 @@ class Qwen4ExpForCausalLM(
self,
ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
)
- return loader.load_weights(weights, mapper=mapper)
+ loaded = loader.load_weights(weights, mapper=mapper)
+ _maybe_load_qsa_static_kv_scales(self)
+ return loaded
class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo):
@@ -1050,7 +1155,9 @@ class Qwen4ExpForConditionalGeneration(
self,
ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
)
- return loader.load_weights(weights, mapper=mapper)
+ loaded = loader.load_weights(weights, mapper=mapper)
+ _maybe_load_qsa_static_kv_scales(self)
+ return loaded
@classmethod
def get_mamba_state_dtype_from_config(
diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
index d3986ba..2acc6f4 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
@@ -10,14 +10,38 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
TritonWarmupTensor,
triton_scalar_specialization_rep,
)
+from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON, tl, triton
+@triton.jit
+def _decode_e4m3_to_bf16(b):
+ # Shifting the 7 magnitude bits left by 7 lands the E4M3 exponent in the
+ # FP16 exponent field and the 3 mantissa bits at the top of the FP16
+ # mantissa field, giving the right number under the wrong exponent bias
+ # (E4M3 7 vs FP16 15). That is a constant factor of 2^8. E4M3 subnormals
+ # land on FP16 subnormals and the same factor renormalizes them, which is
+ # only true because sm_86 does not flush FP16 subnormals -- verified on
+ # device, not assumed.
+ #
+ # Bit-exact against torch on all 254 finite codepoints: every E4M3 value
+ # fits in BF16's 7 mantissa bits, and x256 is exact in FP16.
+ #
+ # Deviation: 0x7F/0xFF are NaN in e4m3fn but decode to +-480 here. The CUDA
+ # writer clamps to +-448 and never emits them, so this is unreachable from
+ # a cache written by reshape_and_cache_flash.
+ b = b.to(tl.uint16)
+ bits = ((b & 0x80) << 8) | ((b & 0x7F) << 7)
+ return (bits.to(tl.uint16).to(tl.float16, bitcast=True) * 256.0).to(tl.bfloat16)
+
+
@triton.jit(do_not_specialize=["num_rows", "num_requests"])
def _qsa_sparse_paged_gqa_splitk_kernel(
q_ptr,
k_cache_ptr,
v_cache_ptr,
+ k_scale_ptr,
+ v_scale_ptr,
indices_ptr,
block_table_ptr,
token_to_req_ptr,
@@ -50,6 +74,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
NUM_QUERY_HEADS: tl.constexpr,
NUM_SPLITS: tl.constexpr,
NUM_TILES: tl.constexpr,
+ KV_DTYPE: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
) -> None:
@@ -83,6 +108,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)
softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634
+ if KV_DTYPE != 0:
+ k_scale_value = tl.load(k_scale_ptr)
+ v_scale_value = tl.load(v_scale_ptr)
+
tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N))
for tile in range(split_id, tile_end, NUM_SPLITS):
@@ -111,33 +140,87 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
valid &= (physical_page >= 0) & (physical_page < num_cache_blocks)
# physical_page * block stride can overflow int32 for large caches.
safe_page = tl.maximum(physical_page, 0).to(tl.int64)
- keys = tl.load(
- k_cache_ptr
- + safe_page[None, :] * stride_k_block
- + page_offset[None, :] * stride_k_token
- + kv_head * stride_k_head
- + dim_offsets[:, None],
- mask=valid[None, :],
- other=0.0,
- )
- values = tl.load(
- v_cache_ptr
- + safe_page[:, None] * stride_v_block
- + page_offset[:, None] * stride_v_token
- + kv_head * stride_v_head
- + dim_offsets[None, :],
- mask=valid[:, None],
- other=0.0,
- )
+ if KV_DTYPE == 0:
+ keys = tl.load(
+ k_cache_ptr
+ + safe_page[None, :] * stride_k_block
+ + page_offset[None, :] * stride_k_token
+ + kv_head * stride_k_head
+ + dim_offsets[:, None],
+ mask=valid[None, :],
+ other=0.0,
+ )
+ values = tl.load(
+ v_cache_ptr
+ + safe_page[:, None] * stride_v_block
+ + page_offset[:, None] * stride_v_token
+ + kv_head * stride_v_head
+ + dim_offsets[None, :],
+ mask=valid[:, None],
+ other=0.0,
+ )
+ elif KV_DTYPE == 1:
+ key_bytes = tl.load(
+ k_cache_ptr
+ + safe_page[None, :] * stride_k_block
+ + page_offset[None, :] * stride_k_token
+ + kv_head * stride_k_head
+ + dim_offsets[:, None],
+ mask=valid[None, :],
+ other=0,
+ )
+ # UNSCALED: k_scale is folded into the score scaling below.
+ keys = _decode_e4m3_to_bf16(key_bytes)
+ else:
+ keys = tl.load(
+ k_cache_ptr
+ + safe_page[None, :] * stride_k_block
+ + page_offset[None, :] * stride_k_token
+ + kv_head * stride_k_head
+ + dim_offsets[:, None],
+ mask=valid[None, :],
+ other=0.0,
+ ).to(tl.bfloat16)
scores = tl.dot(query, keys)
# Scaling scores avoids re-quantizing a scaled query to BF16.
- scores *= softmax_scale_log2
+ if KV_DTYPE == 0:
+ scores *= softmax_scale_log2
+ else:
+ # k_scale is a scalar, so scaling each SCORE is equivalent to
+ # scaling every key element -- but costs one multiply per score
+ # instead of one per element, and skips the fp32 round trip and
+ # the BF16 re-quantization of the scaled key (which for a
+ # non-power-of-two scale also lost mantissa bits). Strictly fewer
+ # roundings than the per-element form it replaces.
+ scores *= softmax_scale_log2 * k_scale_value
scores = tl.where(valid[None, :], scores, -1.0e20)
next_max = tl.maximum(max_value, tl.max(scores, axis=1))
alpha = tl.math.exp2(max_value - next_max)
probabilities = tl.where(
valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0
)
+ if KV_DTYPE == 1:
+ value_bytes = tl.load(
+ v_cache_ptr
+ + safe_page[:, None] * stride_v_block
+ + page_offset[:, None] * stride_v_token
+ + kv_head * stride_v_head
+ + dim_offsets[None, :],
+ mask=valid[:, None],
+ other=0,
+ )
+ # UNSCALED: v_scale is folded into the epilogue below.
+ values = _decode_e4m3_to_bf16(value_bytes)
+ elif KV_DTYPE == 2:
+ values = tl.load(
+ v_cache_ptr
+ + safe_page[:, None] * stride_v_block
+ + page_offset[:, None] * stride_v_token
+ + kv_head * stride_v_head
+ + dim_offsets[None, :],
+ mask=valid[:, None],
+ other=0.0,
+ ).to(tl.bfloat16)
accumulator = tl.dot(
probabilities.to(values.dtype),
values,
@@ -152,6 +235,13 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
accumulator / tl.maximum(normalizer[:, None], 1.0e-20),
0.0,
)
+ if KV_DTYPE != 0:
+ # v_scale folded out of the per-element V decode: the attention output
+ # is linear in V, so one multiply per output element replaces one per
+ # cache element per tile. Safe under split-K too -- every partial is
+ # scaled by the same constant before the LSE merge, which is itself
+ # linear in the partial outputs.
+ normalized_output *= v_scale_value
output_mask = head_offsets[:, None] < GROUP_SIZE
if NUM_SPLITS == 1:
# Preserve the unfused path's BF16 attention-output rounding before
@@ -483,10 +573,18 @@ def qsa_sparse_paged_attention(
token_to_req: torch.Tensor,
use_prefill_config: bool,
out: torch.Tensor | None = None,
+ k_scale: torch.Tensor | None = None,
+ v_scale: torch.Tensor | None = None,
+ _force_num_warps: int | None = None,
*,
output_gate: torch.Tensor,
) -> torch.Tensor:
- """Run sparse GQA directly over paged BF16 K/V caches.
+ """Run sparse GQA over paged BF16 or FP8 (E4M3 via static scalar scales) K/V caches.
+
+ ``_force_num_warps`` overrides the tuned launch profile for the warp-count
+ equivalence test only; production callers leave it None. Hardware policy
+ (E4M3 only on CUDA sm_86) is enforced once at Qwen4ExpQSAAttention /
+ Qwen4ExpQSAFlashAttentionImpl construction, not per call.
logical_indices is the PACKED selection buffer: [rows, selection_width + 1]
with the trailing column holding each row's valid-entry count (written by
@@ -510,7 +608,26 @@ def qsa_sparse_paged_attention(
raise ValueError("QSA sparse attention requires valid grouped-query heads")
head_dim = q.shape[2]
assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0
- assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16
+ if q.dtype != torch.bfloat16:
+ raise ValueError("QSA sparse attention requires BF16 queries")
+ if k_cache.dtype != v_cache.dtype:
+ raise ValueError("QSA sparse attention requires matching K/V dtypes")
+ fp8_dtype = current_platform.fp8_dtype()
+ if k_cache.dtype == torch.bfloat16:
+ kv_dtype = 0
+ elif k_cache.dtype in (torch.uint8, fp8_dtype):
+ kv_dtype = 1
+ # Present platform E4M3 storage as bytes so Triton never materializes
+ # the unsupported fp8e4nv type on SM86.
+ if k_cache.dtype == fp8_dtype:
+ k_cache = k_cache.view(torch.uint8)
+ v_cache = v_cache.view(torch.uint8)
+ elif k_cache.dtype == torch.float8_e5m2:
+ kv_dtype = 2
+ else:
+ raise ValueError(
+ "QSA supports only BF16, uint8/E4M3, and E5M2 K/V caches"
+ )
assert logical_indices.dtype == block_table.dtype == torch.int32
assert token_to_req.dtype == torch.int32
assert q.device == k_cache.device == v_cache.device
@@ -531,12 +648,37 @@ def qsa_sparse_paged_attention(
if not q.shape[0]:
return out
+ if kv_dtype == 0:
+ # Scale pointers are compile-time dead in the unchanged BF16 kernel.
+ k_scale_ptr = q
+ v_scale_ptr = q
+ else:
+ if k_scale is None:
+ k_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device)
+ if v_scale is None:
+ v_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device)
+ if (
+ k_scale.numel() != 1
+ or v_scale.numel() != 1
+ or k_scale.dtype != torch.float32
+ or v_scale.dtype != torch.float32
+ or k_scale.device != q.device
+ or v_scale.device != q.device
+ ):
+ raise ValueError(
+ "QSA K/V scales must be scalar float32 tensors on the query device"
+ )
+ k_scale_ptr = k_scale
+ v_scale_ptr = v_scale
+
group_size = q.shape[1] // k_cache.shape[2]
block_m = triton.next_power_of_2(group_size)
selection_width = logical_indices.shape[1] - 1 # trailing column is the count
block_n, partial_warps, num_tiles, num_splits = _select_config(
q.shape[0], k_cache.shape[2], use_prefill_config, selection_width
)
+ if _force_num_warps is not None:
+ partial_warps = _force_num_warps
# Split=1 writes output directly and compiles out all workspace accesses.
if num_splits == 1:
@@ -559,6 +701,8 @@ def qsa_sparse_paged_attention(
q,
k_cache,
v_cache,
+ k_scale_ptr,
+ v_scale_ptr,
logical_indices,
block_table,
token_to_req,
@@ -591,6 +735,7 @@ def qsa_sparse_paged_attention(
NUM_QUERY_HEADS=q.shape[1],
NUM_SPLITS=num_splits,
NUM_TILES=num_tiles,
+ KV_DTYPE=kv_dtype,
BLOCK_M=block_m,
BLOCK_N=block_n,
num_warps=partial_warps,
@@ -649,13 +794,32 @@ def warmup_qsa_sparse_paged_attention(
q_ptr = TritonWarmupTensor(
torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)
)
+ # Mirror qsa_sparse_paged_attention's dtype dispatch: E4M3 storage is launched
+ # as a uint8 view with scalar float32 scale pointers.
+ if key_cache.dtype == torch.bfloat16:
+ warm_kv_dtype = 0
+ warm_cache_dtype = key_cache.dtype
+ elif key_cache.dtype in (torch.uint8, current_platform.fp8_dtype()):
+ warm_kv_dtype = 1
+ warm_cache_dtype = torch.uint8
+ elif key_cache.dtype == torch.float8_e5m2:
+ warm_kv_dtype = 2
+ warm_cache_dtype = key_cache.dtype
+ else:
+ raise ValueError("QSA warmup: unsupported K/V cache dtype")
+ if warm_kv_dtype == 0:
+ k_scale_ptr = q_ptr
+ v_scale_ptr = q_ptr
+ else:
+ k_scale_ptr = TritonWarmupTensor(torch.float32)
+ v_scale_ptr = TritonWarmupTensor(torch.float32)
k_cache_ptr = TritonWarmupTensor(
- key_cache.dtype,
+ warm_cache_dtype,
shape=tuple(key_cache.shape),
strides=tuple(key_cache.stride()),
)
v_cache_ptr = TritonWarmupTensor(
- value_cache.dtype,
+ warm_cache_dtype,
shape=tuple(value_cache.shape),
strides=tuple(value_cache.stride()),
)
@@ -695,6 +859,8 @@ def warmup_qsa_sparse_paged_attention(
q_ptr,
k_cache_ptr,
v_cache_ptr,
+ k_scale_ptr,
+ v_scale_ptr,
indices_ptr,
block_table_ptr,
token_to_req_ptr,
@@ -727,6 +893,7 @@ def warmup_qsa_sparse_paged_attention(
NUM_QUERY_HEADS=num_query_heads,
NUM_SPLITS=num_splits,
NUM_TILES=num_tiles,
+ KV_DTYPE=warm_kv_dtype,
BLOCK_M=block_m,
BLOCK_N=block_n,
num_warps=warps,
diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py
index 2ec8790..bec35e6 100644
--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py
+++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py
@@ -4,6 +4,8 @@
from collections.abc import Sequence
+import os
+
import torch
from torch import nn
@@ -63,6 +65,886 @@ class Qwen4ExpPLEGroupedNorm(nn.Module):
return (normalized * (1.0 + self.weight.float())).to(input_dtype)
+class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase):
+ """FP8 PLE embedding with one global checkpoint scale."""
+
+ def create_weights(
+ self,
+ layer: nn.Module,
+ input_size_per_partition: int,
+ output_partition_sizes: list[int],
+ input_size: int,
+ output_size: int,
+ params_dtype: torch.dtype,
+ **extra_weight_attrs,
+ ) -> None:
+ del input_size, output_size, params_dtype
+ weight_loader = extra_weight_attrs.get("weight_loader")
+ weight = create_fp8_weight_parameter(
+ sum(output_partition_sizes), input_size_per_partition, weight_loader
+ )
+ layer.register_parameter("weight", weight)
+
+ weight_scale = create_fp8_scale_parameter(
+ PerTensorScaleParameter,
+ output_partition_sizes,
+ input_size_per_partition,
+ None,
+ weight_loader,
+ scale_dtype=torch.float32,
+ )
+ layer.register_parameter("weight_scale", weight_scale)
+
+ def process_weights_after_loading(self, layer: nn.Module) -> None:
+ """Reject FP8 PLE checkpoints without a global scale."""
+ sentinel = torch.finfo(torch.float32).min
+ if torch.any(layer.weight_scale == sentinel):
+ raise ValueError("FP8 PLE checkpoint is missing its global scale")
+
+ def apply(
+ self,
+ layer: nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ raise NotImplementedError("PLE FP8 weights only support embedding lookup")
+
+ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
+ return F.embedding(input_, layer.weight)
+
+
+_NVFP4_BLOCK_SIZE = 16
+
+
+class Qwen4ExpPLENVFp4EmbeddingMethod(QuantizeMethodBase):
+ """NVFP4 PLE embedding kept packed at runtime.
+
+ The table stays as uint8 codes with FP8 block scales (16 values per
+ block) and one FP32 global scale. A lookup only gathers the packed rows;
+ the PLE layer dequantizes them after lookup.
+ """
+
+ def create_weights(
+ self,
+ layer: nn.Module,
+ input_size_per_partition: int,
+ output_partition_sizes: list[int],
+ input_size: int,
+ output_size: int,
+ params_dtype: torch.dtype,
+ **extra_weight_attrs,
+ ) -> None:
+ del input_size, output_size, extra_weight_attrs
+ if input_size_per_partition % _NVFP4_BLOCK_SIZE:
+ raise ValueError(
+ "NVFP4 PLE embedding requires the embedding dim to be a "
+ f"multiple of {_NVFP4_BLOCK_SIZE}, got {input_size_per_partition}"
+ )
+ self.params_dtype = params_dtype
+ self.head_dim = input_size_per_partition
+ self.packed_row_width = (
+ input_size_per_partition // 2
+ + input_size_per_partition // _NVFP4_BLOCK_SIZE
+ )
+ rows = sum(output_partition_sizes)
+ weight = nn.Parameter(
+ torch.empty(rows, input_size_per_partition // 2, dtype=torch.uint8),
+ requires_grad=False,
+ )
+ set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+ layer.register_parameter("weight", weight)
+
+ weight_scale = nn.Parameter(
+ torch.empty(
+ rows,
+ input_size_per_partition // _NVFP4_BLOCK_SIZE,
+ dtype=torch.float8_e4m3fn,
+ ),
+ requires_grad=False,
+ )
+ set_weight_attrs(weight_scale, {"input_dim": 1, "output_dim": 0})
+ layer.register_parameter("weight_scale", weight_scale)
+
+ weight_scale_2 = nn.Parameter(
+ torch.empty((), dtype=torch.float32), requires_grad=False
+ )
+ layer.register_parameter("weight_scale_2", weight_scale_2)
+ self._lut: torch.Tensor | None = None
+
+ def process_weights_after_loading(self, layer: nn.Module) -> None:
+ # Build the lookup table before CUDA graph capture.
+ self._lut = torch.tensor(
+ _FP4_VALUES, dtype=torch.float32, device=layer.weight.device
+ )
+
+ def apply(
+ self,
+ layer: nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ raise NotImplementedError("PLE NVFP4 weights only support embedding lookup")
+
+ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
+ """Gather packed rows for deferred dequantization."""
+ codes = F.embedding(input_, layer.weight)
+ scales = F.embedding(input_, layer.weight_scale)
+ return torch.cat((codes, scales.view(torch.uint8)), dim=-1)
+
+ def dequantize(
+ self,
+ packed_rows: torch.Tensor,
+ scale_2: torch.Tensor,
+ output_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """Dequantize rows containing NVFP4 codes and block scales."""
+ lut = self._lut
+ if lut is None or lut.device != packed_rows.device:
+ lut = torch.tensor(
+ _FP4_VALUES, dtype=torch.float32, device=packed_rows.device
+ )
+ self._lut = lut
+ return _dequant_nvfp4_rows(
+ packed_rows, self.head_dim, scale_2, output_dtype, lut
+ )
+
+
+def _ple_dtype_is_fp8(ple_embedding_dtype: object) -> bool:
+ """Return whether the model config declares FP8 PLE checkpoint weights."""
+
+ if ple_embedding_dtype is None:
+ return False
+ if isinstance(ple_embedding_dtype, torch.dtype):
+ return ple_embedding_dtype == torch.float8_e4m3fn
+ return str(ple_embedding_dtype).rsplit(".", 1)[-1] == "float8_e4m3fn"
+
+
+def _get_ple_embedding_quant_method(
+ quant_config: QuantizationConfig | None,
+ prefix: str,
+ ple_embedding_dtype: object = None,
+) -> QuantizeMethodBase | None:
+ """Select a packed PLE embedding method for quantized checkpoint shards."""
+
+ if isinstance(quant_config, ModelOptMixedPrecisionConfig):
+ if quant_config._resolve_quant_algo(prefix) == "FP8":
+ return Qwen4ExpPLEFp8EmbeddingMethod()
+ return None
+
+ if isinstance(quant_config, Fp8Config):
+ if not quant_config.is_checkpoint_fp8_serialized:
+ return None
+ ignored_layers = quant_config.ignored_layers
+ if is_layer_skipped(
+ prefix,
+ ignored_layers,
+ quant_config.packed_modules_mapping,
+ match_mode=quant_config.ignored_layers_match_mode,
+ ):
+ return None
+ # PLE checkpoint shards form one runtime embedding parameter.
+ shard_prefix = f"{prefix}.shard_"
+ if any(name.startswith(shard_prefix) for name in ignored_layers):
+ return None
+ return Qwen4ExpPLEFp8EmbeddingMethod()
+
+ if isinstance(quant_config, ModelOptNvFp4Config):
+ if not quant_config.is_checkpoint_nvfp4_serialized:
+ return None
+ if not quant_config.is_layer_excluded(prefix):
+ logger.info_once("PLE embedding %s uses the runtime NVFP4 method", prefix)
+ return Qwen4ExpPLENVFp4EmbeddingMethod()
+ if _ple_dtype_is_fp8(ple_embedding_dtype):
+ logger.info_once(
+ "Excluded ModelOpt PLE embedding %s uses the runtime FP8 method",
+ prefix,
+ )
+ return Qwen4ExpPLEFp8EmbeddingMethod()
+
+ # Checkpoints pairing quantized experts with a global-scale FP8 PLE table
+ # (the RadixArk Flash-Next table grafted onto a CompressedTensors checkpoint)
+ # list the PLE under exclude_modules, so the FP8-ness is visible only in the
+ # tensors' dtype / the config's ple_embedding_dtype, not in the quant config
+ # type dispatched above. Without this branch the offload worker treats the
+ # table as BF16 and materializes ~96 GiB in host RAM (OOM-killed on the
+ # 96 GB host). Honour either the config key (the -e1
+ # checkpoint copy) or the explicit VLLM_PLE_FP8_GLOBAL_SCALE=1 opt-in the
+ # previous release uses. Placement after the typed dispatch is the
+ # position already proven in service.
+ if _ple_dtype_is_fp8(ple_embedding_dtype) or os.environ.get(
+ "VLLM_PLE_FP8_GLOBAL_SCALE"
+ ) == "1":
+ logger.info_once(
+ "PLE embedding %s: config declares an FP8 table with quant_config %s; "
+ "using the runtime FP8 method (global weight_scale)",
+ prefix,
+ type(quant_config).__name__,
+ )
+ return Qwen4ExpPLEFp8EmbeddingMethod()
+
+ return None
+
+
+_FP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+
+
+def _dequant_nvfp4_codes(
+ packed: torch.Tensor,
+ scale: torch.Tensor,
+ scale_2: torch.Tensor,
+ lut: torch.Tensor | None = None,
+) -> torch.Tensor:
+ """Unpack NVFP4 codes with FP8 block scales and an FP32 global scale."""
+ half = packed.shape[-1]
+ codes = torch.stack([packed & 0xF, (packed >> 4) & 0xF], dim=-1).reshape(
+ *packed.shape[:-1], half * 2
+ )
+ if lut is None:
+ lut = torch.tensor(_FP4_VALUES, dtype=torch.float32, device=packed.device)
+ magnitude = lut[(codes & 0x7).long()]
+ sign = ((codes >> 3) & 1).to(torch.float32)
+ fp4 = (magnitude * (1 - 2 * sign)).reshape(
+ *packed.shape[:-1], -1, _NVFP4_BLOCK_SIZE
+ )
+ output = fp4 * scale.float().unsqueeze(-1) * scale_2.float()
+ return output.reshape(*packed.shape[:-1], half * 2)
+
+
+def _dequant_nvfp4_rows(
+ packed_rows: torch.Tensor,
+ head_dim: int,
+ scale_2: torch.Tensor,
+ output_dtype: torch.dtype,
+ lut: torch.Tensor,
+) -> torch.Tensor:
+ """Dequantize rows containing NVFP4 codes followed by block scales."""
+ half = head_dim // 2
+ codes = packed_rows[..., :half]
+ scales = packed_rows[..., half:].contiguous().view(torch.float8_e4m3fn)
+ return _dequant_nvfp4_codes(codes, scales, scale_2, lut).to(output_dtype)
+
+
+def _get_shared_nvfp4_outer_scale(
+ outer_scales: dict[int, torch.Tensor],
+) -> torch.Tensor:
+ """Validate and return the global scale shared by NVFP4 PLE shards."""
+ first_index, reference = next(iter(outer_scales.items()))
+ reference = reference.reshape(())
+ for shard_index, outer_scale in outer_scales.items():
+ if not torch.equal(outer_scale.reshape(()), reference):
+ raise ValueError(
+ "NVFP4 PLE shards must share the same global scale, but "
+ f"shards {first_index} and {shard_index} differ"
+ )
+ return reference
+
+
+class Qwen4ExpNGramEmbedding(PleOffloadLayer):
+ _offload_quant_method: QuantizeMethodBase | None = None
+ _MASK64 = (1 << 64) - 1
+ _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
+ _SPLITMIX_M1 = 0xBF58476D1CE4E5B9
+ _SPLITMIX_M2 = 0x94D049BB133111EB
+ _PLE_LAYER_PRIME = 10007
+
+ @classmethod
+ def _splitmix64(cls, value: int) -> int:
+ """Mix an integer into a deterministic unsigned 64-bit value."""
+ value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64
+ value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64
+ value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64
+ return (value ^ (value >> 31)) & cls._MASK64
+
+ @staticmethod
+ def _is_prime_64(value: int) -> bool:
+ """Return whether a 64-bit integer is prime."""
+ if value < 2:
+ return False
+ for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
+ if value % prime == 0:
+ return value == prime
+ exponent = value - 1
+ shifts = 0
+ while exponent % 2 == 0:
+ exponent //= 2
+ shifts += 1
+ for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022):
+ if base % value == 0:
+ continue
+ witness = pow(base, exponent, value)
+ if witness in (1, value - 1):
+ continue
+ for _ in range(shifts - 1):
+ witness = pow(witness, 2, value)
+ if witness == value - 1:
+ break
+ else:
+ return False
+ return True
+
+ @classmethod
+ def _nth_prime_after(cls, start: int, count: int) -> int:
+ """Return the ``count``-th prime strictly greater than ``start``."""
+ prime = int(start)
+ for _ in range(count):
+ candidate = prime + 1
+ if candidate <= 2:
+ prime = 2
+ continue
+ if candidate % 2 == 0:
+ candidate += 1
+ while not cls._is_prime_64(candidate):
+ candidate += 2
+ prime = candidate
+ return prime
+
+ @classmethod
+ def _make_layer_multipliers(
+ cls,
+ *,
+ ngram_size: int,
+ unigram_vocab_size: int,
+ seed: int,
+ ple_dense_layer_id: int,
+ ) -> list[int]:
+ """Build deterministic hash multipliers for one PLE layer."""
+ max_multiplier = ((1 << 63) - 1) // unigram_vocab_size
+ half_bound = max(1, max_multiplier // 2)
+ base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id
+ multipliers = []
+ for index in range(ngram_size):
+ value = base_seed + cls._SPLITMIX_GAMMA * (index + 1)
+ multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1)
+ return multipliers
+
+ @classmethod
+ def _make_vocab_layout(
+ cls,
+ *,
+ ngram_vocab_size_base: int,
+ ngram_heads: int,
+ ple_dense_layer_id: int,
+ ) -> tuple[list[int], list[int], int]:
+ """Build per-head vocabulary sizes, offsets, and total row count."""
+ sizes: list[int] = []
+ offsets: list[int] = []
+ offset = 0
+ for local_head in range(ngram_heads):
+ global_head = ple_dense_layer_id * ngram_heads + local_head
+ size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1)
+ sizes.append(size)
+ offsets.append(offset)
+ offset += size
+ return sizes, offsets, offset
+
+ def __init__(
+ self,
+ config: Qwen4ExpTextConfig,
+ embedding_dim: int,
+ ple_dense_layer_id: int,
+ prefix: str,
+ layer_name: str,
+ quant_config: QuantizationConfig | None = None,
+ params_dtype: torch.dtype | None = None,
+ ) -> None:
+ super().__init__()
+ self.layer_name = layer_name
+ self.embedding_dim = embedding_dim
+ self.ngram_size = int(config.ngram_size)
+ self.heads_per_ngram = int(config.heads_per_ngram)
+ self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
+ if self.ngram_size < 2:
+ raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}")
+ if self.heads_per_ngram <= 0:
+ raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}")
+ if embedding_dim % self.ngram_heads:
+ raise ValueError(
+ "ple_embed_dim must be divisible by total ngram heads: "
+ f"{embedding_dim} % {self.ngram_heads} != 0"
+ )
+ self.head_dim = embedding_dim // self.ngram_heads
+ self.eos_token_id = int(config.eos_token_id)
+ self.unigram_vocab_size = int(config.vocab_size)
+ self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512))
+ if self.split_ngram_parts <= 0:
+ raise ValueError("split_ngram_parts must be positive")
+
+ multipliers = self._make_layer_multipliers(
+ ngram_size=self.ngram_size,
+ unigram_vocab_size=self.unigram_vocab_size,
+ seed=int(getattr(config, "seed", 1234)),
+ ple_dense_layer_id=ple_dense_layer_id,
+ )
+ self.register_buffer(
+ "layer_multipliers",
+ torch.tensor(multipliers, dtype=torch.long),
+ persistent=True,
+ )
+
+ sizes, offsets, total_vocab_size = self._make_vocab_layout(
+ ngram_vocab_size_base=int(config.ngram_vocab_size_base),
+ ngram_heads=self.ngram_heads,
+ ple_dense_layer_id=ple_dense_layer_id,
+ )
+ self.register_buffer(
+ "ngram_heads_vocab_sizes",
+ torch.tensor(sizes, dtype=torch.long),
+ persistent=True,
+ )
+ self.register_buffer(
+ "ngram_heads_offsets",
+ torch.tensor(offsets, dtype=torch.long),
+ persistent=True,
+ )
+ divisor = int(config.make_ngram_vocab_size_divisible_by)
+ padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
+ self.ngram_embedding = PLEVocabParallelEmbedding(
+ padded_vocab_size,
+ self.head_dim,
+ params_dtype=params_dtype,
+ padding_size=divisor,
+ prefix=f"{prefix}.ngram_embedding",
+ quant_method=_get_ple_embedding_quant_method(
+ quant_config,
+ f"{prefix}.ngram_embedding",
+ getattr(config, "ple_embedding_dtype", None),
+ ),
+ )
+
+ @staticmethod
+ def _shift_precompute(
+ tokens: torch.Tensor, eos_token_id: int
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ if tokens.dim() != 2:
+ raise ValueError("tokens must be a 2D tensor")
+ batch_size, seq_len = tokens.shape
+ positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64)
+ eos_positions = torch.where(tokens == eos_token_id, positions, -1)
+ previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
+ previous_eos = torch.cat(
+ [
+ eos_positions.new_full((batch_size, 1), -1),
+ previous_eos_inclusive[:, :-1],
+ ],
+ dim=1,
+ )
+ return positions, positions.unsqueeze(0) - previous_eos - 1
+
+ @staticmethod
+ def _shift_apply(
+ tokens: torch.Tensor,
+ positions: torch.Tensor,
+ position_in_segment: torch.Tensor,
+ shift: int,
+ eos_token_id: int,
+ ) -> torch.Tensor:
+ if shift == 0:
+ return tokens
+ source = positions - shift
+ gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1)
+ shifted = tokens.gather(1, gather_indices)
+ valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift)
+ return torch.where(valid, shifted, tokens.new_full((), eos_token_id))
+
+ def compute_ngram_ids(
+ self,
+ input_ids: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ ngram_context: torch.Tensor,
+ output: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """Compute n-gram embedding indices for the current request layout."""
+ input_ids = input_ids.reshape(-1)
+ num_reqs = query_start_loc.numel() - 1
+ num_tokens = input_ids.shape[0]
+
+ if input_ids.is_cuda:
+ return ple_ngram_ids(
+ input_ids=input_ids,
+ query_start_loc=query_start_loc,
+ ngram_context=ngram_context,
+ layer_multipliers=self.layer_multipliers,
+ ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes,
+ ngram_heads_offsets=self.ngram_heads_offsets,
+ eos_token_id=self.eos_token_id,
+ heads_per_ngram=self.heads_per_ngram,
+ output=output,
+ )
+ input_ids = input_ids.long()
+ query_start_loc = query_start_loc.long()
+ if num_reqs <= 0:
+ raise ValueError("PLE requires at least one request")
+
+ if is_offload_process():
+ max_seq_len = max(
+ 1,
+ int((query_start_loc[1:] - query_start_loc[:-1]).max().item()),
+ )
+ # CUDA graphs can pad input_ids beyond the valid request layout.
+ # Keep those rows from overwriting the final real token.
+ num_valid_tokens = min(int(query_start_loc[-1].item()), num_tokens)
+ else:
+ max_seq_len = num_tokens
+ num_valid_tokens = num_tokens
+
+ positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64)
+ packed = torch.full(
+ (num_reqs, max_seq_len),
+ self.eos_token_id,
+ device=input_ids.device,
+ dtype=torch.int64,
+ )
+ request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
+ request_indices.clamp_(max=num_reqs - 1)
+ columns = (positions - query_start_loc[request_indices]).clamp(
+ 0, packed.shape[1] - 1
+ )
+ packed[request_indices[:num_valid_tokens], columns[:num_valid_tokens]] = (
+ input_ids[:num_valid_tokens]
+ )
+ ngram_context = ngram_context[:num_reqs].to(
+ device=input_ids.device, dtype=torch.long
+ )
+
+ context = torch.cat([ngram_context, packed], dim=-1)
+ positions_2d, position_in_segment = self._shift_precompute(
+ context, self.eos_token_id
+ )
+ shifted = [context]
+ for shift in range(1, self.ngram_size):
+ shifted.append(
+ self._shift_apply(
+ context,
+ positions_2d,
+ position_in_segment,
+ shift,
+ self.eos_token_id,
+ )
+ )
+ adjusted_columns = columns + self.ngram_size - 1
+ id_blocks = []
+ for ngram in range(2, self.ngram_size + 1):
+ start = (ngram - 2) * self.heads_per_ngram
+ end = start + self.heads_per_ngram
+ mixed = shifted[0] * self.layer_multipliers[0]
+ for index in range(1, ngram):
+ mixed = torch.bitwise_xor(
+ mixed, shifted[index] * self.layer_multipliers[index]
+ )
+ sizes = self.ngram_heads_vocab_sizes[start:end]
+ offsets = self.ngram_heads_offsets[start:end]
+ ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
+ id_blocks.append(ids[request_indices, adjusted_columns])
+ return torch.cat(id_blocks, dim=-1)
+
+ def forward_impl( # type: ignore[override]
+ self,
+ hidden_states: torch.Tensor,
+ input_ids: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ ngram_context: torch.Tensor,
+ output_buffer: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ del hidden_states
+ if input_ids.is_cuda:
+ # Keep request-dependent ID generation outside piecewise graphs.
+ ngram_ids = input_ids.new_empty(
+ (input_ids.numel(), self.ngram_heads), dtype=torch.long
+ )
+ torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids(
+ input_ids,
+ query_start_loc,
+ ngram_context,
+ ngram_ids,
+ self.layer_name,
+ )
+ else:
+ ngram_ids = self.compute_ngram_ids(
+ input_ids,
+ query_start_loc,
+ ngram_context,
+ )
+ num_tokens = input_ids.reshape(-1).shape[0]
+ if output_buffer is not None:
+ quant_method = getattr(self.ngram_embedding, "quant_method", None)
+ if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod):
+ output_dim = self.get_offload_output_dim(self.embedding_dim)
+ output = output_buffer[:num_tokens, :output_dim]
+ output.copy_(self.ngram_embedding(ngram_ids).flatten(-2))
+ return output
+ output = output_buffer[:num_tokens, : self.embedding_dim]
+ torch.index_select(
+ self.ngram_embedding.weight,
+ 0,
+ ngram_ids.reshape(-1),
+ out=output.reshape(-1, self.head_dim),
+ )
+ return output
+ return self.ngram_embedding(ngram_ids).flatten(-2)
+
+ def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype:
+ """Keep quantized lookup results in their embedding storage dtype."""
+ embedding = getattr(self, "ngram_embedding", None)
+ weight = getattr(embedding, "weight", None)
+ if weight is not None:
+ return weight.dtype
+ if isinstance(self._offload_quant_method, Qwen4ExpPLENVFp4EmbeddingMethod):
+ return torch.uint8
+ if isinstance(self._offload_quant_method, Qwen4ExpPLEFp8EmbeddingMethod):
+ return torch.float8_e4m3fn
+ return default_dtype
+
+ def get_offload_output_dim(self, default_dim: int) -> int:
+ """Keep NVFP4 lookup rows packed while transferring them to the GPU."""
+ quant_method = getattr(
+ getattr(self, "ngram_embedding", None), "quant_method", None
+ )
+ if quant_method is None:
+ quant_method = self._offload_quant_method
+ if isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod):
+ if default_dim % _NVFP4_BLOCK_SIZE:
+ raise ValueError(
+ "NVFP4 PLE output dim must be a multiple of "
+ f"{_NVFP4_BLOCK_SIZE}, got {default_dim}"
+ )
+ return default_dim // 2 + default_dim // _NVFP4_BLOCK_SIZE
+ return default_dim
+
+ def initialize_dummy_offload_metadata(self, device: torch.device) -> None:
+ """Initialize quantization metadata skipped by the dummy loader."""
+ quant_method = self._offload_quant_method
+ if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod):
+ self.register_buffer(
+ "_offload_weight_scale",
+ torch.ones((), dtype=torch.float32, device=device),
+ persistent=False,
+ )
+ elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod):
+ self.register_buffer(
+ "_offload_weight_scale_2",
+ torch.ones((), dtype=torch.float32, device=device),
+ persistent=False,
+ )
+ self.register_buffer(
+ "_offload_nvfp4_lut",
+ torch.tensor(_FP4_VALUES, dtype=torch.float32, device=device),
+ persistent=False,
+ )
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Load hash buffers and checkpoint-split embedding rows."""
+
+ # Prefix grouping may invoke this method repeatedly for one module.
+ # Treat the global FP8 scale as incrementally loaded state.
+ # GPU workers retain only dequantization metadata. The CPU process owns
+ # the embedding table and transfers quantized lookup rows unchanged.
+ if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process():
+ retained: set[str] = set()
+ quant_method = self._offload_quant_method
+ if isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod):
+ for name, loaded_weight in weights:
+ if name != "ngram_embedding.weight_scale":
+ continue
+ self.register_buffer(
+ "_offload_weight_scale",
+ loaded_weight.to(
+ device=torch.accelerator.current_accelerator()
+ ),
+ persistent=False,
+ )
+ retained.add(name)
+ elif isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod):
+ outer_scales: dict[int, torch.Tensor] = {}
+ for name, loaded_weight in weights:
+ if not name.startswith(
+ "ngram_embedding.shard_"
+ ) or not name.endswith(".weight_scale_2"):
+ continue
+ shard_text = name[
+ len("ngram_embedding.shard_") : -len(".weight_scale_2")
+ ]
+ shard_index = int(shard_text)
+ outer_scales[shard_index] = loaded_weight
+ retained.add(name)
+ if not retained:
+ raise ValueError(
+ "NVFP4 PLE offload checkpoint is missing its global scale"
+ )
+ scale_2 = _get_shared_nvfp4_outer_scale(outer_scales).to(
+ device=torch.accelerator.current_accelerator()
+ )
+ self.register_buffer(
+ "_offload_weight_scale_2", scale_2, persistent=False
+ )
+ self.register_buffer(
+ "_offload_nvfp4_lut",
+ torch.tensor(
+ _FP4_VALUES, dtype=torch.float32, device=scale_2.device
+ ),
+ persistent=False,
+ )
+ else:
+ for _ in weights:
+ pass
+ return retained
+
+ persistent_buffers = {
+ "layer_multipliers": self.layer_multipliers,
+ "ngram_heads_offsets": self.ngram_heads_offsets,
+ "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
+ }
+ loaded: set[str] = set()
+ regular_weights: list[tuple[str, torch.Tensor]] = []
+ shard_prefix = "ngram_embedding.shard_"
+ quant_method = getattr(self.ngram_embedding, "quant_method", None)
+ nvfp4_runtime = isinstance(quant_method, Qwen4ExpPLENVFp4EmbeddingMethod)
+ packed_codes: dict[int, torch.Tensor] = {}
+ packed_scales: dict[int, torch.Tensor] = {}
+ packed_outer_scales: dict[int, torch.Tensor] = {}
+
+ for name, loaded_weight in weights:
+ leaf_name = name.rsplit(".", 1)[-1]
+ if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
+ continue
+ if name in persistent_buffers:
+ buffer = persistent_buffers[name]
+ if buffer.shape != loaded_weight.shape:
+ raise ValueError(
+ f"Shape mismatch for {name}: expected "
+ f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
+ )
+ buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
+ loaded.add(name)
+ continue
+
+ if nvfp4_runtime and name.startswith(shard_prefix):
+ suffix = name[len(shard_prefix) :]
+ for ending, sink in (
+ (".weight_scale_2", packed_outer_scales),
+ (".weight_scale", packed_scales),
+ (".weight", packed_codes),
+ ):
+ if not suffix.endswith(ending):
+ continue
+ shard_text = suffix[: -len(ending)]
+ shard_index = int(shard_text)
+ if shard_index >= self.split_ngram_parts:
+ raise ValueError(
+ f"PLE embedding shard index {shard_index} exceeds "
+ f"split_ngram_parts={self.split_ngram_parts}"
+ )
+ sink[shard_index] = loaded_weight
+ break
+ else:
+ regular_weights.append((name, loaded_weight))
+ continue
+
+ if (
+ not nvfp4_runtime
+ and name.startswith(shard_prefix)
+ and name.endswith(".weight")
+ ):
+ shard_text = name[len(shard_prefix) : -len(".weight")]
+ if not shard_text.isdigit():
+ regular_weights.append((name, loaded_weight))
+ continue
+ shard_index = int(shard_text)
+ if shard_index >= self.split_ngram_parts:
+ raise ValueError(
+ f"PLE embedding shard index {shard_index} exceeds "
+ f"split_ngram_parts={self.split_ngram_parts}"
+ )
+ embedding = self.ngram_embedding
+ shard_size = (
+ embedding.org_vocab_size + self.split_ngram_parts - 1
+ ) // self.split_ngram_parts
+ checkpoint_start = shard_index * shard_size
+ expected_rows = max(
+ 0,
+ min(shard_size, embedding.org_vocab_size - checkpoint_start),
+ )
+ expected_shape = (expected_rows, embedding.embedding_dim)
+ if tuple(loaded_weight.shape) != expected_shape:
+ raise ValueError(
+ f"Shape mismatch for PLE embedding shard {shard_index}: "
+ f"expected {expected_shape}, got "
+ f"{tuple(loaded_weight.shape)}"
+ )
+ embedding.weight.weight_loader(
+ embedding.weight,
+ loaded_weight,
+ checkpoint_start=checkpoint_start,
+ )
+ loaded.add("ngram_embedding.weight")
+ continue
+ regular_weights.append((name, loaded_weight))
+
+ if nvfp4_runtime:
+ if not packed_codes:
+ raise ValueError("NVFP4 PLE checkpoint contains no packed shards")
+ for shard_index in packed_codes:
+ if (
+ shard_index not in packed_scales
+ or shard_index not in packed_outer_scales
+ ):
+ raise ValueError(
+ f"NVFP4 PLE shard {shard_index} is missing its scale tensors"
+ )
+ shared_outer_scale = _get_shared_nvfp4_outer_scale(
+ {
+ shard_index: packed_outer_scales[shard_index]
+ for shard_index in packed_codes
+ }
+ )
+
+ for shard_index, codes in packed_codes.items():
+ embedding = self.ngram_embedding
+ shard_size = (
+ embedding.org_vocab_size + self.split_ngram_parts - 1
+ ) // self.split_ngram_parts
+ checkpoint_start = shard_index * shard_size
+ expected_rows = max(
+ 0, min(shard_size, embedding.org_vocab_size - checkpoint_start)
+ )
+ expected_codes_shape = (expected_rows, embedding.embedding_dim // 2)
+ if tuple(codes.shape) != expected_codes_shape:
+ raise ValueError(
+ f"Shape mismatch for packed PLE shard {shard_index}: "
+ f"expected {expected_codes_shape}, got {tuple(codes.shape)}"
+ )
+ tp_bounds = {
+ "checkpoint_start": checkpoint_start,
+ "tp_start": embedding.shard_indices.org_vocab_start_index,
+ "tp_end": embedding.shard_indices.org_vocab_end_index,
+ }
+ scales = packed_scales[shard_index]
+ expected_scales_shape = (
+ expected_rows,
+ embedding.embedding_dim // _NVFP4_BLOCK_SIZE,
+ )
+ if tuple(scales.shape) != expected_scales_shape:
+ raise ValueError(
+ f"Scale shape mismatch for packed PLE shard {shard_index}: "
+ f"expected {expected_scales_shape}, got {tuple(scales.shape)}"
+ )
+ copy_ple_embedding_shard_(embedding.weight.data, codes, **tp_bounds)
+ copy_ple_embedding_shard_(embedding.weight_scale.data, scales, **tp_bounds)
+
+ if nvfp4_runtime:
+ self.ngram_embedding.weight_scale_2.data.copy_(shared_outer_scale)
+ loaded.update(
+ {
+ "ngram_embedding.weight",
+ "ngram_embedding.weight_scale",
+ "ngram_embedding.weight_scale_2",
+ }
+ )
+
+ if regular_weights:
+ loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
+ return loaded
+
+
class Qwen4ExpPLELayer(nn.Module, MambaBase):
def __init__(
self,
diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py
index 1c2808a..1b7d55f 100644
--- a/vllm/models/qwen4_exp/nvidia/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/qsa.py
@@ -4,6 +4,14 @@
from __future__ import annotations
+import atexit
+import json
+import math
+import os
+import threading
+import time
+from collections.abc import Collection, Mapping
+from pathlib import Path
from typing import ClassVar, cast
import torch
@@ -21,6 +29,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.logger import init_logger
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
from vllm.model_executor.models.qwen3_next import Qwen3NextAttention
from vllm.platforms import current_platform
@@ -54,6 +63,235 @@ from . import model
from .indexer_qsa import QSAIndexer
+logger = init_logger(__name__)
+
+# --- FP8 K/V scale calibration collector -----------------------
+# Active ONLY when VLLM_QSA_KV_COLLECT names a writable directory. Records a
+# GPU-side running absmax of the pre-quantization BF16 K/V per layer and dumps
+# per-rank JSON (periodically + at exit) for an offline max-merge into the
+# VLLM_QSA_KV_SCALES sidecar. Calibration must run with --enforce-eager: the
+# collector mutates module-external state per step, which cudagraph capture
+# would freeze and torch.compile may not retrace. With the env unset this is a
+# constant-false branch and prod behaviour is unchanged.
+_QSA_KV_COLLECT_ENV = "VLLM_QSA_KV_COLLECT"
+_qsa_collect_dir = os.environ.get(_QSA_KV_COLLECT_ENV, "").strip()
+_qsa_collect_state: dict[str, list[torch.Tensor]] = {}
+_qsa_collect_calls = 0
+
+
+_qsa_collect_ident: str | None = None
+
+
+def _qsa_collect_identity() -> str:
+ """Stable per-process dump name: (pp, tp) rank, resolved once and cached so
+ a teardown-time lookup failure cannot move the final flush to a PID-named
+ file. TP-only runs keep the historical name."""
+ global _qsa_collect_ident
+ if _qsa_collect_ident is not None:
+ return _qsa_collect_ident
+ try:
+ from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank
+
+ rank = get_tensor_model_parallel_rank()
+ pp_group = get_pp_group()
+ pp_tag = f"pp{pp_group.rank_in_group}_" if pp_group.world_size > 1 else ""
+ _qsa_collect_ident = f"{pp_tag}rank{rank}"
+ except Exception:
+ return f"pid{os.getpid()}" # not cached: retry once the groups exist
+ return _qsa_collect_ident
+
+
+def _qsa_collect_dump() -> None:
+ if not _qsa_collect_state:
+ return
+ out = {
+ name: {"k_absmax": float(st[0].item()), "v_absmax": float(st[1].item())}
+ for name, st in _qsa_collect_state.items()
+ }
+ path = Path(_qsa_collect_dir) / f"qsa_absmax_{_qsa_collect_identity()}.json"
+ tmp = path.with_suffix(".tmp")
+ tmp.write_text(json.dumps(out, indent=1), encoding="utf-8")
+ tmp.replace(path)
+
+
+def _qsa_collect_absmax(
+ layer_name: str, key: torch.Tensor, value: torch.Tensor
+) -> None:
+ global _qsa_collect_calls
+ k_abs = key.detach().abs().amax().float()
+ v_abs = value.detach().abs().amax().float()
+ st = _qsa_collect_state.get(layer_name)
+ if st is None:
+ if not _qsa_collect_state:
+ atexit.register(_qsa_collect_dump)
+ _qsa_collect_state[layer_name] = [k_abs, v_abs]
+ else:
+ torch.maximum(st[0], k_abs, out=st[0])
+ torch.maximum(st[1], v_abs, out=st[1])
+ _qsa_collect_calls += 1
+ if _qsa_collect_calls % 2000 == 0:
+ _qsa_collect_dump()
+# ------------------------------------------------------------------------------
+
+# --- FP8 K/V runtime clipping counter --------------------------
+# Under-coverage guard for the calibrated scales sidecar: counts K/V elements
+# that would saturate E4M3 after scaling (|x| > 448 * scale), i.e. inputs that
+# exceed the calibration-time absmax. Active ONLY when VLLM_QSA_KV_CLIP_COUNT
+# is set to a positive read interval in seconds.
+#
+# Graph-safety: unlike the calibration collector above, this MUST work under
+# FULL cudagraph replay, where the Python forward does not run per step. The
+# counters are persistent per-layer GPU int64 tensors allocated on the first
+# (eager, pre-capture) forward; the increment is pure captured GPU ops, so
+# graph replay keeps counting. A daemon reader thread .item()s the counters on
+# a cadence and logs: WARNING on any new clips, INFO heartbeat every 10th pass
+# so a quiet log still proves the instrument is alive.
+_QSA_CLIP_ENV = "VLLM_QSA_KV_CLIP_COUNT"
+_qsa_clip_raw = os.environ.get(_QSA_CLIP_ENV, "").strip()
+try:
+ _qsa_clip_interval = float(_qsa_clip_raw) if _qsa_clip_raw else 0.0
+except ValueError as exc: # fail closed: a set-but-broken gate must not be silent
+ raise ValueError(
+ f"{_QSA_CLIP_ENV} must be a read interval in seconds, got {_qsa_clip_raw!r}"
+ ) from exc
+if _qsa_clip_raw and not (math.isfinite(_qsa_clip_interval) and _qsa_clip_interval > 0.0):
+ raise ValueError(
+ f"{_QSA_CLIP_ENV} must be a finite positive number of seconds, got {_qsa_clip_raw!r}"
+ )
+_qsa_clip_counters: dict[str, torch.Tensor] = {}
+# v2 reader fix: the reader must never block on the compute
+# stream; see guard_telemetry.NonBlockingReader.
+from vllm.v1.worker.guard_telemetry import NonBlockingReader as _NonBlockingReader # noqa: E402
+_qsa_clip_nb = _NonBlockingReader()
+_qsa_clip_reader_started = False
+_qsa_clip_lock = threading.Lock()
+
+
+def _qsa_clip_reader() -> None:
+ last: dict[str, tuple[int, int]] = {}
+ passes = 0
+ while True:
+ time.sleep(_qsa_clip_interval)
+ passes += 1
+ try:
+ totals_k = totals_v = 0
+ for name, ctr in list(_qsa_clip_counters.items()):
+ if not _qsa_clip_nb.registered(name):
+ _qsa_clip_nb.register(name, ctr) # reader thread: never capturing
+ ready = _qsa_clip_nb.poll()
+ for name, ctr in list(_qsa_clip_counters.items()):
+ row = ready.get(name)
+ if row is None:
+ # copy still in flight: report on the next pass
+ prev_k, prev_v = last.get(name, (0, 0))
+ totals_k += prev_k
+ totals_v += prev_v
+ continue
+ k_total = int(row[0])
+ v_total = int(row[1])
+ totals_k += k_total
+ totals_v += v_total
+ prev_k, prev_v = last.get(name, (0, 0))
+ if k_total > prev_k or v_total > prev_v:
+ logger.warning(
+ "QSA KV CLIPPING on %s: k +%d (total %d), v +%d (total %d)"
+ " -- inputs exceeded the calibration absmax; the scales"
+ " sidecar may need re-calibration",
+ name,
+ k_total - prev_k,
+ k_total,
+ v_total - prev_v,
+ v_total,
+ )
+ last[name] = (k_total, v_total)
+ if passes % 10 == 1:
+ logger.info(
+ "QSA clip counters alive: %d layer(s), totals k=%d v=%d",
+ len(_qsa_clip_counters),
+ totals_k,
+ totals_v,
+ )
+ except Exception:
+ logger.exception("QSA clip reader pass failed; thread continues")
+
+
+def _qsa_clip_count(
+ layer: "Qwen4ExpQSAAttention",
+ key: torch.Tensor,
+ value: torch.Tensor,
+ slot_mapping: torch.Tensor | None = None,
+) -> None:
+ global _qsa_clip_reader_started
+ ctr = _qsa_clip_counters.get(layer.layer_name)
+ if ctr is None:
+ # First forward is eager (profiling/warmup precede capture), so the
+ # allocation never happens inside cudagraph capture.
+ ctr = torch.zeros(2, dtype=torch.int64, device=key.device)
+ _qsa_clip_counters[layer.layer_name] = ctr # registered with the reader by the reader thread (never during capture)
+ with _qsa_clip_lock:
+ if not _qsa_clip_reader_started:
+ _qsa_clip_reader_started = True
+ threading.Thread(
+ target=_qsa_clip_reader,
+ name="qsa-clip-reader",
+ daemon=True,
+ ).start()
+ logger.info(
+ "QSA clip counter armed (read interval %.0fs)",
+ _qsa_clip_interval,
+ )
+ # Thresholds are 1-D so the comparison promotes to float32; a 0-dim float32
+ # tensor promotes like a Python scalar and the threshold would be rounded to
+ # the BF16 grid of K/V, undercounting values just above the calibrated ceiling.
+ k_ceiling = (448.0 * layer._k_scale).to(torch.float32).view(1)
+ v_ceiling = (448.0 * layer._v_scale).to(torch.float32).view(1)
+ k_over = key.detach().abs() > k_ceiling
+ v_over = value.detach().abs() > v_ceiling
+ if slot_mapping is not None:
+ # The CUDA writer (reshape_and_cache_flash) skips rows with a negative
+ # slot: padding rows of a captured batch size and, under the V2 model
+ # runner, stale rows of the persistent input buffers. Count only what
+ # is actually written, with fixed-shape masking so the op is capturable.
+ n = min(slot_mapping.shape[0], key.shape[0])
+ written = (slot_mapping[:n] >= 0).view(-1, 1, 1)
+ k_over = k_over[:n] & written
+ v_over = v_over[:n] & written
+ ctr[0] += k_over.sum()
+ ctr[1] += v_over.sum()
+# ------------------------------------------------------------------------------
+
+_QSA_FP8_CACHE_DTYPES = ("fp8", "fp8_e4m3")
+_QSA_SUPPORTED_CACHE_DTYPES = ("auto", "bfloat16", *_QSA_FP8_CACHE_DTYPES)
+
+
+def _is_qsa_fp8_cache_dtype(cache_dtype: str) -> bool:
+ return cache_dtype in _QSA_FP8_CACHE_DTYPES
+
+
+def _validated_qsa_fp8_dtype(cache_dtype: str) -> torch.dtype | None:
+ """Return the platform E4M3 dtype after enforcing the validated HW gate."""
+
+ if not _is_qsa_fp8_cache_dtype(cache_dtype):
+ return None
+ if not current_platform.is_cuda():
+ raise ValueError(
+ "Qwen4Exp QSA E4M3 KV cache is supported only on CUDA; "
+ "ROCm and other platforms have not been validated"
+ )
+
+ capability = current_platform.get_device_capability()
+ if capability is None or capability.to_int() != 86:
+ cap_str = (
+ capability.as_version_str() if capability is not None else "unknown"
+ )
+ raise ValueError(
+ "Qwen4Exp QSA E4M3 KV cache is validated only on SM86, but "
+ f"{current_platform.get_device_name()} has compute capability "
+ f"{cap_str}. Re-run with --kv-cache-dtype bfloat16."
+ )
+ return current_platform.fp8_dtype()
+
+
class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
"""Flash metadata supporting uniform decode and target-verify graphs."""
@@ -64,7 +302,12 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
"""FullAttentionSpec backend used by the merged QSA owner."""
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
- supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+ "auto",
+ "bfloat16",
+ "fp8",
+ "fp8_e4m3",
+ ]
@staticmethod
def get_name() -> str:
@@ -99,15 +342,36 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
supports_pcp: bool = False
def __init__(self, *args, **kwargs) -> None:
+ # QSA never dispatches to a FlashAttention kernel: attention runs through
+ # forward_qsa -> qsa_sparse_paged_attention (our Triton sparse-GQA kernel)
+ # and the cache write goes through the CUDA reshape_and_cache_flash op.
+ # FlashAttentionImpl.__init__ nevertheless rejects any quantized KV dtype
+ # that FA itself cannot read, and sm_86 has no FP8 FA kernel. Present
+ # "auto" to the base class, then restore the real dtype. The only other
+ # kv_cache_dtype use in that constructor is an SM90/FA4 dequant path that
+ # cannot trigger here, and supports_quant_query_input is set below anyway.
+ requested_kv_cache_dtype = None
+ if "kv_cache_dtype" in kwargs:
+ requested_kv_cache_dtype = kwargs["kv_cache_dtype"]
+ kwargs = {**kwargs, "kv_cache_dtype": "auto"}
+ elif len(args) > 6:
+ requested_kv_cache_dtype = args[6]
+ args = (*args[:6], "auto", *args[7:])
super().__init__(*args, **kwargs)
+ if requested_kv_cache_dtype is not None:
+ self.kv_cache_dtype = requested_kv_cache_dtype
if not is_flash_attn_varlen_func_available():
raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")
if self.dcp_world_size != 1:
raise NotImplementedError(
"Qwen4Exp QSA does not support decode context parallelism"
)
- if self.kv_cache_dtype not in ("auto", "bfloat16"):
- raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
+ if self.kv_cache_dtype not in _QSA_SUPPORTED_CACHE_DTYPES:
+ raise NotImplementedError(
+ "Qwen4Exp QSA supports BF16 or static per-tensor E4M3 main "
+ f"KV caches, not {self.kv_cache_dtype!r}"
+ )
+ self.qsa_fp8_dtype = _validated_qsa_fp8_dtype(self.kv_cache_dtype)
self.supports_quant_query_input = False
def forward_qsa(
@@ -144,11 +408,42 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
logical_indices = topk_buffer[:num_tokens]
token_to_req = token_to_req[:num_tokens]
key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
- if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:
- raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")
+ if query.dtype != torch.bfloat16:
+ raise NotImplementedError("Qwen4Exp QSA requires a BF16 query")
+ if key_cache.dtype != value_cache.dtype:
+ raise TypeError("Qwen4Exp QSA K and V cache dtypes must match")
+ if _is_qsa_fp8_cache_dtype(self.kv_cache_dtype):
+ if self.qsa_fp8_dtype is None:
+ raise RuntimeError("QSA FP8 dtype was not initialized")
+ if key_cache.dtype not in (torch.uint8, self.qsa_fp8_dtype):
+ raise NotImplementedError(
+ "Qwen4Exp QSA E4M3 cache storage must use raw uint8 or "
+ f"the platform FP8 dtype, not {key_cache.dtype}"
+ )
+ elif key_cache.dtype != torch.bfloat16:
+ raise NotImplementedError(
+ "Qwen4Exp QSA BF16 mode requires a BF16 main KV cache"
+ )
from .ops.qsa import qsa_sparse_paged_attention
+ # Announce the effective scales once per layer. load_qsa_static_kv_scales()
+ # is an explicit integration call, so a deployment that forgets it runs at
+ # the vLLM default of 1.0 -- which is a legitimate configuration but must
+ # never be a silent one. This makes the actual values greppable per rank.
+ if _is_qsa_fp8_cache_dtype(self.kv_cache_dtype) and not getattr(
+ self, "_qsa_scales_logged", False
+ ):
+ self._qsa_scales_logged = True
+ logger.info(
+ "Qwen4Exp QSA E4M3 KV active on %s: k_scale=%.6g v_scale=%.6g%s",
+ getattr(layer, "layer_name", "<unnamed>"),
+ float(layer._k_scale),
+ float(layer._v_scale),
+ "" if float(layer._k_scale) != 1.0 or float(layer._v_scale) != 1.0
+ else " (DEFAULT 1.0 -- no static calibration loaded)",
+ )
+
qsa_sparse_paged_attention(
query[:num_tokens],
key_cache,
@@ -157,7 +452,9 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
attn_metadata.block_table,
token_to_req,
use_prefill_config,
- output[:num_tokens],
+ out=output[:num_tokens],
+ k_scale=layer._k_scale,
+ v_scale=layer._v_scale,
output_gate=output_gate[:num_tokens],
)
return output
@@ -185,10 +482,20 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
raise ValueError("Qwen4Exp QSA requires a paged KV cache")
if model_config.dtype != torch.bfloat16:
raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
- if cache_config.cache_dtype not in ("auto", "bfloat16"):
- raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
- if getattr(quant_config, "kv_cache_scheme", None) is not None:
- raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
+ if cache_config.cache_dtype not in _QSA_SUPPORTED_CACHE_DTYPES:
+ raise NotImplementedError(
+ "Qwen4Exp QSA supports BF16 or static per-tensor E4M3 main "
+ f"KV caches, not {cache_config.cache_dtype!r}"
+ )
+ qsa_fp8_dtype = _validated_qsa_fp8_dtype(cache_config.cache_dtype)
+ kv_cache_scheme = getattr(quant_config, "kv_cache_scheme", None)
+ if kv_cache_scheme is not None:
+ raise NotImplementedError(
+ "Qwen4Exp QSA does not implement kv_cache_scheme variants. "
+ "Dynamic, per-head, per-token, and other non-scalar KV scales "
+ "are unsupported; use static scalar per-layer K/V scales from "
+ "load_qsa_static_kv_scales()."
+ )
parallel_config = vllm_config.parallel_config
if (
parallel_config.prefill_context_parallel_size > 1
@@ -285,10 +592,23 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
self.kv_cache_dtype, model_config
)
- if self.kv_cache_torch_dtype != torch.bfloat16:
- raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")
+ if _is_qsa_fp8_cache_dtype(self.kv_cache_dtype):
+ if qsa_fp8_dtype is None:
+ raise RuntimeError("QSA FP8 dtype was not initialized")
+ if self.kv_cache_torch_dtype not in (torch.uint8, qsa_fp8_dtype):
+ raise NotImplementedError(
+ "Qwen4Exp QSA E4M3 cache storage resolved to unsupported "
+ f"dtype {self.kv_cache_torch_dtype}"
+ )
+ elif self.kv_cache_torch_dtype != torch.bfloat16:
+ raise NotImplementedError(
+ "Qwen4Exp QSA BF16 mode requires BF16 cache storage"
+ )
self.kv_sharing_target_layer_name = None
self.kv_cache = torch.tensor([])
+ # The checkpoint has no KV scale scheme. These registered scalar
+ # buffers therefore remain 1.0 unless explicitly replaced after
+ # weight loading by load_qsa_static_kv_scales().
set_default_quant_scales(self, register_buffer=True)
self.attn_backend = Qwen4ExpQSAFlashAttentionBackend
@@ -381,6 +701,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
)
if selected.shape != (num_tokens, self.indexer.packed_output_width):
raise RuntimeError("QSA indexer returned an invalid selection shape")
+ if _qsa_clip_interval > 0.0 and _is_qsa_fp8_cache_dtype(
+ self.kv_cache_dtype
+ ):
+ _qsa_clip_count(self, key, value, main_metadata.slot_mapping)
impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
impl.do_kv_cache_update(
self,
@@ -414,6 +738,8 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
query = q.view(num_tokens, self.num_heads, self.head_dim)
key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
value = v.view(num_tokens, self.num_kv_heads, self.head_dim)
+ if _qsa_collect_dir:
+ _qsa_collect_absmax(self.layer_name, key, value)
attn_output = torch.empty_like(query)
# Keep the index projection outside the eager break.
projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
@@ -431,9 +757,146 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
return output
+def load_qsa_static_kv_scales(
+ layers: nn.Module | Mapping[str, nn.Module],
+ sidecar_path: str | Path,
+ *,
+ strict: bool = True,
+ expected_layer_names: Collection[str] | None = None,
+) -> list[str]:
+ """Load static scalar QSA K/V scales into already-loaded model layers.
+
+ Call this on every model worker after checkpoint weight loading finishes
+ and before any profiling, warmup, or cudagraph capture. ``layers`` may be
+ the loaded root module or ``compilation_config.static_forward_context``.
+
+ The JSON schema is::
+
+ {
+ "model.layers.3.self_attn.attn": {
+ "k_scale": 0.25,
+ "v_scale": 2.0
+ }
+ }
+
+ ``strict=False`` permits a deliberate subset; omitted layers keep the
+ vLLM default scale of 1.0. Unknown names and non-scalar values are always
+ errors.
+
+ Pipeline parallelism: a rank holds only its stage's layers, so the sidecar
+ legitimately names layers this rank does not own. Pass
+ ``expected_layer_names`` = the QSA layer names of the WHOLE model; names are
+ validated against that set, ``strict`` requires the sidecar to cover all
+ of it (checkable on every rank), and the scales are applied to the local
+ intersection only. Without it (TP-only callers) the local layers are the
+ universe, as before.
+ """
+
+ path = Path(sidecar_path)
+ raw = json.loads(path.read_text(encoding="utf-8"))
+ if not isinstance(raw, dict):
+ raise ValueError(f"{path} must contain a JSON object keyed by layer name")
+
+ if isinstance(layers, nn.Module):
+ qsa_layers = {
+ layer.layer_name: layer
+ for layer in layers.modules()
+ if isinstance(layer, Qwen4ExpQSAAttention)
+ }
+ else:
+ qsa_layers = {
+ name: layer
+ for name, layer in layers.items()
+ if isinstance(name, str)
+ and isinstance(layer, Qwen4ExpQSAAttention)
+ }
+
+ if not qsa_layers and expected_layer_names is None:
+ raise ValueError("No Qwen4Exp QSA attention layers were supplied")
+ non_fp8_layers = sorted(
+ name
+ for name, layer in qsa_layers.items()
+ if not _is_qsa_fp8_cache_dtype(layer.kv_cache_dtype)
+ )
+ if non_fp8_layers:
+ raise ValueError(
+ "Static QSA KV scales may only be injected into E4M3 layers; "
+ f"non-FP8 layers: {non_fp8_layers}"
+ )
+
+ if expected_layer_names is None:
+ universe = set(qsa_layers)
+ else:
+ universe = set(expected_layer_names)
+ not_expected = sorted(name for name in qsa_layers if name not in universe)
+ if not_expected:
+ raise ValueError(
+ "Local QSA layers outside expected_layer_names: "
+ f"{not_expected}"
+ )
+ unknown = sorted(name for name in raw if name not in universe)
+ if unknown:
+ raise ValueError(f"Unknown QSA layer names in {path}: {unknown}")
+ missing = sorted(name for name in universe if name not in raw)
+ if strict and missing:
+ raise ValueError(f"Missing QSA layer scales in {path}: {missing}")
+ remote = sorted(name for name in raw if name not in qsa_layers)
+ if remote:
+ logger.info(
+ "Qwen4Exp QSA: %d sidecar layer(s) belong to other pipeline stages "
+ "and are skipped on this rank: %s%s",
+ len(remote),
+ ", ".join(remote[:4]),
+ " ..." if len(remote) > 4 else "",
+ )
+
+ validated: dict[str, tuple[float, float]] = {}
+ for layer_name, entry in raw.items():
+ if not isinstance(entry, dict):
+ raise ValueError(f"{layer_name} scales must be a JSON object")
+ if set(entry) != {"k_scale", "v_scale"}:
+ raise ValueError(
+ f"{layer_name} must contain exactly k_scale and v_scale"
+ )
+
+ values: list[float] = []
+ for field in ("k_scale", "v_scale"):
+ value = entry[field]
+ if type(value) not in (int, float):
+ raise ValueError(
+ f"{layer_name}.{field} must be a scalar JSON number"
+ )
+ scale = float(value)
+ if not math.isfinite(scale) or scale <= 0.0:
+ raise ValueError(
+ f"{layer_name}.{field} must be finite and positive"
+ )
+ values.append(scale)
+
+ layer = qsa_layers.get(layer_name)
+ if layer is None:
+ continue # another pipeline stage's layer: schema-checked only
+ for field in ("_k_scale", "_v_scale"):
+ tensor = getattr(layer, field, None)
+ if not isinstance(tensor, torch.Tensor) or tensor.numel() != 1:
+ raise ValueError(
+ f"{layer_name}.{field} must be an existing scalar tensor"
+ )
+ validated[layer_name] = (values[0], values[1])
+
+ with torch.no_grad():
+ for layer_name, (k_scale, v_scale) in validated.items():
+ layer = qsa_layers[layer_name]
+ layer._k_scale.fill_(k_scale)
+ layer._v_scale.fill_(v_scale)
+
+ return sorted(validated)
+
+
__all__ = [
"QSAIndexer",
"Qwen4ExpQSAAttention",
"Qwen4ExpQSAFlashAttentionBackend",
"Qwen4ExpQSAFlashAttentionImpl",
+ "load_qsa_static_kv_scales",
]
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink