| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import Any, ClassVar, cast |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config |
| 11 | from vllm.forward_context import get_forward_context |
| 12 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 13 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 14 | from vllm.model_executor.layers.linear import MergedColumnParallelLinear |
| 15 | from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( |
| 16 | compress_norm_rope_store_triton, |
| 17 | compress_norm_rope_store_two_stage_triton, |
| 18 | ) |
| 19 | from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE |
| 20 | from vllm.models.deepseek_v4.common.ops.save_partial_states import ( |
| 21 | _SAVE_PARTIAL_STATES_KERNEL, |
| 22 | ) |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.v1.attention.backend import ( |
| 25 | AttentionBackend, |
| 26 | AttentionCGSupport, |
| 27 | AttentionMetadataBuilder, |
| 28 | CommonAttentionMetadata, |
| 29 | MultipleOf, |
| 30 | ) |
| 31 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 32 | from vllm.v1.kv_cache_interface import ( |
| 33 | KVCacheSpec, |
| 34 | MLAAttentionSpec, |
| 35 | SlidingWindowMLASpec, |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | def _prefer_two_stage_compressor() -> bool: |
| 40 | # Platforms that favor the triton variant of two-stage compressor split. |
| 41 | # Currently only tested on ROCm |
| 42 | return current_platform.is_rocm() |
| 43 | |
| 44 | |
| 45 | def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None: |
| 46 | seq_lens_cpu = metadata.seq_lens_cpu_upper_bound |
| 47 | if seq_lens_cpu is None: |
| 48 | return None |
| 49 | |
| 50 | query_lens = metadata.query_start_loc_cpu[1:] - metadata.query_start_loc_cpu[:-1] |
| 51 | starts = seq_lens_cpu - query_lens |
| 52 | starts_list = starts.tolist() |
| 53 | query_start_loc = metadata.query_start_loc_cpu.tolist() |
| 54 | return any( |
| 55 | start % 128 + query_start_loc[i + 1] - query_start_loc[i] >= 128 |
| 56 | for i, start in enumerate(starts_list) |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | class CompressorBackend(AttentionBackend): |
| 61 | def __init__(self): |
| 62 | super().__init__() |
| 63 | |
| 64 | @staticmethod |
| 65 | def get_name() -> str: |
| 66 | return "CompressorBackend" |
| 67 | |
| 68 | @staticmethod |
| 69 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 70 | return [MultipleOf(1)] |
| 71 | |
| 72 | @classmethod |
| 73 | def get_supported_head_sizes(cls) -> list[int]: |
| 74 | return [512, 1024] |
| 75 | |
| 76 | @staticmethod |
| 77 | def get_builder_cls() -> type["CompressorMetadataBuilder"]: |
| 78 | return CompressorMetadataBuilder |
| 79 | |
| 80 | |
| 81 | @dataclass |
| 82 | class CompressorMetadata: |
| 83 | block_table: torch.Tensor |
| 84 | slot_mapping: torch.Tensor |
| 85 | block_size: int |
| 86 | |
| 87 | token_to_req_indices: torch.Tensor | None = None # [num_tokens] |
| 88 | num_decode_tokens: int | None = None |
| 89 | c128_boundary: bool | None = None |
| 90 | |
| 91 | |
| 92 | class CompressorMetadataBuilder(AttentionMetadataBuilder): |
| 93 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 94 | |
| 95 | def __init__(self, *args, **kwargs): |
| 96 | super().__init__(*args, **kwargs) |
| 97 | assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) |
| 98 | mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) |
| 99 | self.block_size = mla_spec.block_size |
| 100 | |
| 101 | self.token_to_req_indices = torch.zeros( |
| 102 | self.vllm_config.scheduler_config.max_num_batched_tokens, |
| 103 | dtype=torch.int32, |
| 104 | device=self.device, |
| 105 | ) |
| 106 | |
| 107 | def build( |
| 108 | self, |
| 109 | common_prefix_len: int, |
| 110 | common_attn_metadata: CommonAttentionMetadata, |
| 111 | fast_build: bool = False, |
| 112 | ) -> CompressorMetadata: |
| 113 | token_to_req_indices = common_attn_metadata.token_to_req_indices( |
| 114 | self.token_to_req_indices |
| 115 | ) |
| 116 | num_decode_tokens = None |
| 117 | if _prefer_two_stage_compressor(): |
| 118 | _, _, num_decode_tokens, _ = split_decodes_and_prefills( |
| 119 | common_attn_metadata, decode_threshold=1 |
| 120 | ) |
| 121 | async_spec_decode = ( |
| 122 | self.vllm_config.scheduler_config.async_scheduling |
| 123 | and self.vllm_config.speculative_config is not None |
| 124 | ) |
| 125 | return CompressorMetadata( |
| 126 | block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), |
| 127 | slot_mapping=common_attn_metadata.slot_mapping, |
| 128 | block_size=self.block_size, |
| 129 | token_to_req_indices=token_to_req_indices, |
| 130 | num_decode_tokens=num_decode_tokens, |
| 131 | c128_boundary=( |
| 132 | _get_c128_boundary(common_attn_metadata) |
| 133 | if self.block_size == 8 and not async_spec_decode |
| 134 | else None |
| 135 | ), |
| 136 | ) |
| 137 | |
| 138 | |
| 139 | class CompressorStateCache(torch.nn.Module, AttentionLayerBase): |
| 140 | def __init__( |
| 141 | self, |
| 142 | state_dim: int, |
| 143 | dtype: torch.dtype, |
| 144 | compress_ratio: int, |
| 145 | prefix: str, |
| 146 | ): |
| 147 | super().__init__() |
| 148 | self.state_dim = state_dim |
| 149 | self.dtype = dtype |
| 150 | self.prefix = prefix |
| 151 | self.kv_cache = torch.tensor([]) |
| 152 | compilation_config = get_current_vllm_config().compilation_config |
| 153 | if prefix in compilation_config.static_forward_context: |
| 154 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 155 | compilation_config.static_forward_context[prefix] = self |
| 156 | |
| 157 | assert self.dtype == torch.float32 |
| 158 | assert compress_ratio in [4, 128] |
| 159 | coff = 1 + (compress_ratio == 4) |
| 160 | self.sliding_window = coff * compress_ratio |
| 161 | # Block size is constrained by tensor sharing between compressor states |
| 162 | # and KV blocks. Since compressor states share the same physical tensor |
| 163 | # as KV blocks, they must use the same page size. |
| 164 | # The KV block shape [256//4, head_dim] = [64, 584] determines: |
| 165 | # - C4 compressor block shape [4, 2*512*2*4] -> block_size = 4 |
| 166 | # - C128 compressor block shape [8, 512*2*4] -> block_size = 8 |
| 167 | # TODO(yifan): make block size automatically determined and configurable. |
| 168 | if compress_ratio == 4: |
| 169 | self.block_size = 4 |
| 170 | elif compress_ratio == 128: |
| 171 | self.block_size = 8 |
| 172 | else: |
| 173 | raise ValueError(f"Invalid compress ratio: {compress_ratio}") |
| 174 | |
| 175 | def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: |
| 176 | # [B, H=1, N, C] -> [B, N, C] |
| 177 | self.kv_cache = kv_cache.squeeze(1) |
| 178 | |
| 179 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 180 | # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain |
| 181 | # full-cache rows share state pages with contiguous KV pages, so padding |
| 182 | # would break page matching. |
| 183 | uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" |
| 184 | return SlidingWindowMLASpec( # only has one vector instead of K + V |
| 185 | block_size=self.block_size, |
| 186 | num_kv_heads=1, |
| 187 | max_tp_shards=1, |
| 188 | head_size=self.state_dim, |
| 189 | dtype=self.dtype, |
| 190 | sliding_window=self.sliding_window, |
| 191 | alignment=576 if uses_fp8_ds_mla_layout else 512, |
| 192 | ) |
| 193 | |
| 194 | def forward(self): ... |
| 195 | |
| 196 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 197 | return CompressorBackend |
| 198 | |
| 199 | |
| 200 | class DeepseekCompressor(nn.Module): |
| 201 | """DeepSeek V4 KV/score compressor. |
| 202 | |
| 203 | Owns the linear / norm / state-cache / ape state and the shared forward |
| 204 | prologue (kv/score split, save_partial_states launch). The |
| 205 | compress → norm → RoPE → store step is dispatched to a triton kernel |
| 206 | (``compress_norm_rope_store_triton``) by default, except for the NVIDIA |
| 207 | head_dim=512 path which uses the CuTeDSL compressor kernels for better |
| 208 | performance. |
| 209 | """ |
| 210 | |
| 211 | def __init__( |
| 212 | self, |
| 213 | vllm_config: VllmConfig, |
| 214 | compress_ratio: int, |
| 215 | hidden_size: int, |
| 216 | head_dim: int, |
| 217 | rotate: bool = False, |
| 218 | prefix: str = "", |
| 219 | k_cache_prefix="", |
| 220 | use_fp4_cache: bool = False, |
| 221 | ): |
| 222 | super().__init__() |
| 223 | self.compress_ratio = compress_ratio |
| 224 | self.hidden_size = hidden_size |
| 225 | self.head_dim = head_dim |
| 226 | self.rotate = rotate |
| 227 | self.prefix = prefix |
| 228 | self.k_cache_prefix = k_cache_prefix |
| 229 | self.use_fp4_cache = use_fp4_cache |
| 230 | |
| 231 | config = vllm_config.model_config.hf_config |
| 232 | self.rope_head_dim = config.qk_rope_head_dim |
| 233 | self.nope_head_dim = self.head_dim - self.rope_head_dim |
| 234 | self.rms_norm_eps = config.rms_norm_eps |
| 235 | self.device = current_platform.device_type |
| 236 | self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 237 | self.max_model_len = vllm_config.model_config.max_model_len |
| 238 | |
| 239 | self.overlap = compress_ratio == 4 |
| 240 | self.coff = 1 + self.overlap |
| 241 | |
| 242 | # The head=512 cr>=128 no-overlap deep gather uses the two-stage |
| 243 | # compressor, which needs an fp32 scratch [max_batched, 512] for |
| 244 | # the intermediate compressed_kv. |
| 245 | # Currently only tested on ROCm |
| 246 | self._use_two_stage_fused_compressor = ( |
| 247 | _prefer_two_stage_compressor() and head_dim == 512 and not self.overlap |
| 248 | ) |
| 249 | self.max_num_batched_tokens = ( |
| 250 | vllm_config.scheduler_config.max_num_batched_tokens |
| 251 | ) |
| 252 | self._compress_scratch: torch.Tensor | None = None |
| 253 | if self._use_two_stage_fused_compressor: |
| 254 | self._compress_scratch = torch.empty( |
| 255 | self.max_num_batched_tokens, |
| 256 | self.head_dim, |
| 257 | dtype=torch.float32, |
| 258 | device=self.device, |
| 259 | ) |
| 260 | |
| 261 | state_dtype = torch.float32 |
| 262 | self.ape = nn.Parameter( |
| 263 | torch.empty( |
| 264 | (compress_ratio, self.coff * self.head_dim), |
| 265 | dtype=state_dtype, |
| 266 | device=self.device, |
| 267 | ), |
| 268 | requires_grad=False, |
| 269 | ) |
| 270 | |
| 271 | self.fused_wkv_wgate = MergedColumnParallelLinear( |
| 272 | self.hidden_size, |
| 273 | [self.coff * self.head_dim, self.coff * self.head_dim], |
| 274 | bias=False, |
| 275 | return_bias=False, |
| 276 | quant_config=None, |
| 277 | disable_tp=True, |
| 278 | prefix=f"{prefix}.fused_wkv_wgate", |
| 279 | ) |
| 280 | self.norm = RMSNorm(self.head_dim, self.rms_norm_eps) |
| 281 | |
| 282 | self.state_cache = CompressorStateCache( |
| 283 | state_dim=2 * self.coff * self.head_dim, # kv_state + score_state |
| 284 | dtype=state_dtype, |
| 285 | compress_ratio=compress_ratio, |
| 286 | prefix=f"{prefix}.state_cache", |
| 287 | ) |
| 288 | |
| 289 | # Save reference to static_forward_context for forward-time KV cache lookup. |
| 290 | # get_current_vllm_config() is only available during __init__, not forward. |
| 291 | self._static_forward_context = ( |
| 292 | vllm_config.compilation_config.static_forward_context |
| 293 | ) |
| 294 | |
| 295 | if self.head_dim == 512: |
| 296 | assert not use_fp4_cache, ( |
| 297 | "MXFP4 cache is only supported for indexer (head=128)" |
| 298 | ) |
| 299 | self._quant_block = 64 |
| 300 | self._token_stride = self.nope_head_dim + self.rope_head_dim * 2 |
| 301 | self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad |
| 302 | elif self.head_dim == 128: |
| 303 | if use_fp4_cache: |
| 304 | self._quant_block = MXFP4_BLOCK_SIZE |
| 305 | self._token_stride = self.head_dim // 2 |
| 306 | self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE |
| 307 | else: |
| 308 | self._quant_block = 128 |
| 309 | self._token_stride = self.head_dim |
| 310 | self._scale_dim = 4 # single float32 scale |
| 311 | else: |
| 312 | raise ValueError( |
| 313 | f"Unsupported head_dim for fused quant+cache: {self.head_dim}" |
| 314 | ) |
| 315 | |
| 316 | if vllm_config.kernel_config.enable_jit_warmup: |
| 317 | _SAVE_PARTIAL_STATES_KERNEL.register_warmup( |
| 318 | head_dim=self.head_dim, |
| 319 | compress_ratio=self.compress_ratio, |
| 320 | ) |
| 321 | if current_platform.is_cuda() and self.head_dim == 512: |
| 322 | from vllm.models.deepseek_v4.nvidia.ops.sparse_attn_compress_cutedsl import ( # noqa: E501 |
| 323 | _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL, |
| 324 | _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL, |
| 325 | _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL, |
| 326 | _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL, |
| 327 | _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL, |
| 328 | ) |
| 329 | |
| 330 | store_full_kv = vllm_config.cache_config.cache_dtype != "fp8_ds_mla" |
| 331 | if self.compress_ratio == 4: |
| 332 | ( |
| 333 | _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL |
| 334 | if store_full_kv |
| 335 | else _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL |
| 336 | ).register_warmup() |
| 337 | else: |
| 338 | _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL.register_warmup() |
| 339 | if store_full_kv: |
| 340 | _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL.register_warmup() |
| 341 | else: |
| 342 | _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL.register_warmup( |
| 343 | vllm_config, |
| 344 | k_cache_prefix=self.k_cache_prefix, |
| 345 | compress_ratio=self.compress_ratio, |
| 346 | ) |
| 347 | else: |
| 348 | from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( # noqa: E501 |
| 349 | _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL, |
| 350 | ) |
| 351 | |
| 352 | _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL.register_warmup() |
| 353 | |
| 354 | def forward( |
| 355 | self, |
| 356 | # [num_tokens, 2 * self.coff * self.head_dim] |
| 357 | kv_score: torch.Tensor, |
| 358 | # [num_tokens] |
| 359 | positions: torch.Tensor, |
| 360 | rotary_emb, |
| 361 | ) -> None: |
| 362 | # Each of shape [num_tokens, coff * self.head_dim] |
| 363 | # input bf16, output are fp32 |
| 364 | kv, score = kv_score.split( |
| 365 | [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1 |
| 366 | ) |
| 367 | |
| 368 | # Get the metadata and handle dummy profiling run. |
| 369 | forward_context = get_forward_context() |
| 370 | attn_metadata = forward_context.attn_metadata |
| 371 | if not isinstance(attn_metadata, dict): |
| 372 | return |
| 373 | |
| 374 | state_metadata = cast( |
| 375 | CompressorMetadata, attn_metadata[self.state_cache.prefix] |
| 376 | ) |
| 377 | token_to_req_indices = state_metadata.token_to_req_indices |
| 378 | slot_mapping = state_metadata.slot_mapping |
| 379 | num_actual = slot_mapping.shape[0] |
| 380 | block_table = state_metadata.block_table |
| 381 | block_size = state_metadata.block_size |
| 382 | |
| 383 | # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim |
| 384 | state_cache = self.state_cache.kv_cache |
| 385 | # kv_state stored in first half, score_state stored in second half |
| 386 | state_width = state_cache.shape[-1] // 2 |
| 387 | pdl_kwargs = ( |
| 388 | {} |
| 389 | if current_platform.is_rocm() or current_platform.is_xpu() |
| 390 | else {"launch_pdl": False} |
| 391 | ) |
| 392 | |
| 393 | # Store the KV and score (with fused APE addition) in the state. |
| 394 | # NOTE: PDL is disabled — both this kernel and the compress kernels |
| 395 | # below depend on preceding kernel outputs (kv/score from the cublas |
| 396 | # GEMM; state_cache from this kernel) but neither emits/waits on PDL |
| 397 | # grid dependency primitives, so launch_pdl=True caused a |
| 398 | # read-after-write race and non-deterministic output. |
| 399 | _SAVE_PARTIAL_STATES_KERNEL( |
| 400 | kv=kv, |
| 401 | score=score, |
| 402 | ape=self.ape, |
| 403 | positions=positions, |
| 404 | state_cache=state_cache, |
| 405 | slot_mapping=slot_mapping, |
| 406 | block_size=block_size, |
| 407 | state_width=state_width, |
| 408 | compress_ratio=self.compress_ratio, |
| 409 | pdl_kwargs=pdl_kwargs, |
| 410 | ) |
| 411 | |
| 412 | # full graph cannot branch on per-step CPU metadata after capture |
| 413 | if ( |
| 414 | current_platform.is_cuda() |
| 415 | and self.head_dim == 512 |
| 416 | and self.compress_ratio == 128 |
| 417 | and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL |
| 418 | and state_metadata.c128_boundary is False |
| 419 | ): |
| 420 | return |
| 421 | |
| 422 | # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write. |
| 423 | # RoPE requirements (kernel applies forward GPT-J style rotation): |
| 424 | # - is_neox_style=False (interleaved pairs, NOT split-half) |
| 425 | # - cos_sin_cache layout: [max_pos, rope_head_dim] with first half cos, |
| 426 | # second half sin (per-pair, length rope_head_dim // 2 each) |
| 427 | # - applied to LAST rope_head_dim elements of head_dim |
| 428 | # - position used: (positions // compress_ratio) * compress_ratio |
| 429 | cos_sin_cache = rotary_emb.cos_sin_cache |
| 430 | k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) |
| 431 | k_cache_layer = self._static_forward_context[self.k_cache_prefix] |
| 432 | kv_cache = k_cache_layer.kv_cache |
| 433 | |
| 434 | # Plain-row V4 reads a contiguous bf16 / per-tensor fp8 cache row; the |
| 435 | # fp8_ds_mla path uses the UE8M0 paged uint8 layout. |
| 436 | store_full_kv = self.head_dim == 512 and kv_cache.dtype != torch.uint8 |
| 437 | store_full_fp8 = kv_cache.dtype == torch.float8_e4m3fn |
| 438 | fp8_scale = ( |
| 439 | getattr(k_cache_layer, "_flashinfer_fp8_kv_scale", None) |
| 440 | if store_full_fp8 |
| 441 | else None |
| 442 | ) |
| 443 | |
| 444 | # cutedsl (head=512) accepts the full-cache flags; triton (indexer/AMD) |
| 445 | # does not, so the two callables have different signatures. |
| 446 | compress_norm_rope_store_fn: Any |
| 447 | if current_platform.is_cuda() and self.head_dim == 512: |
| 448 | from .nvidia.ops.sparse_attn_compress_cutedsl import ( |
| 449 | _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL, |
| 450 | ) |
| 451 | |
| 452 | # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla |
| 453 | # layout and the plain full-cache layout. The full-cache flags |
| 454 | # are consumed only here. |
| 455 | compress_norm_rope_store_fn = _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL |
| 456 | extra_kwargs: dict[str, Any] = dict( |
| 457 | store_full_kv=store_full_kv, |
| 458 | store_full_fp8=store_full_fp8, |
| 459 | fp8_scale=fp8_scale, |
| 460 | ) |
| 461 | elif self._use_two_stage_fused_compressor: |
| 462 | # head=512 cr>=128 (no overlap): two-pass split compressor on the |
| 463 | # prefill suffix, single-pass on the decode prefix. |
| 464 | assert state_metadata.num_decode_tokens is not None |
| 465 | compress_norm_rope_store_fn = compress_norm_rope_store_two_stage_triton |
| 466 | extra_kwargs = { |
| 467 | "num_decode_tokens": state_metadata.num_decode_tokens, |
| 468 | "compress_scratch": self._compress_scratch, |
| 469 | } |
| 470 | else: |
| 471 | # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.). |
| 472 | compress_norm_rope_store_fn = compress_norm_rope_store_triton |
| 473 | extra_kwargs = {} |
| 474 | |
| 475 | compress_norm_rope_store_fn( |
| 476 | state_cache=state_cache, |
| 477 | num_actual=num_actual, |
| 478 | token_to_req_indices=token_to_req_indices, |
| 479 | positions=positions, |
| 480 | slot_mapping=slot_mapping, |
| 481 | block_table=block_table, |
| 482 | block_size=block_size, |
| 483 | state_width=state_width, |
| 484 | cos_sin_cache=cos_sin_cache, |
| 485 | kv_cache=kv_cache, |
| 486 | k_cache_metadata=k_cache_metadata, |
| 487 | pdl_kwargs=pdl_kwargs, |
| 488 | head_dim=self.head_dim, |
| 489 | rope_head_dim=self.rope_head_dim, |
| 490 | compress_ratio=self.compress_ratio, |
| 491 | overlap=self.overlap, |
| 492 | use_fp4_cache=self.use_fp4_cache, |
| 493 | rms_norm_weight=self.norm.weight, |
| 494 | rms_norm_eps=self.rms_norm_eps, |
| 495 | quant_block=self._quant_block, |
| 496 | token_stride=self._token_stride, |
| 497 | scale_dim=self._scale_dim, |
| 498 | **extra_kwargs, |
| 499 | ) |