diff --git a/docs_new/docs/advanced_features/attention_backend.mdx b/docs_new/docs/advanced_features/attention_backend.mdx index 8057f803b..f31adc78d 100644 --- a/docs_new/docs/advanced_features/attention_backend.mdx +++ b/docs_new/docs/advanced_features/attention_backend.mdx @@ -387,7 +387,8 @@ On SM100/SM103 with CUDA 13+, SGLang automatically selects FlashInfer for GDN pr GDN models are hybrid: the full-attention layers still require a standard `--attention-backend`. Platform constraints for the full-attention backend on hybrid GDN models: -- **Blackwell (e.g., B200)**: `triton`, `trtllm_mha`, or `fa4` only. +- **Blackwell SM120 (e.g., RTX PRO 6000 Blackwell)**: `triton` or `flashinfer` for prefill/full attention; `trtllm_mha` is supported for `--decode-attention-backend` only. +- **Other Blackwell variants (including SM100 B200/GB200)**: `triton`, `trtllm_mha`, or `fa4` only. - **NPU (Ascend)**: `ascend` only. - **AMD (ROCm)**: `triton` recommended. - **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints. diff --git a/docs_new/docs/advanced_features/quantized_kv_cache.mdx b/docs_new/docs/advanced_features/quantized_kv_cache.mdx index 034741fe0..a8b5d82f6 100644 --- a/docs_new/docs/advanced_features/quantized_kv_cache.mdx +++ b/docs_new/docs/advanced_features/quantized_kv_cache.mdx @@ -51,10 +51,15 @@ python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-R1-0528 \ --kv-cache-dtype fp8_e4m3 \ -# Enable FP4 E2M1 KV cache +# Enable NVFP4 FP4 E2M1 KV cache python3 -m sglang.launch_server \ --model-path nvidia/DeepSeek-R1-0528-NVFP4 \ - --kv-cache-dtype fp4_e2m1 \ + --kv-cache-dtype nvfp4 \ + +# Enable block-size-16 FP4 E2M1 KV cache +python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-0528 \ + --kv-cache-dtype fp4_mx_block16 \ ``` ### Scaling Factors @@ -245,7 +250,7 @@ Evaluate FP4 accuracy on your specific model and workload. Large models on simpl ## Best Practices - **Use pre-quantized models**: Prefer models quantized offline with scaling factors included in the checkpoint. -- **Choose the right format**: Use `fp8_e4m3` for better accuracy (recommended), `fp8_e5m2` for larger dynamic range, or `fp4_e2m1` for maximum memory savings (experimental) +- **Choose the right format**: Use `fp8_e4m3` for better accuracy (recommended), `fp8_e5m2` for larger dynamic range, or `nvfp4` / `fp4_mx_block16` for maximum memory savings (experimental) - **Check backend compatibility**: Verify that your chosen attention backend supports quantized KV cache diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index c9815da81..12df59899 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -342,9 +342,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--kv-cache-dtype` - Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+ + Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "nvfp4" selects the NVFP4 FP4 E2M1 KV cache recipe; "fp4_mx_block16" selects the block-size-16 FP4 E2M1 KV cache recipe. Both require CUDA 12.8+ and PyTorch 2.8.0+ `auto` - auto, fp8_e5m2, fp8_e4m3, bf16, bfloat16, fp4_e2m1 + auto, fp8_e5m2, fp8_e4m3, bf16, bfloat16, nvfp4, fp4_mx_block16 `--enable-fp32-lm-head` diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 49821659b..4a36578ca 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -291,7 +291,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac from sglang.srt.layers.attention.linear.utils import ( initialize_linear_attn_config, ) - from sglang.srt.utils import is_blackwell, is_npu + from sglang.srt.utils import ( + is_blackwell, + is_npu, + is_sm120_supported, + ) if not is_npu(): from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( @@ -320,13 +324,25 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac hybrid_backend_cls = HybridLinearAttnBackend if hybrid_gdn_config(runner.model_config) is not None: if is_blackwell(): - assert ( - runner.server_args.attention_backend == "triton" - or runner.server_args.attention_backend == "trtllm_mha" - or runner.server_args.attention_backend == "fa4" - or runner.server_args.attention_backend == "flashinfer" - ), "triton, trtllm_mha, fa4, or flashinfer backend are the only supported backends on Blackwell GPUs for hybrid GDN models, use --attention-backend to specify the backend." - if is_npu(): + if is_sm120_supported(): + allowed = {"triton", "trtllm_mha", "flashinfer"} + else: + allowed = {"triton", "trtllm_mha", "fa4"} + attn_be = runner.server_args.attention_backend + prefill_be = runner.server_args.prefill_attention_backend + decode_be = runner.server_args.decode_attention_backend + # When using split prefill/decode backends, check each individually + if prefill_be and decode_be: + assert prefill_be in allowed and decode_be in allowed, ( + f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. " + f"Got prefill={prefill_be}, decode={decode_be}." + ) + else: + assert attn_be in allowed, ( + f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. " + f"Got attention_backend={attn_be}." + ) + elif is_npu(): assert ( runner.server_args.attention_backend == "ascend" ), "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend." diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 2cffc0ae4..29b26d01f 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -26,6 +26,9 @@ from sglang.kernels.ops.attention.utils import ( from sglang.srt.dllm.config import DllmConfig from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + KVCacheAttentionAccessKind, +) from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.mem_cache.memory_pool import KVWriteLoc @@ -321,9 +324,52 @@ class FlashInferAttnBackend(AttentionBackend): self.dllm_config = DllmConfig.from_server_args(model_runner.server_args) self.is_dllm_model = self.dllm_config is not None + self.kv_cache_quant_method = self.token_to_kv_pool.get_kv_cache_quant_method() + self.prefill_kv_access = self.kv_cache_quant_method.resolve_attention_access( + "prefill", "flashinfer" + ) + self.decode_kv_access = self.kv_cache_quant_method.resolve_attention_access( + "decode", "flashinfer" + ) + prefill_backend, decode_backend = ( + model_runner.server_args.get_attention_backends() + ) + if self.__class__ is FlashInferAttnBackend: + if prefill_backend == "flashinfer": + self._check_kv_attention_access("prefill", self.prefill_kv_access) + if decode_backend == "flashinfer": + self._check_kv_attention_access("decode", self.decode_kv_access) + + self.prefill_uses_dequant_workspace = ( + self.prefill_kv_access is not None + and self.prefill_kv_access.kind + == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE + ) + self.decode_uses_dequant_workspace = ( + self.decode_kv_access is not None + and self.decode_kv_access.kind + == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE + ) + self.is_nvfp4_kvcache = any( + access is not None and access.scale_recipe == "nvfp4" + for access in (self.prefill_kv_access, self.decode_kv_access) + ) + self.dq_page_table = None + self.dq_paged_kernel_lens = None + self.cpu_req_pool_indices = None + # FP4 fake-quant prefill/decode exposes an FP8 workspace to FlashInfer. + self.flashinfer_kv_cache_dtype = ( + torch.float8_e4m3fn + if ( + self.prefill_uses_dequant_workspace + or self.decode_uses_dequant_workspace + ) + else model_runner.kv_cache_dtype + ) + # Parse constants self.decode_use_tensor_cores = should_use_tensor_core( - kv_cache_dtype=model_runner.kv_cache_dtype, + kv_cache_dtype=self.flashinfer_kv_cache_dtype, num_attention_heads=model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size, num_kv_heads=model_runner.model_config.get_num_kv_heads( @@ -331,6 +377,7 @@ class FlashInferAttnBackend(AttentionBackend): ), ) self.max_context_len = model_runner.model_config.context_len + self.page_size = model_runner.page_size self.skip_prefill = skip_prefill self.is_multimodal = model_runner.model_config.is_multimodal assert not ( @@ -490,6 +537,16 @@ class FlashInferAttnBackend(AttentionBackend): List[BatchPrefillWithPagedKVCacheWrapper] ] = None + def _check_kv_attention_access(self, phase: str, access) -> None: + if access is not None: + return + method_name = getattr(self.kv_cache_quant_method, "name", "unknown") + available = self.kv_cache_quant_method.describe_attention_accesses(phase) + raise ValueError( + f"KV cache method {method_name!r} does not support {phase} with " + f"flashinfer attention backend. Available {phase} accesses: {available}." + ) + @staticmethod def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[BaseSWAKVPool]: """Return the SWA KV pool to translate against, or None for non-SWA models. @@ -776,6 +833,86 @@ class FlashInferAttnBackend(AttentionBackend): self.cuda_graph_swa_out_cache_loc[:n] ) + def _prepare_dequant_workspace_metadata_for_extend( + self, forward_batch: ForwardBatch, use_ragged: bool = False + ): + """Prepare FlashInfer metadata for an FP4 dequant workspace. + + Some FP4 recipes store packed KV but expose an FP8 workspace to + FlashInfer prefill. This builds the workspace page table, exact paged + lengths, and CPU request ids needed to populate that workspace before + the prefill kernel runs. + """ + self.dq_page_table = None + self.dq_paged_kernel_lens = None + self.cpu_req_pool_indices = None + if not ( + self.prefill_uses_dequant_workspace + and forward_batch.forward_mode.is_extend_without_speculative() + ): + return + + # Ragged prefill handles current-chunk K/V with raw tensors, so the + # paged side only contains cached prefix lengths. Non-ragged prefill + # uses the dequant workspace for prefix + current chunk, so it needs + # full sequence lengths. These CPU length containers may arrive as + # Python lists or CPU tensors depending on the metadata builder. + paged_seq_lens_cpu = ( + forward_batch.extend_prefix_lens_cpu + if use_ragged + else forward_batch.seq_lens_cpu + ) + raw_paged_seq_lens = ( + paged_seq_lens_cpu + if isinstance(paged_seq_lens_cpu, list) + else paged_seq_lens_cpu.tolist() + ) + paged_seq_lens = [ + int(seq_len.item()) if isinstance(seq_len, torch.Tensor) else int(seq_len) + for seq_len in raw_paged_seq_lens + ] + if sum(paged_seq_lens) <= 0: + self.cpu_req_pool_indices = forward_batch.req_pool_indices.to( + "cpu", non_blocking=True + ) + return + + # dq_buffer layout is page-aligned: each request occupies + # ceil(seq_len/page_size)*page_size slots, starting after a page_size + # dummy prefix. dq_page_table maps only actual token positions and skips + # padding gaps; dq_paged_kernel_lens stores real lengths so FlashInfer + # causal offsets use seq_len - q_len, not page_align(seq_len) - q_len. + seq_lens_with_scratch = paged_seq_lens + [256] + starts = [] + next_start = self.page_size + for seq_len in seq_lens_with_scratch: + starts.append(next_start) + padded_len = ( + (seq_len + self.page_size - 1) // self.page_size + ) * self.page_size + next_start += padded_len + + device = forward_batch.req_pool_indices.device + indices = [ + torch.arange(start, start + seq_len, device=device, dtype=torch.int32) + for start, seq_len in zip(starts, seq_lens_with_scratch) + if seq_len > 0 + ] + self.dq_page_table = torch.cat(indices) if indices else None + self.dq_paged_kernel_lens = torch.tensor( + paged_seq_lens, + dtype=torch.int32, + device=device, + ) + self.cpu_req_pool_indices = forward_batch.req_pool_indices.to( + "cpu", non_blocking=True + ) + + def _kv_write_scales(self, layer: RadixAttention): + if self.kv_cache_quant_method.needs_global_scale(): + return None, None + return layer.k_scale, layer.v_scale + def init_forward_metadata(self, forward_batch: ForwardBatch): swa_out_cache_loc = None if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None: @@ -844,6 +981,10 @@ class FlashInferAttnBackend(AttentionBackend): # Use new backend-specific implementation multi_item_params = self._process_multi_item_scoring(forward_batch) + self._prepare_dequant_workspace_metadata_for_extend( + forward_batch, use_ragged + ) + self.indices_updater_prefill.update( forward_batch.req_pool_indices, forward_batch.seq_lens, @@ -858,6 +999,7 @@ class FlashInferAttnBackend(AttentionBackend): multi_item_params=multi_item_params, cross_attention_custom_mask=forward_batch.cross_attention_custom_mask, extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu, + custom_kv_indices=self.dq_page_table, ) self.forward_metadata = PrefillMetadata( self.prefill_wrappers_paged, @@ -1128,18 +1270,40 @@ class FlashInferAttnBackend(AttentionBackend): logits_soft_cap = layer.logit_cap q = q.contiguous() + + assert not ( + self.prefill_uses_dequant_workspace and layer.is_cross_attention + ), "FP4 dequant KV cache is not supported for cross-attention" + + # We perform dequant for chunk prefill/cache reuse. + pool = self.token_to_kv_pool + if self.prefill_uses_dequant_workspace: + kv_cache = pool.get_flashinfer_dequant_workspace_kv_buffer( + layer, + self.req_to_token_pool.req_to_token, + self.cpu_req_pool_indices, + forward_batch.extend_prefix_lens_cpu, + forward_batch.extend_seq_lens_cpu, + self.page_size, + prepare_workspace=self.dq_page_table is not None, + use_ragged=self.forward_metadata.use_ragged, + k_cur=k, + v_cur=v, + ) + else: + kv_cache = pool.get_kv_buffer(layer.layer_id) + + # use paged attention if not self.forward_metadata.use_ragged: - if k is not None: + if k is not None and save_kv_cache: assert v is not None - if save_kv_cache: - self.token_to_kv_pool.set_kv_buffer( - layer, - KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), - k, - v, - layer.k_scale, - layer.v_scale, - ) + self.token_to_kv_pool.set_kv_buffer( + layer, + KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), + k, + v, + *self._kv_write_scales(layer), + ) causal = ( not layer.is_cross_attention @@ -1147,7 +1311,7 @@ class FlashInferAttnBackend(AttentionBackend): ) o = prefill_wrapper_paged.forward( q.view(-1, layer.tp_q_head_num, layer.head_dim), - self.token_to_kv_pool.get_kv_buffer(layer.layer_id), + kv_cache, causal=causal, sm_scale=layer.scaling, # Disable sliding window attention for multi-item scoring: @@ -1175,6 +1339,9 @@ class FlashInferAttnBackend(AttentionBackend): # previously cached context without re-materializing KV tensors (e.g., the # IQuestLoopCoder path uses token_to_kv_pool as the KV source). if k is None and v is None: + assert ( + not self.prefill_uses_dequant_workspace + ), "KV cache must be provided for ragged attention when using FP4 dequant KV cache" k = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)[0] v = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)[1] causal = True @@ -1219,7 +1386,7 @@ class FlashInferAttnBackend(AttentionBackend): ) o2, s2 = prefill_wrapper_paged.forward_return_lse( q.view(-1, layer.tp_q_head_num, layer.head_dim), - self.token_to_kv_pool.get_kv_buffer(layer.layer_id), + kv_cache, causal=False, sm_scale=layer.scaling, window_left=swa_window_left, @@ -1234,8 +1401,7 @@ class FlashInferAttnBackend(AttentionBackend): KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), k, v, - layer.k_scale, - layer.v_scale, + *self._kv_write_scales(layer), ) return o.view(-1, layer.tp_q_head_num * layer.head_dim) @@ -1267,14 +1433,29 @@ class FlashInferAttnBackend(AttentionBackend): KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), k, v, - layer.k_scale, - layer.v_scale, + *self._kv_write_scales(layer), ) + if self.decode_uses_dequant_workspace: + kv_cache = ( + self.token_to_kv_pool.get_flashinfer_decode_dequant_workspace_kv_buffer( + layer, + self.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + ( + forward_batch.seq_lens_cpu + if forward_batch.seq_lens_cpu is not None + else forward_batch.seq_lens + ), + ) + ) + else: + kv_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id) + # Call the wrapped function o = decode_wrapper.forward( q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), - self.token_to_kv_pool.get_kv_buffer(layer.layer_id), + kv_cache, sm_scale=layer.scaling, logits_soft_cap=layer.logit_cap, # Must use _float to avoid device-to-host copy that breaks cuda graph capture. @@ -1306,7 +1487,7 @@ class FlashInferIndicesUpdaterDecode: get_parallel().attn_tp_size ) self.head_dim = model_runner.model_config.head_dim - self.data_type = model_runner.kv_cache_dtype + self.data_type = attn_backend.flashinfer_kv_cache_dtype self.q_data_type = model_runner.dtype self.sliding_window_size = model_runner.sliding_window_size self.attn_backend = attn_backend @@ -1574,7 +1755,7 @@ class FlashInferIndicesUpdaterPrefill: get_parallel().attn_tp_size ) self.head_dim = model_runner.model_config.head_dim - self.data_type = model_runner.kv_cache_dtype + self.data_type = attn_backend.flashinfer_kv_cache_dtype self.q_data_type = model_runner.dtype self.sliding_window_size = model_runner.sliding_window_size self.attn_backend = attn_backend @@ -1610,6 +1791,7 @@ class FlashInferIndicesUpdaterPrefill: multi_item_params: Optional[MultiItemScoringParams] = None, cross_attention_custom_mask: Optional[torch.Tensor] = None, extend_prefix_lens_cpu: Optional[List[int]] = None, + custom_kv_indices: Optional[torch.Tensor] = None, ): # Keep the signature for type checking. It will be assigned during runtime. raise NotImplementedError() @@ -1629,6 +1811,7 @@ class FlashInferIndicesUpdaterPrefill: multi_item_params: Optional[MultiItemScoringParams] = None, cross_attention_custom_mask: Optional[torch.Tensor] = None, extend_prefix_lens_cpu: Optional[List[int]] = None, + custom_kv_indices: Optional[torch.Tensor] = None, ): if use_ragged: assert prefix_lens is not None @@ -1658,6 +1841,7 @@ class FlashInferIndicesUpdaterPrefill: fixed_split_size=fixed_split_size, multi_item_params=multi_item_params, seq_lens_cpu=seq_lens_cpu, + custom_kv_indices=custom_kv_indices, ) def update_sliding_window( @@ -1675,7 +1859,12 @@ class FlashInferIndicesUpdaterPrefill: multi_item_params: Optional[MultiItemScoringParams] = None, cross_attention_custom_mask: Optional[torch.Tensor] = None, extend_prefix_lens_cpu: Optional[List[int]] = None, + custom_kv_indices: Optional[torch.Tensor] = None, ): + if custom_kv_indices is not None: + raise RuntimeError( + "NVFP4 custom KV indices are only supported by the single-wrapper FlashInfer path." + ) if prefix_lens is None: num_accept_tokens = getattr(spec_info, "num_accept_tokens", None) prefix_lens = ( @@ -1796,7 +1985,12 @@ class FlashInferIndicesUpdaterPrefill: multi_item_params: Optional[MultiItemScoringParams] = None, cross_attention_custom_mask: Optional[torch.Tensor] = None, extend_prefix_lens_cpu: Optional[List[int]] = None, + custom_kv_indices: Optional[torch.Tensor] = None, ): + if custom_kv_indices is not None: + raise RuntimeError( + "NVFP4 custom KV indices are not supported for cross-attention." + ) for wrapper_id in range(2): if wrapper_id == 0: # normal attention @@ -1848,28 +2042,43 @@ class FlashInferIndicesUpdaterPrefill: multi_item_params: Optional[MultiItemScoringParams] = None, cross_attention_custom_mask: Optional[torch.Tensor] = None, seq_lens_cpu: Optional[torch.Tensor] = None, + custom_kv_indices: Optional[torch.Tensor] = None, ): bs = len(seq_lens) if spec_info is None: assert prefix_lens is not None assert len(seq_lens) == len(req_pool_indices) # Normal extend - kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) + # custom_kv_indices uses exact dq_paged_kernel_lens so FlashInfer causal + # offsets are based on real token counts, not page-aligned padding. + if ( + custom_kv_indices is not None + and self.attn_backend.dq_paged_kernel_lens is not None + ): + kv_indptr[1 : bs + 1] = torch.cumsum( + self.attn_backend.dq_paged_kernel_lens, dim=0 + ) + else: + kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) kv_indptr = kv_indptr[: bs + 1] - kv_indices = torch.empty( - paged_kernel_lens_sum + 256, - dtype=torch.int32, - device=req_pool_indices.device, - ) - create_flashinfer_kv_indices_triton[(bs,)]( - self.req_to_token, - req_pool_indices, - paged_kernel_lens, - kv_indptr, - kv_start_idx, - kv_indices, - self.req_to_token.shape[1], - ) + + if custom_kv_indices is not None: + kv_indices = custom_kv_indices + else: + kv_indices = torch.empty( + paged_kernel_lens_sum + 256, + dtype=torch.int32, + device=req_pool_indices.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + kv_start_idx, + kv_indices, + self.req_to_token.shape[1], + ) qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0) qo_indptr = qo_indptr[: bs + 1] diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index c6df78213..cbd1a31a3 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -25,6 +25,9 @@ from sglang.srt.layers.attention.flashinfer_backend import ( FlashInferAttnBackend, FlashInferMultiStepDraftBackend, ) +from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + KVCacheAttentionAccessKind, +) from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -104,6 +107,17 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): super().__init__( model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf ) + self.decode_kv_access = self.kv_cache_quant_method.resolve_attention_access( + "decode", "trtllm_mha" + ) + self._check_decode_kv_access() + self.decode_uses_native_fp4 = ( + self.decode_kv_access.kind == KVCacheAttentionAccessKind.NATIVE_FP4 + ) + self.is_nvfp4_kvcache = ( + self.decode_uses_native_fp4 + and self.decode_kv_access.scale_recipe == "nvfp4" + ) config = model_runner.model_config @@ -177,6 +191,24 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): # KV fp8: q_type = fp8, out_type=model_runner.dtype self.is_xqa_impl = is_sm90_supported() or is_sm120_supported() + def _check_decode_kv_access(self) -> None: + supported_kinds = { + KVCacheAttentionAccessKind.PLAIN, + KVCacheAttentionAccessKind.NATIVE_FP4, + } + if ( + self.decode_kv_access is not None + and self.decode_kv_access.kind in supported_kinds + ): + return + + method_name = getattr(self.kv_cache_quant_method, "name", "unknown") + available = self.kv_cache_quant_method.describe_attention_accesses("decode") + raise ValueError( + f"KV cache method {method_name!r} does not support decode with " + f"trtllm_mha. Available decode accesses: {available}." + ) + @staticmethod def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[SWAKVPool]: """Return the SWAKVPool to translate against, or None for non-SWA models. @@ -847,6 +879,45 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): self.forward_metadata = metadata + def _reshape_paged_kv_cache( + self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + layer: RadixAttention, + head_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + k_cache = k_cache.view( + -1, self.page_size, layer.tp_k_head_num, head_dim + ).permute(0, 2, 1, 3) + v_cache = v_cache.view( + -1, self.page_size, layer.tp_v_head_num, head_dim + ).permute(0, 2, 1, 3) + if layer.tp_k_head_num == 1: + k_cache = canonicalize_stride(k_cache) + if layer.tp_v_head_num == 1: + v_cache = canonicalize_stride(v_cache) + return k_cache, v_cache + + def _get_nvfp4_bmm_scales(self, layer: RadixAttention) -> tuple[float, float]: + assert self.is_nvfp4_kvcache + return self.kv_cache_quant_method.get_bmm_scales(layer.layer_id) + + def _get_nvfp4_decode_kv_cache(self, layer: RadixAttention) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + assert self.is_nvfp4_kvcache + k_fp4, v_fp4, k_scale, v_scale = self.token_to_kv_pool.get_raw_kv_buffer( + layer.layer_id + ) + kv_cache = self._reshape_paged_kv_cache( + k_fp4, v_fp4, layer, layer.head_dim // 2 + ) + kv_cache_block_scales = self._reshape_paged_kv_cache( + k_scale, v_scale, layer, layer.head_dim // 16 + ) + return kv_cache, kv_cache_block_scales + def forward_decode( self, q: torch.Tensor, @@ -862,6 +933,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): use_fused_fp8_path = self._should_use_fused_fp8_path(save_kv_cache, k) use_fused_qkv = use_fused_fp8_path and not self.is_xqa_impl + pool = self.token_to_kv_pool if use_fused_fp8_path: fused_q = self._fused_fp8_qkv_kv_cache( @@ -872,15 +944,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): k = None v = None else: - # Use original set_kv_buffer path if save_kv_cache and k is not None: self.token_to_kv_pool.set_kv_buffer( layer, KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc), k, v, - layer.k_scale, - layer.v_scale, + *self._kv_write_scales(layer), ) # For XQA, q_dtype should be bf16. For trtllm-gen, @@ -893,31 +963,26 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): ): q = q.to(torch.float8_e4m3fn) q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) - k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id) - # shape conversion: - # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] - k_cache = k_cache.view( - -1, self.page_size, layer.tp_k_head_num, layer.head_dim - ).permute(0, 2, 1, 3) - v_cache = v_cache.view( - -1, self.page_size, layer.tp_v_head_num, layer.head_dim - ).permute(0, 2, 1, 3) - if layer.tp_k_head_num == 1: - k_cache = canonicalize_stride(k_cache) - if layer.tp_v_head_num == 1: - v_cache = canonicalize_stride(v_cache) + if self.is_nvfp4_kvcache: + kv_cache, kv_cache_block_scales = self._get_nvfp4_decode_kv_cache(layer) + else: + k_cache, v_cache = pool.get_kv_buffer(layer.layer_id) + kv_cache = self._reshape_paged_kv_cache( + k_cache, v_cache, layer, layer.head_dim + ) + kv_cache_block_scales = None - kv_cache = (k_cache, v_cache) - - bmm1_scale, bmm2_scale = self._get_bmm_scales(layer, q_scale) - # sink: additional value per head in the denominator of the softmax. + if self.is_nvfp4_kvcache: + k_scale, v_scale = self._get_nvfp4_bmm_scales(layer) + bmm1_scale = q_scale * k_scale * layer.scaling + bmm2_scale = v_scale + else: + bmm1_scale, bmm2_scale = self._get_bmm_scales(layer, q_scale) attention_sink = kwargs.get("sinks", None) page_table = self._get_layer_page_table(layer, forward_batch) - # Call TRT-LLM kernel - # raw_out: like q, [bs, acc_q_len, num_q_heads, head_dim] but with output dtype o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q, kv_cache=kv_cache, @@ -931,7 +996,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): sinks=attention_sink, skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), out_dtype=self.q_data_type, # model_runner.dtype + kv_cache_sf=kv_cache_block_scales, ) + if self.is_nvfp4_kvcache and o.dtype != self.q_data_type: + o = o.to(self.q_data_type) return o.view(-1, layer.tp_q_head_num * layer.head_dim) @@ -945,6 +1013,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): save_kv_cache=True, **kwargs, ): + if self.decode_uses_native_fp4: + raise RuntimeError( + "TRTLLM MHA with native FP4 KV cache supports decode only; " + "use a separate prefill backend such as flashinfer or triton." + ) + cache_loc = forward_batch.out_cache_loc use_fused_fp8_path = self._should_use_fused_fp8_path(save_kv_cache, k) diff --git a/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py b/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py index b8ec60db2..d859396c6 100644 --- a/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py +++ b/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py @@ -16,38 +16,211 @@ KV cache quantization strategy pattern. Three-player design: quant_method (pure compute) ► Pool (buffer + batch dequant) ► Backend (view adaptation) + +Why do we need attention access rules? +- Problem: torch.float4_e2m1fn_x2 only describes packed FP4 storage. It does + not say whether the recipe is NVFP4 or fp4_mx_block16, nor how scales are + interpreted. +- Problem: prefill and decode may use different KV views for the same recipe. + For example, NVFP4 uses a FlashInfer dequant workspace for prefill, but TRTLLM + MHA consumes native packed FP4 plus scales for decode. +- Problem: putting these recipe/backend combinations directly into each backend + as dtype checks makes unsupported paths hard to spot and future recipes hard + to add safely. +- Approach: the bottom registry declares KVCacheAttentionAccess entries for + each recipe. A backend resolves one entry by (phase, backend_name, tags), + then either uses the declared access pattern or fails fast if that + combination is unsupported. """ from abc import ABC, abstractmethod -from typing import Optional +from dataclasses import dataclass +from enum import Enum +from typing import Iterable, Optional import torch from torch import Tensor from sglang.srt.layers.quantization.kvfp4_tensor import E2M1_MAX +from sglang.srt.utils.common import is_sm100_supported -class FP4KVCacheQuantMethod(ABC): - """Abstract base for FP4 KV cache quantization strategies. +class KVCacheAttentionPhase(str, Enum): + PREFILL = "prefill" + DECODE = "decode" + + +class KVCacheAttentionAccessKind(str, Enum): + # KV cache is already in the dtype/layout expected by the attention backend. + PLAIN = "plain" + # KV cache is stored quantized, then dequantized and decompressed into a + # temporary workspace before attention. + DEQUANT_WORKSPACE = "dequant_workspace" + # Attention backend directly consumes FP4 KV cache storage and scales. + NATIVE_FP4 = "native_fp4" + + +@dataclass(frozen=True) +class KVCacheBackendMatcher: + exact: frozenset[str] = frozenset() + tags: frozenset[str] = frozenset() + any_backend: bool = False + + def matches(self, backend_name: str, backend_tags: Iterable[str]) -> bool: + backend_tags = frozenset(backend_tags) + return ( + self.any_backend + or backend_name in self.exact + or (bool(self.tags) and self.tags.issubset(backend_tags)) + ) + + +@dataclass(frozen=True) +class KVCacheAttentionAccess: + """Describes how one attention backend reads KV cache for one phase. + + Fields: + - phase: prefill or decode stage where this rule applies. + - kind: access mode, such as plain KV, dequant workspace, or native FP4. + - backend_matcher: backend names/tags that select this rule, e.g. + exact={"trtllm_mha"} or a FlashInfer dequant-workspace tag. + - storage_dtype: dtype stored in the KV pool, e.g. torch.float4_e2m1fn_x2. + - attention_kv_dtype: dtype consumed by attention after any conversion, e.g. + FP8 workspace for FlashInfer prefill or packed FP4 for TRTLLM decode. + - scale_recipe: scale semantics for this FP4 recipe, e.g. "nvfp4" or + "fp4_mx_block16". + - workspace_dtype: temporary workspace dtype for dequant/decompress paths; + None means no temporary workspace is needed. + """ + + phase: KVCacheAttentionPhase + kind: KVCacheAttentionAccessKind + backend_matcher: KVCacheBackendMatcher + storage_dtype: Optional[torch.dtype] = None + attention_kv_dtype: Optional[torch.dtype] = None + scale_recipe: Optional[str] = None + workspace_dtype: Optional[torch.dtype] = None + + def matches(self, phase, backend_name: str, backend_tags: Iterable[str]) -> bool: + return self.phase == KVCacheAttentionPhase( + phase + ) and self.backend_matcher.matches(backend_name, backend_tags) + + +class KVCacheQuantMethodBase(ABC): + """Abstract base for KV cache quantization strategies. Owns the quantize/dequantize computation. The Pool owns the buffers and orchestrates the batch dequant loop. Backends only do view/reshape. - - All operations (quantize_and_store, dequantize_prev_kv) use FlashInfer - kernels or pure tensor ops, so they are CUDA-graph compatible. """ name: str SCALE_BLOCK_SIZE: int = 1 + def attention_accesses(self) -> tuple[KVCacheAttentionAccess, ...]: + return KV_CACHE_ATTENTION_ACCESS_REGISTRY.get(self.name, ()) + + def resolve_attention_access( + self, phase, backend_name: str, backend_tags: Iterable[str] = () + ) -> Optional[KVCacheAttentionAccess]: + for access in self.attention_accesses(): + if access.matches(phase, backend_name, backend_tags): + return access + return None + + def describe_attention_accesses(self, phase=None) -> str: + accesses = self.attention_accesses() + if phase is not None: + phase = KVCacheAttentionPhase(phase) + accesses = tuple(access for access in accesses if access.phase == phase) + if not accesses: + return "none" + return "; ".join( + f"{access.phase.value}:{access.kind.value}:" + f"exact={sorted(access.backend_matcher.exact)}:" + f"tags={sorted(access.backend_matcher.tags)}" + for access in accesses + ) + def needs_dequant_workspace(self) -> bool: - """Whether the pool should allocate dq_k_buffer / dq_v_buffer for prefill.""" - return False + """Whether the pool should allocate dq_k_buffer / dq_v_buffer.""" + return any( + access.kind == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE + for access in self.attention_accesses() + ) + + def needs_plain_kv_dequant_read(self) -> bool: + """Whether plain attention reads require dequantizing packed KV first.""" + return any( + access.kind == KVCacheAttentionAccessKind.PLAIN + and access.storage_dtype is not None + for access in self.attention_accesses() + ) + + def dequant_workspace_dtype(self) -> Optional[torch.dtype]: + """Workspace dtype required by DEQUANT_WORKSPACE access rules.""" + workspace_dtypes = set() + for access in self.attention_accesses(): + if access.kind != KVCacheAttentionAccessKind.DEQUANT_WORKSPACE: + continue + if access.workspace_dtype is None: + raise ValueError( + f"KV cache method {self.name!r} declares DEQUANT_WORKSPACE " + "without workspace_dtype." + ) + workspace_dtypes.add(access.workspace_dtype) + + if not workspace_dtypes: + return None + if len(workspace_dtypes) != 1: + raise ValueError( + f"KV cache method {self.name!r} declares multiple dequant " + f"workspace dtypes: {sorted(str(dtype) for dtype in workspace_dtypes)}." + ) + return next(iter(workspace_dtypes)) + + def kv_storage_dtype(self) -> torch.dtype: + """Packed KV storage dtype declared by attention access rules.""" + storage_dtypes = { + access.storage_dtype + for access in self.attention_accesses() + if access.storage_dtype is not None + } + if not storage_dtypes: + return torch.uint8 + if len(storage_dtypes) != 1: + raise ValueError( + f"KV cache method {self.name!r} declares multiple storage " + f"dtypes: {sorted(str(dtype) for dtype in storage_dtypes)}." + ) + return next(iter(storage_dtypes)) + + def plain_attention_kv_dtype(self) -> Optional[torch.dtype]: + """Dtype produced when packed KV is dequantized for PLAIN attention.""" + attention_dtypes = { + access.attention_kv_dtype + for access in self.attention_accesses() + if access.kind == KVCacheAttentionAccessKind.PLAIN + and access.storage_dtype is not None + and access.attention_kv_dtype is not None + } + if not attention_dtypes: + return None + if len(attention_dtypes) != 1: + raise ValueError( + f"KV cache method {self.name!r} declares multiple plain attention " + f"KV dtypes: {sorted(str(dtype) for dtype in attention_dtypes)}." + ) + return next(iter(attention_dtypes)) def needs_global_scale(self) -> bool: """Whether this method uses a per-layer global FP32 scale.""" return False + def scale_buffer_view_dtype(self) -> Optional[torch.dtype]: + """Optional dtype view for stored per-block scales.""" + return None + @abstractmethod def create_buffers( self, size: int, head_num: int, head_dim: int, layer_num: int, device: str @@ -91,23 +264,73 @@ class FP4KVCacheQuantMethod(ABC): """Dequantize stored FP4 KV (selected token indices already applied). Returns: - (k_fp8, v_fp8): Both in torch.float8_e4m3fn dtype with shape - matching the input (after unpacking). These are written into the - shared dequant workspace buffer for the FlashInfer FP8 prefill kernel. + Dequantized K/V tensors with shape matching the input after unpacking. """ + def dequantize_kv_tensor( + self, + fp4_tensor: Tensor, + scales: Tensor, + layer_id: int, + dtype: Optional[torch.dtype] = None, + ) -> Tensor: + """Dequantize one packed FP4 KV tensor for plain attention reads.""" + raise NotImplementedError( + f"KV cache method {self.name!r} does not support plain KV dequant reads." + ) + @abstractmethod def compute_cell_size( self, head_num: int, head_dim: int, num_layers: int, kv_size: int ) -> int: """Per-token memory footprint in bytes (for capacity estimation).""" - def load_scales_from_model(self, model_runner, sm_version: int = None) -> None: + def load_scales_from_model(self, model_runner) -> None: """Load per-layer global scales from model weights (no-op by default).""" pass -class NVFP4KVMethod(FP4KVCacheQuantMethod): +class UnquantizedKVCacheMethod(KVCacheQuantMethodBase): + """Identity method for BF16 / FP8 KV cache: no extra quantization.""" + + name = "unquantized" + SCALE_BLOCK_SIZE = 1 + + def create_buffers(self, size, head_num, head_dim, layer_num, device) -> dict: + pass + + def quantize_and_store( + self, + k_buffer, + v_buffer, + k_scale_buffer, + v_scale_buffer, + loc, + cache_k, + cache_v, + k_scale=None, + v_scale=None, + ) -> None: + raise RuntimeError( + "Unquantized KV cache writes are handled by MHATokenToKVPool.set_kv_buffer." + ) + + def dequantize_prev_kv( + self, k_fp4, k_scales, v_fp4, v_scales, layer_id + ) -> tuple[Tensor, Tensor]: + raise NotImplementedError( + "Unquantized KV cache does not support FP4 KV dequantization." + ) + + def compute_cell_size( + self, head_num: int, head_dim: int, num_layers: int, kv_size: int + ) -> int: + raise NotImplementedError( + "Unquantized KV cache capacity is computed by the default pool configurator." + ) + + +class NVFP4KVCacheMethod(KVCacheQuantMethodBase): """NVFP4 two-level scaling: global FP32 + per-block FP8 E4M3. Supported on SM100 and SM120. @@ -116,29 +339,28 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod): name = "nvfp4" SCALE_BLOCK_SIZE = 16 - def __init__(self, num_layers: int, device: str, sm_version: int = 120): + def __init__(self, num_layers: int, device: str): self.num_layers = num_layers self.device = device - self.sm_version = sm_version # Per-layer global FP32 scales; filled by load_scales_from_model() self.k_scales_gpu = torch.ones(num_layers, dtype=torch.float32, device=device) self.v_scales_gpu = torch.ones(num_layers, dtype=torch.float32, device=device) - - def needs_dequant_workspace(self) -> bool: - return ( - True # prefill uses FP8 dequant workspace; future native FP4 kernel → False - ) + self.k_scales_float = [1.0] * num_layers + self.v_scales_float = [1.0] * num_layers def needs_global_scale(self) -> bool: return True - def load_scales_from_model(self, model_runner, sm_version: int = None) -> None: - if sm_version is not None: - self.sm_version = sm_version + def scale_buffer_view_dtype(self) -> Optional[torch.dtype]: + return torch.float8_e4m3fn + def load_scales_from_model(self, model_runner) -> None: from sglang.srt.model_loader.utils import resolve_language_model - language_model = resolve_language_model(model_runner.model) + model = getattr(model_runner, "model", model_runner) + language_model = ( + model if hasattr(model, "layers") else resolve_language_model(model) + ) attention_layers = [] for layer in language_model.layers: @@ -158,17 +380,19 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod): # k_scales_gpu is indexed by global (absolute) layer_id. Resize if the model # has layers with global IDs larger than what was pre-allocated. - # This happens in hybrid models (e.g., GDN) where only a subset of layers - # are full-attention, but their layer_ids are non-contiguous. max_global_id = max(layer.layer_id for layer in attention_layers) required_size = max_global_id + 1 if required_size > len(self.k_scales_gpu): + old_size = len(self.k_scales_gpu) self.k_scales_gpu = torch.ones( required_size, dtype=torch.float32, device=self.device ) self.v_scales_gpu = torch.ones( required_size, dtype=torch.float32, device=self.device ) + extra_layers = required_size - old_size + self.k_scales_float.extend([1.0] * extra_layers) + self.v_scales_float.extend([1.0] * extra_layers) k_scales_cpu = self.k_scales_gpu.cpu().clone() v_scales_cpu = self.v_scales_gpu.cpu().clone() @@ -192,23 +416,28 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod): # The FP4 data type itself is identical on both architectures. # Reference: TRT-LLM FP8QDQLinearMethod.process_weights_after_loading_fused_qkv_linear # https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/linear.py - if self.sm_version == 100: + if is_sm100_supported(): k_scale *= E2M1_MAX v_scale *= E2M1_MAX k_scales_cpu[layer_id] = k_scale v_scales_cpu[layer_id] = v_scale + self.k_scales_float[layer_id] = k_scale + self.v_scales_float[layer_id] = v_scale self.k_scales_gpu.copy_(k_scales_cpu, non_blocking=True) self.v_scales_gpu.copy_(v_scales_cpu, non_blocking=True) + def get_bmm_scales(self, layer_id: int) -> tuple[float, float]: + return self.k_scales_float[layer_id], self.v_scales_float[layer_id] + def create_buffers( self, size: int, head_num: int, head_dim: int, layer_num: int, device: str ) -> dict: m = size n = head_num k = head_dim - store_dtype = torch.uint8 - dq_dtype = torch.float8_e4m3fn + store_dtype = self.kv_storage_dtype() + dq_dtype = self.dequant_workspace_dtype() k_buffer = [ torch.zeros((m, n, k // 2), dtype=store_dtype, device=device) @@ -230,9 +459,17 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod): ) for _ in range(layer_num) ] - # Shared dequant workspace — one copy, reused per layer during prefill - dq_k_buffer = torch.zeros((m, n, k), dtype=dq_dtype, device=device) - dq_v_buffer = torch.zeros((m, n, k), dtype=dq_dtype, device=device) + # Shared dequant workspace: one copy, reused per layer during prefill. + dq_k_buffer = ( + torch.zeros((m, n, k), dtype=dq_dtype, device=device) + if dq_dtype is not None + else None + ) + dq_v_buffer = ( + torch.zeros((m, n, k), dtype=dq_dtype, device=device) + if dq_dtype is not None + else None + ) return { "k_buffer": k_buffer, @@ -265,10 +502,15 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod): cache_v.contiguous(), v_scale ) - k_buffer[loc] = cache_k.view(torch.uint8) - v_buffer[loc] = cache_v.view(torch.uint8) - k_scale_buffer[loc] = cache_k_fp4_sf.view(torch.uint8) - v_scale_buffer[loc] = cache_v_fp4_sf.view(torch.uint8) + cache_k = cache_k.view(torch.uint8) + cache_v = cache_v.view(torch.uint8) + cache_k_fp4_sf = cache_k_fp4_sf.view(torch.uint8) + cache_v_fp4_sf = cache_v_fp4_sf.view(torch.uint8) + + k_buffer[loc] = cache_k + v_buffer[loc] = cache_v + k_scale_buffer[loc] = cache_k_fp4_sf + v_scale_buffer[loc] = cache_v_fp4_sf def dequantize_prev_kv( self, @@ -300,26 +542,43 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod): scale_size = ( head_num * (head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size ) - # Dequant workspace: shared across layers (not multiplied by num_layers), FP8 - dq_size = head_num * head_dim * 2 * kv_size + # Dequant workspace is shared across layers, not multiplied by num_layers. + dq_dtype = self.dequant_workspace_dtype() + dq_size = ( + head_num + * head_dim + * 2 + * kv_size + * torch.empty((), dtype=dq_dtype).element_size() + if dq_dtype is not None + else 0 + ) return fp4_size + scale_size + dq_size -class BlockFP4KVMethod(FP4KVCacheQuantMethod): - """Block-wise FP4 single-level scaling (similar to MXFP4 but block_size=16).""" +class FP4MXBlock16KVCacheMethod(KVCacheQuantMethodBase): + """Block-16 FP4 E2M1 single-level scaling. - name = "blockfp4" + This is intentionally not called MXFP4: standard MXFP4 uses a block size of + 32, while this KV cache recipe stores one scale per 16 FP4 values. + """ + + name = "fp4_mx_block16" SCALE_BLOCK_SIZE = 16 - def needs_dequant_workspace(self) -> bool: - return True + def __init__( + self, + num_layers: Optional[int] = None, + device: Optional[str] = None, + ): + pass def create_buffers( self, size: int, head_num: int, head_dim: int, layer_num: int, device: str ) -> dict: m = size - store_dtype = torch.uint8 - dq_dtype = torch.float8_e4m3fn + store_dtype = self.kv_storage_dtype() + dq_dtype = self.dequant_workspace_dtype() k_buffer = [ torch.zeros((m, head_num, head_dim // 2), dtype=store_dtype, device=device) @@ -329,7 +588,7 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod): torch.zeros((m, head_num, head_dim // 2), dtype=store_dtype, device=device) for _ in range(layer_num) ] - # MXFP4 flattens head dimensions for scale storage + # Block-16 FP4 flattens head dimensions for scale storage k_scale_buffer = [ torch.zeros( (m, (head_num * head_dim) // self.SCALE_BLOCK_SIZE), @@ -346,11 +605,15 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod): ) for _ in range(layer_num) ] - dq_k_buffer = torch.zeros( - (m, head_num, head_dim), dtype=dq_dtype, device=device + dq_k_buffer = ( + torch.zeros((m, head_num, head_dim), dtype=dq_dtype, device=device) + if dq_dtype is not None + else None ) - dq_v_buffer = torch.zeros( - (m, head_num, head_dim), dtype=dq_dtype, device=device + dq_v_buffer = ( + torch.zeros((m, head_num, head_dim), dtype=dq_dtype, device=device) + if dq_dtype is not None + else None ) return { @@ -375,15 +638,34 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod): k_scale=None, v_scale=None, ) -> None: - from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + FP4MXBlock16KVQuantizeUtil, + ) + + cache_k_fp4, cache_k_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k) + cache_v_fp4, cache_v_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_v) - cache_k_fp4, cache_k_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k) - cache_v_fp4, cache_v_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v) k_buffer[loc] = cache_k_fp4 v_buffer[loc] = cache_v_fp4 k_scale_buffer[loc] = cache_k_sf v_scale_buffer[loc] = cache_v_sf + def dequantize_kv_tensor( + self, + fp4_tensor: Tensor, + scales: Tensor, + layer_id: int, + dtype: Optional[torch.dtype] = None, + ) -> Tensor: + from sglang.srt.layers.quantization.kvfp4_tensor import ( + FP4MXBlock16KVQuantizeUtil, + ) + + target_dtype = dtype or self.plain_attention_kv_dtype() or torch.bfloat16 + return FP4MXBlock16KVQuantizeUtil.batched_dequantize( + fp4_tensor, scales, dtype=target_dtype + ) + def dequantize_prev_kv( self, k_fp4: Tensor, @@ -392,11 +674,10 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod): v_scales: Tensor, layer_id: int, ) -> tuple[Tensor, Tensor]: - from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil - - k_bf16 = BlockFP4KVQuantizeUtil.batched_dequantize(k_fp4, k_scales) - v_bf16 = BlockFP4KVQuantizeUtil.batched_dequantize(v_fp4, v_scales) - return k_bf16.to(torch.float8_e4m3fn), v_bf16.to(torch.float8_e4m3fn) + return ( + self.dequantize_kv_tensor(k_fp4, k_scales, layer_id), + self.dequantize_kv_tensor(v_fp4, v_scales, layer_id), + ) def compute_cell_size( self, head_num: int, head_dim: int, num_layers: int, kv_size: int @@ -405,22 +686,152 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod): scale_size = ( (head_num * head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size ) - dq_size = head_num * head_dim * 2 * kv_size + dq_dtype = self.dequant_workspace_dtype() + dq_size = ( + head_num + * head_dim + * 2 + * kv_size + * torch.empty((), dtype=dq_dtype).element_size() + if dq_dtype is not None + else 0 + ) return fp4_size + scale_size + dq_size -# Registry: name → class. Only classes for fp4_e2m1 dtype need to be listed. -FP4_KV_CACHE_QUANT_REGISTRY: dict[str, type[FP4KVCacheQuantMethod]] = { - "nvfp4": NVFP4KVMethod, - "blockfp4": BlockFP4KVMethod, +# Registry: method name -> attention access rules. +_PREFILL = KVCacheAttentionPhase.PREFILL +_DECODE = KVCacheAttentionPhase.DECODE +_PLAIN_KIND = KVCacheAttentionAccessKind.PLAIN +_DQ_WORKSPACE_KIND = KVCacheAttentionAccessKind.DEQUANT_WORKSPACE +_NATIVE_FP4_KIND = KVCacheAttentionAccessKind.NATIVE_FP4 +_ANY_BACKEND = KVCacheBackendMatcher(any_backend=True) +_NVFP4_SCALE = "nvfp4" +_FP4_MX_SCALE = "fp4_mx_block16" +_FP8_E4M3 = torch.float8_e4m3fn +_TORCH_FP4 = getattr(torch, "float4_e2m1fn_x2", None) +_BF16 = torch.bfloat16 +_NVFP4_PREFILL_BACKENDS = frozenset({"flashinfer"}) +_NVFP4_DECODE_BACKENDS = frozenset({"trtllm_mha"}) +_FP4_MX_MHA_BACKENDS = frozenset( + {"triton", "torch_native", "flex_attention", "trtllm_mha"} +) +_FP4_MX_PREFILL_BACKENDS = _FP4_MX_MHA_BACKENDS | frozenset({"fa4"}) + + +def _backend_matcher(backends) -> KVCacheBackendMatcher: + if isinstance(backends, KVCacheBackendMatcher): + return backends + return KVCacheBackendMatcher(exact=backends) + + +def _plain( + phase: KVCacheAttentionPhase, + backends, + scale: Optional[str] = None, + attention_dtype: Optional[torch.dtype] = None, +) -> KVCacheAttentionAccess: + return KVCacheAttentionAccess( + phase, + _PLAIN_KIND, + _backend_matcher(backends), + storage_dtype=torch.uint8 if scale is not None else None, + attention_kv_dtype=attention_dtype, + scale_recipe=scale, + ) + + +def _dq_workspace( + phase: KVCacheAttentionPhase, + backends, + scale: str, + attention_dtype: torch.dtype, +) -> KVCacheAttentionAccess: + return KVCacheAttentionAccess( + phase, + _DQ_WORKSPACE_KIND, + _backend_matcher(backends), + storage_dtype=torch.uint8, + attention_kv_dtype=attention_dtype, + scale_recipe=scale, + workspace_dtype=attention_dtype, + ) + + +def _native_fp4( + phase: KVCacheAttentionPhase, + backends, + scale: str, + attention_dtype: Optional[torch.dtype], +) -> KVCacheAttentionAccess: + return KVCacheAttentionAccess( + phase, + _NATIVE_FP4_KIND, + _backend_matcher(backends), + storage_dtype=torch.uint8, + attention_kv_dtype=attention_dtype, + scale_recipe=scale, + ) + + +KV_CACHE_ATTENTION_ACCESS_REGISTRY: dict[str, tuple[KVCacheAttentionAccess, ...]] = { + UnquantizedKVCacheMethod.name: ( + _plain(_PREFILL, _ANY_BACKEND), + _plain(_DECODE, _ANY_BACKEND), + ), + NVFP4KVCacheMethod.name: ( + _dq_workspace(_PREFILL, _NVFP4_PREFILL_BACKENDS, _NVFP4_SCALE, _FP8_E4M3), + _native_fp4(_DECODE, _NVFP4_DECODE_BACKENDS, _NVFP4_SCALE, _TORCH_FP4), + ), + FP4MXBlock16KVCacheMethod.name: ( + _plain(_PREFILL, _FP4_MX_PREFILL_BACKENDS, _FP4_MX_SCALE, _BF16), + _plain(_DECODE, _FP4_MX_MHA_BACKENDS, _FP4_MX_SCALE, _BF16), + ), } -def get_fp4_kv_cache_quant_method(name: str, **kwargs) -> FP4KVCacheQuantMethod: - """Instantiate a FP4KVCacheQuantMethod by recipe name.""" - if name not in FP4_KV_CACHE_QUANT_REGISTRY: +# Registry: explicit --kv-cache-dtype value -> method class. +KV_CACHE_QUANT_REGISTRY: dict[str, type[KVCacheQuantMethodBase]] = { + "nvfp4": NVFP4KVCacheMethod, + "fp4_mx_block16": FP4MXBlock16KVCacheMethod, +} + + +def resolve_kv_cache_quant(kv_cache_dtype) -> Optional[str]: + """Resolve the explicit FP4 KV cache recipe from ``--kv-cache-dtype``.""" + + if not isinstance(kv_cache_dtype, str): + if ( + hasattr(torch, "float4_e2m1fn_x2") + and kv_cache_dtype == torch.float4_e2m1fn_x2 + ): + raise ValueError( + "FP4 KV cache storage dtype does not identify the recipe. " + "Pass the explicit --kv-cache-dtype value: 'nvfp4' or 'fp4_mx_block16'." + ) + return None + + if kv_cache_dtype == "fp4_e2m1": raise ValueError( - f"Unknown fp4_kv_cache_recipe: '{name}'. " - f"Available: {list(FP4_KV_CACHE_QUANT_REGISTRY)}" + "--kv-cache-dtype=fp4_e2m1 is deprecated. " + "Use --kv-cache-dtype=fp4_mx_block16." ) - return FP4_KV_CACHE_QUANT_REGISTRY[name](**kwargs) + if kv_cache_dtype == "mxfp4": + raise ValueError( + "--kv-cache-dtype=mxfp4 is reserved for true MXFP4 block-size-32 " + "semantics. Use --kv-cache-dtype=fp4_mx_block16 for the current " + "block-size-16 FP4 KV recipe." + ) + if kv_cache_dtype in KV_CACHE_QUANT_REGISTRY: + return kv_cache_dtype + return None + + +def get_kv_cache_quant_method(name: str, **kwargs) -> KVCacheQuantMethodBase: + """Instantiate a KVCacheQuantMethodBase by internal method name.""" + if name not in KV_CACHE_QUANT_REGISTRY: + raise ValueError( + f"Unknown KV cache quantization method: '{name}'. " + f"Available: {list(KV_CACHE_QUANT_REGISTRY)}" + ) + return KV_CACHE_QUANT_REGISTRY[name](**kwargs) diff --git a/python/sglang/srt/layers/quantization/kvfp4_tensor.py b/python/sglang/srt/layers/quantization/kvfp4_tensor.py index 2d5eb45d1..553cb5f2e 100644 --- a/python/sglang/srt/layers/quantization/kvfp4_tensor.py +++ b/python/sglang/srt/layers/quantization/kvfp4_tensor.py @@ -25,40 +25,35 @@ class FP4KVCacheRecipe(Enum): E2M1_MAX = 6.0 MAX_BLOCK_SCALE_FP8 = 448.0 # Maximum FP8 E4M3 value -# Put constants directly on CUDA if available -_device = "cuda" if torch.cuda.is_available() else "cpu" # E2M1 format: 1 sign bit + 2 exponent bits + 1 mantissa bit = 4 bits # 16 possible values: 0x0-0xF # Negative values: 0x8-0xF (sign bit = 1) # Positive values: 0x0-0x7 (sign bit = 0) -E2M1_VALUES = torch.tensor( - [ - 0, - 0.5, - 1, - 1.5, - 2, - 3, - 4, - 6, # 0x0-0x7: positive values - -0, - -0.5, - -1, - -1.5, - -2, - -3, - -4, - -6, - ], # 0x8-0xF: negative values - dtype=torch.float32, - device=_device, -) -E2M1_BOUNDS = torch.tensor( - [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5], dtype=torch.float32, device=_device +# Keep constants as Python literals. Compiled helpers materialize them with +# input.new_tensor(), so they follow the caller device without a global GPU tensor +# or a CPU tensor .to(device) in the hot path. +E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, ) +E2M1_BOUNDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) -class BlockFP4KVQuantizeUtil: +class FP4MXBlock16KVQuantizeUtil: """Block-wise FP4 (E2M1) quantization for KV cache. Similar to MXFP4 but uses block_size=16 (MXFP4 spec defines block_size=32). @@ -69,6 +64,7 @@ class BlockFP4KVQuantizeUtil: @torch.compile def batched_quantize(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ + Quantize tensor to KVFP4 format Args: tensor: Input tensor of shape [B, M, N] @@ -95,7 +91,8 @@ class BlockFP4KVQuantizeUtil: abs_vals = scaled.abs() # Pure tensor version (CUDA Graph safe) - magnitude_bits = torch.sum(abs_vals.unsqueeze(-1) >= E2M1_BOUNDS, dim=-1) + bounds = tensor.new_tensor(E2M1_BOUNDS, dtype=torch.float32) + magnitude_bits = torch.sum(abs_vals.unsqueeze(-1) >= bounds, dim=-1) # Combine sign and magnitude fp4_vals = sign_bits + magnitude_bits.to(torch.uint8) @@ -136,7 +133,8 @@ class BlockFP4KVQuantizeUtil: magnitude_idx = fp4_vals & 0x07 # Convert to float values - float_vals = E2M1_VALUES[magnitude_idx.long()] + values = quant_tensor.new_tensor(E2M1_VALUES[:8], dtype=torch.float32) + float_vals = values[magnitude_idx.long()] float_vals = torch.where(sign_mask, -float_vals, float_vals) # Reshape for block-wise scaling @@ -177,21 +175,34 @@ class NVFP4KVQuantizeUtil: block_scales: shape [B, M, N/16], dtype float8_e4m3fn global_scale: passthrough """ - from sglang.srt.utils import is_sm90_supported, is_sm100_supported + from sglang.srt.utils import ( + is_sm90_supported, + is_sm100_supported, + is_sm120_supported, + ) - assert is_sm90_supported(), "NVFP4 KV cache quantize requires SM90+ GPU" + assert ( + is_sm100_supported() or is_sm120_supported() or is_sm90_supported() + ), "NVFP4 KV cache quantize requires SM100/SM120 or SM90 fallback GPU" b, m, n = tensor.shape tensor_2d = tensor.reshape(b * m, n) + # The KV cache path passes preloaded per-layer scales already on device. + # Keep scalar/0-d support for tests and future fallback paths, but do not + # silently move tensor scales here. if isinstance(global_scale, (int, float)): global_scale = torch.tensor( [global_scale], dtype=torch.float32, device=tensor.device ) elif global_scale.dim() == 0: global_scale = global_scale.unsqueeze(0) + elif global_scale.device != tensor.device: + raise ValueError( + "NVFP4 global scale tensor must already be on the KV tensor device." + ) - if is_sm100_supported(): + if is_sm100_supported() or is_sm120_supported(): from flashinfer import nvfp4_kv_quantize # nvfp4_kv_quantize takes global_scale directly (not inverted) @@ -238,18 +249,29 @@ class NVFP4KVQuantizeUtil: Returns: Dequantized tensor of shape [B, M, N] """ - from sglang.srt.utils import is_sm100_supported + from sglang.srt.utils import ( + is_sm90_supported, + is_sm100_supported, + is_sm120_supported, + ) b, m, n_half = quant_tensor.shape + # The KV cache path passes preloaded per-layer scales already on device. + # Keep scalar/0-d support for tests and future fallback paths, but do not + # silently move tensor scales here. if isinstance(global_scale, (int, float)): global_scale = torch.tensor( [global_scale], dtype=torch.float32, device=quant_tensor.device ) elif global_scale.dim() == 0: global_scale = global_scale.unsqueeze(0) + elif global_scale.device != quant_tensor.device: + raise ValueError( + "NVFP4 global scale tensor must already be on the KV tensor device." + ) - if is_sm100_supported(): + if is_sm100_supported() or is_sm120_supported(): from flashinfer import nvfp4_kv_dequantize quant_2d = quant_tensor.view(torch.uint8).reshape(b * m, n_half) @@ -259,6 +281,9 @@ class NVFP4KVQuantizeUtil: ) return output_2d.reshape(b, m, -1) else: + assert ( + is_sm90_supported() + ), "NVFP4 KV cache dequantize requires SM100/SM120 or SM90 fallback GPU" # Pure PyTorch fallback for SM90 n = n_half * 2 fp4_vals = torch.empty( @@ -266,7 +291,8 @@ class NVFP4KVQuantizeUtil: ) fp4_vals[..., 0::2] = quant_tensor & 0x0F fp4_vals[..., 1::2] = (quant_tensor >> 4) & 0x0F - float_vals = E2M1_VALUES[fp4_vals.long()] + values = quant_tensor.new_tensor(E2M1_VALUES, dtype=torch.float32) + float_vals = values[fp4_vals.long()] reshaped = float_vals.view(b, m * n // 16, 16) block_scales_float = block_scales.float().unsqueeze(-1) scaled = reshaped * block_scales_float diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 29db2d9f6..998609f82 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -21,6 +21,10 @@ from sglang.srt.configs.model_config import ( ) from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.environ import envs +from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + get_kv_cache_quant_method, + resolve_kv_cache_quant, +) from sglang.srt.mem_cache.allocation_sizing import get_req_to_token_extra_context_len from sglang.srt.mem_cache.allocator import ( BaseTokenToKVPoolAllocator, @@ -153,6 +157,7 @@ class KVCacheConfigurator: gpu_id: int ps: ParallelState pp_group: Any + model: Any model_config: ModelConfig server_args: ServerArgs kv_cache_dtype: torch.dtype @@ -178,6 +183,20 @@ class KVCacheConfigurator: self.mambaish_config = mambaish_config(self.model_config) self.hybrid_gdn_config = hybrid_gdn_config(self.model_config) + def _build_fp4_quant_method(self, *, num_layers: int): + if not is_float4_e2m1fn_x2(self.kv_cache_dtype): + return None + quant_name = resolve_kv_cache_quant(self.server_args.kv_cache_dtype) + if quant_name is None: + return None + quant_method = get_kv_cache_quant_method( + quant_name, + num_layers=num_layers, + device=self.device, + ) + quant_method.load_scales_from_model(self.model) + return quant_method + def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: """Apply a resolved MemoryPoolConfig and initialize pools.""" if not self.spec_algorithm.is_none() and self.is_draft_worker: @@ -349,7 +368,7 @@ class KVCacheConfigurator: "Supported configurations today: plain MHA models on CUDA with the FA " "(fa3/fa4) prefill backend, --is-embedding, --chunked-prefill-size=-1, " "--disable-radix-cache, no context-parallel attention, no HiSparse, " - "and --kv-cache-dtype != fp4_e2m1." + "and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}." ) return _InitializedPools( req_to_token_pool=req_to_token_pool, @@ -552,7 +571,7 @@ class KVCacheConfigurator: f"{unsupported_pool_family}. Supported configurations today: plain MHA " "models on CUDA with the FA (fa3/fa4) prefill backend, --is-embedding, " "--chunked-prefill-size=-1, --disable-radix-cache, no context-parallel " - "attention, no HiSparse, and --kv-cache-dtype != fp4_e2m1." + "attention, no HiSparse, and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}." ) def _build_req_to_token_pool(self, *, max_num_reqs: int) -> ReqToTokenPool: @@ -786,18 +805,19 @@ class KVCacheConfigurator: mha_pool_class=mha_pool_class, ) else: + quant_method = None if is_float4_e2m1fn_x2(self.kv_cache_dtype): assert ( not enable_page_major ), "page-major KV layout is not supported with fp4 KV cache" - token_to_kv_pool = self._build_mha_fp4_kv_pool( - max_total_num_tokens=sizes.max_total_num_tokens, - ) - else: - token_to_kv_pool = self._build_mha_kv_pool( - max_total_num_tokens=sizes.max_total_num_tokens, - mha_pool_class=mha_pool_class, + quant_method = self._build_fp4_quant_method( + num_layers=self.layer_info.num_effective_layers ) + token_to_kv_pool = self._build_mha_kv_pool( + max_total_num_tokens=sizes.max_total_num_tokens, + mha_pool_class=mha_pool_class, + quant_method=quant_method, + ) return token_to_kv_pool def _build_dsv4_kv_pool( @@ -1179,6 +1199,18 @@ class KVCacheConfigurator: "kv_lora_rank": self.model_config.kv_lora_rank, "qk_rope_head_dim": self.model_config.qk_rope_head_dim, } + full_attention_layer_ids = ( + [0] + if self.is_draft_worker + else [ + i + for i in self.mambaish_config.full_attention_layer_ids + if self.layer_info.start_layer <= i < self.layer_info.end_layer + ] + ) + quant_method = self._build_fp4_quant_method( + num_layers=len(full_attention_layer_ids) + ) token_to_kv_pool = HybridLinearKVPool( page_size=self.server_args.page_size, size=max_total_num_tokens, @@ -1186,15 +1218,7 @@ class KVCacheConfigurator: head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, # if draft worker, we only need 1 attention layer's kv pool - full_attention_layer_ids=( - [0] - if self.is_draft_worker - else [ - i - for i in self.mambaish_config.full_attention_layer_ids - if self.layer_info.start_layer <= i < self.layer_info.end_layer - ] - ), + full_attention_layer_ids=full_attention_layer_ids, device=self.device, mamba_pool=req_to_token_pool.mamba_pool, enable_memory_saver=self.server_args.enable_memory_saver, @@ -1202,7 +1226,8 @@ class KVCacheConfigurator: use_mla=self.use_mla_backend, start_layer=self.layer_info.start_layer, full_kv_pool_class=mha_pool_class, - post_capture_active=self.post_capture_kv_active, + quant_method=quant_method, + post_capture_active=self.post_capture_kv_active and quant_method is None, **extra_args, ) return token_to_kv_pool @@ -1226,13 +1251,18 @@ class KVCacheConfigurator: return token_to_kv_pool def _build_mha_kv_pool( - self, *, max_total_num_tokens: int, mha_pool_class: type + self, *, max_total_num_tokens: int, mha_pool_class: type, quant_method=None ) -> KVCache: pool_cls = ( NoOpMHATokenToKVPool if self.server_args.prefill_only_disable_kv_cache else mha_pool_class ) + pool_kwargs = {} + if quant_method is not None: + pool_kwargs["quant_method"] = quant_method + else: + pool_kwargs["post_capture_active"] = self.post_capture_kv_active token_to_kv_pool = pool_cls( max_total_num_tokens, page_size=self.server_args.page_size, @@ -1247,7 +1277,7 @@ class KVCacheConfigurator: end_layer=self.layer_info.end_layer, enable_alt_stream=not self.server_args.enable_pdmux, enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None), - post_capture_active=self.post_capture_kv_active, + **pool_kwargs, ) return token_to_kv_pool diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py index 46baf38e4..9635ea5ce 100644 --- a/python/sglang/srt/mem_cache/kv_cache_dtype.py +++ b/python/sglang/srt/mem_cache/kv_cache_dtype.py @@ -53,14 +53,23 @@ def configure_kv_cache_dtype( elif server_args_kv_cache_dtype in ("bf16", "bfloat16"): kv_cache_dtype = torch.bfloat16 elif server_args_kv_cache_dtype == "fp4_e2m1": + raise ValueError( + "--kv-cache-dtype=fp4_e2m1 is deprecated. " + "Use --kv-cache-dtype=fp4_mx_block16." + ) + elif server_args_kv_cache_dtype in ("nvfp4", "fp4_mx_block16"): if hasattr(torch, "float4_e2m1fn_x2"): kv_cache_dtype = torch.float4_e2m1fn_x2 - logger.warning(f"FP4 (E2M1) KV Cache might lead to a accuracy drop!") - else: logger.warning( - f"--kv-cache-dtype falls back to 'auto' because this torch version does not support torch.float4_e2m1fn_x2" + "%s KV Cache might lead to an accuracy drop!", + server_args_kv_cache_dtype.upper(), + ) + else: + raise ValueError( + f"--kv-cache-dtype={server_args_kv_cache_dtype} requires " + "torch.float4_e2m1fn_x2 support. Please use PyTorch 2.8.0+ " + "with CUDA 12.8+." ) - kv_cache_dtype = model_dtype else: raise ValueError(f"Unsupported kv_cache_dtype: {server_args_kv_cache_dtype}.") diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 844fd175d..0599ec7a1 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -51,6 +51,9 @@ from sglang.srt.configs.mamba_utils import BaseLinearStateParams from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa +from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + UnquantizedKVCacheMethod, +) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator from sglang.srt.mem_cache.kv_vmm_backing import KvVmmBufferOwner @@ -73,6 +76,7 @@ from sglang.srt.utils import ( cpu_has_amx_support, is_cpu, is_cuda, + is_float4_e2m1fn_x2, is_hip, is_npu, next_power_of_2, @@ -1387,6 +1391,24 @@ class KVCache(abc.ABC): def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): raise NotImplementedError() + def get_kv_cache_quant_method(self) -> Any: + """Return the concrete KV quant method, unwrapping composite KV pools.""" + fallback = None + for pool in ( + self, + getattr(self, "full_kv_pool", None), + getattr(self, "swa_kv_pool", None), + ): + if pool is None: + continue + quant_method = getattr(pool, "quant_method", None) + if quant_method is None: + continue + if getattr(quant_method, "name", None) != "unquantized": + return quant_method + fallback = quant_method + return fallback + def maybe_get_custom_mem_pool(self): return self.custom_mem_pool @@ -1411,6 +1433,7 @@ class MHATokenToKVPool(KVCache): enable_alt_stream: bool = True, enable_kv_cache_copy: bool = False, kv_cache_layout: Optional[str] = None, + quant_method=None, post_capture_active: bool = False, ): if post_capture_active: @@ -1482,6 +1505,10 @@ class MHATokenToKVPool(KVCache): assert self.head_dim % self._kv_vector_x == 0 assert self.v_head_dim % self._kv_vector_x == 0 + self.quant_method = ( + quant_method if quant_method is not None else UnquantizedKVCacheMethod() + ) + self._create_buffers() self.device_module = torch.get_device_module(self.device) @@ -1554,12 +1581,100 @@ class MHATokenToKVPool(KVCache): self._kv_copy_config, ) + @property + def is_quantized_kv_cache(self) -> bool: + return not isinstance(self.quant_method, UnquantizedKVCacheMethod) + def _create_buffers(self): - if self.post_capture_active: - self._alloc_post_capture_buffers() + if self.is_quantized_kv_cache: + if self.post_capture_active: + raise NotImplementedError( + "Post-capture KV backing is not supported for quantized KV cache." + ) + self._create_quantized_buffers() else: - self._create_buffers_normal() + self.k_scale_buffer = None + self.v_scale_buffer = None + self.dq_k_buffer = None + self.dq_v_buffer = None + if self.post_capture_active: + self._alloc_post_capture_buffers() + else: + self._create_buffers_normal() self._kv_buffer_descs = self._build_kv_buffer_descs() + self._init_data_ptrs_and_strides() + + def _create_quantized_buffers(self): + # Quantized recipes own packed-data, scale, and workspace shapes. + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + with ( + torch.cuda.use_mem_pool(self.custom_mem_pool) + if self.enable_custom_mem_pool + else nullcontext() + ): + buf = self.quant_method.create_buffers( + self.size + self.page_size, + self.head_num, + self.head_dim, + self.layer_num, + self.device, + ) + self.k_buffer = buf["k_buffer"] + self.v_buffer = buf["v_buffer"] + self.k_scale_buffer = buf.get("k_scale_buffer") + self.v_scale_buffer = buf.get("v_scale_buffer") + self.dq_k_buffer = buf.get("dq_k_buffer") + self.dq_v_buffer = buf.get("dq_v_buffer") + self.store_dtype = buf.get("store_dtype", torch.uint8) + self._check_quantized_buffer_access_requirements() + + def _check_quantized_buffer_access_requirements(self): + expected_workspace_dtype = self.quant_method.dequant_workspace_dtype() + has_k_workspace = self.dq_k_buffer is not None + has_v_workspace = self.dq_v_buffer is not None + if has_k_workspace != has_v_workspace: + raise RuntimeError( + f"KV cache method {self.quant_method.name!r} created only one " + "dequant workspace buffer." + ) + + if expected_workspace_dtype is None: + if has_k_workspace: + raise RuntimeError( + f"KV cache method {self.quant_method.name!r} does not declare " + "DEQUANT_WORKSPACE access but created dequant buffers." + ) + return + + if not has_k_workspace: + raise RuntimeError( + f"KV cache method {self.quant_method.name!r} declares " + "DEQUANT_WORKSPACE access but did not create dequant buffers." + ) + + if ( + self.dq_k_buffer.dtype != expected_workspace_dtype + or self.dq_v_buffer.dtype != expected_workspace_dtype + ): + raise RuntimeError( + f"KV cache method {self.quant_method.name!r} declares dequant " + f"workspace dtype {expected_workspace_dtype}, but created " + f"{self.dq_k_buffer.dtype}/{self.dq_v_buffer.dtype}." + ) + + def _slot_move_pointer_buffers(self): + """Buffers whose pointers/strides are used when KV slots are remapped. + + FP4 KV cache stores data and per-block scales separately, so slot moves + must update both. This list feeds data_ptrs/data_strides; it does not + copy tensor contents by itself. + """ + buffers = [*self.k_buffer, *self.v_buffer] + if getattr(self, "k_scale_buffer", None) is not None: + buffers.extend([*self.k_scale_buffer, *self.v_scale_buffer]) + return buffers + + def _init_data_ptrs_and_strides(self): self.k_data_ptrs = torch.tensor( [x.data_ptr() for x in self.k_buffer], dtype=torch.uint64, @@ -1570,11 +1685,16 @@ class MHATokenToKVPool(KVCache): dtype=torch.uint64, device=self.device, ) - self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0) + slot_move_pointer_buffers = self._slot_move_pointer_buffers() + self.data_ptrs = torch.tensor( + [x.data_ptr() for x in slot_move_pointer_buffers], + dtype=torch.uint64, + device=self.device, + ) self.data_strides = torch.tensor( [ np.prod(x.shape[1:]) * x.dtype.itemsize - for x in self.k_buffer + self.v_buffer + for x in slot_move_pointer_buffers ], device=self.device, ) @@ -1716,6 +1836,14 @@ class MHATokenToKVPool(KVCache): def _clear_buffers(self): del self.k_buffer del self.v_buffer + if hasattr(self, "k_scale_buffer") and self.k_scale_buffer is not None: + del self.k_scale_buffer + if hasattr(self, "v_scale_buffer") and self.v_scale_buffer is not None: + del self.v_scale_buffer + if hasattr(self, "dq_k_buffer") and self.dq_k_buffer is not None: + del self.dq_k_buffer + if hasattr(self, "dq_v_buffer") and self.dq_v_buffer is not None: + del self.dq_v_buffer if self._post_capture_owner is not None: self._post_capture_owner.close() self._post_capture_owner = None @@ -1723,12 +1851,14 @@ class MHATokenToKVPool(KVCache): def get_kv_size_bytes(self): assert hasattr(self, "k_buffer") assert hasattr(self, "v_buffer") - k_size_bytes = 0 - for k_cache in self.k_buffer: - k_size_bytes += get_tensor_size_bytes(k_cache) - v_size_bytes = 0 - for v_cache in self.v_buffer: - v_size_bytes += get_tensor_size_bytes(v_cache) + k_size_bytes = get_tensor_size_bytes(self.k_buffer) + v_size_bytes = get_tensor_size_bytes(self.v_buffer) + if getattr(self, "k_scale_buffer", None) is not None: + k_size_bytes += get_tensor_size_bytes(self.k_scale_buffer) + v_size_bytes += get_tensor_size_bytes(self.v_scale_buffer) + if getattr(self, "dq_k_buffer", None) is not None: + k_size_bytes += get_tensor_size_bytes(self.dq_k_buffer) + v_size_bytes += get_tensor_size_bytes(self.dq_v_buffer) return k_size_bytes, v_size_bytes # for disagg @@ -1798,9 +1928,19 @@ class MHATokenToKVPool(KVCache): def _get_key_buffer(self, layer_id: int): # for internal use of referencing + local_layer_id = layer_id - self.start_layer + if ( + self.is_quantized_kv_cache + and self.quant_method.needs_plain_kv_dequant_read() + ): + return self.quant_method.dequantize_kv_tensor( + self.k_buffer[local_layer_id], + self.k_scale_buffer[local_layer_id], + layer_id, + ) if self.store_dtype != self.dtype: - return self.k_buffer[layer_id - self.start_layer].view(self.dtype) - return self.k_buffer[layer_id - self.start_layer] + return self.k_buffer[local_layer_id].view(self.dtype) + return self.k_buffer[local_layer_id] def get_key_buffer(self, layer_id: int): # note: get_key_buffer is hooked with synchronization for layer-wise KV cache loading @@ -1812,9 +1952,19 @@ class MHATokenToKVPool(KVCache): def _get_value_buffer(self, layer_id: int): # for internal use of referencing + local_layer_id = layer_id - self.start_layer + if ( + self.is_quantized_kv_cache + and self.quant_method.needs_plain_kv_dequant_read() + ): + return self.quant_method.dequantize_kv_tensor( + self.v_buffer[local_layer_id], + self.v_scale_buffer[local_layer_id], + layer_id, + ) if self.store_dtype != self.dtype: - return self.v_buffer[layer_id - self.start_layer].view(self.dtype) - return self.v_buffer[layer_id - self.start_layer] + return self.v_buffer[local_layer_id].view(self.dtype) + return self.v_buffer[local_layer_id] def get_value_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: @@ -1839,10 +1989,25 @@ class MHATokenToKVPool(KVCache): # Catch stale slot ids here instead of as illegal-addr / silent KV # corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT). maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)") - if layer_id_override is not None: - layer_id = layer_id_override - else: - layer_id = layer.layer_id + layer_id = ( + layer_id_override if layer_id_override is not None else layer.layer_id + ) + global_layer_id = layer.layer_id if layer is not None else layer_id + + if self.is_quantized_kv_cache: + if dcp_kv_mask is not None: + raise RuntimeError("dcp_kv_mask is not supported for FP4 KV cache.") + self._set_quantized_kv_buffer( + layer_id, + global_layer_id, + loc, + cache_k, + cache_v, + k_scale, + v_scale, + ) + return + if cache_k.dtype != self.dtype: if k_scale is not None: cache_k.div_(k_scale) @@ -1936,6 +2101,255 @@ class MHATokenToKVPool(KVCache): same_kv_dim=self.same_kv_dim, ) + def _quantized_scales(self, global_layer_id: int, k_scale, v_scale): + if k_scale is None and hasattr(self.quant_method, "k_scales_gpu"): + k_scale = self.quant_method.k_scales_gpu[ + global_layer_id : global_layer_id + 1 + ] + v_scale = self.quant_method.v_scales_gpu[ + global_layer_id : global_layer_id + 1 + ] + return k_scale, v_scale + + def _set_quantized_kv_buffer( + self, + layer_id: int, + global_layer_id: int, + loc_info, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale=None, + v_scale=None, + ) -> None: + loc, _, _ = unwrap_write_loc(loc_info) + local_layer_id = layer_id - self.start_layer + k_scale, v_scale = self._quantized_scales(global_layer_id, k_scale, v_scale) + self.quant_method.quantize_and_store( + self.k_buffer[local_layer_id], + self.v_buffer[local_layer_id], + ( + self.k_scale_buffer[local_layer_id] + if self.k_scale_buffer is not None + else None + ), + ( + self.v_scale_buffer[local_layer_id] + if self.v_scale_buffer is not None + else None + ), + loc, + cache_k, + cache_v, + k_scale, + v_scale, + ) + + def get_raw_kv_buffer( + self, layer_id: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + local_layer_id = layer_id - self.start_layer + if self.k_scale_buffer is None or self.v_scale_buffer is None: + raise RuntimeError("Raw FP4 KV cache requested from a non-FP4 KV pool.") + k_scale = self.k_scale_buffer[local_layer_id] + v_scale = self.v_scale_buffer[local_layer_id] + scale_view_dtype = self.quant_method.scale_buffer_view_dtype() + if scale_view_dtype is not None: + k_scale = k_scale.view(scale_view_dtype) + v_scale = v_scale.view(scale_view_dtype) + return ( + self.k_buffer[local_layer_id], + self.v_buffer[local_layer_id], + k_scale, + v_scale, + ) + + def get_dequant_workspace(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.dq_k_buffer is None or self.dq_v_buffer is None: + raise RuntimeError( + "Dequant workspace requested from a KV pool without FP4 dequant buffers." + ) + return self.dq_k_buffer, self.dq_v_buffer + + def get_flashinfer_dequant_workspace_kv_buffer( + self, + layer: RadixAttention, + req_to_token: torch.Tensor, + req_pool_indices_cpu, + extend_prefix_lens_cpu, + extend_seq_lens_cpu, + page_size: int, + *, + prepare_workspace: bool, + use_ragged: bool, + k_cur: Optional[torch.Tensor] = None, + v_cur: Optional[torch.Tensor] = None, + layer_id_override: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return the FlashInfer FP8 KV view for a quantized KV cache. + + FlashInfer prefill consumes FP8 KV. Quantized pools store packed FP4 plus + per-block scales, so the pool owns the dequant workspace and returns the + view shape expected by FlashInfer. + """ + if not self.is_quantized_kv_cache: + raise RuntimeError( + "FlashInfer quantized KV buffer requested from a non-quantized KV pool." + ) + + if prepare_workspace: + transfer_cur_kv = not use_ragged + k_cur_fp8 = ( + k_cur.to(torch.float8_e4m3fn) + if k_cur is not None and transfer_cur_kv + else None + ) + v_cur_fp8 = ( + v_cur.to(torch.float8_e4m3fn) + if v_cur is not None and transfer_cur_kv + else None + ) + self._prepare_dequant_extend_workspace( + layer.layer_id if layer_id_override is None else layer_id_override, + layer.layer_id, + req_to_token, + req_pool_indices_cpu, + extend_prefix_lens_cpu, + extend_seq_lens_cpu, + page_size, + k_cur_fp8=k_cur_fp8, + v_cur_fp8=v_cur_fp8, + ) + + k_buffer_dq, v_buffer_dq = self.get_dequant_workspace() + return ( + k_buffer_dq.view(-1, layer.tp_k_head_num, layer.head_dim), + v_buffer_dq.view(-1, layer.tp_v_head_num, layer.head_dim), + ) + + def get_flashinfer_decode_dequant_workspace_kv_buffer( + self, + layer: RadixAttention, + req_to_token: torch.Tensor, + req_pool_indices, + seq_lens, + *, + layer_id_override: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not self.is_quantized_kv_cache: + raise RuntimeError( + "FlashInfer dequant workspace requested from a non-quantized KV pool." + ) + + self._prepare_dequant_decode_workspace( + layer.layer_id if layer_id_override is None else layer_id_override, + layer.layer_id, + req_to_token, + req_pool_indices, + seq_lens, + ) + k_buffer_dq, v_buffer_dq = self.get_dequant_workspace() + return ( + k_buffer_dq.view(-1, layer.tp_k_head_num, layer.head_dim), + v_buffer_dq.view(-1, layer.tp_v_head_num, layer.head_dim), + ) + + @staticmethod + def _to_cpu_int_list(values) -> list[int]: + if isinstance(values, list): + return [int(value) for value in values] + if isinstance(values, torch.Tensor): + return [int(value) for value in values.cpu().tolist()] + return [int(value) for value in values] + + def _prepare_dequant_extend_workspace( + self, + layer_id: int, + global_layer_id: int, + req_to_token: torch.Tensor, + req_pool_indices_cpu, + extend_prefix_lens_cpu, + extend_seq_lens_cpu, + page_size: int, + k_cur_fp8: Optional[torch.Tensor] = None, + v_cur_fp8: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build the shared FP8 workspace used by FlashInfer extend attention. + + Cached prefix tokens are stored as packed FP4 plus per-block scales, so + paged prefill dequantizes those prefix tokens into the FP8 workspace. + The current extend chunk can already be FP8 and is copied into the same + workspace after the prefix region. + """ + k_fp4, v_fp4, k_scales, v_scales = self.get_raw_kv_buffer(layer_id) + dq_k, dq_v = self.get_dequant_workspace() + + cur_batch_start_loc_cpu = 0 + cur_token_idx_dq = page_size + + for i in range(len(req_pool_indices_cpu)): + req_idx = int(req_pool_indices_cpu[i]) + prev_len = int(extend_prefix_lens_cpu[i]) + extend_len = int(extend_seq_lens_cpu[i]) + + if prev_len > 0: + prev_indices = req_to_token[req_idx, :prev_len] + k_prev_fp8, v_prev_fp8 = self.quant_method.dequantize_prev_kv( + k_fp4[prev_indices], + k_scales[prev_indices], + v_fp4[prev_indices], + v_scales[prev_indices], + global_layer_id, + ) + dq_k[cur_token_idx_dq : cur_token_idx_dq + prev_len] = k_prev_fp8 + dq_v[cur_token_idx_dq : cur_token_idx_dq + prev_len] = v_prev_fp8 + + if k_cur_fp8 is not None: + cur_end = cur_batch_start_loc_cpu + extend_len + dst_start = cur_token_idx_dq + prev_len + dst_end = dst_start + extend_len + dq_k[dst_start:dst_end] = k_cur_fp8[cur_batch_start_loc_cpu:cur_end] + dq_v[dst_start:dst_end] = v_cur_fp8[cur_batch_start_loc_cpu:cur_end] + cur_batch_start_loc_cpu = cur_end + + workspace_len = prev_len + (extend_len if k_cur_fp8 is not None else 0) + cur_token_idx_dq = ( + (cur_token_idx_dq + workspace_len + page_size - 1) + // page_size + * page_size + ) + + return dq_k, dq_v + + def _prepare_dequant_decode_workspace( + self, + layer_id: int, + global_layer_id: int, + req_to_token: torch.Tensor, + req_pool_indices, + seq_lens, + ) -> tuple[torch.Tensor, torch.Tensor]: + k_fp4, v_fp4, k_scales, v_scales = self.get_raw_kv_buffer(layer_id) + dq_k, dq_v = self.get_dequant_workspace() + + req_pool_indices_cpu = self._to_cpu_int_list(req_pool_indices) + seq_lens_cpu = self._to_cpu_int_list(seq_lens) + + for req_idx, seq_len in zip(req_pool_indices_cpu, seq_lens_cpu): + if seq_len <= 0: + continue + kv_indices = req_to_token[req_idx, :seq_len] + k_prev_fp8, v_prev_fp8 = self.quant_method.dequantize_prev_kv( + k_fp4[kv_indices], + k_scales[kv_indices], + v_fp4[kv_indices], + v_scales[kv_indices], + global_layer_id, + ) + dq_k[kv_indices] = k_prev_fp8 + dq_v[kv_indices] = v_prev_fp8 + + return dq_k, dq_v + def set_kv_buffer_prefix_valid( self, layer: RadixAttention, @@ -2047,6 +2461,10 @@ class MHATokenToKVPool(KVCache): # per-layer buffers here ignore page_size in move_kv_cache_native. if self.use_native_move_kv_cache: move_kv_cache_native(self.k_buffer, self.v_buffer, tgt_loc, src_loc) + if getattr(self, "k_scale_buffer", None) is not None: + move_kv_cache_native( + self.k_scale_buffer, self.v_scale_buffer, tgt_loc, src_loc + ) return N = tgt_loc.numel() @@ -2266,10 +2684,10 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): cache_k_nope_fp4_sf = self.k_scale_buffer[layer_id - self.start_layer] from sglang.srt.layers.quantization.kvfp4_tensor import ( - BlockFP4KVQuantizeUtil, + FP4MXBlock16KVQuantizeUtil, ) - cache_k_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( + cache_k_nope_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize( cache_k_nope_fp4, cache_k_nope_fp4_sf ) return cache_k_nope_fp4_dequant @@ -2284,10 +2702,10 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): cache_v_nope_fp4_sf = self.v_scale_buffer[layer_id - self.start_layer] from sglang.srt.layers.quantization.kvfp4_tensor import ( - BlockFP4KVQuantizeUtil, + FP4MXBlock16KVQuantizeUtil, ) - cache_v_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( + cache_v_nope_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize( cache_v_nope_fp4, cache_v_nope_fp4_sf ) return cache_v_nope_fp4_dequant @@ -2318,11 +2736,15 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): cache_v.div_(v_scale) from sglang.srt.layers.quantization.kvfp4_tensor import ( - BlockFP4KVQuantizeUtil, + FP4MXBlock16KVQuantizeUtil, ) - cache_k, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k) - cache_v, cache_v_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v) + cache_k, cache_k_fp4_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize( + cache_k + ) + cache_v, cache_v_fp4_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize( + cache_v + ) if self.store_dtype != self.dtype: cache_k = cache_k.view(self.store_dtype) @@ -2526,6 +2948,7 @@ class HybridLinearKVPool(KVCache): qk_rope_head_dim: int = None, start_layer: Optional[int] = None, full_kv_pool_class: Optional[type] = None, + quant_method=None, # When provided (shared-KV-pool path), use this pool for the # full-attention layers instead of constructing one internally. full_kv_pool: Optional[KVCache] = None, @@ -2551,20 +2974,28 @@ class HybridLinearKVPool(KVCache): self.full_kv_pool = full_kv_pool elif not use_mla: TokenToKVPoolClass = MHATokenToKVPool + quant_method_kwarg = {"quant_method": quant_method} if current_platform.is_out_of_tree(): TokenToKVPoolClass = current_platform.get_mha_kv_pool_cls() + quant_method_kwarg = {} elif _is_npu: + assert not is_float4_e2m1fn_x2( + dtype + ), "FP4 is not supported on NPU yet." from sglang.srt.hardware_backend.npu.memory_pool_npu import ( NPUMHATokenToKVPool, ) TokenToKVPoolClass = NPUMHATokenToKVPool + quant_method_kwarg = {} elif full_kv_pool_class is not None: # Caller-selected MHA layout variant (e.g. the page-major # PageMajorMHATokenToKVPool). NPU / out-of-tree classes keep # priority since they don't understand alternate layouts. TokenToKVPoolClass = full_kv_pool_class + else: + TokenToKVPoolClass = MHATokenToKVPool post_capture_kwargs = ( {"post_capture_active": True} if post_capture_active else {} @@ -2579,6 +3010,7 @@ class HybridLinearKVPool(KVCache): device=device, enable_memory_saver=enable_memory_saver, enable_kv_cache_copy=enable_kv_cache_copy, + **quant_method_kwarg, **post_capture_kwargs, ) else: @@ -2665,14 +3097,18 @@ class HybridLinearKVPool(KVCache): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - def get_key_buffer(self, layer_id: int): + def get_key_buffer(self, layer_id: int, scale: Optional[float] = None): self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) + if scale is not None: + return self.full_kv_pool.get_key_buffer(layer_id, scale) return self.full_kv_pool.get_key_buffer(layer_id) - def get_value_buffer(self, layer_id: int): + def get_value_buffer(self, layer_id: int, scale: Optional[float] = None): self._wait_for_layer(layer_id) layer_id = self._transfer_full_attention_id(layer_id) + if scale is not None: + return self.full_kv_pool.get_value_buffer(layer_id, scale) return self.full_kv_pool.get_value_buffer(layer_id) def get_kv_buffer(self, layer_id: int): @@ -2680,6 +3116,30 @@ class HybridLinearKVPool(KVCache): layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_kv_buffer(layer_id) + def get_raw_kv_buffer( + self, layer_id: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + self._wait_for_layer(layer_id) + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_raw_kv_buffer(layer_id) + + def get_dequant_workspace(self) -> tuple[torch.Tensor, torch.Tensor]: + return self.full_kv_pool.get_dequant_workspace() + + def get_flashinfer_dequant_workspace_kv_buffer(self, layer, *args, **kwargs): + self._wait_for_layer(layer.layer_id) + local_layer_id = self._transfer_full_attention_id(layer.layer_id) + return self.full_kv_pool.get_flashinfer_dequant_workspace_kv_buffer( + layer, *args, layer_id_override=local_layer_id, **kwargs + ) + + def get_flashinfer_decode_dequant_workspace_kv_buffer(self, layer, *args, **kwargs): + self._wait_for_layer(layer.layer_id) + local_layer_id = self._transfer_full_attention_id(layer.layer_id) + return self.full_kv_pool.get_flashinfer_decode_dequant_workspace_kv_buffer( + layer, *args, layer_id_override=local_layer_id, **kwargs + ) + @contextmanager def _transfer_id_context(self, layer: RadixAttention): @contextmanager @@ -2712,7 +3172,7 @@ class HybridLinearKVPool(KVCache): if not self.use_mla: write_loc = full_loc if full_loc is not None else loc self.full_kv_pool.set_kv_buffer( - None, + layer, write_loc, cache_k, cache_v, @@ -3095,10 +3555,10 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): cache_k_nope_fp4_sf = self.kv_scale_buffer[layer_id - self.start_layer] from sglang.srt.layers.quantization.kvfp4_tensor import ( - BlockFP4KVQuantizeUtil, + FP4MXBlock16KVQuantizeUtil, ) - cache_k_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( + cache_k_nope_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize( cache_k_nope_fp4, cache_k_nope_fp4_sf ) return cache_k_nope_fp4_dequant @@ -3119,10 +3579,10 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): assert not self.dsa_kv_cache_store_fp8 if cache_k.dtype != self.dtype: from sglang.srt.layers.quantization.kvfp4_tensor import ( - BlockFP4KVQuantizeUtil, + FP4MXBlock16KVQuantizeUtil, ) - cache_k_fp4, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize( + cache_k_fp4, cache_k_fp4_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize( cache_k ) @@ -3158,14 +3618,14 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): else: if cache_k_nope.dtype != self.dtype: from sglang.srt.layers.quantization.kvfp4_tensor import ( - BlockFP4KVQuantizeUtil, + FP4MXBlock16KVQuantizeUtil, ) cache_k_nope_fp4, cache_k_nope_fp4_sf = ( - BlockFP4KVQuantizeUtil.batched_quantize(cache_k_nope) + FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_nope) ) cache_k_rope_fp4, cache_k_rope_fp4_sf = ( - BlockFP4KVQuantizeUtil.batched_quantize(cache_k_rope) + FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_rope) ) if self.store_dtype != self.dtype: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c622a33de..568a83a92 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -516,6 +516,7 @@ class ModelRunner: gpu_id=self.gpu_id, ps=self.ps, pp_group=self.pp_group, + model=self.model, model_config=self.model_config, server_args=self.server_args, kv_cache_dtype=self.kv_cache_dtype, @@ -1073,14 +1074,19 @@ class ModelRunner: ) def configure_kv_cache_dtype(self): + spec_algorithm = getattr(self, "spec_algorithm", None) resolved_kv_cache_dtype, self.kv_cache_dtype = ( kv_cache_dtype.configure_kv_cache_dtype( server_args_kv_cache_dtype=self.server_args.kv_cache_dtype, - model=self.model, - model_dtype=self.dtype, - is_draft_worker=self.is_draft_worker, - is_dflash=self.spec_algorithm.is_dflash(), - speculative_draft_attention_backend=self.server_args.speculative_draft_attention_backend, + model=getattr(self, "model", None), + model_dtype=getattr(self, "dtype", torch.bfloat16), + is_draft_worker=getattr(self, "is_draft_worker", False), + is_dflash=( + spec_algorithm.is_dflash() if spec_algorithm is not None else False + ), + speculative_draft_attention_backend=getattr( + self.server_args, "speculative_draft_attention_backend", None + ), ) ) if resolved_kv_cache_dtype is not None: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index e728a531f..341326718 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -273,6 +273,8 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): cell_size = (cell_size // 2) + ( (n * k * effective_num_layers * 2 * kv_size) // scale_block_size ) + # FP4 prefill uses one shared FP8 dequant workspace across layers. + cell_size += n * k * 2 * kv_size return cell_size diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d84093c4f..599baa3eb 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -605,10 +605,21 @@ class ServerArgs: help=( 'Data type for kv cache storage. "auto" will use model data type. ' '"bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and ' - '"fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only ' - "mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+" + '"fp8_e4m3" are supported for CUDA 11.8+. "nvfp4" selects ' + 'the NVFP4 FP4 E2M1 KV cache recipe; "fp4_mx_block16" ' + "selects the MX-style block-size-16 FP4 E2M1 KV cache " + "recipe. Both require CUDA 12.8+ and PyTorch 2.8.0+" ), - choices=["auto", "fp8_e5m2", "fp8_e4m3", "bf16", "bfloat16", "fp4_e2m1"], + choices=[ + "auto", + "fp8_e5m2", + "fp8_e4m3", + "bf16", + "bfloat16", + "nvfp4", + "fp4_mx_block16", + "fp4_e2m1", + ], resolvable=True, ), ] = "auto" @@ -4964,7 +4975,7 @@ class ServerArgs: """Check FP4 KV cache compatibility with the attention backend""" from sglang.srt.arg_groups.overrides import resolved_view - if self.kv_cache_dtype != "fp4_e2m1": + if self.kv_cache_dtype not in ("nvfp4", "fp4_mx_block16"): return use_mla_backend = self.use_mla_backend() @@ -4972,6 +4983,13 @@ class ServerArgs: attention_backend = resolved_view(self).attention_backend if is_cuda(): + if self.kv_cache_dtype == "nvfp4" and not ( + is_sm100_supported() or is_sm120_supported() + ): + raise RuntimeError( + "--kv-cache-dtype=nvfp4 requires Blackwell SM100 or SM120. " + "Use --kv-cache-dtype=fp4_mx_block16 for the block-size-16 FP4 recipe." + ) if ( prefill_backend != decode_backend and prefill_backend != "fa4" ): # Take care of prefill=fa4 later @@ -5008,7 +5026,6 @@ class ServerArgs: "cutlass_mla", "flashinfer", "trtllm_mla", - "flashmla", ] assert attention_backend in KV4_ATTENTION_MLA_BACKEND_CHOICES, ( f"KV4 MLA expects attention_backend to be one of " @@ -5895,11 +5912,11 @@ class ServerArgs: "Other prefill-only workloads may be supported in a future change once " "their attention paths stop reading or writing the paged KV cache." ) - if self.kv_cache_dtype == "fp4_e2m1": + if self.kv_cache_dtype in ("nvfp4", "fp4_mx_block16"): raise ValueError( "--prefill-only-disable-kv-cache does not currently support " - "--kv-cache-dtype=fp4_e2m1 because the FP4 pool uses a separate " - "allocation path." + "--kv-cache-dtype=nvfp4 or --kv-cache-dtype=fp4_mx_block16 because " + "the FP4 pool uses a separate allocation path." ) # Structural preconditions for the FA backend's fa_skip_kv_cache path, diff --git a/test/registered/models_e2e/test_llama8b_nvfp4_kv_cache_sm120.py b/test/registered/models_e2e/test_llama8b_nvfp4_kv_cache_sm120.py new file mode 100644 index 000000000..d516ef121 --- /dev/null +++ b/test/registered/models_e2e/test_llama8b_nvfp4_kv_cache_sm120.py @@ -0,0 +1,64 @@ +import unittest + +from sglang.srt.utils.common import is_sm120_supported +from sglang.test.accuracy_test_runner import AccuracyTestParams +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_combined_tests import run_combined_tests +from sglang.test.test_utils import CustomTestCase, ModelLaunchSettings + +register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-small") + +LLAMA8B_NVFP4_MODEL = "nvidia/Llama-3.1-8B-Instruct-NVFP4" +TP_SIZE = 1 + + +@unittest.skipUnless( + is_sm120_supported(), "requires at least 1 SM120 GPU with CUDA 12.8+" +) +class TestLlama8BNVFP4KVCacheSM120(CustomTestCase): + """Llama-3.1-8B-Instruct-NVFP4 with NVFP4 KV cache on SM120.""" + + def test_gsm8k(self): + variants = [ + ModelLaunchSettings( + LLAMA8B_NVFP4_MODEL, + tp_size=TP_SIZE, + extra_args=[ + "--quantization", + "modelopt_fp4", + "--fp4-gemm-backend", + "auto", + "--kv-cache-dtype", + "nvfp4", + "--prefill-attention-backend", + "flashinfer", + "--decode-attention-backend", + "trtllm_mha", + "--page-size", + "64", + "--cuda-graph-backend-prefill=disabled", + ], + variant="NVFP4-GEMM+NVFP4-KV+SM120-XQA", + ) + ] + + run_combined_tests( + models=variants, + test_name="Llama-3.1-8B-Instruct-NVFP4-KV-SM120", + accuracy_params=AccuracyTestParams( + dataset="gsm8k", + # Full GSM8K measured locally with 1319 requested / 1314 scored: + # - FP8 KV: 0.6461187214611872 + # - NVFP4 KV: 0.632420091324201 + # Keep the threshold 0.015 below the NVFP4 KV score. + baseline_accuracy=0.632420091324201 - 0.015, + num_examples=1319, + num_threads=200, + max_tokens=512, + api="completion", + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/quant/test_kvfp4_quant_dequant.py b/test/registered/quant/test_kvfp4_quant_dequant.py similarity index 78% rename from test/manual/quant/test_kvfp4_quant_dequant.py rename to test/registered/quant/test_kvfp4_quant_dequant.py index 3f7d5b6a8..69371ac82 100755 --- a/test/manual/quant/test_kvfp4_quant_dequant.py +++ b/test/registered/quant/test_kvfp4_quant_dequant.py @@ -1,12 +1,16 @@ #!/usr/bin/env python3 +import sys import time import numpy as np import pytest import torch -from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil +from sglang.srt.layers.quantization.kvfp4_tensor import FP4MXBlock16KVQuantizeUtil +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large") def calculate_accuracy_metrics( @@ -28,7 +32,7 @@ def calculate_accuracy_metrics( return {"MSE": mse, "MAE": mae, "PSNR": psnr, "Relative Error": rel_error} -def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]: +def run_benchmark(m, n, k, num_runs=10) -> dict[str, dict[str, float]]: """Run FP8 vs KVFP4 quantization benchmark and return metrics.""" tensor_bf16 = torch.randn(m, n, k, dtype=torch.bfloat16, device="cuda") @@ -52,18 +56,20 @@ def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]: fp8_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp8_dequant) # --- KVFP4 --- - tensor_fp4, scale_factors = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16) - _ = BlockFP4KVQuantizeUtil.batched_dequantize(tensor_fp4, scale_factors) + tensor_fp4, scale_factors = FP4MXBlock16KVQuantizeUtil.batched_quantize(tensor_bf16) + _ = FP4MXBlock16KVQuantizeUtil.batched_dequantize(tensor_fp4, scale_factors) start = time.time() for _ in range(num_runs): - tensor_fp4, scale_factors = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16) + tensor_fp4, scale_factors = FP4MXBlock16KVQuantizeUtil.batched_quantize( + tensor_bf16 + ) torch.cuda.synchronize() fp4_quant_time = (time.time() - start) / num_runs start = time.time() for _ in range(num_runs): - tensor_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( + tensor_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize( tensor_fp4, scale_factors ) torch.cuda.synchronize() @@ -91,14 +97,8 @@ def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]: MNK_FACTORS = [ (64, 1, 576), (512, 1, 576), - (1024, 1, 576), - (4096, 1, 576), - (2868672, 1, 576), (64, 8, 64), (512, 8, 64), - (1024, 8, 64), - (4096, 8, 64), - (2868672, 8, 64), ] @@ -112,5 +112,9 @@ def test_kvfp4_quant_dequant(m, n, k): print("FP4:", results["fp4"]) # Basic assertions to make sure metrics are reasonable - assert results["fp4"]["MSE"] < 1.0 - assert results["fp8"]["MSE"] < 1.0 + assert results["fp4"]["MSE"] < 0.1 + assert results["fp8"]["MSE"] < 0.1 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py index 9726d79f2..316456a5e 100644 --- a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py +++ b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py @@ -1,4 +1,4 @@ -"""Unit tests for FP4 KV cache quantization strategy pattern — no server, no model loading.""" +"""Unit tests for FP4 KV cache quantization strategy pattern - no server, no model loading.""" from sglang.test.ci.ci_register import register_cpu_ci @@ -23,54 +23,87 @@ def skip_if_no_blackwell_nvfp4(func): class TestKVCacheQuantRegistry(CustomTestCase): """Test the registry and factory function.""" - def test_registry_contains_nvfp4_and_mxfp4(self): + def test_registry_contains_nvfp4_and_blockfp4(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - FP4_KV_CACHE_QUANT_REGISTRY, + KV_CACHE_QUANT_REGISTRY, ) - self.assertIn("nvfp4", FP4_KV_CACHE_QUANT_REGISTRY) - self.assertIn("blockfp4", FP4_KV_CACHE_QUANT_REGISTRY) + self.assertIn("nvfp4", KV_CACHE_QUANT_REGISTRY) + self.assertIn("fp4_mx_block16", KV_CACHE_QUANT_REGISTRY) def test_factory_nvfp4(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - NVFP4KVMethod, - get_fp4_kv_cache_quant_method, + NVFP4KVCacheMethod, + get_kv_cache_quant_method, ) - method = get_fp4_kv_cache_quant_method( - "nvfp4", num_layers=4, device="cpu", sm_version=120 - ) - self.assertIsInstance(method, NVFP4KVMethod) + method = get_kv_cache_quant_method("nvfp4", num_layers=4, device="cpu") + self.assertIsInstance(method, NVFP4KVCacheMethod) self.assertEqual(method.name, "nvfp4") - def test_factory_mxfp4(self): + def test_factory_blockfp4(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - BlockFP4KVMethod, - get_fp4_kv_cache_quant_method, + FP4MXBlock16KVCacheMethod, + get_kv_cache_quant_method, ) - method = get_fp4_kv_cache_quant_method("blockfp4") - self.assertIsInstance(method, BlockFP4KVMethod) - self.assertEqual(method.name, "blockfp4") + method = get_kv_cache_quant_method("fp4_mx_block16") + self.assertIsInstance(method, FP4MXBlock16KVCacheMethod) + self.assertEqual(method.name, "fp4_mx_block16") - def test_factory_unknown_raises(self): + def test_resolve_explicit_recipes(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - get_fp4_kv_cache_quant_method, + resolve_kv_cache_quant, + ) + + self.assertEqual(resolve_kv_cache_quant("nvfp4"), "nvfp4") + self.assertEqual(resolve_kv_cache_quant("fp4_mx_block16"), "fp4_mx_block16") + self.assertIsNone(resolve_kv_cache_quant("fp8_e4m3")) + + def test_resolve_legacy_fp4_alias_raises(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + resolve_kv_cache_quant, + ) + + with self.assertRaisesRegex(ValueError, "fp4_mx_block16"): + resolve_kv_cache_quant("fp4_e2m1") + + def test_model_runner_rejects_legacy_fp4_alias(self): + from types import SimpleNamespace + + from sglang.srt.model_executor.model_runner import ModelRunner + + runner = object.__new__(ModelRunner) + runner.server_args = SimpleNamespace(kv_cache_dtype="fp4_e2m1") + with self.assertRaisesRegex(ValueError, "fp4_mx_block16"): + runner.configure_kv_cache_dtype() + + def test_resolve_mxfp4_name_raises(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + resolve_kv_cache_quant, ) with self.assertRaises(ValueError): - get_fp4_kv_cache_quant_method("unknown_method") + resolve_kv_cache_quant("mxfp4") + + def test_factory_unknown_raises(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + get_kv_cache_quant_method, + ) + + with self.assertRaises(ValueError): + get_kv_cache_quant_method("unknown_method") -class TestNVFP4KVMethod(CustomTestCase): - """Test NVFP4KVMethod buffer creation and properties.""" +class TestNVFP4KVCacheMethod(CustomTestCase): + """Test NVFP4KVCacheMethod buffer creation and properties.""" def test_properties(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - NVFP4KVMethod, + NVFP4KVCacheMethod, ) - m = NVFP4KVMethod(num_layers=4, device="cpu", sm_version=120) + m = NVFP4KVCacheMethod(num_layers=4, device="cpu") self.assertEqual(m.name, "nvfp4") self.assertEqual(m.SCALE_BLOCK_SIZE, 16) self.assertTrue(m.needs_dequant_workspace()) @@ -78,10 +111,10 @@ class TestNVFP4KVMethod(CustomTestCase): def test_create_buffers_shapes(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - NVFP4KVMethod, + NVFP4KVCacheMethod, ) - m = NVFP4KVMethod(num_layers=4, device="cpu", sm_version=120) + m = NVFP4KVCacheMethod(num_layers=4, device="cpu") size, heads, dim, layers = 64, 8, 128, 4 bufs = m.create_buffers(size, heads, dim, layers, "cpu") @@ -101,20 +134,20 @@ class TestNVFP4KVMethod(CustomTestCase): def test_compute_cell_size(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - NVFP4KVMethod, + NVFP4KVCacheMethod, ) - m = NVFP4KVMethod(num_layers=4, device="cpu") + m = NVFP4KVCacheMethod(num_layers=4, device="cpu") cell = m.compute_cell_size(head_num=8, head_dim=128, num_layers=4, kv_size=1) # FP4: 8*64*4*2 = 4096, scales: 8*8*4*2 = 512, dq: 8*128*2 = 2048 self.assertEqual(cell, 4096 + 512 + 2048) def test_scales_init(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - NVFP4KVMethod, + NVFP4KVCacheMethod, ) - m = NVFP4KVMethod(num_layers=4, device="cpu") + m = NVFP4KVCacheMethod(num_layers=4, device="cpu") # Default scales should be 1.0 self.assertTrue(torch.all(m.k_scales_gpu == 1.0)) self.assertTrue(torch.all(m.v_scales_gpu == 1.0)) @@ -122,13 +155,13 @@ class TestNVFP4KVMethod(CustomTestCase): @skip_if_no_blackwell_nvfp4 def test_quantize_dequantize_roundtrip(self): - """Test NVFP4 quantize→dequantize roundtrip on CUDA.""" + """Test NVFP4 quantize->dequantize roundtrip on CUDA.""" from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - NVFP4KVMethod, + NVFP4KVCacheMethod, ) major, minor = torch.cuda.get_device_capability() - m = NVFP4KVMethod(num_layers=1, device="cuda", sm_version=major * 10 + minor) + m = NVFP4KVCacheMethod(num_layers=1, device="cuda") size, heads, dim = 32, 8, 128 bufs = m.create_buffers(size, heads, dim, 1, "cuda") @@ -171,40 +204,55 @@ class TestNVFP4KVMethod(CustomTestCase): ) -class TestBlockFP4KVMethod(CustomTestCase): - """Test BlockFP4KVMethod buffer creation and roundtrip.""" +class TestFP4MXBlock16KVCacheMethod(CustomTestCase): + """Test FP4MXBlock16KVCacheMethod buffer creation and roundtrip.""" def test_properties(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - BlockFP4KVMethod, + FP4MXBlock16KVCacheMethod, + KVCacheAttentionAccessKind, ) - m = BlockFP4KVMethod() - self.assertEqual(m.name, "blockfp4") - self.assertTrue(m.needs_dequant_workspace()) + m = FP4MXBlock16KVCacheMethod() + self.assertEqual(m.name, "fp4_mx_block16") + self.assertFalse(m.needs_dequant_workspace()) + self.assertTrue(m.needs_plain_kv_dequant_read()) self.assertFalse(m.needs_global_scale()) + self.assertEqual(m.plain_attention_kv_dtype(), torch.bfloat16) + self.assertEqual( + m.resolve_attention_access("prefill", "triton").kind, + KVCacheAttentionAccessKind.PLAIN, + ) + self.assertEqual( + m.resolve_attention_access("decode", "trtllm_mha").kind, + KVCacheAttentionAccessKind.PLAIN, + ) + self.assertIsNone(m.resolve_attention_access("prefill", "flashinfer")) + self.assertIsNone(m.resolve_attention_access("decode", "flashinfer")) def test_create_buffers_shapes(self): from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - BlockFP4KVMethod, + FP4MXBlock16KVCacheMethod, ) - m = BlockFP4KVMethod() + m = FP4MXBlock16KVCacheMethod() size, heads, dim, layers = 64, 8, 128, 4 bufs = m.create_buffers(size, heads, dim, layers, "cpu") self.assertEqual(len(bufs["k_buffer"]), layers) self.assertEqual(bufs["k_buffer"][0].shape, (size, heads, dim // 2)) - # MXFP4 flattens head dims for scales + # Block-16 FP4 flattens head dims for scales self.assertEqual(bufs["k_scale_buffer"][0].shape, (size, (heads * dim) // 16)) + self.assertIsNone(bufs["dq_k_buffer"]) + self.assertIsNone(bufs["dq_v_buffer"]) def test_quantize_dequantize_roundtrip_cpu(self): - """Test MXFP4 quantize→dequantize roundtrip on CPU.""" + """Test block-16 FP4 quantize->dequantize roundtrip on CPU.""" from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( - BlockFP4KVMethod, + FP4MXBlock16KVCacheMethod, ) - m = BlockFP4KVMethod() + m = FP4MXBlock16KVCacheMethod() size, heads, dim = 32, 8, 128 bufs = m.create_buffers(size, heads, dim, 1, "cpu") @@ -231,18 +279,22 @@ class TestBlockFP4KVMethod(CustomTestCase): k_out, v_out = m.dequantize_prev_kv(k_fp4, k_scales, v_fp4, v_scales, 0) self.assertEqual(k_out.shape, (4, heads, dim)) - self.assertEqual(k_out.dtype, torch.float8_e4m3fn) + self.assertEqual(v_out.shape, (4, heads, dim)) + self.assertEqual(k_out.dtype, torch.bfloat16) + self.assertEqual(v_out.dtype, torch.bfloat16) -class TestBlockFP4KVQuantizeUtil(CustomTestCase): - """Test the existing MXFP4 BlockFP4KVQuantizeUtil roundtrip.""" +class TestFP4MXBlock16KVQuantizeUtil(CustomTestCase): + """Test the existing block-16 FP4 FP4MXBlock16KVQuantizeUtil roundtrip.""" def test_roundtrip_cpu(self): - from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + FP4MXBlock16KVQuantizeUtil, + ) x = torch.randn(4, 8, 128, dtype=torch.bfloat16) - packed, scales = BlockFP4KVQuantizeUtil.batched_quantize(x) - reconstructed = BlockFP4KVQuantizeUtil.batched_dequantize(packed, scales) + packed, scales = FP4MXBlock16KVQuantizeUtil.batched_quantize(x) + reconstructed = FP4MXBlock16KVQuantizeUtil.batched_dequantize(packed, scales) self.assertEqual(reconstructed.shape, x.shape) rel_error = ( @@ -251,15 +303,5 @@ class TestBlockFP4KVQuantizeUtil(CustomTestCase): self.assertLess(rel_error, 0.5) -class TestFP4KVCacheRecipe(CustomTestCase): - """Test enum.""" - - def test_enum_values(self): - from sglang.srt.layers.quantization.kvfp4_tensor import FP4KVCacheRecipe - - self.assertEqual(FP4KVCacheRecipe.MXFP4.value, 1) - self.assertEqual(FP4KVCacheRecipe.NVFP4.value, 2) - - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_quantized_kv_pool.py b/test/registered/unit/mem_cache/test_quantized_kv_pool.py new file mode 100644 index 000000000..efb71885a --- /dev/null +++ b/test/registered/unit/mem_cache/test_quantized_kv_pool.py @@ -0,0 +1,124 @@ +# Copyright 2026 SGLang Team +# Licensed under the Apache License, Version 2.0 + +import types +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class _FakeQuantMethod: + name = "fake_quant" + + def __init__(self): + self.k_scales_gpu = torch.tensor([2.0], dtype=torch.float32) + self.v_scales_gpu = torch.tensor([3.0], dtype=torch.float32) + self.store_calls = [] + + def dequant_workspace_dtype(self): + return torch.float32 + + def create_buffers(self, size, head_num, head_dim, layer_num, device): + return { + "k_buffer": [ + torch.zeros( + (size, head_num, head_dim), dtype=torch.uint8, device=device + ) + for _ in range(layer_num) + ], + "v_buffer": [ + torch.zeros( + (size, head_num, head_dim), dtype=torch.uint8, device=device + ) + for _ in range(layer_num) + ], + "k_scale_buffer": [ + torch.zeros((size, head_num, 1), dtype=torch.uint8, device=device) + for _ in range(layer_num) + ], + "v_scale_buffer": [ + torch.zeros((size, head_num, 1), dtype=torch.uint8, device=device) + for _ in range(layer_num) + ], + "dq_k_buffer": torch.zeros( + (size, head_num, head_dim), dtype=torch.float32, device=device + ), + "dq_v_buffer": torch.zeros( + (size, head_num, head_dim), dtype=torch.float32, device=device + ), + "store_dtype": torch.uint8, + } + + def quantize_and_store( + self, + k_buffer, + v_buffer, + k_scale_buffer, + v_scale_buffer, + loc, + cache_k, + cache_v, + k_scale=None, + v_scale=None, + ): + self.store_calls.append( + { + "loc": loc, + "k_scale": k_scale, + "v_scale": v_scale, + "k_scale_buffer": k_scale_buffer, + "v_scale_buffer": v_scale_buffer, + } + ) + k_buffer[loc] = 1 + v_buffer[loc] = 2 + k_scale_buffer[loc] = 3 + v_scale_buffer[loc] = 4 + + +class TestQuantizedKVPool(unittest.TestCase): + def test_quant_method_owns_buffers_and_store_path(self): + from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool + + quant_method = _FakeQuantMethod() + pool = MHATokenToKVPool( + size=4, + page_size=1, + dtype=torch.bfloat16, + head_num=1, + head_dim=8, + layer_num=1, + device="cpu", + enable_memory_saver=False, + quant_method=quant_method, + ) + + self.assertTrue(pool.is_quantized_kv_cache) + self.assertIs(pool.quant_method, quant_method) + self.assertIsNotNone(pool.k_scale_buffer) + self.assertIs(pool.get_dequant_workspace()[0], pool.dq_k_buffer) + + loc = torch.tensor([0, 1], dtype=torch.int64) + layer = types.SimpleNamespace(layer_id=0) + pool.set_kv_buffer( + layer, + loc, + torch.zeros((2, 1, 8), dtype=torch.bfloat16), + torch.zeros((2, 1, 8), dtype=torch.bfloat16), + ) + + self.assertEqual(len(quant_method.store_calls), 1) + call = quant_method.store_calls[0] + self.assertIs(call["loc"], loc) + self.assertTrue(torch.equal(call["k_scale"], quant_method.k_scales_gpu[0:1])) + self.assertTrue(torch.equal(call["v_scale"], quant_method.v_scales_gpu[0:1])) + self.assertEqual(pool.k_buffer[0][loc].unique().tolist(), [1]) + self.assertEqual(pool.v_buffer[0][loc].unique().tolist(), [2]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 3367a2147..3887cbb7c 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1136,7 +1136,7 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase): - disable_radix_cache (radix cache otherwise indexes empty pool slots), - no context-parallel attention (CP writes to the pool via set_kv_buffer), - no HiSparse (uses a different pool family), - - kv_cache_dtype != fp4_e2m1 (FP4 pool is a separate allocation path). + - kv_cache_dtype is not nvfp4/fp4_mx_block16 (FP4 pool is a separate allocation path). All other configurations must be rejected before model load. """ @@ -1186,8 +1186,10 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase): self._validate_prefill_only_args(enable_hisparse=True) def test_rejects_fp4_kv_cache(self): - with self.assertRaisesRegex(ValueError, "fp4_e2m1"): - self._validate_prefill_only_args(kv_cache_dtype="fp4_e2m1") + for kv_cache_dtype in ("nvfp4", "fp4_mx_block16"): + with self.subTest(kv_cache_dtype=kv_cache_dtype): + with self.assertRaisesRegex(ValueError, "nvfp4.*fp4_mx_block16"): + self._validate_prefill_only_args(kv_cache_dtype=kv_cache_dtype) class TestSessionRadixCacheServerArgs(unittest.TestCase):