diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 5b6648908dd..9bb5a38272f 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -715,6 +715,7 @@ class CompilationConfig: "vllm::kda_attention", "vllm::sparse_attn_indexer", "vllm::rocm_aiter_sparse_attn_indexer", + "vllm::monolithic_attn", ] def compute_hash(self) -> str: diff --git a/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py b/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py index 559e6f0a6db..fb8d49fbdac 100644 --- a/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py +++ b/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Monolithic decoder layer for DeepSeek V3.2 on SM100 (Blackwell). -Direct kernel calls, no module wrappers for norms. -Gate weight inlined, FusedMoE kept for quantized expert kernels. """ from __future__ import annotations @@ -11,24 +9,161 @@ from __future__ import annotations import torch from torch import nn -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size -from .allreduce_rms import AllReduceRMSParams, allreduce_add_rms_norm from .attention import MonolithicMLAAttention -from .ops import fused_norm_rope, fused_q, rms_norm +from .ops import fused_norm_rope, fused_q from .sparse_indexer import sparse_attn_indexer -_side_stream: torch.cuda.Stream | None = None + +def monolithic_attn( + positions: torch.Tensor, + q_c: torch.Tensor, + kv_c: torch.Tensor, + k_pe: torch.Tensor, + index_k: torch.Tensor, + index_weights: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + layer = get_forward_context().no_compile_layers[layer_name] + attn = layer.attn + mla = attn.mla_attn + + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + output.zero_() + return output + + mla_attn_metadata = attn_metadata.get(mla.layer_name) + if mla_attn_metadata is None: + output.zero_() + return output + + num_actual_toks = mla_attn_metadata.num_actual_tokens + if num_actual_toks == 0: + output.zero_() + return output + + # Step 2. fused norm + rope + cache writes + slot_mapping = None + indexer_k_cache = None + mla_kv_cache = None + mla_k_scale = None + idx_meta = attn_metadata.get(attn.indexer_k_cache.prefix) + if idx_meta is not None: + slot_mapping = idx_meta.slot_mapping + indexer_k_cache = attn.indexer_k_cache.kv_cache + mla_kv_cache = attn.mla_attn.kv_cache + mla_k_scale = attn.mla_attn._k_scale + + q_c = fused_norm_rope( + positions, + q_c, + attn.q_a_layernorm_weight, + layer.rms_norm_eps, + kv_c, + attn.kv_a_layernorm_weight, + attn.rms_norm_eps, + k_pe, + attn.rotary_emb.cos_sin_cache, + index_k, + attn.indexer_k_norm.weight, + attn.indexer_k_norm.bias, + attn.rms_norm_eps, + attn.indexer_rope_emb.cos_sin_cache, + attn.topk_indices_buffer, + slot_mapping=slot_mapping, + indexer_k_cache=indexer_k_cache, + mla_kv_cache=mla_kv_cache, + mla_kv_cache_dtype=attn.mla_attn.kv_cache_dtype, + mla_k_scale=mla_k_scale, + ) + + # Step 3. q_c -> index_q, q + step3_out = torch.mm(q_c, layer._fused_step3_q_w.T) + index_q, q = step3_out.split( + [layer._step3_index_q_dim, step3_out.shape[-1] - layer._step3_index_q_dim], + dim=-1, + ) + index_q = index_q.view(-1, attn.index_n_heads, attn.index_head_dim) + q = q.view(-1, attn.num_local_heads, attn.qk_head_dim) + + # Step 4. Q RoPE + W_UK_T absorption + FP8 packing + q_nope, q_pe = q.split( + [mla.qk_nope_head_dim, mla.qk_rope_head_dim], + dim=-1, + ) + q_nope = q_nope.transpose(0, 1) + ql_nope = torch.bmm(q_nope, mla.W_UK_T) + ql_nope = ql_nope.transpose(0, 1) + + index_q_fp8, index_weights, mqa_q = fused_q( + positions, + q_pe, + attn.rotary_emb.cos_sin_cache, + index_q, + attn.indexer_rope_emb.cos_sin_cache, + ql_nope, + mla._q_scale, + index_weights, + attn.indexer_softmax_scale, + attn.index_n_heads**-0.5, + ) + + # Steps 5-6. Sparse indexer + MLA sparse decode attention + sparse_attn_indexer( + attn.indexer_k_cache.prefix, + attn.indexer_k_cache.kv_cache, + index_q_fp8, + index_weights, + attn.topk_tokens, + attn.index_head_dim, + layer.max_model_len, + layer.indexer_workspace_size, + attn.topk_indices_buffer, + ) + + mqa_q = mqa_q[:num_actual_toks] + kv_cache = mla.kv_cache + if mla.kv_cache_dtype.startswith("fp8") and mla.kv_cache_dtype != "fp8_ds_mla": + kv_cache = kv_cache.view(torch.float8_e4m3fn) + attn_out, _ = mla.impl.forward_mqa(mqa_q, kv_cache, mla_attn_metadata, mla) + x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1) + + out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim) + out = out.transpose(0, 1) + torch.bmm(x, mla.W_UV, out=out) + return output -def _get_side_stream() -> torch.cuda.Stream: - """Lazily created CUDA stream shared by all decoder layers.""" - global _side_stream - if _side_stream is None: - _side_stream = torch.cuda.Stream() - return _side_stream +def monolithic_attn_fake( + positions: torch.Tensor, + q_c: torch.Tensor, + kv_c: torch.Tensor, + k_pe: torch.Tensor, + index_k: torch.Tensor, + index_weights: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + del positions, q_c, kv_c, k_pe, index_k, index_weights, layer_name + return output + + +direct_register_custom_op( + op_name="monolithic_attn", + op_func=monolithic_attn, + fake_impl=monolithic_attn_fake, + mutates_args=["output"], + dispatch_key=current_platform.dispatch_key, +) class MonolithicDecoderLayer(nn.Module): @@ -45,9 +180,14 @@ class MonolithicDecoderLayer(nn.Module): layer_idx: int, topk_indices_buffer: torch.Tensor, prefix: str = "", - fi_params: AllReduceRMSParams | None = None, ) -> None: super().__init__() + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + self.layer_name = prefix self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.rms_norm_eps = config.rms_norm_eps @@ -55,7 +195,6 @@ class MonolithicDecoderLayer(nn.Module): self.kv_lora_rank = config.kv_lora_rank self.qk_rope_head_dim = config.qk_rope_head_dim self.tp_size = get_tensor_model_parallel_world_size() - self._fi_params = fi_params cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config @@ -63,13 +202,18 @@ class MonolithicDecoderLayer(nn.Module): self.indexer_workspace_size = get_max_prefill_buffer_size(vllm_config) self.max_model_len = vllm_config.model_config.max_model_len - # LayerNorm weights (raw) + # Use the regular vLLM RMSNorm modules so the compiler sees the + # canonical residual-add + RMSNorm pattern. dtype = torch.get_default_dtype() - self.input_layernorm_weight = nn.Parameter( - torch.ones(config.hidden_size, dtype=dtype) + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=dtype, ) - self.post_attention_layernorm_weight = nn.Parameter( - torch.ones(config.hidden_size, dtype=dtype) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=dtype, ) # Fused QKV A-projection lives inside self_attn namespace @@ -103,7 +247,6 @@ class MonolithicDecoderLayer(nn.Module): topk_indices_buffer=topk_indices_buffer, prefix=f"{prefix}.self_attn", ) - self.attn.o_proj.reduce_results = False # MoE or Dense MLP moe_layer_freq = getattr(config, "moe_layer_freq", 1) @@ -126,15 +269,12 @@ class MonolithicDecoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.mlp", ) - self.mlp.skip_final_allreduce = True - self.mlp.skip_scale_and_add = True else: self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, - reduce_results=False, prefix=f"{prefix}.mlp", ) @@ -183,56 +323,11 @@ class MonolithicDecoderLayer(nn.Module): hidden_states: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: - from vllm.forward_context import get_forward_context - - fwd_ctx = get_forward_context() - attn_metadata = fwd_ctx.attn_metadata - mla = self.attn.mla_attn - mla_attn_metadata = None - slot_mapping = None - if isinstance(attn_metadata, dict): - idx_meta = attn_metadata[self.attn.indexer_k_cache.prefix] - mla_attn_metadata = attn_metadata[mla.layer_name] - # Indexer and MLA caches share the same block_size and track - # the same requests, so their slot_mappings are identical. - slot_mapping = idx_meta.slot_mapping - - if slot_mapping is not None: - indexer_k_cache = self.attn.indexer_k_cache.kv_cache - mla_kv_cache = mla.kv_cache - mla_k_scale = mla._k_scale - else: - indexer_k_cache = None - mla_kv_cache = None - mla_k_scale = None - - # Input norm + residual - # When fused_allreduce_rms is enabled, hidden_states arriving from - # the previous layer is the *unreduced* MLP/MoE output. We fuse - # AllReduce + residual-add + RMSNorm into a single kernel. if residual is None: - # First layer: hidden_states is from embed_tokens (already - # fully materialised), no allreduce needed. residual = hidden_states - hidden_states = rms_norm( - hidden_states, self.input_layernorm_weight, self.rms_norm_eps - ) + hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = allreduce_add_rms_norm( - hidden_states, - residual, - self.input_layernorm_weight, - self.rms_norm_eps, - self._fi_params, - ) - - if not hasattr(self, "_fused_step1_hidden_w") or not hasattr( - self, "_fused_step3_q_w" - ): - raise RuntimeError( - "Monolithic decoder fused weights are not initialized. " - "Call fuse_indexer_weights() after weight loading." - ) + hidden_states, residual = self.input_layernorm(hidden_states, residual) # Step 1. hidden_states -> q_c, kv_c, k_pe, index_k, index_weights step1_out = torch.mm(hidden_states, self._fused_step1_hidden_w.T) @@ -241,202 +336,27 @@ class MonolithicDecoderLayer(nn.Module): dim=-1, ) - # Step 2. Q RMS norm - # + KV RMS norm + KV RoPE + MLA cache write - # + Index K layer norm + RoPE + FP8 quant + cache write - # + Init topk indices - q_c = fused_norm_rope( - positions, - # Q RMS norm - q_c, - self.attn.q_a_layernorm_weight, - self.rms_norm_eps, - # KV RMS norm - kv_c, - self.attn.kv_a_layernorm_weight, - self.attn.rms_norm_eps, - # KV RoPE - k_pe, - self.attn.rotary_emb.cos_sin_cache, - # Index K layer norm + RoPE - index_k, - self.attn.indexer_k_norm.weight, - self.attn.indexer_k_norm.bias, - self.attn.rms_norm_eps, - self.attn.indexer_rope_emb.cos_sin_cache, - # Top k indices - self.attn.topk_indices_buffer, - # Fused cache writes (single slot_mapping for both caches) - slot_mapping=slot_mapping, - indexer_k_cache=indexer_k_cache, - mla_kv_cache=mla_kv_cache, - mla_kv_cache_dtype=self.attn.mla_attn.kv_cache_dtype, - mla_k_scale=mla_k_scale, - ) - - # Step 3. q_c -> index_q, q - step3_out = torch.mm(q_c, self._fused_step3_q_w.T) - index_q, q = step3_out.split( - [self._step3_index_q_dim, step3_out.shape[-1] - self._step3_index_q_dim], - dim=-1, - ) - index_q = index_q.view(-1, self.attn.index_n_heads, self.attn.index_head_dim) - q = q.view(-1, self.attn.num_local_heads, self.attn.qk_head_dim) - - # Step 4. Second fused stage: - # Q RoPE + Index Q RoPE + Index Q FP8 + index-weight scaling - # + W_UK_T absorption + MQA FP8 query packing. - q_nope, q_pe = q.split([mla.qk_nope_head_dim, mla.qk_rope_head_dim], dim=-1) - q_nope = q_nope.transpose(0, 1) - ql_nope = torch.bmm(q_nope, mla.W_UK_T) - ql_nope = ql_nope.transpose(0, 1) - - assert mla.kv_cache_dtype.startswith("fp8") - assert mla.impl.supports_quant_query_input - - index_q_fp8, index_weights, mqa_q = fused_q( - positions, - q_pe, - self.attn.rotary_emb.cos_sin_cache, - index_q, - self.attn.indexer_rope_emb.cos_sin_cache, - ql_nope, - mla._q_scale, - index_weights, - self.attn.indexer_softmax_scale, - self.attn.index_n_heads**-0.5, - ) - - # Step 5. Sparse indexer. - # The FP8 quant + cache write for index_k is already done in - # fused_norm_rope (step 2) when slot_mapping is available. - sparse_attn_indexer( - self.attn.indexer_k_cache.prefix, - self.attn.indexer_k_cache.kv_cache, - index_q_fp8, - index_weights, - self.attn.topk_tokens, - self.attn.index_head_dim, - self.max_model_len, - self.indexer_workspace_size, - self.attn.topk_indices_buffer, - ) - - # Step 6. MLA sparse decode attention (inlined). - # The KV cache update was already done in fused_norm_rope (step 2). + # Steps 2-6. Combined: fused norm/rope + Q projections + sparse MLA. + mla = self.attn.mla_attn output_shape = (hidden_states.shape[0], mla.num_heads * mla.v_head_dim) output_dtype = mla.W_UV.dtype - if mla_attn_metadata is None or slot_mapping is None: - attn_out = torch.zeros( - output_shape, - dtype=output_dtype, - device=hidden_states.device, - ) - else: - num_actual_toks = mla_attn_metadata.num_actual_tokens - mqa_q = mqa_q[:num_actual_toks] - kv_cache = mla.kv_cache - if ( - mla.kv_cache_dtype.startswith("fp8") - and mla.kv_cache_dtype != "fp8_ds_mla" - ): - kv_cache = kv_cache.view(torch.float8_e4m3fn) + attn_out = torch.empty( + output_shape, + dtype=output_dtype, + device=hidden_states.device, + ) + attn_out = torch.ops.vllm.monolithic_attn( + positions, + q_c, + kv_c, + k_pe, + index_k, + index_weights, + attn_out, + self.layer_name, + ) - attn_out, _ = mla.impl.forward_mqa( - mqa_q, - kv_cache, - mla_attn_metadata, - mla, - ) - - output = torch.empty( - output_shape, - dtype=output_dtype, - device=kv_cache.device, - ) - x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1) - out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim) - out = out.transpose(0, 1) - torch.bmm(x, mla.W_UV, out=out) - attn_out = output - - # Step 7. Output projection (AllReduce disabled when fused). hidden_states, _ = self.attn.o_proj(attn_out) - - # Post-attn norm + residual - # Fuse the o_proj AllReduce with post-attention RMSNorm. - hidden_states, residual = allreduce_add_rms_norm( - hidden_states, - residual, - self.post_attention_layernorm_weight, - self.rms_norm_eps, - self._fi_params, - ) - - # MLP / MoE - # When fused_allreduce_rms is enabled, the MLP/MoE AllReduce is - # deferred — it will be fused with the next layer's input norm. - if self.is_moe: - # MoE returns raw (shared_output, routed_output) without - # applying scale + add. torch.compile fuses the elementwise ops. - shared_output, routed_output = self.mlp(hidden_states) - hidden_states = scale_and_add( - routed_output, self.routed_scaling_factor, shared_output - ) - else: - hidden_states = self.mlp(hidden_states) - + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) return hidden_states, residual - - def fuse_shared_expert_act_quant(self) -> None: - """Fuse SiLU-and-Mul + NVFP4 quantize in the shared expert MLP. - - Replaces the shared expert forward so that the activation and the - FP4 input quantisation happen in a single kernel. The monolithic - path always runs on Blackwell with FLASHINFER_CUTLASS for linear - NVFP4, so we hard-code that backend. - """ - if not self.is_moe: - return - shared_experts = self.mlp.shared_experts - if shared_experts is None: - return - - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4LinearMethod, - ) - - if not isinstance( - shared_experts.down_proj.quant_method, ModelOptNvFp4LinearMethod - ): - return - - from vllm._custom_ops import silu_and_mul_nvfp4_quant - from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm - - dp = shared_experts.down_proj - - def _fused_forward(x: torch.Tensor) -> torch.Tensor: - gate_up, _ = shared_experts.gate_up_proj(x) - x_fp4, x_bs = silu_and_mul_nvfp4_quant(gate_up, dp.input_global_scale_inv) - out = flashinfer_scaled_fp4_mm( - x_fp4, - dp.weight, - x_bs, - dp.weight_scale, - dp.alpha, - gate_up.dtype, - backend="cutlass", - ) - return out - - shared_experts.forward = _fused_forward - - -@torch.compile -def scale_and_add(x: torch.Tensor, scale: float, y: torch.Tensor) -> torch.Tensor: - orig_dtype = x.dtype - x = x.to(torch.float32) - x = x * scale - z = x + y.to(torch.float32) - return z.to(orig_dtype) diff --git a/vllm/model_executor/models/deepseek_v3_2_monolithic/model.py b/vllm/model_executor/models/deepseek_v3_2_monolithic/model.py index 4d3819878c6..6cf8961aa6c 100644 --- a/vllm/model_executor/models/deepseek_v3_2_monolithic/model.py +++ b/vllm/model_executor/models/deepseek_v3_2_monolithic/model.py @@ -10,28 +10,24 @@ from collections.abc import Iterable import torch from torch import nn +from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_gather, -) +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from vllm.platforms import current_platform -from .allreduce_rms import ( - AllReduceRMSParams, - allreduce_add_rms_norm, - should_use_allreduce_rms, -) from .decoder_layer import MonolithicDecoderLayer logger = init_logger(__name__) +@support_torch_compile class MonolithicDeepseekV32Model(nn.Module): """Transformer backbone.""" @@ -43,13 +39,6 @@ class MonolithicDeepseekV32Model(nn.Module): quant_config = vllm_config.quant_config self.config = config self.device = current_platform.device_type - self.fi_params: AllReduceRMSParams | None = None - if should_use_allreduce_rms(): - self.fi_params = AllReduceRMSParams(vllm_config, config.hidden_size) - logger.info( - "Enabling fused AllReduce + RMSNorm in monolithic " - "DeepSeek V3.2 decoder layers." - ) topk_tokens = config.index_topk self.topk_indices_buffer = torch.empty( @@ -74,16 +63,16 @@ class MonolithicDeepseekV32Model(nn.Module): layer_idx=i, topk_indices_buffer=self.topk_indices_buffer, prefix=f"{prefix}.layers.{i}", - fi_params=self.fi_params, ) for i in range(config.num_hidden_layers) ] ) - self.norm_weight = nn.Parameter( - torch.ones(config.hidden_size, dtype=torch.get_default_dtype()) + self.norm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=torch.get_default_dtype(), ) - self.rms_norm_eps = config.rms_norm_eps def forward( self, @@ -94,15 +83,7 @@ class MonolithicDeepseekV32Model(nn.Module): residual = None for layer in self.layers: hidden_states, residual = layer(positions, hidden_states, residual) - # After the last layer, hidden_states is unreduced when fused. - # AllReduce before the final norm. - hidden_states, _ = allreduce_add_rms_norm( - hidden_states, - residual, - self.norm_weight, - self.rms_norm_eps, - self.fi_params, - ) + hidden_states, _ = self.norm(hidden_states, residual) return hidden_states @@ -134,6 +115,7 @@ class DeepseekV32MonolithicForCausalLM(nn.Module): quant_config=quant_config, prefix=f"{prefix}.lm_head" if prefix else "lm_head", ) + self.logits_processor = LogitsProcessor(config.vocab_size) self.num_redundant_experts = 0 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -149,10 +131,7 @@ class DeepseekV32MonolithicForCausalLM(nn.Module): return self.model(input_ids, positions) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: - logits = self.lm_head.quant_method.apply(self.lm_head, hidden_states) - logits = tensor_model_parallel_all_gather(logits) - logits = logits[..., : self.config.vocab_size] - return logits + return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Delegate to the original DeepSeek V2 weight loader. @@ -176,21 +155,15 @@ class DeepseekV32MonolithicForCausalLM(nn.Module): # Fuse indexer linear weights after loading. for layer in self.model.layers: layer.fuse_indexer_weights() - layer.fuse_shared_expert_act_quant() return loaded def _remap(self, name: str) -> str: """Remap only names that differ from original model structure.""" - # Only remap layernorms (raw params) and indexer (underscore prefix). + # Only remap layernorms and indexer (underscore prefix). # Everything else (fused_qkv_a_proj, experts, gate, etc.) uses the # same module paths as the original model. replacements = [ - ("input_layernorm.weight", "input_layernorm_weight"), - ( - "post_attention_layernorm.weight", - "post_attention_layernorm_weight", - ), ( "self_attn.q_a_layernorm.weight", "attn.q_a_layernorm_weight", @@ -203,7 +176,6 @@ class DeepseekV32MonolithicForCausalLM(nn.Module): ("self_attn.kv_b_proj", "attn.kv_b_proj"), ("self_attn.o_proj", "attn.o_proj"), ("self_attn.indexer.", "attn.indexer_"), - ("model.norm.weight", "model.norm_weight"), ] for old, new in replacements: if old in name: diff --git a/vllm/model_executor/models/deepseek_v3_2_monolithic/silu_mul_nvfp4_quant.py b/vllm/model_executor/models/deepseek_v3_2_monolithic/silu_mul_nvfp4_quant.py new file mode 100644 index 00000000000..363d9aef8e6 --- /dev/null +++ b/vllm/model_executor/models/deepseek_v3_2_monolithic/silu_mul_nvfp4_quant.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Triton fused SiLU-and-Mul + NVFP4 (e2m1) quantization. + +Uses PTX inline assembly to match the CUDA kernel's fast-math intrinsics +and e2m1 conversion bitwise-exactly: + - rcp.approx.ftz.f32 (reciprocal_approximate_ftz) + - ex2.approx.ftz.f32 (__expf via base-2 fast exp) + - cvt.rn.satfinite.e2m1x2.f32 (float32 → packed e2m1) +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from vllm._custom_ops import create_fp4_output_tensors + +# ── PTX helpers ────────────────────────────────────────────────────────── + + +@triton.jit +def _rcp_approx_ftz(x): + """rcp.approx.ftz.f32 — fast reciprocal (~1 mantissa-bit precision).""" + return tl.inline_asm_elementwise( + asm="rcp.approx.ftz.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _ex2_approx_ftz(x): + """ex2.approx.ftz.f32 — fast 2^x (used to implement __expf).""" + return tl.inline_asm_elementwise( + asm="ex2.approx.ftz.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _cvt_e2m1x2(even, odd): + """cvt.rn.satfinite.e2m1x2.f32 — pack two f32 into one byte of e2m1.""" + return tl.inline_asm_elementwise( + asm=( + "{ .reg .b8 tmp;" + " cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1;" + " cvt.u32.u8 $0, tmp; }" + ), + constraints="=r,f,f", + args=[even, odd], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + + +# ── kernel ─────────────────────────────────────────────────────────────── + + +@triton.jit +def _silu_mul_nvfp4_quant_kernel( + input_ptr, # [M, 2*N] bfloat16 + output_ptr, # [M, N//2] uint8 (packed e2m1 pairs) + sf_out_ptr, # swizzled scale factors (viewed as uint8*) + sf_scale_ptr, # global scale (float32 scalar) + M, + N, # output dim (= input last dim / 2) + stride_in_m, + stride_out_m, + num_k_tiles, # ceil(N / 64) + PAIRS: tl.constexpr, # = 8 (16 values → 8 packed bytes) +): + """One program = one (row, quantisation-group-of-16) pair.""" + LOG2E: tl.constexpr = 1.4426950408889634 + + pid_m = tl.program_id(0) + pid_g = tl.program_id(1) + + col_base = pid_g * (PAIRS * 2) # = pid_g * 16 + pair_off = tl.arange(0, PAIRS) # [0..7] + even_cols = col_base + 2 * pair_off # [0,2,4,...,14] + col_base + odd_cols = even_cols + 1 # [1,3,5,...,15] + col_base + + row_off = pid_m * stride_in_m + + # ── load gate / up (BF16 → float32) ── + gate_e = tl.load(input_ptr + row_off + even_cols).to(tl.float32) + up_e = tl.load(input_ptr + row_off + N + even_cols).to(tl.float32) + gate_o = tl.load(input_ptr + row_off + odd_cols).to(tl.float32) + up_o = tl.load(input_ptr + row_off + N + odd_cols).to(tl.float32) + + # ── fast-math SiLU: __fdividef(x, 1 + __expf(-x)) ── + # __expf(y) = ex2.approx.ftz(y * log2(e)) + # __fdividef(a, b) = a * rcp.approx.ftz(b) + exp_e = _ex2_approx_ftz((-gate_e) * LOG2E) + exp_o = _ex2_approx_ftz((-gate_o) * LOG2E) + silu_e = gate_e * _rcp_approx_ftz(1.0 + exp_e) + silu_o = gate_o * _rcp_approx_ftz(1.0 + exp_o) + + res_e = silu_e * up_e + res_o = silu_o * up_o + + # ── BF16 round-trip (matches compute_silu_mul returning PackedVec) ── + res_e = res_e.to(tl.bfloat16).to(tl.float32) + res_o = res_o.to(tl.bfloat16).to(tl.float32) + + # ── per-group abs-max (over all 16 values) ── + amax = tl.maximum(tl.max(tl.abs(res_e)), tl.max(tl.abs(res_o))) + + # ── scale factor: sf_scale * (amax * rcp.approx.ftz(6.0)) ── + sf_scale_val = tl.load(sf_scale_ptr).to(tl.float32) + sf_raw = sf_scale_val * (amax * _rcp_approx_ftz(6.0)) + + # quantise sf → float8_e4m3fn → float32 (exact round-trip) + sf_fp8 = sf_raw.to(tl.float8e4nv) + sf_rounded = sf_fp8.to(tl.float32) + + # ── write scale byte into swizzled layout ── + sf_col = pid_g + m_tile = pid_m // 128 + outer_m = pid_m % 32 + inner_m = (pid_m // 32) % 4 + k_tile = sf_col // 4 + inner_k = sf_col % 4 + sf_offset = ( + (m_tile * num_k_tiles + k_tile) * 512 + outer_m * 16 + inner_m * 4 + inner_k + ) + sf_byte = sf_fp8.to(tl.uint8, bitcast=True) + tl.store(sf_out_ptr + sf_offset, sf_byte) + + # ── output scale: rcp(sf_rounded * rcp(sf_scale)) ── + rcp_sf_scale = _rcp_approx_ftz(sf_scale_val) + out_scale = tl.where( + sf_rounded != 0.0, + _rcp_approx_ftz(sf_rounded * rcp_sf_scale), + 0.0, + ) + + # ── scale values ── + scaled_e = res_e * out_scale + scaled_o = res_o * out_scale + + # ── PTX e2m1 conversion + pack ── + packed = _cvt_e2m1x2(scaled_e, scaled_o).to(tl.uint8) + + # ── store 8 packed bytes ── + out_off = pid_m * stride_out_m + col_base // 2 + pair_off + tl.store(output_ptr + out_off, packed) + + +# ── Python entry point ─────────────────────────────────────────────────── + + +def silu_and_mul_nvfp4_quant( + input: torch.Tensor, + input_global_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Drop-in replacement for vllm._custom_ops.silu_and_mul_nvfp4_quant.""" + if input.ndim == 1: + input = input.unsqueeze(0) + else: + input = input.reshape(-1, input.shape[-1]) + + M, two_N = input.shape + N = two_N // 2 + assert N % 16 == 0, f"N must be a multiple of 16, got {N}" + + output, output_scale = create_fp4_output_tensors( + M, + N, + input.device, + is_sf_swizzled_layout=True, + ) + + num_groups = N // 16 + num_k_tiles = (num_groups + 3) // 4 # ceil(N / 64) + + grid = (M, num_groups) + _silu_mul_nvfp4_quant_kernel[grid]( + input, + output, + output_scale.view(torch.uint8), + input_global_scale, + M, + N, + input.stride(0), + output.stride(0), + num_k_tiles, + PAIRS=8, + ) + + output_scale = output_scale.view(torch.float8_e4m3fn) + return output, output_scale diff --git a/vllm/v1/attention/ops/concat_and_cache_mla.py b/vllm/v1/attention/ops/concat_and_cache_mla.py new file mode 100644 index 00000000000..d8cf4f88f6c --- /dev/null +++ b/vllm/v1/attention/ops/concat_and_cache_mla.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton implementation of concat_and_cache_mla_kernel. + +Concatenates compressed KV (kv_c) and positional-encoding keys (k_pe) +into a packed MLA KV cache, with optional FP8 quantization. + +Cache layout per slot: [kv_lora_rank | pe_dim] elements, contiguous. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _concat_and_cache_mla_kernel( + # Pointers + kv_c_ptr, # [num_tokens, kv_lora_rank] + k_pe_ptr, # [num_tokens, pe_dim] + kv_cache_ptr, # [num_blocks, block_size, entry_stride] + slot_mapping_ptr, # [num_tokens], int64 + scale_ptr, # [1], float32 + # Strides (in elements, not bytes) + kv_c_stride, + k_pe_stride, + block_stride, + entry_stride, + # Dimensions + kv_lora_rank: tl.constexpr, + pe_dim: tl.constexpr, + block_size, + # Mode + QUANTIZE_FP8: tl.constexpr, + KV_C_BLOCK: tl.constexpr, # next_power_of_2(kv_lora_rank) + PE_BLOCK: tl.constexpr, # next_power_of_2(pe_dim) +): + """Copy one token's kv_c and k_pe into the KV cache. + + Grid: (num_tokens,) + """ + token_idx = tl.program_id(0) + + slot_idx = tl.load(slot_mapping_ptr + token_idx) + # slot_idx == -1 for padded tokens + if slot_idx < 0: + return + + block_idx = slot_idx // block_size + block_offset = slot_idx % block_size + dst_base = block_idx * block_stride + block_offset * entry_stride + + # -- Copy kv_c (kv_lora_rank elements) -- + kv_c_offsets = tl.arange(0, KV_C_BLOCK) + kv_c_mask = kv_c_offsets < kv_lora_rank + kv_c_vals = tl.load( + kv_c_ptr + token_idx * kv_c_stride + kv_c_offsets, + mask=kv_c_mask, + other=0.0, + ) + if QUANTIZE_FP8: + scale = tl.load(scale_ptr) + kv_c_fp8 = (kv_c_vals.to(tl.float32) / scale).to(tl.float8e4nv) + tl.store(kv_cache_ptr + dst_base + kv_c_offsets, kv_c_fp8, mask=kv_c_mask) + else: + tl.store(kv_cache_ptr + dst_base + kv_c_offsets, kv_c_vals, mask=kv_c_mask) + + # -- Copy k_pe (pe_dim elements) at offset kv_lora_rank -- + pe_offsets = tl.arange(0, PE_BLOCK) + pe_mask = pe_offsets < pe_dim + k_pe_vals = tl.load( + k_pe_ptr + token_idx * k_pe_stride + pe_offsets, + mask=pe_mask, + other=0.0, + ) + if QUANTIZE_FP8: + k_pe_fp8 = (k_pe_vals.to(tl.float32) / scale).to(tl.float8e4nv) + tl.store( + kv_cache_ptr + dst_base + kv_lora_rank + pe_offsets, + k_pe_fp8, + mask=pe_mask, + ) + else: + tl.store( + kv_cache_ptr + dst_base + kv_lora_rank + pe_offsets, + k_pe_vals, + mask=pe_mask, + ) + + +def concat_and_cache_mla( + kv_c: torch.Tensor, # [num_tokens, kv_lora_rank] + k_pe: torch.Tensor, # [num_tokens, pe_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, entry_stride] + slot_mapping: torch.Tensor, # [num_tokens] + kv_cache_dtype: str, + scale: torch.Tensor, # [1] +) -> None: + """Concatenate kv_c and k_pe into an MLA KV cache. + + Drop-in replacement for the CUDA ``concat_and_cache_mla`` op. + + Args: + kv_c: Compressed KV of shape ``[num_tokens, kv_lora_rank]``. + k_pe: Positional-encoding keys of shape ``[num_tokens, pe_dim]``. + kv_cache: Cache of shape ``[num_blocks, block_size, entry_stride]``. + slot_mapping: Slot index per token (``-1`` = padding). + kv_cache_dtype: ``"auto"`` for direct copy, ``"fp8_e4m3"`` for + FP8 quantization with the provided ``scale``. + scale: Scalar FP8 scale factor (used only when + ``kv_cache_dtype != "auto"``). + """ + num_tokens = slot_mapping.shape[0] + kv_lora_rank = kv_c.shape[1] + pe_dim = k_pe.shape[1] + block_size = kv_cache.shape[1] + + kv_c_stride = kv_c.stride(0) + k_pe_stride = k_pe.stride(0) + block_stride = kv_cache.stride(0) + entry_stride = kv_cache.stride(1) + + quantize_fp8 = kv_cache_dtype != "auto" + + # When quantizing to FP8, view the (uint8) cache as float8_e4m3fn so + # that Triton's store sees an fp8-typed pointer. + if quantize_fp8 and kv_cache.dtype == torch.uint8: + kv_cache_view = kv_cache.view(torch.float8_e4m3fn) + else: + kv_cache_view = kv_cache + + grid = (num_tokens,) + _concat_and_cache_mla_kernel[grid]( + kv_c, + k_pe, + kv_cache_view, + slot_mapping, + scale, + kv_c_stride, + k_pe_stride, + block_stride, + entry_stride, + kv_lora_rank=kv_lora_rank, + pe_dim=pe_dim, + block_size=block_size, + QUANTIZE_FP8=quantize_fp8, + KV_C_BLOCK=triton.next_power_of_2(kv_lora_rank), + PE_BLOCK=triton.next_power_of_2(pe_dim), + )