diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py index 38a34bc923a5..2717c4f4e470 100644 --- a/tests/models/kimi_k3/test_attn_res.py +++ b/tests/models/kimi_k3/test_attn_res.py @@ -194,6 +194,135 @@ def test_attn_res_without_output_norm(): torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) +def test_attn_res_reuses_delta_for_output(): + prefix = torch.randn(17, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.randn_like(prefix) + blocks = torch.randn( + 17, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = torch.ones_like(norm_weight) + expected, expected_prefix = _reference( + prefix.clone(), + delta.clone(), + blocks, + norm_weight, + qk_weight, + output_norm_weight, + MAX_BLOCKS, + ) + output_pointer = delta.data_ptr() + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + MAX_BLOCKS, + -1, + EPS, + EPS, + output=delta, + ) + + assert actual.data_ptr() == output_pointer + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) + + +def test_attn_res_uses_an_inactive_workspace_block_for_output(): + num_tokens = 17 + prefix = torch.randn(num_tokens, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + block_storage = torch.randn( + MAX_BLOCKS, + num_tokens, + HIDDEN_SIZE, + device="cuda", + dtype=torch.bfloat16, + ) + blocks = block_storage.permute(1, 0, 2) + output = blocks[:, -1] + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = torch.ones_like(norm_weight) + expected, expected_prefix = _reference( + prefix.clone(), + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + 0, + ) + output_pointer = output.data_ptr() + + actual = attn_res( + prefix, + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + 0, + 0, + EPS, + EPS, + output=output, + ) + + assert actual.data_ptr() == output_pointer + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) + torch.testing.assert_close(blocks[:, 0], expected_prefix, atol=0, rtol=0) + + +def test_attn_res_reused_output_supports_cuda_graph_replay(): + prefix = torch.randn(17, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.randn_like(prefix) + blocks = torch.randn( + 17, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = torch.ones_like(norm_weight) + original_prefix = prefix.clone() + original_delta = delta.clone() + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + MAX_BLOCKS, + -1, + EPS, + EPS, + output=delta, + ) + + prefix.copy_(original_prefix) + delta.copy_(original_delta) + graph.replay() + first = actual.clone() + prefix.copy_(original_prefix) + delta.copy_(original_delta) + graph.replay() + torch.testing.assert_close(actual, first, atol=0, rtol=0) + + @pytest.mark.parametrize("num_tokens", [0, 1, 17]) def test_fused_mtp_input(num_tokens: int): positions = torch.arange(num_tokens, device="cuda") diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 61a24a83e041..e9ae509291c3 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -9,6 +9,7 @@ from vllm.model_executor.models.interfaces import supports_eagle3 from vllm.models.kimi_k3.nvidia import model as kimi_model from vllm.models.kimi_k3.nvidia.model import ( + KimiDecoderLayer, KimiK3ForConditionalGeneration, KimiLinearModel, ) @@ -16,8 +17,11 @@ def _make_kimi_linear_model() -> KimiLinearModel: model = object.__new__(KimiLinearModel) + torch.nn.Module.__init__(model) + model.register_buffer("_attn_res_workspace", None, persistent=False) object.__setattr__(model, "aux_hidden_state_layers", (2,)) object.__setattr__(model, "use_sequence_parallel", False) + object.__setattr__(model, "reuse_attn_res_output", True) return model @@ -128,3 +132,176 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) assert final_attn_res.call_args.args[2] is block_residual + + +def test_kimi_attn_res_workspace_is_reused_and_sliced(): + model = _make_kimi_linear_model() + object.__setattr__(model, "num_attn_res_blocks", 3) + + largest = model._get_attn_res_workspace(torch.empty(8, 4)) + smaller = model._get_attn_res_workspace(torch.empty(2, 4)) + + assert largest.shape == (8, 3, 4) + assert smaller.shape == (2, 3, 4) + assert smaller.untyped_storage().data_ptr() == largest.untyped_storage().data_ptr() + + grown = model._get_attn_res_workspace(torch.empty(16, 4)) + assert grown.shape == (16, 3, 4) + assert grown.untyped_storage().data_ptr() != largest.untyped_storage().data_ptr() + + +def test_kimi_retained_workspace_releases_cuda_cache(monkeypatch): + events = [] + device = torch.device("cuda", 3) + monkeypatch.setattr( + torch.accelerator, + "synchronize", + lambda actual_device: events.append(("synchronize", actual_device)), + ) + monkeypatch.setattr( + torch.accelerator, + "empty_cache", + lambda: events.append(("empty_cache", None)), + ) + + kimi_model._release_cuda_cache_before_retained_allocation(device) + + assert events == [("synchronize", device), ("empty_cache", None)] + + +def test_kimi_attn_res_workspace_can_be_reserved_before_prefill(monkeypatch): + model = _make_kimi_linear_model() + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 3) + object.__setattr__(model, "_max_num_batched_tokens", 16) + object.__setattr__(model, "_model_dtype", torch.bfloat16) + object.__setattr__(model, "config", SimpleNamespace(hidden_size=4)) + model.register_parameter( + "reference_weight", + torch.nn.Parameter(torch.empty(1, dtype=torch.bfloat16)), + ) + + events = [] + original_empty = torch.empty + + def record_allocator_release(device): + events.append(("release", device.type)) + + def record_workspace_allocation(*args, **kwargs): + device = kwargs.get("device") + if device is not None: + events.append(("allocate", device.type)) + return original_empty(*args, **kwargs) + + monkeypatch.setattr( + kimi_model, + "_release_cuda_cache_before_retained_allocation", + record_allocator_release, + ) + monkeypatch.setattr(torch, "empty", record_workspace_allocation) + model.reserve_attn_res_workspace() + + assert events == [("release", "cpu"), ("allocate", "cpu")] + workspace = model._attn_res_workspace + assert workspace is not None + assert workspace.shape == (16, 3, 4) + assert workspace.dtype == torch.bfloat16 + pointer = workspace.untyped_storage().data_ptr() + sliced = model._get_attn_res_workspace(torch.empty(8, 4, dtype=torch.bfloat16)) + assert sliced.untyped_storage().data_ptr() == pointer + + +def _make_attn_res_decoder_layer(*, block_write: bool): + layer = object.__new__(KimiDecoderLayer) + torch.nn.Module.__init__(layer) + object.__setattr__(layer, "use_attn_res", True) + object.__setattr__(layer, "reuse_attn_res_output", True) + object.__setattr__(layer, "is_block_write_layer", block_write) + object.__setattr__(layer, "block_write_idx", 0) + object.__setattr__(layer, "prev_valid_blocks", 0) + object.__setattr__( + layer, + "self_attention_res_norm", + SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5), + ) + object.__setattr__( + layer, + "self_attention_res_proj", + SimpleNamespace(weight=torch.ones(1, 4)), + ) + object.__setattr__( + layer, + "input_layernorm", + SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5), + ) + object.__setattr__( + layer, + "mlp_res_norm", + SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5), + ) + object.__setattr__( + layer, + "mlp_res_proj", + SimpleNamespace(weight=torch.ones(1, 4)), + ) + object.__setattr__( + layer, + "post_attention_layernorm", + SimpleNamespace(weight=torch.ones(4), variance_epsilon=1e-5), + ) + return layer + + +def test_kimi_pre_attn_norm_uses_consumed_or_workspace_output(monkeypatch): + layer = _make_attn_res_decoder_layer(block_write=True) + prefix = torch.randn(2, 4) + delta = torch.randn(2, 4) + blocks = torch.randn(2, 3, 4) + scratch = blocks[:, -1] + calls = [] + + def record_attn_res(*args, **kwargs): + calls.append(kwargs["output"]) + output = kwargs["output"] + return args[0].clone() if output is None else output + + monkeypatch.setattr(kimi_model, "attn_res", record_attn_res) + + first, _, _ = layer._pre_attn_norm(None, blocks, prefix, scratch) + later, _, _ = layer._pre_attn_norm(delta, blocks, prefix, scratch) + + assert first is scratch + assert later is delta + assert calls == [scratch, delta] + + object.__setattr__(layer, "reuse_attn_res_output", False) + layer._pre_attn_norm(delta, blocks, prefix, scratch) + assert calls[-1] is None + + +def test_kimi_post_attn_norm_reuses_dead_input(monkeypatch): + prefix = torch.randn(2, 4) + attention_output = torch.randn(2, 4) + blocks = torch.randn(2, 3, 4) + outputs = [] + + def record_attn_res(*args, **kwargs): + outputs.append(kwargs["output"]) + return kwargs["output"] + + monkeypatch.setattr(kimi_model, "attn_res", record_attn_res) + + block_layer = _make_attn_res_decoder_layer(block_write=True) + block_hidden, block_prefix, _ = block_layer._post_attn_norm( + attention_output, blocks, prefix + ) + regular_layer = _make_attn_res_decoder_layer(block_write=False) + regular_hidden, regular_prefix, _ = regular_layer._post_attn_norm( + attention_output, blocks, prefix + ) + + assert block_hidden is prefix + assert block_prefix is attention_output + assert regular_hidden is attention_output + assert regular_prefix is prefix + assert outputs == [prefix, attention_output] diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 752caab89ae6..8691f25f87b5 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -142,6 +142,19 @@ _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256 +def _release_cuda_cache_before_retained_allocation(device: torch.device) -> None: + """Give retained post-load storage a dedicated allocator segment. + + A persistent allocation must not pin the unused part of a large cached + segment left by quantization repacking. KV cache allocation follows this + hook and needs every otherwise-inactive segment to be releasable. + """ + if device.type != "cuda": + return + torch.accelerator.synchronize(device) + torch.accelerator.empty_cache() + + def _uses_native_b12x_mxfp4_intermediate_size( vllm_config: VllmConfig, ) -> bool: @@ -1226,6 +1239,9 @@ def __init__( attn_res_block_size = config.attn_res_block_size self.use_attn_res = attn_res_block_size is not None + self.reuse_attn_res_output = ( + self.use_attn_res and current_platform.is_device_capability_family(120) + ) if self.use_attn_res: assert attn_res_block_size is not None self.attn_res_block_size = attn_res_block_size @@ -1274,6 +1290,7 @@ def _pre_attn_norm( hidden_states: torch.Tensor | None, residual: torch.Tensor | None, prefix_sum: torch.Tensor | None, + attn_res_scratch: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: if not self.use_attn_res: assert hidden_states is not None @@ -1297,6 +1314,9 @@ def _pre_attn_norm( block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), eps=self.self_attention_res_norm.variance_epsilon, output_norm_eps=self.input_layernorm.variance_epsilon, + output=(hidden_states if hidden_states is not None else attn_res_scratch) + if self.reuse_attn_res_output + else None, ) return hidden_states, prefix_sum, residual @@ -1314,10 +1334,12 @@ def _post_attn_norm( assert prefix_sum is not None if self.is_block_write_layer: + output = prefix_sum if self.reuse_attn_res_output else None prefix_sum = hidden_states prefix_delta = None else: prefix_delta = hidden_states + output = prefix_delta if self.reuse_attn_res_output else None mlp_valid_blocks = self.prev_valid_blocks + self.is_block_write_layer hidden_states = attn_res( prefix_sum, @@ -1330,6 +1352,7 @@ def _post_attn_norm( block_write_idx=-1, eps=self.mlp_res_norm.variance_epsilon, output_norm_eps=self.post_attention_layernorm.variance_epsilon, + output=output, ) return hidden_states, prefix_sum, residual @@ -1339,10 +1362,11 @@ def forward( hidden_states: torch.Tensor | None, residual: torch.Tensor | None, prefix_sum: torch.Tensor | None = None, + attn_res_scratch: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: hidden_states, prefix_sum, residual = self._pre_attn_norm( - hidden_states, residual, prefix_sum + hidden_states, residual, prefix_sum, attn_res_scratch ) assert hidden_states is not None @@ -1384,6 +1408,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.attn_res_block_size: int | None = config.attn_res_block_size self.use_attn_res = self.attn_res_block_size is not None + self.reuse_attn_res_output = ( + self.use_attn_res and current_platform.is_device_capability_family(120) + ) parallel_config = vllm_config.parallel_config use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" self.use_sequence_parallel = ( @@ -1427,6 +1454,12 @@ def get_layer(prefix: str): if self.attn_res_block_size is not None else 0 ) + self._max_num_batched_tokens = int( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self._model_dtype = vllm_config.model_config.dtype + self._attn_res_workspace: torch.Tensor | None + self.register_buffer("_attn_res_workspace", None, persistent=False) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -1478,6 +1511,71 @@ def make_empty_intermediate_tensors( def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + def _get_attn_res_workspace(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return retained AttnRes storage for the active token rows.""" + shape = ( + hidden_states.size(0), + self.num_attn_res_blocks, + hidden_states.size(1), + ) + workspace = self._attn_res_workspace + if ( + workspace is None + or workspace.device != hidden_states.device + or workspace.dtype != hidden_states.dtype + or workspace.size(0) < shape[0] + or workspace.size(1) != shape[1] + or workspace.size(2) != shape[2] + ): + # Physical block-major storage keeps every token-major block view + # contiguous for AttnRes kernels and later buffer reuse. + workspace = hidden_states.new_empty( + shape[1], + shape[0], + shape[2], + ).permute(1, 0, 2) + self._attn_res_workspace = workspace + return workspace[: shape[0]] + + def reserve_attn_res_workspace(self) -> None: + """Reserve maximum-size AttnRes storage before KV cache allocation. + + Chunked prefill reuses one allocation for every scheduler chunk. Early + reservation prevents model-load and CUDA-graph allocations from + fragmenting the contiguous block required by a maximum-size chunk. + """ + if not self.use_attn_res or self.num_attn_res_blocks == 0: + return + shape = ( + self._max_num_batched_tokens, + self.num_attn_res_blocks, + self.config.hidden_size, + ) + workspace = self._attn_res_workspace + parameter = next(self.parameters()) + if ( + workspace is None + or workspace.device != parameter.device + or workspace.dtype != self._model_dtype + or tuple(workspace.shape) != shape + ): + _release_cuda_cache_before_retained_allocation(parameter.device) + self._attn_res_workspace = torch.empty( + shape[1], + shape[0], + shape[2], + dtype=self._model_dtype, + device=parameter.device, + ).permute(1, 0, 2) + logger.info_once( + "Kimi-K3 retained %.2f MiB/rank for the %d-token AttnRes " + "prefill workspace.", + self._attn_res_workspace.numel() + * self._attn_res_workspace.element_size() + / (1024**2), + self._max_num_batched_tokens, + ) + def forward( self, input_ids: torch.Tensor | None, @@ -1518,16 +1616,21 @@ def forward( prefix_sum = None if self.use_attn_res: - block_residual = hidden_states.new_empty( - hidden_states.size(0), - self.num_attn_res_blocks, - hidden_states.size(1), - ) + block_residual = self._get_attn_res_workspace(hidden_states) if residual is not None: block_residual[:, : residual.size(1), :].copy_(residual) prefix_sum = hidden_states hidden_states = None residual = block_residual + # The final block is not an AttnRes input until the final block + # boundary, so its storage can hold the first normalized output. + attn_res_scratch = ( + block_residual[:, -1] + if self.reuse_attn_res_output and self.num_attn_res_blocks > 1 + else None + ) + else: + attn_res_scratch = None for layer_idx, layer in enumerate( self.layers[self.start_layer : self.end_layer], @@ -1538,6 +1641,7 @@ def forward( hidden_states=hidden_states, prefix_sum=prefix_sum, residual=residual, + attn_res_scratch=attn_res_scratch, ) if (layer_idx + 1) in self.aux_hidden_state_layers: if self.use_attn_res: @@ -1574,6 +1678,7 @@ def forward( block_write_idx=-1, eps=self.output_attn_res_norm.variance_epsilon, output_norm_eps=0.0, + output=(hidden_states if self.reuse_attn_res_output else None), ) else: hidden_states = hidden_states + residual @@ -1910,6 +2015,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # hook for any AttentionLayerBase, so no manual trigger is needed here. return loaded + def process_weights_after_loading(self) -> None: + self.model.reserve_attn_res_workspace() + def get_spec_layer_idx_from_weight_name( config: KimiLinearConfig, weight_name: str @@ -2374,3 +2482,6 @@ def get_mamba_state_copy_func(cls): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: + self.language_model.process_weights_after_loading() diff --git a/vllm/models/kimi_k3/nvidia/ops/attn_res.py b/vllm/models/kimi_k3/nvidia/ops/attn_res.py index 06f73ce8f7bd..f74e12f8d6ac 100644 --- a/vllm/models/kimi_k3/nvidia/ops/attn_res.py +++ b/vllm/models/kimi_k3/nvidia/ops/attn_res.py @@ -178,7 +178,15 @@ def attn_res( block_write_idx: int, eps: float, output_norm_eps: float, + output: torch.Tensor | None = None, ) -> torch.Tensor: + """Mix AttnRes sources into caller-owned or newly allocated storage. + + ``output`` may alias ``delta`` because each token row is read completely + before that row is written. If it shares storage with ``blocks``, the + caller must select a block outside both the active source range and + ``block_write_idx``. + """ num_tokens, hidden_size = prefix.shape assert prefix.stride(-1) == 1 assert delta is None or delta.stride(-1) == 1 @@ -186,6 +194,11 @@ def attn_res( assert norm_weight.stride(-1) == 1 assert qk_weight.stride(-1) == 1 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1 + if output is not None: + assert output.shape == prefix.shape + assert output.device == prefix.device + assert output.dtype == prefix.dtype + assert output.stride(-1) == 1 # The in-tree NVIDIA kernel covers the common fused-add + output-norm path; # Triton handles block boundaries and final pre-norm output. The native op # is only compiled for SM100 under CUDA >= 13, so a device check alone is @@ -198,6 +211,9 @@ def attn_res( and block_write_idx < 0 and current_platform.is_device_capability_family(100) and hasattr(torch.ops._C, "kimi_k3_attn_res") + # The SM100 operation does not declare caller-output aliasing. An + # explicit output therefore selects the Triton implementation. + and output is None ): return ops.kimi_k3_attn_res( prefix, @@ -210,7 +226,8 @@ def attn_res( eps, output_norm_eps, ) - output = prefix.new_empty(prefix.shape) + if output is None: + output = prefix.new_empty(prefix.shape) # Tuned on GB300: source tiling helps decode, while one-source tiles scale # better for prefill. # Keep get_attn_res_triton_warmup_profiles in sync with these fallbacks.