diff --git a/.gitignore b/.gitignore index 72089b06e7e..b3c01d41ea2 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ htmlcov/ .coverage.* coverage.xml .pytest_cache/ +.cache/ # Sphinx documentation docs/build diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4580a045f39..588529cc6ea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -10,6 +10,8 @@ Changelog - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add fused NVFP4/FP8 Q/K/P/V and 2:4 attention quantization for MLA models (DeepSeek-family) on the vLLM ``TRITON_MLA`` backend, served through ``examples/vllm_serve`` (``QuantSparseAttnWorker`` / ``install_vllm_nvfp4_attention``). Each attention BMM operand is fake-quantized in-kernel from a raw latent cache — decode quantizes K (feature axis) and V (token axis) on read and P after the softmax row-sum; prefill quantizes the projected operands. Optional checkpoint N:M score sparsity applies to prefill new-token attention. Requires ``--attention-backend TRITON_MLA`` and ``--enforce-eager``. + *Megatron Framework (M-LM / M-Bridge)* - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root `` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..30c26df1efb 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -219,7 +219,28 @@ K is QDQ before its cache write, while V is written pristine. Complete 16-token Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. The default `FULL_AND_PIECEWISE` mode remains enabled for fixed N:M and attention-only NVFP4; checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar. -Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs. +Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs. + +#### MLA attention (DeepSeek-family) + +The quantized worker also supports dense MLA text attention on the `TRITON_MLA` backend (pass `--attention-backend TRITON_MLA` when it is not the platform default): + +```bash +python vllm_serve_sparse_attn.py -tp 8 \ + --no-enable-prefix-caching --enforce-eager \ + --attention-backend TRITON_MLA \ + --worker-cls sparse_attn_worker.QuantSparseAttnWorker +``` + +The MLA operand mapping differs from regular attention because the latent cache is shared by both BMMs: + +- `q_format`: prefill quantizes the projected 192-d query in-kernel; decode quantizes the absorbed `kv_lora_rank + rope`-d query (FP32 QDQ carrier). With `q_format=fp8` the module-level quantizer QDQs the pre-projection query instead. +- `k_format`: governs the write-once latent-cache QDQ (`kv_c`, `k_pe`, applied at cache-write time by the impl) and the prefill projected K (in-kernel). +- `v_format`: governs the prefill projected V (in-kernel) **only**. **Decode V is not independently quantized** — decode BMM2 consumes the write-once quantized latent as-is, so decode V inherits `k_format` (the latent's feature-axis quantization used for BMM1-K), not `v_format`. This is a deliberate consequence of the write-once, no-on-read-re-quant design (stable across steps/splits): a shared latent cannot be stored quantized along both K's feature axis and V's token axis at once. Reading it independently as a token-axis-quantized V would require an on-read re-quant (the MNI-style raw-cache decode model), trading away that stability. So for decode, treat the contract as "V inherits k_format," not independent V quantization. +- The latent QDQ is applied at the cache-write hook (not the module forward), so the **new-tokens** prefill projection reads the bf16 latent and its K/V operands are quantized exactly once. Note one residual: cached-**context** prefill chunks gather from the already-quantized paged cache and re-quantize the projected operands, so those chunks are double-quantized — inherent to reading a stored quantized latent, and absent for short prompts that fit a single chunk. +- `p_format`: fused into both the prefill and decode kernels; the softmax denominator stays unquantized and P amax defaults to 1.0. + +MLA decode uses a fixed 32-split, 32-key-tile schedule with tile boundaries at absolute token positions, so quantized decode results are stable as the sequence grows and reproducible across batch shapes and devices. Optional checkpoint N:M sparsity applies to prefill new-token attention only (cached-context chunks run dense); skip-softmax is rejected on MLA layers. Sparse-only installation ignores MLA layers. DeepSeek V3.2-style sparse-indexer MLA and FP8 latent caches are unsupported. ## Known Problems diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 709d6532fb3..fe050492599 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -141,12 +141,14 @@ def update_kv_cfg_for_mla(model: torch.nn.Module, kv_quant_cfg: list) -> list: `k_bmm_quantizer` and `v_bmm_quantizer`. This function copies the config from `*[kv]_bmm_quantizer` to also cover `*kv_c_bmm_quantizer`. """ - try: - from vllm.attention.layer import MLAAttention - except ImportError: - return kv_quant_cfg - - if not any(isinstance(m, MLAAttention) for m in model.modules()): + # Resolved via the quant plugin, which handles the MLAAttention module + # moving across vLLM releases (vllm.attention.layer no longer exists in + # vLLM >= 0.26). + from modelopt.torch.quantization.plugins.vllm import VllmMLAAttention + + if VllmMLAAttention is None or not any( + isinstance(m, VllmMLAAttention) for m in model.modules() + ): return kv_quant_cfg kv_entry = next( diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 1bfc7f95fd4..362283058de 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -143,6 +143,14 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]: ) return ("group", group_key, value) + # MLA latent quantizers: self_attn.kv_c/k_pe_bmm_quantizer -> + # self_attn.mla_attn.mla_attn.* (DeepSeek-style wrapper nesting in vLLM). + # vLLM-native keys that already carry the mla_attn prefix copy as-is below. + mla_bmm_match = re.search(r"(.*\.self_attn)\.((?:kv_c|k_pe)_bmm_quantizer.*)$", key) + if mla_bmm_match: + new_key = mla_bmm_match.group(1) + ".mla_attn.mla_attn." + mla_bmm_match.group(2) + return ("copy", new_key, value) + # Transform bmm_quantizer keys: self_attn.q/k/v_bmm_quantizer -> self_attn.attn.q/k/v_bmm_quantizer bmm_match = re.search(r"(.*\.self_attn)\.([qkv]_bmm_quantizer.*)$", key) or re.search( r"(.*\.mixer)\.([qkv]_bmm_quantizer.*)$", key diff --git a/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py b/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py index b33900546eb..70802a9b999 100644 --- a/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py +++ b/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py @@ -13,12 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NVFP4 operand helpers for the attention ``P @ V`` matmul (BMM2). +"""NVFP4 operand helpers for the attention BMM matmuls. -P and V share the low-level ``nvfp4_scalar_qdq`` primitive, but retain thin -operand-specific wrappers because their layouts and amax reductions differ. -P is nonnegative with layout ``[M, K]``; V is signed with layout ``[K, N]``. -Both use block-16 scaling along the BMM2 contraction axis. +P, V, and the signed A-side share the low-level ``nvfp4_scalar_qdq`` +primitive, but retain thin operand-specific wrappers because their layouts +and amax reductions differ. P is nonnegative with layout ``[M, K]``; V is +signed with layout ``[K, N]``; the signed A-side (Q of BMM1) is ``[M, K]``. +All use block-16 scaling along the BMM contraction axis. """ import math @@ -70,6 +71,21 @@ def _v_qdq_nvfp4(v, global_scale, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr): return tl.reshape(nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16), (BLOCK_N, BLOCK_D)) +@triton.jit +def _a_qdq_nvfp4(x, global_scale, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr): + """Fake-quantize a signed A-side operand ``[M, K]`` in block-16 groups along K. + + The BMM1 Q-operand counterpart of :func:`_p_qdq_nvfp4`: same ``[M, K]`` + layout and contraction-axis blocking, but signed, so the block amax uses + ``abs``. Zero-padded lanes (masked loads) form all-zero blocks that + ``nvfp4_scalar_qdq`` guards to zero. + """ + tl.static_assert(BLOCK_K % 16 == 0, "BLOCK_K must be divisible by 16 for NVFP4") + grouped = tl.reshape(x, (BLOCK_M, BLOCK_K // 16, 16)) + block_amax = tl.expand_dims(tl.max(tl.abs(grouped), axis=2), 2) + return tl.reshape(nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16), (BLOCK_M, BLOCK_K)) + + @triton.jit def _fake_quant_v_onwrite_kernel( V_cache, diff --git a/modelopt/torch/kernels/quantization/attention/mla/__init__.py b/modelopt/torch/kernels/quantization/attention/mla/__init__.py new file mode 100644 index 00000000000..e83482d0eee --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/mla/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MLA attention kernels with fused fake quantization (prefill and decode).""" + +from .mla_decode import mla_attention_decode +from .mla_prefill import mla_prefill_attention + +__all__ = ["mla_attention_decode", "mla_prefill_attention"] diff --git a/modelopt/torch/kernels/quantization/attention/mla/mla_decode.py b/modelopt/torch/kernels/quantization/attention/mla/mla_decode.py new file mode 100644 index 00000000000..3c7e9f4a7ff --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/mla/mla_decode.py @@ -0,0 +1,459 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Split-K absorbed-MLA decode over a paged latent cache, with fused P QDQ. + +Absorbed MLA decode is MQA: every query head attends to the single latent KV +"head". The absorbed query carries ``kv_lora_rank + qk_rope_head_dim`` +features (576 for DeepSeek-family models); K is the full latent cache row and +V is its first ``kv_lora_rank`` features — the same memory read once per tile +and reused for both BMMs. + +Quantization contract (quantized-BMM model): the latent cache holds RAW +(bf16/fp16) values. Each BMM operand is fake-quantized once, on read, along +its own contraction axis, and K and V are quantized INDEPENDENTLY from the +raw latent — K along the feature axis (``k_qdq``), V along the token axis +(``v_qdq``). So decode honors ``k_format`` and ``v_format`` separately, the +same faithful independent-operand model the prefill kernel and the dense +Q/K/P/V path use. The absorbed Q is fake-quantized by the caller (dynamic +NVFP4 uses an FP32 carrier, like ``triton_fa``'s ``Q_IS_FP32``); P is +quantized in-kernel after the row-sum, so the softmax denominator stays +unquantized. + +Split determinism: P and V groups are block-16 at absolute token positions +(``kv_start`` is a multiple of ``BLOCK_N``) under a fixed split/tile schedule, +so partial results are reproducible across batch shapes and devices. The +token-axis V QDQ does re-quantize the open (incomplete) tail block as the +sequence grows — the inherent tradeoff of on-read operand quantization versus +an immutable write-once cache. Inference-only. +""" + +import torch +import triton +import triton.language as tl + +from modelopt.torch.kernels.common.attention.decode_attention import _qdq_scale +from modelopt.torch.kernels.common.attention.triton_fa import LOG2E +from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _p_qdq_nvfp4, _v_qdq_nvfp4 +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq + +__all__ = ["mla_attention_decode"] + +# Referenced inside @triton.jit code, so it must be a tl.constexpr global. +LN2 = tl.constexpr(0.6931471805599453) + +_BLOCK_N = 32 +_DEFAULT_KV_SPLITS = 32 +_MAX_KV_SPLITS = 32 +_P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} + + +@triton.jit +def _mla_decode_split_kernel( + Q, # [batch, num_heads, KV_LORA_RANK + QK_ROPE_DIM] absorbed query + Latent_cache, # [num_blocks, page_size, KV_LORA_RANK + QK_ROPE_DIM] + Block_table, # [batch, max_blocks_per_seq] + B_seq_len, # [batch] + M_partial, # [batch, num_heads, NUM_KV_SPLITS] + L_partial, # [batch, num_heads, NUM_KV_SPLITS] + Acc_partial, # [batch, num_heads, NUM_KV_SPLITS, BLOCK_DL] + qk_scale, # softmax_scale * log2(e) + stride_qb, + stride_qh, + stride_lc_block, + stride_lc_pos, + stride_mb, + stride_mh, + stride_ab, + stride_ah, + stride_as, + p_qdq_scale, + k_qdq_scale, + v_qdq_scale, + max_blocks_per_seq, + H: tl.constexpr, # number of query heads + BLOCK_H: tl.constexpr, # query heads per program (grouped MQA) + BLOCK_N: tl.constexpr, + BLOCK_DL: tl.constexpr, # next_power_of_2(KV_LORA_RANK) + BLOCK_DPE: tl.constexpr, # next_power_of_2(QK_ROPE_DIM) + KV_LORA_RANK: tl.constexpr, + QK_ROPE_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + K_QDQ: tl.constexpr, # 0=off, 1=FP8 E4M3, 2=NVFP4 (on-read, feature axis) + V_QDQ: tl.constexpr, # 0=off, 1=FP8 E4M3, 2=NVFP4 (on-read, token axis) + P_QDQ: tl.constexpr, # 0=off, 1=FP8 E4M3, 2=NVFP4 + Q_IS_FP32: tl.constexpr, # dynamic NVFP4 QDQ carrier uses FP32 +): + """Compute one partial softmax for one head group, request, and KV split.""" + # Head groups on axis 0 (fastest varying) so programs sharing the same + # latent KV tiles are co-scheduled for L2 reuse (MQA: kv_group_num == H). + head_group = tl.program_id(0) + batch_idx = tl.program_id(1) + split_idx = tl.program_id(2) + + head_ids = head_group * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = head_ids < H + seq_len = tl.load(B_seq_len + batch_idx) + + dl_pos = tl.arange(0, BLOCK_DL) + dpe_pos = tl.arange(0, BLOCK_DPE) + dl_mask = dl_pos < KV_LORA_RANK + dpe_mask = dpe_pos < QK_ROPE_DIM + kv_pos = tl.arange(0, BLOCK_N) + + q_base = Q + batch_idx * stride_qb + head_ids[:, None] * stride_qh + q_nope = tl.load(q_base + dl_pos[None, :], mask=mask_h[:, None] & dl_mask[None, :], other=0.0) + q_pe = tl.load( + q_base + KV_LORA_RANK + dpe_pos[None, :], + mask=mask_h[:, None] & dpe_mask[None, :], + other=0.0, + ) + if Q_IS_FP32: + q_nope = q_nope.to(tl.float32) + q_pe = q_pe.to(tl.float32) + + # Tile-aligned split bounds: quant-relevant tile boundaries sit at + # absolute token positions, so numerics are stable as the sequence grows. + num_tiles = tl.cdiv(seq_len, BLOCK_N) + tiles_per_split = tl.cdiv(num_tiles, NUM_KV_SPLITS) + kv_lo = split_idx * tiles_per_split * BLOCK_N + kv_hi = tl.minimum(kv_lo + tiles_per_split * BLOCK_N, seq_len) + + running_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + running_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DL], dtype=tl.float32) + + for kv_start in range(kv_lo, kv_hi, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) + kv_abs = kv_start + kv_pos + kv_valid = kv_abs < seq_len + + page = tl.load( + Block_table + batch_idx * max_blocks_per_seq + kv_abs // PAGE_SIZE, + mask=kv_valid, + other=0, + ).to(tl.int64) + pos_ptrs = page * stride_lc_block + (kv_abs % PAGE_SIZE) * stride_lc_pos + + # K^T tiles from the RAW latent cache: NOPE [BLOCK_DL, BLOCK_N], PE [BLOCK_DPE, BLOCK_N] + k_nope = tl.load( + Latent_cache + pos_ptrs[None, :] + dl_pos[:, None], + mask=kv_valid[None, :] & dl_mask[:, None], + other=0.0, + ) + k_pe = tl.load( + Latent_cache + pos_ptrs[None, :] + KV_LORA_RANK + dpe_pos[:, None], + mask=kv_valid[None, :] & dpe_mask[:, None], + other=0.0, + ) + + # Capture RAW V (tokens x features) before the K on-read quant reassigns + # k_nope. K and V are quantized independently from the raw latent, each + # along its own BMM contraction axis (K: feature; V: token), so decode + # honors k_format and v_format separately (quantized-BMM model). + v = tl.trans(k_nope) + + # On-read K quant along the feature (BMM1 contraction) axis. + if K_QDQ == 1: + k_nope = fp8_scalar_qdq(k_nope, k_qdq_scale).to(k_nope.dtype) + k_pe = fp8_scalar_qdq(k_pe, k_qdq_scale).to(k_pe.dtype) + elif K_QDQ == 2: + k_nope = _v_qdq_nvfp4(k_nope.to(tl.float32), k_qdq_scale, BLOCK_DL, BLOCK_N) + k_pe = _v_qdq_nvfp4(k_pe.to(tl.float32), k_qdq_scale, BLOCK_DPE, BLOCK_N) + + if Q_IS_FP32 or K_QDQ == 2: + scores = tl.dot(q_nope.to(tl.float32), k_nope.to(tl.float32), input_precision="ieee") + scores += tl.dot(q_pe.to(tl.float32), k_pe.to(tl.float32), input_precision="ieee") + else: + scores = tl.dot(q_nope, k_nope) + tl.dot(q_pe, k_pe) + scores = scores * qk_scale + scores = tl.where(kv_valid[None, :], scores, float("-inf")) + + # --- Online softmax update: the denominator uses unquantized p --- + m_new = tl.maximum(running_max, tl.max(scores, 1)) + p = tl.math.exp2(scores - m_new[:, None]) + p = tl.where(kv_valid[None, :], p, 0.0) + correction = tl.math.exp2(running_max - m_new) + running_sum = running_sum * correction + tl.sum(p, 1) + acc = acc * correction[:, None] + + if P_QDQ == 1: + p = fp8_scalar_qdq(p, p_qdq_scale) + elif P_QDQ == 2: + # Native packing consumes the cache dtype; the QDQ value stays FP32. + p = p.to(Latent_cache.dtype.element_ty).to(tl.float32) + p = _p_qdq_nvfp4(p, p_qdq_scale, BLOCK_H, BLOCK_N) + + # On-read V quant along the token (BMM2 contraction) axis, independent of + # K. Groups are block-16 at absolute token positions (kv_start is a + # multiple of BLOCK_N), so split results are deterministic; the open tail + # block re-quantizes as the sequence grows (the on-read tradeoff). + if V_QDQ == 1: + v = fp8_scalar_qdq(v, v_qdq_scale).to(v.dtype) + elif V_QDQ == 2: + v = _v_qdq_nvfp4(v.to(tl.float32), v_qdq_scale, BLOCK_N, BLOCK_DL) + + if P_QDQ == 2 or V_QDQ == 2: + acc = tl.dot(p.to(tl.float32), v.to(tl.float32), acc, input_precision="ieee") + else: + acc = tl.dot(p.to(v.dtype), v, acc) + running_max = m_new + + partial_offset = batch_idx * stride_mb + head_ids * stride_mh + split_idx + tl.store(M_partial + partial_offset, running_max, mask=mask_h) + tl.store(L_partial + partial_offset, running_sum, mask=mask_h) + acc_ptrs = ( + batch_idx * stride_ab + + head_ids[:, None] * stride_ah + + split_idx * stride_as + + dl_pos[None, :] + ) + tl.store(Acc_partial + acc_ptrs, acc, mask=mask_h[:, None] & dl_mask[None, :]) + + +@triton.jit +def _mla_decode_combine_kernel( + M_partial, + L_partial, + Acc_partial, + Out, # [batch, num_heads, KV_LORA_RANK] + Lse, # [batch, num_heads] natural-log LSE (dummy when not STORE_LSE) + stride_mb, + stride_mh, + stride_ab, + stride_ah, + stride_as, + stride_ob, + stride_oh, + stride_lse_b, + BLOCK_DL: tl.constexpr, + KV_LORA_RANK: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + STORE_LSE: tl.constexpr, +): + """Merge split-local online-softmax states.""" + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + dl_pos = tl.arange(0, BLOCK_DL) + dl_mask = dl_pos < KV_LORA_RANK + base_ml = batch_idx * stride_mb + head_idx * stride_mh + base_acc = batch_idx * stride_ab + head_idx * stride_ah + + running_max = -float("inf") + running_sum = 0.0 + acc = tl.zeros([BLOCK_DL], dtype=tl.float32) + for split_idx in range(NUM_KV_SPLITS): + split_sum = tl.load(L_partial + base_ml + split_idx) + if split_sum > 0.0: + split_max = tl.load(M_partial + base_ml + split_idx) + split_acc = tl.load( + Acc_partial + base_acc + split_idx * stride_as + dl_pos, + mask=dl_mask, + other=0.0, + ) + new_max = tl.maximum(running_max, split_max) + correction = tl.math.exp2(running_max - new_max) + split_correction = tl.math.exp2(split_max - new_max) + acc = acc * correction + split_acc * split_correction + running_sum = running_sum * correction + split_sum * split_correction + running_max = new_max + + output = acc / tl.maximum(running_sum, 1e-6) + tl.store( + Out + batch_idx * stride_ob + head_idx * stride_oh + dl_pos, + output, + mask=dl_mask, + ) + if STORE_LSE: + lse = LN2 * (running_max + tl.math.log2(running_sum)) + lse = tl.where(running_sum == 0.0, float("-inf"), lse) + tl.store(Lse + batch_idx * stride_lse_b + head_idx, lse) + + +def mla_attention_decode( + q: torch.Tensor, + latent_cache: torch.Tensor, + block_table: torch.Tensor, + b_seq_len: torch.Tensor, + *, + softmax_scale: float, + kv_lora_rank: int = 512, + qk_rope_head_dim: int = 64, + page_size: int | None = None, + num_kv_splits: int = _DEFAULT_KV_SPLITS, + k_qdq: str | None = None, + k_qdq_amax: float | None = None, + v_qdq: str | None = None, + v_qdq_amax: float | None = None, + p_qdq: str | None = None, + p_qdq_amax: float = 1.0, + return_lse: bool = True, + out_dtype: torch.dtype | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Decode one absorbed query token per request over a paged latent cache. + + Quantized-BMM model: the latent cache holds RAW (bf16/fp16) values, and each + BMM operand is fake-quantized once, on read, along its own contraction axis + — K along the feature axis (``k_qdq``), V along the token axis (``v_qdq``), + independently. So decode honors ``k_format`` and ``v_format`` separately. + Q is fake-quantized by the caller (dynamic NVFP4 uses an FP32 carrier); P is + quantized in-kernel after the row-sum (denominator stays unquantized). + + Args: + q: ``[batch, num_heads, kv_lora_rank + qk_rope_head_dim]`` absorbed + query. Pass FP32 for the dynamic-NVFP4 QDQ carrier (Q is expected + to be fake-quantized by the caller); BF16/FP16 otherwise. + latent_cache: ``[num_blocks, page_size, kv_lora_rank + qk_rope_head_dim]`` + paged latent cache holding raw (unquantized) values. + block_table: ``[batch, max_blocks_per_seq]`` page table. + b_seq_len: ``[batch]`` KV sequence lengths. + softmax_scale: Softmax scale (required; MLA layers fold in mscale). + kv_lora_rank: Latent width (V/output width). + qk_rope_head_dim: RoPE feature width appended to the latent. + page_size: Tokens per page; defaults to ``latent_cache.shape[1]``. + num_kv_splits: Fixed split count. P/V QDQ numerics follow the + split-local schedule; kept fixed by default for reproducibility + across batch shapes and devices. + k_qdq: K fake quant-dequant: ``None``, ``"fp8"``, ``"nvfp4"`` (feature + axis, block-16 for NVFP4). Covers both the NOPE and RoPE slices. + k_qdq_amax: Per-tensor K amax (``None`` = scale 1.0). + v_qdq: V fake quant-dequant, same modes; token/contraction axis. The + open 16-token tail block re-quantizes as the sequence grows. + v_qdq_amax: Per-tensor V amax (``None`` = scale 1.0). + p_qdq: Softmax-P fake quant-dequant: ``None``, ``"fp8"``, ``"nvfp4"``. + p_qdq_amax: Per-tensor P amax (default 1.0, the theoretical bound). + return_lse: Also return the natural-log LSE ``[batch, num_heads]``. + out_dtype: Output dtype (default ``latent_cache.dtype``, the model + compute dtype expected by the V up-projection). + + Returns: + ``(out [batch, num_heads, kv_lora_rank], lse [batch, num_heads] | None)``. + """ + if q.ndim != 3: + raise ValueError(f"q must be [batch, heads, head_dim], got {tuple(q.shape)}") + if latent_cache.ndim != 3: + raise ValueError( + f"latent_cache must be [num_blocks, page_size, head_dim], " + f"got {tuple(latent_cache.shape)}" + ) + head_dim = kv_lora_rank + qk_rope_head_dim + if q.shape[2] != head_dim or latent_cache.shape[2] != head_dim: + raise ValueError( + f"q and latent_cache feature dims must equal kv_lora_rank + qk_rope_head_dim " + f"({head_dim}), got {q.shape[2]} and {latent_cache.shape[2]}" + ) + if page_size is None: + page_size = latent_cache.shape[1] + if page_size != latent_cache.shape[1]: + raise ValueError(f"page_size {page_size} must match latent_cache.shape[1]") + if not 1 <= num_kv_splits <= _MAX_KV_SPLITS: + raise ValueError(f"num_kv_splits must be in [1, {_MAX_KV_SPLITS}], got {num_kv_splits}") + nvfp4_active = "nvfp4" in (k_qdq, v_qdq, p_qdq) + if nvfp4_active and (kv_lora_rank % 16 or qk_rope_head_dim % 16): + raise ValueError("NVFP4 decode requires dimensions divisible by 16") + if v_qdq == "nvfp4" and _BLOCK_N % 16: + raise ValueError("NVFP4 V decode requires the KV tile (BLOCK_N) divisible by 16") + batch, num_heads = q.shape[0], q.shape[1] + if b_seq_len.shape != (batch,) or block_table.shape[0] != batch: + raise ValueError("decode metadata batch dimension must match q") + + # Operand "p" permits None|fp8|nvfp4 with the standard amax/448 (fp8) and + # amax/(6*448) (nvfp4) scales — shared by K, V, and P here. + k_qdq_scale = _qdq_scale(k_qdq, k_qdq_amax, "p") + v_qdq_scale = _qdq_scale(v_qdq, v_qdq_amax, "p") + p_qdq_scale = _qdq_scale(p_qdq, p_qdq_amax, "p") + q = q.contiguous() + if latent_cache.stride(-1) != 1: + raise ValueError("latent_cache last dim must be contiguous") + + block_dl = triton.next_power_of_2(kv_lora_rank) + block_dpe = triton.next_power_of_2(qk_rope_head_dim) + block_h = 16 if num_heads <= 16 else 32 + qk_scale = softmax_scale * LOG2E + if out_dtype is None: + out_dtype = latent_cache.dtype + + m_partial = torch.empty(batch, num_heads, num_kv_splits, dtype=torch.float32, device=q.device) + l_partial = torch.empty_like(m_partial) + acc_partial = torch.empty( + batch, num_heads, num_kv_splits, block_dl, dtype=torch.float32, device=q.device + ) + out = torch.empty(batch, num_heads, kv_lora_rank, dtype=out_dtype, device=q.device) + if return_lse: + lse = torch.empty(batch, num_heads, dtype=torch.float32, device=q.device) + else: + lse = torch.empty(1, dtype=torch.float32, device=q.device) + + with torch.cuda.device(q.device): + _mla_decode_split_kernel[(triton.cdiv(num_heads, block_h), batch, num_kv_splits)]( + q, + latent_cache, + block_table, + b_seq_len, + m_partial, + l_partial, + acc_partial, + qk_scale, + q.stride(0), + q.stride(1), + latent_cache.stride(0), + latent_cache.stride(1), + m_partial.stride(0), + m_partial.stride(1), + acc_partial.stride(0), + acc_partial.stride(1), + acc_partial.stride(2), + p_qdq_scale, + k_qdq_scale, + v_qdq_scale, + block_table.shape[1], + H=num_heads, + BLOCK_H=block_h, + BLOCK_N=_BLOCK_N, + BLOCK_DL=block_dl, + BLOCK_DPE=block_dpe, + KV_LORA_RANK=kv_lora_rank, + QK_ROPE_DIM=qk_rope_head_dim, + PAGE_SIZE=page_size, + NUM_KV_SPLITS=num_kv_splits, + K_QDQ=_P_QDQ_MODES[k_qdq], + V_QDQ=_P_QDQ_MODES[v_qdq], + P_QDQ=_P_QDQ_MODES[p_qdq], + Q_IS_FP32=q.dtype == torch.float32, + num_warps=4, + num_stages=2, + ) + _mla_decode_combine_kernel[(batch, num_heads)]( + m_partial, + l_partial, + acc_partial, + out, + lse, + m_partial.stride(0), + m_partial.stride(1), + acc_partial.stride(0), + acc_partial.stride(1), + acc_partial.stride(2), + out.stride(0), + out.stride(1), + lse.stride(0) if return_lse else 0, + BLOCK_DL=block_dl, + KV_LORA_RANK=kv_lora_rank, + NUM_KV_SPLITS=num_kv_splits, + STORE_LSE=return_lse, + num_warps=4, + ) + return out, (lse if return_lse else None) diff --git a/modelopt/torch/kernels/quantization/attention/mla/mla_prefill.py b/modelopt/torch/kernels/quantization/attention/mla/mla_prefill.py new file mode 100644 index 00000000000..dbd134812ee --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/mla/mla_prefill.py @@ -0,0 +1,461 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Varlen MLA prefill attention with fused Q/K/P/V fake quantization. + +MLA prefill runs standard multi-head attention over the up-projected per-head +K/V, but with asymmetric head dims: Q and K carry ``qk_nope_head_dim + +qk_rope_head_dim`` features (192 for DeepSeek-family models) while V and the +output carry ``v_head_dim`` (128). The kernel keeps two independent feature +axes so no V padding is required. + +Quantization follows the ModelOpt conventions: NVFP4 uses 1x16 blocks along +the BMM contraction axis (Q/K along the feature axis, P/V along the key/token +axis) with E4M3 block scales and a per-tensor global scale ``amax / (6*448)``; +FP8 is per-tensor E4M3 with scale ``amax / 448``. The softmax denominator +accumulates the unquantized P; only the quantized P feeds ``P @ V``. NVFP4 +operands run their dots on FP32 carriers with IEEE precision. RoPE feature +dims of Q/K participate in the same single quantization pass (no PE +special-casing). + +Causal masking treats Q as the suffix of the KV span (``k_len >= q_len``). +Cached-context chunks run with ``causal=False`` and ``return_lse=True``; the +caller (e.g. vLLM's MLA layer) merges chunk states via ``merge_attn_states`` +using the returned natural-log LSE. Inference-only: no autograd support. +""" + +import math + +import torch +import triton +import triton.language as tl + +from modelopt.torch.kernels.common.attention.triton_fa import LOG2E, _apply_mask +from modelopt.torch.kernels.quantization.attention.bmm2_qdq import ( + _a_qdq_nvfp4, + _p_qdq_nvfp4, + _v_qdq_nvfp4, +) +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq +from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( + _apply_sparse_nm_to_qk_tile, +) + +__all__ = ["mla_prefill_attention"] + +# Referenced inside @triton.jit code, so it must be a tl.constexpr global. +LN2 = tl.constexpr(0.6931471805599453) + +# Maps public operand quant options to kernel constexpr values. +_OPERAND_MODES = {None: 0, "fp8": 1, "nvfp4": 2} + + +def _operand_scale(mode: int, amax: float | None, name: str) -> float: + """Convert a per-tensor amax to the kernel scale for one operand.""" + if mode == 0 or amax is None: + return 1.0 + if not (math.isfinite(amax) and amax > 0): + raise ValueError(f"{name} must be a finite positive value, got {amax}") + return amax / 448.0 if mode == 1 else amax / (6.0 * 448.0) + + +def _resolve_mode(mode: str | None, name: str) -> int: + if mode not in _OPERAND_MODES: + raise ValueError( + f"{name} must be one of {sorted(m for m in _OPERAND_MODES if m)} or None, got {mode!r}" + ) + return _OPERAND_MODES[mode] + + +@triton.jit +def _qdq_a(x, MODE: tl.constexpr, scale, M: tl.constexpr, K: tl.constexpr): + """A-side (Q) operand QDQ: FP8 per-tensor or NVFP4 1x16 along K (axis 1).""" + if MODE == 1: + x = fp8_scalar_qdq(x, scale).to(x.dtype) + elif MODE == 2: + x = _a_qdq_nvfp4(x, scale, M, K) + return x + + +@triton.jit +def _qdq_b(x, MODE: tl.constexpr, scale, K: tl.constexpr, N: tl.constexpr): + """B-side (K^T / V) operand QDQ: FP8 per-tensor or NVFP4 1x16 along K (axis 0).""" + if MODE == 1: + x = fp8_scalar_qdq(x, scale).to(x.dtype) + elif MODE == 2: + x = _v_qdq_nvfp4(x, scale, K, N) + return x + + +@triton.jit +def _apply_sparse_nm_dense_tokens( + scores, + kv_start, + q_pos, + kv_pos, + seq_len_q, + seq_len_kv, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + SPARSITY_N: tl.constexpr, + SPARSITY_M: tl.constexpr, + DENSE_SINK_TOKENS: tl.constexpr, + DENSE_RECENT_TOKENS: tl.constexpr, +): + """Apply N:M sparsity outside token-exact sink and recent regions. + + Same semantics as ``triton_fa._apply_sparse_nm_with_dense_tokens`` (Q is + the suffix of the KV span), duplicated here because that helper resolves + the sparsity primitive through lazily-populated module globals. + """ + sparse_scores = _apply_sparse_nm_to_qk_tile(scores, BLOCK_M, BLOCK_N, SPARSITY_N, SPARSITY_M) + q_abs_pos = q_pos[:, None] + seq_len_kv - seq_len_q + kv_abs_pos = kv_start + kv_pos[None, :] + token_distance = q_abs_pos - kv_abs_pos + dense_tokens = ( + (seq_len_q <= 1) + | (kv_abs_pos < DENSE_SINK_TOKENS) + | ((token_distance >= 0) & (token_distance < DENSE_RECENT_TOKENS)) + ) + return tl.where(dense_tokens, scores, sparse_scores) + + +@triton.jit +def _mla_prefill_kernel( + Q, # [total_q, num_heads, LQK] + K, # [total_k, num_kv_heads, LQK] + V, # [total_k, num_kv_heads, LV] + Out, # [total_q, num_heads, LV] + Lse, # [num_heads, total_q] natural-log LSE (dummy when not RETURN_LSE) + qk_scale, # softmax_scale * log2(e) + Cu_seqlens_q, # [batch + 1] + Cu_seqlens_k, # [batch + 1] + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + stride_lse_h, + stride_lse_s, + q_scale, # runtime per-tensor scales (host-converted from amax) + k_scale, + p_scale, + v_scale, + kv_group_num: tl.constexpr, # num_heads // num_kv_heads + IS_CAUSAL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DQK: tl.constexpr, # next_power_of_2(LQK) + BLOCK_DV: tl.constexpr, # next_power_of_2(LV) + LQK: tl.constexpr, + LV: tl.constexpr, + RETURN_LSE: tl.constexpr, + Q_QUANT: tl.constexpr, # 0=off, 1=FP8 E4M3, 2=NVFP4 + K_QUANT: tl.constexpr, + P_QUANT: tl.constexpr, + V_QUANT: tl.constexpr, + IEEE_QK: tl.constexpr, # FP32 carriers + IEEE dot for BMM1 + IEEE_PV: tl.constexpr, # FP32 carriers + IEEE dot for BMM2 + SPARSITY_N: tl.constexpr = 0, # N:M sparsity - keep top-N of every M (0 = off) + SPARSITY_M: tl.constexpr = 4, + DENSE_SINK_TOKENS: tl.constexpr = 0, + DENSE_RECENT_TOKENS: tl.constexpr = 64, +): + # --- Grid: (batch, num_heads, num_q_tiles) --- + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + tile_q = tl.program_id(2) + kv_head_idx = head_idx // kv_group_num + + q_start = tl.load(Cu_seqlens_q + batch_idx) + q_end = tl.load(Cu_seqlens_q + batch_idx + 1) + k_start = tl.load(Cu_seqlens_k + batch_idx) + k_end = tl.load(Cu_seqlens_k + batch_idx + 1) + q_len = q_end - q_start + k_len = k_end - k_start + + if tile_q * BLOCK_M >= q_len: + return + + q_pos = tile_q * BLOCK_M + tl.arange(0, BLOCK_M) + kv_pos = tl.arange(0, BLOCK_N) + offs_dqk = tl.arange(0, BLOCK_DQK) + offs_dv = tl.arange(0, BLOCK_DV) + mask_dqk = offs_dqk < LQK + mask_dv = offs_dv < LV + q_mask = q_pos < q_len + + # --- Load Q tile [BLOCK_M, BLOCK_DQK]: stays in registers for the KV loop --- + q_ptrs = (q_start + q_pos)[:, None] * stride_qbs + head_idx * stride_qh + offs_dqk[None, :] + q = tl.load(Q + q_ptrs, mask=q_mask[:, None] & mask_dqk[None, :], other=0.0) + if IEEE_QK: + q = q.to(tl.float32) + q = _qdq_a(q, Q_QUANT, q_scale, BLOCK_M, BLOCK_DQK) + + # --- Online softmax state --- + row_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + row_sum = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32) + + # Causal: Q is the suffix of the KV span (k_len >= q_len). + causal_offset = k_len - q_len + kv_bound = k_len if not IS_CAUSAL else tl.minimum(causal_offset + (tile_q + 1) * BLOCK_M, k_len) + + for kv_start in range(0, kv_bound, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) + kv_valid = (kv_start + kv_pos) < k_len + + # K^T tile [BLOCK_DQK, BLOCK_N] + k_ptrs = ( + (k_start + kv_start + kv_pos)[None, :] * stride_kbs + + kv_head_idx * stride_kh + + offs_dqk[:, None] + ) + k = tl.load(K + k_ptrs, mask=kv_valid[None, :] & mask_dqk[:, None], other=0.0) + if IEEE_QK: + k = k.to(tl.float32) + k = _qdq_b(k, K_QUANT, k_scale, BLOCK_DQK, BLOCK_N) + + if IEEE_QK: + scores = tl.dot(q, k, input_precision="ieee") * qk_scale + else: + scores = tl.dot(q, k) * qk_scale + scores = _apply_mask(scores, q_pos, kv_pos, q_len, k_len, kv_start, IS_CAUSAL) + + if SPARSITY_N > 0: + scores = _apply_sparse_nm_dense_tokens( + scores, + kv_start, + q_pos, + kv_pos, + q_len, + k_len, + BLOCK_M, + BLOCK_N, + SPARSITY_N, + SPARSITY_M, + DENSE_SINK_TOKENS, + DENSE_RECENT_TOKENS, + ) + + # --- Online softmax update: the denominator uses unquantized p --- + m_new = tl.maximum(row_max, tl.max(scores, 1)) + p = tl.math.exp2(scores - m_new[:, None]) + l_new = tl.sum(p, 1) + correction = tl.math.exp2(row_max - m_new) + row_sum = row_sum * correction + l_new + acc = acc * correction[:, None] + + if P_QUANT == 1: + p = fp8_scalar_qdq(p, p_scale) + elif P_QUANT == 2: + # Native packing consumes the model dtype; the QDQ value stays FP32. + p = p.to(V.dtype.element_ty).to(tl.float32) + p = _p_qdq_nvfp4(p, p_scale, BLOCK_M, BLOCK_N) + + # V tile [BLOCK_N, BLOCK_DV] + v_ptrs = ( + (k_start + kv_start + kv_pos)[:, None] * stride_vbs + + kv_head_idx * stride_vh + + offs_dv[None, :] + ) + v = tl.load(V + v_ptrs, mask=kv_valid[:, None] & mask_dv[None, :], other=0.0) + if IEEE_PV: + v = v.to(tl.float32) + v = _qdq_b(v, V_QUANT, v_scale, BLOCK_N, BLOCK_DV) + + if IEEE_PV: + acc = tl.dot(p.to(tl.float32), v, acc, input_precision="ieee") + else: + acc = tl.dot(p.to(v.dtype), v, acc) + row_max = m_new + + # Clamp the denominator: empty context chunks (k_len == 0) leave acc at 0. + acc = acc / tl.maximum(row_sum[:, None], 1e-6) + + if RETURN_LSE: + lse = LN2 * (row_max + tl.math.log2(row_sum)) + lse = tl.where(row_sum == 0.0, float("-inf"), lse) + lse_ptrs = head_idx * stride_lse_h + (q_start + q_pos) * stride_lse_s + tl.store(Lse + lse_ptrs, lse, mask=q_mask) + + o_ptrs = (q_start + q_pos)[:, None] * stride_obs + head_idx * stride_oh + offs_dv[None, :] + tl.store(Out + o_ptrs, acc, mask=q_mask[:, None] & mask_dv[None, :]) + + +def mla_prefill_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float | None = None, + causal: bool = True, + return_lse: bool = False, + *, + q_quant: str | None = None, + k_quant: str | None = None, + p_quant: str | None = None, + v_quant: str | None = None, + q_amax: float | None = None, + k_amax: float | None = None, + p_amax: float = 1.0, + v_amax: float | None = None, + sparsity_n: int = 0, + sparsity_m: int = 4, + dense_sink_tokens: int = 0, + dense_recent_tokens: int = 64, + block_m: int = 64, + block_n: int = 64, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Varlen MLA prefill attention with fused Q/K/P/V fake quantization. + + Args: + q: ``[total_q, num_heads, qk_head_dim]`` packed queries. + k: ``[total_k, num_kv_heads, qk_head_dim]`` packed keys (NOPE ++ RoPE). + v: ``[total_k, num_kv_heads, v_head_dim]`` packed values. + cu_seqlens_q: ``[batch + 1]`` cumulative Q sequence lengths. + cu_seqlens_k: ``[batch + 1]`` cumulative K/V sequence lengths. + max_seqlen_q: Maximum Q sequence length (grid sizing). + softmax_scale: Scale factor (default ``qk_head_dim ** -0.5``). + causal: Causal masking; Q is treated as the suffix of the KV span. + Cached-context chunks use ``causal=False``. + return_lse: Also return the natural-log LSE ``[num_heads, total_q]`` + for chunk-state merging (``merge_attn_states``). + q_quant: Q fake quant-dequant: ``None``, ``"fp8"`` (per-tensor E4M3), + or ``"nvfp4"`` (1x16 blocks along the feature/contraction axis). + k_quant: K fake quant-dequant, same modes; blocks along features. + p_quant: Softmax-P fake quant-dequant, blocks along the key axis. + The softmax denominator stays unquantized. + v_quant: V fake quant-dequant, blocks along the key/token axis. + q_amax: Per-tensor amax for Q (``None`` = scale 1.0). + k_amax: Per-tensor amax for K (``None`` = scale 1.0). + p_amax: Per-tensor amax for P; defaults to 1.0, the theoretical upper + bound of the unnormalized P's amax. + v_amax: Per-tensor amax for V (``None`` = scale 1.0). + sparsity_n: N:M score sparsity along the key axis (0 = off). + sparsity_m: N:M group size (4 or 8). + dense_sink_tokens: Leading KV tokens kept dense (token-exact). + dense_recent_tokens: Recent KV tokens kept dense (token-exact). + block_m: Q tile size (multiple of 16). + block_n: KV tile size (multiple of 16). + + Returns: + Output ``[total_q, num_heads, v_head_dim]`` in ``v.dtype``; with + ``return_lse`` a tuple ``(out, lse)``. + """ + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("q, k, v must be [tokens, heads, head_dim] tensors") + total_q, num_heads, lqk = q.shape + total_k, num_kv_heads, lv = v.shape[0], v.shape[1], v.shape[2] + if k.shape != (total_k, num_kv_heads, lqk): + raise ValueError(f"k shape {tuple(k.shape)} != ({total_k}, {num_kv_heads}, {lqk})") + if num_heads % num_kv_heads: + raise ValueError("num_heads must be divisible by num_kv_heads") + if cu_seqlens_q.numel() != cu_seqlens_k.numel(): + raise ValueError("cu_seqlens_q and cu_seqlens_k must have the same length") + if block_m % 16 or block_n % 16: + raise ValueError("block_m and block_n must be multiples of 16") + + q_mode = _resolve_mode(q_quant, "q_quant") + k_mode = _resolve_mode(k_quant, "k_quant") + p_mode = _resolve_mode(p_quant, "p_quant") + v_mode = _resolve_mode(v_quant, "v_quant") + if (q_mode == 2 or k_mode == 2) and lqk % 16: + raise ValueError(f"NVFP4 Q/K requires qk_head_dim % 16 == 0, got {lqk}") + q_scale = _operand_scale(q_mode, q_amax, "q_amax") + k_scale = _operand_scale(k_mode, k_amax, "k_amax") + p_scale = _operand_scale(p_mode, p_amax, "p_amax") + v_scale = _operand_scale(v_mode, v_amax, "v_amax") + + if q.stride(-1) != 1: + q = q.contiguous() + if k.stride(-1) != 1: + k = k.contiguous() + if v.stride(-1) != 1: + v = v.contiguous() + + sm_scale = lqk**-0.5 if softmax_scale is None else softmax_scale + block_dqk = triton.next_power_of_2(lqk) + block_dv = triton.next_power_of_2(lv) + if block_dqk >= 512: + # Very wide QK tiles (e.g. absorbed 576-d shapes) blow the shared + # memory budget at 64x64 tiles (q + k^T staging ~ 3 * BLOCK_DQK KB); + # shrink to ~96 KB so the kernel fits 100 KB-class SMs. + block_m = min(block_m, 16) + block_n = min(block_n, 32) + batch = cu_seqlens_q.numel() - 1 + + out = torch.empty(total_q, num_heads, lv, dtype=v.dtype, device=q.device) + if return_lse: + lse = torch.empty(num_heads, total_q, dtype=torch.float32, device=q.device) + else: + lse = torch.empty(1, dtype=torch.float32, device=q.device) + + grid = (batch, num_heads, triton.cdiv(max(1, max_seqlen_q), block_m)) + with torch.cuda.device(q.device): + _mla_prefill_kernel[grid]( + q, + k, + v, + out, + lse, + sm_scale * LOG2E, + cu_seqlens_q, + cu_seqlens_k, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + out.stride(0), + out.stride(1), + lse.stride(0) if return_lse else 0, + lse.stride(1) if return_lse else 0, + q_scale, + k_scale, + p_scale, + v_scale, + kv_group_num=num_heads // num_kv_heads, + IS_CAUSAL=causal, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_DQK=block_dqk, + BLOCK_DV=block_dv, + LQK=lqk, + LV=lv, + RETURN_LSE=return_lse, + Q_QUANT=q_mode, + K_QUANT=k_mode, + P_QUANT=p_mode, + V_QUANT=v_mode, + IEEE_QK=(q_mode == 2 or k_mode == 2 or q.dtype == torch.float32), + IEEE_PV=(p_mode == 2 or v_mode == 2), + SPARSITY_N=sparsity_n, + SPARSITY_M=sparsity_m, + DENSE_SINK_TOKENS=dense_sink_tokens, + DENSE_RECENT_TOKENS=dense_recent_tokens, + # 256-wide FP32 tiles are register-heavy; revisit if tuning. # tune + num_warps=8 if block_dqk >= 256 else 4, + num_stages=1, + ) + if return_lse: + return out, lse + return out diff --git a/modelopt/torch/kernels/quantization/attention/mla/reference.py b/modelopt/torch/kernels/quantization/attention/mla/reference.py new file mode 100644 index 00000000000..a19f7f6c431 --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/mla/reference.py @@ -0,0 +1,389 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pure-torch oracles for the MLA attention kernels (no Triton imports). + +Independent eager re-implementations of the fake-quant math and the kernels' +tile schedules, used by the GPU golden tests. The FP4 rounding ladder mirrors +``fp4_round_magnitude`` (RNE ties-to-even on the E2M1 grid) and the two-level +NVFP4 scale mirrors ``fp8_quantize_scale``, but no production QDQ helper is +called here. +""" + +import math + +import torch + +LOG2E: float = 1.4426950408889634 +LN2: float = 0.6931471805599453 + +_E2M1_LEVELS = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def round_to_e2m1(x: torch.Tensor) -> torch.Tensor: + """Round to the nearest E2M1 value, ties to even (matches the kernels).""" + a = x.abs() + q = torch.where( + a <= 0.25, + torch.zeros_like(a), + torch.where( + a < 0.75, + torch.full_like(a, 0.5), + torch.where( + a <= 1.25, + torch.full_like(a, 1.0), + torch.where( + a < 1.75, + torch.full_like(a, 1.5), + torch.where( + a <= 2.5, + torch.full_like(a, 2.0), + torch.where( + a < 3.5, + torch.full_like(a, 3.0), + torch.where(a <= 5.0, torch.full_like(a, 4.0), torch.full_like(a, 6.0)), + ), + ), + ), + ), + ), + ) + return torch.where(x >= 0, q, -q) + + +def _quant_e4m3(x: torch.Tensor) -> torch.Tensor: + """Round-trip through FP8 E4M3 with saturation at +-448.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def nvfp4_fake_quant( + x: torch.Tensor, + amax: float | None = None, + block_axis: int = -1, + block: int = 16, +) -> torch.Tensor: + """Two-level NVFP4 fake quant: E2M1 elements, E4M3 block scales along one axis.""" + x = x.float() + x = x.movedim(block_axis, -1) + shape = x.shape + assert shape[-1] % block == 0, f"axis size {shape[-1]} not divisible by {block}" + g = x.reshape(*shape[:-1], shape[-1] // block, block) + block_amax = g.abs().amax(dim=-1, keepdim=True) + global_scale = 1.0 if amax is None else amax / (6.0 * 448.0) + scale = _quant_e4m3(block_amax / (6.0 * global_scale)) * global_scale + scale_safe = torch.where(scale == 0, torch.ones_like(scale), scale) + q = round_to_e2m1(g / scale_safe) * scale_safe + q = torch.where(scale == 0, torch.zeros_like(q), q) + return q.reshape(shape).movedim(-1, block_axis) + + +def fp8_tensor_fake_quant(x: torch.Tensor, amax: float | None = None) -> torch.Tensor: + """Per-tensor FP8 E4M3 fake quant with scale ``amax / 448``.""" + scale = 1.0 if amax is None else amax / 448.0 + return _quant_e4m3(x.float() / scale) * scale + + +def apply_operand_quant( + x: torch.Tensor, + mode: str | None, + amax: float | None, + block_axis: int = -1, +) -> torch.Tensor: + """Dispatch one operand's fake quant by mode name.""" + if mode is None: + return x.float() + if mode == "fp8": + return fp8_tensor_fake_quant(x, amax) + if mode == "nvfp4": + return nvfp4_fake_quant(x, amax, block_axis=block_axis) + raise ValueError(f"unknown quant mode {mode!r}") + + +def apply_sparse_nm( + scores: torch.Tensor, + sparsity_n: int, + sparsity_m: int, + q_abs_pos: torch.Tensor, + kv_abs_pos: torch.Tensor, + seq_len_q: int, + dense_sink_tokens: int, + dense_recent_tokens: int, +) -> torch.Tensor: + """Token-exact N:M score sparsity with dense sink/recent regions. + + ``scores`` is ``[..., Lq, Lk]`` with the key axis last; ``q_abs_pos`` and + ``kv_abs_pos`` are the absolute positions matching those two axes. + """ + shape = scores.shape + grouped = scores.reshape(*shape[:-1], shape[-1] // sparsity_m, sparsity_m) + topk = grouped.topk(sparsity_n, dim=-1).indices + keep = torch.zeros_like(grouped, dtype=torch.bool).scatter_(-1, topk, True) + sparse = torch.where(keep, grouped, torch.full_like(grouped, float("-inf"))) + sparse = sparse.reshape(shape) + distance = q_abs_pos[:, None] - kv_abs_pos[None, :] + dense = ( + (seq_len_q <= 1) + | (kv_abs_pos[None, :] < dense_sink_tokens) + | ((distance >= 0) & (distance < dense_recent_tokens)) + ) + return torch.where(dense, scores, sparse) + + +def _quantize_p_tile( + p: torch.Tensor, + mode: str | None, + amax: float, + carrier_dtype: torch.dtype, +) -> torch.Tensor: + """Quantize an unnormalized softmax tile ``[..., block_n]`` like the kernels do.""" + if mode is None: + return p + if mode == "fp8": + return fp8_tensor_fake_quant(p, amax) + if mode == "nvfp4": + p = p.to(carrier_dtype).float() + return nvfp4_fake_quant(p, amax, block_axis=-1) + raise ValueError(f"unknown quant mode {mode!r}") + + +def mla_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + softmax_scale: float | None = None, + causal: bool = True, + *, + q_mode: str | None = None, + k_mode: str | None = None, + p_mode: str | None = None, + v_mode: str | None = None, + q_amax: float | None = None, + k_amax: float | None = None, + p_amax: float = 1.0, + v_amax: float | None = None, + sparsity_n: int = 0, + sparsity_m: int = 4, + dense_sink_tokens: int = 0, + dense_recent_tokens: int = 64, + block_n: int = 64, + return_lse: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Eager oracle for :func:`mla_prefill_attention`. + + Replicates the kernel schedule: Q/K quantized along the feature axis and + V along the token axis up front; P quantized per ``block_n`` KV tile on + the unnormalized online-softmax probabilities, with the denominator kept + unquantized. Returns FP32 output ``[total_q, num_heads, LV]`` (and + natural-log LSE ``[num_heads, total_q]`` with ``return_lse``). + """ + total_q, num_heads, lqk = q.shape + num_kv_heads, lv = k.shape[1], v.shape[2] + scale = lqk**-0.5 if softmax_scale is None else softmax_scale + carrier_dtype = v.dtype + batch = cu_seqlens_q.numel() - 1 + + out = torch.zeros(total_q, num_heads, lv, dtype=torch.float32, device=q.device) + lse_out = torch.full((num_heads, total_q), float("-inf"), dtype=torch.float32, device=q.device) + + for b in range(batch): + q0, q1 = int(cu_seqlens_q[b]), int(cu_seqlens_q[b + 1]) + k0, k1 = int(cu_seqlens_k[b]), int(cu_seqlens_k[b + 1]) + lq, lk = q1 - q0, k1 - k0 + if lq == 0: + continue + qb = apply_operand_quant(q[q0:q1], q_mode, q_amax, block_axis=-1) # [Lq, H, LQK] + kb = apply_operand_quant(k[k0:k1], k_mode, k_amax, block_axis=-1) # [Lk, Hkv, LQK] + # V groups along tokens: the kernel's masked (zero) tail rows are + # equivalent to zero-padding the sequence to a multiple of 16. + vb = v[k0:k1].float() + if v_mode is not None: + pad = (-lk) % 16 + vb = torch.nn.functional.pad(vb, (0, 0, 0, 0, 0, pad)) + vb = apply_operand_quant(vb, v_mode, v_amax, block_axis=0)[:lk] + if num_kv_heads != num_heads: + kb = kb.repeat_interleave(num_heads // num_kv_heads, dim=1) + vb = vb.repeat_interleave(num_heads // num_kv_heads, dim=1) + + qh = qb.permute(1, 0, 2) # [H, Lq, LQK] + kh = kb.permute(1, 2, 0) # [H, LQK, Lk] + vh = vb.permute(1, 0, 2) # [H, Lk, LV] + + q_pos = torch.arange(lq, device=q.device) + row_max = torch.full((num_heads, lq), float("-inf"), device=q.device) + row_sum = torch.zeros(num_heads, lq, device=q.device) + acc = torch.zeros(num_heads, lq, lv, device=q.device) + + for kv_start in range(0, lk, block_n): + kv_end = min(kv_start + block_n, lk) + kv_abs = torch.arange(kv_start, kv_start + block_n, device=q.device) + scores = torch.matmul(qh, kh[:, :, kv_start:kv_end]) * scale * LOG2E + if kv_end - kv_start < block_n: # pad tile to block_n like the kernel + pad = block_n - (kv_end - kv_start) + scores = torch.nn.functional.pad(scores, (0, pad), value=float("-inf")) + valid = kv_abs < lk + if causal: + allowed = (q_pos[:, None] + (lk - lq)) >= kv_abs[None, :] + mask = valid[None, :] & allowed + else: + mask = valid[None, :].expand(lq, block_n) + scores = scores.masked_fill(~mask[None], float("-inf")) + if sparsity_n > 0: + scores = apply_sparse_nm( + scores, + sparsity_n, + sparsity_m, + q_pos + (lk - lq), + kv_abs, + lq, + dense_sink_tokens, + dense_recent_tokens, + ) + scores = scores.masked_fill(~mask[None], float("-inf")) + + m_new = torch.maximum(row_max, scores.amax(dim=-1)) + # Rows with no valid keys yet keep -inf; guard exp2 of (-inf) - (-inf). + shifted = scores - m_new.unsqueeze(-1) + p = torch.where(torch.isneginf(scores), torch.zeros_like(scores), torch.exp2(shifted)) + correction = torch.where( + torch.isneginf(row_max), torch.zeros_like(row_max), torch.exp2(row_max - m_new) + ) + row_sum = row_sum * correction + p.sum(dim=-1) + acc = acc * correction.unsqueeze(-1) + p_q = _quantize_p_tile(p, p_mode, p_amax, carrier_dtype) + v_tile = vh[:, kv_start:kv_end] + if kv_end - kv_start < block_n: + v_tile = torch.nn.functional.pad(v_tile, (0, 0, 0, block_n - (kv_end - kv_start))) + if p_mode != "nvfp4" and v_mode != "nvfp4": + # Non-IEEE kernel path dots p (and unquantized v) in the + # compute dtype; the IEEE path keeps p in FP32. + p_q = p_q.to(carrier_dtype).float() + v_tile = v_tile.to(carrier_dtype).float() if v_mode is None else v_tile + acc = acc + torch.matmul(p_q, v_tile) + row_max = m_new + + out[q0:q1] = (acc / row_sum.clamp_min(1e-6).unsqueeze(-1)).permute(1, 0, 2) + lse = LN2 * (row_max + torch.log2(row_sum.clamp_min(1e-30))) + lse = torch.where(row_sum == 0, torch.full_like(lse, float("-inf")), lse) + lse_out[:, q0:q1] = lse + + if return_lse: + return out, lse_out + return out + + +def mla_decode_reference( + q: torch.Tensor, + latent: torch.Tensor, + seq_lens: torch.Tensor, + softmax_scale: float, + kv_lora_rank: int, + *, + k_mode: str | None = None, + k_amax: float | None = None, + v_mode: str | None = None, + v_amax: float | None = None, + p_mode: str | None = None, + p_amax: float = 1.0, + num_kv_splits: int = 32, + block_n: int = 32, + return_lse: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Eager split-local oracle for :func:`mla_attention_decode`. + + ``latent`` is the dense (unpaged) RAW cache ``[batch, max_seq, head_dim]``; + ``q`` is the absorbed query ``[batch, num_heads, head_dim]``. Replicates the + fixed-split, tile-aligned stage-1 schedule and stage-2 merge, with on-read + K (feature axis) and V (token axis) quant applied independently per tile. + """ + batch, num_heads, head_dim = q.shape + carrier_dtype = latent.dtype + out = torch.zeros(batch, num_heads, kv_lora_rank, dtype=torch.float32, device=q.device) + lse_out = torch.zeros(batch, num_heads, dtype=torch.float32, device=q.device) + + qf = q.float() + q_nope, q_pe = qf[..., :kv_lora_rank], qf[..., kv_lora_rank:] + + for b in range(batch): + s = int(seq_lens[b]) + k_all = latent[b, :s].float() # [s, head_dim] + k_nope, k_pe = k_all[:, :kv_lora_rank], k_all[:, kv_lora_rank:] + num_tiles = math.ceil(s / block_n) + tiles_per_split = math.ceil(num_tiles / num_kv_splits) + + split_m, split_l, split_acc = [], [], [] + for split in range(num_kv_splits): + kv_lo = split * tiles_per_split * block_n + kv_hi = min(kv_lo + tiles_per_split * block_n, s) + m = torch.full((num_heads,), float("-inf"), device=q.device) + lsum = torch.zeros(num_heads, device=q.device) + acc = torch.zeros(num_heads, kv_lora_rank, device=q.device) + for kv_start in range(kv_lo, kv_hi, block_n): + kv_end = min(kv_start + block_n, s) + valid = kv_end - kv_start + kn_raw = k_nope[kv_start:kv_end] # [valid, rank] raw + kp_raw = k_pe[kv_start:kv_end] # [valid, rope] raw + if valid < block_n: + # Pad to block_n like the kernel (masked -> 0); the padded + # rows get -inf scores below so they never contribute. + pad = block_n - valid + kn_raw = torch.nn.functional.pad(kn_raw, (0, 0, 0, pad)) + kp_raw = torch.nn.functional.pad(kp_raw, (0, 0, 0, pad)) + # Independent on-read operand quant from the raw latent: + # K along the feature axis, V along the token axis. + kn_k = apply_operand_quant(kn_raw, k_mode, k_amax, block_axis=1) + kp_k = apply_operand_quant(kp_raw, k_mode, k_amax, block_axis=1) + v_op = apply_operand_quant(kn_raw, v_mode, v_amax, block_axis=0) + scores = (q_nope[b] @ kn_k.T + q_pe[b] @ kp_k.T) * softmax_scale * LOG2E + if valid < block_n: + scores[:, valid:] = float("-inf") + m_new = torch.maximum(m, scores.amax(dim=-1)) + shifted = scores - m_new.unsqueeze(-1) + p = torch.where( + torch.isneginf(scores), torch.zeros_like(scores), torch.exp2(shifted) + ) + correction = torch.where( + torch.isneginf(m), torch.zeros_like(m), torch.exp2(m - m_new) + ) + lsum = lsum * correction + p.sum(dim=-1) + acc = acc * correction.unsqueeze(-1) + p_q = _quantize_p_tile(p, p_mode, p_amax, carrier_dtype) + if p_mode != "nvfp4": + p_q = p_q.to(carrier_dtype).float() + acc = acc + p_q @ v_op + m = m_new + split_m.append(m) + split_l.append(lsum) + split_acc.append(acc) + + m = torch.full((num_heads,), float("-inf"), device=q.device) + lsum = torch.zeros(num_heads, device=q.device) + acc = torch.zeros(num_heads, kv_lora_rank, device=q.device) + for sm, sl, sa in zip(split_m, split_l, split_acc): + has = sl > 0 + new_max = torch.where(has, torch.maximum(m, sm), m) + corr = torch.where(torch.isneginf(m), torch.zeros_like(m), torch.exp2(m - new_max)) + scorr = torch.where(has, torch.exp2(sm - new_max), torch.zeros_like(sm)) + acc = acc * corr.unsqueeze(-1) + sa * scorr.unsqueeze(-1) + lsum = lsum * corr + sl * scorr + m = new_max + out[b] = acc / lsum.clamp_min(1e-6).unsqueeze(-1) + lse = LN2 * (m + torch.log2(lsum.clamp_min(1e-30))) + lse_out[b] = torch.where(lsum == 0, torch.full_like(lse, float("-inf")), lse) + + if return_lse: + return out, lse_out + return out diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 749d6b10ee5..b2a15293e7e 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -84,6 +84,52 @@ def build_vllm_attention_quant_cfg( ] +# MLA operand -> quantizer-name mapping: k_format governs the write-once +# latent-cache QDQ (kv_c, k_pe) and the prefill projected K; v_format governs +# the prefill projected V (decode BMM2 consumes the quantized cache as-is). +_MLA_FORMAT_QUANTIZERS = { + "q": ("q",), + "k": ("kv_c", "k_pe", "k_mha"), + "p": ("p",), + "v": ("v_mha",), +} + + +def build_vllm_mla_attention_quant_cfg( + *, + q_format: str = "nvfp4", + k_format: str = "nvfp4", + p_format: str = "nvfp4", + v_format: str = "nvfp4", +) -> list: + """Build the MLA attention quantizer config with per-operand formats. + + Same format set as :func:`build_vllm_attention_quant_cfg`; the MLA operand + knobs fan out per :data:`_MLA_FORMAT_QUANTIZERS`. + """ + formats = {"q": q_format, "k": k_format, "p": p_format, "v": v_format} + for name, fmt in formats.items(): + if fmt not in _BMM2_FORMAT_CFGS: + raise ValueError( + f"{name}_format must be one of {sorted(_BMM2_FORMAT_CFGS)}, got {fmt!r}" + ) + return [ + {"quantizer_name": "*_bmm_quantizer", "enable": False}, + *( + { + "quantizer_name": f"*{quantizer}_bmm_quantizer", + "cfg": _BMM2_FORMAT_CFGS[fmt], + "enable": True, + } + for name, fmt in formats.items() + for quantizer in _MLA_FORMAT_QUANTIZERS[name] + ), + ] + + +_VLLM_NVFP4_MLA_ATTENTION_QUANT_CFG = build_vllm_mla_attention_quant_cfg() + + def _import_attention_module(): """Import a vLLM module that exports the concrete ``Attention`` class.""" for module_name in ( @@ -324,7 +370,14 @@ def vllm_replace_quant_module_hook(model: torch.nn.Module) -> None: def _set_vllm_attention_kv_default_amax(module, device: torch.device) -> None: """Set a global-scale-one amax on uncalibrated block-16 NVFP4 K/V quantizers.""" - for name in ("k_bmm_quantizer", "v_bmm_quantizer"): + for name in ( + "k_bmm_quantizer", + "v_bmm_quantizer", + "kv_c_bmm_quantizer", + "k_pe_bmm_quantizer", + "k_mha_bmm_quantizer", + "v_mha_bmm_quantizer", + ): quantizer = getattr(module, name, None) if ( not isinstance(quantizer, TensorQuantizer) @@ -344,6 +397,10 @@ def _set_vllm_attention_fp8_bmm2_default_amax(module, device: torch.device) -> N ("k_bmm_quantizer", 448.0), ("p_bmm_quantizer", 1.0), ("v_bmm_quantizer", 448.0), + ("kv_c_bmm_quantizer", 448.0), + ("k_pe_bmm_quantizer", 448.0), + ("k_mha_bmm_quantizer", 448.0), + ("v_mha_bmm_quantizer", 448.0), ): quantizer = getattr(module, name, None) if ( @@ -398,6 +455,51 @@ def configure_vllm_nvfp4_attention_quantizers( return module +def configure_vllm_nvfp4_mla_quantizers( + module: torch.nn.Module, + *, + device: torch.device | str, + dtype: torch.dtype, + cfg: list | None = None, +) -> torch.nn.Module: + """Configure one vLLM ``MLAAttention`` module for fused fake quantization. + + MLA counterpart of :func:`configure_vllm_nvfp4_attention_quantizers`. + ``kv_c``/``k_pe`` quantizers stay module-level (write-once latent-cache QDQ + — the single representation both decode BMMs consume); ``k``/``v`` + quantizers drive the prefill kernel's projected-K/V QDQ; ``q`` and ``p`` + drive both phases. The caller remains responsible for installing the + ModelOpt MLA impl and setting ``_query_quant_in_kernel``. + + Args: + module: A vLLM ``MLAAttention`` module to convert and configure in place. + device: Device on which the attention quantizer state should reside. + dtype: Model compute dtype associated with the attention module. + cfg: Optional quantizer config (default: all-NVFP4). + + Returns: + The supplied module, converted in place to ``_QuantVLLMMLAAttention``. + """ + if VllmMLAAttention is None or not isinstance(module, VllmMLAAttention): + raise TypeError(f"Expected vLLM MLAAttention, got {type(module).__name__}") + if not isinstance(dtype, torch.dtype): + raise TypeError(f"Expected torch.dtype, got {type(dtype).__name__}") + + device = torch.device(device) + module.device, module.dtype = device, dtype + if not isinstance(module, _QuantVLLMMLAAttention): + module = QuantModuleRegistry.convert(module) + if not hasattr(module, "p_bmm_quantizer"): + module.p_bmm_quantizer = TensorQuantizer() + + set_quantizer_by_cfg(module, _VLLM_NVFP4_MLA_ATTENTION_QUANT_CFG if cfg is None else cfg) + for name in ("q", "kv_c", "k_pe", "k_mha", "p", "v_mha"): + getattr(module, f"{name}_bmm_quantizer").to(device=device) + _set_vllm_attention_kv_default_amax(module, device) + _set_vllm_attention_fp8_bmm2_default_amax(module, device) + return module + + def _vllm_attention_modelopt_post_restore(self) -> None: """Move Attention module to its correct device after ModelOpt state restore.""" device, dtype = _get_device_dtype(self) @@ -765,12 +867,30 @@ def _setup(self): self.q_bmm_quantizer = TensorQuantizer() self.kv_c_bmm_quantizer = TensorQuantizer() self.k_pe_bmm_quantizer = TensorQuantizer() + # Prefill-only quantizers for the up-projected per-head K/V; the + # ModelOpt MLA kernels consume them in-kernel. Disabled by default + # so pre-existing MLA checkpoints restore unchanged. + self.k_mha_bmm_quantizer = TensorQuantizer() + self.k_mha_bmm_quantizer.disable() + self.v_mha_bmm_quantizer = TensorQuantizer() + self.v_mha_bmm_quantizer.disable() self.parallel_state = create_parallel_state() def forward(self, query, kv_c, k_pe, *args, **kwargs): - query = self.q_bmm_quantizer(query) - kv_c = self.kv_c_bmm_quantizer(kv_c) - k_pe = self.k_pe_bmm_quantizer(k_pe) + # With in-kernel Q quantization, prefill quantizes the projected + # 192-d q inside the kernel and decode quantizes the absorbed + # 576-d q (FP32 carrier) — skip the module-level QDQ entirely. + if not getattr(self, "_query_quant_in_kernel", False): + query = self.q_bmm_quantizer(query) + # Latent-cache QDQ (models a quantized latent cache; used by the + # eager fakequant / calibration path). The fused quantized-BMM MLA + # impl instead keeps the cache RAW and quantizes each BMM operand + # in-kernel (prefill projected K/V; decode on-read latent K/V), so + # it sets ``_skip_module_kv_quant`` to bypass this and let the bf16 + # latent reach both the cache write and vLLM's prefill projection. + if not getattr(self, "_skip_module_kv_quant", False): + kv_c = self.kv_c_bmm_quantizer(kv_c) + k_pe = self.k_pe_bmm_quantizer(k_pe) return super().forward(query, kv_c, k_pe, *args, **kwargs) def modelopt_post_restore(self, prefix: str = "") -> None: diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mla.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mla.py new file mode 100644 index 00000000000..d42156411f6 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mla.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ModelOpt MLA attention adapter for vLLM's TRITON_MLA backend. + +Same integration philosophy as the regular-attention adapter in +``plugins/vllm.py``: the ``MLAAttention`` module stays intact and only +``layer.impl`` is reclassed. vLLM keeps owning projections, RoPE, cache writes, +metadata, chunked-context gathering, state merging, and the V up-projection. + +Quantized-BMM model (aligned with the MNI reference and the dense Q/K/P/V +path): the latent cache stays RAW, and every BMM operand is fake-quantized +once, in-kernel, along its own contraction axis. The module sets +``_skip_module_kv_quant`` so the module-level ``kv_c``/``k_pe`` quantizers do +not touch the cache. + +- ``forward_mha`` (all prefill) temporarily swaps the per-layer ModelOpt + prefill backend into the prefill metadata and delegates to the inherited + implementation, so the ``kv_b_proj`` projection and chunked-context plumbing + are reused rather than forked. The backend routes the two attention calls + (causal new tokens; non-causal context chunks with LSE) to + :func:`mla_prefill_attention`, which quantizes the projected Q/K/P/V operands + once (from the bf16 latent) with optional 2:4 score sparsity (prefill only). +- ``forward_mqa`` (decode) fake-quantizes the absorbed query (FP32 QDQ + carrier) and calls :func:`mla_attention_decode`, which reads the raw latent + and quantizes K (feature axis), V (token axis), and P independently in the + kernel — so decode honors ``k_format`` and ``v_format`` separately. +""" + +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +from vllm.v1.attention.backends.mla.triton_mla import TritonMLAImpl + +from modelopt.torch.kernels.quantization.attention.mla import ( + mla_attention_decode, + mla_prefill_attention, +) + +from . import vllm as attention_plugin + +__all__ = [ + "ModelOptMLAImpl", + "clone_mla_impl", + "mla_quant_kw_from_layer", + "select_mla_impl_cls", +] + +# Fixed split count: decode P QDQ numerics follow the split-local schedule, so +# a fixed value keeps results reproducible across batch shapes and devices. +_MLA_DECODE_NUM_KV_SPLITS = 32 + + +@dataclass(frozen=True, slots=True) +class _MLAQuantKw: + """Kernel quantization kwargs resolved once per layer at install time.""" + + prefill: dict[str, Any] + decode: dict[str, Any] + + @property + def any_active(self) -> bool: + return ( + any(self.prefill[f"{op}_quant"] is not None for op in ("q", "k", "p", "v")) + or self.decode["p_qdq"] is not None + ) + + +def mla_quant_kw_from_layer(layer, *, query_in_kernel: bool) -> _MLAQuantKw: + """Resolve the MLA kernels' quantization kwargs from the layer's quantizers. + + Quantized-BMM model: the latent cache is RAW, and ``k_format``/``v_format`` + (the ``k_mha``/``v_mha`` quantizers) drive the in-kernel operand quant for + BOTH prefill (projected K/V) and decode (on-read latent K feature-axis / V + token-axis). The module-level ``kv_c``/``k_pe`` quantizers do not quantize + the cache on this path (see ``_skip_module_kv_quant``). With + ``query_in_kernel`` False (FP8 Q), the module already QDQ'd the query, so + neither kernel applies a Q transform. + """ + q_qdq, q_amax = attention_plugin._bmm_qdq_from_layer(layer, "q_bmm_quantizer", None) + k_qdq, k_amax = attention_plugin._bmm_qdq_from_layer(layer, "k_mha_bmm_quantizer", None) + v_qdq, v_amax = attention_plugin._bmm_qdq_from_layer(layer, "v_mha_bmm_quantizer", None) + p_qdq, p_amax = attention_plugin._p_qdq_from_layer(layer) + return _MLAQuantKw( + prefill={ + "q_quant": q_qdq if query_in_kernel else None, + "q_amax": q_amax, + "k_quant": k_qdq, + "k_amax": k_amax, + "p_quant": p_qdq, + "p_amax": p_amax, + "v_quant": v_qdq, + "v_amax": v_amax, + }, + decode={ + "p_qdq": p_qdq, + "p_qdq_amax": p_amax, + "k_qdq": k_qdq, + "k_qdq_amax": k_amax, + "v_qdq": v_qdq, + "v_qdq_amax": v_amax, + }, + ) + + +class _ModelOptMLAPrefillBackend(MLAPrefillBackend): + """Per-layer prefill backend carrying this layer's quant/sparse kwargs. + + Swapped into ``prefill_metadata.prefill_backend`` for the duration of one + ``forward_mha`` call (the builder-stamped backend is shared across layers, + so per-layer state cannot live there permanently). + """ + + def __init__(self, base_backend: MLAPrefillBackend, quant_kw: dict, sparse_kw: dict): + super().__init__( + num_heads=base_backend.num_heads, + scale=base_backend.scale, + kv_lora_rank=base_backend.kv_lora_rank, + qk_nope_head_dim=base_backend.qk_nope_head_dim, + qk_rope_head_dim=base_backend.qk_rope_head_dim, + v_head_dim=base_backend.v_head_dim, + vllm_config=base_backend.vllm_config, + ) + self._quant_kw = dict(quant_kw) + self._sparse_kw = dict(sparse_kw) + + @staticmethod + def get_name() -> str: + return "MODELOPT_MLA" + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + out: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if output_scale is not None: + raise NotImplementedError( + "ModelOpt MLA attention does not support fused FP8 output quantization" + ) + pm = self._prefill_metadata + result = mla_prefill_attention( + q, + k, + v, + cu_seqlens_q=pm.query_start_loc, + cu_seqlens_k=pm.query_start_loc, + max_seqlen_q=pm.max_query_len, + softmax_scale=self.scale, + causal=True, + return_lse=return_softmax_lse, + **self._quant_kw, + **self._sparse_kw, + ) + if out is not None and not return_softmax_lse: + out.copy_(result) + return out + return result + + def run_prefill_context_chunk( + self, + chunk_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + pm = self._prefill_metadata + assert pm.chunked_context is not None + # No sparse kwargs here: the 2:4 dense-window semantics are + # suffix-relative, so cached-context chunks run dense. + return mla_prefill_attention( + q, + k, + v, + cu_seqlens_q=pm.query_start_loc, + cu_seqlens_k=pm.chunked_context.cu_seq_lens[chunk_idx], + max_seqlen_q=pm.max_query_len, + softmax_scale=self.scale, + causal=False, + return_lse=True, + **self._quant_kw, + ) + + +class ModelOptMLAImpl(TritonMLAImpl): + """TRITON_MLA impl adapter routing prefill and decode to ModelOpt kernels. + + Instances are created by :func:`clone_mla_impl` (state-preserving reclass); + ``quant_kw`` (an :class:`_MLAQuantKw`) and ``sparse_kw`` are stashed by the + installer before the impl is published on the layer. + """ + + quant_kw: _MLAQuantKw + sparse_kw: dict[str, Any] + + # do_kv_cache_update is intentionally NOT overridden: the latent cache stays + # RAW (bf16). Decode quantizes K/V on read and prefill quantizes the + # projected operands in-kernel — the quantized-BMM model — so nothing + # quantizes the stored latent. + + def _get_prefill_backend(self, prefill_metadata) -> _ModelOptMLAPrefillBackend: + backend = self.__dict__.get("_modelopt_prefill_backend") + if backend is None: + backend = _ModelOptMLAPrefillBackend( + prefill_metadata.prefill_backend, + quant_kw=self.quant_kw.prefill, + sparse_kw=self.sparse_kw, + ) + self._modelopt_prefill_backend = backend + return backend + + def forward_mha( + self, + q, + kv_c_normed, + k_pe, + kv_c_and_k_pe_cache, + attn_metadata, + k_scale, + output, + output_scale=None, + ) -> None: + """Run MLA prefill with the ModelOpt prefill backend swapped in.""" + prefill_metadata = attn_metadata.prefill + assert prefill_metadata is not None + if not self.quant_kw.any_active and not self.sparse_kw: + return super().forward_mha( + q, + kv_c_normed, + k_pe, + kv_c_and_k_pe_cache, + attn_metadata, + k_scale, + output, + output_scale, + ) + backend = self._get_prefill_backend(prefill_metadata) + backend.prepare_metadata(prefill_metadata) + saved = prefill_metadata.prefill_backend + prefill_metadata.prefill_backend = backend + try: + return super().forward_mha( + q, + kv_c_normed, + k_pe, + kv_c_and_k_pe_cache, + attn_metadata, + k_scale, + output, + output_scale, + ) + finally: + prefill_metadata.prefill_backend = saved + + def forward_mqa(self, q, kv_c_and_k_pe_cache, attn_metadata, layer): + """Run absorbed-MLA decode through the ModelOpt split-K kernel.""" + query_in_kernel = getattr(layer, "_query_quant_in_kernel", False) + dec = self.quant_kw.decode + if ( + dec["p_qdq"] is None + and dec["k_qdq"] is None + and dec["v_qdq"] is None + and not query_in_kernel + ): + return super().forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) + if isinstance(q, tuple): + q = torch.cat(q, dim=-1) + if query_in_kernel: + # FP32 QDQ carrier on the absorbed query: the hardware-faithful + # emulation of the decode BMM1 A-operand. + q = layer.q_bmm_quantizer(q.float()) + decode_meta = attn_metadata.decode + assert decode_meta is not None + return mla_attention_decode( + q, + kv_c_and_k_pe_cache, + decode_meta.block_table, + decode_meta.seq_lens, + softmax_scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + num_kv_splits=_MLA_DECODE_NUM_KV_SPLITS, + out_dtype=kv_c_and_k_pe_cache.dtype, + return_lse=True, + **self.quant_kw.decode, + ) + + +def select_mla_impl_cls(impl) -> type | None: + """Return the ModelOpt MLA adapter class matching a native implementation.""" + if isinstance(impl, ModelOptMLAImpl): + return type(impl) + if isinstance(impl, TritonMLAImpl): + return ModelOptMLAImpl + return None + + +def clone_mla_impl(old_impl) -> ModelOptMLAImpl: + """Create the MLA adapter while preserving vLLM's initialized impl state.""" + new_cls = select_mla_impl_cls(old_impl) + if new_cls is None: + raise TypeError( + f"MLA backend {type(old_impl).__name__} is not supported; launch vLLM with " + "--attention-backend TRITON_MLA" + ) + new_impl = object.__new__(new_cls) + new_impl.__dict__.update(vars(old_impl)) + # A re-install over an existing adapter must not inherit the cached prefill + # backend, which froze the previous install's quant/sparse kwargs. + new_impl.__dict__.pop("_modelopt_prefill_backend", None) + return new_impl diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index b4141740b64..a0c9c739547 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -55,6 +55,30 @@ def _import_attention_type() -> type: _VLLM_ATTENTION = _import_attention_type() +def _import_mla_attention_type() -> type | None: + """Import the concrete vLLM MLAAttention type, or None when unavailable.""" + for module_name in ( + "vllm.attention.layer", + "vllm.model_executor.layers.attention", + "vllm.attention", + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + if hasattr(module, "MLAAttention"): + return module.MLAAttention + return None + + +def _load_mla_plugin(): + # MLA support is optional: imported lazily so regular FlashAttention and + # FlashInfer users never acquire the TRITON_MLA import requirements. + from . import vllm_mla + + return vllm_mla + + @dataclass(frozen=True, slots=True) class VllmAttentionInstallReport: """Summary of attention modules changed by a vLLM runtime installation.""" @@ -86,6 +110,7 @@ class _AttentionPlan: device: object | None dtype: torch.dtype | None requires_flashinfer_patch: bool + is_mla: bool = False @dataclass(frozen=True, slots=True) @@ -233,6 +258,37 @@ def _layer_errors(module) -> list[str]: return errors +def _layer_errors_mla(module, sparse_kw: dict[str, Any]) -> list[str]: + """Per-layer gates for MLA attention (decoder-only latent attention).""" + errors = [] + if str(getattr(module, "kv_cache_dtype", "auto")).startswith("fp8"): + errors.append("FP8 KV cache is unsupported for MLA") + if getattr(module, "use_sparse", False) or getattr(module, "indexer", None) is not None: + errors.append("sparse-indexer MLA (DeepSeek V3.2-style) is unsupported") + num_heads = getattr(module, "num_heads", None) + q_pad = getattr(module, "q_pad_num_heads", None) + if q_pad not in (None, num_heads): + errors.append(f"q_pad_num_heads={q_pad!r} padding is unsupported") + for dim_name in ("kv_lora_rank", "qk_rope_head_dim", "qk_head_dim", "v_head_dim"): + dim = getattr(module, dim_name, None) + if not isinstance(dim, int) or dim <= 0 or dim % 16: + errors.append(f"{dim_name}={dim!r} must be a positive multiple of 16") + if "skip_softmax_threshold" in sparse_kw or "threshold_scale_factor" in sparse_kw: + errors.append("skip-softmax is unsupported on MLA layers (N:M sparsity only)") + return errors + + +def _select_new_mla_impl(module) -> tuple[object | None, str | None]: + try: + mla_plugin = _load_mla_plugin() + except ImportError as err: + return None, f"MLA attention support requires a TRITON_MLA-capable vLLM: {err}" + try: + return mla_plugin.clone_mla_impl(module.impl), None + except (NotImplementedError, TypeError) as err: + return None, str(err) + + def _device_capability_error(device) -> str | None: if device is None: return None @@ -331,15 +387,19 @@ def _plan_vllm_attention( ) -> _InstallPlan: model = _unwrapped_model(model_runner) resolved_sparse_cfg, sparse_algorithm = _resolve_sparse_config(model_runner, sparse_cfg) + # MLA layers are quantize-only: the sparse-only entry point keeps ignoring + # them (MLAAttention is a sibling of Attention, so they never matched). + mla_type = _import_mla_attention_type() if quantize else None candidates = [] attention_count = 0 for name, module in model.named_modules(): - if not isinstance(module, _VLLM_ATTENTION): + is_mla = mla_type is not None and isinstance(module, mla_type) + if not is_mla and not isinstance(module, _VLLM_ATTENTION): continue attention_count += 1 sparse_kw = _sparse_kwargs(name, resolved_sparse_cfg) if quantize or sparse_kw: - candidates.append((name, module, sparse_kw)) + candidates.append((name, module, sparse_kw, is_mla)) if not candidates and not quantize: return _InstallPlan( @@ -351,8 +411,8 @@ def _plan_vllm_attention( mode = _cudagraph_mode(model_runner) if quantize else None quant_plugin: Any = _load_quant_plugin() if quantize else None plans = [] - for name, module, sparse_kw in candidates: - reasons = _layer_errors(module) + for name, module, sparse_kw, is_mla in candidates: + reasons = _layer_errors_mla(module, sparse_kw) if is_mla else _layer_errors(module) device = dtype = None if quantize: device, dtype = quant_plugin._get_device_dtype(module) @@ -368,7 +428,11 @@ def _plan_vllm_attention( if quantize: if graph_error := _sparse_graph_error(sparse_kw, mode): reasons.append(graph_error) - new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) + requires_flashinfer_patch = False + if is_mla: + new_impl, backend_error = _select_new_mla_impl(module) + else: + new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) if backend_error: reasons.append(backend_error) if reasons: @@ -383,10 +447,11 @@ def _plan_vllm_attention( device, dtype, requires_flashinfer_patch, + is_mla, ) ) if quantize and attention_count == 0: - errors.append("no regular attention layers were found") + errors.append("no attention layers were found") _raise_unsupported(errors, "NVFP4 attention" if quantize else "sparse attention") return _InstallPlan( model_runner, @@ -428,11 +493,43 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor plan.model_runner.cascade_attn_enabled = False for layer in plan.layers: layer.new_impl.sparse_kw = layer.sparse_kw - if plan.quantize: + default_formats = (plan.q_format, plan.k_format, plan.p_format, plan.v_format) == ( + "nvfp4", + "nvfp4", + "nvfp4", + "nvfp4", + ) + if plan.quantize and layer.is_mla: + _cfg_kwargs = ( + {} + if default_formats + else { + "cfg": quant_plugin.build_vllm_mla_attention_quant_cfg( + q_format=plan.q_format, + k_format=plan.k_format, + p_format=plan.p_format, + v_format=plan.v_format, + ) + } + ) + converted = quant_plugin.configure_vllm_nvfp4_mla_quantizers( + layer.module, + device=layer.device, + dtype=layer.dtype, + **_cfg_kwargs, + ) + if converted is not None and converted is not layer.module: + raise RuntimeError("vLLM attention quantization must convert modules in place") + layer.new_impl.quant_kw = _load_mla_plugin().mla_quant_kw_from_layer( + layer.module, query_in_kernel=plan.q_format != "fp8" + ) + elif plan.quantize: # Pass cfg only for non-default formats: keeps the default call # signature stable for callers/fakes that predate the cfg parameter. _cfg_kwargs = ( - { + {} + if default_formats + else { "cfg": quant_plugin.build_vllm_attention_quant_cfg( q_format=plan.q_format, k_format=plan.k_format, @@ -440,9 +537,6 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor v_format=plan.v_format, ) } - if (plan.q_format, plan.k_format, plan.p_format, plan.v_format) - != ("nvfp4", "nvfp4", "nvfp4", "nvfp4") - else {} ) converted = quant_plugin.configure_vllm_nvfp4_attention_quantizers( layer.module, @@ -469,11 +563,20 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor missing = object() old_query_flag = getattr(layer.module, "_query_quant_in_kernel", missing) old_value_flag = getattr(layer.module, "_value_quant_in_kernel", missing) + old_kv_flag = getattr(layer.module, "_skip_module_kv_quant", missing) if plan.quantize: # fp8 Q is module-level (bf16 losslessly carries E4M3 QDQ values); # the kernel then runs a plain bf16 BMM1 with no Q transform. layer.module._query_quant_in_kernel = plan.q_format != "fp8" - layer.module._value_quant_in_kernel = plan.v_format != "fp8" + if not layer.is_mla: + # MLA V is either in-kernel (prefill projected V) or the + # write-once quantized cache (decode) — no module-level flag. + layer.module._value_quant_in_kernel = plan.v_format != "fp8" + else: + # Quantized-BMM model: keep the latent cache RAW by skipping the + # module-level kv_c/k_pe QDQ. Prefill projects bf16 latent and + # decode quantizes K/V on read; both quantize operands in-kernel. + layer.module._skip_module_kv_quant = True try: # Publish the adapter last so a native impl never runs with in-kernel # quantization flags that only the ModelOpt adapter understands. @@ -483,9 +586,11 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor for name, value in ( ("_query_quant_in_kernel", old_query_flag), ("_value_quant_in_kernel", old_value_flag), + ("_skip_module_kv_quant", old_kv_flag), ): if value is missing: - delattr(layer.module, name) + if hasattr(layer.module, name): + delattr(layer.module, name) else: setattr(layer.module, name, value) raise diff --git a/tests/gpu/torch/kernels/quantization/attention/mla/test_mla_decode.py b/tests/gpu/torch/kernels/quantization/attention/mla/test_mla_decode.py new file mode 100644 index 00000000000..f70248e1087 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/attention/mla/test_mla_decode.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Golden tests for the split-K absorbed-MLA decode kernel.""" + +import pytest +import torch + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels.quantization.attention.mla import mla_attention_decode + from modelopt.torch.kernels.quantization.attention.mla.reference import ( + mla_decode_reference, + nvfp4_fake_quant, + ) + +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) + +pytestmark = pytest.mark.skipif( + not TRITON_KERNEL_AVAILABLE, reason="Triton attention kernel requires CUDA + triton" +) + +_RANK = 128 # kv_lora_rank kept small so tests run fast; 16-divisible +_ROPE = 64 +_DIM = _RANK + _ROPE + + +def _make_decode_inputs( + seq_lens: list[int], + num_heads: int = 32, + page_size: int = 16, + dtype: torch.dtype = torch.float16, + seed: int = 0, +): + """Build an absorbed query, a dense latent, and its paged copy.""" + torch.manual_seed(seed) + batch = len(seq_lens) + max_seq = max(seq_lens) + q = torch.randn(batch, num_heads, _DIM, device="cuda", dtype=dtype) + latent_dense = torch.randn(batch, max_seq, _DIM, device="cuda", dtype=dtype) + + max_blocks = (max_seq + page_size - 1) // page_size + cache = torch.zeros(batch * max_blocks, page_size, _DIM, device="cuda", dtype=dtype) + block_table = torch.zeros(batch, max_blocks, dtype=torch.int32, device="cuda") + # Shuffled page assignment to exercise the block-table walk. + perm = torch.randperm(batch * max_blocks) + next_page = 0 + for b, s in enumerate(seq_lens): + for blk in range((s + page_size - 1) // page_size): + page = int(perm[next_page]) + next_page += 1 + block_table[b, blk] = page + lo = blk * page_size + hi = min(lo + page_size, s) + cache[page, : hi - lo] = latent_dense[b, lo:hi] + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") + return q, latent_dense, cache, block_table, seq_lens_t + + +def _dense_decode(q, latent_dense, seq_lens, scale): + """Straightforward fp32 decode oracle (no schedule replication).""" + batch, num_heads, _ = q.shape + out = torch.zeros(batch, num_heads, _RANK, device=q.device, dtype=torch.float32) + lse = torch.zeros(batch, num_heads, device=q.device, dtype=torch.float32) + for b in range(int(seq_lens.shape[0])): + s = int(seq_lens[b]) + kv = latent_dense[b, :s].float() # [s, DIM] + scores = q[b].float() @ kv.T * scale # [H, s] + p = torch.softmax(scores, dim=-1) + out[b] = p @ kv[:, :_RANK] + lse[b] = torch.logsumexp(scores, dim=-1) + return out, lse + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> float: + return torch.nn.functional.cosine_similarity( + a.flatten().float(), b.flatten().float(), dim=0 + ).item() + + +class TestMLADecodeBaseline: + @pytest.mark.parametrize("num_kv_splits", [1, 32]) + @pytest.mark.parametrize("page_size", [16, 32, 64]) + def test_no_quant_matches_dense(self, num_kv_splits, page_size): + seq_lens = [200, 7, 64] + q, latent, cache, block_table, seq_t = _make_decode_inputs(seq_lens, page_size=page_size) + scale = _DIM**-0.5 + out, lse = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + page_size=page_size, + num_kv_splits=num_kv_splits, + ) + ref, ref_lse = _dense_decode(q, latent, seq_t, scale) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(lse, ref_lse, rtol=5e-3, atol=5e-3) + + def test_small_head_count(self): + """Head counts below BLOCK_H exercise the head-group masking.""" + q, latent, cache, block_table, seq_t = _make_decode_inputs([33], num_heads=2, seed=1) + scale = 0.13 + out, _ = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + ) + ref, _ = _dense_decode(q, latent, seq_t, scale) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + def test_output_dtype_follows_cache(self): + q, _, cache, block_table, seq_t = _make_decode_inputs([16], dtype=torch.bfloat16, seed=2) + out, lse = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=0.1, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + ) + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + + +class TestMLADecodeQuant: + @requires_native_e4m3 + @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) + def test_p_qdq_matches_split_local_oracle(self, mode): + seq_lens = [150, 40] + q, latent, cache, block_table, seq_t = _make_decode_inputs(seq_lens, seed=3) + scale = _DIM**-0.5 + out, _ = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + p_qdq=mode, + ) + dense_out, _ = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + ) + assert not torch.equal(out, dense_out) # P quant actually applied + ref = mla_decode_reference( + q, latent, seq_t, scale, _RANK, p_mode=mode, num_kv_splits=32, block_n=32 + ) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=2e-2) + + @requires_native_e4m3 + @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) + def test_kv_qdq_matches_oracle(self, mode): + """On-read K (feature axis) + V (token axis) quant vs the split-local oracle.""" + seq_lens = [150, 40] + q, latent, cache, block_table, seq_t = _make_decode_inputs(seq_lens, seed=6) + scale = _DIM**-0.5 + out, _ = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + k_qdq=mode, + v_qdq=mode, + ) + dense_out, _ = mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + ) + assert not torch.equal(out, dense_out) # K/V quant actually applied + ref = mla_decode_reference( + q, latent, seq_t, scale, _RANK, k_mode=mode, v_mode=mode, num_kv_splits=32, block_n=32 + ) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=2e-2) + + @requires_native_e4m3 + def test_v_format_honored_independently(self): + """Decode V is quantized independently along the token axis — toggling + v_qdq must change the output (the quantized-BMM contract), and it must + differ from the K-only quantization of the same latent.""" + seq_lens = [96] + q, _, cache, block_table, seq_t = _make_decode_inputs(seq_lens, seed=7) + scale = _DIM**-0.5 + common = { + "softmax_scale": scale, + "kv_lora_rank": _RANK, + "qk_rope_head_dim": _ROPE, + "num_kv_splits": 32, + } + k_only, _ = mla_attention_decode(q, cache, block_table, seq_t, k_qdq="nvfp4", **common) + kv_both, _ = mla_attention_decode( + q, cache, block_table, seq_t, k_qdq="nvfp4", v_qdq="nvfp4", **common + ) + # v_qdq has an independent effect (V is not merely inheriting k_qdq). + assert not torch.equal(k_only, kv_both) + + @requires_native_e4m3 + def test_full_qkpv_nvfp4_cosine(self): + """fp32-carrier Q + on-read K/V + fused P, all NVFP4, from a RAW cache.""" + seq_lens = [96] + q, latent, cache, block_table, seq_t = _make_decode_inputs( + seq_lens, dtype=torch.bfloat16, seed=8 + ) + scale = _DIM**-0.5 + q_carrier = nvfp4_fake_quant(q.float(), block_axis=-1) # caller-side Q QDQ (fp32) + out, _ = mla_attention_decode( + q_carrier, + cache, + block_table, + seq_t, + softmax_scale=scale, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + k_qdq="nvfp4", + v_qdq="nvfp4", + p_qdq="nvfp4", + ) + ref, _ = _dense_decode(q, latent, seq_t, scale) + assert torch.isfinite(out.float()).all() + assert _cos(out, ref) > 0.98 + + +class TestMLADecodeErrors: + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_kv_splits": 0}, "num_kv_splits"), + ({"num_kv_splits": 64}, "num_kv_splits"), + ({"page_size": 8}, "page_size"), + ({"p_qdq": "int8"}, "p_qdq must be one of"), + ], + ) + def test_invalid_config(self, kwargs, match): + q, _, cache, block_table, seq_t = _make_decode_inputs([16]) + with pytest.raises(ValueError, match=match): + mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=0.1, + kv_lora_rank=_RANK, + qk_rope_head_dim=_ROPE, + **kwargs, + ) + + def test_nvfp4_requires_divisible_dims(self): + q = torch.randn(1, 4, 88 + 8, device="cuda", dtype=torch.float16) + cache = torch.randn(4, 16, 96, device="cuda", dtype=torch.float16) + block_table = torch.zeros(1, 4, dtype=torch.int32, device="cuda") + seq_t = torch.tensor([16], dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="divisible by 16"): + mla_attention_decode( + q, + cache, + block_table, + seq_t, + softmax_scale=0.1, + kv_lora_rank=88, + qk_rope_head_dim=8, + p_qdq="nvfp4", + ) diff --git a/tests/gpu/torch/kernels/quantization/attention/mla/test_mla_prefill.py b/tests/gpu/torch/kernels/quantization/attention/mla/test_mla_prefill.py new file mode 100644 index 00000000000..292aaffef4a --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/attention/mla/test_mla_prefill.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Golden tests for the varlen MLA prefill kernel against eager torch oracles.""" + +import pytest +import torch + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels.quantization.attention.mla import mla_prefill_attention + from modelopt.torch.kernels.quantization.attention.mla.reference import mla_attention_reference + +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) + +pytestmark = pytest.mark.skipif( + not TRITON_KERNEL_AVAILABLE, reason="Triton attention kernel requires CUDA + triton" +) + + +def _cu_seqlens(seq_lens: list[int], device: str = "cuda") -> torch.Tensor: + cu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device=device) + cu[1:] = torch.cumsum(torch.tensor(seq_lens, device=device), dim=0) + return cu + + +def _make_mla_qkv( + q_lens: list[int], + k_lens: list[int], + num_heads: int, + num_kv_heads: int, + lqk: int, + lv: int, + dtype: torch.dtype = torch.float16, + seed: int = 0, +): + torch.manual_seed(seed) + total_q, total_k = sum(q_lens), sum(k_lens) + q = torch.randn(total_q, num_heads, lqk, device="cuda", dtype=dtype) + k = torch.randn(total_k, num_kv_heads, lqk, device="cuda", dtype=dtype) + v = torch.randn(total_k, num_kv_heads, lv, device="cuda", dtype=dtype) + return q, k, v, _cu_seqlens(q_lens), _cu_seqlens(k_lens) + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> float: + return torch.nn.functional.cosine_similarity( + a.flatten().float(), b.flatten().float(), dim=0 + ).item() + + +class TestMLAPrefillBaseline: + def test_no_quant_matches_reference(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([37, 64], [37, 64], 4, 4, 192, 128) + out = mla_prefill_attention(q, k, v, cu_q, cu_k, 64) + ref = mla_attention_reference(q, k, v, cu_q, cu_k) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + def test_gqa_no_quant(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([48, 33], [48, 33], 8, 2, 192, 128, seed=1) + out = mla_prefill_attention(q, k, v, cu_q, cu_k, 48) + ref = mla_attention_reference(q, k, v, cu_q, cu_k) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + @pytest.mark.parametrize( + ("num_heads", "num_kv_heads", "lqk", "lv"), + [(8, 8, 192, 128), (8, 2, 64, 32), (4, 4, 576, 512)], + ) + def test_dim_sweep_no_quant(self, num_heads, num_kv_heads, lqk, lv): + q, k, v, cu_q, cu_k = _make_mla_qkv([50], [50], num_heads, num_kv_heads, lqk, lv, seed=2) + out = mla_prefill_attention(q, k, v, cu_q, cu_k, 50) + ref = mla_attention_reference(q, k, v, cu_q, cu_k) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + def test_causal_suffix_qk(self): + """Causal with k_len > q_len: Q is the suffix of the KV span.""" + q, k, v, cu_q, cu_k = _make_mla_qkv([16, 8], [80, 40], 4, 4, 192, 128, seed=3) + out = mla_prefill_attention(q, k, v, cu_q, cu_k, 16) + ref = mla_attention_reference(q, k, v, cu_q, cu_k) + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + def test_lse_matches_logsumexp(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([40], [40], 4, 4, 192, 128, seed=4) + _, lse = mla_prefill_attention(q, k, v, cu_q, cu_k, 40, causal=False, return_lse=True) + scale = 192**-0.5 + scores = torch.einsum("qhd,khd->hqk", q.float(), k.float()) * scale + expected = torch.logsumexp(scores, dim=-1) # [H, total_q] + torch.testing.assert_close(lse, expected, rtol=5e-3, atol=5e-3) + + def test_two_chunk_lse_merge_matches_single_shot(self): + """Non-causal chunked K/V merged via LSE equals the single-shot result.""" + q, k, v, cu_q, cu_k = _make_mla_qkv([32], [128], 4, 4, 192, 128, seed=5) + full = mla_prefill_attention(q, k, v, cu_q, cu_k, 32, causal=False) + + outs, lses = [], [] + for lo, hi in ((0, 64), (64, 128)): + o, lse = mla_prefill_attention( + q, + k[lo:hi], + v[lo:hi], + cu_q, + _cu_seqlens([hi - lo]), + 32, + causal=False, + return_lse=True, + ) + outs.append(o.float()) + lses.append(lse) + lse_all = torch.logaddexp(lses[0], lses[1]) + w0 = torch.exp(lses[0] - lse_all).transpose(0, 1).unsqueeze(-1) # [total_q, H, 1] + w1 = torch.exp(lses[1] - lse_all).transpose(0, 1).unsqueeze(-1) + merged = outs[0] * w0 + outs[1] * w1 + torch.testing.assert_close(merged, full.float(), rtol=5e-3, atol=5e-3) + + def test_empty_kv_rows_are_zero_with_neg_inf_lse(self): + q, k, v, cu_q, _ = _make_mla_qkv([8], [8], 4, 4, 192, 128, seed=6) + cu_k_empty = _cu_seqlens([0]) + out, lse = mla_prefill_attention( + q, k[:0], v[:0], cu_q, cu_k_empty, 8, causal=False, return_lse=True + ) + assert torch.all(out == 0) + assert torch.all(torch.isneginf(lse)) + + +class TestMLAPrefillSparse: + def test_sparse24_no_quant(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([96], [96], 4, 4, 192, 128, seed=7) + dense = mla_prefill_attention(q, k, v, cu_q, cu_k, 96) + out = mla_prefill_attention( + q, k, v, cu_q, cu_k, 96, sparsity_n=2, sparsity_m=4, dense_recent_tokens=16 + ) + ref = mla_attention_reference( + q, k, v, cu_q, cu_k, sparsity_n=2, sparsity_m=4, dense_recent_tokens=16 + ) + assert not torch.equal(out, dense) # sparsity actually applied + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + @requires_native_e4m3 + def test_sparse24_nvfp4_composition(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([96], [96], 4, 4, 192, 128, seed=8) + out = mla_prefill_attention( + q, + k, + v, + cu_q, + cu_k, + 96, + q_quant="nvfp4", + k_quant="nvfp4", + p_quant="nvfp4", + v_quant="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=16, + ) + ref = mla_attention_reference( + q, + k, + v, + cu_q, + cu_k, + q_mode="nvfp4", + k_mode="nvfp4", + p_mode="nvfp4", + v_mode="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=16, + ) + assert _cos(out, ref) > 0.98 + + +class TestMLAPrefillQuant: + @requires_native_e4m3 + @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) + def test_qkpv_quant_matches_reference(self, mode): + q, k, v, cu_q, cu_k = _make_mla_qkv([64, 41], [64, 41], 4, 4, 192, 128, seed=9) + kwargs = {f"{op}_quant": mode for op in ("q", "k", "p", "v")} + out = mla_prefill_attention(q, k, v, cu_q, cu_k, 64, **kwargs) + dense = mla_prefill_attention(q, k, v, cu_q, cu_k, 64) + assert not torch.equal(out, dense) # quant actually applied + ref = mla_attention_reference( + q, k, v, cu_q, cu_k, **{f"{op}_mode": mode for op in ("q", "k", "p", "v")} + ) + if mode == "fp8": + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=2e-2) + else: + assert _cos(out, ref) > 0.98 + + @requires_native_e4m3 + def test_p_only_nvfp4(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([80], [80], 4, 4, 192, 128, seed=10) + out = mla_prefill_attention(q, k, v, cu_q, cu_k, 80, p_quant="nvfp4") + ref = mla_attention_reference(q, k, v, cu_q, cu_k, p_mode="nvfp4") + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=2e-2) + + @requires_native_e4m3 + def test_amax_scales_change_results(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([64], [64], 4, 4, 192, 128, seed=11) + out_default = mla_prefill_attention(q, k, v, cu_q, cu_k, 64, q_quant="fp8") + out_amax = mla_prefill_attention(q, k, v, cu_q, cu_k, 64, q_quant="fp8", q_amax=4.0) + assert not torch.equal(out_default, out_amax) + ref = mla_attention_reference(q, k, v, cu_q, cu_k, q_mode="fp8", q_amax=4.0) + torch.testing.assert_close(out_amax.float(), ref, rtol=5e-3, atol=2e-2) + + +class TestMLAPrefillErrors: + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"q_quant": "int8"}, "q_quant must be one of"), + ({"block_n": 24}, "multiples of 16"), + ({"p_quant": "fp8", "p_amax": -1.0}, "finite positive"), + ], + ) + def test_invalid_config(self, kwargs, match): + q, k, v, cu_q, cu_k = _make_mla_qkv([16], [16], 4, 4, 192, 128) + with pytest.raises(ValueError, match=match): + mla_prefill_attention(q, k, v, cu_q, cu_k, 16, **kwargs) + + def test_nvfp4_requires_divisible_qk_dim(self): + q, k, v, cu_q, cu_k = _make_mla_qkv([16], [16], 4, 4, 40, 32) + with pytest.raises(ValueError, match="qk_head_dim % 16"): + mla_prefill_attention(q, k, v, cu_q, cu_k, 16, q_quant="nvfp4") diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 038bbd8e978..62ab3f9b523 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -498,6 +498,71 @@ def test_tiny_deepseek_mla_quantize(tiny_deepseek_llm): assert vllm_key.rsplit("._amax", 1)[0] in summary["quantizer_names"], vllm_key +@pytest.fixture(scope="module") +def tiny_deepseek_triton_mla_llm(tmp_path_factory): + """Tiny DeepSeek pinned to TRITON_MLA for the fused NVFP4 attention path.""" + tmp = tmp_path_factory.mktemp("tiny_deepseek_triton_mla") + model_dir = create_tiny_deepseek_v3_dir( + tmp, qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128 + ) + llm = _boot_llm( + model_dir, + moe_backend="triton", + enable_expert_parallel=True, + attention_backend="TRITON_MLA", + enable_prefix_caching=False, # rejected by the attention installer + ) + try: + yield llm + finally: + _shutdown_llm(llm) + + +def _install_mla_nvfp4_attention(self): + """Worker-side: install fused NVFP4 MLA attention and summarize the result.""" + from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_mla + from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import ( + install_vllm_nvfp4_attention, + ) + + report = install_vllm_nvfp4_attention(self.model_runner, sparse_cfg=None) + model = self.model_runner.model + mla_modules = [m for m in model.modules() if isinstance(m, VllmMLAAttention)] + return { + "installed_layers": list(report.installed_layers), + "backend_counts": dict(report.backend_counts), + "mla_count": len(mla_modules), + "all_modelopt_impl": all(isinstance(m.impl, vllm_mla.ModelOptMLAImpl) for m in mla_modules), + "query_in_kernel": all(getattr(m, "_query_quant_in_kernel", False) for m in mla_modules), + } + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9), + reason="NVFP4 attention install requires compute capability >= 8.9", +) +def test_tiny_deepseek_mla_nvfp4_attention_install_and_generate(tiny_deepseek_triton_mla_llm): + """End-to-end: install fused NVFP4 Q/K/P/V MLA attention, then prefill + decode.""" + from vllm import SamplingParams + from vllm.inputs import TokensPrompt + + llm = tiny_deepseek_triton_mla_llm + summaries = llm.collective_rpc(_install_mla_nvfp4_attention) + summary = summaries[0] + assert summary["mla_count"] >= 2, summary + assert summary["all_modelopt_impl"], summary + assert summary["query_in_kernel"], summary + assert summary["backend_counts"] == {"ModelOptMLAImpl": summary["mla_count"]}, summary + + # Prefill (17 tokens) + decode (8 steps) through the ModelOpt MLA kernels. + outputs = llm.generate( + [TokensPrompt(prompt_token_ids=list(range(1, 18)))], + SamplingParams(max_tokens=8, temperature=0.0), + ) + token_ids = outputs[0].outputs[0].token_ids + assert len(token_ids) == 8, outputs + + def test_configure_vllm_attention_quantizers_fp8_bmm2(monkeypatch): monkeypatch.setattr( vllm_plugin, diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_mla_runtime.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_mla_runtime.py new file mode 100644 index 00000000000..f2440b37d39 --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_mla_runtime.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the vLLM MLA attention runtime installer and impl adapter.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +# The MLA adapter targets the vLLM 0.26-style prefill-backend architecture. +pytest.importorskip("vllm.v1.attention.backends.mla.prefill.base") + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.mla.triton_mla import TritonMLAImpl + +from modelopt.torch.quantization.plugins import vllm as quant_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_mla, vllm_runtime + +_MLA_ATTENTION = vllm_runtime._import_mla_attention_type() + +pytestmark = pytest.mark.skipif( + _MLA_ATTENTION is None, reason="vLLM MLAAttention type is unavailable" +) + + +def _bare_mla_attention(impl_cls=TritonMLAImpl, num_heads=8): + module = object.__new__(_MLA_ATTENTION) + nn.Module.__init__(module) + module.kv_cache_dtype = "auto" + module.use_sparse = False + module.indexer = None + module.q_pad_num_heads = None + module.num_heads = num_heads + module.kv_lora_rank = 128 + module.qk_rope_head_dim = 64 + module.qk_nope_head_dim = 128 + module.qk_head_dim = 192 + module.v_head_dim = 128 + module.device = torch.device("cpu") + module.dtype = torch.float16 + impl = object.__new__(impl_cls) + impl.scale = 192**-0.5 + impl.num_heads = num_heads + impl.kv_lora_rank = module.kv_lora_rank + impl.qk_rope_head_dim = module.qk_rope_head_dim + module.impl = impl + return module + + +def _model_runner(model): + model_config = SimpleNamespace( + hf_config=SimpleNamespace(sparse_attention_config=None), dtype=torch.float16 + ) + return SimpleNamespace( + model=model, + model_config=model_config, + cascade_attn_enabled=True, + vllm_config=SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_dbo=False, + use_ubatching=False, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False, cache_dtype="auto"), + compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), + kv_transfer_config=None, + speculative_config=None, + ), + ) + + +@pytest.fixture +def patched_parallel_state(monkeypatch): + monkeypatch.setattr( + quant_plugin, + "create_parallel_state", + lambda: quant_plugin.ParallelState(data_parallel_group=None), + ) + + +class TestMLAInstall: + def test_nvfp4_install_configures_quantizers_and_impl(self, patched_parallel_state): + attention = _bare_mla_attention() + runner = _model_runner(nn.ModuleDict({"mla_attn": attention})) + + report = vllm_runtime.install_vllm_nvfp4_attention(runner, sparse_cfg=None) + + assert isinstance(attention, quant_plugin._QuantVLLMMLAAttention) + assert type(attention.impl) is vllm_mla.ModelOptMLAImpl + for name in ("q", "kv_c", "k_pe", "k_mha", "p", "v_mha"): + quantizer = getattr(attention, f"{name}_bmm_quantizer") + assert quantizer.is_enabled + assert quantizer.is_nvfp4_dynamic + for name in ("kv_c", "k_pe", "k_mha", "v_mha"): + amax = getattr(attention, f"{name}_bmm_quantizer")._amax + assert float(amax) == 6.0 * 448.0 + assert attention._query_quant_in_kernel is True + assert not hasattr(attention, "_value_quant_in_kernel") + # Quantized-BMM model: the module flag keeps the latent cache raw by + # skipping the module-side kv_c/k_pe quant; operands are quantized + # in-kernel (prefill projected K/V; decode on-read latent K/V). + assert attention._skip_module_kv_quant is True + assert attention.impl.quant_kw.prefill == { + "q_quant": "nvfp4", + "q_amax": None, + "k_quant": "nvfp4", + "k_amax": 6.0 * 448.0, + "p_quant": "nvfp4", + "p_amax": 1.0, + "v_quant": "nvfp4", + "v_amax": 6.0 * 448.0, + } + assert attention.impl.quant_kw.decode == { + "p_qdq": "nvfp4", + "p_qdq_amax": 1.0, + "k_qdq": "nvfp4", + "k_qdq_amax": 6.0 * 448.0, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + } + assert attention.impl.sparse_kw == {} + assert report.installed_layers == ("mla_attn",) + assert report.quantized_layers == ("mla_attn",) + assert report.backend_counts == {"ModelOptMLAImpl": 1} + assert runner.cascade_attn_enabled is False + + def test_fp8_q_and_v_formats(self, patched_parallel_state): + attention = _bare_mla_attention() + runner = _model_runner(nn.ModuleDict({"mla_attn": attention})) + + vllm_runtime.install_vllm_nvfp4_attention( + runner, sparse_cfg=None, q_format="fp8", v_format="fp8" + ) + + assert attention._query_quant_in_kernel is False + # Module-level FP8 Q means neither kernel applies a Q transform. + assert attention.impl.quant_kw.prefill["q_quant"] is None + assert attention.impl.quant_kw.prefill["v_quant"] == "fp8" + assert attention.impl.quant_kw.prefill["v_amax"] == 448.0 + assert float(attention.q_bmm_quantizer._amax) == 448.0 + + def test_unsupported_mla_backend_rejected_before_mutation(self, patched_parallel_state): + attention = _bare_mla_attention(impl_cls=FlashAttentionImpl) + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"mla_attn": attention})) + + with pytest.raises(NotImplementedError, match="TRITON_MLA"): + vllm_runtime.install_vllm_nvfp4_attention(runner, sparse_cfg=None) + + assert attention.impl is original_impl + assert not hasattr(attention, "q_bmm_quantizer") + assert runner.cascade_attn_enabled is True + + @pytest.mark.parametrize( + ("mutation", "match"), + [ + ({"kv_cache_dtype": "fp8"}, "FP8 KV cache"), + ({"indexer": object()}, "sparse-indexer"), + ({"use_sparse": True}, "sparse-indexer"), + ({"q_pad_num_heads": 256}, "q_pad_num_heads"), + ({"v_head_dim": 100}, "v_head_dim"), + ], + ) + def test_layer_gates(self, patched_parallel_state, mutation, match): + attention = _bare_mla_attention() + for name, value in mutation.items(): + setattr(attention, name, value) + runner = _model_runner(nn.ModuleDict({"mla_attn": attention})) + + with pytest.raises(NotImplementedError, match=match): + vllm_runtime.install_vllm_nvfp4_attention(runner, sparse_cfg=None) + + def test_skip_softmax_sparse_config_rejected(self, patched_parallel_state, monkeypatch): + attention = _bare_mla_attention() + runner = _model_runner(nn.ModuleDict({"mla_attn": attention})) + monkeypatch.setattr( + vllm_runtime, + "_sparse_kwargs", + lambda name, cfg: {"skip_softmax_threshold": 0.5}, + ) + + with pytest.raises(NotImplementedError, match="skip-softmax is unsupported on MLA"): + vllm_runtime.install_vllm_nvfp4_attention(runner, sparse_cfg={"*": {"enable": True}}) + + def test_sparse_only_install_ignores_mla(self): + attention = _bare_mla_attention() + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"mla_attn": attention})) + runner.model_config.hf_config.sparse_attention_config = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + + assert attention.impl is original_impl + assert not hasattr(attention, "q_bmm_quantizer") + assert report.installed_layers == () + + +def _bare_mla_impl(*, p_qdq="nvfp4", prefill_active=True, sparse_kw=None): + impl = object.__new__(vllm_mla.ModelOptMLAImpl) + impl.scale = 0.25 + impl.kv_lora_rank = 128 + impl.qk_rope_head_dim = 64 + mode = "nvfp4" if prefill_active else None + impl.quant_kw = vllm_mla._MLAQuantKw( + prefill={ + "q_quant": mode, + "q_amax": None, + "k_quant": mode, + "k_amax": None, + "p_quant": mode, + "p_amax": 1.0, + "v_quant": mode, + "v_amax": None, + }, + decode={ + "p_qdq": p_qdq, + "p_qdq_amax": 1.0, + "k_qdq": mode, + "k_qdq_amax": None, + "v_qdq": mode, + "v_qdq_amax": None, + }, + ) + impl.sparse_kw = dict(sparse_kw or {}) + return impl + + +class TestMLAImplDispatch: + def test_forward_mqa_dispatches_to_decode_kernel(self, monkeypatch): + impl = _bare_mla_impl() + recorded = {} + + def _record_decode(q, latent_cache, block_table, b_seq_len, **kwargs): + recorded["q"] = q + recorded["latent_cache"] = latent_cache + recorded["block_table"] = block_table + recorded["b_seq_len"] = b_seq_len + recorded.update(kwargs) + return torch.zeros(q.shape[0], q.shape[1], 128), torch.zeros(q.shape[0], q.shape[1]) + + monkeypatch.setattr(vllm_mla, "mla_attention_decode", _record_decode) + + def _poison(self, *args, **kwargs): + raise AssertionError("native TritonMLAImpl.forward_mqa must not run") + + monkeypatch.setattr(TritonMLAImpl, "forward_mqa", _poison) + + layer = SimpleNamespace(_query_quant_in_kernel=False) + cache = torch.zeros(4, 16, 192, dtype=torch.float16) + ql_nope = torch.zeros(2, 8, 128, dtype=torch.float16) + q_pe = torch.zeros(2, 8, 64, dtype=torch.float16) + metadata = SimpleNamespace( + decode=SimpleNamespace( + block_table=torch.zeros(2, 4, dtype=torch.int32), + seq_lens=torch.tensor([5, 9], dtype=torch.int32), + ) + ) + + o, lse = impl.forward_mqa((ql_nope, q_pe), cache, metadata, layer) + + assert recorded["q"].shape == (2, 8, 192) # tuple q concatenated + assert recorded["latent_cache"] is cache + assert recorded["block_table"] is metadata.decode.block_table + assert recorded["b_seq_len"] is metadata.decode.seq_lens + assert recorded["softmax_scale"] == impl.scale + assert recorded["num_kv_splits"] == 32 + assert recorded["p_qdq"] == "nvfp4" + assert recorded["p_qdq_amax"] == 1.0 + # Decode quantizes K (feature) and V (token) on read from the raw cache. + assert recorded["k_qdq"] == "nvfp4" + assert recorded["v_qdq"] == "nvfp4" + assert recorded["kv_lora_rank"] == 128 + assert recorded["qk_rope_head_dim"] == 64 + assert recorded["out_dtype"] == cache.dtype + assert o.shape == (2, 8, 128) + assert lse.shape == (2, 8) + + def test_forward_mqa_no_transform_uses_native_path(self, monkeypatch): + impl = _bare_mla_impl(p_qdq=None, prefill_active=False) + called = {} + + def _native(self, q, cache, metadata, layer): + called["native"] = True + return "native-result" + + monkeypatch.setattr(TritonMLAImpl, "forward_mqa", _native) + layer = SimpleNamespace(_query_quant_in_kernel=False) + + result = impl.forward_mqa(torch.zeros(1, 8, 192), None, None, layer) + + assert called.get("native") is True + assert result == "native-result" + + @pytest.mark.parametrize("super_raises", [False, True]) + def test_forward_mha_swaps_and_restores_prefill_backend(self, monkeypatch, super_raises): + impl = _bare_mla_impl() + base_backend = SimpleNamespace( + num_heads=8, + scale=impl.scale, + kv_lora_rank=128, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + vllm_config=None, + ) + prefill_metadata = SimpleNamespace(prefill_backend=base_backend) + metadata = SimpleNamespace(prefill=prefill_metadata) + seen = {} + + def _fake_super(self, q, kv_c, k_pe, cache, attn_metadata, k_scale, output, output_scale): + seen["backend_during_call"] = attn_metadata.prefill.prefill_backend + if super_raises: + raise RuntimeError("boom") + + monkeypatch.setattr(TritonMLAImpl, "forward_mha", _fake_super) + + args = (None, None, None, None, metadata, None, None, None) + if super_raises: + with pytest.raises(RuntimeError, match="boom"): + impl.forward_mha(*args) + else: + impl.forward_mha(*args) + + assert isinstance(seen["backend_during_call"], vllm_mla._ModelOptMLAPrefillBackend) + assert seen["backend_during_call"]._prefill_metadata is prefill_metadata + assert prefill_metadata.prefill_backend is base_backend # restored + + def test_forward_mha_no_transform_uses_native_backend(self, monkeypatch): + impl = _bare_mla_impl(p_qdq=None, prefill_active=False) + base_backend = SimpleNamespace(prefill_backend=None) + prefill_metadata = SimpleNamespace(prefill_backend=base_backend) + metadata = SimpleNamespace(prefill=prefill_metadata) + seen = {} + + def _fake_super(self, q, kv_c, k_pe, cache, attn_metadata, k_scale, output, output_scale): + seen["backend_during_call"] = attn_metadata.prefill.prefill_backend + + monkeypatch.setattr(TritonMLAImpl, "forward_mha", _fake_super) + + impl.forward_mha(None, None, None, None, metadata, None, None, None) + + assert seen["backend_during_call"] is base_backend # untouched + + +def test_impl_keeps_cache_raw_no_update_override(): + """Quantized-BMM model: the latent cache stays RAW, so the impl must NOT + override do_kv_cache_update (which would re-introduce a cache-write quant). + Decode instead quantizes K/V on read; prefill quantizes projected operands. + """ + assert vllm_mla.ModelOptMLAImpl.do_kv_cache_update is TritonMLAImpl.do_kv_cache_update